feat: in-session model switching and cost awareness (#32)

src-tauri/src/commands/preferences.rs

@@ -1,3 +1,6 @@
+use std::collections::HashMap;
+use std::sync::{Mutex, OnceLock};
+
 use serde::{Deserialize, Serialize};
 
 use crate::config_io::{read_json, write_json};
@@ -89,6 +92,19 @@ pub fn get_zeroclaw_usage_stats() -> Result<ZeroclawUsageStatsResponse, String>
     })
 }
 
+#[tauri::command]
+pub fn get_session_usage_stats(session_id: String) -> Result<ZeroclawUsageStatsResponse, String> {
+    let stats = crate::runtime::zeroclaw::process::get_session_usage(&session_id);
+    Ok(ZeroclawUsageStatsResponse {
+        total_calls: stats.total_calls,
+        usage_calls: stats.usage_calls,
+        prompt_tokens: stats.prompt_tokens,
+        completion_tokens: stats.completion_tokens,
+        total_tokens: stats.total_tokens,
+        last_updated_ms: stats.last_updated_ms,
+    })
+}
+
 #[tauri::command]
 pub fn get_zeroclaw_runtime_target() -> Result<ZeroclawRuntimeTargetResponse, String> {
     let target = crate::runtime::zeroclaw::process::get_zeroclaw_runtime_target();
@@ -101,6 +117,52 @@ pub fn get_zeroclaw_runtime_target() -> Result<ZeroclawRuntimeTargetResponse, St
     })
 }
 
+// ---------------------------------------------------------------------------
+// Per-session model overrides (in-memory only)
+// ---------------------------------------------------------------------------
+
+fn session_model_overrides() -> &'static Mutex<HashMap<String, String>> {
+    static STORE: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
+    STORE.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+/// Look up a session model override without going through Tauri command dispatch.
+pub fn lookup_session_model_override(session_id: &str) -> Option<String> {
+    session_model_overrides()
+        .lock()
+        .ok()?
+        .get(session_id)
+        .cloned()
+}
+
+#[tauri::command]
+pub fn set_session_model_override(session_id: String, model: String) -> Result<(), String> {
+    let trimmed = model.trim().to_string();
+    if trimmed.is_empty() {
+        return Err("model must not be empty".into());
+    }
+    if let Ok(mut map) = session_model_overrides().lock() {
+        map.insert(session_id, trimmed);
+    }
+    Ok(())
+}
+
+#[tauri::command]
+pub fn get_session_model_override(session_id: String) -> Result<Option<String>, String> {
+    let map = session_model_overrides()
+        .lock()
+        .map_err(|e| e.to_string())?;
+    Ok(map.get(&session_id).cloned())
+}
+
+#[tauri::command]
+pub fn clear_session_model_override(session_id: String) -> Result<(), String> {
+    if let Ok(mut map) = session_model_overrides().lock() {
+        map.remove(&session_id);
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src-tauri/src/lib.rs

@@ -8,12 +8,13 @@ use crate::cli_runner::{
 };
 use crate::commands::{
     analyze_sessions, apply_config_patch, backup_before_upgrade, chat_via_openclaw,
-    check_openclaw_update, clear_all_sessions, connect_docker_instance, connect_local_instance,
-    connect_ssh_instance, create_agent, delete_agent, delete_backup, delete_cron_job,
-    delete_local_instance_home, delete_model_profile, delete_registered_instance,
-    delete_sessions_by_ids, delete_ssh_host, deploy_watchdog, diagnose_primary_via_rescue,
-    discover_local_instances, ensure_access_profile, extract_model_profiles_from_config,
-    fix_issues, get_app_preferences, get_cached_model_catalog, get_cron_runs, get_status_extra,
+    check_openclaw_update, clear_all_sessions, clear_session_model_override,
+    connect_docker_instance, connect_local_instance, connect_ssh_instance, create_agent,
+    delete_agent, delete_backup, delete_cron_job, delete_local_instance_home, delete_model_profile,
+    delete_registered_instance, delete_sessions_by_ids, delete_ssh_host, deploy_watchdog,
+    diagnose_primary_via_rescue, discover_local_instances, ensure_access_profile,
+    extract_model_profiles_from_config, fix_issues, get_app_preferences, get_cached_model_catalog,
+    get_cron_runs, get_session_model_override, get_session_usage_stats, get_status_extra,
     get_status_light, get_system_status, get_watchdog_status, get_zeroclaw_runtime_target,
     get_zeroclaw_usage_stats, list_agents_overview, list_backups, list_bindings,
     list_channels_minimal, list_cron_jobs, list_discord_guild_channels, list_history,
@@ -42,11 +43,11 @@ use crate::commands::{
     remote_upsert_model_profile, remote_write_raw_config, repair_primary_via_rescue,
     resolve_api_keys, resolve_provider_auth, restart_gateway, restore_from_backup, rollback,
     run_doctor_command, run_openclaw_upgrade, set_active_clawpal_data_dir,
-    set_active_openclaw_home, set_agent_model, set_global_model, set_zeroclaw_model_preference,
-    setup_agent_identity, sftp_list_dir, sftp_read_file, sftp_remove_file, sftp_write_file,
-    ssh_connect, ssh_connect_with_passphrase, ssh_disconnect, ssh_exec, ssh_status, start_watchdog,
-    stop_watchdog, test_model_profile, trigger_cron_job, uninstall_watchdog, upsert_model_profile,
-    upsert_ssh_host,
+    set_active_openclaw_home, set_agent_model, set_global_model, set_session_model_override,
+    set_zeroclaw_model_preference, setup_agent_identity, sftp_list_dir, sftp_read_file,
+    sftp_remove_file, sftp_write_file, ssh_connect, ssh_connect_with_passphrase, ssh_disconnect,
+    ssh_exec, ssh_status, start_watchdog, stop_watchdog, test_model_profile, trigger_cron_job,
+    uninstall_watchdog, upsert_model_profile, upsert_ssh_host,
 };
 use crate::doctor_commands::{
     collect_doctor_context, collect_doctor_context_remote, doctor_approve_invoke, doctor_connect,
@@ -59,6 +60,7 @@ use crate::install::commands::{
 use crate::install::session_store::InstallSessionStore;
 use crate::install_commands::{install_send_message, install_start_session};
 use crate::node_client::NodeClient;
+use crate::runtime::zeroclaw::cost::estimate_query_cost;
 use crate::ssh::SshConnectionPool;
 
 pub mod access_discovery;
@@ -119,7 +121,12 @@ pub fn run() {
             get_status_extra,
             get_app_preferences,
             get_zeroclaw_usage_stats,
+            get_session_usage_stats,
             get_zeroclaw_runtime_target,
+            set_session_model_override,
+            get_session_model_override,
+            clear_session_model_override,
+            estimate_query_cost,
             list_recipes,
             list_model_profiles,
             get_cached_model_catalog,

src-tauri/src/runtime/zeroclaw/cost.rs

@@ -0,0 +1,78 @@
+use serde::{Deserialize, Serialize};
+
+/// Model pricing: (prompt_price_per_1k_tokens, completion_price_per_1k_tokens)
+fn model_pricing(model: &str) -> Option<(f64, f64)> {
+    let lower = model.trim().to_ascii_lowercase();
+    // Match by substring to handle provider prefixes like "openrouter/anthropic/claude-3.7-sonnet"
+    if lower.contains("claude-3.7-sonnet") || lower.contains("claude-3-7-sonnet") {
+        Some((0.003, 0.015))
+    } else if lower.contains("claude-3.5-haiku") || lower.contains("claude-3-5-haiku") {
+        Some((0.0008, 0.004))
+    } else if lower.contains("gpt-4o-mini") {
+        Some((0.00015, 0.0006))
+    } else if lower.contains("gpt-4o") {
+        Some((0.0025, 0.01))
+    } else if lower.contains("gpt-4.1") {
+        Some((0.002, 0.008))
+    } else if lower.contains("gemini-2.0-flash") {
+        Some((0.0001, 0.0004))
+    } else if lower.contains("kimi-k2") {
+        Some((0.0006, 0.002))
+    } else {
+        None
+    }
+}
+
+/// Estimate cost in USD for a given model and token counts.
+pub fn estimate_cost(model: &str, prompt_tokens: u64, completion_tokens: u64) -> Option<f64> {
+    let (prompt_price, completion_price) = model_pricing(model)?;
+    let cost = (prompt_tokens as f64 / 1000.0) * prompt_price
+        + (completion_tokens as f64 / 1000.0) * completion_price;
+    Some(cost)
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct CostEstimate {
+    pub model: String,
+    pub prompt_tokens: u64,
+    pub completion_tokens: u64,
+    pub estimated_cost_usd: Option<f64>,
+}
+
+#[tauri::command]
+pub fn estimate_query_cost(
+    model: String,
+    prompt_tokens: u64,
+    completion_tokens: u64,
+) -> Result<CostEstimate, String> {
+    let cost = estimate_cost(&model, prompt_tokens, completion_tokens);
+    Ok(CostEstimate {
+        model,
+        prompt_tokens,
+        completion_tokens,
+        estimated_cost_usd: cost,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn estimate_cost_for_known_model() {
+        let cost = estimate_cost("gpt-4o", 1000, 1000).unwrap();
+        assert!((cost - 0.0125).abs() < 0.0001);
+    }
+
+    #[test]
+    fn estimate_cost_for_unknown_model() {
+        assert!(estimate_cost("unknown-model", 1000, 1000).is_none());
+    }
+
+    #[test]
+    fn estimate_cost_with_provider_prefix() {
+        let cost = estimate_cost("openrouter/anthropic/claude-3.7-sonnet", 1000, 500).unwrap();
+        assert!(cost > 0.0);
+    }
+}

src-tauri/src/runtime/zeroclaw/mod.rs

@@ -1,4 +1,5 @@
 pub mod adapter;
+pub mod cost;
 pub mod install_adapter;
 pub mod process;
 pub mod sanitize;

src-tauri/src/runtime/zeroclaw/process.rs

@@ -266,6 +266,43 @@ pub fn get_zeroclaw_usage_stats() -> ZeroclawUsageStats {
     usage_store().lock().map(|stats| *stats).unwrap_or_default()
 }
 
+// ---------------------------------------------------------------------------
+// Per-session usage tracking
+// ---------------------------------------------------------------------------
+
+fn session_usage_store() -> &'static Mutex<std::collections::HashMap<String, ZeroclawUsageStats>> {
+    static STORE: OnceLock<Mutex<std::collections::HashMap<String, ZeroclawUsageStats>>> =
+        OnceLock::new();
+    STORE.get_or_init(|| Mutex::new(std::collections::HashMap::new()))
+}
+
+pub fn record_session_usage(session_id: &str, prompt_tokens: u64, completion_tokens: u64) {
+    if session_id.is_empty() {
+        return;
+    }
+    if let Ok(mut map) = session_usage_store().lock() {
+        let stats = map
+            .entry(session_id.to_string())
+            .or_insert_with(ZeroclawUsageStats::default);
+        stats.total_calls = stats.total_calls.saturating_add(1);
+        stats.usage_calls = stats.usage_calls.saturating_add(1);
+        stats.prompt_tokens = stats.prompt_tokens.saturating_add(prompt_tokens);
+        stats.completion_tokens = stats.completion_tokens.saturating_add(completion_tokens);
+        stats.total_tokens = stats
+            .total_tokens
+            .saturating_add(prompt_tokens.saturating_add(completion_tokens));
+        stats.last_updated_ms = now_ms();
+    }
+}
+
+pub fn get_session_usage(session_id: &str) -> ZeroclawUsageStats {
+    session_usage_store()
+        .lock()
+        .ok()
+        .and_then(|map| map.get(session_id).copied())
+        .unwrap_or_default()
+}
+
 fn sanitize_instance_namespace(raw: &str) -> String {
     let trimmed = raw.trim();
     if trimmed.is_empty() {
@@ -804,7 +841,9 @@ pub fn run_zeroclaw_message(
         "-m".to_string(),
         message,
     ];
-    let preferred_model = crate::commands::load_zeroclaw_model_preference();
+    // Per-session model override takes priority over global preference.
+    let preferred_model = crate::commands::preferences::lookup_session_model_override(instance_id)
+        .or_else(|| crate::commands::load_zeroclaw_model_preference());
     let provider_order = provider_order_for_runtime(&env_pairs, preferred_model.as_deref());
     if provider_order.is_empty() {
         return Err(
@@ -821,7 +860,13 @@ pub fn run_zeroclaw_message(
         let stdout = sanitize_output(&String::from_utf8_lossy(&output.stdout));
         let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
         record_zeroclaw_usage(&stdout, &stderr);
-        if parse_usage_from_text(&stdout).is_none() && parse_usage_from_text(&stderr).is_none() {
+        // Also record per-session usage.
+        let session_usage =
+            parse_usage_from_text(&stdout).or_else(|| parse_usage_from_text(&stderr));
+        if let Some((prompt, completion, _total)) = session_usage {
+            record_session_usage(instance_id, prompt, completion);
+        }
+        if session_usage.is_none() {
             if let Ok(mut stats) = usage_store().lock() {
                 if let Some((prompt, completion, total)) =
                     read_usage_from_builtin_traces(&cmd, &cfg, &env_pairs)
@@ -831,6 +876,8 @@ pub fn run_zeroclaw_message(
                     stats.completion_tokens = stats.completion_tokens.saturating_add(completion);
                     stats.total_tokens = stats.total_tokens.saturating_add(total);
                     stats.last_updated_ms = now_ms();
+                    // Record per-session usage from traces as well.
+                    record_session_usage(instance_id, prompt, completion);
                 }
             }
         }