mirror of
https://github.com/block/buzz.git
synced 2026-10-27 04:30:52 +00:00
feat(acp): report standard adapter usage (#4950)
## Why Claude Code and Codex expose standard ACP prompt-response usage, but Buzz only consumed Goose’s private cumulative usage notification. Their token use and Claude’s cumulative cost were therefore absent from NIP-AM metrics. ## What - Read per-turn `session/prompt` response usage for known Claude and Codex adapters - Publish Claude’s raw cumulative cost separately from per-turn tokens without changing the NIP-AM schema - Keep Goose usage exclusive and cover Claude/Codex wire serialization ## Risk Assessment Low-to-medium: changes best-effort observability only and does not affect prompt execution. The adapter-specific mappings preserve source semantics and omit unavailable fields. ## References - Validated with `cargo fmt --check`, `cargo test -p buzz-acp --no-run`, and full `cargo test -p buzz-acp` (678 passed at `652e373a` before merge-trailer amendment). Generated with Codex --------- Signed-off-by: Atish Patel <atish@squareup.com> Signed-off-by: Brainy Bumble <0ed7657b57c0e8a9f5288390dd6c8d5d0a3a06abe9b01b9006814f52077d6cdf@buzz.block.builderlab.xyz> Co-authored-by: WorkerBeeGPT <c5c455215c2506cb8ba776518cec804af62d3a0526e32d496a22072e395042b9@buzz.block.builderlab.xyz> Co-authored-by: Brainy Bumble <0ed7657b57c0e8a9f5288390dd6c8d5d0a3a06abe9b01b9006814f52077d6cdf@buzz.block.builderlab.xyz>
This commit is contained in:
@@ -14,7 +14,9 @@ use tokio::process::{Child, ChildStdin, ChildStdout};
|
||||
use tokio_util::codec::{FramedRead, LinesCodec, LinesCodecError};
|
||||
|
||||
use crate::observer::{ObserverContext, ObserverHandle};
|
||||
use crate::usage::{TurnUsage, UsageTracker};
|
||||
use crate::usage::{
|
||||
PromptResponseUsage, StandardAdapterKind, StandardUsageTracker, TurnUsage, UsageTracker,
|
||||
};
|
||||
|
||||
/// Maximum allowed size of a single NDJSON line from the agent's stdout.
|
||||
/// Lines exceeding this limit are rejected to prevent OOM from rogue agents.
|
||||
@@ -206,11 +208,12 @@ pub struct AcpClient {
|
||||
/// outside of a goose-native turn — the read loop's steer arm is
|
||||
/// disabled in that case.
|
||||
steer_rx: Option<tokio::sync::mpsc::Receiver<crate::pool::SteerRequest>>,
|
||||
/// Usage tracker — accumulates cumulative token counts from
|
||||
/// `_goose/unstable/session/update` notifications and computes per-turn
|
||||
/// deltas. Both goose and buzz-agent emit this notification; goose gates
|
||||
/// on client capability advertisement, buzz-agent emits unconditionally.
|
||||
/// Usage tracker for goose/buzz-agent's cumulative notification format.
|
||||
goose_usage: UsageTracker,
|
||||
/// Per-turn prompt-response usage and Claude's optional cumulative cost.
|
||||
standard_usage: StandardUsageTracker,
|
||||
/// Known adapter identity for prompt-response usage mapping.
|
||||
standard_adapter: Option<StandardAdapterKind>,
|
||||
}
|
||||
|
||||
/// Recursively merge `overlay` into `base`, with `overlay` winning on scalar/shape
|
||||
@@ -523,6 +526,14 @@ impl AcpClient {
|
||||
// console-subsystem child process spawned from a GUI/non-console parent.
|
||||
configure_no_window(&mut cmd);
|
||||
|
||||
let standard_adapter =
|
||||
match crate::config::normalize_agent_command_identity(command).as_str() {
|
||||
"claude-agent-acp" | "claude-code-acp" | "claude-code" | "claudecode" => {
|
||||
Some(StandardAdapterKind::Claude)
|
||||
}
|
||||
"codex" | "codex-acp" => Some(StandardAdapterKind::Codex),
|
||||
_ => None,
|
||||
};
|
||||
let mut child = cmd.spawn()?;
|
||||
|
||||
let stdin = child
|
||||
@@ -550,6 +561,8 @@ impl AcpClient {
|
||||
steering_supported: false,
|
||||
steer_rx: None,
|
||||
goose_usage: UsageTracker::default(),
|
||||
standard_usage: StandardUsageTracker::default(),
|
||||
standard_adapter,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -776,6 +789,7 @@ impl AcpClient {
|
||||
// prompt so that any setup notifications recorded earlier are not
|
||||
// misattributed to this turn.
|
||||
self.goose_usage.begin_turn(session_id);
|
||||
self.standard_usage.begin_turn(session_id);
|
||||
|
||||
self.last_prompt_id = Some(self.next_id);
|
||||
let id = self.next_id;
|
||||
@@ -821,7 +835,7 @@ impl AcpClient {
|
||||
self.current_hard_deadline = None;
|
||||
}
|
||||
}
|
||||
self.parse_stop_reason(&result?)
|
||||
self.parse_prompt_response(session_id, &result?)
|
||||
}
|
||||
|
||||
/// Send a `session/cancel` **notification** (no `id` field, no response expected).
|
||||
@@ -867,18 +881,13 @@ impl AcpClient {
|
||||
self.steering_supported
|
||||
}
|
||||
|
||||
/// Consume and return the per-turn usage record computed from the most
|
||||
/// recent `_goose/unstable/session/update` notification.
|
||||
///
|
||||
/// Returns `None` if no usage update arrived since the last call (i.e.
|
||||
/// the harness did not emit one for this turn, or this is not a goose
|
||||
/// agent). Must be called at most once per turn; subsequent calls return
|
||||
/// `None` until the next `usage_update` notification is recorded.
|
||||
///
|
||||
/// Intended for consumption by `publish_agent_turn_metric` in `pool.rs` to
|
||||
/// publish a kind 44200 NIP-AM event.
|
||||
/// Consume per-turn usage for NIP-AM publishing. Goose/buzz-agent is an
|
||||
/// exclusive cumulative path; standard ACP prompt usage is used only when
|
||||
/// goose emitted nothing for this turn.
|
||||
pub fn take_turn_usage(&mut self) -> Option<TurnUsage> {
|
||||
self.goose_usage.take()
|
||||
let goose_usage = self.goose_usage.take();
|
||||
let standard_usage = self.standard_usage.take();
|
||||
goose_usage.or(standard_usage)
|
||||
}
|
||||
|
||||
/// Notify the usage tracker that buzz-acp just spawned a new session.
|
||||
@@ -889,6 +898,7 @@ impl AcpClient {
|
||||
/// never when attaching to a pre-existing session.
|
||||
pub(crate) fn notify_session_spawned(&mut self, session_id: &str) {
|
||||
self.goose_usage.seed_zero_baseline(session_id);
|
||||
self.standard_usage.seed_zero_baseline(session_id);
|
||||
}
|
||||
|
||||
/// Install a per-turn steer request channel for goose-native
|
||||
@@ -1048,7 +1058,7 @@ impl AcpClient {
|
||||
remaining,
|
||||
)
|
||||
.await?;
|
||||
self.parse_stop_reason(&result)
|
||||
self.parse_prompt_response(session_id, &result)
|
||||
}
|
||||
|
||||
/// Serialize `value` as a single NDJSON line and flush to the agent's stdin.
|
||||
@@ -1831,6 +1841,10 @@ impl AcpClient {
|
||||
}
|
||||
false
|
||||
}
|
||||
"usage_update" => {
|
||||
self.handle_standard_usage_update(msg);
|
||||
false
|
||||
}
|
||||
"keepalive" => false,
|
||||
other => {
|
||||
tracing::debug!(target: "acp::update", "session/update: {other}");
|
||||
@@ -1839,6 +1853,30 @@ impl AcpClient {
|
||||
}
|
||||
}
|
||||
|
||||
/// Record the standard ACP cumulative cost notification when emitted by
|
||||
/// Claude. Unlike Goose's payload, `used`/`size` are context occupancy and
|
||||
/// are intentionally not mapped to token accounting.
|
||||
fn handle_standard_usage_update(&mut self, msg: &serde_json::Value) {
|
||||
if self.standard_adapter != Some(StandardAdapterKind::Claude) {
|
||||
return;
|
||||
}
|
||||
let session_id = match msg
|
||||
.pointer("/params/sessionId")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
{
|
||||
Some(session_id) => session_id,
|
||||
None => return,
|
||||
};
|
||||
let cost = match msg
|
||||
.pointer("/params/update/cost/amount")
|
||||
.and_then(serde_json::Value::as_f64)
|
||||
{
|
||||
Some(cost) => cost,
|
||||
None => return,
|
||||
};
|
||||
self.standard_usage.record_cost(session_id, cost);
|
||||
}
|
||||
|
||||
/// Parse a `_goose/unstable/session/update` notification and record the
|
||||
/// usage snapshot in the per-session tracker.
|
||||
///
|
||||
@@ -1970,6 +2008,28 @@ impl AcpClient {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Parse a completed prompt response and retain its optional per-turn usage.
|
||||
fn parse_prompt_response(
|
||||
&mut self,
|
||||
session_id: &str,
|
||||
result: &serde_json::Value,
|
||||
) -> Result<StopReason, AcpError> {
|
||||
let stop_reason = self.parse_stop_reason(result)?;
|
||||
if let Some(adapter) = self.standard_adapter {
|
||||
match serde_json::from_value::<PromptResponseUsage>(result["usage"].clone()) {
|
||||
Ok(usage) => self
|
||||
.standard_usage
|
||||
.record_prompt_usage(session_id, usage, adapter),
|
||||
Err(_) if result.get("usage").is_some() => tracing::debug!(
|
||||
target: "acp::usage",
|
||||
"session/prompt response contained malformed standard usage"
|
||||
),
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
Ok(stop_reason)
|
||||
}
|
||||
|
||||
/// Parse `stopReason` from a `session/prompt` result value.
|
||||
fn parse_stop_reason(&self, result: &serde_json::Value) -> Result<StopReason, AcpError> {
|
||||
let raw = result["stopReason"].as_str().ok_or_else(|| {
|
||||
@@ -2899,6 +2959,30 @@ mod tests {
|
||||
.expect("failed to spawn test script")
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn spawn_named_script(name: &str, script: &str) -> (AcpClient, std::path::PathBuf) {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
let dir = std::env::temp_dir().join(format!(
|
||||
"buzz-acp-{name}-{}-{}",
|
||||
std::process::id(),
|
||||
uuid::Uuid::new_v4()
|
||||
));
|
||||
std::fs::create_dir_all(&dir).expect("create temp adapter dir");
|
||||
let path = dir.join(name);
|
||||
std::fs::write(&path, format!("#!/usr/bin/env bash\n{script}\n"))
|
||||
.expect("write fake adapter");
|
||||
let mut permissions = std::fs::metadata(&path)
|
||||
.expect("adapter metadata")
|
||||
.permissions();
|
||||
permissions.set_mode(0o755);
|
||||
std::fs::set_permissions(&path, permissions).expect("chmod fake adapter");
|
||||
let client = AcpClient::spawn(path.to_str().expect("utf8 path"), &[], &[], false)
|
||||
.await
|
||||
.expect("spawn named fake adapter");
|
||||
(client, dir)
|
||||
}
|
||||
|
||||
/// Spawn a probe script whose file name carries a runtime identity (e.g.
|
||||
/// `hermes-acp`) and return the value of `var` as the child observed it.
|
||||
/// `<unset>` means the child did not receive the var.
|
||||
@@ -4275,6 +4359,254 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Standard ACP prompt-response usage ─────────────────────────────────
|
||||
|
||||
fn prompt_response_usage(
|
||||
input: u64,
|
||||
output: u64,
|
||||
total: u64,
|
||||
cached_read: Option<u64>,
|
||||
cached_write: Option<u64>,
|
||||
) -> serde_json::Value {
|
||||
let mut usage = serde_json::json!({
|
||||
"inputTokens": input,
|
||||
"outputTokens": output,
|
||||
"totalTokens": total,
|
||||
});
|
||||
if let Some(cached_read) = cached_read {
|
||||
usage["cachedReadTokens"] = serde_json::json!(cached_read);
|
||||
}
|
||||
if let Some(cached_write) = cached_write {
|
||||
usage["cachedWriteTokens"] = serde_json::json!(cached_write);
|
||||
}
|
||||
serde_json::json!({"stopReason": "end_turn", "usage": usage})
|
||||
}
|
||||
|
||||
fn standard_cost_update(session_id: &str, cost: f64) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "session/update",
|
||||
"params": {
|
||||
"sessionId": session_id,
|
||||
"update": {
|
||||
"sessionUpdate": "usage_update",
|
||||
"cost": {"amount": cost, "currency": "USD"}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn claude_prompt_response_usage_merges_with_cumulative_cost() {
|
||||
let mut client = spawn_inert_client().await;
|
||||
client.standard_adapter = Some(StandardAdapterKind::Claude);
|
||||
client.notify_session_spawned("claude-session");
|
||||
client.standard_usage.begin_turn("claude-session");
|
||||
client.handle_session_update(&standard_cost_update("claude-session", 0.042));
|
||||
assert_eq!(
|
||||
client
|
||||
.parse_prompt_response(
|
||||
"claude-session",
|
||||
&prompt_response_usage(100, 20, 175, Some(30), Some(25)),
|
||||
)
|
||||
.unwrap(),
|
||||
StopReason::EndTurn
|
||||
);
|
||||
|
||||
let usage = client.take_turn_usage().expect("prompt usage");
|
||||
assert!(usage.delta_reliable, "response tokens need no baseline");
|
||||
assert_eq!(usage.turn_input_tokens, Some(155));
|
||||
assert_eq!(usage.turn_output_tokens, Some(20));
|
||||
assert_eq!(
|
||||
usage.turn_total_tokens, None,
|
||||
"Claude total is adapter-derived"
|
||||
);
|
||||
assert_eq!(usage.turn_cache_read_tokens, Some(30));
|
||||
assert_eq!(usage.turn_cache_write_tokens, Some(25));
|
||||
assert_eq!(usage.turn_cost_usd, Some(0.042));
|
||||
assert_eq!(usage.cumulative_cost_usd, Some(0.042));
|
||||
assert_eq!(usage.cumulative_input_tokens, None);
|
||||
assert_eq!(usage.cumulative_output_tokens, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_prompt_response_usage_preserves_provider_total_without_cost() {
|
||||
let mut client = spawn_inert_client().await;
|
||||
client.standard_adapter = Some(StandardAdapterKind::Codex);
|
||||
client.standard_usage.begin_turn("codex-session");
|
||||
client.handle_session_update(&standard_cost_update("codex-session", 0.042));
|
||||
client
|
||||
.parse_prompt_response(
|
||||
"codex-session",
|
||||
&prompt_response_usage(90, 10, 140, Some(40), None),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let usage = client.take_turn_usage().expect("prompt usage");
|
||||
assert!(usage.delta_reliable);
|
||||
assert_eq!(usage.turn_input_tokens, Some(130));
|
||||
assert_eq!(usage.turn_output_tokens, Some(10));
|
||||
assert_eq!(usage.turn_total_tokens, Some(140));
|
||||
assert_eq!(usage.turn_cache_read_tokens, Some(40));
|
||||
assert_eq!(usage.turn_cache_write_tokens, None);
|
||||
assert_eq!(
|
||||
usage.cumulative_cost_usd, None,
|
||||
"Codex cost update is ignored"
|
||||
);
|
||||
assert_eq!(usage.cumulative_input_tokens, None);
|
||||
assert_eq!(usage.cumulative_output_tokens, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standard_prompt_input_overflow_fails_closed() {
|
||||
let mut client = spawn_inert_client().await;
|
||||
client.standard_adapter = Some(StandardAdapterKind::Claude);
|
||||
client.standard_usage.begin_turn("overflow-session");
|
||||
client
|
||||
.parse_prompt_response(
|
||||
"overflow-session",
|
||||
&prompt_response_usage(u64::MAX, 10, u64::MAX, Some(1), None),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
client.take_turn_usage().is_none(),
|
||||
"overflow without another valid signal must not emit all-null usage"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn claude_named_adapter_wire_lifecycle_records_prompt_and_cost() {
|
||||
let script = r#"
|
||||
read -r REQ
|
||||
ID=$(printf '%s' "$REQ" | sed -E 's/.*"id":([0-9]+).*/\1/')
|
||||
echo '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"wire-session","update":{"sessionUpdate":"usage_update","cost":{"amount":0.5,"currency":"USD"}}}}'
|
||||
echo '{"jsonrpc":"2.0","id":'"$ID"',"result":{"stopReason":"end_turn","usage":{"inputTokens":7,"outputTokens":3,"totalTokens":10,"cachedReadTokens":2}}}'
|
||||
sleep 1
|
||||
"#;
|
||||
let (mut client, dir) = spawn_named_script("claude-code", script).await;
|
||||
assert_eq!(client.standard_adapter, Some(StandardAdapterKind::Claude));
|
||||
client.notify_session_spawned("wire-session");
|
||||
|
||||
let stop = client
|
||||
.session_prompt_with_idle_timeout(
|
||||
"wire-session",
|
||||
"hello",
|
||||
std::time::Duration::from_secs(2),
|
||||
std::time::Duration::from_secs(5),
|
||||
)
|
||||
.await
|
||||
.expect("wire prompt");
|
||||
assert_eq!(stop, StopReason::EndTurn);
|
||||
|
||||
let usage = client.take_turn_usage().expect("wire usage");
|
||||
assert_eq!(usage.turn_seq, 1);
|
||||
assert_eq!(usage.turn_input_tokens, Some(9));
|
||||
assert_eq!(usage.turn_output_tokens, Some(3));
|
||||
assert_eq!(usage.turn_cost_usd, Some(0.5));
|
||||
assert_eq!(usage.cumulative_cost_usd, Some(0.5));
|
||||
drop(client);
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn claude_cost_only_record_survives_missing_prompt_usage() {
|
||||
let mut client = spawn_inert_client().await;
|
||||
client.standard_adapter = Some(StandardAdapterKind::Claude);
|
||||
client.notify_session_spawned("cost-only-session");
|
||||
client.standard_usage.begin_turn("cost-only-session");
|
||||
client.handle_session_update(&standard_cost_update("cost-only-session", 0.125));
|
||||
|
||||
let usage = client.take_turn_usage().expect("cost-only usage");
|
||||
assert_eq!(usage.turn_seq, 1);
|
||||
assert!(usage.delta_reliable);
|
||||
assert_eq!(usage.turn_input_tokens, None);
|
||||
assert_eq!(usage.turn_cost_usd, Some(0.125));
|
||||
assert_eq!(usage.cumulative_cost_usd, Some(0.125));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn attached_claude_session_does_not_invent_first_cost_delta() {
|
||||
let mut client = spawn_inert_client().await;
|
||||
client.standard_adapter = Some(StandardAdapterKind::Claude);
|
||||
client.standard_usage.begin_turn("attached-session");
|
||||
client.handle_session_update(&standard_cost_update("attached-session", 1.25));
|
||||
client
|
||||
.parse_prompt_response(
|
||||
"attached-session",
|
||||
&prompt_response_usage(10, 2, 12, None, None),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let usage = client.take_turn_usage().expect("attached usage");
|
||||
assert_eq!(usage.turn_cost_usd, None);
|
||||
assert_eq!(usage.cumulative_cost_usd, Some(1.25));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standard_usage_two_prompts_preserve_both_monotonic_sequences() {
|
||||
let mut client = spawn_inert_client().await;
|
||||
client.standard_adapter = Some(StandardAdapterKind::Claude);
|
||||
client.notify_session_spawned("two-prompt-session");
|
||||
|
||||
client.standard_usage.begin_turn("two-prompt-session");
|
||||
client.handle_session_update(&standard_cost_update("two-prompt-session", 0.1));
|
||||
client
|
||||
.parse_prompt_response(
|
||||
"two-prompt-session",
|
||||
&prompt_response_usage(10, 2, 12, None, None),
|
||||
)
|
||||
.unwrap();
|
||||
let initial = client.take_turn_usage().expect("initial prompt usage");
|
||||
|
||||
client.standard_usage.begin_turn("two-prompt-session");
|
||||
client.handle_session_update(&standard_cost_update("two-prompt-session", 0.25));
|
||||
client
|
||||
.parse_prompt_response(
|
||||
"two-prompt-session",
|
||||
&prompt_response_usage(20, 3, 23, None, None),
|
||||
)
|
||||
.unwrap();
|
||||
let user = client.take_turn_usage().expect("user prompt usage");
|
||||
|
||||
assert_eq!((initial.turn_seq, user.turn_seq), (1, 2));
|
||||
assert_eq!(
|
||||
(initial.turn_input_tokens, user.turn_input_tokens),
|
||||
(Some(10), Some(20))
|
||||
);
|
||||
assert_eq!(
|
||||
(initial.turn_cost_usd, user.turn_cost_usd),
|
||||
(Some(0.1), Some(0.15))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn goose_usage_stays_exclusive_and_drains_standard_usage() {
|
||||
let mut client = spawn_inert_client().await;
|
||||
client.standard_adapter = Some(StandardAdapterKind::Claude);
|
||||
client.goose_usage.begin_turn("goose-session");
|
||||
client.standard_usage.begin_turn("goose-session");
|
||||
client.handle_goose_usage_update(&goose_usage_update_msg("goose-session", 1000, 200, None));
|
||||
client
|
||||
.parse_prompt_response(
|
||||
"goose-session",
|
||||
&prompt_response_usage(100, 20, 120, None, None),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let usage = client.take_turn_usage().expect("goose usage");
|
||||
assert_eq!(usage.cumulative_input_tokens, Some(1000));
|
||||
assert_eq!(
|
||||
usage.turn_input_tokens, None,
|
||||
"goose first delta remains exclusive"
|
||||
);
|
||||
assert!(
|
||||
client.take_turn_usage().is_none(),
|
||||
"standard usage was drained"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Goose usage notification integration ──────────────────────────────
|
||||
|
||||
/// Build a `_goose/unstable/session/update` JSON-RPC notification.
|
||||
|
||||
@@ -1850,6 +1850,16 @@ pub async fn run_prompt_task(
|
||||
if !agent.has_system_prompt_support() {
|
||||
agent.state.mark_channel_delivery_success(*cid, true, []);
|
||||
}
|
||||
let usage = agent.acp.take_turn_usage();
|
||||
publish_agent_turn_metric(
|
||||
&ctx,
|
||||
usage,
|
||||
Some(*cid),
|
||||
&session_id,
|
||||
&format!("{turn_id}:initial"),
|
||||
Some(acp_stop_to_core(&stop_reason)),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(AcpError::AgentExited) => {
|
||||
agent.state.invalidate_all();
|
||||
@@ -1874,7 +1884,17 @@ pub async fn run_prompt_task(
|
||||
.cancel_with_cleanup(&session_id, ctx.idle_timeout)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
Ok(stop_reason) => {
|
||||
let usage = agent.acp.take_turn_usage();
|
||||
publish_agent_turn_metric(
|
||||
&ctx,
|
||||
usage,
|
||||
Some(*cid),
|
||||
&session_id,
|
||||
&format!("{turn_id}:initial"),
|
||||
Some(acp_stop_to_core(&stop_reason)),
|
||||
)
|
||||
.await;
|
||||
agent.state.invalidate(&source);
|
||||
}
|
||||
Err(AcpError::AgentExited) => {
|
||||
|
||||
@@ -244,6 +244,158 @@ pub struct TurnUsage {
|
||||
pub pricing_identity: Option<buzz_core::agent_turn_metric::PricingIdentity>,
|
||||
}
|
||||
|
||||
/// Per-turn usage carried by a standard ACP `session/prompt` response.
|
||||
/// Adapter input excludes cache reads and writes, so NIP-AM input must add
|
||||
/// those subsets with checked arithmetic.
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct PromptResponseUsage {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub cached_read_tokens: Option<u64>,
|
||||
pub cached_write_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum StandardAdapterKind {
|
||||
Claude,
|
||||
Codex,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct StandardSessionState {
|
||||
published_seq: u64,
|
||||
last_cost: Option<f64>,
|
||||
cost_poisoned: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct StandardUsageTracker {
|
||||
sessions: HashMap<String, StandardSessionState>,
|
||||
in_flight_session: Option<String>,
|
||||
pending_cost: Option<(String, f64)>,
|
||||
pending_prompt: Option<(String, PromptResponseUsage, StandardAdapterKind)>,
|
||||
}
|
||||
|
||||
impl StandardUsageTracker {
|
||||
pub(crate) fn seed_zero_baseline(&mut self, session_id: &str) {
|
||||
self.sessions
|
||||
.entry(session_id.to_string())
|
||||
.or_insert_with(|| StandardSessionState {
|
||||
published_seq: 0,
|
||||
last_cost: Some(0.0),
|
||||
cost_poisoned: false,
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn begin_turn(&mut self, session_id: &str) {
|
||||
self.in_flight_session = Some(session_id.to_string());
|
||||
self.pending_cost = None;
|
||||
self.pending_prompt = None;
|
||||
}
|
||||
|
||||
/// Claude's `usage_update.cost.amount` is a raw session-cumulative total.
|
||||
pub(crate) fn record_cost(&mut self, session_id: &str, cost: f64) {
|
||||
if cost.is_finite() && cost >= 0.0 && self.in_flight_session.as_deref() == Some(session_id)
|
||||
{
|
||||
self.pending_cost = Some((session_id.to_string(), cost));
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn record_prompt_usage(
|
||||
&mut self,
|
||||
session_id: &str,
|
||||
usage: PromptResponseUsage,
|
||||
adapter: StandardAdapterKind,
|
||||
) {
|
||||
if self.in_flight_session.as_deref() == Some(session_id) {
|
||||
self.pending_prompt = Some((session_id.to_string(), usage, adapter));
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn take(&mut self) -> Option<TurnUsage> {
|
||||
self.in_flight_session = None;
|
||||
let prompt = self.pending_prompt.take();
|
||||
let cost = self.pending_cost.take();
|
||||
let session_id = prompt
|
||||
.as_ref()
|
||||
.map(|(session_id, _, _)| session_id.clone())
|
||||
.or_else(|| cost.as_ref().map(|(session_id, _)| session_id.clone()))?;
|
||||
|
||||
let (inclusive_input, output_tokens, total_tokens, cache_read, cache_write) = match prompt {
|
||||
Some((_, usage, adapter)) => {
|
||||
let inclusive_input = usage
|
||||
.input_tokens
|
||||
.checked_add(usage.cached_read_tokens.unwrap_or(0))
|
||||
.and_then(|input| input.checked_add(usage.cached_write_tokens.unwrap_or(0)));
|
||||
let total_tokens =
|
||||
(adapter == StandardAdapterKind::Codex).then_some(usage.total_tokens);
|
||||
(
|
||||
inclusive_input,
|
||||
inclusive_input.map(|_| usage.output_tokens),
|
||||
inclusive_input.and(total_tokens),
|
||||
inclusive_input.and(usage.cached_read_tokens),
|
||||
inclusive_input.and(usage.cached_write_tokens),
|
||||
)
|
||||
}
|
||||
None => (None, None, None, None, None),
|
||||
};
|
||||
|
||||
let state = self.sessions.entry(session_id.clone()).or_default();
|
||||
let cumulative_cost = cost.map(|(_, cost)| cost);
|
||||
let turn_cost = match (state.cost_poisoned, state.last_cost, cumulative_cost) {
|
||||
(false, Some(previous), Some(current)) if current >= previous => {
|
||||
let delta = current - previous;
|
||||
delta.is_finite().then_some(delta)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
if let Some(current) = cumulative_cost {
|
||||
// A decrease means the cumulative series restarted or is corrupt.
|
||||
// Poison the baseline rather than deriving a later delta across the
|
||||
// discontinuity. The raw cumulative value still remains observable.
|
||||
if state.last_cost.is_some_and(|previous| current < previous) {
|
||||
state.cost_poisoned = true;
|
||||
state.last_cost = None;
|
||||
} else if !state.cost_poisoned {
|
||||
state.last_cost = Some(current);
|
||||
}
|
||||
}
|
||||
|
||||
// Input overflow invalidates the standard prompt counters. Emit only if
|
||||
// another valid signal (normally Claude cost) remains; NIP-AM forbids an
|
||||
// otherwise all-null usage record.
|
||||
if inclusive_input.is_none() && cumulative_cost.is_none() {
|
||||
return None;
|
||||
}
|
||||
|
||||
state.published_seq += 1;
|
||||
Some(TurnUsage {
|
||||
session_id,
|
||||
turn_seq: state.published_seq,
|
||||
// Standard prompt counters are per-turn already. A cost-only record
|
||||
// is reliable only when a seeded/previous cumulative baseline made
|
||||
// the cost delta provable.
|
||||
delta_reliable: inclusive_input.is_some() || turn_cost.is_some(),
|
||||
turn_input_tokens: inclusive_input,
|
||||
turn_output_tokens: output_tokens,
|
||||
turn_total_tokens: total_tokens,
|
||||
turn_cost_usd: turn_cost,
|
||||
turn_cache_read_tokens: cache_read,
|
||||
turn_cache_write_tokens: cache_write,
|
||||
cumulative_input_tokens: None,
|
||||
cumulative_output_tokens: None,
|
||||
cumulative_total_tokens: None,
|
||||
cumulative_cost_usd: cumulative_cost,
|
||||
cumulative_cache_read_tokens: None,
|
||||
cumulative_cache_write_tokens: None,
|
||||
model: None,
|
||||
pricing_identity: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Tracks per-session cumulative usage state across turns.
|
||||
///
|
||||
/// Cheap to construct. Usage lifecycle per turn:
|
||||
|
||||
Reference in New Issue
Block a user