diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 10f8eda1..d231d901 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -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, + last_block_token_usage_by_model: HashMap, /// Fallback title used when no task description or initial query exists. fallback_display_title: Option, @@ -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 { + pub fn bedrock_message_history_mut( + &mut self, + ) -> &mut Vec { &mut self.bedrock_message_history } - pub fn append_to_bedrock_history(&mut self, messages: Vec) { + pub fn append_to_bedrock_history( + &mut self, + messages: Vec, + ) { 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, 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 { diff --git a/app/src/terminal/block_list_element.rs b/app/src/terminal/block_list_element.rs index d73798a3..30136981 100644 --- a/app/src/terminal/block_list_element.rs +++ b/app/src/terminal/block_list_element.rs @@ -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" + ) } }