Refactor /Usage into smaller modules (#3324)

Co-authored-by: Akhilesh Bontala <[email protected]>

Author: akhileshb

crates/chat-cli/src/cli/chat/cli/usage.rs

@@ -0,0 +1,37 @@
+use clap::Args;
+
+use crate::cli::chat::token_counter::TokenCount;
+use crate::cli::chat::{ChatError, ChatSession, ChatState};
+use crate::os::Os;
+
+pub mod usage_data_provider;
+pub mod usage_renderer;
+
+/// Detailed usage data for context window analysis
+#[derive(Debug)]
+pub struct DetailedUsageData {
+    pub total_tokens: TokenCount,
+    pub context_tokens: TokenCount,
+    pub assistant_tokens: TokenCount,
+    pub user_tokens: TokenCount,
+    pub tools_tokens: TokenCount,
+    pub context_window_size: usize,
+    pub dropped_context_files: Vec<(String, String)>,
+}
+
+/// Arguments for the usage command that displays token usage statistics and context window
+/// information.
+///
+/// This command shows how many tokens are being used by different components (context files, tools,
+/// assistant responses, and user prompts) within the current chat session's context window.
+#[deny(missing_docs)]
+#[derive(Debug, PartialEq, Args)]
+pub struct UsageArgs;
+
+impl UsageArgs {
+    pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
+        let usage_data = usage_data_provider::get_detailed_usage_data(session, os).await?;
+        usage_renderer::render_context_window(&usage_data, session).await?;
+        Ok(ChatState::PromptUser { skip_printing_tools: true })
+    }
+}

crates/chat-cli/src/cli/chat/cli/usage/mod.rs

@@ -0,0 +1,41 @@
+use crate::cli::chat::cli::model::context_window_tokens;
+use crate::cli::chat::token_counter::{CharCount, TokenCount};
+use crate::cli::chat::{ChatError, ChatSession};
+use crate::os::Os;
+
+/// Get detailed usage data for context window analysis
+pub(super) async fn get_detailed_usage_data(session: &mut ChatSession, os: &Os) -> Result<super::DetailedUsageData, ChatError> {
+    let context_window_size = context_window_tokens(session.conversation.model_info.as_ref());
+
+    let state = session
+        .conversation
+        .backend_conversation_state(os, true, &mut std::io::stderr())
+        .await?;
+
+    let data = state.calculate_conversation_size();
+    let tool_specs_json: String = state
+        .tools
+        .values()
+        .filter_map(|s| serde_json::to_string(s).ok())
+        .collect::<Vec<String>>()
+        .join("");
+    let tools_char_count: CharCount = tool_specs_json.len().into();
+    let total_tokens: TokenCount =
+        (data.context_messages + data.user_messages + data.assistant_messages + tools_char_count).into();
+
+    Ok(super::DetailedUsageData {
+        total_tokens,
+        context_tokens: data.context_messages.into(),
+        assistant_tokens: data.assistant_messages.into(),
+        user_tokens: data.user_messages.into(),
+        tools_tokens: tools_char_count.into(),
+        context_window_size,
+        dropped_context_files: state.dropped_context_files,
+    })
+}
+
+/// Get total usage percentage (external API)
+pub async fn get_total_usage_percentage(session: &mut ChatSession, os: &Os) -> Result<f32, ChatError> {
+    let data = get_detailed_usage_data(session, os).await?;
+    Ok((data.total_tokens.value() as f32 / data.context_window_size as f32) * 100.0)
+}