Skip to content

Commit fe94368

Browse files
evanliu048evanliu048
authored
fix: Split /usage and /context commands for better UX (#3374)
* split out billin info * recover * refactor
1 parent 2866262 commit fe94368

File tree

8 files changed

+313
-330
lines changed

8 files changed

+313
-330
lines changed

‎crates/chat-cli/src/cli/chat/cli/context/context_data_provider.rs‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
//! Provides context window usage data and analysis
2+
3+
use crate::cli::chat::cli::model::context_window_tokens;
4+
use crate::cli::chat::token_counter::{
5+
CharCount,
6+
TokenCount,
7+
};
8+
use crate::cli::chat::{
9+
ChatError,
10+
ChatSession,
11+
};
12+
use crate::os::Os;
13+
14+
/// Detailed context window usage data
15+
#[derive(Debug)]
16+
pub struct ContextWindowData {
17+
pub total_tokens: TokenCount,
18+
pub context_tokens: TokenCount,
19+
pub assistant_tokens: TokenCount,
20+
pub user_tokens: TokenCount,
21+
pub tools_tokens: TokenCount,
22+
pub context_window_size: usize,
23+
pub dropped_context_files: Vec<(String, String)>,
24+
}
25+
26+
/// Get detailed usage data for context window analysis
27+
pub(super) async fn get_detailed_context_data(
28+
session: &mut ChatSession,
29+
os: &Os,
30+
) -> Result<ContextWindowData, ChatError> {
31+
let context_window_size = context_window_tokens(session.conversation.model_info.as_ref());
32+
33+
let state = session
34+
.conversation
35+
.backend_conversation_state(os, true, &mut std::io::stderr())
36+
.await?;
37+
38+
let data = state.calculate_conversation_size();
39+
let tool_specs_json: String = state
40+
.tools
41+
.values()
42+
.filter_map(|s| serde_json::to_string(s).ok())
43+
.collect::<Vec<String>>()
44+
.join("");
45+
let tools_char_count: CharCount = tool_specs_json.len().into();
46+
let total_tokens: TokenCount =
47+
(data.context_messages + data.user_messages + data.assistant_messages + tools_char_count).into();
48+
49+
Ok(ContextWindowData {
50+
total_tokens,
51+
context_tokens: data.context_messages.into(),
52+
assistant_tokens: data.assistant_messages.into(),
53+
user_tokens: data.user_messages.into(),
54+
tools_tokens: tools_char_count.into(),
55+
context_window_size,
56+
dropped_context_files: state.dropped_context_files,
57+
})
58+
}
59+
60+
/// Get total context usage percentage (external API)
61+
pub async fn get_total_context_usage_percentage(session: &mut ChatSession, os: &Os) -> Result<f32, ChatError> {
62+
let data = get_detailed_context_data(session, os).await?;
63+
Ok((data.total_tokens.value() as f32 / data.context_window_size as f32) * 100.0)
64+
}

‎crates/chat-cli/src/cli/chat/cli/context/context_renderer.rs‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
//! Rendering for context window information
2+
3+
use crossterm::style::Attribute;
4+
use crossterm::{
5+
execute,
6+
queue,
7+
style,
8+
};
9+
10+
use super::context_data_provider::ContextWindowData;
11+
use crate::cli::chat::token_counter::TokenCount;
12+
use crate::cli::chat::{
13+
ChatError,
14+
ChatSession,
15+
};
16+
use crate::theme::StyledText;
17+
18+
/// Calculate context percentage from token counts (private utility)
19+
fn calculate_context_percentage(tokens: TokenCount, context_window_size: usize) -> f32 {
20+
(tokens.value() as f32 / context_window_size as f32) * 100.0
21+
}
22+
23+
/// Render context window information section
24+
pub async fn render_context_window(
25+
context_data: &ContextWindowData,
26+
session: &mut ChatSession,
27+
) -> Result<(), ChatError> {
28+
if !context_data.dropped_context_files.is_empty() {
29+
execute!(
30+
session.stderr,
31+
StyledText::warning_fg(),
32+
style::Print("\nSome context files are dropped due to size limit, please run "),
33+
StyledText::success_fg(),
34+
style::Print("/context show "),
35+
StyledText::warning_fg(),
36+
style::Print("to learn more.\n"),
37+
StyledText::reset(),
38+
)?;
39+
}
40+
41+
let window_width = session.terminal_width();
42+
// set a max width for the progress bar for better aesthetic
43+
let progress_bar_width = std::cmp::min(window_width, 80);
44+
45+
let context_width = ((context_data.context_tokens.value() as f64 / context_data.context_window_size as f64)
46+
* progress_bar_width as f64) as usize;
47+
let assistant_width = ((context_data.assistant_tokens.value() as f64 / context_data.context_window_size as f64)
48+
* progress_bar_width as f64) as usize;
49+
let tools_width = ((context_data.tools_tokens.value() as f64 / context_data.context_window_size as f64)
50+
* progress_bar_width as f64) as usize;
51+
let user_width = ((context_data.user_tokens.value() as f64 / context_data.context_window_size as f64)
52+
* progress_bar_width as f64) as usize;
53+
54+
let left_over_width = progress_bar_width
55+
- std::cmp::min(
56+
context_width + assistant_width + user_width + tools_width,
57+
progress_bar_width,
58+
);
59+
60+
let is_overflow = (context_width + assistant_width + user_width + tools_width) > progress_bar_width;
61+
62+
let total_percentage = calculate_context_percentage(context_data.total_tokens, context_data.context_window_size);
63+
64+
if is_overflow {
65+
queue!(
66+
session.stderr,
67+
style::Print(format!(
68+
"\nCurrent context window ({} of {}k tokens used)\n",
69+
context_data.total_tokens,
70+
context_data.context_window_size / 1000
71+
)),
72+
StyledText::error_fg(),
73+
style::Print("█".repeat(progress_bar_width)),
74+
StyledText::reset(),
75+
style::Print(" "),
76+
style::Print(format!("{total_percentage:.2}%")),
77+
)?;
78+
} else {
79+
queue!(
80+
session.stderr,
81+
style::Print(format!(
82+
"\nCurrent context window ({} of {}k tokens used)\n",
83+
context_data.total_tokens,
84+
context_data.context_window_size / 1000
85+
)),
86+
// Context files
87+
StyledText::brand_fg(),
88+
// add a nice visual to mimic "tiny" progress, so the overrall progress bar doesn't look too
89+
// empty
90+
style::Print(
91+
"|".repeat(if context_width == 0 && context_data.context_tokens.value() > 0 {
92+
1
93+
} else {
94+
0
95+
})
96+
),
97+
style::Print("█".repeat(context_width)),
98+
// Tools
99+
StyledText::error_fg(),
100+
style::Print(
101+
"|".repeat(if tools_width == 0 && context_data.tools_tokens.value() > 0 {
102+
1
103+
} else {
104+
0
105+
})
106+
),
107+
style::Print("█".repeat(tools_width)),
108+
// Assistant responses
109+
StyledText::info_fg(),
110+
style::Print(
111+
"|".repeat(if assistant_width == 0 && context_data.assistant_tokens.value() > 0 {
112+
1
113+
} else {
114+
0
115+
})
116+
),
117+
style::Print("█".repeat(assistant_width)),
118+
// User prompts
119+
StyledText::emphasis_fg(),
120+
style::Print("|".repeat(if user_width == 0 && context_data.user_tokens.value() > 0 {
121+
1
122+
} else {
123+
0
124+
})),
125+
style::Print("█".repeat(user_width)),
126+
StyledText::secondary_fg(),
127+
style::Print("█".repeat(left_over_width)),
128+
style::Print(" "),
129+
StyledText::reset(),
130+
style::Print(format!("{total_percentage:.2}%")),
131+
)?;
132+
}
133+
134+
execute!(session.stderr, style::Print("\n\n"))?;
135+
136+
queue!(
137+
session.stderr,
138+
StyledText::brand_fg(),
139+
style::Print("█ Context files: "),
140+
StyledText::reset(),
141+
style::Print(format!(
142+
"~{} tokens ({:.2}%)\n",
143+
context_data.context_tokens,
144+
calculate_context_percentage(context_data.context_tokens, context_data.context_window_size)
145+
)),
146+
StyledText::error_fg(),
147+
style::Print("█ Tools: "),
148+
StyledText::reset(),
149+
style::Print(format!(
150+
" ~{} tokens ({:.2}%)\n",
151+
context_data.tools_tokens,
152+
calculate_context_percentage(context_data.tools_tokens, context_data.context_window_size)
153+
)),
154+
StyledText::info_fg(),
155+
style::Print("█ Kiro responses: "),
156+
StyledText::reset(),
157+
style::Print(format!(
158+
" ~{} tokens ({:.2}%)\n",
159+
context_data.assistant_tokens,
160+
calculate_context_percentage(context_data.assistant_tokens, context_data.context_window_size)
161+
)),
162+
StyledText::emphasis_fg(),
163+
style::Print("█ Your prompts: "),
164+
StyledText::reset(),
165+
style::Print(format!(
166+
" ~{} tokens ({:.2}%)\n\n",
167+
context_data.user_tokens,
168+
calculate_context_percentage(context_data.user_tokens, context_data.context_window_size)
169+
)),
170+
)?;
171+
172+
queue!(
173+
session.stderr,
174+
style::SetAttribute(Attribute::Bold),
175+
style::Print("\n💡 Pro Tips:\n"),
176+
StyledText::reset_attributes(),
177+
StyledText::secondary_fg(),
178+
style::Print("Run "),
179+
StyledText::success_fg(),
180+
style::Print("/compact"),
181+
StyledText::secondary_fg(),
182+
style::Print(" to replace the conversation history with its summary\n"),
183+
style::Print("Run "),
184+
StyledText::success_fg(),
185+
style::Print("/clear"),
186+
StyledText::secondary_fg(),
187+
style::Print(" to erase the entire chat history\n"),
188+
style::Print("Run "),
189+
StyledText::success_fg(),
190+
style::Print("/context show"),
191+
StyledText::secondary_fg(),
192+
style::Print(" to see tokens per context file\n\n"),
193+
StyledText::reset(),
194+
)?;
195+
196+
Ok(())
197+
}

crates/chat-cli/src/cli/chat/cli/context.rs renamed to crates/chat-cli/src/cli/chat/cli/context/mod.rs

Lines changed: 40 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,18 @@
11
use std::collections::HashSet;
22

3-
use clap::Subcommand;
3+
use clap::{
4+
Args,
5+
Subcommand,
6+
};
47
use crossterm::style::Attribute;
58
use crossterm::{
69
execute,
710
style,
811
};
912

13+
pub mod context_data_provider;
14+
pub mod context_renderer;
15+
1016
use crate::cli::chat::consts::AGENT_FORMAT_HOOKS_DOC_URL;
1117
use crate::cli::chat::context::{
1218
ContextFilePath,
@@ -27,12 +33,43 @@ use crate::constants::help_text::{
2733
use crate::os::Os;
2834
use crate::theme::StyledText;
2935

30-
#[deny(missing_docs)]
31-
#[derive(Debug, PartialEq, Subcommand)]
36+
/// Context command for viewing token usage and managing context files.
37+
/// Without subcommands: displays context window token usage breakdown by component
38+
/// (context files, tools, assistant responses, user prompts).
39+
/// With subcommands: manages context file rules (show, add, remove, clear).
40+
#[derive(Debug, PartialEq, Args)]
3241
#[command(
3342
about = context_description(),
3443
before_long_help = context_long_help()
3544
)]
45+
pub struct ContextArgs {
46+
#[command(subcommand)]
47+
pub subcommand: Option<ContextSubcommand>,
48+
}
49+
50+
impl ContextArgs {
51+
pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
52+
match self.subcommand {
53+
Some(subcommand) => subcommand.execute(os, session).await,
54+
None => {
55+
// No subcommand - show context window usage (replaces /usage --context)
56+
let contexet_data = context_data_provider::get_detailed_context_data(session, os).await?;
57+
context_renderer::render_context_window(&contexet_data, session).await?;
58+
59+
Ok(ChatState::PromptUser {
60+
skip_printing_tools: true,
61+
})
62+
},
63+
}
64+
}
65+
66+
pub fn subcommand_name(&self) -> Option<&'static str> {
67+
self.subcommand.as_ref().map(|s| s.name())
68+
}
69+
}
70+
71+
#[deny(missing_docs)]
72+
#[derive(Debug, PartialEq, Subcommand)]
3673
/// Context subcommands
3774
pub enum ContextSubcommand {
3875
/// Display the context rule configuration and matched files

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@ use changelog::ChangelogArgs;
2626
use clap::Parser;
2727
use clear::ClearArgs;
2828
use compact::CompactArgs;
29-
use context::ContextSubcommand;
3029
use editor::EditorArgs;
3130
use experiment::ExperimentArgs;
3231
use hooks::HooksArgs;
@@ -44,6 +43,7 @@ use todos::TodoSubcommand;
4443
use tools::ToolsArgs;
4544

4645
use crate::cli::chat::cli::checkpoint::CheckpointSubcommand;
46+
use crate::cli::chat::cli::context::ContextArgs;
4747
use crate::cli::chat::cli::subscribe::SubscribeArgs;
4848
use crate::cli::chat::cli::usage::UsageArgs;
4949
use crate::cli::chat::consts::AGENT_MIGRATION_DOC_URL;
@@ -72,9 +72,8 @@ pub enum SlashCommand {
7272
Agent(AgentSubcommand),
7373
#[command(hide = true)]
7474
Profile,
75-
/// Manage context files for the chat session
76-
#[command(subcommand)]
77-
Context(ContextSubcommand),
75+
/// Manage context files and view context window usage
76+
Context(ContextArgs),
7877
/// (Beta) Manage knowledge base for persistent context storage. Requires "q settings
7978
/// chat.enableKnowledge true"
8079
#[command(subcommand, hide = true)]
@@ -99,7 +98,7 @@ pub enum SlashCommand {
9998
Prompts(PromptsArgs),
10099
/// View context hooks
101100
Hooks(HooksArgs),
102-
/// Show current session's context window usage
101+
/// Show billing and credits information
103102
Usage(UsageArgs),
104103
/// See mcp server loaded
105104
Mcp(McpArgs),
@@ -237,7 +236,7 @@ impl SlashCommand {
237236
pub fn subcommand_name(&self) -> Option<&'static str> {
238237
match self {
239238
SlashCommand::Agent(sub) => Some(sub.name()),
240-
SlashCommand::Context(sub) => Some(sub.name()),
239+
SlashCommand::Context(args) => args.subcommand_name(),
241240
SlashCommand::Knowledge(sub) => Some(sub.name()),
242241
SlashCommand::Tools(arg) => arg.subcommand_name(),
243242
SlashCommand::Prompts(arg) => arg.subcommand_name(),

0 commit comments

Comments
 (0)