public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Replace debug memory stats footer with AI request token count

Browse Source 
When 'Show Memory Stats' debug mode is enabled, the block footer
now displays 'Request tokens: N' (or 'Request tokens: --') instead
of the old grid/flat storage size breakdown.

- Add last_block_token_usage_by_model to AIConversation, tracking
  token usage accumulated since the most recent user-initiated request
- Add last_block_total_tokens() accessor summing across all models
- BlockListElement reads this value per block and renders it in the
  debug footer text element
This commit is contained in:
Ryan Ward
2026-05-13 09:55:04 -05:00
parent 029e90be2b
commit ed53aa99eb
2 changed files with 147 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/agent/conversation.rs
+106 -52
View File
@@ -14,7 +14,7 @@ use crate::terminal::model::block::{
};
use crate::ai::agent::api::convert_conversation::{
compute_time_to_first_token_ms_from_messages, ConvertToExchanges,
ConvertToExchanges, compute_time_to_first_token_ms_from_messages,
};
use ai::document::AIDocumentId;
use chrono::{DateTime, Local, TimeZone};
@@ -24,52 +24,51 @@ use std::collections::HashSet;
use std::{collections::HashMap, fmt::Display};
use super::task_store::TaskStore;
use uuid::Uuid;
use vec1::{Size0Error, Vec1};
use galaxy_core::command::ExitCode;
use galaxy_core::execution_mode::AppExecutionMode;
use galaxy_core::features::FeatureFlag;
use galaxy_core::send_telemetry_from_ctx;
use galaxy_core::ui::appearance::Appearance;
use galaxy_core::ui::theme::color::internal_colors;
use galaxy_core::ui::theme::WarpTheme;
use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, response_event::stream_finished::TokenUsage};
use galaxy_core::ui::theme::color::internal_colors;
use galaxyui::color::ColorU;
use galaxyui::{EntityId, ModelContext, SingletonEntity};
use uuid::Uuid;
use vec1::{Size0Error, Vec1};
use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, response_event::stream_finished::TokenUsage};
use crate::ai::agent::{AIIdentifiers, CancellationReason};
use crate::{
BlocklistAIHistoryModel, GlobalResourceHandlesProvider,
ai::{
agent::{
AIAgentOutputMessage, AIAgentOutputMessageType, MessageToAIAgentOutputMessageError,
icons::{
failed_icon, gray_stop_icon, in_progress_icon, succeeded_icon, yellow_stop_icon,
},
todos::AIAgentTodoList,
AIAgentOutputMessage, AIAgentOutputMessageType, MessageToAIAgentOutputMessageError,
},
blocklist::BlocklistAIHistoryEvent,
},
persistence::{
model::{AgentConversationData, PersistedAutoexecuteMode},
ModelEvent,
model::{AgentConversationData, PersistedAutoexecuteMode},
},
ui_components::icons::Icon,
BlocklistAIHistoryModel, GlobalResourceHandlesProvider,
};
use super::task::{ExtractMessagesError, UpdateTaskError, UpgradeOptimisticTaskError};
use super::{
api::ServerConversationToken,
task::{
derive_todo_lists_from_root_task,
helper::*,
transaction::{SavedTask, Transaction},
Task, TaskId,
},
AIAgentAction, AIAgentActionId, AIAgentContext, AIAgentExchange, AIAgentExchangeId,
AIAgentInput, AIAgentOutputStatus, AIAgentTodo, AIAgentTodoId, FinishedAIAgentOutput,
MessageId, RenderableAIError, RequestCost,
api::ServerConversationToken,
task::{
Task, TaskId, derive_todo_lists_from_root_task,
helper::*,
transaction::{SavedTask, Transaction},
},
};
use super::{
AIAgentOutput, OutputModelInfo, ServerOutputId, Shared, SuggestedLoggingId, Suggestions,
@@ -203,6 +202,7 @@ pub struct AIConversation {
total_request_cost: RequestCost,
total_token_usage_by_model: HashMap<String, TokenUsage>,
last_block_token_usage_by_model: HashMap<String, TokenUsage>,
/// Fallback title used when no task description or initial query exists.
fallback_display_title: Option<String>,
@@ -277,6 +277,7 @@ impl AIConversation {
dismissed_suggestion_ids: Default::default(),
total_request_cost: RequestCost::new(0.),
total_token_usage_by_model: Default::default(),
last_block_token_usage_by_model: Default::default(),
fallback_display_title: None,
artifacts: Vec::new(),
parent_agent_id: None,
@@ -458,6 +459,7 @@ impl AIConversation {
dismissed_suggestion_ids: Default::default(),
total_request_cost: RequestCost::new(0.),
total_token_usage_by_model: Default::default(),
last_block_token_usage_by_model: Default::default(),
optimistic_cli_subagent_subtask_id: None,
fallback_display_title: None,
artifacts,
@@ -478,11 +480,16 @@ impl AIConversation {
&self.bedrock_message_history
}
pub fn bedrock_message_history_mut(&mut self) -> &mut Vec<crate::ai::bedrock::convert::ConversationMessage> {
pub fn bedrock_message_history_mut(
&mut self,
) -> &mut Vec<crate::ai::bedrock::convert::ConversationMessage> {
&mut self.bedrock_message_history
}
pub fn append_to_bedrock_history(&mut self, messages: Vec<crate::ai::bedrock::convert::ConversationMessage>) {
pub fn append_to_bedrock_history(
&mut self,
messages: Vec<crate::ai::bedrock::convert::ConversationMessage>,
) {
self.bedrock_message_history.extend(messages);
}
@@ -525,6 +532,13 @@ impl AIConversation {
.map(|credits| (credits * 10.0).round() / 10.0)
}
pub fn last_block_total_tokens(&self) -> u32 {
self.last_block_token_usage_by_model
.values()
.map(Self::total_tokens_for_usage)
.sum()
}
/// Time to first token for the last completed set of agent responses
/// since the most recent user query
pub fn time_to_first_token_for_last_user_query_ms(&self) -> i64 {
@@ -1548,6 +1562,9 @@ impl AIConversation {
usage_metadata: Option<stream_finished::ConversationUsageMetadata>,
was_user_initiated_request: bool,
) -> Result<(), UpdateConversationError> {
if was_user_initiated_request {
self.last_block_token_usage_by_model.clear();
}
for usage in token_usage.into_iter() {
let entry = self
.total_token_usage_by_model
@@ -1566,6 +1583,24 @@ impl AIConversation {
entry.input_cache_read += usage.input_cache_read;
entry.input_cache_write += usage.input_cache_write;
entry.cost_in_cents += usage.cost_in_cents;
let last_block_entry = self
.last_block_token_usage_by_model
.entry(usage.model_id.clone())
.or_insert_with(|| TokenUsage {
model_id: usage.model_id.clone(),
total_input: 0,
output: 0,
input_cache_read: 0,
input_cache_write: 0,
cost_in_cents: 0.0,
});
last_block_entry.total_input += usage.total_input;
last_block_entry.output += usage.output;
last_block_entry.input_cache_read += usage.input_cache_read;
last_block_entry.input_cache_write += usage.input_cache_write;
last_block_entry.cost_in_cents += usage.cost_in_cents;
}
if let Some(request_cost) = request_cost {
@@ -2497,63 +2532,78 @@ impl AIConversation {
mask: Some(mask),
}) => {
let task_id = TaskId::new(task_id);
log::info!("[bedrock-debug] AppendToMessageContent: task_id={:?}, message_id={:?}", task_id, message.id);
let exchange_id = match self
.added_exchanges_by_response
.get(response_stream_id)
{
log::info!(
"[bedrock-debug] AppendToMessageContent: task_id={:?}, message_id={:?}",
task_id,
message.id
);
let exchange_id = match self.added_exchanges_by_response.get(response_stream_id) {
Some(exchanges) => {
log::info!("[bedrock-debug] AppendToMessageContent: found {} exchanges for stream", exchanges.len());
log::info!(
"[bedrock-debug] AppendToMessageContent: found {} exchanges for stream",
exchanges.len()
);
for ex in exchanges.iter() {
log::info!("[bedrock-debug] exchange: task_id={:?}, exchange_id={:?}", ex.task_id, ex.exchange_id);
log::info!(
"[bedrock-debug] exchange: task_id={:?}, exchange_id={:?}",
ex.task_id,
ex.exchange_id
);
}
match exchanges.iter().find_map(|new_exchange| {
(new_exchange.task_id == task_id).then_some(new_exchange.exchange_id)
}) {
Some(id) => id,
None => {
log::error!("[bedrock-debug] AppendToMessageContent: ExchangeNotFound - no exchange with matching task_id");
log::error!(
"[bedrock-debug] AppendToMessageContent: ExchangeNotFound - no exchange with matching task_id"
);
return Err(UpdateConversationError::ExchangeNotFound);
}
}
}
None => {
log::error!("[bedrock-debug] AppendToMessageContent: NoPendingRequest - no exchanges for this stream_id");
log::error!(
"[bedrock-debug] AppendToMessageContent: NoPendingRequest - no exchanges for this stream_id"
);
return Err(UpdateConversationError::NoPendingRequest);
}
};
log::info!("[bedrock-debug] AppendToMessageContent: found exchange_id={:?}", exchange_id);
log::info!(
"[bedrock-debug] AppendToMessageContent: found exchange_id={:?}",
exchange_id
);
let current_todo_list = self.todo_lists.last().cloned();
let current_comment_state = self.code_review.as_ref().cloned();
// Update the message and get the updated todos op, if any.
let todos_op = match self
.task_store
.modify_task(&task_id, |task| {
task.append_to_message_content(
message,
exchange_id,
current_todo_list.as_ref(),
current_comment_state.as_ref(),
mask,
)
.map(|msg| msg.todos_op().cloned())
}) {
Some(result) => {
match result {
Ok(todos_op) => {
log::info!("[bedrock-debug] AppendToMessageContent: append succeeded");
todos_op
}
Err(e) => {
log::error!("[bedrock-debug] AppendToMessageContent: append_to_message_content failed: {e:?}");
return Err(e.into());
}
let todos_op = match self.task_store.modify_task(&task_id, |task| {
task.append_to_message_content(
message,
exchange_id,
current_todo_list.as_ref(),
current_comment_state.as_ref(),
mask,
)
.map(|msg| msg.todos_op().cloned())
}) {
Some(result) => match result {
Ok(todos_op) => {
log::info!("[bedrock-debug] AppendToMessageContent: append succeeded");
todos_op
}
}
Err(e) => {
log::error!(
"[bedrock-debug] AppendToMessageContent: append_to_message_content failed: {e:?}"
);
return Err(e.into());
}
},
None => {
log::error!("[bedrock-debug] AppendToMessageContent: TaskNotFound in task_store");
log::error!(
"[bedrock-debug] AppendToMessageContent: TaskNotFound in task_store"
);
return Err(UpdateConversationError::TaskNotFound);
}
};
@@ -3011,6 +3061,10 @@ impl AIConversation {
self.total_token_usage_by_model.values().cloned().collect()
}
fn total_tokens_for_usage(usage: &TokenUsage) -> u32 {
usage.total_input + usage.output + usage.input_cache_read + usage.input_cache_write
}
/// Normalize all newlines to CRLF so restored blocks render lines starting at column 0,
/// which is consistent with how we serialize real terminal blocks.
fn to_stylized_bytes(s: &str) -> Vec<u8> {
app/src/terminal/block_list_element.rs
+41 -52
View File
@@ -1,5 +1,6 @@
use crate::ai::blocklist::agent_view::{agent_view_bg_fill, AgentViewState};
use crate::ai::blocklist::{ai_brand_color, ATTACH_AS_AGENT_MODE_CONTEXT_TEXT};
use crate::BlocklistAIHistoryModel;
use crate::ai::blocklist::agent_view::{AgentViewState, agent_view_bg_fill};
use crate::ai::blocklist::{ATTACH_AS_AGENT_MODE_CONTEXT_TEXT, ai_brand_color};
use crate::ai_assistant::{AI_ASSISTANT_SVG_PATH, ASK_AI_ASSISTANT_TEXT};
use crate::appearance::Appearance;
use crate::drive::settings::WarpDriveSettings;
@@ -19,30 +20,21 @@ use crate::terminal::model::index::Point as IndexPoint;
use crate::terminal::model::selection::{SelectAction, SelectionPoint};
use crate::terminal::safe_mode_settings::get_secret_obfuscation_mode;
use crate::terminal::view::TerminalAction;
use crate::terminal::{grid_renderer, SizeInfo};
use crate::terminal::{SizeInfo, grid_renderer};
use crate::themes::theme::{Fill, WarpTheme};
use crate::ui_components::{self, icons as UIIcon};
use crate::util::color::Opacity;
use enum_iterator::Sequence;
use itertools::Itertools;
use parking_lot::FairMutex;
use vec1::Vec1;
use galaxy_core::semantic_selection::SemanticSelection;
use galaxy_core::ui::builder::UiBuilder;
use galaxy_core::ui::theme::AnsiColorIdentifier;
use galaxy_util::user_input::UserInput;
use galaxyui::platform::Cursor;
use galaxyui::text::SelectionType;
use itertools::Itertools;
use parking_lot::FairMutex;
use vec1::Vec1;
use pathfinder_color::ColorU;
use session_sharing_protocol::common::{ParticipantId, Selection};
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::mem;
use std::ops::{Deref, Range, RangeInclusive};
use std::rc::Rc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use galaxyui::elements::new_scrollable::{NewScrollableElement, ScrollableAxis};
use galaxyui::elements::{
Axis, Border, ChildAnchor, ClippedScrollStateHandle, ConstrainedBox, Container, CornerRadius,
@@ -52,23 +44,34 @@ use galaxyui::elements::{
use galaxyui::event::{KeyState, ModifiersState};
use galaxyui::fonts::{FamilyId, Properties, Weight};
use galaxyui::geometry::rect::RectF;
use galaxyui::geometry::vector::{vec2f, Vector2F};
use galaxyui::geometry::vector::{Vector2F, vec2f};
use galaxyui::platform::keyboard::KeyCode;
use galaxyui::ui_components::components::UiComponent;
use galaxyui::units::{IntoLines, IntoPixels, Lines, Pixels};
use galaxyui::{elements::Icon, ClipBounds};
use galaxyui::{
elements::SavePosition, event::DispatchedEvent, AfterLayoutContext, AppContext, Element, Event,
EventContext, LayoutContext, PaintContext, SizeConstraint,
AfterLayoutContext, AppContext, Element, Event, EventContext, LayoutContext, PaintContext,
SizeConstraint, elements::SavePosition, event::DispatchedEvent,
};
use galaxyui::{ClipBounds, elements::Icon};
use galaxyui::{EntityId, ModelHandle, SingletonEntity as _};
use pathfinder_color::ColorU;
use session_sharing_protocol::common::{ParticipantId, Selection};
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::mem;
use std::ops::{Deref, Range, RangeInclusive};
use std::rc::Rc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use super::block_list_viewport::{ClampingMode, InputMode, ScrollPosition, ViewportState};
use super::blockgrid_renderer::GridRenderParams;
use super::find::{BlockListFindRun, BlockListMatch, TerminalFindModel};
use super::grid_renderer::CellGlyphCache;
use super::TerminalModel;
use super::meta_shortcuts::handle_keystroke_despite_composing;
use super::model::SecretHandle;
use super::model::block::BlockId;
use super::model::blocks::{RichContentItem, SelectionRange};
use super::model::grid::grid_handler::{Link, TermMode};
@@ -76,24 +79,22 @@ use super::model::image_map::StoredImageMetadata;
use super::model::mouse::{MouseAction, MouseButton, MouseState};
use super::model::session::SessionId;
use super::model::terminal_model::{SelectedBlocks, WithinBlock, WithinModel};
use super::model::SecretHandle;
use super::shared_session::presence_manager::{
text_selection_color, PresenceManager, MUTED_PARTICIPANT_COLOR,
MUTED_PARTICIPANT_COLOR, PresenceManager, text_selection_color,
};
use super::shared_session::render_util::SHARED_SESSION_AVATAR_DIAMETER;
use super::view::{
BlocklistAIRenderContext, InlineBannerId, RichContentMetadata, SeparatorId,
SharedSessionBanners, TerminalEditor, TerminalViewRenderContext, BLOCK_BANNER_HEIGHT,
BLOCK_BANNER_HEIGHT, BlocklistAIRenderContext, InlineBannerId, RichContentMetadata,
SeparatorId, SharedSessionBanners, TerminalEditor, TerminalViewRenderContext,
};
use super::warpify::render::{draw_flag_pole, render_subshell_flag};
use super::TerminalModel;
use super::{heights_approx_eq, HEIGHT_FUDGE_FACTOR_LINES};
use super::{HEIGHT_FUDGE_FACTOR_LINES, heights_approx_eq};
use crate::terminal::blockgrid_renderer::BlockGridParams;
use crate::terminal::model::terminal_model::BlockIndex;
use crate::terminal::warpify::SubshellSource;
use crate::terminal::model::escape_sequences::{
maybe_kitty_keyboard_escape_sequence, KeystrokeWithDetails, ToEscapeSequence,
KeystrokeWithDetails, ToEscapeSequence, maybe_kitty_keyboard_escape_sequence,
};
/// The number of pixels at the bottom of padding where selection scrolling is performed.
@@ -3547,36 +3548,22 @@ impl Element for BlockListElement {
drop(viewport_iter);
if DebugSettings::as_ref(app).should_show_memory_stats() {
let history_model = BlocklistAIHistoryModel::as_ref(app);
for block_index in &visible_block_indices {
if let Some(block) = model.block_list().block_at(*block_index) {
if !block.has_footer() {
continue;
}
fn adjusted_bytes(bytes: usize) -> byte_unit::AdjustedByte {
let unit = if bytes >= 1_000_000 {
byte_unit::Unit::MB
} else {
byte_unit::Unit::KB
};
byte_unit::Byte::from(bytes).get_adjusted_unit(unit)
}
let grid_storage_lines = block.grid_storage_lines();
let grid_storage_bytes = block.grid_storage_bytes();
let flat_storage_lines = block.flat_storage_lines();
let flat_storage_bytes = block.flat_storage_bytes();
let total_lines = grid_storage_lines + flat_storage_lines;
let total_bytes = grid_storage_bytes + flat_storage_bytes;
let text = format!("\
Lines: {total_lines} (grid: {grid_storage_lines}, flat: {flat_storage_lines}); \
Size: {:#.1} (grid: {:#.1}, flat: {:#.1})\
",
adjusted_bytes(total_bytes),
adjusted_bytes(grid_storage_bytes),
adjusted_bytes(flat_storage_bytes),
);
let request_tokens = block
.agent_view_visibility()
.agent_view_conversation_id()
.and_then(|conversation_id| history_model.conversation(&conversation_id))
.map_or(0, |conversation| conversation.last_block_total_tokens());
let text = if request_tokens > 0 {
format!("Request tokens: {request_tokens}")
} else {
"Request tokens: --".to_string()
};
let mut element = Text::new_inline(text, self.ui_font_family, self.font_size)
.with_style(Properties::default().weight(self.font_weight))
@@ -3977,7 +3964,9 @@ impl Element for BlockListElement {
- SPACE_BETWEEN_SELECTED_BLOCK_AVATARS,
);
} else {
log::warn!("Should show avatar for shared session participant at selected block but avatar element was not found")
log::warn!(
"Should show avatar for shared session participant at selected block but avatar element was not found"
)
}
}