1use std::collections::{HashMap, VecDeque};
19use std::path::PathBuf;
20use std::time::SystemTime;
21
22use crate::app::instructions::LoadedInstructions;
23use crate::app::{Config, McpServerConfig};
24use crate::models::ChatMessage;
25use crate::models::tool_call::ToolCall as ModelToolCall;
26use crate::models::{ReasoningLevel, TokenUsage, TokenUsageSource};
27use crate::session::ConversationHistory;
28
29use super::cmd::ChatRequest;
30use super::compaction::CompactionTrigger;
31use super::ids::{IdAllocator, ToolCallId, TurnId};
32use super::msg::Msg;
33use super::runtime::{RuntimeState, ToolArtifact, ToolRunMetadata, ToolStatus};
34
35#[derive(Debug, Clone)]
41pub struct State {
42 pub session: Session,
43 pub turn: TurnState,
44 pub ui: UiState,
45 pub mcp: McpState,
46 pub settings: Config,
47 pub instructions: Option<LoadedInstructions>,
48 pub cwd: PathBuf,
52 pub ids: IdAllocatorBundle,
53 pub confirm: Option<Confirmation>,
57 pub status: Option<StatusLine>,
61 pub runtime: RuntimeState,
65 pub should_exit: bool,
68}
69
70impl State {
71 pub fn new(settings: Config, cwd: PathBuf, model_id: String) -> Self {
74 let project_path = cwd.display().to_string();
75 let conversation = ConversationHistory::new(project_path, model_id.clone());
76 let initial_title = conversation.title.clone();
77 let mcp = {
83 let mut m = McpState::default();
84 for (name, cfg) in &settings.mcp_servers {
85 m.servers.insert(
86 name.clone(),
87 McpServerEntry {
88 config: cfg.clone(),
89 status: McpServerStatus::Starting,
90 tools: Vec::new(),
91 },
92 );
93 }
94 m
95 };
96 let reasoning = settings
100 .reasoning_per_model
101 .get(&model_id)
102 .copied()
103 .unwrap_or(settings.default_model.reasoning);
104 let runtime = RuntimeState::new(&model_id);
105 Self {
106 session: Session {
107 conversation,
108 model_id,
109 reasoning,
110 cumulative_tokens: 0,
111 last_token_usage: None,
112 cumulative_token_usage: TokenUsageTotals::default(),
113 context_usage: None,
114 },
115 turn: TurnState::Idle,
116 ui: UiState {
117 last_title_dispatched: Some(initial_title),
118 ..UiState::default()
119 },
120 mcp,
121 settings,
122 instructions: None,
123 cwd,
124 ids: IdAllocatorBundle::default(),
125 confirm: None,
126 status: None,
127 runtime,
128 should_exit: false,
129 }
130 }
131
132 pub fn is_busy(&self) -> bool {
135 !matches!(self.turn, TurnState::Idle)
136 }
137
138 pub fn current_turn_id(&self) -> Option<TurnId> {
143 self.turn.id()
144 }
145}
146
147#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
152pub struct TokenUsageTotals {
153 pub prompt_tokens: usize,
154 pub completion_tokens: usize,
155 pub total_tokens: usize,
156 pub cached_input_tokens: usize,
157 pub cache_creation_input_tokens: usize,
158 pub reasoning_output_tokens: usize,
159}
160
161impl TokenUsageTotals {
162 pub fn from_usage(usage: &TokenUsage) -> Self {
163 Self {
164 prompt_tokens: usage.prompt_tokens,
165 completion_tokens: usage.completion_tokens,
166 total_tokens: usage.total_tokens,
167 cached_input_tokens: usage.cached_input_tokens,
168 cache_creation_input_tokens: usage.cache_creation_input_tokens,
169 reasoning_output_tokens: usage.reasoning_output_tokens,
170 }
171 }
172
173 pub fn add_assign(&mut self, other: Self) {
174 self.prompt_tokens = self.prompt_tokens.saturating_add(other.prompt_tokens);
175 self.completion_tokens = self
176 .completion_tokens
177 .saturating_add(other.completion_tokens);
178 self.total_tokens = self.total_tokens.saturating_add(other.total_tokens);
179 self.cached_input_tokens = self
180 .cached_input_tokens
181 .saturating_add(other.cached_input_tokens);
182 self.cache_creation_input_tokens = self
183 .cache_creation_input_tokens
184 .saturating_add(other.cache_creation_input_tokens);
185 self.reasoning_output_tokens = self
186 .reasoning_output_tokens
187 .saturating_add(other.reasoning_output_tokens);
188 }
189
190 pub fn input_total_tokens(&self) -> usize {
191 self.prompt_tokens
192 .saturating_add(self.cached_input_tokens)
193 .saturating_add(self.cache_creation_input_tokens)
194 }
195
196 pub fn output_total_tokens(&self) -> usize {
197 self.completion_tokens
198 .saturating_add(self.reasoning_output_tokens)
199 }
200}
201
202#[derive(Debug, Clone, Default, PartialEq, Eq)]
205pub struct PromptTokenBreakdown {
206 pub system_tokens: usize,
207 pub instructions_tokens: usize,
208 pub message_tokens: usize,
209 pub tool_schema_tokens: usize,
210 pub image_count: usize,
211 pub message_count: usize,
212 pub tool_count: usize,
213}
214
215impl PromptTokenBreakdown {
216 pub fn total_tokens(&self) -> usize {
217 self.system_tokens
218 .saturating_add(self.instructions_tokens)
219 .saturating_add(self.message_tokens)
220 .saturating_add(self.tool_schema_tokens)
221 }
222}
223
224#[derive(Debug, Clone, PartialEq, Eq)]
227pub struct ContextUsageSnapshot {
228 pub used_tokens: usize,
229 pub max_tokens: Option<usize>,
230 pub remaining_tokens: Option<usize>,
231 pub used_percent: Option<u8>,
232 pub source: TokenUsageSource,
233 pub prompt_tokens: usize,
234 pub cached_input_tokens: usize,
235 pub cache_creation_input_tokens: usize,
236 pub completion_tokens: usize,
237 pub reasoning_output_tokens: usize,
238 pub breakdown: Option<PromptTokenBreakdown>,
239}
240
241impl ContextUsageSnapshot {
242 pub fn from_usage(usage: &TokenUsage, max_tokens: Option<usize>) -> Self {
243 Self::new(
244 usage.total_tokens,
245 max_tokens,
246 usage.source,
247 usage.prompt_tokens,
248 usage.cached_input_tokens,
249 usage.cache_creation_input_tokens,
250 usage.completion_tokens,
251 usage.reasoning_output_tokens,
252 None,
253 )
254 }
255
256 pub fn from_estimate(breakdown: PromptTokenBreakdown, max_tokens: Option<usize>) -> Self {
257 let used = breakdown.total_tokens();
258 Self::new(
259 used,
260 max_tokens,
261 TokenUsageSource::Estimate,
262 used,
263 0,
264 0,
265 0,
266 0,
267 Some(breakdown),
268 )
269 }
270
271 #[allow(clippy::too_many_arguments)]
272 fn new(
273 used_tokens: usize,
274 max_tokens: Option<usize>,
275 source: TokenUsageSource,
276 prompt_tokens: usize,
277 cached_input_tokens: usize,
278 cache_creation_input_tokens: usize,
279 completion_tokens: usize,
280 reasoning_output_tokens: usize,
281 breakdown: Option<PromptTokenBreakdown>,
282 ) -> Self {
283 let remaining_tokens = max_tokens.map(|max| max.saturating_sub(used_tokens));
284 let used_percent = max_tokens
285 .filter(|max| *max > 0)
286 .map(|max| ((used_tokens.saturating_mul(100)) / max).min(100) as u8);
287 Self {
288 used_tokens,
289 max_tokens,
290 remaining_tokens,
291 used_percent,
292 source,
293 prompt_tokens,
294 cached_input_tokens,
295 cache_creation_input_tokens,
296 completion_tokens,
297 reasoning_output_tokens,
298 breakdown,
299 }
300 }
301
302 pub fn is_estimate(&self) -> bool {
303 self.source == TokenUsageSource::Estimate
304 }
305}
306
307pub fn estimate_context_usage_for_request(
308 request: &ChatRequest,
309 max_tokens: Option<usize>,
310) -> ContextUsageSnapshot {
311 let system_tokens = approx_tokens(&request.system_prompt);
312 let instructions_tokens = request
313 .instructions
314 .as_deref()
315 .map(approx_tokens)
316 .unwrap_or(0);
317 let message_tokens = request
318 .messages
319 .iter()
320 .map(|msg| {
321 let image_chars = msg
322 .images
323 .as_ref()
324 .map(|imgs| imgs.iter().map(|img| img.len()).sum::<usize>())
325 .unwrap_or(0);
326 let tool_call_chars = msg
329 .tool_calls
330 .as_ref()
331 .map(|calls| {
332 calls
333 .iter()
334 .map(|tc| {
335 tc.function.name.len()
336 + tc.function.arguments.to_string().len()
337 + tc.id.as_deref().map(str::len).unwrap_or(0)
338 })
339 .sum::<usize>()
340 })
341 .unwrap_or(0);
342 approx_tokens(&msg.content)
343 .saturating_add(approx_tokens(&format!(
344 "{:?}{}{}",
345 msg.role,
346 msg.tool_name.as_deref().unwrap_or(""),
347 msg.tool_call_id.as_deref().unwrap_or("")
348 )))
349 .saturating_add(image_chars.div_ceil(4))
350 .saturating_add(tool_call_chars.div_ceil(4))
351 })
352 .sum();
353 let tool_schema: Vec<_> = request
354 .tools
355 .iter()
356 .map(|tool| tool.to_openai_json())
357 .collect();
358 let tool_schema_tokens = serde_json::to_string(&tool_schema)
359 .map(|s| approx_tokens(&s))
360 .unwrap_or(0);
361 let image_count = request
362 .messages
363 .iter()
364 .filter_map(|msg| msg.images.as_ref())
365 .map(Vec::len)
366 .sum();
367 ContextUsageSnapshot::from_estimate(
368 PromptTokenBreakdown {
369 system_tokens,
370 instructions_tokens,
371 message_tokens,
372 tool_schema_tokens,
373 image_count,
374 message_count: request.messages.len(),
375 tool_count: request.tools.len(),
376 },
377 max_tokens,
378 )
379}
380
381fn approx_tokens(text: &str) -> usize {
382 text.len().div_ceil(4)
383}
384
385#[derive(Debug, Clone)]
391pub struct Session {
392 pub conversation: ConversationHistory,
393 pub model_id: String,
394 pub reasoning: ReasoningLevel,
395 pub cumulative_tokens: usize,
399 pub last_token_usage: Option<TokenUsageTotals>,
402 pub cumulative_token_usage: TokenUsageTotals,
404 pub context_usage: Option<ContextUsageSnapshot>,
408}
409
410impl Session {
411 pub fn messages(&self) -> &[ChatMessage] {
415 &self.conversation.messages
416 }
417
418 pub fn append(&mut self, msg: ChatMessage) {
422 self.conversation.add_messages(&[msg]);
423 }
424}
425
426#[derive(Debug, Clone)]
436pub enum TurnState {
437 Idle,
438 Generating {
439 id: TurnId,
440 started: SystemTime,
441 partial_text: String,
442 partial_reasoning: String,
443 tokens: usize,
445 phase: GenPhase,
447 thinking_signature: Option<String>,
451 pending_tool_calls: Vec<ModelToolCall>,
457 },
458 ExecutingTools {
459 id: TurnId,
460 calls: Vec<PendingToolCall>,
461 outcomes: Vec<Option<ToolOutcome>>,
462 },
463 Compacting {
468 id: TurnId,
469 started: SystemTime,
470 trigger: CompactionTrigger,
471 },
472 Cancelling {
480 id: TurnId,
481 since: SystemTime,
482 },
483}
484
485impl TurnState {
486 pub fn id(&self) -> Option<TurnId> {
487 match self {
488 TurnState::Idle => None,
489 TurnState::Generating { id, .. }
490 | TurnState::ExecutingTools { id, .. }
491 | TurnState::Compacting { id, .. }
492 | TurnState::Cancelling { id, .. } => Some(*id),
493 }
494 }
495
496 pub fn accepts(&self, event_turn: TurnId) -> bool {
500 self.id() == Some(event_turn)
501 }
502}
503
504#[derive(Debug, Clone, Copy, PartialEq, Eq)]
508pub enum GenPhase {
509 Sending,
511 Thinking,
514 Streaming,
517}
518
519#[derive(Debug, Clone)]
523pub struct PendingToolCall {
524 pub call_id: ToolCallId,
525 pub source: ModelToolCall,
529}
530
531#[derive(Debug, Clone, PartialEq)]
538pub struct ToolOutcome {
539 pub status: ToolStatus,
540 pub summary: String,
541 pub model_content: String,
542 pub error: Option<String>,
543 pub metadata: Box<ToolRunMetadata>,
544 pub artifacts: Vec<ToolArtifact>,
545 pub duration_secs: Option<f64>,
546}
547
548impl ToolOutcome {
549 pub fn success(
550 model_content: impl Into<String>,
551 summary: impl Into<String>,
552 duration_secs: f64,
553 ) -> Self {
554 let duration = Some(duration_secs);
555 let metadata = ToolRunMetadata {
556 duration_secs: duration,
557 ..ToolRunMetadata::default()
558 };
559 Self {
560 status: ToolStatus::Success,
561 summary: summary.into(),
562 model_content: model_content.into(),
563 error: None,
564 metadata: Box::new(metadata),
565 artifacts: Vec::new(),
566 duration_secs: duration,
567 }
568 }
569
570 pub fn error(error: impl Into<String>, duration_secs: f64) -> Self {
571 let error = error.into();
572 let duration = Some(duration_secs);
573 Self {
574 status: ToolStatus::Error,
575 summary: error.clone(),
576 model_content: format!("Error: {}", error),
577 error: Some(error),
578 metadata: Box::new(ToolRunMetadata {
579 duration_secs: duration,
580 ..ToolRunMetadata::default()
581 }),
582 artifacts: Vec::new(),
583 duration_secs: duration,
584 }
585 }
586
587 pub fn cancelled() -> Self {
588 Self {
589 status: ToolStatus::Cancelled,
590 summary: "[cancelled]".to_string(),
591 model_content: "[Tool call skipped: the user cancelled before execution]".to_string(),
592 error: None,
593 metadata: Box::new(ToolRunMetadata::default()),
594 artifacts: Vec::new(),
595 duration_secs: None,
596 }
597 }
598
599 pub fn with_metadata(mut self, mut metadata: ToolRunMetadata) -> Self {
600 metadata.duration_secs = self.duration_secs;
601 self.metadata = Box::new(metadata);
602 self
603 }
604
605 pub fn with_artifacts(mut self, artifacts: Vec<ToolArtifact>) -> Self {
606 self.artifacts = artifacts.clone();
607 self.metadata.artifacts = artifacts;
608 self
609 }
610
611 pub fn with_images(self, images: Vec<String>) -> Self {
612 self.with_artifacts(
613 images
614 .into_iter()
615 .map(|data| ToolArtifact::Image { data })
616 .collect(),
617 )
618 }
619
620 pub fn was_cancelled(&self) -> bool {
621 self.status == ToolStatus::Cancelled
622 }
623
624 pub fn is_success(&self) -> bool {
625 self.status == ToolStatus::Success
626 }
627
628 pub fn output(&self) -> &str {
629 &self.model_content
630 }
631
632 pub fn error_message(&self) -> Option<&str> {
633 self.error.as_deref()
634 }
635
636 pub fn images(&self) -> Option<Vec<String>> {
637 let images: Vec<String> = self
638 .artifacts
639 .iter()
640 .filter_map(|artifact| match artifact {
641 ToolArtifact::Image { data } => Some(data.clone()),
642 _ => None,
643 })
644 .collect();
645 if images.is_empty() {
646 None
647 } else {
648 Some(images)
649 }
650 }
651
652 pub fn as_tool_message_content(&self) -> String {
657 self.model_content.clone()
658 }
659}
660
661#[derive(Debug, Clone, Default)]
664pub struct UiState {
665 pub mode: UiMode,
666 pub input_buffer: String,
667 pub input_cursor: usize,
671 pub attachments: Vec<Attachment>,
673 pub attachment_focused: bool,
676 pub attachment_selected: usize,
679 pub chat_scroll: usize,
681 pub palette_filter: String,
685 pub palette_cursor: Option<usize>,
688 pub queued_messages: VecDeque<String>,
692 pub last_title_dispatched: Option<String>,
696 pub pending_msgs: VecDeque<Msg>,
702 pub input_history_cursor: Option<usize>,
708 pub history_draft: String,
712 pub mouse_scroll_accum: i32,
719 pub show_reasoning: bool,
723}
724
725#[derive(Debug, Clone, PartialEq, Eq, Default)]
728pub enum UiMode {
729 #[default]
730 EditingInput,
731 Palette,
733 ConversationList {
737 candidates: Vec<ConversationSummary>,
738 cursor: usize,
739 },
740 ModelList,
742}
743
744#[derive(Debug, Clone, PartialEq, Eq)]
747pub struct ConversationSummary {
748 pub id: String,
749 pub title: String,
750 pub message_count: usize,
751 pub updated_at: String,
752}
753
754#[derive(Debug, Clone)]
757pub struct Attachment {
758 pub id: u64,
759 pub base64_data: String,
760 pub temp_path: PathBuf,
763 pub size_bytes: usize,
764 pub format: String,
765}
766
767#[derive(Debug, Clone, Default)]
771pub struct McpState {
772 pub servers: HashMap<String, McpServerEntry>,
773}
774
775#[derive(Debug, Clone)]
776pub struct McpServerEntry {
777 pub config: McpServerConfig,
778 pub status: McpServerStatus,
779 pub tools: Vec<McpToolSpec>,
783}
784
785#[derive(Debug, Clone, PartialEq, Eq)]
786pub enum McpServerStatus {
787 Starting,
789 Ready,
790 Errored {
791 reason: String,
792 },
793 Stopped,
794}
795
796#[derive(Debug, Clone)]
801pub struct McpToolSpec {
802 pub name: String,
803 pub description: String,
804 pub input_schema: serde_json::Value,
805}
806
807#[derive(Debug, Clone)]
810pub struct Confirmation {
811 pub prompt: String,
812 pub accept_msg_token: ConfirmationTarget,
813}
814
815#[derive(Debug, Clone)]
818pub enum ConfirmationTarget {
819 ClearConversation,
820}
821
822#[derive(Debug, Clone)]
826pub struct StatusLine {
827 pub text: String,
828 pub kind: StatusKind,
829 pub shown_at: SystemTime,
830}
831
832#[derive(Debug, Clone, Copy, PartialEq, Eq)]
833pub enum StatusKind {
834 Info,
835 Warn,
836 Error,
837 Persistent,
839}
840
841#[derive(Debug, Clone, Copy, Default)]
844pub struct IdAllocatorBundle {
845 pub turn: IdAllocator,
846 pub tool_call: IdAllocator,
847}
848
849impl IdAllocatorBundle {
850 pub fn fresh_turn(&mut self) -> TurnId {
851 TurnId(self.turn.next())
852 }
853
854 pub fn fresh_tool_call(&mut self) -> ToolCallId {
855 ToolCallId(self.tool_call.next())
856 }
857}
858
859#[cfg(test)]
860mod tests {
861 use super::*;
862
863 fn mock_state() -> State {
864 State::new(
865 Config::default(),
866 PathBuf::from("/tmp/project"),
867 "ollama/test".to_string(),
868 )
869 }
870
871 #[test]
872 fn fresh_state_is_idle() {
873 let s = mock_state();
874 assert!(matches!(s.turn, TurnState::Idle));
875 assert!(!s.is_busy());
876 assert!(s.current_turn_id().is_none());
877 }
878
879 #[test]
880 fn turn_state_accepts_matches_id() {
881 let s = TurnState::Generating {
882 id: TurnId(7),
883 started: SystemTime::now(),
884 partial_text: String::new(),
885 partial_reasoning: String::new(),
886 tokens: 0,
887 phase: GenPhase::Sending,
888 thinking_signature: None,
889 pending_tool_calls: Vec::new(),
890 };
891 assert!(s.accepts(TurnId(7)));
892 assert!(!s.accepts(TurnId(6)));
893 assert!(!s.accepts(TurnId(8)));
894 }
895
896 #[test]
897 fn idle_rejects_all_turn_ids() {
898 let s = TurnState::Idle;
899 assert!(!s.accepts(TurnId(1)));
900 assert!(!s.accepts(TurnId(999)));
901 }
902
903 #[test]
904 fn fresh_id_allocators_monotonic() {
905 let mut bundle = IdAllocatorBundle::default();
906 assert_eq!(bundle.fresh_turn(), TurnId(1));
907 assert_eq!(bundle.fresh_turn(), TurnId(2));
908 assert_eq!(bundle.fresh_tool_call(), ToolCallId(1));
909 }
912
913 #[test]
914 fn tool_outcome_cancelled_content_is_placeholder() {
915 let o = ToolOutcome::cancelled();
916 assert!(o.was_cancelled());
917 let content = o.as_tool_message_content();
918 assert!(content.contains("cancelled"));
919 }
920
921 #[test]
922 fn tool_outcome_finished_returns_output_verbatim() {
923 let o = ToolOutcome::success("hello world", "hello world", 0.1);
924 assert_eq!(o.as_tool_message_content(), "hello world");
925 assert!(!o.was_cancelled());
926 }
927
928 #[test]
929 fn session_append_records_message() {
930 let mut s = mock_state();
931 s.session.append(ChatMessage::user("hi"));
932 assert_eq!(s.session.messages().len(), 1);
933 assert_eq!(s.session.messages()[0].content, "hi");
934 }
935}