Publish harness and TUI open-source

initial sync from the monorepo

crates/codegen/xai-chat-state/src/handle.rs

@@ -0,0 +1,641 @@
+//! Handle to communicate with ChatStateActor.
+
+use std::collections::BTreeSet;
+
+use tokio::sync::{mpsc, oneshot};
+use xai_grok_sampling_types::{
+    ConversationItem, ConversationRequest, DanglingToolCallReason, SamplingConfig, TokenUsage,
+    ToolSpec, TraceContext,
+};
+
+use crate::commands::{ChatStateCommand, RepairHistoryBlocked};
+use crate::types::{
+    AutoCompactTrigger, ChatStateSnapshot, ConversationCounts, Credentials, NotificationMeta,
+    TurnCapture,
+};
+
+/// Handle to communicate with ChatStateActor.
+/// This is cheap to clone and can be shared across tasks.
+#[derive(Clone)]
+pub struct ChatStateHandle {
+    cmd_tx: mpsc::UnboundedSender<ChatStateCommand>,
+}
+
+impl ChatStateHandle {
+    /// Create a new handle with the given command sender.
+    pub(crate) fn new(cmd_tx: mpsc::UnboundedSender<ChatStateCommand>) -> Self {
+        Self { cmd_tx }
+    }
+
+    /// Create a no-op handle that discards all commands.
+    /// Useful for tests and situations where chat state tracking is not needed.
+    pub fn noop() -> Self {
+        let (cmd_tx, _cmd_rx) = mpsc::unbounded_channel();
+        Self { cmd_tx }
+    }
+
+    // ═══ Fire-and-forget mutations ═══
+
+    /// Push a user message into the conversation.
+    pub fn push_user_message(&self, item: ConversationItem) {
+        let _ = self.cmd_tx.send(ChatStateCommand::PushUserMessage { item });
+    }
+
+    /// Push a user message and await acknowledgement that the chat-state actor
+    /// has accepted and processed it.
+    pub async fn push_user_message_and_ack(&self, item: ConversationItem) -> Option<()> {
+        self.query("PushUserMessageAndAck", |reply| {
+            ChatStateCommand::PushUserMessageAndAck { item, reply }
+        })
+        .await
+    }
+
+    /// Push a user message with an explicit dangling-repair reason.
+    pub fn push_user_message_with_repair_reason(
+        &self,
+        item: ConversationItem,
+        reason: DanglingToolCallReason,
+    ) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::PushUserMessageWithRepairReason { item, reason });
+    }
+
+    /// Record the assistant's response.
+    pub fn push_assistant_response(&self, item: ConversationItem) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::PushAssistantResponse { item });
+    }
+
+    /// Record a tool result.
+    pub fn push_tool_result(&self, item: ConversationItem) {
+        let _ = self.cmd_tx.send(ChatStateCommand::PushToolResult { item });
+    }
+
+    /// Record accumulated token usage.
+    pub fn record_token_usage(&self, total_tokens: u64) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RecordTokenUsage { total_tokens });
+    }
+
+    /// Stash the per-turn `TokenUsage` from the most recent model response.
+    /// Fire-and-forget — no ack returned.
+    pub fn record_last_turn_usage(&self, usage: TokenUsage) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RecordLastTurnUsage { usage });
+    }
+
+    pub fn record_model_call_usage(
+        &self,
+        model_id: Option<String>,
+        usage: TokenUsage,
+        api_duration_ms: Option<u64>,
+        cost_usd_ticks: Option<i64>,
+    ) {
+        let _ = self.cmd_tx.send(ChatStateCommand::RecordModelCallUsage {
+            model_id,
+            usage,
+            api_duration_ms,
+            cost_usd_ticks,
+        });
+    }
+
+    /// Apply subagent usage; returns false if the actor did not acknowledge.
+    pub async fn record_subagent_usage(
+        &self,
+        by_model: Vec<(String, crate::usage::UsageTotals)>,
+        attribute_to_prompt: bool,
+        incomplete: bool,
+    ) -> bool {
+        self.query("RecordSubagentUsage", |reply| {
+            ChatStateCommand::RecordSubagentUsage {
+                by_model,
+                attribute_to_prompt,
+                incomplete,
+                reply,
+            }
+        })
+        .await
+        .is_some()
+    }
+
+    /// Mark open prompt and/or session ledgers incomplete.
+    pub async fn mark_usage_incomplete(&self, prompt: bool, session: bool) -> bool {
+        self.query("MarkUsageIncomplete", |reply| {
+            ChatStateCommand::MarkUsageIncomplete {
+                prompt,
+                session,
+                reply,
+            }
+        })
+        .await
+        .is_some()
+    }
+
+    /// Increment prompt index (called at start of each user turn).
+    pub fn increment_prompt_index(&self) {
+        let _ = self.cmd_tx.send(ChatStateCommand::IncrementPromptIndex);
+    }
+
+    /// Update the sampling config (e.g., model switch).
+    pub fn update_sampling_config(&self, config: SamplingConfig) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::UpdateSamplingConfig { config });
+    }
+
+    /// Track that the agent edited a file path.
+    pub fn record_agent_edited_path(&self, path: String) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RecordAgentEditedPath { path });
+    }
+
+    /// Record stream timing metadata.
+    pub fn record_stream_start(&self, timestamp_ms: i64) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RecordStreamStart { timestamp_ms });
+    }
+
+    /// Record turn timing metadata.
+    pub fn record_turn_start(&self, timestamp_ms: i64) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RecordTurnStart { timestamp_ms });
+    }
+
+    /// Replace conversation history.
+    pub fn replace_conversation(&self, items: Vec<ConversationItem>) {
+        self.send_replace(items, false);
+    }
+
+    /// Replace conversation history for compaction.
+    /// Sets `compaction_occurred` on the active turn capture.
+    pub fn replace_conversation_for_compaction(&self, items: Vec<ConversationItem>) {
+        self.send_replace(items, true);
+    }
+
+    fn send_replace(&self, items: Vec<ConversationItem>, is_compaction: bool) {
+        let _ = self.cmd_tx.send(ChatStateCommand::ReplaceConversation {
+            items,
+            is_compaction,
+        });
+    }
+
+    /// Out-of-band history repair (`x.ai/session/repair`); see
+    /// [`ChatStateCommand::RepairHistory`]. Returns `None` if the actor is
+    /// dead, `Some(Err(_))` if a turn was in flight at processing time.
+    pub async fn repair_history(
+        &self,
+        dry_run: bool,
+        turn_active: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
+    ) -> Option<Result<crate::compaction_utils::HistoryRepairReport, RepairHistoryBlocked>> {
+        self.query("RepairHistory", |reply| ChatStateCommand::RepairHistory {
+            dry_run,
+            turn_active,
+            reply,
+        })
+        .await
+    }
+
+    /// Atomically align the leading `System` message with `prompt` (insert one
+    /// if absent), persisting when changed. Serializes with turn pushes inside
+    /// the actor, so a mid-turn reconnect can't drop concurrent updates.
+    /// Returns `Some(changed)`, or `None` if the actor is dead.
+    pub async fn replace_system_head(&self, prompt: &str) -> Option<bool> {
+        let prompt = prompt.to_owned();
+        self.query("ReplaceSystemHead", |reply| {
+            ChatStateCommand::ReplaceSystemHead { prompt, reply }
+        })
+        .await
+    }
+
+    /// Cache prompt text for rewind preview.
+    pub fn cache_prompt_text(&self, text: String) {
+        let _ = self.cmd_tx.send(ChatStateCommand::CachePromptText { text });
+    }
+
+    /// Record compaction boundary for rewind.
+    pub fn record_compaction_at(&self, prompt_index: usize) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RecordCompactionAt { prompt_index });
+    }
+
+    /// Flush pending persistence writes to disk.
+    pub fn flush(&self) {
+        let _ = self.cmd_tx.send(ChatStateCommand::Flush);
+    }
+
+    /// Update opaque credential secrets held by the actor.
+    pub fn update_credentials(&self, credentials: Credentials) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::UpdateCredentials { credentials });
+    }
+
+    /// Restore from a snapshot.
+    pub fn restore_snapshot(&self, snapshot: ChatStateSnapshot) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RestoreSnapshot(Box::new(snapshot)));
+    }
+
+    /// Begin capturing turn messages. Call at the start of a real user turn
+    /// (in `handle_prompt`), before `push_user_message`.
+    pub fn begin_turn_capture(&self) {
+        let _ = self.cmd_tx.send(ChatStateCommand::BeginTurnCapture);
+    }
+
+    /// Append synthetic `task` pairs for a harness-spawned subagent (goal
+    /// planner / verifier skeptic) to the in-progress harness trace phase. They
+    /// are sealed into a standalone trace turn by [`Self::flush_harness_trace_turn`]
+    /// and never enter the live `conversation` sent to the model. No-op on
+    /// empty input.
+    pub fn append_harness_trace_items(&self, items: Vec<ConversationItem>) {
+        if items.is_empty() {
+            return;
+        }
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::AppendHarnessTraceItems { items });
+    }
+
+    /// Seal the harness items accumulated since the last flush into one trace
+    /// turn. Call once per harness phase (after the planner, after a verifier
+    /// panel) so each phase becomes its own uploaded `turn_{N}` artifact. No-op
+    /// when nothing was recorded since the last flush.
+    pub fn flush_harness_trace_turn(&self) {
+        let _ = self.cmd_tx.send(ChatStateCommand::FlushHarnessTraceTurn);
+    }
+
+    /// Repair dangling tool calls after a harness-initiated halt.
+    pub fn repair_dangling_after_harness_halt(&self, class: &'static str) {
+        let _ = self
+            .cmd_tx
+            .send(ChatStateCommand::RepairDanglingAfterHarnessHalt { class });
+    }
+
+    // ═══ Async queries (via oneshot) ═══
+
+    /// Send a query to the actor and await the reply.
+    ///
+    /// Returns `None` when the actor is dead (channel send failure or reply
+    /// dropped due to panic/cancellation). Both failure modes are logged at
+    /// `error` level with `cmd_name` for post-mortem diagnostics.
+    async fn query<T>(
+        &self,
+        cmd_name: &str,
+        make_cmd: impl FnOnce(oneshot::Sender<T>) -> ChatStateCommand,
+    ) -> Option<T> {
+        let (tx, rx) = oneshot::channel();
+        if self.cmd_tx.send(make_cmd(tx)).is_err() {
+            tracing::error!(cmd_name, "ChatStateActor dead: send failed");
+            return None;
+        }
+        match rx.await {
+            Ok(v) => Some(v),
+            Err(_) => {
+                tracing::error!(cmd_name, "ChatStateActor dead: reply dropped");
+                None
+            }
+        }
+    }
+
+    /// Build a ConversationRequest from the current state.
+    /// Prunes, repairs, injects memory, and returns a ready-to-send request.
+    pub async fn build_request(
+        &self,
+        tool_definitions: Vec<ToolSpec>,
+        memory_reminder: Option<String>,
+        persist_memory_reminder: bool,
+        trace: Option<Box<dyn TraceContext>>,
+        conv_id: String,
+        req_id: String,
+    ) -> Option<ConversationRequest> {
+        self.query("BuildConversationRequest", |reply| {
+            ChatStateCommand::BuildConversationRequest {
+                tool_definitions,
+                memory_reminder,
+                persist_memory_reminder,
+                trace,
+                conv_id,
+                req_id,
+                reply,
+            }
+        })
+        .await
+    }
+
+    /// Get a clone of the full conversation.
+    pub async fn get_conversation(&self) -> Vec<ConversationItem> {
+        self.query("GetConversation", |reply| {
+            ChatStateCommand::GetConversation { reply }
+        })
+        .await
+        .unwrap_or_default()
+    }
+
+    /// Get current prompt index.
+    pub async fn get_prompt_index(&self) -> usize {
+        self.query("GetPromptIndex", |reply| ChatStateCommand::GetPromptIndex {
+            reply,
+        })
+        .await
+        .unwrap_or(0)
+    }
+
+    /// Get the prompt index at which the last compaction occurred.
+    /// `Some` means the context currently holds a compaction summary.
+    pub async fn get_last_compaction_prompt_index(&self) -> Option<usize> {
+        self.query("GetLastCompactionPromptIndex", |reply| {
+            ChatStateCommand::GetLastCompactionPromptIndex { reply }
+        })
+        .await
+        .flatten()
+    }
+
+    /// Get total accumulated tokens.
+    pub async fn get_total_tokens(&self) -> u64 {
+        self.query("GetTotalTokens", |reply| ChatStateCommand::GetTotalTokens {
+            reply,
+        })
+        .await
+        .unwrap_or(0)
+    }
+
+    /// Retrieve the most recent stashed per-turn `TokenUsage`. Returns
+    /// `None` if no model turn has completed in this session yet, or if
+    /// the actor channel is closed.
+    pub async fn get_last_turn_usage(&self) -> Option<TokenUsage> {
+        self.query("GetLastTurnUsage", |reply| {
+            ChatStateCommand::GetLastTurnUsage { reply }
+        })
+        .await
+        .flatten()
+    }
+
+    /// Fail-closed prompt bill read.
+    /// `Ok(None)` means the actor answered "no ledger"; `Err(())` means it did
+    /// not answer at all. Never collapse `Err` to `None`: an unreadable bill
+    /// must not be mistaken for a free prompt.
+    pub async fn try_get_prompt_usage(&self) -> Result<Option<crate::usage::UsageLedger>, ()> {
+        self.query("GetPromptUsage", |reply| ChatStateCommand::GetPromptUsage {
+            reply,
+        })
+        .await
+        .ok_or(())
+    }
+
+    /// Fail-closed session bill read. `Err(())` if the actor is dead.
+    pub async fn try_get_session_usage(&self) -> Result<crate::usage::UsageLedger, ()> {
+        self.query("GetSessionUsage", |reply| {
+            ChatStateCommand::GetSessionUsage { reply }
+        })
+        .await
+        .ok_or(())
+    }
+
+    /// `total_tokens` plus bytes/4 estimate of tool results pushed since the
+    /// last model response. Used by `check_preflight_overflow`.
+    pub async fn get_estimated_total_tokens(&self) -> u64 {
+        self.query("GetEstimatedTotalTokens", |reply| {
+            ChatStateCommand::GetEstimatedTotalTokens { reply }
+        })
+        .await
+        .unwrap_or(0)
+    }
+
+    /// Bytes/4 estimate of all non-system conversation items.
+    pub async fn get_estimated_messages_tokens(&self) -> u64 {
+        self.query("GetEstimatedMessagesTokens", |reply| {
+            ChatStateCommand::GetEstimatedMessagesTokens { reply }
+        })
+        .await
+        .unwrap_or(0)
+    }
+
+    /// Get sampling config.
+    pub async fn get_sampling_config(&self) -> Option<SamplingConfig> {
+        self.query("GetSamplingConfig", |reply| {
+            ChatStateCommand::GetSamplingConfig { reply }
+        })
+        .await
+    }
+
+    /// Get the set of agent-edited file paths.
+    pub async fn get_agent_edited_paths(&self) -> BTreeSet<String> {
+        self.query("GetAgentEditedPaths", |reply| {
+            ChatStateCommand::GetAgentEditedPaths { reply }
+        })
+        .await
+        .unwrap_or_default()
+    }
+
+    /// Get notification meta (timing info).
+    pub async fn get_notification_meta(&self) -> Option<NotificationMeta> {
+        self.query("GetNotificationMeta", |reply| {
+            ChatStateCommand::GetNotificationMeta { reply }
+        })
+        .await
+    }
+
+    /// Snapshot state for forking or rewind.
+    pub async fn snapshot(&self) -> Option<ChatStateSnapshot> {
+        self.query("Snapshot", |reply| ChatStateCommand::Snapshot { reply })
+            .await
+    }
+
+    /// Truncate conversation to a target prompt index (for rewind).
+    pub async fn truncate_to_prompt_index(&self, target: usize) {
+        self.query("TruncateToPromptIndex", |reply| {
+            ChatStateCommand::TruncateToPromptIndex {
+                target_prompt_index: target,
+                reply,
+            }
+        })
+        .await;
+    }
+
+    /// Get credential secrets.
+    pub async fn get_credentials(&self) -> Credentials {
+        self.query("GetCredentials", |reply| ChatStateCommand::GetCredentials {
+            reply,
+        })
+        .await
+        .unwrap_or_default()
+    }
+
+    pub async fn get_last_model_metadata(&self) -> crate::commands::ModelMetadata {
+        self.query("GetLastModelMetadata", |reply| {
+            ChatStateCommand::GetLastModelMetadata { reply }
+        })
+        .await
+        .unwrap_or_default()
+    }
+
+    /// Take the accumulated turn messages and end the capture.
+    /// Returns `None` if no capture was active.
+    pub async fn take_turn_messages(&self) -> Option<TurnCapture> {
+        self.query("TakeTurnMessages", |reply| {
+            ChatStateCommand::TakeTurnMessages { reply }
+        })
+        .await
+        .flatten()
+    }
+
+    /// Drain the sealed harness trace turns (goal planner + verifier panels).
+    /// Each returned `Vec` is one turn's worth of synthetic `task` pairs,
+    /// destined to be uploaded as its own sibling `turn_{N}` artifact. A
+    /// trailing un-flushed accumulator is sealed defensively before draining.
+    /// Returns empty when nothing was recorded (the common, non-goal case).
+    pub async fn take_harness_trace_turns(&self) -> Vec<Vec<ConversationItem>> {
+        self.query("TakeHarnessTraceTurns", |reply| {
+            ChatStateCommand::TakeHarnessTraceTurns { reply }
+        })
+        .await
+        .unwrap_or_default()
+    }
+
+    /// Check if auto-compact is needed.
+    pub async fn check_auto_compact_needed(
+        &self,
+        threshold_percent: u8,
+    ) -> Option<AutoCompactTrigger> {
+        self.query("CheckAutoCompactNeeded", |reply| {
+            ChatStateCommand::CheckAutoCompactNeeded {
+                threshold_percent,
+                reply,
+            }
+        })
+        .await
+        .flatten()
+    }
+
+    // ═══ Narrow targeted queries ═══
+
+    /// Get the number of items in the conversation.
+    ///
+    /// Cheaper than [`get_conversation`] when only the length is needed —
+    /// the actor returns a single `usize` without cloning any items.
+    pub async fn get_conversation_len(&self) -> usize {
+        self.query("GetConversationLen", |reply| {
+            ChatStateCommand::GetConversationLen { reply }
+        })
+        .await
+        .unwrap_or(0)
+    }
+
+    /// Whether any assistant tool call lacks a matching `ToolResult` (the
+    /// dangling-tool-call repair would fire on the next request build).
+    ///
+    /// Returns `false` if the actor is dead. Cheaper than [`get_conversation`]
+    /// — the actor scans in place and returns a single `bool`.
+    pub async fn has_dangling_tool_calls(&self) -> bool {
+        self.query("HasDanglingToolCalls", |reply| {
+            ChatStateCommand::HasDanglingToolCalls { reply }
+        })
+        .await
+        .unwrap_or(false)
+    }
+
+    /// Get the text content of the last assistant message with non-empty text.
+    ///
+    /// Returns `None` if no such message exists or the actor is dead.
+    /// Cheaper than [`get_conversation`] when only the final assistant
+    /// response text is needed.
+    pub async fn get_last_assistant_text(&self) -> Option<String> {
+        self.query("GetLastAssistantText", |reply| {
+            ChatStateCommand::GetLastAssistantText { reply }
+        })
+        .await
+        .flatten()
+    }
+
+    /// Get the text of the first `Text` content part in the first `User` message.
+    ///
+    /// Returns `None` if no user message with text content exists or the actor
+    /// is dead. Cheaper than [`get_conversation`] when only the initial user
+    /// query text is needed (e.g. for memory context search).
+    pub async fn get_first_user_text(&self) -> Option<String> {
+        self.query("GetFirstUserText", |reply| {
+            ChatStateCommand::GetFirstUserText { reply }
+        })
+        .await
+        .flatten()
+    }
+
+    /// Get a single conversation item by index (0-based).
+    ///
+    /// Returns `None` if the index is out of bounds or the actor is dead.
+    /// Cheaper than [`get_conversation`] when only one specific item is needed
+    /// (e.g. item[1] for the original user-info block after compaction).
+    pub async fn get_conversation_item_at(&self, index: usize) -> Option<ConversationItem> {
+        self.query("GetConversationItemAt", |reply| {
+            ChatStateCommand::GetConversationItemAt { index, reply }
+        })
+        .await
+        .flatten()
+    }
+
+    /// Get the processed text of the last user query (metadata tags stripped).
+    ///
+    /// Equivalent to `extract_last_user_query(&full_conv)` but without cloning
+    /// the full conversation. Returns `None` if there are no user messages or
+    /// the last user message is empty after processing.
+    pub async fn get_last_user_query_text(&self) -> Option<String> {
+        self.query("GetLastUserQueryText", |reply| {
+            ChatStateCommand::GetLastUserQueryText { reply }
+        })
+        .await
+        .flatten()
+    }
+
+    /// Get item counts for the conversation by role.
+    ///
+    /// Returns a [`ConversationCounts`] struct without cloning any items.
+    /// Suitable for telemetry / logging that only needs totals.
+    pub async fn get_conversation_counts(&self) -> ConversationCounts {
+        self.query("GetConversationCounts", |reply| {
+            ChatStateCommand::GetConversationCounts { reply }
+        })
+        .await
+        .unwrap_or_default()
+    }
+
+    /// Get the first `System` message in the conversation, if any.
+    ///
+    /// Cheaper than [`get_conversation`] when only the system prompt is needed
+    /// (e.g. for compaction setup or error validation).
+    pub async fn get_system_message(&self) -> Option<ConversationItem> {
+        self.query("GetSystemMessage", |reply| {
+            ChatStateCommand::GetSystemMessage { reply }
+        })
+        .await
+        .flatten()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn noop_handle_does_not_panic() {
+        let handle = ChatStateHandle::noop();
+        handle.push_user_message(ConversationItem::user("test"));
+        handle.flush();
+        drop(handle);
+    }
+
+    #[test]
+    fn handle_is_clone() {
+        let handle = ChatStateHandle::noop();
+        let clone = handle.clone();
+        clone.push_user_message(ConversationItem::user("from clone"));
+    }
+}