//! Commands sent to the ChatStateActor. use std::collections::BTreeSet; use kigi_sampling_types::{ ConversationItem, ConversationRequest, DanglingToolCallReason, SamplingConfig, TokenUsage, ToolSpec, TraceContext, }; use tokio::sync::oneshot; use crate::types::{ AutoCompactTrigger, ChatStateSnapshot, ConversationCounts, Credentials, NotificationMeta, TurnCapture, }; #[derive(Debug, Clone, Default)] pub struct ModelMetadata { pub resolved_model_id: Option, pub model_fingerprint: Option, } /// Refusal reply for [`ChatStateCommand::RepairHistory`]: a turn was in /// flight, and in-flight tool calls must not be treated as dangling. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct RepairHistoryBlocked; impl std::fmt::Display for RepairHistoryBlocked { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, "cannot repair history while a turn is in flight; stop the turn first" ) } } impl std::error::Error for RepairHistoryBlocked {} /// Commands sent to the ChatStateActor via mpsc channel. pub enum ChatStateCommand { // ═══ Mutations (fire-and-forget) ═══ /// Push a user message into the conversation. PushUserMessage { item: ConversationItem }, /// Push a user message and acknowledge once the chat-state actor has /// accepted and processed it. PushUserMessageAndAck { item: ConversationItem, reply: oneshot::Sender<()>, }, /// Push a user message with an explicit dangling-repair reason. PushUserMessageWithRepairReason { item: ConversationItem, reason: DanglingToolCallReason, }, /// Record the assistant's response (text + tool calls). PushAssistantResponse { item: ConversationItem }, /// Record a tool result. PushToolResult { item: ConversationItem }, /// Record accumulated token usage from a streaming response. RecordTokenUsage { total_tokens: u64 }, /// Stash the per-turn `TokenUsage` from the most recent model response. /// Overwrites any previously stashed value. RecordLastTurnUsage { usage: TokenUsage }, RecordModelCallUsage { model_id: Option, usage: TokenUsage, api_duration_ms: Option, cost_usd_ticks: Option, }, /// Subagent usage into session (and prompt when attributable). Replies when applied. RecordSubagentUsage { by_model: Vec<(String, crate::usage::UsageTotals)>, attribute_to_prompt: bool, /// Nested subagent bill may under-count. incomplete: bool, reply: oneshot::Sender<()>, }, /// Mark open prompt and/or session ledgers incomplete. MarkUsageIncomplete { prompt: bool, session: bool, reply: oneshot::Sender<()>, }, /// Increment prompt_index (called at start of each user turn). IncrementPromptIndex, /// Update the sampling config (e.g., model switch). UpdateSamplingConfig { config: SamplingConfig }, /// Track that the agent edited a file path. RecordAgentEditedPath { path: String }, /// Record stream timing metadata. RecordStreamStart { timestamp_ms: i64 }, /// Record turn timing metadata. RecordTurnStart { timestamp_ms: i64 }, /// Replace conversation history. ReplaceConversation { items: Vec, is_compaction: bool, }, /// Out-of-band history repair (`kigi/session/repair`): run /// [`crate::compaction_utils::repair_history`] and persist when changed; /// `dry_run` only reports. /// /// `turn_active` (the session's shared flag, set at turn start BEFORE the /// turn pushes anything here) is re-checked inside the command handler: /// a caller-side check alone races turn start, whereas at processing time /// the command is either refused or runs on pre-turn state with the /// turn's pushes serialized after it. RepairHistory { dry_run: bool, turn_active: Option>, reply: oneshot::Sender< Result, >, }, /// Atomically align the leading `System` message with `prompt` (inserting /// one if absent), persisting the conversation. Executed inside the actor so /// it serializes with concurrent turn pushes (`PushAssistantResponse` / /// `PushToolResult`) — a mid-turn reconnect cannot lose those updates the /// way a read-modify-write via `GetConversation` + `ReplaceConversation` /// would. Replies `true` iff the conversation changed (no-op when the head /// already matches modulo trailing newlines). A changed head goes through /// `replace_conversation`, which re-bases `total_tokens` to a fresh static /// estimate — acceptable because a changed head invalidates the KV prefix /// anyway. ReplaceSystemHead { prompt: String, reply: oneshot::Sender, }, /// Cache prompt text for rewind preview. CachePromptText { text: String }, /// Record compaction boundary for rewind. RecordCompactionAt { prompt_index: usize }, /// Flush pending persistence writes to disk (end of turn). Flush, /// Update opaque credential secrets held by the actor. UpdateCredentials { credentials: Credentials }, /// Restore from a snapshot. RestoreSnapshot(Box), /// Start capturing turn messages. Clears any previous buffer. BeginTurnCapture, /// Append synthetic `task` pairs for a harness-spawned subagent (goal /// planner / verifier skeptic) to the in-progress harness trace phase. /// Accumulated independently of the live `conversation` and of /// `turn_capture`; sealed into a standalone trace turn by /// `FlushHarnessTraceTurn`. AppendHarnessTraceItems { items: Vec }, /// Seal the harness items accumulated since the last flush into one /// standalone trace turn. Issued once per harness phase (after the planner, /// after each verifier panel). No-op when nothing was recorded. FlushHarnessTraceTurn, /// Repair dangling tool calls after a harness-initiated halt. RepairDanglingAfterHarnessHalt { class: &'static str }, // ═══ Queries (request/response via oneshot) ═══ /// Build a ConversationRequest ready to send to the API. /// Clones the conversation, prunes old tool results, repairs dangling /// tool calls, injects memory reminder, and assembles the request. BuildConversationRequest { tool_definitions: Vec, memory_reminder: Option, persist_memory_reminder: bool, trace: Option>, conv_id: String, req_id: String, reply: oneshot::Sender, }, /// Get a clone of the full conversation. GetConversation { reply: oneshot::Sender>, }, /// Get current prompt index. GetPromptIndex { reply: oneshot::Sender }, /// Get the prompt index at which the last compaction occurred. /// `Some` means the context currently holds a compaction summary. GetLastCompactionPromptIndex { reply: oneshot::Sender>, }, /// Get total accumulated tokens. GetTotalTokens { reply: oneshot::Sender }, /// Retrieve the most recent stashed per-turn `TokenUsage`. Returns /// `None` until at least one `RecordLastTurnUsage` has been processed. GetLastTurnUsage { reply: oneshot::Sender>, }, GetPromptUsage { reply: oneshot::Sender>, }, GetSessionUsage { reply: oneshot::Sender, }, /// `total_tokens` + bytes/4 delta from tool results since last model response. GetEstimatedTotalTokens { reply: oneshot::Sender }, /// Bytes/4 estimate of all non-system conversation items. GetEstimatedMessagesTokens { reply: oneshot::Sender }, /// Get sampling config. GetSamplingConfig { reply: oneshot::Sender, }, /// Get the set of agent-edited file paths. GetAgentEditedPaths { reply: oneshot::Sender>, }, /// Get notification meta (timing info). GetNotificationMeta { reply: oneshot::Sender, }, /// Snapshot state for forking or rewind. Snapshot { reply: oneshot::Sender, }, /// Truncate conversation to a target prompt index (for rewind). TruncateToPromptIndex { target_prompt_index: usize, reply: oneshot::Sender<()>, }, /// Check if auto-compact is needed (returns token info). CheckAutoCompactNeeded { threshold_percent: u8, reply: oneshot::Sender>, }, /// Get credential secrets. GetCredentials { reply: oneshot::Sender }, GetLastModelMetadata { reply: oneshot::Sender, }, /// Take the accumulated turn messages and end the capture. /// Returns `None` if no capture was active. TakeTurnMessages { reply: oneshot::Sender>, }, /// Drain the sealed harness trace turns (goal planner + verifier panels). /// Each `Vec` is one turn's synthetic `task` pairs, uploaded by the agent /// as its own sibling `turn_{N}` artifact. Seals a trailing un-flushed /// accumulator before draining. TakeHarnessTraceTurns { reply: oneshot::Sender>>, }, // ═══ Narrow targeted queries (avoid full-conversation clone) ═══ /// Get the number of items in the conversation. /// Cheaper than `GetConversation` when only the length is needed. GetConversationLen { reply: oneshot::Sender }, /// Whether any assistant tool call lacks a matching `ToolResult` (i.e. the /// dangling-tool-call repair would fire on the next request build). /// Cheaper than `GetConversation` when only this predicate is needed. HasDanglingToolCalls { reply: oneshot::Sender }, /// Get the text content of the last assistant message with non-empty text. /// Returns `None` if no such message exists. /// Cheaper than `GetConversation` when only the final assistant response is needed. GetLastAssistantText { reply: oneshot::Sender>, }, /// Get the text of the first `Text` content part in the first `User` message. /// Returns `None` if the conversation has no user messages or the first user /// message has no text content part. /// Cheaper than `GetConversation` when only the initial user query is needed. GetFirstUserText { reply: oneshot::Sender>, }, /// Get a single conversation item by index (0-based). /// Returns `None` if the index is out of bounds. /// Cheaper than `GetConversation` when only one item is needed. GetConversationItemAt { index: usize, reply: oneshot::Sender>, }, /// Get the processed text of the last user query (metadata tags stripped). /// /// Equivalent to `extract_last_user_query(&conversation)` but without /// cloning the full conversation on the caller side. GetLastUserQueryText { reply: oneshot::Sender>, }, /// Get item counts for the conversation by role. /// /// Returns a `ConversationCounts` struct without cloning any items. /// Suitable for telemetry / logging that only needs totals. GetConversationCounts { reply: oneshot::Sender, }, /// Get the first `System` message in the conversation, if any. /// /// Cheaper than `GetConversation` when only the system prompt is needed /// (e.g. for compaction setup or error guards). GetSystemMessage { reply: oneshot::Sender>, }, } #[cfg(test)] mod tests { use super::*; /// Verify that every command variant is constructible (compile-time check). #[test] fn command_variants_are_constructible() { // Mutations let _ = ChatStateCommand::PushUserMessage { item: ConversationItem::user("hello"), }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::PushUserMessageAndAck { item: ConversationItem::user("hello"), reply: tx, }; let _ = ChatStateCommand::PushAssistantResponse { item: ConversationItem::assistant("hi"), }; let _ = ChatStateCommand::PushToolResult { item: ConversationItem::tool_result("call-1", "result"), }; let _ = ChatStateCommand::RecordTokenUsage { total_tokens: 100 }; let _ = ChatStateCommand::IncrementPromptIndex; let _ = ChatStateCommand::UpdateSamplingConfig { config: SamplingConfig { base_url: String::new(), model: String::new(), max_completion_tokens: None, temperature: None, top_p: None, api_backend: Default::default(), chat_compat: Default::default(), extra_headers: Default::default(), context_window: std::num::NonZeroU64::new(128_000).unwrap(), reasoning_effort: None, stream_tool_calls: None, }, }; let _ = ChatStateCommand::RecordAgentEditedPath { path: "src/main.rs".to_string(), }; let _ = ChatStateCommand::RecordStreamStart { timestamp_ms: 12345, }; let _ = ChatStateCommand::RecordTurnStart { timestamp_ms: 12345, }; let _ = ChatStateCommand::ReplaceConversation { items: vec![], is_compaction: false, }; let _ = ChatStateCommand::CachePromptText { text: "prompt".to_string(), }; let _ = ChatStateCommand::RecordCompactionAt { prompt_index: 0 }; let _ = ChatStateCommand::Flush; // Queries let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetConversation { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetPromptIndex { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetLastCompactionPromptIndex { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetTotalTokens { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetEstimatedTotalTokens { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetSamplingConfig { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetAgentEditedPaths { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::BuildConversationRequest { tool_definitions: vec![], memory_reminder: None, persist_memory_reminder: false, trace: None, conv_id: String::new(), req_id: String::new(), reply: tx, }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetNotificationMeta { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::Snapshot { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::TruncateToPromptIndex { target_prompt_index: 0, reply: tx, }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::CheckAutoCompactNeeded { threshold_percent: 85, reply: tx, }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetLastModelMetadata { reply: tx }; let _ = ChatStateCommand::BeginTurnCapture; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::TakeTurnMessages { reply: tx }; // Narrow targeted queries let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetConversationLen { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetLastAssistantText { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetFirstUserText { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetConversationItemAt { index: 0, reply: tx, }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetLastUserQueryText { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetConversationCounts { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetSystemMessage { reply: tx }; let (tx, _rx) = oneshot::channel(); let _ = ChatStateCommand::GetEstimatedMessagesTokens { reply: tx }; } }