diff --git a/src-tauri/src/bin/codex_monitor_daemon.rs b/src-tauri/src/bin/codex_monitor_daemon.rs index 3a2cd0ce3..6a863d7c5 100644 --- a/src-tauri/src/bin/codex_monitor_daemon.rs +++ b/src-tauri/src/bin/codex_monitor_daemon.rs @@ -82,13 +82,13 @@ use shared::process_core::kill_child_process_tree; use shared::prompts_core::{self, CustomPromptEntry}; use shared::{ agents_config_core, codex_aux_core, codex_core, files_core, git_core, git_ui_core, - local_usage_core, settings_core, workspaces_core, worktree_core, + local_usage_core, settings_core, thread_usage_core, workspaces_core, worktree_core, }; use storage::{read_settings, read_workspaces}; use types::{ AppSettings, GitCommitDiff, GitFileDiff, GitHubIssuesResponse, GitHubPullRequestComment, - GitHubPullRequestDiff, GitHubPullRequestsResponse, GitLogResponse, LocalUsageSnapshot, - WorkspaceEntry, WorkspaceInfo, WorkspaceSettings, WorktreeSetupStatus, + GitHubPullRequestDiff, GitHubPullRequestsResponse, GitLogResponse, LocalThreadUsageSnapshot, + LocalUsageSnapshot, WorkspaceEntry, WorkspaceInfo, WorkspaceSettings, WorktreeSetupStatus, }; use workspace_settings::apply_workspace_settings_update; @@ -1325,6 +1325,19 @@ impl DaemonState { local_usage_core::local_usage_snapshot_core(&self.workspaces, days, workspace_path).await } + async fn local_thread_usage_snapshot( + &self, + thread_ids: Vec, + workspace_path: Option, + ) -> Result { + thread_usage_core::local_thread_usage_snapshot_core( + &self.workspaces, + thread_ids, + workspace_path, + ) + .await + } + async fn menu_set_accelerators(&self, _updates: Vec) -> Result<(), String> { // Daemon has no native menu runtime; treat as no-op for remote parity. Ok(()) @@ -1754,6 +1767,27 @@ mod tests { }); } + #[test] + fn rpc_local_thread_usage_snapshot_returns_snapshot_shape() { + run_async_test(async { + let tmp = make_temp_dir("rpc-local-thread-usage"); + let state = test_state(&tmp); + + let result = rpc::handle_rpc_request( + &state, + "local_thread_usage_snapshot", + json!({ "threadIds": ["thread-1"] }), + "daemon-test".to_string(), + ) + .await + .expect("local_thread_usage_snapshot should succeed"); + + assert!(result.get("updatedAt").and_then(Value::as_i64).is_some()); + assert!(result.get("usageByThread").and_then(Value::as_object).is_some()); + let _ = std::fs::remove_dir_all(&tmp); + }); + } + #[test] fn rpc_daemon_info_reports_identity() { run_async_test(async { diff --git a/src-tauri/src/bin/codex_monitor_daemon/rpc/workspace.rs b/src-tauri/src/bin/codex_monitor_daemon/rpc/workspace.rs index 093a6baf7..c11f77c16 100644 --- a/src-tauri/src/bin/codex_monitor_daemon/rpc/workspace.rs +++ b/src-tauri/src/bin/codex_monitor_daemon/rpc/workspace.rs @@ -266,6 +266,14 @@ pub(super) async fn try_handle( let workspace_path = parse_optional_string(params, "workspacePath"); Some(serialize_result(state.local_usage_snapshot(days, workspace_path)).await) } + "local_thread_usage_snapshot" => { + let thread_ids = parse_optional_string_array(params, "threadIds").unwrap_or_default(); + let workspace_path = parse_optional_string(params, "workspacePath"); + Some( + serialize_result(state.local_thread_usage_snapshot(thread_ids, workspace_path)) + .await, + ) + } _ => None, } } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 318aa5f94..650912155 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -15,6 +15,7 @@ mod files; mod git; mod git_utils; mod local_usage; +mod local_thread_usage; #[cfg(desktop)] mod menu; #[cfg(not(desktop))] @@ -286,6 +287,7 @@ pub fn run() { dictation::dictation_stop, dictation::dictation_cancel, local_usage::local_usage_snapshot, + local_thread_usage::local_thread_usage_snapshot, notifications::is_macos_debug_build, notifications::app_build_type, notifications::send_notification_fallback, diff --git a/src-tauri/src/local_thread_usage.rs b/src-tauri/src/local_thread_usage.rs new file mode 100644 index 000000000..2d1529241 --- /dev/null +++ b/src-tauri/src/local_thread_usage.rs @@ -0,0 +1,33 @@ +use serde_json::json; +use tauri::{AppHandle, State}; + +use crate::remote_backend; +use crate::shared::thread_usage_core; +use crate::state::AppState; +use crate::types::LocalThreadUsageSnapshot; + +#[tauri::command] +pub(crate) async fn local_thread_usage_snapshot( + thread_ids: Vec, + workspace_path: Option, + state: State<'_, AppState>, + app: AppHandle, +) -> Result { + if remote_backend::is_remote_mode(&*state).await { + let response = remote_backend::call_remote( + &*state, + app, + "local_thread_usage_snapshot", + json!({ "threadIds": thread_ids, "workspacePath": workspace_path }), + ) + .await?; + return serde_json::from_value(response).map_err(|err| err.to_string()); + } + + thread_usage_core::local_thread_usage_snapshot_core( + &state.workspaces, + thread_ids, + workspace_path, + ) + .await +} diff --git a/src-tauri/src/remote_backend/mod.rs b/src-tauri/src/remote_backend/mod.rs index 4ae3868d9..ab569f46a 100644 --- a/src-tauri/src/remote_backend/mod.rs +++ b/src-tauri/src/remote_backend/mod.rs @@ -170,6 +170,7 @@ fn can_retry_after_disconnect(method: &str) -> bool { | "list_mcp_server_status" | "list_threads" | "local_usage_snapshot" + | "local_thread_usage_snapshot" | "list_workspace_files" | "list_workspaces" | "model_list" @@ -266,6 +267,7 @@ mod tests { assert!(can_retry_after_disconnect("resume_thread")); assert!(can_retry_after_disconnect("list_threads")); assert!(can_retry_after_disconnect("local_usage_snapshot")); + assert!(can_retry_after_disconnect("local_thread_usage_snapshot")); assert!(!can_retry_after_disconnect("send_user_message")); assert!(!can_retry_after_disconnect("start_thread")); assert!(!can_retry_after_disconnect("remove_workspace")); diff --git a/src-tauri/src/shared/mod.rs b/src-tauri/src/shared/mod.rs index 87e639fbf..d51e79573 100644 --- a/src-tauri/src/shared/mod.rs +++ b/src-tauri/src/shared/mod.rs @@ -9,6 +9,7 @@ pub(crate) mod git_core; pub(crate) mod git_rpc; pub(crate) mod git_ui_core; pub(crate) mod local_usage_core; +pub(crate) mod thread_usage_core; pub(crate) mod process_core; pub(crate) mod prompts_core; pub(crate) mod settings_core; diff --git a/src-tauri/src/shared/thread_usage_core.rs b/src-tauri/src/shared/thread_usage_core.rs new file mode 100644 index 000000000..c60dcf8b5 --- /dev/null +++ b/src-tauri/src/shared/thread_usage_core.rs @@ -0,0 +1,915 @@ +use chrono::DateTime; +use ignore::WalkBuilder; +use serde_json::Value; +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; +use std::sync::{Mutex as StdMutex, OnceLock}; +use std::time::{SystemTime, UNIX_EPOCH}; +use tokio::sync::Mutex; + +use crate::codex::home::{resolve_default_codex_home, resolve_workspace_codex_home}; +use crate::types::{ + LocalThreadUsageSnapshot, ThreadTokenUsageBreakdown, ThreadTokenUsageSnapshot, WorkspaceEntry, +}; + +#[derive(Default, Clone, Copy)] +struct UsageValues { + total_tokens: i64, + input_tokens: i64, + cached_input_tokens: i64, + output_tokens: i64, + reasoning_output_tokens: i64, +} + +impl UsageValues { + fn add_assign(&mut self, other: UsageValues) { + self.total_tokens += other.total_tokens; + self.input_tokens += other.input_tokens; + self.cached_input_tokens += other.cached_input_tokens; + self.output_tokens += other.output_tokens; + self.reasoning_output_tokens += other.reasoning_output_tokens; + } + + fn saturating_delta(self, previous: UsageValues) -> UsageValues { + UsageValues { + total_tokens: (self.total_tokens - previous.total_tokens).max(0), + input_tokens: (self.input_tokens - previous.input_tokens).max(0), + cached_input_tokens: (self.cached_input_tokens - previous.cached_input_tokens).max(0), + output_tokens: (self.output_tokens - previous.output_tokens).max(0), + reasoning_output_tokens: (self.reasoning_output_tokens + - previous.reasoning_output_tokens) + .max(0), + } + } + + fn from_map(map: &serde_json::Map) -> UsageValues { + let input_tokens = read_i64(map, &["input_tokens", "inputTokens"]); + let cached_input_tokens = read_i64( + map, + &[ + "cached_input_tokens", + "cache_read_input_tokens", + "cachedInputTokens", + "cacheReadInputTokens", + ], + ); + let output_tokens = read_i64(map, &["output_tokens", "outputTokens"]); + let reasoning_output_tokens = + read_i64(map, &["reasoning_output_tokens", "reasoningOutputTokens"]); + let total_tokens = map + .get("total_tokens") + .or_else(|| map.get("totalTokens")) + .and_then(|value| { + value + .as_i64() + .or_else(|| value.as_f64().map(|value| value as i64)) + }) + .unwrap_or_else(|| input_tokens + output_tokens); + UsageValues { + total_tokens, + input_tokens, + cached_input_tokens, + output_tokens, + reasoning_output_tokens, + } + } + + fn to_breakdown(self) -> ThreadTokenUsageBreakdown { + ThreadTokenUsageBreakdown { + total_tokens: self.total_tokens, + input_tokens: self.input_tokens, + cached_input_tokens: self.cached_input_tokens, + output_tokens: self.output_tokens, + reasoning_output_tokens: self.reasoning_output_tokens, + } + } + + fn is_zero(self) -> bool { + self.total_tokens == 0 + && self.input_tokens == 0 + && self.cached_input_tokens == 0 + && self.output_tokens == 0 + && self.reasoning_output_tokens == 0 + } +} + +#[derive(Default)] +struct ThreadUsageAggregate { + total: UsageValues, + last: UsageValues, + model_context_window: Option, + latest_timestamp_ms: i64, +} + +impl ThreadUsageAggregate { + fn absorb(&mut self, update: ThreadUsageUpdate) { + self.total.add_assign(update.total_delta); + if update.timestamp_ms >= self.latest_timestamp_ms { + self.latest_timestamp_ms = update.timestamp_ms; + self.last = update.last; + if update.model_context_window.is_some() { + self.model_context_window = update.model_context_window; + } + } + } + + fn into_snapshot(self) -> ThreadTokenUsageSnapshot { + ThreadTokenUsageSnapshot { + total: self.total.to_breakdown(), + last: self.last.to_breakdown(), + model_context_window: self.model_context_window, + } + } +} + +#[derive(Default)] +struct ThreadUsageUpdate { + total_delta: UsageValues, + last: UsageValues, + model_context_window: Option, + timestamp_ms: i64, +} + +// Bound full index refreshes so repeated snapshot calls do not walk the full sessions tree. +const SESSION_INDEX_REFRESH_INTERVAL_MS: i64 = 30_000; + +static SESSION_FILE_INDEX: OnceLock> = OnceLock::new(); + +#[derive(Default)] +struct SessionFileIndex { + by_root: HashMap, +} + +#[derive(Default)] +struct RootSessionIndex { + by_thread: HashMap>, + missing_thread_checked_at: HashMap, + last_full_scan_ms: i64, +} + +#[derive(Clone)] +struct IndexedSessionFile { + path: PathBuf, + cwd: Option, +} + +pub(crate) async fn local_thread_usage_snapshot_core( + workspaces: &Mutex>, + thread_ids: Vec, + workspace_path: Option, +) -> Result { + let thread_ids = sanitize_thread_ids(thread_ids); + let updated_at = now_timestamp_ms(); + if thread_ids.is_empty() { + return Ok(LocalThreadUsageSnapshot { + updated_at, + usage_by_thread: HashMap::new(), + }); + } + + let workspace_path = workspace_path.and_then(|value| { + let trimmed = value.trim(); + if trimmed.is_empty() { + None + } else { + Some(PathBuf::from(trimmed)) + } + }); + + let sessions_roots = { + let workspaces = workspaces.lock().await; + resolve_sessions_roots(&workspaces, workspace_path.as_deref()) + }; + + let usage_by_thread = tokio::task::spawn_blocking(move || { + scan_thread_usage(&thread_ids, workspace_path.as_deref(), &sessions_roots) + }) + .await + .map_err(|err| err.to_string())??; + + Ok(LocalThreadUsageSnapshot { + updated_at, + usage_by_thread, + }) +} + +fn sanitize_thread_ids(thread_ids: Vec) -> Vec { + let mut seen = HashSet::new(); + thread_ids + .into_iter() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .filter(|value| seen.insert(value.clone())) + .collect() +} + +fn scan_thread_usage( + thread_ids: &[String], + workspace_path: Option<&Path>, + sessions_roots: &[PathBuf], +) -> Result, String> { + let requested: HashSet = thread_ids.iter().cloned().collect(); + if requested.is_empty() || sessions_roots.is_empty() { + return Ok(HashMap::new()); + } + + let candidates = collect_candidate_session_files(&requested, workspace_path, sessions_roots); + + let mut aggregate_by_thread: HashMap = HashMap::new(); + for candidate in candidates { + if let Some((thread_id, usage)) = + scan_session_file(&candidate.path, &requested, workspace_path)? + { + aggregate_by_thread + .entry(thread_id) + .or_default() + .absorb(usage); + } + } + + let usage_by_thread = aggregate_by_thread + .into_iter() + .map(|(thread_id, aggregate)| (thread_id, aggregate.into_snapshot())) + .collect(); + + Ok(usage_by_thread) +} + +fn collect_candidate_session_files( + requested_ids: &HashSet, + workspace_path: Option<&Path>, + sessions_roots: &[PathBuf], +) -> Vec { + if requested_ids.is_empty() || sessions_roots.is_empty() { + return Vec::new(); + } + + let now_ms = now_timestamp_ms(); + let mut seen_paths: HashSet = HashSet::new(); + let mut candidates = Vec::new(); + + let index_lock = SESSION_FILE_INDEX.get_or_init(|| StdMutex::new(SessionFileIndex::default())); + let mut index = match index_lock.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + }; + + for root in sessions_roots { + if !root.exists() { + continue; + } + + let root_index = index.by_root.entry(root.clone()).or_default(); + let mut refreshed = false; + + if root_index.last_full_scan_ms == 0 { + rebuild_root_session_index(root_index, root); + refreshed = true; + } + + let missing_requested: Vec<&String> = requested_ids + .iter() + .filter(|thread_id| !root_index.by_thread.contains_key(*thread_id)) + .collect(); + let missing_refresh_due = missing_requested.iter().any(|thread_id| { + match root_index.missing_thread_checked_at.get(*thread_id) { + None => true, + Some(last_checked) => { + now_ms.saturating_sub(*last_checked) >= SESSION_INDEX_REFRESH_INTERVAL_MS + } + } + }); + + if missing_refresh_due && !refreshed { + rebuild_root_session_index(root_index, root); + refreshed = true; + } + + let refresh_due = now_ms.saturating_sub(root_index.last_full_scan_ms) + >= SESSION_INDEX_REFRESH_INTERVAL_MS; + + if refresh_due && !refreshed { + rebuild_root_session_index(root_index, root); + refreshed = true; + } + + if refreshed { + let known_present: HashSet = root_index.by_thread.keys().cloned().collect(); + root_index + .missing_thread_checked_at + .retain(|thread_id, _| !known_present.contains(thread_id)); + } + + for thread_id in requested_ids { + if let Some(files) = root_index.by_thread.get(thread_id) { + root_index.missing_thread_checked_at.remove(thread_id); + for file in files { + if let Some(workspace_path) = workspace_path { + if let Some(cwd) = file.cwd.as_deref() { + if !path_matches_workspace(cwd, workspace_path) { + continue; + } + } + } + if seen_paths.insert(file.path.clone()) { + candidates.push(file.clone()); + } + } + } else if refreshed { + root_index + .missing_thread_checked_at + .insert(thread_id.clone(), now_ms); + } else { + root_index + .missing_thread_checked_at + .entry(thread_id.clone()) + .or_insert(now_ms); + } + } + } + + candidates +} + +fn rebuild_root_session_index(root_index: &mut RootSessionIndex, root: &Path) { + let mut by_thread: HashMap> = HashMap::new(); + + let walker = WalkBuilder::new(root) + .hidden(false) + .follow_links(false) + .require_git(false) + .build(); + + for entry in walker { + let Ok(entry) = entry else { + continue; + }; + let path = entry.path(); + if !entry + .file_type() + .is_some_and(|file_type| file_type.is_file()) + { + continue; + } + if path.extension().and_then(|ext| ext.to_str()) != Some("jsonl") { + continue; + } + + let Some((thread_id, cwd)) = read_session_index_metadata(path) else { + continue; + }; + + by_thread + .entry(thread_id) + .or_default() + .push(IndexedSessionFile { + path: path.to_path_buf(), + cwd, + }); + } + + root_index.by_thread = by_thread; + root_index.last_full_scan_ms = now_timestamp_ms(); +} + +fn read_session_index_metadata(path: &Path) -> Option<(String, Option)> { + let file = File::open(path).ok()?; + let reader = BufReader::new(file); + + for (line_index, line) in reader.lines().enumerate() { + let line = match line { + Ok(line) => line, + Err(_) => continue, + }; + if line.len() > 512_000 { + continue; + } + let value = match serde_json::from_str::(&line) { + Ok(value) => value, + Err(_) => continue, + }; + let Some(thread_id) = extract_session_thread_id(&value) else { + if line_index > 32 { + return None; + } + continue; + }; + return Some((thread_id, extract_cwd(&value))); + } + + None +} + +fn scan_session_file( + path: &Path, + requested_ids: &HashSet, + workspace_path: Option<&Path>, +) -> Result, String> { + let file = match File::open(path) { + Ok(file) => file, + Err(_) => return Ok(None), + }; + + let reader = BufReader::new(file); + let mut matched_thread_id: Option = None; + let mut match_known = false; + + let mut workspace_match_known = workspace_path.is_none(); + let mut matches_workspace = workspace_path.is_none(); + + let mut previous_totals: UsageValues = UsageValues::default(); + let mut update = ThreadUsageUpdate::default(); + + for (line_index, line) in reader.lines().enumerate() { + let line = match line { + Ok(line) => line, + Err(_) => continue, + }; + if line.len() > 512_000 { + continue; + } + let value = match serde_json::from_str::(&line) { + Ok(value) => value, + Err(_) => continue, + }; + + if !match_known { + if let Some(thread_id) = extract_session_thread_id(&value) { + match_known = true; + if requested_ids.contains(&thread_id) { + matched_thread_id = Some(thread_id); + } else { + return Ok(None); + } + } else if line_index > 32 { + // Session metadata should appear near the top; if not, skip unknown files. + return Ok(None); + } + } + + if workspace_path.is_some() { + let entry_type = value + .get("type") + .and_then(|value| value.as_str()) + .unwrap_or(""); + if (entry_type == "session_meta" || entry_type == "turn_context") + && !workspace_match_known + { + if let Some(cwd) = extract_cwd(&value) { + workspace_match_known = true; + matches_workspace = workspace_path + .map(|workspace_path| path_matches_workspace(&cwd, workspace_path)) + .unwrap_or(true); + if !matches_workspace { + return Ok(None); + } + } + } + } + + if !match_known || matched_thread_id.is_none() { + continue; + } + if !workspace_match_known || !matches_workspace { + continue; + } + + let Some(token_info) = extract_token_usage(&value) else { + continue; + }; + + let mut total_delta = UsageValues::default(); + let mut next_last = UsageValues::default(); + + if let Some(total_usage) = token_info.total { + total_delta = total_usage.saturating_delta(previous_totals); + previous_totals = total_usage; + next_last = token_info.last.unwrap_or(total_delta); + } else if let Some(last_usage) = token_info.last { + total_delta = last_usage; + next_last = last_usage; + // Keep cumulative totals in sync so future total snapshots compute correct deltas. + previous_totals.add_assign(last_usage); + } + + if !total_delta.is_zero() { + update.total_delta.add_assign(total_delta); + } + if !next_last.is_zero() { + update.last = next_last; + } + if token_info.model_context_window.is_some() { + update.model_context_window = token_info.model_context_window; + } + let timestamp_ms = token_info.timestamp_ms; + if timestamp_ms > update.timestamp_ms { + update.timestamp_ms = timestamp_ms; + } + } + + let Some(thread_id) = matched_thread_id else { + return Ok(None); + }; + + if update.total_delta.is_zero() && update.last.is_zero() { + return Ok(None); + } + + Ok(Some((thread_id, update))) +} + +fn now_timestamp_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as i64 +} + +fn extract_session_thread_id(value: &Value) -> Option { + let entry_type = value.get("type")?.as_str()?; + if entry_type != "session_meta" { + return None; + } + value + .get("payload") + .and_then(|payload| payload.get("id")) + .and_then(|id| id.as_str()) + .map(|id| id.to_string()) +} + +struct TokenUsageInfo { + total: Option, + last: Option, + model_context_window: Option, + timestamp_ms: i64, +} + +fn extract_token_usage(value: &Value) -> Option { + let entry_type = value + .get("type") + .and_then(|value| value.as_str()) + .unwrap_or(""); + if entry_type != "event_msg" && !entry_type.is_empty() { + return None; + } + + let payload = value.get("payload")?.as_object()?; + let payload_type = payload + .get("type") + .and_then(|value| value.as_str()) + .unwrap_or(""); + if payload_type != "token_count" { + return None; + } + + let info = payload.get("info")?.as_object()?; + let total = + find_usage_map(info, &["total_token_usage", "totalTokenUsage"]).map(UsageValues::from_map); + let last = + find_usage_map(info, &["last_token_usage", "lastTokenUsage"]).map(UsageValues::from_map); + + if total.is_none() && last.is_none() { + return None; + } + + let model_context_window = info + .get("model_context_window") + .or_else(|| info.get("modelContextWindow")) + .and_then(|value| { + value + .as_i64() + .or_else(|| value.as_f64().map(|value| value as i64)) + }); + + let timestamp_ms = read_timestamp_ms(value).unwrap_or_default(); + + Some(TokenUsageInfo { + total, + last, + model_context_window, + timestamp_ms, + }) +} + +fn find_usage_map<'a>( + info: &'a serde_json::Map, + keys: &[&str], +) -> Option<&'a serde_json::Map> { + keys.iter() + .find_map(|key| info.get(*key).and_then(|value| value.as_object())) +} + +fn read_i64(map: &serde_json::Map, keys: &[&str]) -> i64 { + keys.iter() + .find_map(|key| map.get(*key)) + .and_then(|value| { + value + .as_i64() + .or_else(|| value.as_f64().map(|value| value as i64)) + }) + .unwrap_or(0) +} + +fn read_timestamp_ms(value: &Value) -> Option { + let raw = value.get("timestamp")?; + if let Some(text) = raw.as_str() { + return DateTime::parse_from_rfc3339(text) + .map(|value| value.timestamp_millis()) + .ok(); + } + let numeric = raw + .as_i64() + .or_else(|| raw.as_f64().map(|value| value as i64))?; + if numeric > 0 && numeric < 1_000_000_000_000 { + return Some(numeric * 1000); + } + Some(numeric) +} + +fn extract_cwd(value: &Value) -> Option { + value + .get("payload") + .and_then(|payload| payload.get("cwd")) + .and_then(|cwd| cwd.as_str()) + .map(|cwd| cwd.to_string()) +} + +fn path_matches_workspace(cwd: &str, workspace_path: &Path) -> bool { + let cwd_path = Path::new(cwd); + cwd_path == workspace_path || cwd_path.starts_with(workspace_path) +} + +fn resolve_codex_sessions_root(codex_home_override: Option) -> Option { + codex_home_override + .or_else(resolve_default_codex_home) + .map(|home| home.join("sessions")) +} + +fn resolve_sessions_roots( + workspaces: &HashMap, + workspace_path: Option<&Path>, +) -> Vec { + if let Some(workspace_path) = workspace_path { + let codex_home_override = + resolve_workspace_codex_home_for_path(workspaces, Some(workspace_path)); + return resolve_codex_sessions_root(codex_home_override) + .into_iter() + .collect(); + } + + let mut roots = Vec::new(); + let mut seen = HashSet::new(); + + if let Some(root) = resolve_codex_sessions_root(None) { + if seen.insert(root.clone()) { + roots.push(root); + } + } + + for entry in workspaces.values() { + let parent_entry = entry + .parent_id + .as_ref() + .and_then(|parent_id| workspaces.get(parent_id)); + let Some(codex_home) = resolve_workspace_codex_home(entry, parent_entry) else { + continue; + }; + if let Some(root) = resolve_codex_sessions_root(Some(codex_home)) { + if seen.insert(root.clone()) { + roots.push(root); + } + } + } + + roots +} + +fn resolve_workspace_codex_home_for_path( + workspaces: &HashMap, + workspace_path: Option<&Path>, +) -> Option { + let workspace_path = workspace_path?; + let entry = workspaces + .values() + .filter(|entry| { + let entry_path = Path::new(&entry.path); + workspace_path == entry_path || workspace_path.starts_with(entry_path) + }) + .max_by_key(|entry| entry.path.len())?; + + let parent_entry = entry + .parent_id + .as_ref() + .and_then(|parent_id| workspaces.get(parent_id)); + + resolve_workspace_codex_home(entry, parent_entry) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use std::sync::OnceLock; + use std::time::Duration; + use uuid::Uuid; + + fn session_index_test_guard() -> std::sync::MutexGuard<'static, ()> { + static TEST_GUARD: OnceLock> = OnceLock::new(); + match TEST_GUARD.get_or_init(|| StdMutex::new(())).lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } + } + + fn clear_session_index_cache() { + let Some(index_lock) = SESSION_FILE_INDEX.get() else { + return; + }; + match index_lock.lock() { + Ok(mut index) => index.by_root.clear(), + Err(poisoned) => poisoned.into_inner().by_root.clear(), + } + } + + fn last_full_scan_ms_for_root(root: &Path) -> Option { + let index_lock = SESSION_FILE_INDEX.get()?; + let index = match index_lock.lock() { + Ok(index) => index, + Err(poisoned) => poisoned.into_inner(), + }; + index.by_root.get(root).map(|root_index| root_index.last_full_scan_ms) + } + + fn make_temp_sessions_root() -> PathBuf { + let mut root = std::env::temp_dir(); + root.push(format!("codexmonitor-thread-usage-root-{}", Uuid::new_v4())); + std::fs::create_dir_all(&root).expect("create temp root"); + root + } + + fn write_session_file(root: &Path, file_name: &str, lines: &[&str]) -> PathBuf { + let day_dir = root.join("2026").join("02").join("01"); + std::fs::create_dir_all(&day_dir).expect("create day dir"); + let path = day_dir.join(file_name); + let mut file = File::create(&path).expect("create session file"); + for line in lines { + writeln!(file, "{line}").expect("write line"); + } + path + } + + #[test] + fn scan_thread_usage_aggregates_total_without_double_counting_last() { + let _guard = session_index_test_guard(); + clear_session_index_cache(); + + let root = make_temp_sessions_root(); + let thread_id = "thread-abc"; + write_session_file( + &root, + "rollout-2026-02-01-thread-abc.jsonl", + &[ + r#"{"timestamp":"2026-02-01T10:00:00.000Z","type":"session_meta","payload":{"id":"thread-abc","cwd":"/tmp/project"}}"#, + r#"{"timestamp":"2026-02-01T10:00:01.000Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":2,"total_tokens":15}}}}"#, + r#"{"timestamp":"2026-02-01T10:00:02.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":2,"total_tokens":15},"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":2,"total_tokens":15},"model_context_window":200000}}}"#, + r#"{"timestamp":"2026-02-01T10:00:03.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":12,"cached_input_tokens":0,"output_tokens":6,"reasoning_output_tokens":2,"total_tokens":18},"last_token_usage":{"input_tokens":2,"cached_input_tokens":0,"output_tokens":1,"reasoning_output_tokens":0,"total_tokens":3}}}}"#, + ], + ); + + let usage = scan_thread_usage(&[thread_id.to_string()], None, &[root]).expect("scan usage"); + let thread_usage = usage.get(thread_id).expect("thread usage"); + + assert_eq!(thread_usage.total.input_tokens, 12); + assert_eq!(thread_usage.total.output_tokens, 6); + assert_eq!(thread_usage.total.total_tokens, 18); + assert_eq!(thread_usage.last.input_tokens, 2); + assert_eq!(thread_usage.last.output_tokens, 1); + assert_eq!(thread_usage.model_context_window, Some(200000)); + } + + #[test] + fn scan_thread_usage_respects_workspace_filter() { + let _guard = session_index_test_guard(); + clear_session_index_cache(); + + let root = make_temp_sessions_root(); + write_session_file( + &root, + "rollout-2026-02-01-thread-mismatch.jsonl", + &[ + r#"{"timestamp":"2026-02-01T10:00:00.000Z","type":"session_meta","payload":{"id":"thread-mismatch","cwd":"/tmp/other"}}"#, + r#"{"timestamp":"2026-02-01T10:00:01.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"total_tokens":15}}}}"#, + ], + ); + + let usage = scan_thread_usage( + &["thread-mismatch".to_string()], + Some(Path::new("/tmp/project")), + &[root], + ) + .expect("scan usage"); + + assert!(usage.is_empty()); + } + + #[test] + fn scan_thread_usage_uses_session_meta_id_not_filename_substring() { + let _guard = session_index_test_guard(); + clear_session_index_cache(); + + let root = make_temp_sessions_root(); + write_session_file( + &root, + "rollout-2026-02-01-thread-12.jsonl", + &[ + r#"{"timestamp":"2026-02-01T10:00:00.000Z","type":"session_meta","payload":{"id":"thread-12","cwd":"/tmp/project"}}"#, + r#"{"timestamp":"2026-02-01T10:00:01.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"total_tokens":15}}}}"#, + ], + ); + + let usage = + scan_thread_usage(&["thread-1".to_string()], None, &[root]).expect("scan usage"); + + assert!(usage.is_empty()); + } + + #[test] + fn scan_thread_usage_refreshes_index_when_requested_thread_is_missing() { + let _guard = session_index_test_guard(); + clear_session_index_cache(); + + let root = make_temp_sessions_root(); + write_session_file( + &root, + "rollout-2026-02-01-thread-a.jsonl", + &[ + r#"{"timestamp":"2026-02-01T10:00:00.000Z","type":"session_meta","payload":{"id":"thread-a","cwd":"/tmp/project"}}"#, + r#"{"timestamp":"2026-02-01T10:00:01.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"total_tokens":15}}}}"#, + ], + ); + + let initial_usage = + scan_thread_usage(&["thread-a".to_string()], None, std::slice::from_ref(&root)) + .expect("scan usage"); + assert_eq!( + initial_usage + .get("thread-a") + .expect("thread-a usage") + .total + .total_tokens, + 15 + ); + + write_session_file( + &root, + "rollout-2026-02-01-thread-b.jsonl", + &[ + r#"{"timestamp":"2026-02-01T11:00:00.000Z","type":"session_meta","payload":{"id":"thread-b","cwd":"/tmp/project"}}"#, + r#"{"timestamp":"2026-02-01T11:00:01.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":4,"cached_input_tokens":0,"output_tokens":2,"total_tokens":6}}}}"#, + ], + ); + + let usage = + scan_thread_usage(&["thread-b".to_string()], None, std::slice::from_ref(&root)) + .expect("scan usage"); + assert_eq!( + usage.get("thread-b").expect("thread-b usage").total.total_tokens, + 6 + ); + } + + #[test] + fn scan_thread_usage_throttles_repeated_missing_thread_rebuilds() { + let _guard = session_index_test_guard(); + clear_session_index_cache(); + + let root = make_temp_sessions_root(); + write_session_file( + &root, + "rollout-2026-02-01-thread-a.jsonl", + &[ + r#"{"timestamp":"2026-02-01T10:00:00.000Z","type":"session_meta","payload":{"id":"thread-a","cwd":"/tmp/project"}}"#, + r#"{"timestamp":"2026-02-01T10:00:01.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"total_tokens":15}}}}"#, + ], + ); + + let missing_thread_id = "thread-missing".to_string(); + let first_usage = + scan_thread_usage(std::slice::from_ref(&missing_thread_id), None, std::slice::from_ref(&root)) + .expect("scan usage"); + assert!(first_usage.is_empty()); + + let first_scan_ms = last_full_scan_ms_for_root(&root).expect("first scan timestamp"); + std::thread::sleep(Duration::from_millis(5)); + + let second_usage = + scan_thread_usage(std::slice::from_ref(&missing_thread_id), None, std::slice::from_ref(&root)) + .expect("scan usage"); + assert!(second_usage.is_empty()); + + let second_scan_ms = last_full_scan_ms_for_root(&root).expect("second scan timestamp"); + assert_eq!(second_scan_ms, first_scan_ms); + } +} diff --git a/src-tauri/src/types.rs b/src-tauri/src/types.rs index cbb8afbaa..36f9dacb4 100644 --- a/src-tauri/src/types.rs +++ b/src-tauri/src/types.rs @@ -1,4 +1,5 @@ use serde::{Deserialize, Serialize}; +use std::collections::HashMap; #[derive(Debug, Serialize, Deserialize, Clone)] pub(crate) struct GitFileStatus { @@ -187,6 +188,32 @@ pub(crate) struct LocalUsageSnapshot { pub(crate) top_models: Vec, } +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ThreadTokenUsageBreakdown { + pub(crate) total_tokens: i64, + pub(crate) input_tokens: i64, + pub(crate) cached_input_tokens: i64, + pub(crate) output_tokens: i64, + pub(crate) reasoning_output_tokens: i64, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ThreadTokenUsageSnapshot { + pub(crate) total: ThreadTokenUsageBreakdown, + pub(crate) last: ThreadTokenUsageBreakdown, + pub(crate) model_context_window: Option, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(rename_all = "camelCase")] +pub(crate) struct LocalThreadUsageSnapshot { + pub(crate) updated_at: i64, + #[serde(default)] + pub(crate) usage_by_thread: HashMap, +} + #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub(crate) enum TcpDaemonState { @@ -492,6 +519,21 @@ pub(crate) struct AppSettings { rename = "usageShowRemaining" )] pub(crate) usage_show_remaining: bool, + #[serde( + default = "default_show_thread_token_usage", + rename = "showThreadTokenUsage" + )] + pub(crate) show_thread_token_usage: bool, + #[serde( + default = "default_thread_token_usage_show_full", + rename = "threadTokenUsageShowFull" + )] + pub(crate) thread_token_usage_show_full: bool, + #[serde( + default = "default_thread_token_usage_exclude_cache", + rename = "threadTokenUsageExcludeCache" + )] + pub(crate) thread_token_usage_exclude_cache: bool, #[serde( default = "default_show_message_file_path", rename = "showMessageFilePath" @@ -702,6 +744,18 @@ fn default_usage_show_remaining() -> bool { false } +fn default_show_thread_token_usage() -> bool { + true +} + +fn default_thread_token_usage_show_full() -> bool { + true +} + +fn default_thread_token_usage_exclude_cache() -> bool { + true +} + fn default_show_message_file_path() -> bool { true } @@ -1143,6 +1197,9 @@ impl Default for AppSettings { ui_scale: 1.0, theme: default_theme(), usage_show_remaining: default_usage_show_remaining(), + show_thread_token_usage: default_show_thread_token_usage(), + thread_token_usage_show_full: default_thread_token_usage_show_full(), + thread_token_usage_exclude_cache: default_thread_token_usage_exclude_cache(), show_message_file_path: default_show_message_file_path(), chat_history_scrollback_items: default_chat_history_scrollback_items(), thread_title_autogeneration_enabled: false, @@ -1307,6 +1364,9 @@ mod tests { assert!((settings.ui_scale - 1.0).abs() < f64::EPSILON); assert_eq!(settings.theme, "system"); assert!(!settings.usage_show_remaining); + assert!(settings.show_thread_token_usage); + assert!(settings.thread_token_usage_show_full); + assert!(settings.thread_token_usage_exclude_cache); assert!(settings.show_message_file_path); assert_eq!(settings.chat_history_scrollback_items, Some(200)); assert!(!settings.thread_title_autogeneration_enabled); diff --git a/src/App.tsx b/src/App.tsx index e9ac6d8b5..8804b5f5a 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -113,6 +113,13 @@ import { useWorkspaceLaunchScripts } from "@app/hooks/useWorkspaceLaunchScripts" import { useWorktreeSetupScript } from "@app/hooks/useWorktreeSetupScript"; import { useGitCommitController } from "@app/hooks/useGitCommitController"; import { effectiveCommitMessageModelId } from "@/features/git/utils/commitMessageModelSelection"; +import { formatCompactTokenCount } from "@utils/tokenUsage"; +import { + estimateThreadUsageCost, + formatUsageCostLabel, + mergeUsageCostSummaries, + type UsageCostSummary, +} from "@app/utils/usageCost"; import { WorkspaceHome } from "@/features/workspaces/components/WorkspaceHome"; import { MobileServerSetupWizard } from "@/features/mobile/components/MobileServerSetupWizard"; import { useMobileServerSetup } from "@/features/mobile/hooks/useMobileServerSetup"; @@ -120,6 +127,7 @@ import { useWorkspaceHome } from "@/features/workspaces/hooks/useWorkspaceHome"; import { useWorkspaceAgentMd } from "@/features/workspaces/hooks/useWorkspaceAgentMd"; import type { ComposerEditorSettings, + ThreadTokenUsage, WorkspaceInfo, } from "@/types"; import { computePlanFollowupState } from "@/features/messages/utils/messageRenderUtils"; @@ -153,6 +161,7 @@ import { useAppShellOrchestration } from "@app/orchestration/useLayoutOrchestrat import { buildCodexArgsOptions } from "@threads/utils/codexArgsProfiles"; import { normalizeCodexArgsInput } from "@/utils/codexArgsInput"; import { + NO_THREAD_SCOPE_SUFFIX, resolveWorkspaceRuntimeCodexArgsBadgeLabel, resolveWorkspaceRuntimeCodexArgsOverride, } from "@threads/utils/threadCodexParamsSeed"; @@ -176,6 +185,27 @@ const GitHubPanelData = lazy(() => })), ); +function normalizeThreadModelId(value: string | null | undefined): string | null { + if (typeof value !== "string") { + return null; + } + const trimmed = value.trim(); + if (!trimmed) { + return null; + } + const normalized = trimmed.toLowerCase(); + if ( + normalized === "unknown" || + normalized === "default" || + normalized === "auto" || + normalized === "current" || + normalized === "inherit" + ) { + return null; + } + return trimmed; +} + function MainApp() { const { appSettings, @@ -535,9 +565,7 @@ function MainApp() { return; } const modelId = - typeof metadata.modelId === "string" && metadata.modelId.trim().length > 0 - ? metadata.modelId.trim() - : null; + normalizeThreadModelId(metadata.modelId); const effort = typeof metadata.effort === "string" && metadata.effort.trim().length > 0 ? metadata.effort.trim().toLowerCase() @@ -551,7 +579,8 @@ function MainApp() { modelId?: string | null; effort?: string | null; } = {}; - if (modelId && !current?.modelId) { + const currentModelId = normalizeThreadModelId(current?.modelId); + if (modelId && !currentModelId) { patch.modelId = modelId; } if (effort && !current?.effort) { @@ -617,6 +646,7 @@ function MainApp() { threadListCursorByWorkspace, activeTurnIdByThread, tokenUsageByThread, + tokenUsageThreadIdsByWorkspace, rateLimitsByWorkspace, accountByWorkspace, planByThread, @@ -687,6 +717,7 @@ function MainApp() { customPrompts: prompts, onMessageActivity: handleThreadMessageActivity, threadSortKey: threadListSortKey, + enableBackgroundThreadMetadataHydration: true, onThreadCodexMetadataDetected: handleThreadCodexMetadataDetected, }); const { connectionState: remoteThreadConnectionState, reconnectLive } = @@ -1365,6 +1396,164 @@ function MainApp() { const activeTokenUsage = activeThreadId ? tokenUsageByThread[activeThreadId] ?? null : null; + const fullTokenCountFormatter = useMemo(() => new Intl.NumberFormat(), []); + const formatFullTokenCount = useCallback((value: number) => { + if (!Number.isFinite(value) || value <= 0) { + return null; + } + return fullTokenCountFormatter.format(Math.round(value)); + }, [fullTokenCountFormatter]); + const threadModelIdByWorkspace = useMemo(() => { + const byWorkspace: Record> = {}; + Object.entries(threadsByWorkspace).forEach(([workspaceId, threads]) => { + const byThread: Record = {}; + threads.forEach((thread) => { + const modelId = thread.modelId?.trim(); + if (modelId) { + byThread[thread.id] = modelId; + } + }); + byWorkspace[workspaceId] = byThread; + }); + return byWorkspace; + }, [threadsByWorkspace]); + const resolveThreadModelId = useCallback( + (workspaceId: string, threadId: string) => { + const fromList = normalizeThreadModelId( + threadModelIdByWorkspace[workspaceId]?.[threadId], + ); + if (fromList) { + return fromList; + } + const storedModelId = normalizeThreadModelId( + getThreadCodexParams(workspaceId, threadId)?.modelId, + ); + if (storedModelId) { + return storedModelId; + } + const workspaceDefaultModelId = normalizeThreadModelId( + getThreadCodexParams(workspaceId, NO_THREAD_SCOPE_SUFFIX)?.modelId, + ); + if (workspaceDefaultModelId) { + return workspaceDefaultModelId; + } + return normalizeThreadModelId(appSettings.lastComposerModelId); + }, + [appSettings.lastComposerModelId, getThreadCodexParams, threadModelIdByWorkspace], + ); + const getDisplayThreadTokenUsageTotal = useCallback( + (usage: ThreadTokenUsage | null | undefined) => { + const totalTokens = usage?.total.totalTokens ?? 0; + if (!appSettings.threadTokenUsageExcludeCache) { + return totalTokens; + } + const cachedInputTokens = usage?.total.cachedInputTokens ?? 0; + return Math.max(0, totalTokens - cachedInputTokens); + }, + [appSettings.threadTokenUsageExcludeCache], + ); + const getThreadTokenUsageLabel = useCallback( + (_workspaceId: string, threadId: string) => { + if (!appSettings.showThreadTokenUsage) { + return null; + } + const totalTokens = getDisplayThreadTokenUsageTotal(tokenUsageByThread[threadId]); + const formattedTotal = appSettings.threadTokenUsageShowFull + ? formatFullTokenCount(totalTokens) + : formatCompactTokenCount(totalTokens); + if (!formattedTotal) { + return null; + } + const costSummary = estimateThreadUsageCost( + tokenUsageByThread[threadId], + resolveThreadModelId(_workspaceId, threadId), + { excludeCache: appSettings.threadTokenUsageExcludeCache }, + ); + const costLabel = formatUsageCostLabel( + costSummary, + !appSettings.threadTokenUsageShowFull, + ); + return costLabel + ? `${formattedTotal} tokens · ${costLabel}` + : `${formattedTotal} tokens`; + }, + [ + appSettings.showThreadTokenUsage, + appSettings.threadTokenUsageExcludeCache, + appSettings.threadTokenUsageShowFull, + formatFullTokenCount, + getDisplayThreadTokenUsageTotal, + resolveThreadModelId, + tokenUsageByThread, + ], + ); + const tokenUsageSummaryByWorkspace = useMemo(() => { + const totals: Record = {}; + const workspaceIds = new Set([ + ...Object.keys(threadsByWorkspace), + ...Object.keys(tokenUsageThreadIdsByWorkspace), + ]); + workspaceIds.forEach((workspaceId) => { + const threadIds = new Set([ + ...(tokenUsageThreadIdsByWorkspace[workspaceId] ?? []), + ...(threadsByWorkspace[workspaceId] ?? []).map((thread) => thread.id), + ]); + const threadIdList = Array.from(threadIds); + const total = threadIdList.reduce( + (sum, threadId) => sum + getDisplayThreadTokenUsageTotal(tokenUsageByThread[threadId]), + 0, + ); + const cost = mergeUsageCostSummaries( + threadIdList.map((threadId) => + estimateThreadUsageCost( + tokenUsageByThread[threadId], + resolveThreadModelId(workspaceId, threadId), + { excludeCache: appSettings.threadTokenUsageExcludeCache }, + ), + ), + ); + totals[workspaceId] = { + tokens: total, + cost, + }; + }); + return totals; + }, [ + appSettings.threadTokenUsageExcludeCache, + getDisplayThreadTokenUsageTotal, + resolveThreadModelId, + threadsByWorkspace, + tokenUsageByThread, + tokenUsageThreadIdsByWorkspace, + ]); + const getWorkspaceTokenUsageLabel = useCallback( + (workspaceId: string) => { + if (!appSettings.showThreadTokenUsage) { + return null; + } + const usageSummary = tokenUsageSummaryByWorkspace[workspaceId]; + const totalTokens = usageSummary?.tokens ?? 0; + const formattedTotal = appSettings.threadTokenUsageShowFull + ? formatFullTokenCount(totalTokens) + : formatCompactTokenCount(totalTokens); + if (!formattedTotal) { + return null; + } + const costLabel = formatUsageCostLabel( + usageSummary?.cost ?? { knownUsd: 0, unknownTokens: 0, totalTokens: 0 }, + !appSettings.threadTokenUsageShowFull, + ); + return costLabel + ? `${formattedTotal} tokens · ${costLabel}` + : `${formattedTotal} tokens`; + }, + [ + appSettings.showThreadTokenUsage, + appSettings.threadTokenUsageShowFull, + formatFullTokenCount, + tokenUsageSummaryByWorkspace, + ], + ); const activePlan = activeThreadId ? planByThread[activeThreadId] ?? null : null; @@ -2495,6 +2684,8 @@ function MainApp() { onWorkspaceDragLeave: handleWorkspaceDragLeave, onWorkspaceDrop: handleWorkspaceDrop, getThreadArgsBadge, + getThreadTokenUsageLabel, + getWorkspaceTokenUsageLabel, }); const gitRootOverride = activeWorkspace?.settings.gitRoot; diff --git a/src/features/app/components/PinnedThreadList.test.tsx b/src/features/app/components/PinnedThreadList.test.tsx index 704f91bfa..4b476a4d3 100644 --- a/src/features/app/components/PinnedThreadList.test.tsx +++ b/src/features/app/components/PinnedThreadList.test.tsx @@ -121,4 +121,15 @@ describe("PinnedThreadList", () => { expect(row?.querySelector(".thread-status")?.className).toContain("unread"); expect(row?.querySelector(".thread-status")?.className).not.toContain("processing"); }); + + it("renders a token usage label when provided", () => { + render( + "900 tokens"} + />, + ); + + expect(screen.getByText("900 tokens")).toBeTruthy(); + }); }); diff --git a/src/features/app/components/PinnedThreadList.tsx b/src/features/app/components/PinnedThreadList.tsx index 5def5ee6d..1f816c531 100644 --- a/src/features/app/components/PinnedThreadList.tsx +++ b/src/features/app/components/PinnedThreadList.tsx @@ -19,6 +19,7 @@ type PinnedThreadListProps = { getWorkspaceLabel?: (workspaceId: string) => string | null; getThreadTime: (thread: ThreadSummary) => string | null; getThreadArgsBadge?: (workspaceId: string, threadId: string) => string | null; + getThreadTokenUsageLabel?: (workspaceId: string, threadId: string) => string | null; isThreadPinned: (workspaceId: string, threadId: string) => boolean; onSelectThread: (workspaceId: string, threadId: string) => void; onShowThreadMenu: ( @@ -38,6 +39,7 @@ export function PinnedThreadList({ getWorkspaceLabel, getThreadTime, getThreadArgsBadge, + getThreadTokenUsageLabel, isThreadPinned, onSelectThread, onShowThreadMenu, @@ -59,6 +61,7 @@ export function PinnedThreadList({ workspaceLabel={getWorkspaceLabel?.(workspaceId) ?? null} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} isThreadPinned={isThreadPinned} onSelectThread={onSelectThread} onShowThreadMenu={onShowThreadMenu} diff --git a/src/features/app/components/Sidebar.test.tsx b/src/features/app/components/Sidebar.test.tsx index 893e447bb..bd89bef60 100644 --- a/src/features/app/components/Sidebar.test.tsx +++ b/src/features/app/components/Sidebar.test.tsx @@ -560,4 +560,54 @@ describe("Sidebar", () => { const indicator = screen.queryByTitle("Streaming updates in progress"); expect(indicator).toBeNull(); }); + + it("renders per-project token usage labels for workspace and worktree cards", () => { + render( + + workspaceId === "ws-1" ? "12.4K tokens" : "3.1K tokens" + } + />, + ); + + expect(screen.getByText("12.4K tokens")).toBeTruthy(); + expect(screen.getByText("3.1K tokens")).toBeTruthy(); + }); }); diff --git a/src/features/app/components/Sidebar.tsx b/src/features/app/components/Sidebar.tsx index 53ce98a06..2083bfed8 100644 --- a/src/features/app/components/Sidebar.tsx +++ b/src/features/app/components/Sidebar.tsx @@ -103,6 +103,8 @@ type SidebarProps = { isThreadPinned: (workspaceId: string, threadId: string) => boolean; getPinTimestamp: (workspaceId: string, threadId: string) => number | null; getThreadArgsBadge?: (workspaceId: string, threadId: string) => string | null; + getThreadTokenUsageLabel?: (workspaceId: string, threadId: string) => string | null; + getWorkspaceTokenUsageLabel?: (workspaceId: string) => string | null; onRenameThread: (workspaceId: string, threadId: string) => void; onDeleteWorkspace: (workspaceId: string) => void; onDeleteWorktree: (workspaceId: string) => void; @@ -164,6 +166,8 @@ export const Sidebar = memo(function Sidebar({ isThreadPinned, getPinTimestamp, getThreadArgsBadge, + getThreadTokenUsageLabel, + getWorkspaceTokenUsageLabel, onRenameThread, onDeleteWorkspace, onDeleteWorktree, @@ -870,6 +874,7 @@ export const Sidebar = memo(function Sidebar({ pendingUserInputKeys={pendingUserInputKeys} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} isThreadPinned={isThreadPinned} onSelectThread={onSelectThread} onShowThreadMenu={showThreadMenu} @@ -903,6 +908,7 @@ export const Sidebar = memo(function Sidebar({ pendingUserInputKeys={pendingUserInputKeys} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} isThreadPinned={isThreadPinned} onSelectThread={onSelectThread} onShowThreadMenu={showThreadMenu} @@ -1002,6 +1008,9 @@ export const Sidebar = memo(function Sidebar({ key={entry.id} workspace={entry} workspaceName={renderHighlightedName(entry.name)} + workspaceTokenUsageLabel={ + getWorkspaceTokenUsageLabel?.(entry.id) ?? null + } isActive={entry.id === activeWorkspaceId} isCollapsed={isCollapsed} addMenuOpen={addMenuOpen} @@ -1094,6 +1103,8 @@ export const Sidebar = memo(function Sidebar({ getThreadRows={getThreadRows} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} + getWorkspaceTokenUsageLabel={getWorkspaceTokenUsageLabel} isThreadPinned={isThreadPinned} getPinTimestamp={getPinTimestamp} pinnedThreadsVersion={pinnedThreadsVersion} @@ -1128,6 +1139,8 @@ export const Sidebar = memo(function Sidebar({ getThreadRows={getThreadRows} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} + getWorkspaceTokenUsageLabel={getWorkspaceTokenUsageLabel} isThreadPinned={isThreadPinned} getPinTimestamp={getPinTimestamp} pinnedThreadsVersion={pinnedThreadsVersion} @@ -1156,6 +1169,7 @@ export const Sidebar = memo(function Sidebar({ pendingUserInputKeys={pendingUserInputKeys} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} isThreadPinned={isThreadPinned} onToggleExpanded={handleToggleExpanded} onLoadOlderThreads={onLoadOlderThreads} diff --git a/src/features/app/components/ThreadList.test.tsx b/src/features/app/components/ThreadList.test.tsx index 0451cab03..647e71fe7 100644 --- a/src/features/app/components/ThreadList.test.tsx +++ b/src/features/app/components/ThreadList.test.tsx @@ -152,4 +152,15 @@ describe("ThreadList", () => { expect(row?.querySelector(".thread-status")?.className).toContain("unread"); expect(row?.querySelector(".thread-status")?.className).not.toContain("processing"); }); + + it("renders a token usage label when provided", () => { + render( + "1.2k tokens"} + />, + ); + + expect(screen.getByText("1.2k tokens")).toBeTruthy(); + }); }); diff --git a/src/features/app/components/ThreadList.tsx b/src/features/app/components/ThreadList.tsx index 732b9c65d..7744e892f 100644 --- a/src/features/app/components/ThreadList.tsx +++ b/src/features/app/components/ThreadList.tsx @@ -25,6 +25,7 @@ type ThreadListProps = { pendingUserInputKeys?: Set; getThreadTime: (thread: ThreadSummary) => string | null; getThreadArgsBadge?: (workspaceId: string, threadId: string) => string | null; + getThreadTokenUsageLabel?: (workspaceId: string, threadId: string) => string | null; isThreadPinned: (workspaceId: string, threadId: string) => boolean; onToggleExpanded: (workspaceId: string) => void; onLoadOlderThreads: (workspaceId: string) => void; @@ -53,6 +54,7 @@ export function ThreadList({ pendingUserInputKeys, getThreadTime, getThreadArgsBadge, + getThreadTokenUsageLabel, isThreadPinned, onToggleExpanded, onLoadOlderThreads, @@ -76,6 +78,7 @@ export function ThreadList({ pendingUserInputKeys={pendingUserInputKeys} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} isThreadPinned={isThreadPinned} onSelectThread={onSelectThread} onShowThreadMenu={onShowThreadMenu} @@ -97,6 +100,7 @@ export function ThreadList({ pendingUserInputKeys={pendingUserInputKeys} getThreadTime={getThreadTime} getThreadArgsBadge={getThreadArgsBadge} + getThreadTokenUsageLabel={getThreadTokenUsageLabel} isThreadPinned={isThreadPinned} onSelectThread={onSelectThread} onShowThreadMenu={onShowThreadMenu} diff --git a/src/features/app/components/ThreadRow.tsx b/src/features/app/components/ThreadRow.tsx index 05650582f..430278ab7 100644 --- a/src/features/app/components/ThreadRow.tsx +++ b/src/features/app/components/ThreadRow.tsx @@ -15,6 +15,7 @@ type ThreadRowProps = { workspaceLabel?: string | null; getThreadTime: (thread: ThreadSummary) => string | null; getThreadArgsBadge?: (workspaceId: string, threadId: string) => string | null; + getThreadTokenUsageLabel?: (workspaceId: string, threadId: string) => string | null; isThreadPinned: (workspaceId: string, threadId: string) => boolean; onSelectThread: (workspaceId: string, threadId: string) => void; onShowThreadMenu: ( @@ -37,18 +38,13 @@ export function ThreadRow({ workspaceLabel, getThreadTime, getThreadArgsBadge, + getThreadTokenUsageLabel, isThreadPinned, onSelectThread, onShowThreadMenu, }: ThreadRowProps) { const relativeTime = getThreadTime(thread); const badge = getThreadArgsBadge?.(workspaceId, thread.id) ?? null; - const modelBadge = - thread.modelId && thread.modelId.trim().length > 0 - ? thread.effort && thread.effort.trim().length > 0 - ? `${thread.modelId} · ${thread.effort}` - : thread.modelId - : null; const indentStyle = depth > 0 ? ({ "--thread-indent": `${depth * indentUnit}px` } as CSSProperties) @@ -62,6 +58,13 @@ export function ThreadRow({ ); const canPin = depth === 0; const isPinned = canPin && isThreadPinned(workspaceId, thread.id); + const tokenUsageLabel = getThreadTokenUsageLabel?.(workspaceId, thread.id) ?? null; + const modelBadge = + thread.modelId && thread.modelId.trim().length > 0 + ? thread.effort && thread.effort.trim().length > 0 + ? `${thread.modelId} · ${thread.effort}` + : thread.modelId + : null; return (
{isPinned && 📌} - {thread.name} +
+ {thread.name} + {tokenUsageLabel && ( + {tokenUsageLabel} + )} +
{workspaceLabel && {workspaceLabel}} {modelBadge && ( diff --git a/src/features/app/components/WorkspaceCard.tsx b/src/features/app/components/WorkspaceCard.tsx index 670384ae4..7514fe206 100644 --- a/src/features/app/components/WorkspaceCard.tsx +++ b/src/features/app/components/WorkspaceCard.tsx @@ -5,6 +5,7 @@ import type { WorkspaceInfo } from "../../../types"; type WorkspaceCardProps = { workspace: WorkspaceInfo; workspaceName?: React.ReactNode; + workspaceTokenUsageLabel?: string | null; isActive: boolean; isCollapsed: boolean; addMenuOpen: boolean; @@ -25,6 +26,7 @@ type WorkspaceCardProps = { export function WorkspaceCard({ workspace, workspaceName, + workspaceTokenUsageLabel, isActive, isCollapsed, addMenuOpen, @@ -55,20 +57,25 @@ export function WorkspaceCard({ >
-
- {workspaceName ?? workspace.name} - +
+
+ {workspaceName ?? workspace.name} + +
+ {workspaceTokenUsageLabel && ( +
{workspaceTokenUsageLabel}
+ )}
+
+
+
Show thread token usage & estimated costs
+
+ Display token totals & costs beneath each thread title. (both are estimates and can + differ from the actual values) +
+
+ +
+
+
+
Show full thread token counts & costs
+
+ Display exact totals instead of compact abbreviations. +
+
+ +
+
+
+
Exclude cache from thread token usage
+
+ Subtract cached input tokens from each thread total. +
+
+ +
Split chat and diff center panes
diff --git a/src/features/settings/hooks/useAppSettings.test.ts b/src/features/settings/hooks/useAppSettings.test.ts index c6eb4221e..f9530887b 100644 --- a/src/features/settings/hooks/useAppSettings.test.ts +++ b/src/features/settings/hooks/useAppSettings.test.ts @@ -55,6 +55,7 @@ describe("useAppSettings", () => { expect(result.current.settings.personality).toBe("friendly"); expect(result.current.settings.backendMode).toBe("remote"); expect(result.current.settings.remoteBackendHost).toBe("example:1234"); + expect(result.current.settings.threadTokenUsageExcludeCache).toBe(true); }); it("keeps defaults when getAppSettings fails", async () => { diff --git a/src/features/settings/hooks/useAppSettings.ts b/src/features/settings/hooks/useAppSettings.ts index e4dafaccc..4738adafd 100644 --- a/src/features/settings/hooks/useAppSettings.ts +++ b/src/features/settings/hooks/useAppSettings.ts @@ -168,6 +168,9 @@ function buildDefaultSettings(): AppSettings { uiScale: UI_SCALE_DEFAULT, theme: "system", usageShowRemaining: false, + showThreadTokenUsage: true, + threadTokenUsageShowFull: true, + threadTokenUsageExcludeCache: true, showMessageFilePath: true, chatHistoryScrollbackItems: CHAT_SCROLLBACK_DEFAULT, threadTitleAutogenerationEnabled: false, @@ -245,6 +248,9 @@ function normalizeAppSettings(settings: AppSettings): AppSettings { codexArgs: settings.codexArgs?.trim() ? settings.codexArgs.trim() : null, uiScale: clampUiScale(settings.uiScale), theme: allowedThemes.has(settings.theme) ? settings.theme : "system", + showThreadTokenUsage: settings.showThreadTokenUsage !== false, + threadTokenUsageShowFull: settings.threadTokenUsageShowFull !== false, + threadTokenUsageExcludeCache: settings.threadTokenUsageExcludeCache !== false, uiFontFamily: normalizeFontFamily( settings.uiFontFamily, DEFAULT_UI_FONT_FAMILY, diff --git a/src/features/threads/hooks/threadReducer/threadLifecycleSlice.ts b/src/features/threads/hooks/threadReducer/threadLifecycleSlice.ts index 8fbcecece..115e1e288 100644 --- a/src/features/threads/hooks/threadReducer/threadLifecycleSlice.ts +++ b/src/features/threads/hooks/threadReducer/threadLifecycleSlice.ts @@ -14,6 +14,20 @@ function statusEquals(previous: ThreadStatus, nextStatus: ThreadStatus) { ); } +function normalizeWorkspaceThreadIds( + threadIds: string[], + hiddenByThreadId: Record = {}, +): string[] { + const unique = new Set(); + threadIds.forEach((threadId) => { + if (!threadId || hiddenByThreadId[threadId]) { + return; + } + unique.add(threadId); + }); + return Array.from(unique); +} + export function reduceThreadLifecycle( state: ThreadState, action: ThreadAction, @@ -52,9 +66,25 @@ export function reduceThreadLifecycle( return state; } const list = state.threadsByWorkspace[action.workspaceId] ?? []; - if (list.some((thread) => thread.id === action.threadId)) { + const usageThreadIds = + state.tokenUsageThreadIdsByWorkspace[action.workspaceId] ?? []; + const hasThread = list.some((thread) => thread.id === action.threadId); + const hasUsageThreadId = usageThreadIds.includes(action.threadId); + if (hasThread && hasUsageThreadId) { return state; } + const nextUsageThreadIds = hasUsageThreadId + ? usageThreadIds + : [action.threadId, ...usageThreadIds]; + if (hasThread) { + return { + ...state, + tokenUsageThreadIdsByWorkspace: { + ...state.tokenUsageThreadIdsByWorkspace, + [action.workspaceId]: nextUsageThreadIds, + }, + }; + } const thread: ThreadSummary = { id: action.threadId, name: "New Agent", @@ -66,6 +96,10 @@ export function reduceThreadLifecycle( ...state.threadsByWorkspace, [action.workspaceId]: [thread, ...list], }, + tokenUsageThreadIdsByWorkspace: { + ...state.tokenUsageThreadIdsByWorkspace, + [action.workspaceId]: nextUsageThreadIds, + }, threadStatusById: { ...state.threadStatusById, [action.threadId]: { @@ -97,6 +131,11 @@ export function reduceThreadLifecycle( const list = state.threadsByWorkspace[action.workspaceId] ?? []; const filtered = list.filter((thread) => thread.id !== action.threadId); + const usageThreadIds = + state.tokenUsageThreadIdsByWorkspace[action.workspaceId] ?? []; + const filteredUsageThreadIds = usageThreadIds.filter( + (threadId) => threadId !== action.threadId, + ); const nextActive = state.activeThreadIdByWorkspace[action.workspaceId] === action.threadId ? filtered[0]?.id ?? null @@ -112,6 +151,10 @@ export function reduceThreadLifecycle( ...state.threadsByWorkspace, [action.workspaceId]: filtered, }, + tokenUsageThreadIdsByWorkspace: { + ...state.tokenUsageThreadIdsByWorkspace, + [action.workspaceId]: filteredUsageThreadIds, + }, activeThreadIdByWorkspace: { ...state.activeThreadIdByWorkspace, [action.workspaceId]: nextActive, @@ -121,6 +164,11 @@ export function reduceThreadLifecycle( case "removeThread": { const list = state.threadsByWorkspace[action.workspaceId] ?? []; const filtered = list.filter((thread) => thread.id !== action.threadId); + const usageThreadIds = + state.tokenUsageThreadIdsByWorkspace[action.workspaceId] ?? []; + const filteredUsageThreadIds = usageThreadIds.filter( + (threadId) => threadId !== action.threadId, + ); const nextActive = state.activeThreadIdByWorkspace[action.workspaceId] === action.threadId ? filtered[0]?.id ?? null @@ -131,18 +179,24 @@ export function reduceThreadLifecycle( const { [action.threadId]: ____, ...restDiffs } = state.turnDiffByThread; const { [action.threadId]: _____, ...restPlans } = state.planByThread; const { [action.threadId]: ______, ...restParents } = state.threadParentById; + const { [action.threadId]: _______, ...restTokenUsage } = state.tokenUsageByThread; return { ...state, threadsByWorkspace: { ...state.threadsByWorkspace, [action.workspaceId]: filtered, }, + tokenUsageThreadIdsByWorkspace: { + ...state.tokenUsageThreadIdsByWorkspace, + [action.workspaceId]: filteredUsageThreadIds, + }, itemsByThread: restItems, threadStatusById: restStatus, activeTurnIdByThread: restTurns, turnDiffByThread: restDiffs, planByThread: restPlans, threadParentById: restParents, + tokenUsageByThread: restTokenUsage, activeThreadIdByWorkspace: { ...state.activeThreadIdByWorkspace, [action.workspaceId]: nextActive, @@ -404,6 +458,25 @@ export function reduceThreadLifecycle( }, }; } + case "setWorkspaceTokenUsageThreadIds": { + const hidden = state.hiddenThreadIdsByWorkspace[action.workspaceId] ?? {}; + const nextThreadIds = normalizeWorkspaceThreadIds(action.threadIds, hidden); + const currentThreadIds = + state.tokenUsageThreadIdsByWorkspace[action.workspaceId] ?? []; + if ( + nextThreadIds.length === currentThreadIds.length && + nextThreadIds.every((threadId, index) => threadId === currentThreadIds[index]) + ) { + return state; + } + return { + ...state, + tokenUsageThreadIdsByWorkspace: { + ...state.tokenUsageThreadIdsByWorkspace, + [action.workspaceId]: nextThreadIds, + }, + }; + } case "setThreadListLoading": return { ...state, diff --git a/src/features/threads/hooks/useThreadActions.test.tsx b/src/features/threads/hooks/useThreadActions.test.tsx index fd492429a..5de097c25 100644 --- a/src/features/threads/hooks/useThreadActions.test.tsx +++ b/src/features/threads/hooks/useThreadActions.test.tsx @@ -1,11 +1,12 @@ // @vitest-environment jsdom -import { act, renderHook } from "@testing-library/react"; +import { act, renderHook, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ConversationItem, WorkspaceInfo } from "@/types"; import { archiveThread, forkThread, listThreads, + localThreadUsageSnapshot, listWorkspaces, resumeThread, startThread, @@ -26,6 +27,7 @@ vi.mock("@services/tauri", () => ({ forkThread: vi.fn(), resumeThread: vi.fn(), listThreads: vi.fn(), + localThreadUsageSnapshot: vi.fn(), listWorkspaces: vi.fn(), archiveThread: vi.fn(), })); @@ -63,6 +65,10 @@ describe("useThreadActions", () => { vi.clearAllMocks(); vi.mocked(listWorkspaces).mockResolvedValue([]); vi.mocked(getThreadCreatedTimestamp).mockReturnValue(0); + vi.mocked(localThreadUsageSnapshot).mockResolvedValue({ + updatedAt: 1, + usageByThread: {}, + }); }); function renderActions( @@ -82,6 +88,7 @@ describe("useThreadActions", () => { dispatch, itemsByThread: {}, threadsByWorkspace: {}, + tokenUsageThreadIdsByWorkspace: {}, activeThreadIdByWorkspace: {}, activeTurnIdByThread: {}, threadParentById: {}, @@ -321,6 +328,63 @@ describe("useThreadActions", () => { }); }); + it("hydrates thread token usage from resume payload when available", async () => { + vi.mocked(resumeThread).mockResolvedValue({ + result: { + thread: { + id: "thread-2", + token_usage: { + total: { + total_tokens: 900, + input_tokens: 600, + cached_input_tokens: 120, + output_tokens: 300, + reasoning_output_tokens: 50, + }, + last: { + total_tokens: 120, + input_tokens: 80, + cached_input_tokens: 10, + output_tokens: 40, + reasoning_output_tokens: 6, + }, + model_context_window: 200000, + }, + }, + }, + }); + vi.mocked(buildItemsFromThread).mockReturnValue([]); + vi.mocked(isReviewingFromThread).mockReturnValue(false); + + const { result, dispatch } = renderActions(); + + await act(async () => { + await result.current.resumeThreadForWorkspace("ws-1", "thread-2"); + }); + + expect(dispatch).toHaveBeenCalledWith({ + type: "setThreadTokenUsage", + threadId: "thread-2", + tokenUsage: { + total: { + totalTokens: 900, + inputTokens: 600, + cachedInputTokens: 120, + outputTokens: 300, + reasoningOutputTokens: 50, + }, + last: { + totalTokens: 120, + inputTokens: 80, + cachedInputTokens: 10, + outputTokens: 40, + reasoningOutputTokens: 6, + }, + modelContextWindow: 200000, + }, + }); + }); + it("links resumed spawn subagent to its parent from thread source", async () => { vi.mocked(resumeThread).mockResolvedValue({ result: { @@ -707,6 +771,28 @@ describe("useThreadActions", () => { nextCursor: "cursor-1", }, }); + vi.mocked(localThreadUsageSnapshot).mockResolvedValueOnce({ + updatedAt: 100, + usageByThread: { + "thread-1": { + total: { + totalTokens: 1200, + inputTokens: 700, + cachedInputTokens: 100, + outputTokens: 500, + reasoningOutputTokens: 80, + }, + last: { + totalTokens: 200, + inputTokens: 120, + cachedInputTokens: 20, + outputTokens: 80, + reasoningOutputTokens: 15, + }, + modelContextWindow: 200000, + }, + }, + }); vi.mocked(getThreadTimestamp).mockImplementation((thread) => { const value = (thread as Record).updated_at as number; return value ?? 0; @@ -752,6 +838,31 @@ describe("useThreadActions", () => { workspaceId: "ws-1", cursor: "cursor-1", }); + expect(localThreadUsageSnapshot).toHaveBeenCalledWith( + ["thread-1"], + "/tmp/codex", + ); + expect(dispatch).toHaveBeenCalledWith({ + type: "setThreadTokenUsage", + threadId: "thread-1", + tokenUsage: { + total: { + totalTokens: 1200, + inputTokens: 700, + cachedInputTokens: 100, + outputTokens: 500, + reasoningOutputTokens: 80, + }, + last: { + totalTokens: 200, + inputTokens: 120, + cachedInputTokens: 20, + outputTokens: 80, + reasoningOutputTokens: 15, + }, + modelContextWindow: 200000, + }, + }); expect(saveThreadActivity).toHaveBeenCalledWith({ "ws-1": { "thread-1": 5000 }, }); @@ -760,6 +871,128 @@ describe("useThreadActions", () => { }); }); + it("merges out-of-order thread usage hydration responses across thread sets", async () => { + let resolveFirstHydration: + | ((value: Record) => void) + | null = null; + let resolveSecondHydration: + | ((value: Record) => void) + | null = null; + const firstHydrationPromise = new Promise>((resolve) => { + resolveFirstHydration = resolve; + }); + const secondHydrationPromise = new Promise>((resolve) => { + resolveSecondHydration = resolve; + }); + + vi.mocked(listThreads) + .mockResolvedValueOnce({ + result: { + data: [ + { + id: "thread-a", + cwd: "/tmp/codex", + preview: "Thread A", + updated_at: 3000, + }, + ], + nextCursor: null, + }, + }) + .mockResolvedValueOnce({ + result: { + data: [ + { + id: "thread-b", + cwd: "/tmp/codex", + preview: "Thread B", + updated_at: 4000, + }, + ], + nextCursor: null, + }, + }); + vi.mocked(localThreadUsageSnapshot) + .mockReturnValueOnce(firstHydrationPromise as Promise) + .mockReturnValueOnce(secondHydrationPromise as Promise); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + + const { result, dispatch } = renderActions(); + + await act(async () => { + await result.current.listThreadsForWorkspace(workspace); + await result.current.listThreadsForWorkspace(workspace); + }); + + await act(async () => { + resolveSecondHydration?.({ + updatedAt: 200, + usageByThread: { + "thread-b": { + total: { + totalTokens: 2000, + inputTokens: 1200, + cachedInputTokens: 100, + outputTokens: 800, + reasoningOutputTokens: 80, + }, + last: { + totalTokens: 300, + inputTokens: 180, + cachedInputTokens: 20, + outputTokens: 120, + reasoningOutputTokens: 12, + }, + modelContextWindow: 200000, + }, + }, + }); + await secondHydrationPromise; + }); + + await act(async () => { + resolveFirstHydration?.({ + updatedAt: 100, + usageByThread: { + "thread-a": { + total: { + totalTokens: 1500, + inputTokens: 900, + cachedInputTokens: 120, + outputTokens: 600, + reasoningOutputTokens: 60, + }, + last: { + totalTokens: 250, + inputTokens: 150, + cachedInputTokens: 30, + outputTokens: 100, + reasoningOutputTokens: 10, + }, + modelContextWindow: 200000, + }, + }, + }); + await firstHydrationPromise; + }); + + expect(dispatch).toHaveBeenCalledWith( + expect.objectContaining({ + type: "setThreadTokenUsage", + threadId: "thread-a", + }), + ); + expect(dispatch).toHaveBeenCalledWith( + expect.objectContaining({ + type: "setThreadTokenUsage", + threadId: "thread-b", + }), + ); + }); + it("uses fresh fetched data for active anchors outside top thread target", async () => { const data = Array.from({ length: 21 }, (_, index) => ({ id: `thread-${index + 1}`, @@ -867,6 +1100,45 @@ describe("useThreadActions", () => { }); }); + it("tracks workspace token usage ids from the full fetched thread set", async () => { + const threads = Array.from({ length: 25 }, (_, index) => ({ + id: `thread-${index + 1}`, + cwd: "/tmp/codex", + preview: `Thread ${index + 1}`, + updated_at: 5000 - index, + })); + vi.mocked(listThreads).mockResolvedValue({ + result: { + data: threads, + nextCursor: null, + }, + }); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + + const { result, dispatch } = renderActions(); + + await act(async () => { + await result.current.listThreadsForWorkspace(workspace); + }); + + expect(dispatch).toHaveBeenCalledWith({ + type: "setWorkspaceTokenUsageThreadIds", + workspaceId: "ws-1", + threadIds: threads.map((thread) => thread.id), + }); + const setThreadsCall = dispatch.mock.calls.find( + ([action]) => action?.type === "setThreads" && action?.workspaceId === "ws-1", + ); + expect(setThreadsCall?.[0]?.threads).toHaveLength(20); + expect(localThreadUsageSnapshot).toHaveBeenCalledWith( + threads.map((thread) => thread.id), + "/tmp/codex", + ); + }); + it("assigns shared-root threads to a single target workspace when listing multiple workspaces", async () => { const workspaceAlias: WorkspaceInfo = { ...workspaceTwo, @@ -1374,6 +1646,28 @@ describe("useThreadActions", () => { nextCursor: null, }, }); + vi.mocked(localThreadUsageSnapshot).mockResolvedValueOnce({ + updatedAt: 120, + usageByThread: { + "thread-2": { + total: { + totalTokens: 400, + inputTokens: 250, + cachedInputTokens: 30, + outputTokens: 150, + reasoningOutputTokens: 10, + }, + last: { + totalTokens: 60, + inputTokens: 40, + cachedInputTokens: 5, + outputTokens: 20, + reasoningOutputTokens: 2, + }, + modelContextWindow: null, + }, + }, + }); vi.mocked(getThreadTimestamp).mockImplementation((thread) => { const value = (thread as Record).updated_at as number; return value ?? 0; @@ -1409,6 +1703,10 @@ describe("useThreadActions", () => { workspaceId: "ws-1", cursor: null, }); + expect(localThreadUsageSnapshot).toHaveBeenCalledWith( + ["thread-2"], + "/tmp/codex", + ); }); it("supports snake_case next_cursor when loading older threads", async () => { @@ -1711,6 +2009,362 @@ describe("useThreadActions", () => { ); }); + it("prefetches model metadata in background after listing threads", async () => { + vi.mocked(listThreads).mockResolvedValue({ + result: { + data: [ + { + id: "thread-bg-1", + cwd: "/tmp/codex", + preview: "Background model", + updated_at: 5000, + }, + ], + nextCursor: null, + }, + }); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + vi.mocked(resumeThread).mockResolvedValue({ + result: { + thread: { + id: "thread-bg-1", + model: "gpt-5.2-codex", + reasoning_effort: "high", + }, + }, + }); + + const onThreadCodexMetadataDetected = vi.fn(); + const { result } = renderActions({ + enableBackgroundMetadataHydration: true, + onThreadCodexMetadataDetected, + }); + + await act(async () => { + await result.current.listThreadsForWorkspace(workspace); + }); + + await waitFor(() => { + expect(resumeThread).toHaveBeenCalledWith("ws-1", "thread-bg-1"); + }); + await waitFor(() => { + expect(onThreadCodexMetadataDetected).toHaveBeenCalledWith( + "ws-1", + "thread-bg-1", + { modelId: "gpt-5.2-codex", effort: "high" }, + ); + }); + }); + + it("does not re-prefetch model-less metadata after first background hydration", async () => { + vi.mocked(listThreads).mockResolvedValue({ + result: { + data: [ + { + id: "thread-bg-repeat-1", + cwd: "/tmp/codex", + preview: "Background model-less thread", + updated_at: 5000, + }, + ], + nextCursor: null, + }, + }); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + vi.mocked(resumeThread).mockResolvedValue({ + result: { + thread: { + id: "thread-bg-repeat-1", + preview: "Background model-less thread", + }, + }, + }); + + const onThreadCodexMetadataDetected = vi.fn(); + const { result } = renderActions({ + enableBackgroundMetadataHydration: true, + onThreadCodexMetadataDetected, + }); + + await act(async () => { + await result.current.listThreadsForWorkspace(workspace); + }); + + await waitFor(() => { + expect(resumeThread).toHaveBeenCalledTimes(1); + }); + + await act(async () => { + await result.current.listThreadsForWorkspace(workspace); + }); + + expect(resumeThread).toHaveBeenCalledTimes(1); + expect(onThreadCodexMetadataDetected).not.toHaveBeenCalled(); + }); + + it("does not overwrite snapshot token usage during background metadata hydration", async () => { + vi.mocked(listThreads).mockResolvedValue({ + result: { + data: [ + { + id: "thread-bg-usage-1", + cwd: "/tmp/codex", + preview: "Background usage thread", + updated_at: 5100, + }, + ], + nextCursor: null, + }, + }); + vi.mocked(localThreadUsageSnapshot).mockResolvedValue({ + updatedAt: 500, + usageByThread: { + "thread-bg-usage-1": { + total: { + totalTokens: 1000, + inputTokens: 650, + cachedInputTokens: 100, + outputTokens: 350, + reasoningOutputTokens: 25, + }, + last: { + totalTokens: 100, + inputTokens: 65, + cachedInputTokens: 10, + outputTokens: 35, + reasoningOutputTokens: 2, + }, + modelContextWindow: 200000, + }, + }, + }); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + vi.mocked(resumeThread).mockResolvedValue({ + result: { + thread: { + id: "thread-bg-usage-1", + model: "gpt-5.2-codex", + token_usage: { + total: { + total_tokens: 9999, + input_tokens: 6000, + cached_input_tokens: 1200, + output_tokens: 3999, + reasoning_output_tokens: 100, + }, + last: { + total_tokens: 999, + input_tokens: 600, + cached_input_tokens: 120, + output_tokens: 399, + reasoning_output_tokens: 10, + }, + model_context_window: 200000, + }, + }, + }, + }); + + const { result, dispatch } = renderActions({ + enableBackgroundMetadataHydration: true, + onThreadCodexMetadataDetected: vi.fn(), + }); + + await act(async () => { + await result.current.listThreadsForWorkspace(workspace); + }); + + await waitFor(() => { + expect(resumeThread).toHaveBeenCalledWith("ws-1", "thread-bg-usage-1"); + }); + + const tokenUsageActions = dispatch.mock.calls + .map(([action]) => action) + .filter( + (action) => + action.type === "setThreadTokenUsage" && + action.threadId === "thread-bg-usage-1", + ); + + expect(tokenUsageActions).toHaveLength(1); + expect(tokenUsageActions[0]).toEqual( + expect.objectContaining({ + type: "setThreadTokenUsage", + threadId: "thread-bg-usage-1", + tokenUsage: expect.objectContaining({ + total: expect.objectContaining({ + totalTokens: 1000, + }), + }), + }), + ); + }); + + it("prefetches metadata for all listed threads across multiple hydration batches", async () => { + const listedThreads = Array.from({ length: 10 }, (_, index) => ({ + id: `thread-bg-batch-${index + 1}`, + cwd: "/tmp/codex", + preview: `Background ${index + 1}`, + updated_at: 6_000 - index, + })); + vi.mocked(listThreads).mockResolvedValue({ + result: { + data: listedThreads, + nextCursor: null, + }, + }); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + vi.mocked(resumeThread).mockImplementation(async (_workspaceId, threadId) => ({ + result: { + thread: { + id: threadId, + model: "gpt-5.2-codex", + }, + }, + })); + + const onThreadCodexMetadataDetected = vi.fn(); + const { result } = renderActions({ + enableBackgroundMetadataHydration: true, + onThreadCodexMetadataDetected, + }); + + await act(async () => { + await result.current.listThreadsForWorkspace(workspace); + }); + + await waitFor(() => { + expect(resumeThread).toHaveBeenCalledTimes(10); + }); + listedThreads.forEach((thread) => { + expect(resumeThread).toHaveBeenCalledWith("ws-1", thread.id); + expect(onThreadCodexMetadataDetected).toHaveBeenCalledWith( + "ws-1", + thread.id, + { modelId: "gpt-5.2-codex", effort: null }, + ); + }); + }); + + it("prefetches model metadata in background for older-page additions", async () => { + vi.mocked(listThreads).mockResolvedValue({ + result: { + data: [ + { + id: "thread-older-bg-1", + cwd: "/tmp/codex", + preview: "Older background model", + updated_at: 4000, + }, + ], + nextCursor: null, + }, + }); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + vi.mocked(resumeThread).mockResolvedValue({ + result: { + thread: { + id: "thread-older-bg-1", + model: "gpt-5.1-codex", + reasoning_effort: "medium", + }, + }, + }); + + const onThreadCodexMetadataDetected = vi.fn(); + const { result } = renderActions({ + enableBackgroundMetadataHydration: true, + onThreadCodexMetadataDetected, + threadsByWorkspace: { + "ws-1": [{ id: "thread-1", name: "Agent 1", updatedAt: 6000 }], + }, + threadListCursorByWorkspace: { "ws-1": "cursor-1" }, + }); + + await act(async () => { + await result.current.loadOlderThreadsForWorkspace(workspace); + }); + + await waitFor(() => { + expect(resumeThread).toHaveBeenCalledWith("ws-1", "thread-older-bg-1"); + }); + await waitFor(() => { + expect(onThreadCodexMetadataDetected).toHaveBeenCalledWith( + "ws-1", + "thread-older-bg-1", + { modelId: "gpt-5.1-codex", effort: "medium" }, + ); + }); + }); + + it("prefetches metadata for all older-page additions across batches", async () => { + const olderThreads = Array.from({ length: 11 }, (_, index) => ({ + id: `thread-older-bg-batch-${index + 1}`, + cwd: "/tmp/codex", + preview: `Older background ${index + 1}`, + updated_at: 4_500 - index, + })); + vi.mocked(listThreads).mockResolvedValue({ + result: { + data: olderThreads, + nextCursor: null, + }, + }); + vi.mocked(getThreadTimestamp).mockImplementation((thread) => { + const value = (thread as Record).updated_at as number; + return value ?? 0; + }); + vi.mocked(resumeThread).mockImplementation(async (_workspaceId, threadId) => ({ + result: { + thread: { + id: threadId, + model: "gpt-5.1-codex", + }, + }, + })); + + const onThreadCodexMetadataDetected = vi.fn(); + const { result } = renderActions({ + enableBackgroundMetadataHydration: true, + onThreadCodexMetadataDetected, + threadsByWorkspace: { + "ws-1": [{ id: "thread-1", name: "Agent 1", updatedAt: 6_000 }], + }, + threadListCursorByWorkspace: { "ws-1": "cursor-batch-1" }, + }); + + await act(async () => { + await result.current.loadOlderThreadsForWorkspace(workspace); + }); + + await waitFor(() => { + expect(resumeThread).toHaveBeenCalledTimes(11); + }); + olderThreads.forEach((thread) => { + expect(resumeThread).toHaveBeenCalledWith("ws-1", thread.id); + expect(onThreadCodexMetadataDetected).toHaveBeenCalledWith( + "ws-1", + thread.id, + { modelId: "gpt-5.1-codex", effort: null }, + ); + }); + }); + it("archives threads and reports errors", async () => { vi.mocked(archiveThread).mockRejectedValue(new Error("nope")); const onDebug = vi.fn(); diff --git a/src/features/threads/hooks/useThreadActions.ts b/src/features/threads/hooks/useThreadActions.ts index 817bf5b08..a2b413ec2 100644 --- a/src/features/threads/hooks/useThreadActions.ts +++ b/src/features/threads/hooks/useThreadActions.ts @@ -1,36 +1,38 @@ import { useCallback, useRef } from "react"; import type { Dispatch, MutableRefObject } from "react"; import type { - ConversationItem, - DebugEntry, - ThreadListSortKey, - ThreadSummary, - WorkspaceInfo, + ConversationItem, + DebugEntry, + ThreadListSortKey, + ThreadSummary, + WorkspaceInfo, } from "@/types"; import { - archiveThread as archiveThreadService, - forkThread as forkThreadService, - listThreads as listThreadsService, - listWorkspaces as listWorkspacesService, - resumeThread as resumeThreadService, - startThread as startThreadService, + archiveThread as archiveThreadService, + forkThread as forkThreadService, + localThreadUsageSnapshot as localThreadUsageSnapshotService, + listThreads as listThreadsService, + listWorkspaces as listWorkspacesService, + resumeThread as resumeThreadService, + startThread as startThreadService, } from "@services/tauri"; import { - buildItemsFromThread, - getThreadCreatedTimestamp, - getThreadTimestamp, - isReviewingFromThread, - mergeThreadItems, - previewThreadName, + buildItemsFromThread, + getThreadCreatedTimestamp, + getThreadTimestamp, + isReviewingFromThread, + mergeThreadItems, + previewThreadName, } from "@utils/threadItems"; import { extractThreadCodexMetadata } from "@threads/utils/threadCodexMetadata"; import { - asString, - normalizeRootPath, + asString, + normalizeTokenUsage, + normalizeRootPath, } from "@threads/utils/threadNormalize"; import { - getParentThreadIdFromThread, - getResumedTurnState, + getParentThreadIdFromThread, + getResumedTurnState, } from "@threads/utils/threadRpc"; import { saveThreadActivity } from "@threads/utils/threadStorage"; import type { ThreadAction, ThreadState } from "./useThreadsReducer"; @@ -40,998 +42,1262 @@ const THREAD_LIST_PAGE_SIZE = 100; const THREAD_LIST_MAX_PAGES_OLDER = 6; const THREAD_LIST_MAX_PAGES_DEFAULT = 6; const THREAD_LIST_CURSOR_PAGE_START = "__codex_monitor_page_start__"; +const THREAD_METADATA_PREFETCH_BATCH_SIZE = 8; function isWithinWorkspaceRoot(path: string, workspaceRoot: string) { - if (!path || !workspaceRoot) { - return false; - } - return ( - path === workspaceRoot || - (path.length > workspaceRoot.length && - path.startsWith(workspaceRoot) && - path.charCodeAt(workspaceRoot.length) === 47) - ); + if (!path || !workspaceRoot) { + return false; + } + return ( + path === workspaceRoot || + (path.length > workspaceRoot.length && + path.startsWith(workspaceRoot) && + path.charCodeAt(workspaceRoot.length) === 47) + ); } type WorkspacePathLookup = { - workspaceIdsByPath: Record; - workspacePathsSorted: string[]; + workspaceIdsByPath: Record; + workspacePathsSorted: string[]; }; function buildWorkspacePathLookup(workspaces: WorkspaceInfo[]): WorkspacePathLookup { - const workspaceIdsByPath: Record = {}; - const workspacePathsSorted: string[] = []; - workspaces.forEach((workspace) => { - const workspacePath = normalizeRootPath(workspace.path); - if (!workspacePath) { - return; - } - if (!workspaceIdsByPath[workspacePath]) { - workspaceIdsByPath[workspacePath] = []; - workspacePathsSorted.push(workspacePath); - } - workspaceIdsByPath[workspacePath].push(workspace.id); - }); - workspacePathsSorted.sort((a, b) => b.length - a.length); - return { workspaceIdsByPath, workspacePathsSorted }; + const workspaceIdsByPath: Record = {}; + const workspacePathsSorted: string[] = []; + workspaces.forEach((workspace) => { + const workspacePath = normalizeRootPath(workspace.path); + if (!workspacePath) { + return; + } + if (!workspaceIdsByPath[workspacePath]) { + workspaceIdsByPath[workspacePath] = []; + workspacePathsSorted.push(workspacePath); + } + workspaceIdsByPath[workspacePath].push(workspace.id); + }); + workspacePathsSorted.sort((a, b) => b.length - a.length); + return { workspaceIdsByPath, workspacePathsSorted }; } function resolveWorkspaceIdForThreadPath( - path: string, - lookup: WorkspacePathLookup, - allowedWorkspaceIds?: Set, + path: string, + lookup: WorkspacePathLookup, + allowedWorkspaceIds?: Set, ) { - const normalizedPath = normalizeRootPath(path); - if (!normalizedPath) { - return null; - } - const matchedWorkspacePath = lookup.workspacePathsSorted.find((workspacePath) => - isWithinWorkspaceRoot(normalizedPath, workspacePath), - ); - if (!matchedWorkspacePath) { - return null; - } - const workspaceIds = lookup.workspaceIdsByPath[matchedWorkspacePath] ?? []; - if (!allowedWorkspaceIds) { - return workspaceIds[0] ?? null; - } - return ( - workspaceIds.find((workspaceId) => allowedWorkspaceIds.has(workspaceId)) ?? - null - ); + const normalizedPath = normalizeRootPath(path); + if (!normalizedPath) { + return null; + } + const matchedWorkspacePath = lookup.workspacePathsSorted.find((workspacePath) => + isWithinWorkspaceRoot(normalizedPath, workspacePath), + ); + if (!matchedWorkspacePath) { + return null; + } + const workspaceIds = lookup.workspaceIdsByPath[matchedWorkspacePath] ?? []; + if (!allowedWorkspaceIds) { + return workspaceIds[0] ?? null; + } + return ( + workspaceIds.find((workspaceId) => allowedWorkspaceIds.has(workspaceId)) ?? + null + ); } function getThreadListNextCursor(result: Record): string | null { - if (typeof result.nextCursor === "string") { - return result.nextCursor; - } - if (typeof result.next_cursor === "string") { - return result.next_cursor; - } - return null; + if (typeof result.nextCursor === "string") { + return result.nextCursor; + } + if (typeof result.next_cursor === "string") { + return result.next_cursor; + } + return null; +} + +type ThreadUsageHydrationState = { + inFlightRequestKeys: Set; + latestAppliedUpdatedAtByThread: Record; +}; + +type ThreadMetadataHydrationState = { + inFlightThreadIds: Set; + hydratedThreadIds: Set; +}; + +function toRecord(value: unknown): Record | null { + if (value && typeof value === "object" && !Array.isArray(value)) { + return value as Record; + } + return null; } type UseThreadActionsOptions = { - dispatch: Dispatch; - itemsByThread: ThreadState["itemsByThread"]; - threadsByWorkspace: ThreadState["threadsByWorkspace"]; - activeThreadIdByWorkspace: ThreadState["activeThreadIdByWorkspace"]; - activeTurnIdByThread: ThreadState["activeTurnIdByThread"]; - threadParentById: ThreadState["threadParentById"]; - threadListCursorByWorkspace: ThreadState["threadListCursorByWorkspace"]; - threadStatusById: ThreadState["threadStatusById"]; - threadSortKey: ThreadListSortKey; - onDebug?: (entry: DebugEntry) => void; - getCustomName: (workspaceId: string, threadId: string) => string | undefined; - threadActivityRef: MutableRefObject>>; - loadedThreadsRef: MutableRefObject>; - replaceOnResumeRef: MutableRefObject>; - applyCollabThreadLinksFromThread: ( - workspaceId: string, - threadId: string, - thread: Record, - ) => void; - updateThreadParent: (parentId: string, childIds: string[]) => void; - onSubagentThreadDetected: (workspaceId: string, threadId: string) => void; - onThreadCodexMetadataDetected?: ( - workspaceId: string, - threadId: string, - metadata: { modelId: string | null; effort: string | null }, - ) => void; + dispatch: Dispatch; + itemsByThread: ThreadState["itemsByThread"]; + threadsByWorkspace: ThreadState["threadsByWorkspace"]; + tokenUsageThreadIdsByWorkspace: ThreadState["tokenUsageThreadIdsByWorkspace"]; + activeThreadIdByWorkspace: ThreadState["activeThreadIdByWorkspace"]; + activeTurnIdByThread: ThreadState["activeTurnIdByThread"]; + threadParentById: ThreadState["threadParentById"]; + threadListCursorByWorkspace: ThreadState["threadListCursorByWorkspace"]; + threadStatusById: ThreadState["threadStatusById"]; + threadSortKey: ThreadListSortKey; + onDebug?: (entry: DebugEntry) => void; + getCustomName: (workspaceId: string, threadId: string) => string | undefined; + threadActivityRef: MutableRefObject>>; + loadedThreadsRef: MutableRefObject>; + replaceOnResumeRef: MutableRefObject>; + applyCollabThreadLinksFromThread: ( + workspaceId: string, + threadId: string, + thread: Record, + ) => void; + updateThreadParent: (parentId: string, childIds: string[]) => void; + onSubagentThreadDetected: (workspaceId: string, threadId: string) => void; + onThreadCodexMetadataDetected?: ( + workspaceId: string, + threadId: string, + metadata: { modelId: string | null; effort: string | null }, + ) => void; + enableBackgroundMetadataHydration?: boolean; }; export function useThreadActions({ - dispatch, - itemsByThread, - threadsByWorkspace, - activeThreadIdByWorkspace, - activeTurnIdByThread, - threadParentById, - threadListCursorByWorkspace, - threadStatusById, - threadSortKey, - onDebug, - getCustomName, - threadActivityRef, - loadedThreadsRef, - replaceOnResumeRef, - applyCollabThreadLinksFromThread, - updateThreadParent, - onSubagentThreadDetected, - onThreadCodexMetadataDetected, + dispatch, + itemsByThread, + threadsByWorkspace, + tokenUsageThreadIdsByWorkspace, + activeThreadIdByWorkspace, + activeTurnIdByThread, + threadParentById, + threadListCursorByWorkspace, + threadStatusById, + threadSortKey, + onDebug, + getCustomName, + threadActivityRef, + loadedThreadsRef, + replaceOnResumeRef, + applyCollabThreadLinksFromThread, + updateThreadParent, + onSubagentThreadDetected, + onThreadCodexMetadataDetected, + enableBackgroundMetadataHydration = false, }: UseThreadActionsOptions) { - const resumeInFlightByThreadRef = useRef>({}); - const threadStatusByIdRef = useRef(threadStatusById); - const activeTurnIdByThreadRef = useRef(activeTurnIdByThread); - threadStatusByIdRef.current = threadStatusById; - activeTurnIdByThreadRef.current = activeTurnIdByThread; + const resumeInFlightByThreadRef = useRef>({}); + const threadUsageHydrationByWorkspaceRef = useRef< + Record + >({}); + const threadMetadataHydrationByWorkspaceRef = useRef< + Record + >({}); + const threadStatusByIdRef = useRef(threadStatusById); + const activeTurnIdByThreadRef = useRef(activeTurnIdByThread); + threadStatusByIdRef.current = threadStatusById; + activeTurnIdByThreadRef.current = activeTurnIdByThread; - const extractThreadId = useCallback((response: Record) => { - const thread = response.result?.thread ?? response.thread ?? null; - return String(thread?.id ?? ""); - }, []); + const extractThreadId = useCallback((response: Record) => { + const thread = response.result?.thread ?? response.thread ?? null; + return String(thread?.id ?? ""); + }, []); - const startThreadForWorkspace = useCallback( - async (workspaceId: string, options?: { activate?: boolean }) => { - const shouldActivate = options?.activate !== false; - onDebug?.({ - id: `${Date.now()}-client-thread-start`, - timestamp: Date.now(), - source: "client", - label: "thread/start", - payload: { workspaceId }, - }); - try { - const response = await startThreadService(workspaceId); - onDebug?.({ - id: `${Date.now()}-server-thread-start`, - timestamp: Date.now(), - source: "server", - label: "thread/start response", - payload: response, - }); - const threadId = extractThreadId(response); - if (threadId) { - dispatch({ type: "ensureThread", workspaceId, threadId }); - if (shouldActivate) { - dispatch({ type: "setActiveThreadId", workspaceId, threadId }); - } - loadedThreadsRef.current[threadId] = true; - return threadId; - } - return null; - } catch (error) { - onDebug?.({ - id: `${Date.now()}-client-thread-start-error`, - timestamp: Date.now(), - source: "error", - label: "thread/start error", - payload: error instanceof Error ? error.message : String(error), - }); - throw error; - } - }, - [dispatch, extractThreadId, loadedThreadsRef, onDebug], - ); + const hydrateThreadUsageForWorkspace = useCallback( + async (workspace: WorkspaceInfo, threadIds: string[]) => { + const uniqueThreadIds = Array.from( + new Set( + threadIds + .map((threadId) => threadId.trim()) + .filter((threadId) => threadId.length > 0), + ), + ); + if (uniqueThreadIds.length === 0) { + return; + } - const resumeThreadForWorkspace = useCallback( - async ( - workspaceId: string, - threadId: string, - force = false, - replaceLocal = false, - ) => { - if (!threadId) { - return null; - } - if (!force && loadedThreadsRef.current[threadId]) { - return threadId; - } - const status = threadStatusByIdRef.current[threadId]; - if (status?.isProcessing && loadedThreadsRef.current[threadId] && !force) { - onDebug?.({ - id: `${Date.now()}-client-thread-resume-skipped`, - timestamp: Date.now(), - source: "client", - label: "thread/resume skipped", - payload: { workspaceId, threadId, reason: "active-turn" }, - }); - return threadId; - } - onDebug?.({ - id: `${Date.now()}-client-thread-resume`, - timestamp: Date.now(), - source: "client", - label: "thread/resume", - payload: { workspaceId, threadId }, - }); - const inFlightCount = - (resumeInFlightByThreadRef.current[threadId] ?? 0) + 1; - resumeInFlightByThreadRef.current[threadId] = inFlightCount; - if (inFlightCount === 1) { - dispatch({ type: "setThreadResumeLoading", threadId, isLoading: true }); - } - try { - const response = - (await resumeThreadService(workspaceId, threadId)) as - | Record - | null; - onDebug?.({ - id: `${Date.now()}-server-thread-resume`, - timestamp: Date.now(), - source: "server", - label: "thread/resume response", - payload: response, - }); - const result = (response?.result ?? response) as - | Record - | null; - const thread = (result?.thread ?? response?.thread ?? null) as - | Record - | null; - if (thread) { - const codexMetadata = extractThreadCodexMetadata(thread); - if (codexMetadata.modelId || codexMetadata.effort) { - onThreadCodexMetadataDetected?.(workspaceId, threadId, codexMetadata); - } - dispatch({ type: "ensureThread", workspaceId, threadId }); - applyCollabThreadLinksFromThread(workspaceId, threadId, thread); - const sourceParentId = getParentThreadIdFromThread(thread); - if (sourceParentId) { - updateThreadParent(sourceParentId, [threadId]); - onSubagentThreadDetected(workspaceId, threadId); - } - const items = buildItemsFromThread(thread); - const localItems = itemsByThread[threadId] ?? []; - const shouldReplace = - replaceLocal || replaceOnResumeRef.current[threadId] === true; - if (shouldReplace) { - replaceOnResumeRef.current[threadId] = false; - } - if (localItems.length > 0 && !shouldReplace) { - loadedThreadsRef.current[threadId] = true; - return threadId; - } - const resumedTurnState = getResumedTurnState(thread); - const localStatus = threadStatusByIdRef.current[threadId]; - const localActiveTurnId = - activeTurnIdByThreadRef.current[threadId] ?? null; - const keepLocalProcessing = - (localStatus?.isProcessing ?? false) && - !resumedTurnState.activeTurnId && - !resumedTurnState.confidentNoActiveTurn; - const resumedActiveTurnId = keepLocalProcessing - ? localActiveTurnId - : resumedTurnState.activeTurnId; - const shouldMarkProcessing = keepLocalProcessing || Boolean(resumedActiveTurnId); - const processingTimestamp = - resumedTurnState.activeTurnStartedAtMs ?? Date.now(); - if (keepLocalProcessing) { - onDebug?.({ - id: `${Date.now()}-client-thread-resume-keep-processing`, - timestamp: Date.now(), - source: "client", - label: "thread/resume keep-processing", - payload: { workspaceId, threadId }, - }); - } - dispatch({ - type: "markProcessing", - threadId, - isProcessing: shouldMarkProcessing, - timestamp: processingTimestamp, - }); - dispatch({ - type: "setActiveTurnId", - threadId, - turnId: resumedActiveTurnId, - }); - dispatch({ - type: "markReviewing", - threadId, - isReviewing: isReviewingFromThread(thread), - }); - const hasOverlap = - items.length > 0 && - localItems.length > 0 && - items.some((item) => localItems.some((local) => local.id === item.id)); - const mergedItems = - items.length > 0 - ? shouldReplace - ? items - : localItems.length > 0 && !hasOverlap - ? localItems - : mergeThreadItems(items, localItems) - : localItems; - if (mergedItems.length > 0) { - dispatch({ type: "setThreadItems", threadId, items: mergedItems }); - } - const preview = asString(thread?.preview ?? ""); - const customName = getCustomName(workspaceId, threadId); - if (!customName && preview) { - dispatch({ - type: "setThreadName", - workspaceId, - threadId, - name: previewThreadName(preview, "New Agent"), - }); - } - const lastAgentMessage = [...mergedItems] - .reverse() - .find( - (item) => item.kind === "message" && item.role === "assistant", - ) as ConversationItem | undefined; - const lastText = - lastAgentMessage && lastAgentMessage.kind === "message" - ? lastAgentMessage.text - : preview; - if (lastText) { - dispatch({ - type: "setLastAgentMessage", - threadId, - text: lastText, - timestamp: getThreadTimestamp(thread), - }); - } - } - loadedThreadsRef.current[threadId] = true; - return threadId; - } catch (error) { - onDebug?.({ - id: `${Date.now()}-client-thread-resume-error`, - timestamp: Date.now(), - source: "error", - label: "thread/resume error", - payload: error instanceof Error ? error.message : String(error), - }); - return null; - } finally { - const nextCount = Math.max( - 0, - (resumeInFlightByThreadRef.current[threadId] ?? 1) - 1, - ); - if (nextCount === 0) { - delete resumeInFlightByThreadRef.current[threadId]; - dispatch({ type: "setThreadResumeLoading", threadId, isLoading: false }); - } else { - resumeInFlightByThreadRef.current[threadId] = nextCount; - } - } - }, - [ - applyCollabThreadLinksFromThread, - dispatch, - getCustomName, - itemsByThread, - loadedThreadsRef, - onDebug, - onSubagentThreadDetected, - onThreadCodexMetadataDetected, - replaceOnResumeRef, - updateThreadParent, - ], - ); + const requestKey = [...uniqueThreadIds].sort().join("|"); + const currentState = + threadUsageHydrationByWorkspaceRef.current[workspace.id] ?? { + inFlightRequestKeys: new Set(), + latestAppliedUpdatedAtByThread: {}, + }; + if (currentState.inFlightRequestKeys.has(requestKey)) { + return; + } + currentState.inFlightRequestKeys.add(requestKey); + threadUsageHydrationByWorkspaceRef.current[workspace.id] = currentState; - const forkThreadForWorkspace = useCallback( - async ( - workspaceId: string, - threadId: string, - options?: { activate?: boolean }, - ) => { - if (!threadId) { - return null; - } - const shouldActivate = options?.activate !== false; - onDebug?.({ - id: `${Date.now()}-client-thread-fork`, - timestamp: Date.now(), - source: "client", - label: "thread/fork", - payload: { workspaceId, threadId }, - }); - try { - const response = await forkThreadService(workspaceId, threadId); - onDebug?.({ - id: `${Date.now()}-server-thread-fork`, - timestamp: Date.now(), - source: "server", - label: "thread/fork response", - payload: response, - }); - const forkedThreadId = extractThreadId(response); - if (!forkedThreadId) { - return null; - } - dispatch({ type: "ensureThread", workspaceId, threadId: forkedThreadId }); - if (shouldActivate) { - dispatch({ - type: "setActiveThreadId", - workspaceId, - threadId: forkedThreadId, - }); - } - loadedThreadsRef.current[forkedThreadId] = false; - await resumeThreadForWorkspace(workspaceId, forkedThreadId, true, true); - return forkedThreadId; - } catch (error) { - onDebug?.({ - id: `${Date.now()}-client-thread-fork-error`, - timestamp: Date.now(), - source: "error", - label: "thread/fork error", - payload: error instanceof Error ? error.message : String(error), - }); - return null; - } - }, - [ - dispatch, - extractThreadId, - loadedThreadsRef, - onDebug, - resumeThreadForWorkspace, - ], - ); + try { + const snapshot = await localThreadUsageSnapshotService( + uniqueThreadIds, + workspace.path, + ); + const snapshotRecord = toRecord(snapshot); + if (!snapshotRecord) { + return; + } + const updatedAtValue = snapshotRecord.updatedAt; + const updatedAt = + typeof updatedAtValue === "number" && Number.isFinite(updatedAtValue) + ? updatedAtValue + : null; + const usageByThread = toRecord(snapshotRecord.usageByThread); + if (!usageByThread) { + return; + } + Object.entries(usageByThread).forEach(([threadId, tokenUsage]) => { + const tokenUsageRecord = toRecord(tokenUsage); + if (!threadId || !tokenUsageRecord) { + return; + } + const latestAppliedForThread = + currentState.latestAppliedUpdatedAtByThread[threadId] ?? 0; + if (updatedAt !== null && updatedAt < latestAppliedForThread) { + return; + } + dispatch({ + type: "setThreadTokenUsage", + threadId, + tokenUsage: normalizeTokenUsage(tokenUsageRecord), + }); + if (updatedAt !== null) { + currentState.latestAppliedUpdatedAtByThread[threadId] = Math.max( + latestAppliedForThread, + updatedAt, + ); + } + }); + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-usage-hydrate-error`, + timestamp: Date.now(), + source: "error", + label: "thread/usage hydrate error", + payload: error instanceof Error ? error.message : String(error), + }); + } finally { + currentState.inFlightRequestKeys.delete(requestKey); + } + }, + [dispatch, onDebug], + ); - const refreshThread = useCallback( - async (workspaceId: string, threadId: string) => { - if (!threadId) { - return null; - } - replaceOnResumeRef.current[threadId] = true; - return resumeThreadForWorkspace(workspaceId, threadId, true, true); - }, - [replaceOnResumeRef, resumeThreadForWorkspace], - ); + const hydrateThreadMetadataForWorkspace = useCallback( + async (workspace: WorkspaceInfo, threadIds: string[]) => { + if (!enableBackgroundMetadataHydration || !onThreadCodexMetadataDetected) { + return; + } + + const uniqueThreadIds = Array.from( + new Set( + threadIds + .map((threadId) => threadId.trim()) + .filter((threadId) => threadId.length > 0), + ), + ); + if (uniqueThreadIds.length === 0) { + return; + } - const resetWorkspaceThreads = useCallback( - (workspaceId: string) => { - const threadIds = new Set(); - const list = threadsByWorkspace[workspaceId] ?? []; - list.forEach((thread) => threadIds.add(thread.id)); - const activeThread = activeThreadIdByWorkspace[workspaceId]; - if (activeThread) { - threadIds.add(activeThread); - } - threadIds.forEach((threadId) => { - loadedThreadsRef.current[threadId] = false; - }); - }, - [activeThreadIdByWorkspace, loadedThreadsRef, threadsByWorkspace], - ); + const currentState = + threadMetadataHydrationByWorkspaceRef.current[workspace.id] ?? { + inFlightThreadIds: new Set(), + hydratedThreadIds: new Set(), + }; + threadMetadataHydrationByWorkspaceRef.current[workspace.id] = currentState; - const buildThreadSummary = useCallback( - ( - workspaceId: string, - thread: Record, - fallbackIndex: number, - ): ThreadSummary | null => { - const id = String(thread?.id ?? ""); - if (!id) { - return null; - } - const preview = asString(thread?.preview ?? "").trim(); - const customName = getCustomName(workspaceId, id); - const fallbackName = `Agent ${fallbackIndex + 1}`; - const name = customName - ? customName - : preview.length > 0 - ? preview.length > 38 - ? `${preview.slice(0, 38)}…` - : preview - : fallbackName; - const metadata = extractThreadCodexMetadata(thread); - return { - id, - name, - updatedAt: getThreadTimestamp(thread), - createdAt: getThreadCreatedTimestamp(thread), - ...(metadata.modelId ? { modelId: metadata.modelId } : {}), - ...(metadata.effort ? { effort: metadata.effort } : {}), - }; - }, - [getCustomName], - ); + const pendingThreadIds = uniqueThreadIds.filter((threadId) => { + if (currentState.hydratedThreadIds.has(threadId)) { + return false; + } + if (currentState.inFlightThreadIds.has(threadId)) { + return false; + } + if (threadStatusByIdRef.current[threadId]?.isProcessing) { + return false; + } + return true; + }); + if (pendingThreadIds.length === 0) { + return; + } - const listThreadsForWorkspaces = useCallback( - async ( - workspaces: WorkspaceInfo[], - options?: { - preserveState?: boolean; - sortKey?: ThreadListSortKey; - maxPages?: number; - }, - ) => { - const targets = workspaces.filter((workspace) => workspace.id); - if (targets.length === 0) { - return; - } - const preserveState = options?.preserveState ?? false; - const requestedSortKey = options?.sortKey ?? threadSortKey; - const maxPages = Math.max(1, options?.maxPages ?? THREAD_LIST_MAX_PAGES_DEFAULT); - if (!preserveState) { - targets.forEach((workspace) => { - dispatch({ - type: "setThreadListLoading", - workspaceId: workspace.id, - isLoading: true, - }); - dispatch({ - type: "setThreadListCursor", - workspaceId: workspace.id, - cursor: null, - }); - }); - } - onDebug?.({ - id: `${Date.now()}-client-thread-list`, - timestamp: Date.now(), - source: "client", - label: "thread/list", - payload: { - workspaceIds: targets.map((workspace) => workspace.id), - preserveState, - maxPages, + for ( + let startIndex = 0; + startIndex < pendingThreadIds.length; + startIndex += THREAD_METADATA_PREFETCH_BATCH_SIZE + ) { + const batch = pendingThreadIds.slice( + startIndex, + startIndex + THREAD_METADATA_PREFETCH_BATCH_SIZE, + ); + await Promise.all( + batch.map(async (threadId) => { + currentState.inFlightThreadIds.add(threadId); + try { + const response = + (await resumeThreadService(workspace.id, threadId)) as + | Record + | null; + const result = (response?.result ?? response) as + | Record + | null; + const thread = (result?.thread ?? response?.thread ?? null) as + | Record + | null; + if (!thread) { + return; + } + const codexMetadata = extractThreadCodexMetadata(thread); + if (codexMetadata.modelId || codexMetadata.effort) { + onThreadCodexMetadataDetected(workspace.id, threadId, codexMetadata); + } + // Mark this thread as hydrated even when model metadata is absent to + // avoid repeatedly resuming the same thread on every list refresh. + currentState.hydratedThreadIds.add(threadId); + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-metadata-hydrate-error`, + timestamp: Date.now(), + source: "error", + label: "thread/metadata hydrate error", + payload: error instanceof Error ? error.message : String(error), + }); + } finally { + currentState.inFlightThreadIds.delete(threadId); + } + }), + ); + } }, - }); - try { - const requester = targets.find((workspace) => workspace.connected) ?? targets[0]; - const matchingThreadsByWorkspace: Record[]> = {}; - let workspacePathLookup = buildWorkspacePathLookup(targets); - const targetWorkspaceIds = new Set(targets.map((workspace) => workspace.id)); - try { - const knownWorkspaces = await listWorkspacesService(); - if (knownWorkspaces.length > 0) { - workspacePathLookup = buildWorkspacePathLookup([ - ...targets, - ...knownWorkspaces, - ]); - } - } catch { - workspacePathLookup = buildWorkspacePathLookup(targets); - } - const uniqueThreadIdsByWorkspace: Record> = {}; - const resumeCursorByWorkspace: Record = {}; - targets.forEach((workspace) => { - matchingThreadsByWorkspace[workspace.id] = []; - uniqueThreadIdsByWorkspace[workspace.id] = new Set(); - resumeCursorByWorkspace[workspace.id] = null; - }); - let pagesFetched = 0; - let cursor: string | null = null; - do { - const pageCursor = cursor; - pagesFetched += 1; - const response = - (await listThreadsService( - requester.id, - cursor, - THREAD_LIST_PAGE_SIZE, - requestedSortKey, - )) as Record; - onDebug?.({ - id: `${Date.now()}-server-thread-list`, - timestamp: Date.now(), - source: "server", - label: "thread/list response", - payload: response, - }); - const result = (response.result ?? response) as Record; - const data = Array.isArray(result?.data) - ? (result.data as Record[]) - : []; - const nextCursor = getThreadListNextCursor(result); - data.forEach((thread) => { - const workspaceId = resolveWorkspaceIdForThreadPath( - String(thread?.cwd ?? ""), - workspacePathLookup, - targetWorkspaceIds, - ); - if (!workspaceId) { - return; + [ + dispatch, + enableBackgroundMetadataHydration, + onDebug, + onThreadCodexMetadataDetected, + ], + ); + + const startThreadForWorkspace = useCallback( + async (workspaceId: string, options?: { activate?: boolean }) => { + const shouldActivate = options?.activate !== false; + onDebug?.({ + id: `${Date.now()}-client-thread-start`, + timestamp: Date.now(), + source: "client", + label: "thread/start", + payload: { workspaceId }, + }); + try { + const response = await startThreadService(workspaceId); + onDebug?.({ + id: `${Date.now()}-server-thread-start`, + timestamp: Date.now(), + source: "server", + label: "thread/start response", + payload: response, + }); + const threadId = extractThreadId(response); + if (threadId) { + dispatch({ type: "ensureThread", workspaceId, threadId }); + if (shouldActivate) { + dispatch({ type: "setActiveThreadId", workspaceId, threadId }); + } + loadedThreadsRef.current[threadId] = true; + return threadId; + } + return null; + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-start-error`, + timestamp: Date.now(), + source: "error", + label: "thread/start error", + payload: error instanceof Error ? error.message : String(error), + }); + throw error; } - matchingThreadsByWorkspace[workspaceId]?.push(thread); - const threadId = String(thread?.id ?? ""); + }, + [dispatch, extractThreadId, loadedThreadsRef, onDebug], + ); + + const resumeThreadForWorkspace = useCallback( + async ( + workspaceId: string, + threadId: string, + force = false, + replaceLocal = false, + ) => { if (!threadId) { - return; + return null; } - const uniqueThreadIds = uniqueThreadIdsByWorkspace[workspaceId]; - if (!uniqueThreadIds || uniqueThreadIds.has(threadId)) { - return; + if (!force && loadedThreadsRef.current[threadId]) { + return threadId; } - uniqueThreadIds.add(threadId); - if ( - uniqueThreadIds.size > THREAD_LIST_TARGET_COUNT && - resumeCursorByWorkspace[workspaceId] === null - ) { - resumeCursorByWorkspace[workspaceId] = - pageCursor ?? THREAD_LIST_CURSOR_PAGE_START; + const status = threadStatusByIdRef.current[threadId]; + if (status?.isProcessing && loadedThreadsRef.current[threadId] && !force) { + onDebug?.({ + id: `${Date.now()}-client-thread-resume-skipped`, + timestamp: Date.now(), + source: "client", + label: "thread/resume skipped", + payload: { workspaceId, threadId, reason: "active-turn" }, + }); + return threadId; } - }); - cursor = nextCursor; - if (pagesFetched >= maxPages) { - break; - } - } while (cursor); - - const nextThreadActivity = { ...threadActivityRef.current }; - let didChangeAnyActivity = false; - targets.forEach((workspace) => { - const matchingThreads = matchingThreadsByWorkspace[workspace.id] ?? []; - const uniqueById = new Map>(); - matchingThreads.forEach((thread) => { - const id = String(thread?.id ?? ""); - if (id && !uniqueById.has(id)) { - uniqueById.set(id, thread); + onDebug?.({ + id: `${Date.now()}-client-thread-resume`, + timestamp: Date.now(), + source: "client", + label: "thread/resume", + payload: { workspaceId, threadId }, + }); + const inFlightCount = + (resumeInFlightByThreadRef.current[threadId] ?? 0) + 1; + resumeInFlightByThreadRef.current[threadId] = inFlightCount; + if (inFlightCount === 1) { + dispatch({ type: "setThreadResumeLoading", threadId, isLoading: true }); + } + try { + const response = + (await resumeThreadService(workspaceId, threadId)) as + | Record + | null; + onDebug?.({ + id: `${Date.now()}-server-thread-resume`, + timestamp: Date.now(), + source: "server", + label: "thread/resume response", + payload: response, + }); + const result = (response?.result ?? response) as + | Record + | null; + const thread = (result?.thread ?? response?.thread ?? null) as + | Record + | null; + if (thread) { + const codexMetadata = extractThreadCodexMetadata(thread); + if (codexMetadata.modelId || codexMetadata.effort) { + onThreadCodexMetadataDetected?.(workspaceId, threadId, codexMetadata); + } + dispatch({ type: "ensureThread", workspaceId, threadId }); + applyCollabThreadLinksFromThread(workspaceId, threadId, thread); + const rawTokenUsage = + toRecord(thread.tokenUsage) ?? + toRecord(thread.token_usage) ?? + toRecord(result?.tokenUsage) ?? + toRecord(result?.token_usage); + if (rawTokenUsage) { + dispatch({ + type: "setThreadTokenUsage", + threadId, + tokenUsage: normalizeTokenUsage(rawTokenUsage), + }); + } + const sourceParentId = getParentThreadIdFromThread(thread); + if (sourceParentId) { + updateThreadParent(sourceParentId, [threadId]); + onSubagentThreadDetected(workspaceId, threadId); + } + const items = buildItemsFromThread(thread); + const localItems = itemsByThread[threadId] ?? []; + const shouldReplace = + replaceLocal || replaceOnResumeRef.current[threadId] === true; + if (shouldReplace) { + replaceOnResumeRef.current[threadId] = false; + } + if (localItems.length > 0 && !shouldReplace) { + loadedThreadsRef.current[threadId] = true; + return threadId; + } + const resumedTurnState = getResumedTurnState(thread); + const localStatus = threadStatusByIdRef.current[threadId]; + const localActiveTurnId = + activeTurnIdByThreadRef.current[threadId] ?? null; + const keepLocalProcessing = + (localStatus?.isProcessing ?? false) && + !resumedTurnState.activeTurnId && + !resumedTurnState.confidentNoActiveTurn; + const resumedActiveTurnId = keepLocalProcessing + ? localActiveTurnId + : resumedTurnState.activeTurnId; + const shouldMarkProcessing = keepLocalProcessing || Boolean(resumedActiveTurnId); + const processingTimestamp = + resumedTurnState.activeTurnStartedAtMs ?? Date.now(); + if (keepLocalProcessing) { + onDebug?.({ + id: `${Date.now()}-client-thread-resume-keep-processing`, + timestamp: Date.now(), + source: "client", + label: "thread/resume keep-processing", + payload: { workspaceId, threadId }, + }); + } + dispatch({ + type: "markProcessing", + threadId, + isProcessing: shouldMarkProcessing, + timestamp: processingTimestamp, + }); + dispatch({ + type: "setActiveTurnId", + threadId, + turnId: resumedActiveTurnId, + }); + dispatch({ + type: "markReviewing", + threadId, + isReviewing: isReviewingFromThread(thread), + }); + const hasOverlap = + items.length > 0 && + localItems.length > 0 && + items.some((item) => localItems.some((local) => local.id === item.id)); + const mergedItems = + items.length > 0 + ? shouldReplace + ? items + : localItems.length > 0 && !hasOverlap + ? localItems + : mergeThreadItems(items, localItems) + : localItems; + if (mergedItems.length > 0) { + dispatch({ type: "setThreadItems", threadId, items: mergedItems }); + } + const preview = asString(thread?.preview ?? ""); + const customName = getCustomName(workspaceId, threadId); + if (!customName && preview) { + dispatch({ + type: "setThreadName", + workspaceId, + threadId, + name: previewThreadName(preview, "New Agent"), + }); + } + const lastAgentMessage = [...mergedItems] + .reverse() + .find( + (item) => item.kind === "message" && item.role === "assistant", + ) as ConversationItem | undefined; + const lastText = + lastAgentMessage && lastAgentMessage.kind === "message" + ? lastAgentMessage.text + : preview; + if (lastText) { + dispatch({ + type: "setLastAgentMessage", + threadId, + text: lastText, + timestamp: getThreadTimestamp(thread), + }); + } + } + loadedThreadsRef.current[threadId] = true; + return threadId; + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-resume-error`, + timestamp: Date.now(), + source: "error", + label: "thread/resume error", + payload: error instanceof Error ? error.message : String(error), + }); + return null; + } finally { + const nextCount = Math.max( + 0, + (resumeInFlightByThreadRef.current[threadId] ?? 1) - 1, + ); + if (nextCount === 0) { + delete resumeInFlightByThreadRef.current[threadId]; + dispatch({ type: "setThreadResumeLoading", threadId, isLoading: false }); + } else { + resumeInFlightByThreadRef.current[threadId] = nextCount; + } } - }); - const uniqueThreads = Array.from(uniqueById.values()); - const activityByThread = nextThreadActivity[workspace.id] ?? {}; - const nextActivityByThread = { ...activityByThread }; - let didChangeActivity = false; - uniqueThreads.forEach((thread) => { - const threadId = String(thread?.id ?? ""); + }, + [ + applyCollabThreadLinksFromThread, + dispatch, + getCustomName, + itemsByThread, + loadedThreadsRef, + onDebug, + onSubagentThreadDetected, + onThreadCodexMetadataDetected, + replaceOnResumeRef, + updateThreadParent, + ], + ); + + const forkThreadForWorkspace = useCallback( + async ( + workspaceId: string, + threadId: string, + options?: { activate?: boolean }, + ) => { if (!threadId) { - return; + return null; } - const codexMetadata = extractThreadCodexMetadata(thread); - if (codexMetadata.modelId || codexMetadata.effort) { - onThreadCodexMetadataDetected?.(workspace.id, threadId, codexMetadata); + const shouldActivate = options?.activate !== false; + onDebug?.({ + id: `${Date.now()}-client-thread-fork`, + timestamp: Date.now(), + source: "client", + label: "thread/fork", + payload: { workspaceId, threadId }, + }); + try { + const response = await forkThreadService(workspaceId, threadId); + onDebug?.({ + id: `${Date.now()}-server-thread-fork`, + timestamp: Date.now(), + source: "server", + label: "thread/fork response", + payload: response, + }); + const forkedThreadId = extractThreadId(response); + if (!forkedThreadId) { + return null; + } + dispatch({ type: "ensureThread", workspaceId, threadId: forkedThreadId }); + if (shouldActivate) { + dispatch({ + type: "setActiveThreadId", + workspaceId, + threadId: forkedThreadId, + }); + } + loadedThreadsRef.current[forkedThreadId] = false; + await resumeThreadForWorkspace(workspaceId, forkedThreadId, true, true); + return forkedThreadId; + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-fork-error`, + timestamp: Date.now(), + source: "error", + label: "thread/fork error", + payload: error instanceof Error ? error.message : String(error), + }); + return null; } - const sourceParentId = getParentThreadIdFromThread(thread); - if (sourceParentId) { - updateThreadParent(sourceParentId, [threadId]); - onSubagentThreadDetected(workspace.id, threadId); + }, + [ + dispatch, + extractThreadId, + loadedThreadsRef, + onDebug, + resumeThreadForWorkspace, + ], + ); + + const refreshThread = useCallback( + async (workspaceId: string, threadId: string) => { + if (!threadId) { + return null; } - const timestamp = getThreadTimestamp(thread); - if (timestamp > (nextActivityByThread[threadId] ?? 0)) { - nextActivityByThread[threadId] = timestamp; - didChangeActivity = true; + replaceOnResumeRef.current[threadId] = true; + return resumeThreadForWorkspace(workspaceId, threadId, true, true); + }, + [replaceOnResumeRef, resumeThreadForWorkspace], + ); + + const resetWorkspaceThreads = useCallback( + (workspaceId: string) => { + const threadIds = new Set(); + const list = threadsByWorkspace[workspaceId] ?? []; + list.forEach((thread) => threadIds.add(thread.id)); + const activeThread = activeThreadIdByWorkspace[workspaceId]; + if (activeThread) { + threadIds.add(activeThread); } - }); - if (didChangeActivity) { - nextThreadActivity[workspace.id] = nextActivityByThread; - didChangeAnyActivity = true; - } - if (requestedSortKey === "updated_at") { - uniqueThreads.sort((a, b) => { - const aId = String(a?.id ?? ""); - const bId = String(b?.id ?? ""); - const aCreated = getThreadTimestamp(a); - const bCreated = getThreadTimestamp(b); - const aActivity = Math.max(nextActivityByThread[aId] ?? 0, aCreated); - const bActivity = Math.max(nextActivityByThread[bId] ?? 0, bCreated); - return bActivity - aActivity; + threadIds.forEach((threadId) => { + loadedThreadsRef.current[threadId] = false; }); - } else { - uniqueThreads.sort((a, b) => { - const delta = - getThreadCreatedTimestamp(b) - getThreadCreatedTimestamp(a); - if (delta !== 0) { - return delta; - } - const aId = String(a?.id ?? ""); - const bId = String(b?.id ?? ""); - return aId.localeCompare(bId); - }); - } - const summaryById = new Map(); - uniqueThreads.forEach((thread, index) => { - const summary = buildThreadSummary(workspace.id, thread, index); - if (!summary) { - return; - } - summaryById.set(summary.id, summary); - }); - const summaries = uniqueThreads - .slice(0, THREAD_LIST_TARGET_COUNT) - .map((thread) => summaryById.get(String(thread?.id ?? "")) ?? null) - .filter((entry): entry is ThreadSummary => Boolean(entry)); - const includedIds = new Set(summaries.map((thread) => thread.id)); - const appendFreshAnchor = (threadId: string | null | undefined) => { - if (!threadId || includedIds.has(threadId)) { - return; + }, + [activeThreadIdByWorkspace, loadedThreadsRef, threadsByWorkspace], + ); + + const buildThreadSummary = useCallback( + ( + workspaceId: string, + thread: Record, + fallbackIndex: number, + ): ThreadSummary | null => { + const id = String(thread?.id ?? ""); + if (!id) { + return null; } - const summary = summaryById.get(threadId); - if (!summary) { - return; + const preview = asString(thread?.preview ?? "").trim(); + const customName = getCustomName(workspaceId, id); + const fallbackName = `Agent ${fallbackIndex + 1}`; + const name = customName + ? customName + : preview.length > 0 + ? preview.length > 38 + ? `${preview.slice(0, 38)}…` + : preview + : fallbackName; + const metadata = extractThreadCodexMetadata(thread); + return { + id, + name, + updatedAt: getThreadTimestamp(thread), + createdAt: getThreadCreatedTimestamp(thread), + ...(metadata.modelId ? { modelId: metadata.modelId } : {}), + ...(metadata.effort ? { effort: metadata.effort } : {}), + }; + }, + [getCustomName], + ); + + const listThreadsForWorkspaces = useCallback( + async ( + workspaces: WorkspaceInfo[], + options?: { + preserveState?: boolean; + sortKey?: ThreadListSortKey; + maxPages?: number; + }, + ) => { + const targets = workspaces.filter((workspace) => workspace.id); + if (targets.length === 0) { + return; } - summaries.push(summary); - includedIds.add(threadId); - }; - appendFreshAnchor(activeThreadIdByWorkspace[workspace.id]); - const workspaceThreadIds = new Set([ - ...Array.from(summaryById.keys()), - ...(threadsByWorkspace[workspace.id] ?? []).map((thread) => thread.id), - ]); - const activeThreadId = activeThreadIdByWorkspace[workspace.id]; - if (activeThreadId) { - workspaceThreadIds.add(activeThreadId); - } - workspaceThreadIds.forEach((threadId) => { - if (threadStatusById[threadId]?.isProcessing) { - appendFreshAnchor(threadId); + const preserveState = options?.preserveState ?? false; + const requestedSortKey = options?.sortKey ?? threadSortKey; + const maxPages = Math.max(1, options?.maxPages ?? THREAD_LIST_MAX_PAGES_DEFAULT); + if (!preserveState) { + targets.forEach((workspace) => { + dispatch({ + type: "setThreadListLoading", + workspaceId: workspace.id, + isLoading: true, + }); + dispatch({ + type: "setThreadListCursor", + workspaceId: workspace.id, + cursor: null, + }); + }); } - }); - const seedThreadIds = [...includedIds]; - seedThreadIds.forEach((threadId) => { - const visited = new Set([threadId]); - let parentId = threadParentById[threadId]; - while (parentId && !visited.has(parentId)) { - visited.add(parentId); - appendFreshAnchor(parentId); - parentId = threadParentById[parentId]; + onDebug?.({ + id: `${Date.now()}-client-thread-list`, + timestamp: Date.now(), + source: "client", + label: "thread/list", + payload: { + workspaceIds: targets.map((workspace) => workspace.id), + preserveState, + maxPages, + }, + }); + try { + const requester = targets.find((workspace) => workspace.connected) ?? targets[0]; + const matchingThreadsByWorkspace: Record[]> = {}; + let workspacePathLookup = buildWorkspacePathLookup(targets); + const targetWorkspaceIds = new Set(targets.map((workspace) => workspace.id)); + try { + const knownWorkspaces = await listWorkspacesService(); + if (knownWorkspaces.length > 0) { + workspacePathLookup = buildWorkspacePathLookup([ + ...targets, + ...knownWorkspaces, + ]); + } + } catch { + workspacePathLookup = buildWorkspacePathLookup(targets); + } + const uniqueThreadIdsByWorkspace: Record> = {}; + const resumeCursorByWorkspace: Record = {}; + targets.forEach((workspace) => { + matchingThreadsByWorkspace[workspace.id] = []; + uniqueThreadIdsByWorkspace[workspace.id] = new Set(); + resumeCursorByWorkspace[workspace.id] = null; + }); + let pagesFetched = 0; + let cursor: string | null = null; + do { + const pageCursor = cursor; + pagesFetched += 1; + const response = + (await listThreadsService( + requester.id, + cursor, + THREAD_LIST_PAGE_SIZE, + requestedSortKey, + )) as Record; + onDebug?.({ + id: `${Date.now()}-server-thread-list`, + timestamp: Date.now(), + source: "server", + label: "thread/list response", + payload: response, + }); + const result = (response.result ?? response) as Record; + const data = Array.isArray(result?.data) + ? (result.data as Record[]) + : []; + const nextCursor = getThreadListNextCursor(result); + data.forEach((thread) => { + const workspaceId = resolveWorkspaceIdForThreadPath( + String(thread?.cwd ?? ""), + workspacePathLookup, + targetWorkspaceIds, + ); + if (!workspaceId) { + return; + } + matchingThreadsByWorkspace[workspaceId]?.push(thread); + const threadId = String(thread?.id ?? ""); + if (!threadId) { + return; + } + const uniqueThreadIds = uniqueThreadIdsByWorkspace[workspaceId]; + if (!uniqueThreadIds || uniqueThreadIds.has(threadId)) { + return; + } + uniqueThreadIds.add(threadId); + if ( + uniqueThreadIds.size > THREAD_LIST_TARGET_COUNT && + resumeCursorByWorkspace[workspaceId] === null + ) { + resumeCursorByWorkspace[workspaceId] = + pageCursor ?? THREAD_LIST_CURSOR_PAGE_START; + } + }); + cursor = nextCursor; + if (pagesFetched >= maxPages) { + break; + } + } while (cursor); + + const nextThreadActivity = { ...threadActivityRef.current }; + let didChangeAnyActivity = false; + targets.forEach((workspace) => { + const matchingThreads = matchingThreadsByWorkspace[workspace.id] ?? []; + const uniqueById = new Map>(); + matchingThreads.forEach((thread) => { + const id = String(thread?.id ?? ""); + if (id && !uniqueById.has(id)) { + uniqueById.set(id, thread); + } + }); + const uniqueThreads = Array.from(uniqueById.values()); + const activityByThread = nextThreadActivity[workspace.id] ?? {}; + const nextActivityByThread = { ...activityByThread }; + let didChangeActivity = false; + uniqueThreads.forEach((thread) => { + const threadId = String(thread?.id ?? ""); + if (!threadId) { + return; + } + const codexMetadata = extractThreadCodexMetadata(thread); + if (codexMetadata.modelId || codexMetadata.effort) { + onThreadCodexMetadataDetected?.(workspace.id, threadId, codexMetadata); + } + const sourceParentId = getParentThreadIdFromThread(thread); + if (sourceParentId) { + updateThreadParent(sourceParentId, [threadId]); + onSubagentThreadDetected(workspace.id, threadId); + } + const timestamp = getThreadTimestamp(thread); + if (timestamp > (nextActivityByThread[threadId] ?? 0)) { + nextActivityByThread[threadId] = timestamp; + didChangeActivity = true; + } + }); + if (didChangeActivity) { + nextThreadActivity[workspace.id] = nextActivityByThread; + didChangeAnyActivity = true; + } + if (requestedSortKey === "updated_at") { + uniqueThreads.sort((a, b) => { + const aId = String(a?.id ?? ""); + const bId = String(b?.id ?? ""); + const aCreated = getThreadTimestamp(a); + const bCreated = getThreadTimestamp(b); + const aActivity = Math.max(nextActivityByThread[aId] ?? 0, aCreated); + const bActivity = Math.max(nextActivityByThread[bId] ?? 0, bCreated); + return bActivity - aActivity; + }); + } else { + uniqueThreads.sort((a, b) => { + const delta = + getThreadCreatedTimestamp(b) - getThreadCreatedTimestamp(a); + if (delta !== 0) { + return delta; + } + const aId = String(a?.id ?? ""); + const bId = String(b?.id ?? ""); + return aId.localeCompare(bId); + }); + } + const summaryById = new Map(); + uniqueThreads.forEach((thread, index) => { + const summary = buildThreadSummary(workspace.id, thread, index); + if (!summary) { + return; + } + summaryById.set(summary.id, summary); + }); + const summaries = uniqueThreads + .slice(0, THREAD_LIST_TARGET_COUNT) + .map((thread) => summaryById.get(String(thread?.id ?? "")) ?? null) + .filter((entry): entry is ThreadSummary => Boolean(entry)); + const includedIds = new Set(summaries.map((thread) => thread.id)); + const appendFreshAnchor = (threadId: string | null | undefined) => { + if (!threadId || includedIds.has(threadId)) { + return; + } + const summary = summaryById.get(threadId); + if (!summary) { + return; + } + summaries.push(summary); + includedIds.add(threadId); + }; + appendFreshAnchor(activeThreadIdByWorkspace[workspace.id]); + const workspaceThreadIds = new Set([ + ...Array.from(summaryById.keys()), + ...(threadsByWorkspace[workspace.id] ?? []).map((thread) => thread.id), + ]); + const activeThreadId = activeThreadIdByWorkspace[workspace.id]; + if (activeThreadId) { + workspaceThreadIds.add(activeThreadId); + } + workspaceThreadIds.forEach((threadId) => { + if (threadStatusById[threadId]?.isProcessing) { + appendFreshAnchor(threadId); + } + }); + const seedThreadIds = [...includedIds]; + seedThreadIds.forEach((threadId) => { + const visited = new Set([threadId]); + let parentId = threadParentById[threadId]; + while (parentId && !visited.has(parentId)) { + visited.add(parentId); + appendFreshAnchor(parentId); + parentId = threadParentById[parentId]; + } + }); + const tokenUsageThreadIds = uniqueThreads + .map((thread) => String(thread?.id ?? "").trim()) + .filter((threadId) => threadId.length > 0); + dispatch({ + type: "setWorkspaceTokenUsageThreadIds", + workspaceId: workspace.id, + threadIds: tokenUsageThreadIds, + }); + dispatch({ + type: "setThreads", + workspaceId: workspace.id, + threads: summaries, + sortKey: requestedSortKey, + preserveAnchors: true, + }); + dispatch({ + type: "setThreadListCursor", + workspaceId: workspace.id, + cursor: resumeCursorByWorkspace[workspace.id] ?? cursor, + }); + uniqueThreads.forEach((thread) => { + const threadId = String(thread?.id ?? ""); + const preview = asString(thread?.preview ?? "").trim(); + if (!threadId || !preview) { + return; + } + dispatch({ + type: "setLastAgentMessage", + threadId, + text: preview, + timestamp: getThreadTimestamp(thread), + }); + }); + void hydrateThreadUsageForWorkspace( + workspace, + tokenUsageThreadIds, + ); + const missingModelThreadIds = summaries + .filter((thread) => !thread.modelId) + .map((thread) => thread.id); + void hydrateThreadMetadataForWorkspace( + workspace, + missingModelThreadIds, + ); + }); + if (didChangeAnyActivity) { + threadActivityRef.current = nextThreadActivity; + saveThreadActivity(nextThreadActivity); + } + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-list-error`, + timestamp: Date.now(), + source: "error", + label: "thread/list error", + payload: error instanceof Error ? error.message : String(error), + }); + } finally { + if (!preserveState) { + targets.forEach((workspace) => { + dispatch({ + type: "setThreadListLoading", + workspaceId: workspace.id, + isLoading: false, + }); + }); + } } - }); - dispatch({ - type: "setThreads", - workspaceId: workspace.id, - threads: summaries, - sortKey: requestedSortKey, - preserveAnchors: true, - }); - dispatch({ - type: "setThreadListCursor", - workspaceId: workspace.id, - cursor: resumeCursorByWorkspace[workspace.id] ?? cursor, - }); - uniqueThreads.forEach((thread) => { - const threadId = String(thread?.id ?? ""); - const preview = asString(thread?.preview ?? "").trim(); - if (!threadId || !preview) { - return; + }, + [ + buildThreadSummary, + dispatch, + hydrateThreadMetadataForWorkspace, + hydrateThreadUsageForWorkspace, + onDebug, + onSubagentThreadDetected, + onThreadCodexMetadataDetected, + activeThreadIdByWorkspace, + threadParentById, + threadActivityRef, + threadStatusById, + threadSortKey, + threadsByWorkspace, + updateThreadParent, + ], + ); + + const listThreadsForWorkspace = useCallback( + async ( + workspace: WorkspaceInfo, + options?: { + preserveState?: boolean; + sortKey?: ThreadListSortKey; + maxPages?: number; + }, + ) => { + await listThreadsForWorkspaces([workspace], options); + }, + [listThreadsForWorkspaces], + ); + + const loadOlderThreadsForWorkspace = useCallback( + async (workspace: WorkspaceInfo) => { + const requestedSortKey = threadSortKey; + const cursorValue = threadListCursorByWorkspace[workspace.id] ?? null; + if (!cursorValue) { + return; } + const nextCursor = + cursorValue === THREAD_LIST_CURSOR_PAGE_START ? null : cursorValue; + let workspacePathLookup = buildWorkspacePathLookup([workspace]); + const allowedWorkspaceIds = new Set([workspace.id]); + const existing = threadsByWorkspace[workspace.id] ?? []; dispatch({ - type: "setLastAgentMessage", - threadId, - text: preview, - timestamp: getThreadTimestamp(thread), + type: "setThreadListPaging", + workspaceId: workspace.id, + isLoading: true, }); - }); - }); - if (didChangeAnyActivity) { - threadActivityRef.current = nextThreadActivity; - saveThreadActivity(nextThreadActivity); - } - } catch (error) { - onDebug?.({ - id: `${Date.now()}-client-thread-list-error`, - timestamp: Date.now(), - source: "error", - label: "thread/list error", - payload: error instanceof Error ? error.message : String(error), - }); - } finally { - if (!preserveState) { - targets.forEach((workspace) => { - dispatch({ - type: "setThreadListLoading", - workspaceId: workspace.id, - isLoading: false, + onDebug?.({ + id: `${Date.now()}-client-thread-list-older`, + timestamp: Date.now(), + source: "client", + label: "thread/list older", + payload: { workspaceId: workspace.id, cursor: cursorValue }, }); - }); - } - } - }, - [ - buildThreadSummary, - dispatch, - onDebug, - onSubagentThreadDetected, - onThreadCodexMetadataDetected, - activeThreadIdByWorkspace, - threadParentById, - threadActivityRef, - threadStatusById, - threadSortKey, - threadsByWorkspace, - updateThreadParent, - ], - ); + try { + try { + const knownWorkspaces = await listWorkspacesService(); + if (knownWorkspaces.length > 0) { + workspacePathLookup = buildWorkspacePathLookup([ + workspace, + ...knownWorkspaces, + ]); + } + } catch { + workspacePathLookup = buildWorkspacePathLookup([workspace]); + } + const matchingThreads: Record[] = []; + const maxPagesWithoutMatch = THREAD_LIST_MAX_PAGES_OLDER; + let pagesFetched = 0; + let cursor: string | null = nextCursor; + do { + pagesFetched += 1; + const response = + (await listThreadsService( + workspace.id, + cursor, + THREAD_LIST_PAGE_SIZE, + requestedSortKey, + )) as Record; + onDebug?.({ + id: `${Date.now()}-server-thread-list-older`, + timestamp: Date.now(), + source: "server", + label: "thread/list older response", + payload: response, + }); + const result = (response.result ?? response) as Record; + const data = Array.isArray(result?.data) + ? (result.data as Record[]) + : []; + const next = getThreadListNextCursor(result); + matchingThreads.push( + ...data.filter( + (thread) => { + const workspaceId = resolveWorkspaceIdForThreadPath( + String(thread?.cwd ?? ""), + workspacePathLookup, + allowedWorkspaceIds, + ); + return workspaceId === workspace.id; + }, + ), + ); + cursor = next; + if (matchingThreads.length === 0 && pagesFetched >= maxPagesWithoutMatch) { + break; + } + if (pagesFetched >= THREAD_LIST_MAX_PAGES_OLDER) { + break; + } + } while (cursor && matchingThreads.length < THREAD_LIST_TARGET_COUNT); - const listThreadsForWorkspace = useCallback( - async ( - workspace: WorkspaceInfo, - options?: { - preserveState?: boolean; - sortKey?: ThreadListSortKey; - maxPages?: number; - }, - ) => { - await listThreadsForWorkspaces([workspace], options); - }, - [listThreadsForWorkspaces], - ); + const existingIds = new Set(existing.map((thread) => thread.id)); + const additions: ThreadSummary[] = []; + matchingThreads.forEach((thread) => { + const id = String(thread?.id ?? ""); + if (!id || existingIds.has(id)) { + return; + } + const codexMetadata = extractThreadCodexMetadata(thread); + if (codexMetadata.modelId || codexMetadata.effort) { + onThreadCodexMetadataDetected?.(workspace.id, id, codexMetadata); + } + const sourceParentId = getParentThreadIdFromThread(thread); + if (sourceParentId) { + updateThreadParent(sourceParentId, [id]); + } + const summary = buildThreadSummary( + workspace.id, + thread, + existing.length + additions.length, + ); + if (!summary) { + return; + } + additions.push(summary); + existingIds.add(id); + }); - const loadOlderThreadsForWorkspace = useCallback( - async (workspace: WorkspaceInfo) => { - const requestedSortKey = threadSortKey; - const cursorValue = threadListCursorByWorkspace[workspace.id] ?? null; - if (!cursorValue) { - return; - } - const nextCursor = - cursorValue === THREAD_LIST_CURSOR_PAGE_START ? null : cursorValue; - let workspacePathLookup = buildWorkspacePathLookup([workspace]); - const allowedWorkspaceIds = new Set([workspace.id]); - const existing = threadsByWorkspace[workspace.id] ?? []; - dispatch({ - type: "setThreadListPaging", - workspaceId: workspace.id, - isLoading: true, - }); - onDebug?.({ - id: `${Date.now()}-client-thread-list-older`, - timestamp: Date.now(), - source: "client", - label: "thread/list older", - payload: { workspaceId: workspace.id, cursor: cursorValue }, - }); - try { - try { - const knownWorkspaces = await listWorkspacesService(); - if (knownWorkspaces.length > 0) { - workspacePathLookup = buildWorkspacePathLookup([ - workspace, - ...knownWorkspaces, - ]); - } - } catch { - workspacePathLookup = buildWorkspacePathLookup([workspace]); - } - const matchingThreads: Record[] = []; - const maxPagesWithoutMatch = THREAD_LIST_MAX_PAGES_OLDER; - let pagesFetched = 0; - let cursor: string | null = nextCursor; - do { - pagesFetched += 1; - const response = - (await listThreadsService( - workspace.id, - cursor, - THREAD_LIST_PAGE_SIZE, - requestedSortKey, - )) as Record; - onDebug?.({ - id: `${Date.now()}-server-thread-list-older`, - timestamp: Date.now(), - source: "server", - label: "thread/list older response", - payload: response, - }); - const result = (response.result ?? response) as Record; - const data = Array.isArray(result?.data) - ? (result.data as Record[]) - : []; - const next = getThreadListNextCursor(result); - matchingThreads.push( - ...data.filter( - (thread) => { - const workspaceId = resolveWorkspaceIdForThreadPath( - String(thread?.cwd ?? ""), - workspacePathLookup, - allowedWorkspaceIds, + if (additions.length > 0) { + dispatch({ + type: "setThreads", + workspaceId: workspace.id, + threads: [...existing, ...additions], + sortKey: requestedSortKey, + }); + } + const tokenUsageThreadIds = Array.from( + new Set([ + ...(tokenUsageThreadIdsByWorkspace[workspace.id] ?? []), + ...matchingThreads + .map((thread) => String(thread?.id ?? "").trim()) + .filter((threadId) => threadId.length > 0), + ]), ); - return workspaceId === workspace.id; - }, - ), - ); - cursor = next; - if (matchingThreads.length === 0 && pagesFetched >= maxPagesWithoutMatch) { - break; - } - if (pagesFetched >= THREAD_LIST_MAX_PAGES_OLDER) { - break; - } - } while (cursor && matchingThreads.length < THREAD_LIST_TARGET_COUNT); - - const existingIds = new Set(existing.map((thread) => thread.id)); - const additions: ThreadSummary[] = []; - matchingThreads.forEach((thread) => { - const id = String(thread?.id ?? ""); - if (!id || existingIds.has(id)) { - return; - } - const codexMetadata = extractThreadCodexMetadata(thread); - if (codexMetadata.modelId || codexMetadata.effort) { - onThreadCodexMetadataDetected?.(workspace.id, id, codexMetadata); - } - const sourceParentId = getParentThreadIdFromThread(thread); - if (sourceParentId) { - updateThreadParent(sourceParentId, [id]); - } - const summary = buildThreadSummary( - workspace.id, - thread, - existing.length + additions.length, - ); - if (!summary) { - return; - } - additions.push(summary); - existingIds.add(id); - }); - - if (additions.length > 0) { - dispatch({ - type: "setThreads", - workspaceId: workspace.id, - threads: [...existing, ...additions], - sortKey: requestedSortKey, - }); - } - dispatch({ - type: "setThreadListCursor", - workspaceId: workspace.id, - cursor, - }); - matchingThreads.forEach((thread) => { - const threadId = String(thread?.id ?? ""); - const preview = asString(thread?.preview ?? "").trim(); - if (!threadId || !preview) { - return; - } - dispatch({ - type: "setLastAgentMessage", - threadId, - text: preview, - timestamp: getThreadTimestamp(thread), - }); - }); - } catch (error) { - onDebug?.({ - id: `${Date.now()}-client-thread-list-older-error`, - timestamp: Date.now(), - source: "error", - label: "thread/list older error", - payload: error instanceof Error ? error.message : String(error), - }); - } finally { - dispatch({ - type: "setThreadListPaging", - workspaceId: workspace.id, - isLoading: false, - }); - } - }, - [ - buildThreadSummary, - dispatch, - onDebug, - threadListCursorByWorkspace, - threadsByWorkspace, - threadSortKey, - updateThreadParent, - onThreadCodexMetadataDetected, - ], - ); + dispatch({ + type: "setWorkspaceTokenUsageThreadIds", + workspaceId: workspace.id, + threadIds: tokenUsageThreadIds, + }); + dispatch({ + type: "setThreadListCursor", + workspaceId: workspace.id, + cursor, + }); + matchingThreads.forEach((thread) => { + const threadId = String(thread?.id ?? ""); + const preview = asString(thread?.preview ?? "").trim(); + if (!threadId || !preview) { + return; + } + dispatch({ + type: "setLastAgentMessage", + threadId, + text: preview, + timestamp: getThreadTimestamp(thread), + }); + }); + void hydrateThreadUsageForWorkspace( + workspace, + additions.map((thread) => thread.id), + ); + const missingModelThreadIds = additions + .filter((thread) => !thread.modelId) + .map((thread) => thread.id); + void hydrateThreadMetadataForWorkspace( + workspace, + missingModelThreadIds, + ); + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-list-older-error`, + timestamp: Date.now(), + source: "error", + label: "thread/list older error", + payload: error instanceof Error ? error.message : String(error), + }); + } finally { + dispatch({ + type: "setThreadListPaging", + workspaceId: workspace.id, + isLoading: false, + }); + } + }, + [ + buildThreadSummary, + dispatch, + hydrateThreadMetadataForWorkspace, + hydrateThreadUsageForWorkspace, + onDebug, + threadListCursorByWorkspace, + threadsByWorkspace, + tokenUsageThreadIdsByWorkspace, + threadSortKey, + updateThreadParent, + onThreadCodexMetadataDetected, + ], + ); - const archiveThread = useCallback( - async (workspaceId: string, threadId: string) => { - try { - await archiveThreadService(workspaceId, threadId); - } catch (error) { - onDebug?.({ - id: `${Date.now()}-client-thread-archive-error`, - timestamp: Date.now(), - source: "error", - label: "thread/archive error", - payload: error instanceof Error ? error.message : String(error), - }); - } - }, - [onDebug], - ); + const archiveThread = useCallback( + async (workspaceId: string, threadId: string) => { + try { + await archiveThreadService(workspaceId, threadId); + } catch (error) { + onDebug?.({ + id: `${Date.now()}-client-thread-archive-error`, + timestamp: Date.now(), + source: "error", + label: "thread/archive error", + payload: error instanceof Error ? error.message : String(error), + }); + } + }, + [onDebug], + ); - return { - startThreadForWorkspace, - forkThreadForWorkspace, - resumeThreadForWorkspace, - refreshThread, - resetWorkspaceThreads, - listThreadsForWorkspaces, - listThreadsForWorkspace, - loadOlderThreadsForWorkspace, - archiveThread, - }; + return { + startThreadForWorkspace, + forkThreadForWorkspace, + resumeThreadForWorkspace, + refreshThread, + resetWorkspaceThreads, + listThreadsForWorkspaces, + listThreadsForWorkspace, + loadOlderThreadsForWorkspace, + archiveThread, + }; } diff --git a/src/features/threads/hooks/useThreads.integration.test.tsx b/src/features/threads/hooks/useThreads.integration.test.tsx index a6089db2c..029570f6d 100644 --- a/src/features/threads/hooks/useThreads.integration.test.tsx +++ b/src/features/threads/hooks/useThreads.integration.test.tsx @@ -8,6 +8,7 @@ import { archiveThread, interruptTurn, listThreads, + localThreadUsageSnapshot, resumeThread, sendUserMessage as sendUserMessageService, setThreadName, @@ -38,6 +39,7 @@ vi.mock("@services/tauri", () => ({ startReview: vi.fn(), startThread: vi.fn(), listThreads: vi.fn(), + localThreadUsageSnapshot: vi.fn(), resumeThread: vi.fn(), archiveThread: vi.fn(), setThreadName: vi.fn(), @@ -62,6 +64,10 @@ describe("useThreads UX integration", () => { handlers = null; localStorage.clear(); vi.clearAllMocks(); + vi.mocked(localThreadUsageSnapshot).mockResolvedValue({ + updatedAt: 1, + usageByThread: {}, + }); now = 1000; nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now++); }); diff --git a/src/features/threads/hooks/useThreads.ts b/src/features/threads/hooks/useThreads.ts index 8153a2c7f..9a304ac7a 100644 --- a/src/features/threads/hooks/useThreads.ts +++ b/src/features/threads/hooks/useThreads.ts @@ -53,6 +53,7 @@ type UseThreadsOptions = { customPrompts?: CustomPromptOption[]; onMessageActivity?: () => void; threadSortKey?: ThreadListSortKey; + enableBackgroundThreadMetadataHydration?: boolean; onThreadCodexMetadataDetected?: ( workspaceId: string, threadId: string, @@ -82,6 +83,7 @@ export function useThreads({ customPrompts = [], onMessageActivity, threadSortKey = "updated_at", + enableBackgroundThreadMetadataHydration = false, onThreadCodexMetadataDetected, }: UseThreadsOptions) { const maxItemsPerThread = @@ -546,6 +548,7 @@ export function useThreads({ dispatch, itemsByThread: state.itemsByThread, threadsByWorkspace: state.threadsByWorkspace, + tokenUsageThreadIdsByWorkspace: state.tokenUsageThreadIdsByWorkspace, activeThreadIdByWorkspace: state.activeThreadIdByWorkspace, activeTurnIdByThread: state.activeTurnIdByThread, threadParentById: state.threadParentById, @@ -561,6 +564,7 @@ export function useThreads({ updateThreadParent, onSubagentThreadDetected, onThreadCodexMetadataDetected, + enableBackgroundMetadataHydration: enableBackgroundThreadMetadataHydration, }); const ensureWorkspaceRuntimeCodexArgsBestEffort = useCallback( @@ -848,6 +852,7 @@ export function useThreads({ activeTurnIdByThread: state.activeTurnIdByThread, turnDiffByThread: state.turnDiffByThread, tokenUsageByThread: state.tokenUsageByThread, + tokenUsageThreadIdsByWorkspace: state.tokenUsageThreadIdsByWorkspace, rateLimitsByWorkspace: state.rateLimitsByWorkspace, accountByWorkspace: state.accountByWorkspace, planByThread: state.planByThread, diff --git a/src/features/threads/hooks/useThreadsReducer.test.ts b/src/features/threads/hooks/useThreadsReducer.test.ts index f4077d9a7..74e9947ba 100644 --- a/src/features/threads/hooks/useThreadsReducer.test.ts +++ b/src/features/threads/hooks/useThreadsReducer.test.ts @@ -532,6 +532,28 @@ describe("threadReducer", () => { }, activeThreadIdByWorkspace: { "ws-1": "thread-1" }, turnDiffByThread: { "thread-1": "diff --git a/file.ts b/file.ts" }, + tokenUsageByThread: { + "thread-1": { + total: { + totalTokens: 100, + inputTokens: 60, + cachedInputTokens: 5, + outputTokens: 40, + reasoningOutputTokens: 2, + }, + last: { + totalTokens: 20, + inputTokens: 12, + cachedInputTokens: 1, + outputTokens: 8, + reasoningOutputTokens: 0, + }, + modelContextWindow: 200000, + }, + }, + tokenUsageThreadIdsByWorkspace: { + "ws-1": ["thread-1", "thread-2"], + }, }; const next = threadReducer(base, { @@ -541,6 +563,8 @@ describe("threadReducer", () => { }); expect(next.turnDiffByThread["thread-1"]).toBeUndefined(); + expect(next.tokenUsageByThread["thread-1"]).toBeUndefined(); + expect(next.tokenUsageThreadIdsByWorkspace["ws-1"]).toEqual(["thread-2"]); }); it("hides background threads and keeps them hidden on future syncs", () => { @@ -550,6 +574,7 @@ describe("threadReducer", () => { threadId: "thread-bg", }); expect(withThread.threadsByWorkspace["ws-1"]?.some((t) => t.id === "thread-bg")).toBe(true); + expect(withThread.tokenUsageThreadIdsByWorkspace["ws-1"]).toContain("thread-bg"); const hidden = threadReducer(withThread, { type: "hideThread", @@ -557,6 +582,7 @@ describe("threadReducer", () => { threadId: "thread-bg", }); expect(hidden.threadsByWorkspace["ws-1"]?.some((t) => t.id === "thread-bg")).toBe(false); + expect(hidden.tokenUsageThreadIdsByWorkspace["ws-1"]?.includes("thread-bg")).toBe(false); const synced = threadReducer(hidden, { type: "setThreads", @@ -572,6 +598,23 @@ describe("threadReducer", () => { expect(ids).not.toContain("thread-bg"); }); + it("stores deduplicated workspace usage thread ids and excludes hidden threads", () => { + const hidden = threadReducer(initialState, { + type: "hideThread", + workspaceId: "ws-1", + threadId: "thread-hidden", + }); + const next = threadReducer(hidden, { + type: "setWorkspaceTokenUsageThreadIds", + workspaceId: "ws-1", + threadIds: ["thread-1", "thread-1", "", "thread-hidden", "thread-2"], + }); + expect(next.tokenUsageThreadIdsByWorkspace["ws-1"]).toEqual([ + "thread-1", + "thread-2", + ]); + }); + it("preserves active, processing, and ancestor anchors on partial setThreads payloads", () => { const base: ThreadState = { ...initialState, diff --git a/src/features/threads/hooks/useThreadsReducer.ts b/src/features/threads/hooks/useThreadsReducer.ts index 6ec20bcfa..7cdfa6df3 100644 --- a/src/features/threads/hooks/useThreadsReducer.ts +++ b/src/features/threads/hooks/useThreadsReducer.ts @@ -42,6 +42,7 @@ export type ThreadState = { approvals: ApprovalRequest[]; userInputRequests: RequestUserInputRequest[]; tokenUsageByThread: Record; + tokenUsageThreadIdsByWorkspace: Record; rateLimitsByWorkspace: Record; accountByWorkspace: Record; planByThread: Record; @@ -145,6 +146,11 @@ export type ThreadAction = workspaceId: string; } | { type: "setThreadTokenUsage"; threadId: string; tokenUsage: ThreadTokenUsage } + | { + type: "setWorkspaceTokenUsageThreadIds"; + workspaceId: string; + threadIds: string[]; + } | { type: "setRateLimits"; workspaceId: string; @@ -186,6 +192,7 @@ export const initialState: ThreadState = { approvals: [], userInputRequests: [], tokenUsageByThread: {}, + tokenUsageThreadIdsByWorkspace: {}, rateLimitsByWorkspace: {}, accountByWorkspace: {}, planByThread: {}, diff --git a/src/features/threads/utils/threadCodexMetadata.test.ts b/src/features/threads/utils/threadCodexMetadata.test.ts index 3b0f28598..f358621e1 100644 --- a/src/features/threads/utils/threadCodexMetadata.test.ts +++ b/src/features/threads/utils/threadCodexMetadata.test.ts @@ -63,4 +63,38 @@ describe("extractThreadCodexMetadata", () => { effort: null, }); }); + + it("ignores placeholder model ids and keeps searching nested metadata", () => { + const metadata = extractThreadCodexMetadata({ + model: "unknown", + turns: [ + { + items: [ + { + payload: { + info: { + model_name: "gpt-5.2-codex", + }, + }, + }, + ], + }, + ], + }); + + expect(metadata.modelId).toBe("gpt-5.2-codex"); + }); + + it("returns null model for placeholder-only metadata", () => { + const metadata = extractThreadCodexMetadata({ + model: "default", + turns: [ + { + items: [{ payload: { model: "unknown" } }], + }, + ], + }); + + expect(metadata.modelId).toBeNull(); + }); }); diff --git a/src/features/threads/utils/threadCodexMetadata.ts b/src/features/threads/utils/threadCodexMetadata.ts index 32dc9ba58..817307d1d 100644 --- a/src/features/threads/utils/threadCodexMetadata.ts +++ b/src/features/threads/utils/threadCodexMetadata.ts @@ -24,6 +24,27 @@ function normalizeEffort(value: string | null): string | null { return normalized; } +function normalizeModelId(value: string | null): string | null { + if (!value) { + return null; + } + const trimmed = value.trim(); + if (!trimmed) { + return null; + } + const normalized = trimmed.toLowerCase(); + if ( + normalized === "unknown" || + normalized === "default" || + normalized === "auto" || + normalized === "current" || + normalized === "inherit" + ) { + return null; + } + return trimmed; +} + function pickString( record: Record, keys: readonly string[], @@ -77,7 +98,7 @@ function extractFromRecord(record: Record): { for (const container of containers) { if (!modelId) { - modelId = pickString(container, MODEL_KEYS); + modelId = normalizeModelId(pickString(container, MODEL_KEYS)); } if (!effort) { effort = normalizeEffort(pickString(container, EFFORT_KEYS)); diff --git a/src/services/tauri.test.ts b/src/services/tauri.test.ts index 808cfd213..4295c749b 100644 --- a/src/services/tauri.test.ts +++ b/src/services/tauri.test.ts @@ -18,6 +18,7 @@ import { getOpenAppIcon, listThreads, listMcpServerStatus, + localThreadUsageSnapshot, readGlobalAgentsMd, readGlobalCodexConfigToml, listWorkspaces, @@ -283,6 +284,21 @@ describe("tauri invoke wrappers", () => { }); }); + it("maps threadIds/workspacePath for local_thread_usage_snapshot", async () => { + const invokeMock = vi.mocked(invoke); + invokeMock.mockResolvedValueOnce({ + updatedAt: 123, + usageByThread: {}, + }); + + await localThreadUsageSnapshot(["thread-1", "thread-2"], "/tmp/codex"); + + expect(invokeMock).toHaveBeenCalledWith("local_thread_usage_snapshot", { + threadIds: ["thread-1", "thread-2"], + workspacePath: "/tmp/codex", + }); + }); + it("maps workspaceId/cursor/limit/threadId for apps_list", async () => { const invokeMock = vi.mocked(invoke); invokeMock.mockResolvedValueOnce({}); diff --git a/src/services/tauri.ts b/src/services/tauri.ts index 07919a72f..277ef1729 100644 --- a/src/services/tauri.ts +++ b/src/services/tauri.ts @@ -7,6 +7,7 @@ import type { CodexDoctorResult, DictationModelStatus, DictationSessionState, + LocalThreadUsageSnapshot, LocalUsageSnapshot, TcpDaemonStatus, TailscaleDaemonCommandPreview, @@ -709,6 +710,19 @@ export async function localUsageSnapshot( return invoke("local_usage_snapshot", payload); } +export async function localThreadUsageSnapshot( + threadIds: string[], + workspacePath?: string | null, +): Promise { + const payload: { threadIds: string[]; workspacePath?: string } = { + threadIds, + }; + if (workspacePath) { + payload.workspacePath = workspacePath; + } + return invoke("local_thread_usage_snapshot", payload); +} + export async function getModelList(workspaceId: string) { return invoke("model_list", { workspaceId }); } diff --git a/src/styles/sidebar.css b/src/styles/sidebar.css index 0d319f0cf..bf698c28a 100644 --- a/src/styles/sidebar.css +++ b/src/styles/sidebar.css @@ -813,12 +813,19 @@ .workspace-name-row { display: flex; - align-items: center; + align-items: flex-start; gap: 8px; justify-content: space-between; min-width: 0; } +.workspace-title-block { + display: flex; + flex-direction: column; + gap: 1px; + min-width: 0; +} + .workspace-title { display: flex; align-items: center; @@ -826,6 +833,16 @@ min-width: 0; } +.workspace-token-usage { + font-size: 10px; + color: var(--text-faint); + line-height: 1.2; + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + .workspace-toggle { border: none; background: transparent; @@ -946,7 +963,7 @@ .thread-row { display: flex; - align-items: center; + align-items: flex-start; gap: 6px; padding: 4px 10px 4px calc(6px + var(--thread-indent, 0px)); border-radius: 6px; @@ -965,11 +982,13 @@ border-radius: 999px; display: inline-block; margin-left: 4px; + margin-top: 3px; } .thread-pin-icon { font-size: 10px; margin-right: 2px; + margin-top: 1px; opacity: 0.7; } @@ -1009,9 +1028,26 @@ color: var(--text-strong); } -.thread-name { +.thread-text { flex: 1; min-width: 0; + display: flex; + flex-direction: column; + gap: 1px; +} + +.thread-name { + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.thread-token-usage { + font-size: 10px; + color: var(--text-faint); + line-height: 1.2; + min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; @@ -1022,6 +1058,7 @@ display: inline-flex; align-items: center; gap: 6px; + align-self: flex-start; } .thread-args-badge { @@ -1233,13 +1270,29 @@ text-overflow: ellipsis; } +.worktree-label-block { + display: flex; + flex-direction: column; + gap: 1px; + min-width: 0; +} + +.worktree-token-usage { + font-size: 10px; + color: var(--text-faint); + line-height: 1.2; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} + .worktree-row.deleting .worktree-label { color: var(--text-faint); } .worktree-actions { display: inline-flex; - align-items: center; + align-items: flex-start; gap: 6px; } diff --git a/src/types.ts b/src/types.ts index 6eebb7ebf..6e6c4080a 100644 --- a/src/types.ts +++ b/src/types.ts @@ -236,6 +236,9 @@ export type AppSettings = { uiScale: number; theme: ThemePreference; usageShowRemaining: boolean; + showThreadTokenUsage: boolean; + threadTokenUsageShowFull: boolean; + threadTokenUsageExcludeCache: boolean; showMessageFilePath: boolean; chatHistoryScrollbackItems: number | null; threadTitleAutogenerationEnabled: boolean; @@ -502,6 +505,11 @@ export type ThreadTokenUsage = { modelContextWindow: number | null; }; +export type LocalThreadUsageSnapshot = { + updatedAt: number; + usageByThread: Record; +}; + export type LocalUsageDay = { day: string; inputTokens: number; diff --git a/src/utils/tokenUsage.test.ts b/src/utils/tokenUsage.test.ts new file mode 100644 index 000000000..ce79b2bd0 --- /dev/null +++ b/src/utils/tokenUsage.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, it } from "vitest"; +import { formatCompactTokenCount } from "./tokenUsage"; + +describe("formatCompactTokenCount", () => { + it("returns null for non-positive values", () => { + expect(formatCompactTokenCount(0)).toBeNull(); + expect(formatCompactTokenCount(-1)).toBeNull(); + }); + + it("formats small values without unit suffixes", () => { + expect(formatCompactTokenCount(12)).toBe("12"); + expect(formatCompactTokenCount(999)).toBe("999"); + }); + + it("formats thousands, millions, and billions compactly", () => { + expect(formatCompactTokenCount(1_200)).toBe("1.2K"); + expect(formatCompactTokenCount(2_300_000)).toBe("2.3M"); + expect(formatCompactTokenCount(4_100_000_000)).toBe("4.1B"); + }); + + it("promotes boundary-rounded values to the next unit", () => { + expect(formatCompactTokenCount(999_950)).toBe("1.0M"); + expect(formatCompactTokenCount(999_950_000)).toBe("1.0B"); + }); +}); diff --git a/src/utils/tokenUsage.ts b/src/utils/tokenUsage.ts new file mode 100644 index 000000000..c5bdfd11f --- /dev/null +++ b/src/utils/tokenUsage.ts @@ -0,0 +1,42 @@ +const COMPACT_TOKEN_UNITS = [ + { divisor: 1_000_000_000, suffix: "B" }, + { divisor: 1_000_000, suffix: "M" }, + { divisor: 1_000, suffix: "K" }, +] as const; + +export function formatCompactTokenCount(value: number): string | null { + if (!Number.isFinite(value) || value <= 0) { + return null; + } + + for (let index = 0; index < COMPACT_TOKEN_UNITS.length; index += 1) { + const unit = COMPACT_TOKEN_UNITS[index]; + if (value < unit.divisor) { + continue; + } + + const scaled = value / unit.divisor; + const decimals = scaled >= 10 ? 0 : 1; + const factor = 10 ** decimals; + const rounded = Math.round(scaled * factor) / factor; + + if (rounded >= 1000 && index > 0) { + const promotedUnit = COMPACT_TOKEN_UNITS[index - 1]; + const promotedScaled = value / promotedUnit.divisor; + const promotedDecimals = promotedScaled >= 10 ? 0 : 1; + const promotedFactor = 10 ** promotedDecimals; + const promotedRounded = + Math.round(promotedScaled * promotedFactor) / promotedFactor; + const promotedText = + promotedDecimals === 0 + ? promotedRounded.toFixed(0) + : promotedRounded.toFixed(1); + return `${promotedText}${promotedUnit.suffix}`; + } + + const text = decimals === 0 ? rounded.toFixed(0) : rounded.toFixed(1); + return `${text}${unit.suffix}`; + } + + return String(Math.round(value)); +}