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 approx_tokens(&msg.content).saturating_add(approx_tokens(&format!(
327 "{:?}{}{}",
328 msg.role,
329 msg.tool_name.as_deref().unwrap_or(""),
330 msg.tool_call_id.as_deref().unwrap_or("")
331 ))) + image_chars.div_ceil(4)
332 })
333 .sum();
334 let tool_schema: Vec<_> = request
335 .tools
336 .iter()
337 .map(|tool| tool.to_openai_json())
338 .collect();
339 let tool_schema_tokens = serde_json::to_string(&tool_schema)
340 .map(|s| approx_tokens(&s))
341 .unwrap_or(0);
342 let image_count = request
343 .messages
344 .iter()
345 .filter_map(|msg| msg.images.as_ref())
346 .map(Vec::len)
347 .sum();
348 ContextUsageSnapshot::from_estimate(
349 PromptTokenBreakdown {
350 system_tokens,
351 instructions_tokens,
352 message_tokens,
353 tool_schema_tokens,
354 image_count,
355 message_count: request.messages.len(),
356 tool_count: request.tools.len(),
357 },
358 max_tokens,
359 )
360}
361
362fn approx_tokens(text: &str) -> usize {
363 text.len().div_ceil(4)
364}
365
366#[derive(Debug, Clone)]
372pub struct Session {
373 pub conversation: ConversationHistory,
374 pub model_id: String,
375 pub reasoning: ReasoningLevel,
376 pub cumulative_tokens: usize,
380 pub last_token_usage: Option<TokenUsageTotals>,
383 pub cumulative_token_usage: TokenUsageTotals,
385 pub context_usage: Option<ContextUsageSnapshot>,
389}
390
391impl Session {
392 pub fn messages(&self) -> &[ChatMessage] {
396 &self.conversation.messages
397 }
398
399 pub fn append(&mut self, msg: ChatMessage) {
403 self.conversation.add_messages(&[msg]);
404 }
405}
406
407#[derive(Debug, Clone)]
417pub enum TurnState {
418 Idle,
419 Generating {
420 id: TurnId,
421 started: SystemTime,
422 partial_text: String,
423 partial_reasoning: String,
424 tokens: usize,
426 phase: GenPhase,
428 thinking_signature: Option<String>,
432 pending_tool_calls: Vec<ModelToolCall>,
438 },
439 ExecutingTools {
440 id: TurnId,
441 calls: Vec<PendingToolCall>,
442 outcomes: Vec<Option<ToolOutcome>>,
443 },
444 Compacting {
449 id: TurnId,
450 started: SystemTime,
451 trigger: CompactionTrigger,
452 },
453 Cancelling {
461 id: TurnId,
462 since: SystemTime,
463 },
464}
465
466impl TurnState {
467 pub fn id(&self) -> Option<TurnId> {
468 match self {
469 TurnState::Idle => None,
470 TurnState::Generating { id, .. }
471 | TurnState::ExecutingTools { id, .. }
472 | TurnState::Compacting { id, .. }
473 | TurnState::Cancelling { id, .. } => Some(*id),
474 }
475 }
476
477 pub fn accepts(&self, event_turn: TurnId) -> bool {
481 self.id() == Some(event_turn)
482 }
483}
484
485#[derive(Debug, Clone, Copy, PartialEq, Eq)]
489pub enum GenPhase {
490 Sending,
492 Thinking,
495 Streaming,
498}
499
500#[derive(Debug, Clone)]
504pub struct PendingToolCall {
505 pub call_id: ToolCallId,
506 pub source: ModelToolCall,
510}
511
512#[derive(Debug, Clone, PartialEq)]
519pub struct ToolOutcome {
520 pub status: ToolStatus,
521 pub summary: String,
522 pub model_content: String,
523 pub error: Option<String>,
524 pub metadata: Box<ToolRunMetadata>,
525 pub artifacts: Vec<ToolArtifact>,
526 pub duration_secs: Option<f64>,
527}
528
529impl ToolOutcome {
530 pub fn success(
531 model_content: impl Into<String>,
532 summary: impl Into<String>,
533 duration_secs: f64,
534 ) -> Self {
535 let duration = Some(duration_secs);
536 let metadata = ToolRunMetadata {
537 duration_secs: duration,
538 ..ToolRunMetadata::default()
539 };
540 Self {
541 status: ToolStatus::Success,
542 summary: summary.into(),
543 model_content: model_content.into(),
544 error: None,
545 metadata: Box::new(metadata),
546 artifacts: Vec::new(),
547 duration_secs: duration,
548 }
549 }
550
551 pub fn error(error: impl Into<String>, duration_secs: f64) -> Self {
552 let error = error.into();
553 let duration = Some(duration_secs);
554 Self {
555 status: ToolStatus::Error,
556 summary: error.clone(),
557 model_content: format!("Error: {}", error),
558 error: Some(error),
559 metadata: Box::new(ToolRunMetadata {
560 duration_secs: duration,
561 ..ToolRunMetadata::default()
562 }),
563 artifacts: Vec::new(),
564 duration_secs: duration,
565 }
566 }
567
568 pub fn cancelled() -> Self {
569 Self {
570 status: ToolStatus::Cancelled,
571 summary: "[cancelled]".to_string(),
572 model_content: "[Tool call skipped: the user cancelled before execution]".to_string(),
573 error: None,
574 metadata: Box::new(ToolRunMetadata::default()),
575 artifacts: Vec::new(),
576 duration_secs: None,
577 }
578 }
579
580 pub fn with_metadata(mut self, mut metadata: ToolRunMetadata) -> Self {
581 metadata.duration_secs = self.duration_secs;
582 self.metadata = Box::new(metadata);
583 self
584 }
585
586 pub fn with_artifacts(mut self, artifacts: Vec<ToolArtifact>) -> Self {
587 self.artifacts = artifacts.clone();
588 self.metadata.artifacts = artifacts;
589 self
590 }
591
592 pub fn with_images(self, images: Vec<String>) -> Self {
593 self.with_artifacts(
594 images
595 .into_iter()
596 .map(|data| ToolArtifact::Image { data })
597 .collect(),
598 )
599 }
600
601 pub fn was_cancelled(&self) -> bool {
602 self.status == ToolStatus::Cancelled
603 }
604
605 pub fn is_success(&self) -> bool {
606 self.status == ToolStatus::Success
607 }
608
609 pub fn output(&self) -> &str {
610 &self.model_content
611 }
612
613 pub fn error_message(&self) -> Option<&str> {
614 self.error.as_deref()
615 }
616
617 pub fn images(&self) -> Option<Vec<String>> {
618 let images: Vec<String> = self
619 .artifacts
620 .iter()
621 .filter_map(|artifact| match artifact {
622 ToolArtifact::Image { data } => Some(data.clone()),
623 _ => None,
624 })
625 .collect();
626 if images.is_empty() {
627 None
628 } else {
629 Some(images)
630 }
631 }
632
633 pub fn as_tool_message_content(&self) -> String {
638 self.model_content.clone()
639 }
640}
641
642#[derive(Debug, Clone, Default)]
645pub struct UiState {
646 pub mode: UiMode,
647 pub input_buffer: String,
648 pub input_cursor: usize,
652 pub attachments: Vec<Attachment>,
654 pub attachment_focused: bool,
657 pub attachment_selected: usize,
660 pub chat_scroll: usize,
662 pub palette_filter: String,
666 pub palette_cursor: Option<usize>,
669 pub queued_messages: VecDeque<String>,
673 pub last_title_dispatched: Option<String>,
677 pub pending_msgs: VecDeque<Msg>,
683 pub input_history_cursor: Option<usize>,
689 pub history_draft: String,
693 pub mouse_scroll_accum: i32,
700}
701
702#[derive(Debug, Clone, PartialEq, Eq, Default)]
705pub enum UiMode {
706 #[default]
707 EditingInput,
708 Palette,
710 ConversationList {
714 candidates: Vec<ConversationSummary>,
715 cursor: usize,
716 },
717 ModelList,
719}
720
721#[derive(Debug, Clone, PartialEq, Eq)]
724pub struct ConversationSummary {
725 pub id: String,
726 pub title: String,
727 pub message_count: usize,
728 pub updated_at: String,
729}
730
731#[derive(Debug, Clone)]
734pub struct Attachment {
735 pub id: u64,
736 pub base64_data: String,
737 pub temp_path: PathBuf,
740 pub size_bytes: usize,
741 pub format: String,
742}
743
744#[derive(Debug, Clone, Default)]
748pub struct McpState {
749 pub servers: HashMap<String, McpServerEntry>,
750}
751
752#[derive(Debug, Clone)]
753pub struct McpServerEntry {
754 pub config: McpServerConfig,
755 pub status: McpServerStatus,
756 pub tools: Vec<McpToolSpec>,
760}
761
762#[derive(Debug, Clone, PartialEq, Eq)]
763pub enum McpServerStatus {
764 Starting,
766 Ready,
767 Errored {
768 reason: String,
769 },
770 Stopped,
771}
772
773#[derive(Debug, Clone)]
778pub struct McpToolSpec {
779 pub name: String,
780 pub description: String,
781 pub input_schema: serde_json::Value,
782}
783
784#[derive(Debug, Clone)]
787pub struct Confirmation {
788 pub prompt: String,
789 pub accept_msg_token: ConfirmationTarget,
790}
791
792#[derive(Debug, Clone)]
795pub enum ConfirmationTarget {
796 ClearConversation,
797}
798
799#[derive(Debug, Clone)]
803pub struct StatusLine {
804 pub text: String,
805 pub kind: StatusKind,
806 pub shown_at: SystemTime,
807}
808
809#[derive(Debug, Clone, Copy, PartialEq, Eq)]
810pub enum StatusKind {
811 Info,
812 Warn,
813 Error,
814 Persistent,
816}
817
818#[derive(Debug, Clone, Copy, Default)]
821pub struct IdAllocatorBundle {
822 pub turn: IdAllocator,
823 pub tool_call: IdAllocator,
824}
825
826impl IdAllocatorBundle {
827 pub fn fresh_turn(&mut self) -> TurnId {
828 TurnId(self.turn.next())
829 }
830
831 pub fn fresh_tool_call(&mut self) -> ToolCallId {
832 ToolCallId(self.tool_call.next())
833 }
834}
835
836#[cfg(test)]
837mod tests {
838 use super::*;
839
840 fn mock_state() -> State {
841 State::new(
842 Config::default(),
843 PathBuf::from("/tmp/project"),
844 "ollama/test".to_string(),
845 )
846 }
847
848 #[test]
849 fn fresh_state_is_idle() {
850 let s = mock_state();
851 assert!(matches!(s.turn, TurnState::Idle));
852 assert!(!s.is_busy());
853 assert!(s.current_turn_id().is_none());
854 }
855
856 #[test]
857 fn turn_state_accepts_matches_id() {
858 let s = TurnState::Generating {
859 id: TurnId(7),
860 started: SystemTime::now(),
861 partial_text: String::new(),
862 partial_reasoning: String::new(),
863 tokens: 0,
864 phase: GenPhase::Sending,
865 thinking_signature: None,
866 pending_tool_calls: Vec::new(),
867 };
868 assert!(s.accepts(TurnId(7)));
869 assert!(!s.accepts(TurnId(6)));
870 assert!(!s.accepts(TurnId(8)));
871 }
872
873 #[test]
874 fn idle_rejects_all_turn_ids() {
875 let s = TurnState::Idle;
876 assert!(!s.accepts(TurnId(1)));
877 assert!(!s.accepts(TurnId(999)));
878 }
879
880 #[test]
881 fn fresh_id_allocators_monotonic() {
882 let mut bundle = IdAllocatorBundle::default();
883 assert_eq!(bundle.fresh_turn(), TurnId(1));
884 assert_eq!(bundle.fresh_turn(), TurnId(2));
885 assert_eq!(bundle.fresh_tool_call(), ToolCallId(1));
886 }
889
890 #[test]
891 fn tool_outcome_cancelled_content_is_placeholder() {
892 let o = ToolOutcome::cancelled();
893 assert!(o.was_cancelled());
894 let content = o.as_tool_message_content();
895 assert!(content.contains("cancelled"));
896 }
897
898 #[test]
899 fn tool_outcome_finished_returns_output_verbatim() {
900 let o = ToolOutcome::success("hello world", "hello world", 0.1);
901 assert_eq!(o.as_tool_message_content(), "hello world");
902 assert!(!o.was_cancelled());
903 }
904
905 #[test]
906 fn session_append_records_message() {
907 let mut s = mock_state();
908 s.session.append(ChatMessage::user("hi"));
909 assert_eq!(s.session.messages().len(), 1);
910 assert_eq!(s.session.messages()[0].content, "hi");
911 }
912}