consciousness/src/agent/runner.rs

// agent.rs — Core agent loop
//
// The simplest possible implementation of the agent pattern:
// send messages + tool definitions to the model, if it responds
// with tool calls then dispatch them and loop, if it responds
// with text then display it and wait for the next prompt.
//
// Uses streaming by default so text tokens appear as they're
// generated. Tool calls are accumulated from stream deltas and
// dispatched after the stream completes.
//
// The DMN (dmn.rs) is the outer loop that decides what prompts
// to send here. This module just handles single turns: prompt
// in, response out, tool calls dispatched.

use anyhow::Result;
use tiktoken_rs::CoreBPE;

use crate::agent::api::ApiClient;
use crate::thought::context as journal;
use crate::agent::log::ConversationLog;
use crate::agent::api::StreamEvent;
use crate::agent::tools;
use crate::agent::tools::ProcessTracker;
use crate::agent::types::*;
use crate::agent::ui_channel::{ContextSection, SharedContextState, StatusInfo, StreamTarget, UiMessage, UiSender};

/// Result of a single agent turn.
pub struct TurnResult {
    /// The text response (already sent through UI channel).
    #[allow(dead_code)]
    pub text: String,
    /// Whether the model called yield_to_user during this turn.
    pub yield_requested: bool,
    /// Whether any tools (other than yield_to_user) were called.
    pub had_tool_calls: bool,
    /// Number of tool calls that returned errors this turn.
    pub tool_errors: u32,
    /// Model name to switch to after this turn completes.
    pub model_switch: Option<String>,
    /// Agent requested DMN pause (full stop on autonomous behavior).
    pub dmn_pause: bool,
}

/// Accumulated state across tool dispatches within a single turn.
struct DispatchState {
    yield_requested: bool,
    had_tool_calls: bool,
    tool_errors: u32,
    model_switch: Option<String>,
    dmn_pause: bool,
}

pub struct Agent {
    client: ApiClient,
    tool_defs: Vec<ToolDef>,
    /// Last known prompt token count from the API (tracks context size).
    last_prompt_tokens: u32,
    /// Shared process tracker for bash tool — lets TUI show/kill running commands.
    pub process_tracker: ProcessTracker,
    /// Current reasoning effort level ("none", "low", "high").
    pub reasoning_effort: String,
    /// Persistent conversation log — append-only record of all messages.
    conversation_log: Option<ConversationLog>,
    /// BPE tokenizer for token counting (cl100k_base — close enough
    /// for Claude and Qwen budget allocation, ~85-90% count accuracy).
    tokenizer: CoreBPE,
    /// Mutable context state — personality, working stack, etc.
    pub context: ContextState,
    /// Shared live context summary — TUI reads this directly for debug screen.
    pub shared_context: SharedContextState,
    /// Stable session ID for memory-search dedup across turns.
    session_id: String,
    /// Agent orchestration state (surface-observe, journal, reflect).
    pub agent_cycles: crate::subconscious::subconscious::AgentCycleState,
}

fn render_journal(entries: &[journal::JournalEntry]) -> String {
    if entries.is_empty() { return String::new(); }
    let mut text = String::from("[Earlier — from your journal]\n\n");
    for entry in entries {
        use std::fmt::Write;
        writeln!(text, "## {}\n{}\n", entry.timestamp.format("%Y-%m-%dT%H:%M"), entry.content).ok();
    }
    text
}

impl Agent {
    pub fn new(
        client: ApiClient,
        system_prompt: String,
        personality: Vec<(String, String)>,
        conversation_log: Option<ConversationLog>,
        shared_context: SharedContextState,
    ) -> Self {
        let tool_defs = tools::definitions();
        let tokenizer = tiktoken_rs::cl100k_base()
            .expect("failed to load cl100k_base tokenizer");

        let context = ContextState {
            system_prompt: system_prompt.clone(),
            personality,
            journal: Vec::new(),
            working_stack: Vec::new(),
            entries: Vec::new(),
        };
        let session_id = format!("poc-agent-{}", chrono::Utc::now().format("%Y%m%d-%H%M%S"));
        let agent_cycles = crate::subconscious::subconscious::AgentCycleState::new(&session_id);
        let mut agent = Self {
            client,
            tool_defs,
            last_prompt_tokens: 0,
            process_tracker: ProcessTracker::new(),
            reasoning_effort: "none".to_string(),
            conversation_log,
            tokenizer,
            context,
            shared_context,
            session_id,
            agent_cycles,
        };

        agent.load_startup_journal();
        agent.load_working_stack();
        agent.publish_context_state();
        agent
    }

    /// Assemble the full message list for the API call from typed sources.
    /// System prompt + personality context + journal + conversation messages.
    fn assemble_api_messages(&self) -> Vec<Message> {
        let mut msgs = Vec::new();
        msgs.push(Message::system(&self.context.system_prompt));
        let ctx = self.context.render_context_message();
        if !ctx.is_empty() {
            msgs.push(Message::user(ctx));
        }
        let jnl = render_journal(&self.context.journal);
        if !jnl.is_empty() {
            msgs.push(Message::user(jnl));
        }
        msgs.extend(self.context.entries.iter().map(|e| e.api_message().clone()));
        msgs
    }

    /// Run agent orchestration cycle, returning structured output.
    fn run_agent_cycle(&mut self) -> crate::subconscious::subconscious::AgentCycleOutput {
        let transcript_path = self.conversation_log.as_ref()
            .map(|l| l.path().to_string_lossy().to_string())
            .unwrap_or_default();

        let session = crate::session::HookSession::from_fields(
            self.session_id.clone(),
            transcript_path,
            "UserPromptSubmit".into(),
        );

        self.agent_cycles.trigger(&session);
        std::mem::take(&mut self.agent_cycles.last_output)
    }

    /// Push a conversation message — stamped and logged.
    fn push_message(&mut self, mut msg: Message) {
        msg.stamp();
        let entry = ConversationEntry::Message(msg);
        self.push_entry(entry);
    }

    fn push_entry(&mut self, entry: ConversationEntry) {
        if let Some(ref log) = self.conversation_log {
            if let Err(e) = log.append(&entry) {
                eprintln!("warning: failed to log entry: {:#}", e);
            }
        }
        self.context.entries.push(entry);
    }

    /// Push a context-only message (system prompt, identity context,
    /// journal summaries). Not logged — these are reconstructed on
    /// every startup/compaction.
    pub fn budget(&self) -> ContextBudget {
        let count_str = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
        let count_msg = |m: &Message| crate::thought::context::msg_token_count(&self.tokenizer, m);
        let window = crate::thought::context::model_context_window(&self.client.model);
        self.context.budget(&count_str, &count_msg, window)
    }

    /// Send a user message and run the agent loop until the model
    /// produces a text response (no more tool calls). Streams text
    /// and tool activity through the UI channel.
    pub async fn turn(
        &mut self,
        user_input: &str,
        ui_tx: &UiSender,
        target: StreamTarget,
    ) -> Result<TurnResult> {
        // Run agent orchestration cycle (surface-observe, reflect, journal)
        let cycle = self.run_agent_cycle();

        // Surfaced memories — each as a separate Memory entry
        for key in &cycle.surfaced_keys {
            if let Some(rendered) = crate::cli::node::render_node(
                &crate::store::Store::load().unwrap_or_default(), key,
            ) {
                let mut msg = Message::user(format!(
                    "<system-reminder>\n--- {} (surfaced) ---\n{}\n</system-reminder>",
                    key, rendered,
                ));
                msg.stamp();
                self.push_entry(ConversationEntry::Memory { key: key.clone(), message: msg });
            }
        }

        // Reflection — separate system reminder
        if let Some(ref reflection) = cycle.reflection {
            self.push_message(Message::user(format!(
                "<system-reminder>\n--- subconscious reflection ---\n{}\n</system-reminder>",
                reflection.trim(),
            )));
        }

        // User input — clean, just what was typed
        self.push_message(Message::user(user_input));
        let _ = ui_tx.send(UiMessage::AgentUpdate(self.agent_cycles.snapshots()));

        let mut overflow_retries: u32 = 0;
        let mut empty_retries: u32 = 0;
        let mut ds = DispatchState {
            yield_requested: false,
            had_tool_calls: false,
            tool_errors: 0,
            model_switch: None,
            dmn_pause: false,
        };

        loop {
            let _ = ui_tx.send(UiMessage::Activity("thinking...".into()));

            // Stream events from the API — we route each event to the
            // appropriate UI pane rather than letting the API layer do it.
            let api_messages = self.assemble_api_messages();
            let mut rx = self.client.start_stream(
                &api_messages,
                Some(&self.tool_defs),
                ui_tx,
                &self.reasoning_effort,
                None,
                None, // priority: interactive
            );

            let mut content = String::new();
            let mut tool_calls: Vec<ToolCall> = Vec::new();
            let mut usage = None;
            let mut finish_reason = None;
            let mut in_tool_call = false;
            let mut stream_error = None;
            let mut first_content = true;
            // Buffer for content not yet sent to UI — holds a tail
            // that might be a partial <tool_call> tag.
            let mut display_buf = String::new();

            while let Some(event) = rx.recv().await {
                match event {
                    StreamEvent::Content(text) => {
                        if first_content {
                            let _ = ui_tx.send(UiMessage::Activity("streaming...".into()));
                            first_content = false;
                        }
                        content.push_str(&text);

                        if in_tool_call {
                            // Already inside a tool call — suppress display.
                        } else {
                            display_buf.push_str(&text);

                            if let Some(pos) = display_buf.find("<tool_call>") {
                                // Flush content before the tag, suppress the rest.
                                let before = &display_buf[..pos];
                                if !before.is_empty() {
                                    let _ = ui_tx.send(UiMessage::TextDelta(before.to_string(), target));
                                }
                                display_buf.clear();
                                in_tool_call = true;
                            } else {
                                // Flush display_buf except a tail that could be
                                // a partial "<tool_call>" (10 chars).
                                let safe = display_buf.len().saturating_sub(10);
                                // Find a char boundary at or before safe
                                let safe = display_buf.floor_char_boundary(safe);
                                if safe > 0 {
                                    let flush = display_buf[..safe].to_string();
                                    display_buf = display_buf[safe..].to_string();
                                    let _ = ui_tx.send(UiMessage::TextDelta(flush, target));
                                }
                            }
                        }
                    }
                    StreamEvent::Reasoning(text) => {
                        let _ = ui_tx.send(UiMessage::Reasoning(text));
                    }
                    StreamEvent::ToolCallDelta { index, id, call_type, name, arguments } => {
                        while tool_calls.len() <= index {
                            tool_calls.push(ToolCall {
                                id: String::new(),
                                call_type: "function".to_string(),
                                function: FunctionCall { name: String::new(), arguments: String::new() },
                            });
                        }
                        if let Some(id) = id { tool_calls[index].id = id; }
                        if let Some(ct) = call_type { tool_calls[index].call_type = ct; }
                        if let Some(n) = name { tool_calls[index].function.name = n; }
                        if let Some(a) = arguments { tool_calls[index].function.arguments.push_str(&a); }
                    }
                    StreamEvent::Usage(u) => usage = Some(u),
                    StreamEvent::Finished { reason, .. } => {
                        finish_reason = Some(reason);
                        break;
                    }
                    StreamEvent::Error(e) => {
                        stream_error = Some(e);
                        break;
                    }
                }
            }

            // Handle stream errors with retry logic
            if let Some(e) = stream_error {
                let err = anyhow::anyhow!("{}", e);
                if crate::thought::context::is_context_overflow(&err) && overflow_retries < 2 {
                    overflow_retries += 1;
                    let _ = ui_tx.send(UiMessage::Info(format!(
                        "[context overflow — compacting and retrying ({}/2)]",
                        overflow_retries,
                    )));
                    self.emergency_compact();
                    continue;
                }
                if crate::thought::context::is_stream_error(&err) && empty_retries < 2 {
                    empty_retries += 1;
                    let _ = ui_tx.send(UiMessage::Info(format!(
                        "[stream error: {} — retrying ({}/2)]",
                        e, empty_retries,
                    )));
                    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
                    continue;
                }
                return Err(err);
            }

            if finish_reason.as_deref() == Some("error") {
                let detail = if content.is_empty() { "no details".into() } else { content };
                return Err(anyhow::anyhow!("model stream error: {}", detail));
            }

            // Flush remaining display buffer (normal responses without tool calls).
            if !in_tool_call && !display_buf.is_empty() {
                let _ = ui_tx.send(UiMessage::TextDelta(display_buf, target));
            }
            if !content.is_empty() && !in_tool_call {
                let _ = ui_tx.send(UiMessage::TextDelta("\n".to_string(), target));
            }

            let msg = crate::agent::api::build_response_message(content, tool_calls);



            if let Some(usage) = &usage {
                self.last_prompt_tokens = usage.prompt_tokens;
                
                self.publish_context_state();
                let _ = ui_tx.send(UiMessage::StatusUpdate(StatusInfo {
                    dmn_state: String::new(), // filled by main loop
                    dmn_turns: 0,
                    dmn_max_turns: 0,
                    prompt_tokens: usage.prompt_tokens,
                    completion_tokens: usage.completion_tokens,
                    model: self.client.model.clone(),
                    turn_tools: 0, // tracked by TUI from ToolCall messages
                    context_budget: self.budget().status_string(),
                }));
            }

            // Empty response — model returned finish=stop with no content
            // or tool calls. Inject a nudge so the retry has different input.
            let has_content = msg.content.is_some();
            let has_tools = msg.tool_calls.as_ref().map_or(false, |tc| !tc.is_empty());
            if !has_content && !has_tools {
                if empty_retries < 2 {
                    empty_retries += 1;
                    let _ = ui_tx.send(UiMessage::Debug(format!(
                        "empty response, injecting nudge and retrying ({}/2)",
                        empty_retries,
                    )));
                    self.push_message(Message::user(
                        "[system] Your previous response was empty. \
                         Please respond with text or use a tool."
                    ));
                    continue;
                }
                // After max retries, fall through — return the empty response
            } else {
                empty_retries = 0;
            }

            // Tool calls (structured from API, or recovered from content
            // by build_response_message if the model leaked them as XML).
            if let Some(ref tool_calls) = msg.tool_calls {
                if !tool_calls.is_empty() {
                    self.push_message(msg.clone());
                    for call in tool_calls {
                        self.dispatch_tool_call(call, None, ui_tx, &mut ds)
                            .await;
                    }
                    continue;
                }
            }

            // Genuinely text-only response
            let text = msg.content_text().to_string();
            let _ = ui_tx.send(UiMessage::Activity(String::new()));
            self.push_message(msg);

            return Ok(TurnResult {
                text,
                yield_requested: ds.yield_requested,
                had_tool_calls: ds.had_tool_calls,
                tool_errors: ds.tool_errors,
                model_switch: ds.model_switch,
                dmn_pause: ds.dmn_pause,
            });
        }
    }

    /// Dispatch a single tool call: send UI annotations, run the tool,
    /// push results into the conversation, handle images.
    async fn dispatch_tool_call(
        &mut self,
        call: &ToolCall,
        tag: Option<&str>,
        ui_tx: &UiSender,
        ds: &mut DispatchState,
    ) {
        let args: serde_json::Value =
            serde_json::from_str(&call.function.arguments).unwrap_or_default();

        let args_summary = summarize_args(&call.function.name, &args);
        let label = match tag {
            Some(t) => format!("calling: {} ({})", call.function.name, t),
            None => format!("calling: {}", call.function.name),
        };
        let _ = ui_tx.send(UiMessage::Activity(label));
        let _ = ui_tx.send(UiMessage::ToolCall {
            name: call.function.name.clone(),
            args_summary: args_summary.clone(),
        });
        let _ = ui_tx.send(UiMessage::ToolStarted {
            id: call.id.clone(),
            name: call.function.name.clone(),
            detail: args_summary,
        });

        // Handle working_stack tool — needs &mut self for context state
        if call.function.name == "working_stack" {
            let result = tools::working_stack::handle(&args, &mut self.context.working_stack);
            let output = tools::ToolOutput {
                text: result.clone(),
                is_yield: false,
                images: Vec::new(),
                model_switch: None,
                dmn_pause: false,
            };
            let _ = ui_tx.send(UiMessage::ToolResult {
                name: call.function.name.clone(),
                result: output.text.clone(),
            });
            let _ = ui_tx.send(UiMessage::ToolFinished { id: call.id.clone() });
            self.push_message(Message::tool_result(&call.id, &output.text));
            ds.had_tool_calls = true;
            
            // Re-render the context message so the model sees the updated stack
            if !result.starts_with("Error:") {
                self.refresh_context_state();
            }
            return;
        }

        // Handle memory tools — needs &mut self for node tracking
        if call.function.name.starts_with("memory_") {
            let result = tools::memory::dispatch(&call.function.name, &args, None);
            let text = match &result {
                Ok(s) => s.clone(),
                Err(e) => format!("Error: {:#}", e),
            };

            // Disambiguate memory renders from other tool results
            let memory_key = if result.is_ok() {
                match call.function.name.as_str() {
                    "memory_render" =>
                        args.get("key").and_then(|v| v.as_str()).map(String::from),
                    _ => None,
                }
            } else {
                None
            };

            let output = tools::ToolOutput {
                text,
                is_yield: false,
                images: Vec::new(),
                model_switch: None,
                dmn_pause: false,
            };
            let _ = ui_tx.send(UiMessage::ToolResult {
                name: call.function.name.clone(),
                result: output.text.clone(),
            });
            let _ = ui_tx.send(UiMessage::ToolFinished { id: call.id.clone() });
            let mut msg = Message::tool_result(&call.id, &output.text);
            msg.stamp();
            if let Some(key) = memory_key {
                self.push_entry(ConversationEntry::Memory { key, message: msg });
            } else {
                self.push_entry(ConversationEntry::Message(msg));
            }
            ds.had_tool_calls = true;
            if output.text.starts_with("Error:") {
                ds.tool_errors += 1;
            }
            
            self.publish_context_state();
            return;
        }

        let output =
            tools::dispatch(&call.function.name, &args, &self.process_tracker).await;

        if output.is_yield {
            ds.yield_requested = true;
        } else {
            ds.had_tool_calls = true;
        }
        if output.model_switch.is_some() {
            ds.model_switch = output.model_switch;
        }
        if output.dmn_pause {
            ds.dmn_pause = true;
        }
        if output.text.starts_with("Error:") {
            ds.tool_errors += 1;
        }

        let _ = ui_tx.send(UiMessage::ToolResult {
            name: call.function.name.clone(),
            result: output.text.clone(),
        });
        let _ = ui_tx.send(UiMessage::ToolFinished { id: call.id.clone() });

        self.push_message(Message::tool_result(&call.id, &output.text));

        if !output.images.is_empty() {
            // Only one live image in context at a time — age out any
            // previous ones to avoid accumulating ~90KB+ per image.
            self.age_out_images();
            self.push_message(Message::user_with_images(
                "Here is the image you requested:",
                &output.images,
            ));
        }
    }

    /// Build context state summary for the debug screen.
    pub fn context_state_summary(&self) -> Vec<ContextSection> {
        let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();

        let mut sections = Vec::new();

        // System prompt
        sections.push(ContextSection {
            name: "System prompt".into(),
            tokens: count(&self.context.system_prompt),
            content: self.context.system_prompt.clone(),
            children: Vec::new(),
        });

        // Personality — parent with file children
        let personality_children: Vec<ContextSection> = self.context.personality.iter()
            .map(|(name, content)| ContextSection {
                name: name.clone(),
                tokens: count(content),
                content: content.clone(),
                children: Vec::new(),
            })
            .collect();
        let personality_tokens: usize = personality_children.iter().map(|c| c.tokens).sum();
        sections.push(ContextSection {
            name: format!("Personality ({} files)", personality_children.len()),
            tokens: personality_tokens,
            content: String::new(),
            children: personality_children,
        });

        // Journal
        {
            let journal_children: Vec<ContextSection> = self.context.journal.iter()
                .map(|entry| {
                    let preview: String = entry.content.lines()
                        .find(|l| !l.trim().is_empty())
                        .unwrap_or("").chars().take(60).collect();
                    ContextSection {
                        name: format!("{}: {}", entry.timestamp.format("%Y-%m-%dT%H:%M"), preview),
                        tokens: count(&entry.content),
                        content: entry.content.clone(),
                        children: Vec::new(),
                    }
                })
                .collect();
            let journal_tokens: usize = journal_children.iter().map(|c| c.tokens).sum();
            sections.push(ContextSection {
                name: format!("Journal ({} entries)", journal_children.len()),
                tokens: journal_tokens,
                content: String::new(),
                children: journal_children,
            });
        }

        // Working stack — instructions + items as children
        let instructions = std::fs::read_to_string(WORKING_STACK_INSTRUCTIONS)
            .unwrap_or_default();
        let mut stack_children = vec![ContextSection {
            name: "Instructions".into(),
            tokens: count(&instructions),
            content: instructions,
            children: Vec::new(),
        }];
        for (i, item) in self.context.working_stack.iter().enumerate() {
            let marker = if i == self.context.working_stack.len() - 1 { "→" } else { " " };
            stack_children.push(ContextSection {
                name: format!("{} [{}] {}", marker, i, item),
                tokens: count(item),
                content: String::new(),
                children: Vec::new(),
            });
        }
        let stack_tokens: usize = stack_children.iter().map(|c| c.tokens).sum();
        sections.push(ContextSection {
            name: format!("Working stack ({} items)", self.context.working_stack.len()),
            tokens: stack_tokens,
            content: String::new(),
            children: stack_children,
        });

        // Memory nodes — extracted from Memory entries in the conversation
        let memory_entries: Vec<&ConversationEntry> = self.context.entries.iter()
            .filter(|e| e.is_memory())
            .collect();
        if !memory_entries.is_empty() {
            let node_children: Vec<ContextSection> = memory_entries.iter()
                .map(|entry| {
                    let key = match entry {
                        ConversationEntry::Memory { key, .. } => key.as_str(),
                        _ => unreachable!(),
                    };
                    let text = entry.message().content_text();
                    ContextSection {
                        name: key.to_string(),
                        tokens: count(text),
                        content: String::new(),
                        children: Vec::new(),
                    }
                })
                .collect();
            let node_tokens: usize = node_children.iter().map(|c| c.tokens).sum();
            sections.push(ContextSection {
                name: format!("Memory nodes ({} loaded)", memory_entries.len()),
                tokens: node_tokens,
                content: String::new(),
                children: node_children,
            });
        }

        // Conversation — each message as a child
        let conv_messages = &self.context.entries;
        let conv_children: Vec<ContextSection> = conv_messages.iter().enumerate()
            .map(|(i, entry)| {
                let m = entry.message();
                let text = m.content.as_ref()
                    .map(|c| c.as_text().to_string())
                    .unwrap_or_default();
                let tool_info = m.tool_calls.as_ref().map(|tc| {
                    tc.iter()
                        .map(|c| c.function.name.clone())
                        .collect::<Vec<_>>()
                        .join(", ")
                });
                let label = if entry.is_memory() {
                    if let ConversationEntry::Memory { key, .. } = entry {
                        format!("[memory: {}]", key)
                    } else { unreachable!() }
                } else {
                    match &tool_info {
                        Some(tools) => format!("[tool_call: {}]", tools),
                        None => {
                            let preview: String = text.chars().take(60).collect();
                            let preview = preview.replace('\n', " ");
                            if text.len() > 60 { format!("{}...", preview) } else { preview }
                        }
                    }
                };
                let tokens = count(&text);
                let role_name = if entry.is_memory() { "mem" } else {
                    match m.role {
                        Role::Assistant => "PoC",
                        Role::User => "Kent",
                        Role::Tool => "tool",
                        Role::System => "system",
                    }
                };
                ContextSection {
                    name: format!("[{}] {}: {}", i, role_name, label),
                    tokens,
                    content: text,
                    children: Vec::new(),
                }
            })
            .collect();
        let conv_tokens: usize = conv_children.iter().map(|c| c.tokens).sum();
        sections.push(ContextSection {
            name: format!("Conversation ({} messages)", conv_children.len()),
            tokens: conv_tokens,
            content: String::new(),
            children: conv_children,
        });

        sections
    }

    /// Load recent journal entries at startup for orientation.
    /// Uses the same budget logic as compaction but with empty conversation.
    /// Only parses the tail of the journal file (last 64KB) for speed.
    fn load_startup_journal(&mut self) {
        let store = match crate::store::Store::load() {
            Ok(s) => s,
            Err(_) => return,
        };

        // Find oldest message timestamp in conversation log
        let oldest_msg_ts = self.conversation_log.as_ref()
            .and_then(|log| log.oldest_timestamp());

        // Get journal entries from the memory graph
        let mut journal_nodes: Vec<_> = store.nodes.values()
            .filter(|n| n.node_type == crate::store::NodeType::EpisodicSession)
            .collect();
        let mut dbg = std::fs::OpenOptions::new().create(true).append(true)
                .open("/tmp/poc-journal-debug.log").ok();
        macro_rules! dbg_log {
            ($($arg:tt)*) => {
                if let Some(ref mut f) = dbg { use std::io::Write; let _ = writeln!(f, $($arg)*); }
            }
        }
        dbg_log!("[journal] {} nodes, oldest_msg={:?}", journal_nodes.len(), oldest_msg_ts);

        journal_nodes.sort_by_key(|n| n.created_at);
        if let Some(first) = journal_nodes.first() {
            dbg_log!("[journal] first created_at={}", first.created_at);
        }
        if let Some(last) = journal_nodes.last() {
            dbg_log!("[journal] last created_at={}", last.created_at);
        }

        // Find the cutoff index — entries older than conversation, plus one overlap
        let cutoff_idx = if let Some(cutoff) = oldest_msg_ts {
            let cutoff_ts = cutoff.timestamp();
            dbg_log!("[journal] cutoff timestamp={}", cutoff_ts);
            let mut idx = journal_nodes.len();
            for (i, node) in journal_nodes.iter().enumerate() {
                if node.created_at >= cutoff_ts {
                    idx = i + 1;
                    break;
                }
            }
            idx
        } else {
            journal_nodes.len()
        };
        dbg_log!("[journal] cutoff_idx={}", cutoff_idx);

        // Walk backwards from cutoff, accumulating entries within 5% of context
        let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
        let context_window = crate::thought::context::model_context_window(&self.client.model);
        let journal_budget = context_window * 5 / 100;
        dbg_log!("[journal] budget={} tokens ({}*5%)", journal_budget, context_window);

        let mut entries = Vec::new();
        let mut total_tokens = 0;

        for node in journal_nodes[..cutoff_idx].iter().rev() {
            let tokens = count(&node.content);
            if total_tokens + tokens > journal_budget && !entries.is_empty() {
                break;
            }
            entries.push(journal::JournalEntry {
                timestamp: chrono::DateTime::from_timestamp(node.created_at, 0)
                    .unwrap_or_default(),
                content: node.content.clone(),
            });
            total_tokens += tokens;
        }
        entries.reverse();
        dbg_log!("[journal] loaded {} entries, {} tokens", entries.len(), total_tokens);

        if entries.is_empty() {
            dbg_log!("[journal] no entries!");
            return;
        }

        self.context.journal = entries;
        dbg_log!("[journal] context.journal now has {} entries", self.context.journal.len());
    }

    /// Called after any change to context state (working stack, etc).
    fn refresh_context_state(&mut self) {
        
        self.publish_context_state();
        self.save_working_stack();
    }

    /// Persist working stack to disk.
    fn save_working_stack(&self) {
        if let Ok(json) = serde_json::to_string(&self.context.working_stack) {
            let _ = std::fs::write(WORKING_STACK_FILE, json);
        }
    }

    /// Load working stack from disk.
    fn load_working_stack(&mut self) {
        if let Ok(data) = std::fs::read_to_string(WORKING_STACK_FILE) {
            if let Ok(stack) = serde_json::from_str::<Vec<String>>(&data) {
                self.context.working_stack = stack;
            }
        }
    }

    /// Push the current context summary to the shared state for the TUI to read.
    fn publish_context_state(&self) {
        let summary = self.context_state_summary();
        if let Ok(mut dbg) = std::fs::OpenOptions::new().create(true).append(true)
                .open("/tmp/poc-journal-debug.log") {
            use std::io::Write;
            for s in &summary {
                let _ = writeln!(dbg, "[publish] {} ({} tokens, {} children)", s.name, s.tokens, s.children.len());
            }
        }
        if let Ok(mut state) = self.shared_context.write() {
            *state = summary;
        }
    }

    /// Replace base64 image data in older messages with text placeholders.
    /// Only the most recent image stays live — each new image ages out
    /// all previous ones. The tool result message (right before each image
    /// message) already records what was loaded, so no info is lost.
    fn age_out_images(&mut self) {
        for entry in &mut self.context.entries {
            let msg = entry.message_mut();
            if let Some(MessageContent::Parts(parts)) = &msg.content {
                let has_images = parts.iter().any(|p| matches!(p, ContentPart::ImageUrl { .. }));
                if !has_images {
                    continue;
                }
                let mut replacement = String::new();
                for part in parts {
                    match part {
                        ContentPart::Text { text } => {
                            if !replacement.is_empty() {
                                replacement.push('\n');
                            }
                            replacement.push_str(text);
                        }
                        ContentPart::ImageUrl { .. } => {
                            if !replacement.is_empty() {
                                replacement.push('\n');
                            }
                            replacement.push_str(
                                "[image aged out — see tool result above for details]",
                            );
                        }
                    }
                }
                msg.content = Some(MessageContent::Text(replacement));
            }
        }
    }

    /// Strip ephemeral tool calls from the conversation history.
    ///
    /// Last prompt token count reported by the API.
    pub fn last_prompt_tokens(&self) -> u32 {
        self.last_prompt_tokens
    }

    /// Build context window from conversation messages + journal.
    /// Used by both compact() (in-memory messages) and restore_from_log()
    /// (conversation log). The context window is always:
    ///   identity + journal summaries + raw recent messages
    pub fn compact(&mut self, new_system_prompt: String, new_personality: Vec<(String, String)>) {
        self.context.system_prompt = new_system_prompt;
        self.context.personality = new_personality;
        self.do_compact();
    }

    /// Internal compaction — rebuilds context window from current messages.
    fn do_compact(&mut self) {
        let conversation: Vec<Message> = self.context.entries.iter()
            .map(|e| e.api_message().clone()).collect();
        let messages = crate::thought::context::trim_conversation(
            &self.context,
            &conversation,
            &self.client.model,
            &self.tokenizer,
        );
        self.context.entries = messages.into_iter()
            .map(ConversationEntry::Message).collect();
        self.last_prompt_tokens = 0;

        self.publish_context_state();
    }

    /// Emergency compaction using stored config — called on context overflow.
    fn emergency_compact(&mut self) {
        self.do_compact();
    }

    /// Restore from the conversation log. Builds the context window
    /// the same way compact() does — journal summaries for old messages,
    /// raw recent messages. This is the unified startup path.
    /// Returns true if the log had content to restore.
    pub fn restore_from_log(
        &mut self,
        system_prompt: String,
        personality: Vec<(String, String)>,
    ) -> bool {
        self.context.system_prompt = system_prompt;
        self.context.personality = personality;

        let entries = match &self.conversation_log {
            Some(log) => match log.read_tail(512 * 1024) {
                Ok(entries) if !entries.is_empty() => {
                    dbglog!("[restore] read {} entries from log tail", entries.len());
                    entries
                }
                Ok(_) => {
                    dbglog!("[restore] log exists but is empty");
                    return false;
                }
                Err(e) => {
                    dbglog!("[restore] failed to read log: {}", e);
                    return false;
                }
            },
            None => {
                dbglog!("[restore] no conversation log configured");
                return false;
            }
        };

        // Filter out system messages, keep everything else (including Memory entries)
        let entries: Vec<ConversationEntry> = entries
            .into_iter()
            .filter(|e| e.message().role != Role::System)
            .collect();

        // Trim to fit context budget
        let n = entries.len();
        let conversation: Vec<Message> = entries.iter()
            .map(|e| e.api_message().clone()).collect();
        let trimmed = crate::thought::context::trim_conversation(
            &self.context,
            &conversation,
            &self.client.model,
            &self.tokenizer,
        );
        // Keep only the entries that survived trimming (by count from the end)
        let keep = trimmed.len();
        self.context.entries = entries.into_iter()
            .skip(n.saturating_sub(keep))
            .collect();

        dbglog!("[restore] {} entries, journal: {} entries",
            self.context.entries.len(), self.context.journal.len());
        self.last_prompt_tokens = 0;
        self.publish_context_state();
        true
    }

    /// Replace the API client (for model switching).
    pub fn swap_client(&mut self, new_client: ApiClient) {
        self.client = new_client;
    }

    /// Get the model identifier.
    pub fn model(&self) -> &str {
        &self.client.model
    }

    /// Get the conversation entries for persistence.
    pub fn entries(&self) -> &[ConversationEntry] {
        &self.context.entries
    }

    /// Mutable access to conversation entries (for /retry).
    pub fn entries_mut(&mut self) -> &mut Vec<ConversationEntry> {
        &mut self.context.entries
    }

}

// Context window building, token counting, and error classification
// live in context.rs


/// Create a short summary of tool args for the tools pane header.
fn summarize_args(tool_name: &str, args: &serde_json::Value) -> String {
    match tool_name {
        "read_file" | "write_file" | "edit_file" => args["file_path"]
            .as_str()
            .unwrap_or("")
            .to_string(),
        "bash" => {
            let cmd = args["command"].as_str().unwrap_or("");
            if cmd.len() > 60 {
                let end = cmd.char_indices()
                    .map(|(i, _)| i)
                    .take_while(|&i| i <= 60)
                    .last()
                    .unwrap_or(0);
                format!("{}...", &cmd[..end])
            } else {
                cmd.to_string()
            }
        }
        "grep" => {
            let pattern = args["pattern"].as_str().unwrap_or("");
            let path = args["path"].as_str().unwrap_or(".");
            format!("{} in {}", pattern, path)
        }
        "glob" => args["pattern"]
            .as_str()
            .unwrap_or("")
            .to_string(),
        "view_image" => {
            if let Some(pane) = args["pane_id"].as_str() {
                format!("pane {}", pane)
            } else {
                args["file_path"].as_str().unwrap_or("").to_string()
            }
        }
        "journal" => {
            let entry = args["entry"].as_str().unwrap_or("");
            if entry.len() > 60 {
                format!("{}...", &entry[..60])
            } else {
                entry.to_string()
            }
        }
        "yield_to_user" => args["message"]
            .as_str()
            .unwrap_or("")
            .to_string(),
        "switch_model" => args["model"]
            .as_str()
            .unwrap_or("")
            .to_string(),
        "pause" => String::new(),
        _ => String::new(),
    }
}

// Parsing functions (parse_leaked_tool_calls, strip_leaked_artifacts)
// and their tests live in parsing.rs