use std::collections::{HashMap, VecDeque};
use std::path::PathBuf;
use std::time::SystemTime;
use crate::app::instructions::LoadedInstructions;
use crate::app::{Config, McpServerConfig};
use crate::models::ChatMessage;
use crate::models::tool_call::ToolCall as ModelToolCall;
use crate::models::{ReasoningLevel, TokenUsage, TokenUsageSource};
use crate::session::ConversationHistory;
use super::cmd::ChatRequest;
use super::compaction::CompactionTrigger;
use super::ids::{IdAllocator, ToolCallId, TurnId};
use super::msg::Msg;
use super::runtime::{RuntimeState, ToolArtifact, ToolRunMetadata, ToolStatus};
#[derive(Debug, Clone)]
pub struct State {
pub session: Session,
pub turn: TurnState,
pub ui: UiState,
pub mcp: McpState,
pub settings: Config,
pub instructions: Option<LoadedInstructions>,
pub cwd: PathBuf,
pub ids: IdAllocatorBundle,
pub confirm: Option<Confirmation>,
pub status: Option<StatusLine>,
pub runtime: RuntimeState,
pub should_exit: bool,
}
impl State {
pub fn new(settings: Config, cwd: PathBuf, model_id: String) -> Self {
let project_path = cwd.display().to_string();
let conversation = ConversationHistory::new(project_path, model_id.clone());
let initial_title = conversation.title.clone();
let mcp = {
let mut m = McpState::default();
for (name, cfg) in &settings.mcp_servers {
m.servers.insert(
name.clone(),
McpServerEntry {
config: cfg.clone(),
status: McpServerStatus::Starting,
tools: Vec::new(),
},
);
}
m
};
let reasoning = settings
.reasoning_per_model
.get(&model_id)
.copied()
.unwrap_or(settings.default_model.reasoning);
let runtime = RuntimeState::new(&model_id);
Self {
session: Session {
conversation,
model_id,
reasoning,
cumulative_tokens: 0,
last_token_usage: None,
cumulative_token_usage: TokenUsageTotals::default(),
context_usage: None,
},
turn: TurnState::Idle,
ui: UiState {
last_title_dispatched: Some(initial_title),
..UiState::default()
},
mcp,
settings,
instructions: None,
cwd,
ids: IdAllocatorBundle::default(),
confirm: None,
status: None,
runtime,
should_exit: false,
}
}
pub fn is_busy(&self) -> bool {
!matches!(self.turn, TurnState::Idle)
}
pub fn current_turn_id(&self) -> Option<TurnId> {
self.turn.id()
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenUsageTotals {
pub prompt_tokens: usize,
pub completion_tokens: usize,
pub total_tokens: usize,
pub cached_input_tokens: usize,
pub cache_creation_input_tokens: usize,
pub reasoning_output_tokens: usize,
}
impl TokenUsageTotals {
pub fn from_usage(usage: &TokenUsage) -> Self {
Self {
prompt_tokens: usage.prompt_tokens,
completion_tokens: usage.completion_tokens,
total_tokens: usage.total_tokens,
cached_input_tokens: usage.cached_input_tokens,
cache_creation_input_tokens: usage.cache_creation_input_tokens,
reasoning_output_tokens: usage.reasoning_output_tokens,
}
}
pub fn add_assign(&mut self, other: Self) {
self.prompt_tokens = self.prompt_tokens.saturating_add(other.prompt_tokens);
self.completion_tokens = self
.completion_tokens
.saturating_add(other.completion_tokens);
self.total_tokens = self.total_tokens.saturating_add(other.total_tokens);
self.cached_input_tokens = self
.cached_input_tokens
.saturating_add(other.cached_input_tokens);
self.cache_creation_input_tokens = self
.cache_creation_input_tokens
.saturating_add(other.cache_creation_input_tokens);
self.reasoning_output_tokens = self
.reasoning_output_tokens
.saturating_add(other.reasoning_output_tokens);
}
pub fn input_total_tokens(&self) -> usize {
self.prompt_tokens
.saturating_add(self.cached_input_tokens)
.saturating_add(self.cache_creation_input_tokens)
}
pub fn output_total_tokens(&self) -> usize {
self.completion_tokens
.saturating_add(self.reasoning_output_tokens)
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct PromptTokenBreakdown {
pub system_tokens: usize,
pub instructions_tokens: usize,
pub message_tokens: usize,
pub tool_schema_tokens: usize,
pub image_count: usize,
pub message_count: usize,
pub tool_count: usize,
}
impl PromptTokenBreakdown {
pub fn total_tokens(&self) -> usize {
self.system_tokens
.saturating_add(self.instructions_tokens)
.saturating_add(self.message_tokens)
.saturating_add(self.tool_schema_tokens)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextUsageSnapshot {
pub used_tokens: usize,
pub max_tokens: Option<usize>,
pub remaining_tokens: Option<usize>,
pub used_percent: Option<u8>,
pub source: TokenUsageSource,
pub prompt_tokens: usize,
pub cached_input_tokens: usize,
pub cache_creation_input_tokens: usize,
pub completion_tokens: usize,
pub reasoning_output_tokens: usize,
pub breakdown: Option<PromptTokenBreakdown>,
}
impl ContextUsageSnapshot {
pub fn from_usage(usage: &TokenUsage, max_tokens: Option<usize>) -> Self {
Self::new(
usage.total_tokens,
max_tokens,
usage.source,
usage.prompt_tokens,
usage.cached_input_tokens,
usage.cache_creation_input_tokens,
usage.completion_tokens,
usage.reasoning_output_tokens,
None,
)
}
pub fn from_estimate(breakdown: PromptTokenBreakdown, max_tokens: Option<usize>) -> Self {
let used = breakdown.total_tokens();
Self::new(
used,
max_tokens,
TokenUsageSource::Estimate,
used,
0,
0,
0,
0,
Some(breakdown),
)
}
#[allow(clippy::too_many_arguments)]
fn new(
used_tokens: usize,
max_tokens: Option<usize>,
source: TokenUsageSource,
prompt_tokens: usize,
cached_input_tokens: usize,
cache_creation_input_tokens: usize,
completion_tokens: usize,
reasoning_output_tokens: usize,
breakdown: Option<PromptTokenBreakdown>,
) -> Self {
let remaining_tokens = max_tokens.map(|max| max.saturating_sub(used_tokens));
let used_percent = max_tokens
.filter(|max| *max > 0)
.map(|max| ((used_tokens.saturating_mul(100)) / max).min(100) as u8);
Self {
used_tokens,
max_tokens,
remaining_tokens,
used_percent,
source,
prompt_tokens,
cached_input_tokens,
cache_creation_input_tokens,
completion_tokens,
reasoning_output_tokens,
breakdown,
}
}
pub fn is_estimate(&self) -> bool {
self.source == TokenUsageSource::Estimate
}
}
pub fn estimate_context_usage_for_request(
request: &ChatRequest,
max_tokens: Option<usize>,
) -> ContextUsageSnapshot {
let system_tokens = approx_tokens(&request.system_prompt);
let instructions_tokens = request
.instructions
.as_deref()
.map(approx_tokens)
.unwrap_or(0);
let message_tokens = request
.messages
.iter()
.map(|msg| {
let image_chars = msg
.images
.as_ref()
.map(|imgs| imgs.iter().map(|img| img.len()).sum::<usize>())
.unwrap_or(0);
approx_tokens(&msg.content).saturating_add(approx_tokens(&format!(
"{:?}{}{}",
msg.role,
msg.tool_name.as_deref().unwrap_or(""),
msg.tool_call_id.as_deref().unwrap_or("")
))) + image_chars.div_ceil(4)
})
.sum();
let tool_schema: Vec<_> = request
.tools
.iter()
.map(|tool| tool.to_openai_json())
.collect();
let tool_schema_tokens = serde_json::to_string(&tool_schema)
.map(|s| approx_tokens(&s))
.unwrap_or(0);
let image_count = request
.messages
.iter()
.filter_map(|msg| msg.images.as_ref())
.map(Vec::len)
.sum();
ContextUsageSnapshot::from_estimate(
PromptTokenBreakdown {
system_tokens,
instructions_tokens,
message_tokens,
tool_schema_tokens,
image_count,
message_count: request.messages.len(),
tool_count: request.tools.len(),
},
max_tokens,
)
}
fn approx_tokens(text: &str) -> usize {
text.len().div_ceil(4)
}
#[derive(Debug, Clone)]
pub struct Session {
pub conversation: ConversationHistory,
pub model_id: String,
pub reasoning: ReasoningLevel,
pub cumulative_tokens: usize,
pub last_token_usage: Option<TokenUsageTotals>,
pub cumulative_token_usage: TokenUsageTotals,
pub context_usage: Option<ContextUsageSnapshot>,
}
impl Session {
pub fn messages(&self) -> &[ChatMessage] {
&self.conversation.messages
}
pub fn append(&mut self, msg: ChatMessage) {
self.conversation.add_messages(&[msg]);
}
}
#[derive(Debug, Clone)]
pub enum TurnState {
Idle,
Generating {
id: TurnId,
started: SystemTime,
partial_text: String,
partial_reasoning: String,
tokens: usize,
phase: GenPhase,
thinking_signature: Option<String>,
pending_tool_calls: Vec<ModelToolCall>,
},
ExecutingTools {
id: TurnId,
calls: Vec<PendingToolCall>,
outcomes: Vec<Option<ToolOutcome>>,
},
Compacting {
id: TurnId,
started: SystemTime,
trigger: CompactionTrigger,
},
Cancelling {
id: TurnId,
since: SystemTime,
},
}
impl TurnState {
pub fn id(&self) -> Option<TurnId> {
match self {
TurnState::Idle => None,
TurnState::Generating { id, .. }
| TurnState::ExecutingTools { id, .. }
| TurnState::Compacting { id, .. }
| TurnState::Cancelling { id, .. } => Some(*id),
}
}
pub fn accepts(&self, event_turn: TurnId) -> bool {
self.id() == Some(event_turn)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GenPhase {
Sending,
Thinking,
Streaming,
}
#[derive(Debug, Clone)]
pub struct PendingToolCall {
pub call_id: ToolCallId,
pub source: ModelToolCall,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ToolOutcome {
pub status: ToolStatus,
pub summary: String,
pub model_content: String,
pub error: Option<String>,
pub metadata: Box<ToolRunMetadata>,
pub artifacts: Vec<ToolArtifact>,
pub duration_secs: Option<f64>,
}
impl ToolOutcome {
pub fn success(
model_content: impl Into<String>,
summary: impl Into<String>,
duration_secs: f64,
) -> Self {
let duration = Some(duration_secs);
let metadata = ToolRunMetadata {
duration_secs: duration,
..ToolRunMetadata::default()
};
Self {
status: ToolStatus::Success,
summary: summary.into(),
model_content: model_content.into(),
error: None,
metadata: Box::new(metadata),
artifacts: Vec::new(),
duration_secs: duration,
}
}
pub fn error(error: impl Into<String>, duration_secs: f64) -> Self {
let error = error.into();
let duration = Some(duration_secs);
Self {
status: ToolStatus::Error,
summary: error.clone(),
model_content: format!("Error: {}", error),
error: Some(error),
metadata: Box::new(ToolRunMetadata {
duration_secs: duration,
..ToolRunMetadata::default()
}),
artifacts: Vec::new(),
duration_secs: duration,
}
}
pub fn cancelled() -> Self {
Self {
status: ToolStatus::Cancelled,
summary: "[cancelled]".to_string(),
model_content: "[Tool call skipped: the user cancelled before execution]".to_string(),
error: None,
metadata: Box::new(ToolRunMetadata::default()),
artifacts: Vec::new(),
duration_secs: None,
}
}
pub fn with_metadata(mut self, mut metadata: ToolRunMetadata) -> Self {
metadata.duration_secs = self.duration_secs;
self.metadata = Box::new(metadata);
self
}
pub fn with_artifacts(mut self, artifacts: Vec<ToolArtifact>) -> Self {
self.artifacts = artifacts.clone();
self.metadata.artifacts = artifacts;
self
}
pub fn with_images(self, images: Vec<String>) -> Self {
self.with_artifacts(
images
.into_iter()
.map(|data| ToolArtifact::Image { data })
.collect(),
)
}
pub fn was_cancelled(&self) -> bool {
self.status == ToolStatus::Cancelled
}
pub fn is_success(&self) -> bool {
self.status == ToolStatus::Success
}
pub fn output(&self) -> &str {
&self.model_content
}
pub fn error_message(&self) -> Option<&str> {
self.error.as_deref()
}
pub fn images(&self) -> Option<Vec<String>> {
let images: Vec<String> = self
.artifacts
.iter()
.filter_map(|artifact| match artifact {
ToolArtifact::Image { data } => Some(data.clone()),
_ => None,
})
.collect();
if images.is_empty() {
None
} else {
Some(images)
}
}
pub fn as_tool_message_content(&self) -> String {
self.model_content.clone()
}
}
#[derive(Debug, Clone, Default)]
pub struct UiState {
pub mode: UiMode,
pub input_buffer: String,
pub input_cursor: usize,
pub attachments: Vec<Attachment>,
pub attachment_focused: bool,
pub attachment_selected: usize,
pub chat_scroll: usize,
pub palette_filter: String,
pub palette_cursor: Option<usize>,
pub queued_messages: VecDeque<String>,
pub last_title_dispatched: Option<String>,
pub pending_msgs: VecDeque<Msg>,
pub input_history_cursor: Option<usize>,
pub history_draft: String,
pub mouse_scroll_accum: i32,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub enum UiMode {
#[default]
EditingInput,
Palette,
ConversationList {
candidates: Vec<ConversationSummary>,
cursor: usize,
},
ModelList,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConversationSummary {
pub id: String,
pub title: String,
pub message_count: usize,
pub updated_at: String,
}
#[derive(Debug, Clone)]
pub struct Attachment {
pub id: u64,
pub base64_data: String,
pub temp_path: PathBuf,
pub size_bytes: usize,
pub format: String,
}
#[derive(Debug, Clone, Default)]
pub struct McpState {
pub servers: HashMap<String, McpServerEntry>,
}
#[derive(Debug, Clone)]
pub struct McpServerEntry {
pub config: McpServerConfig,
pub status: McpServerStatus,
pub tools: Vec<McpToolSpec>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum McpServerStatus {
Starting,
Ready,
Errored {
reason: String,
},
Stopped,
}
#[derive(Debug, Clone)]
pub struct McpToolSpec {
pub name: String,
pub description: String,
pub input_schema: serde_json::Value,
}
#[derive(Debug, Clone)]
pub struct Confirmation {
pub prompt: String,
pub accept_msg_token: ConfirmationTarget,
}
#[derive(Debug, Clone)]
pub enum ConfirmationTarget {
ClearConversation,
}
#[derive(Debug, Clone)]
pub struct StatusLine {
pub text: String,
pub kind: StatusKind,
pub shown_at: SystemTime,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StatusKind {
Info,
Warn,
Error,
Persistent,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct IdAllocatorBundle {
pub turn: IdAllocator,
pub tool_call: IdAllocator,
}
impl IdAllocatorBundle {
pub fn fresh_turn(&mut self) -> TurnId {
TurnId(self.turn.next())
}
pub fn fresh_tool_call(&mut self) -> ToolCallId {
ToolCallId(self.tool_call.next())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn mock_state() -> State {
State::new(
Config::default(),
PathBuf::from("/tmp/project"),
"ollama/test".to_string(),
)
}
#[test]
fn fresh_state_is_idle() {
let s = mock_state();
assert!(matches!(s.turn, TurnState::Idle));
assert!(!s.is_busy());
assert!(s.current_turn_id().is_none());
}
#[test]
fn turn_state_accepts_matches_id() {
let s = TurnState::Generating {
id: TurnId(7),
started: SystemTime::now(),
partial_text: String::new(),
partial_reasoning: String::new(),
tokens: 0,
phase: GenPhase::Sending,
thinking_signature: None,
pending_tool_calls: Vec::new(),
};
assert!(s.accepts(TurnId(7)));
assert!(!s.accepts(TurnId(6)));
assert!(!s.accepts(TurnId(8)));
}
#[test]
fn idle_rejects_all_turn_ids() {
let s = TurnState::Idle;
assert!(!s.accepts(TurnId(1)));
assert!(!s.accepts(TurnId(999)));
}
#[test]
fn fresh_id_allocators_monotonic() {
let mut bundle = IdAllocatorBundle::default();
assert_eq!(bundle.fresh_turn(), TurnId(1));
assert_eq!(bundle.fresh_turn(), TurnId(2));
assert_eq!(bundle.fresh_tool_call(), ToolCallId(1));
}
#[test]
fn tool_outcome_cancelled_content_is_placeholder() {
let o = ToolOutcome::cancelled();
assert!(o.was_cancelled());
let content = o.as_tool_message_content();
assert!(content.contains("cancelled"));
}
#[test]
fn tool_outcome_finished_returns_output_verbatim() {
let o = ToolOutcome::success("hello world", "hello world", 0.1);
assert_eq!(o.as_tool_message_content(), "hello world");
assert!(!o.was_cancelled());
}
#[test]
fn session_append_records_message() {
let mut s = mock_state();
s.session.append(ChatMessage::user("hi"));
assert_eq!(s.session.messages().len(), 1);
assert_eq!(s.session.messages()[0].content, "hi");
}
}