Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 44 additions & 1 deletion src/acp/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -755,7 +755,7 @@ impl AcpService {
usage.output,
usage.cache_read,
usage.cache_write,
0.0,
estimate_session_usage_cost(&session, &usage),
));
if usage.output > 0 {
assistant.output_tokens = Some(
Expand Down Expand Up @@ -992,6 +992,27 @@ fn model_reasoning_capability(
.filter(|capability| !capability.values().is_empty())
}

fn estimate_session_usage_cost(
session: &AcpSession,
usage: &crate::aisdk::chunk::TokenUsage,
) -> f64 {
crate::model::discovery::Discovery::new_with_custom(Some(
session.config.merged_config.custom_providers.clone(),
))
.ok()
.map(|discovery| {
discovery.estimate_usage_cost(
&session.provider,
&session.model,
usage.input,
usage.output,
usage.cache_read,
usage.cache_write,
)
})
.unwrap_or(0.0)
}

fn model_context_window(config: &LoadedConfig, provider: &str, model: &str) -> Option<u32> {
let discovery = crate::model::discovery::Discovery::new_with_custom(Some(
config.merged_config.custom_providers.clone(),
Expand Down Expand Up @@ -1884,4 +1905,26 @@ mod tests {
]
);
}

#[test]
fn unknown_model_usage_cost_is_zero() {
let mut session = session_with_config(LoadedConfig {
merged_config: crate::config::configuration::MergedConfig::default(),
raw_merged: serde_json::Value::Null,
diagnostics: Default::default(),
inventory: Default::default(),
project_root: PathBuf::from("/tmp"),
cwd: PathBuf::from("/tmp"),
xdg_config_home: PathBuf::from("/tmp"),
});
session.provider = "no-such-provider".to_string();
session.model = "no-such-model".to_string();
let usage = crate::aisdk::chunk::TokenUsage {
input: 1_000,
output: 1_000,
cache_read: 0,
cache_write: 0,
};
assert_eq!(estimate_session_usage_cost(&session, &usage), 0.0);
}
}
3 changes: 2 additions & 1 deletion src/aisdk/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ In-tree AI SDK used by the host binary (`mod aisdk` in `src/main.rs`).
2. Do not reintroduce a parallel `aisdk/` workspace package until we intentionally extract and publish.
3. Host app imports via `crate::aisdk::...` (binary module path), not an external crate dep.
4. **No hard coupling to the host product.** Code here must not depend on product-specific concepts (TUI, sessions, tools registry, prefs DB, agent loop, product config, host `crate::` app modules, etc.). Keep this layer a generic multi-provider AI SDK — same spirit as Vercel AI SDK for Rust: providers, messages, streams, tools, retries. Product-specific behavior belongs outside this tree (call sites in the app).
5. **Capabilities ≠ host policy.** Expose provider-executed features as ordinary tools (e.g. `openai::tools::web_search()`, `xai::tools::x_search()` via `ToolTransport`). Do **not** encode host preferences like `websearch.native` / `.hosted_web_search(true)` inside aisdk. The host maps policy → `HostedSearchSelection` / tools to pass.
5. **Capabilities ≠ host policy.** Expose provider-executed features as ordinary tools (e.g. `openai::tools::web_search()`, `xai::tools::x_search()` via `ToolTransport`). Do **not** encode host preferences like `websearch.native` / `.hosted_web_search(true)` inside aisdk. The host maps policy → `HostedSearchSelection` / tools to pass. Providers emit billed `ChunkType::Usage(TokenUsage)` (and optional `ResponseCompleted.usage`); pricing stays in the host.
6. **Logging:** use `crate::log::log(...)` (host-injected via `aisdk::log::set_logger`). Never call host `emit_log!` from this tree.
7. **Debug SSE dumps:** feature-gated behind `aisdk-sse-debug` (optional path via `AISDK_SSE_DEBUG_LOG`).
8. **Boundary check:** run `scripts/check-aisdk-boundary.sh` (also documented in root `AGENTS.md`).
Expand All @@ -33,6 +33,7 @@ In-tree AI SDK used by the host binary (`mod aisdk` in `src/main.rs`).
Done for packaging/host hooks:

- Neutral logging (`log` module + host `set_logger`)
- Provider-neutral token usage events (`ChunkType::Usage`)
- No `crate::aisdk::...` inside the tree (`mod.rs` / re-exports use `super::`)
- Absolute `crate::{chunk,error,...}` paths are crate-root-shaped (host re-exports them today)
- Product-leaky debug path renamed/feature-gated
Expand Down
62 changes: 62 additions & 0 deletions src/aisdk/providers/anthropic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -989,6 +989,68 @@ mod tests {
));
}

#[test]
fn max_tokens_delta_emits_final_usage_then_incomplete() {
let value = serde_json::json!({
"type": "message_delta",
"delta": {
"stop_reason": "max_tokens",
},
"usage": {
"input_tokens": 12,
"output_tokens": 34,
"cache_read_input_tokens": 5,
"cache_creation_input_tokens": 2,
},
});
let mut pending = PartialAnthropicUsage::default();
let chunks = anthropic_stream_chunks("message_delta", &value, &mut pending);

assert!(matches!(
chunks.as_slice(),
[
Ok(ChunkType::Usage(crate::chunk::TokenUsage {
input: 12,
output: 34,
cache_read: 5,
cache_write: 2,
})),
Ok(ChunkType::Incomplete(_)),
]
));
}

#[test]
fn end_turn_delta_emits_final_usage_then_terminal_reason() {
let value = serde_json::json!({
"type": "message_delta",
"delta": {
"stop_reason": "end_turn",
},
"usage": {
"input_tokens": 7,
"output_tokens": 11,
},
});
let mut pending = PartialAnthropicUsage::default();
let chunks = anthropic_stream_chunks("message_delta", &value, &mut pending);

assert!(matches!(
chunks.as_slice(),
[
Ok(ChunkType::Usage(crate::chunk::TokenUsage {
input: 7,
output: 11,
cache_read: 0,
cache_write: 0,
})),
Ok(ChunkType::End {
reason: Some(FinishReason::EndTurn)
}),
]
));
}

#[test]
fn groups_adjacent_tool_calls_and_results() {
let messages = vec![
Expand Down
9 changes: 4 additions & 5 deletions src/app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9743,11 +9743,10 @@ impl App {
let cost = self
.discovery
.as_ref()
.and_then(|discovery| {
discovery.get_model_pricing(&self.provider_name.to_lowercase(), &self.model)
})
.map(|pricing| {
pricing.estimate_tokens(
.map(|discovery| {
discovery.estimate_usage_cost(
&self.provider_name,
&self.model,
usage.input,
usage.output,
usage.cache_read,
Expand Down
89 changes: 89 additions & 0 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ mod remote_mcp;
mod session;
mod skill;
mod sound;
mod stats;
mod streaming;
mod terminal_title;
mod theme;
Expand Down Expand Up @@ -790,6 +791,25 @@ enum Command {
target: Option<String>,
},

/// Show token usage and cost statistics
Stats {
/// Show stats for the last N days (default: all time)
#[arg(long)]
days: Option<u64>,

/// Number of tools to show (default: all)
#[arg(long)]
tools: Option<usize>,

/// Show model statistics; optionally limit to the top N
#[arg(long, num_args = 0..=1, default_missing_value = "all")]
models: Option<String>,

/// Filter by project (default: all projects, empty string: current project)
#[arg(long, value_name = "PROJECT", num_args = 0..=1, default_missing_value = "")]
project: Option<String>,
},

/// Manage survive-quit background jobs (list / logs / stop)
Jobs {
#[command(subcommand)]
Expand Down Expand Up @@ -1007,6 +1027,32 @@ async fn main() -> Result<()> {
Some(Command::Upgrade { target }) => {
return crate::upgrade::upgrade(target.as_deref());
}
Some(Command::Stats {
days,
tools,
models,
project,
}) => {
let models = models
.as_deref()
.map(|value| {
if value == "all" {
Ok(None)
} else {
value
.parse::<usize>()
.map(Some)
.context("--models must be a non-negative integer")
}
})
.transpose()?;
return crate::stats::run(crate::stats::StatsOptions {
days: *days,
tools: *tools,
models,
project: project.clone(),
});
}
Some(Command::Jobs { command }) => {
let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
match command {
Expand Down Expand Up @@ -1307,6 +1353,47 @@ mod tests {
}
}

#[test]
fn parses_stats_command_and_compatible_options() {
let args = Args::try_parse_from([
"crabcode",
"stats",
"--days",
"7",
"--tools",
"5",
"--models",
"3",
"--project",
"",
])
.unwrap();

match args.command {
Some(Command::Stats {
days,
tools,
models,
project,
}) => {
assert_eq!(days, Some(7));
assert_eq!(tools, Some(5));
assert_eq!(models.as_deref(), Some("3"));
assert_eq!(project.as_deref(), Some(""));
}
other => panic!("expected stats command, got {other:?}"),
}

let args = Args::try_parse_from(["crabcode", "stats", "--models"]).unwrap();
assert!(matches!(
args.command,
Some(Command::Stats {
models: Some(ref models),
..
}) if models == "all"
));
}

#[test]
fn generates_bash_completion() {
let script = String::from_utf8(crate::completion::generate_script(Shell::Bash)).unwrap();
Expand Down Expand Up @@ -1352,6 +1439,8 @@ mod tests {
assert!(help.contains("Usage: crabcode"));
assert!(help.contains("completion"));
assert!(help.contains("Generate or install shell completions"));
assert!(help.contains("stats"));
assert!(help.contains("Show token usage and cost statistics"));
assert!(
help.contains("serve Host the current workspace for browser and CLI clients")
);
Expand Down
14 changes: 14 additions & 0 deletions src/model/discovery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -796,6 +796,20 @@ impl Discovery {
model.cost.clone()
}

pub fn estimate_usage_cost(
&self,
provider_id: &str,
model_id: &str,
input: u64,
output: u64,
cache_read: u64,
cache_write: u64,
) -> f64 {
self.get_model_pricing(&provider_id.to_lowercase(), model_id)
.map(|pricing| pricing.estimate_tokens(input, output, cache_read, cache_write))
.unwrap_or(0.0)
}

pub fn get_model_limit(&self, provider_id: &str, model_id: &str) -> Option<u32> {
let entry = self.load_cache_entry().ok()??;
let provider = entry.data.get(provider_id)?;
Expand Down
Loading
Loading