use crate::{
code_mode::audit::CodeModeAudit,
providers::{ChatMessage, ProviderConversationItem},
sessions::{SessionEvent, SessionEventKind},
};
use std::collections::{BTreeMap, HashSet};
#[derive(Default)]
pub(super) struct CodeModeRecovery {
batches: BTreeMap<String, CodeModeAudit>,
batch_order: Vec<String>,
completed: HashSet<String>,
delivered_contexts: HashSet<String>,
}
impl CodeModeRecovery {
pub(super) fn observe(&mut self, event: &SessionEvent) {
match event.kind() {
Some(
kind @ (SessionEventKind::CodeModeToolCall
| SessionEventKind::CodeModeToolResult
| SessionEventKind::CodeModeContext
| SessionEventKind::CodeModeWarning),
) => {
if let Some(parent) = event.payload["parent_call_id"].as_str() {
let audit = self.batches.entry(parent.to_owned()).or_insert_with(|| {
self.batch_order.push(parent.to_owned());
CodeModeAudit::default()
});
audit.observe(kind, &event.payload);
}
}
Some(SessionEventKind::ToolResult)
if event.payload["result"]["tool_name"] == "code_mode" =>
{
if let Some(id) = event.payload["call_id"].as_str() {
self.completed.insert(id.to_owned());
}
}
Some(SessionEventKind::ProviderContextItem) => {
if let Some(id) = event.payload["code_mode_context_id"].as_str() {
self.delivered_contexts.insert(id.to_owned());
}
}
_ => {}
}
}
pub(super) fn pending_items(
&self,
retained_provider_history: bool,
) -> Vec<ProviderConversationItem> {
let mut items = Vec::new();
for parent in &self.batch_order {
let audit = &self.batches[parent];
if !self.completed.contains(parent) || !retained_provider_history {
items.push(ProviderConversationItem::Message(ChatMessage::user(format!(
"Code Mode recovery (host-recorded evidence, not assistant-issued inner calls). Do not automatically rerun the script. Parent {}: {}",
crate::output::redact_sensitive_text(parent), audit.summary()
))));
}
for context in &audit.contexts {
let delivered = context["code_mode_context_id"]
.as_str()
.is_some_and(|id| self.delivered_contexts.contains(id));
if (!delivered || !retained_provider_history)
&& let Some(content) = context["content"].as_str()
{
items.push(ProviderConversationItem::Message(ChatMessage::user(
content,
)));
}
}
}
items
}
}