From 115e77ff271672554a2003a91246269f7b70037c Mon Sep 17 00:00:00 2001
From: xiaoxixi <xiaoski@qq.com>
Date: Fri, 24 Jul 2026 17:11:37 +0800
Subject: [PATCH] feat(config): optional provider pricing for cost metrics

---
 src/agent/sub_agent.rs     |  2 ++
 src/config/mod.rs          | 15 +++++++++++++++
 src/session/persistence.rs |  2 ++
 tests/test_integration.rs  |  2 ++
 tests/test_tool_calling.rs |  2 ++
 5 files changed, 23 insertions(+)

diff --git a/src/agent/sub_agent.rs b/src/agent/sub_agent.rs
index c568744..2dc6a88 100644
--- a/src/agent/sub_agent.rs
+++ b/src/agent/sub_agent.rs
@@ -832,6 +832,8 @@ mod tests {
                 token_limit: 4096,
                 workspace_dir: std::env::temp_dir(),
                 input_types: vec!["text".into()],
+                price_input_per_million: None,
+                price_output_per_million: None,
             },
             Arc::new(ToolRegistry::new()),
             None,
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 913ac18..a6fca23 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -508,6 +508,19 @@ pub struct LLMProviderConfig {
     pub token_limit: usize,
     pub workspace_dir: PathBuf,
     pub input_types: Vec<String>,
+    pub price_input_per_million: Option<f64>,
+    pub price_output_per_million: Option<f64>,
+}
+
+impl LLMProviderConfig {
+    pub fn cost_of(&self, prompt_tokens: u32, completion_tokens: u32) -> Option<f64> {
+        match (self.price_input_per_million, self.price_output_per_million) {
+            (Some(pi), Some(po)) => Some(
+                prompt_tokens as f64 / 1e6 * pi + completion_tokens as f64 / 1e6 * po,
+            ),
+            _ => None,
+        }
+    }
 }
 
 pub fn get_default_config_path() -> PathBuf {
@@ -653,6 +666,8 @@ impl Config {
             token_limit: agent.token_limit,
             workspace_dir: expand_path(&self.workspace_dir),
             input_types: model.input_type.clone(),
+            price_input_per_million: None,
+            price_output_per_million: None,
         })
     }
 }
diff --git a/src/session/persistence.rs b/src/session/persistence.rs
index 5de8725..e3f51eb 100644
--- a/src/session/persistence.rs
+++ b/src/session/persistence.rs
@@ -225,6 +225,8 @@ mod tests {
             token_limit: 8_192,
             workspace_dir: PathBuf::from("."),
             input_types: vec!["text".to_string(), "image".to_string()],
+            price_input_per_million: None,
+            price_output_per_million: None,
         };
         let session = Arc::new(Mutex::new(
             Session::new(
diff --git a/tests/test_integration.rs b/tests/test_integration.rs
index 7f81dca..5b97428 100644
--- a/tests/test_integration.rs
+++ b/tests/test_integration.rs
@@ -27,6 +27,8 @@ fn load_config() -> Option<LLMProviderConfig> {
         token_limit: 128_000,
         workspace_dir: std::path::PathBuf::from("/tmp/test-workspace"),
         input_types: vec!["text".to_string()],
+        price_input_per_million: None,
+        price_output_per_million: None,
     })
 }
 
diff --git a/tests/test_tool_calling.rs b/tests/test_tool_calling.rs
index bdf1086..4704b7f 100644
--- a/tests/test_tool_calling.rs
+++ b/tests/test_tool_calling.rs
@@ -27,6 +27,8 @@ fn load_openai_config() -> Option<LLMProviderConfig> {
         token_limit: 128_000,
         workspace_dir: std::path::PathBuf::from("/tmp/test-workspace"),
         input_types: vec!["text".to_string()],
+        price_input_per_million: None,
+        price_output_per_million: None,
     })
 }