fix(cost): accrue review tool LLM usage

This commit is contained in:
Hunter Bown
2026-05-03 13:32:14 -05:00
parent db2f761120
commit 1d315ec3d6
2 changed files with 60 additions and 2 deletions
+42 -2
View File
@@ -10,7 +10,7 @@ use serde_json::{Value, json};
use crate::client::DeepSeekClient;
use crate::llm_client::LlmClient;
use crate::models::{ContentBlock, Message, MessageRequest, SystemPrompt};
use crate::models::{ContentBlock, Message, MessageRequest, SystemPrompt, Usage};
use crate::utils::truncate_with_ellipsis;
use super::spec::{
@@ -240,10 +240,27 @@ impl ToolSpec for ReviewTool {
let response_text = extract_text(&response.content);
let output = ReviewOutput::from_str(&response_text);
ToolResult::json(&output).map_err(|e| ToolError::execution_failed(e.to_string()))
let metadata = review_usage_metadata(&response.model, &response.usage);
let result =
ToolResult::json(&output).map_err(|e| ToolError::execution_failed(e.to_string()))?;
Ok(result.with_metadata(metadata))
}
}
fn review_usage_metadata(model: &str, usage: &Usage) -> Value {
json!({
"tool": "review",
"input_tokens": usage.input_tokens,
"output_tokens": usage.output_tokens,
"child_model": model,
"child_input_tokens": usage.input_tokens,
"child_output_tokens": usage.output_tokens,
"child_prompt_cache_hit_tokens": usage.prompt_cache_hit_tokens,
"child_prompt_cache_miss_tokens": usage.prompt_cache_miss_tokens,
"child_reasoning_tokens": usage.reasoning_tokens,
})
}
enum ReviewSource {
File { display: String, content: String },
Diff { label: String, diff: String },
@@ -537,4 +554,27 @@ mod tests {
assert!(!output.summary.is_empty());
assert!(output.issues.is_empty());
}
#[test]
fn review_usage_metadata_reports_child_tokens_for_cost_accrual() {
let metadata = review_usage_metadata(
"deepseek-v4-flash",
&Usage {
input_tokens: 123,
output_tokens: 45,
prompt_cache_hit_tokens: Some(100),
prompt_cache_miss_tokens: Some(23),
reasoning_tokens: Some(7),
..Default::default()
},
);
assert_eq!(metadata["tool"], "review");
assert_eq!(metadata["child_model"], "deepseek-v4-flash");
assert_eq!(metadata["child_input_tokens"], 123);
assert_eq!(metadata["child_output_tokens"], 45);
assert_eq!(metadata["child_prompt_cache_hit_tokens"], 100);
assert_eq!(metadata["child_prompt_cache_miss_tokens"], 23);
assert_eq!(metadata["child_reasoning_tokens"], 7);
}
}
+18
View File
@@ -1976,6 +1976,24 @@ fn ok_result(
Ok(crate::tools::spec::ToolResult::success(content))
}
#[test]
fn tool_child_usage_metadata_updates_live_cost_counter() {
let mut app = create_test_app();
let result = Ok(crate::tools::spec::ToolResult::success("ok").with_metadata(
serde_json::json!({
"child_model": "deepseek-v4-flash",
"child_input_tokens": 10_000,
"child_output_tokens": 1_000,
"child_prompt_cache_hit_tokens": 7_000,
"child_prompt_cache_miss_tokens": 3_000,
}),
));
handle_tool_call_complete(&mut app, "review-usage", "review", &result);
assert!(app.session.subagent_cost > 0.0);
}
#[test]
fn parallel_exploring_tool_starts_share_one_active_entry() {
// Three exploring tools start in any order; they must collapse into one