mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-09-07 08:27:12 +00:00
Merge 2b6fcdbb80 into c90d195c2f
This commit is contained in:
commit
6d9d7e9925
18 changed files with 450 additions and 12 deletions
|
|
@ -10297,6 +10297,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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1502,6 +1502,7 @@ mod runs {
|
|||
},
|
||||
billing: BilledTokenCounts::default(),
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 0,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
@ -1579,6 +1580,7 @@ mod runs {
|
|||
},
|
||||
billing: BilledTokenCounts::default(),
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 0,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -5137,6 +5137,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,
|
||||
|
|
@ -17174,6 +17175,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(
|
||||
|
|
|
|||
|
|
@ -83,8 +83,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)]
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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>,
|
||||
|
|
@ -1118,6 +1144,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,
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
@ -5112,6 +5122,7 @@ mod tests {
|
|||
model: billed_usage().model().clone(),
|
||||
billing,
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 0,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
@ -5184,6 +5195,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();
|
||||
|
|
|
|||
|
|
@ -660,17 +660,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,
|
||||
|
|
@ -1469,8 +1473,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};
|
||||
|
|
@ -2533,6 +2537,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,
|
||||
|
|
@ -2568,6 +2573,7 @@ mod tests {
|
|||
},
|
||||
cost_usd: None,
|
||||
cost_source: None,
|
||||
session_spend: SessionSpend::default(),
|
||||
tool_call_count: 0,
|
||||
context_window: None,
|
||||
reasoning: None,
|
||||
|
|
@ -2587,6 +2593,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 {
|
||||
|
|
@ -2606,6 +2654,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,
|
||||
|
|
@ -2657,6 +2706,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,
|
||||
|
|
@ -2692,6 +2742,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(
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
||||
|
|
@ -132,6 +132,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(
|
||||
|
|
|
|||
|
|
@ -381,6 +381,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(
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -565,6 +572,7 @@ mod tests {
|
|||
model: sample_model_ref(),
|
||||
billing: BilledTokenCounts::default(),
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 1,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
@ -594,6 +602,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()),
|
||||
|
|
|
|||
|
|
@ -2192,6 +2192,7 @@ mod tests {
|
|||
},
|
||||
billing: BilledTokenCounts::default(),
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 0,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
@ -2223,6 +2224,7 @@ mod tests {
|
|||
},
|
||||
billing: BilledTokenCounts::default(),
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 0,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
@ -2251,6 +2253,7 @@ mod tests {
|
|||
},
|
||||
billing: BilledTokenCounts::default(),
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 1,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
@ -2308,6 +2311,7 @@ mod tests {
|
|||
},
|
||||
billing: BilledTokenCounts::default(),
|
||||
cost_source: None,
|
||||
session_billing: None,
|
||||
tool_call_count: 0,
|
||||
visit: 1,
|
||||
message: None,
|
||||
|
|
|
|||
|
|
@ -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() {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue