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