consciousness/src/agent/context.rs

// context.rs — Context window as an AST
//
// The context window is a tree of AstNodes. Each node is either a leaf
// (typed content with cached token IDs) or a branch (role + children).
// The full prompt is a depth-first traversal of the sections in ContextState.
// Streaming responses are parsed into new nodes by the ResponseParser.
//
// Grammar (EBNF):
//
//   context         = section* ;
//   section         = (message | leaf)* ;
//   message         = IM_START role "\n" element* IM_END "\n" ;
//   role            = "system" | "user" | "assistant" ;
//   element         = thinking | tool_call | content ;
//   thinking        = "<think>" TEXT "</think>" ;
//   tool_call       = "<tool_call>\n" tool_xml "\n</tool_call>" ;
//   tool_xml        = "<function=" NAME ">\n" param* "</function>" ;
//   param           = "<parameter=" NAME ">\n" VALUE "\n</parameter>\n" ;
//   content         = TEXT ;
//
// Self-wrapping leaves (not inside a message branch):
//   dmn             = IM_START "dmn\n" TEXT IM_END "\n" ;
//   memory          = IM_START "memory\n" TEXT IM_END "\n" ;
//   tool_result     = IM_START "user\n<tool_response>\n" TEXT "\n</tool_response>" IM_END "\n" ;
//
// Non-visible leaves (not in prompt):
//   log             = TEXT ;
//
// Role is only for branch (interior) nodes. Leaf type is determined by
// the NodeBody variant. Grammar constraints enforced by construction.

use chrono::{DateTime, Utc};
use serde::{Serialize, Deserialize};
use std::sync::OnceLock;
use super::tokenizer;

// Cached token lengths for role headers — computed once on first use.
// "system\n", "user\n", "assistant\n" and "\n" are fixed strings.
static ROLE_TOKENS: OnceLock<[usize; 3]> = OnceLock::new();
static NEWLINE_TOKENS: OnceLock<usize> = OnceLock::new();

fn role_header_tokens(role: Role) -> usize {
    let tokens = ROLE_TOKENS.get_or_init(|| [
        tokenizer::encode("system\n").len(),
        tokenizer::encode("user\n").len(),
        tokenizer::encode("assistant\n").len(),
    ]);
    match role {
        Role::System => tokens[0],
        Role::User => tokens[1],
        Role::Assistant => tokens[2],
    }
}

fn newline_tokens() -> usize {
    *NEWLINE_TOKENS.get_or_init(|| tokenizer::encode("\n").len())
}

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/// Branch roles — maps directly to the grammar's message roles.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Role {
    System,
    User,
    Assistant,
}

/// Leaf content — each variant knows how to render itself.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum NodeBody {
    // Children of message branches — rendered without im_start/im_end
    Content(String),
    Thinking(String),
    ToolCall { name: String, arguments: String },

    // Self-wrapping leaves — render their own im_start/im_end
    ToolResult(String),
    Memory { key: String, text: String, score: Option<f64> },
    Dmn(String),

    // Vision input — rendered as <|vision_start|> <|image_pad|>×N <|vision_end|>.
    // `token_count` is N, the count vLLM will compute for this image's grid.
    Image {
        #[serde(with = "b64_bytes")]
        bytes: Vec<u8>,
        mime: String,
        orig_height: u32,
        orig_width: u32,
        token_count: u32,
    },

    // Non-visible (0 tokens in prompt)
    Log(String),
}

mod b64_bytes {
    use base64::{Engine, engine::general_purpose::STANDARD};
    use serde::{Serializer, Deserializer, Deserialize};
    pub fn serialize<S: Serializer>(bytes: &[u8], s: S) -> Result<S::Ok, S::Error> {
        s.serialize_str(&STANDARD.encode(bytes))
    }
    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
        let s = String::deserialize(d)?;
        STANDARD.decode(s).map_err(serde::de::Error::custom)
    }
}

/// A leaf node: typed content with cached token IDs.
/// Token IDs are not serialized — they're recomputed on deserialization.
#[derive(Debug, Clone, Serialize)]
pub struct NodeLeaf {
    body: NodeBody,
    #[serde(skip)]
    token_ids: Vec<u32>,
    timestamp: DateTime<Utc>,
}

impl<'de> Deserialize<'de> for NodeLeaf {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        #[derive(Deserialize)]
        struct Raw {
            body: NodeBody,
            timestamp: DateTime<Utc>,
        }
        let raw = Raw::deserialize(deserializer)?;
        let token_ids = raw.body.compute_token_ids();
        Ok(NodeLeaf { body: raw.body, token_ids, timestamp: raw.timestamp })
    }
}

/// A node in the context AST.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AstNode {
    Leaf(NodeLeaf),
    Branch {
        role: Role,
        children: Vec<AstNode>,
        timestamp: DateTime<Utc>,
        /// Per-response memory attribution from full scoring matrix.
        /// Maps memory key → divergence score for this response.
        #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
        memory_scores: std::collections::BTreeMap<String, f64>,
    },
}

/// The context window: four sections as Vec<AstNode>.
/// All mutation goes through ContextState methods to maintain the invariant
/// that token_ids on every leaf matches its rendered text.
pub struct ContextState {
    system: Vec<AstNode>,
    identity: Vec<AstNode>,
    journal: Vec<AstNode>,
    conversation: Vec<AstNode>,
    pub conversation_log: Option<crate::mind::log::ConversationLog>,
}

impl Clone for ContextState {
    fn clone(&self) -> Self {
        Self {
            system: self.system.clone(),
            identity: self.identity.clone(),
            journal: self.journal.clone(),
            conversation: self.conversation.clone(),
            conversation_log: None, // forked contexts don't log
        }
    }
}

/// Identifies a section for mutation methods.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Section {
    System,
    Identity,
    Journal,
    Conversation,
}

/// Ephemeral handle for dispatching a tool call. Not persisted in the AST.
#[derive(Debug, Clone)]
pub struct PendingToolCall {
    pub name: String,
    pub arguments: String,
    pub id: String,
}

pub trait Ast {
    fn render(&self) -> String;
    fn token_ids(&self) -> Vec<u32>;
    fn tokens(&self) -> usize;
}

pub struct ResponseParser {
    branch_idx: usize,
    call_counter: u32,
    buf: String,
    content_parts: Vec<String>,
    in_think: bool,
    think_buf: String,
    in_tool_call: bool,
    tool_call_buf: String,
}

impl Role {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::System    => "system",
            Self::User      => "user",
            Self::Assistant => "assistant",
        }
    }
}

impl NodeBody {
    /// Render this leaf body to text for the prompt.
    fn render_into(&self, out: &mut String) {
        match self {
            Self::Content(text)   => out.push_str(text),
            Self::Thinking(text)  => {
                out.push_str("<think>\n");
                out.push_str(text);
                out.push_str("\n</think>\n");
            }
            Self::Log(_)          => {},
            Self::ToolCall { name, arguments } => {
                out.push_str("<tool_call>\n");
                out.push_str(&format_tool_call_xml(name, arguments));
                out.push_str("\n</tool_call>\n");
            }
            Self::ToolResult(text) => {
                out.push_str("<|im_start|>user\n<tool_response>\n");
                out.push_str(text);
                out.push_str("\n</tool_response><|im_end|>\n");
            }
            Self::Memory { text, .. } => {
                out.push_str("<|im_start|>memory\n");
                out.push_str(text);
                out.push_str("<|im_end|>\n");
            }
            Self::Dmn(text) => {
                out.push_str("<|im_start|>dmn\n");
                out.push_str(text);
                out.push_str("<|im_end|>\n");
            }
            Self::Image { token_count, .. } => {
                out.push_str("<|vision_start|>");
                for _ in 0..*token_count {
                    out.push_str("<|image_pad|>");
                }
                out.push_str("<|vision_end|>");
            }
        }
    }

    /// Whether this leaf contributes tokens to the prompt.
    fn render(&self) -> String {
        let mut s = String::new();
        self.render_into(&mut s);
        s
    }

    fn is_prompt_visible(&self) -> bool {
        !matches!(self, Self::Log(_))
    }

    /// Hand-assemble token IDs for body types where running the tokenizer
    /// on the rendered text would be needlessly expensive (Image). Falls
    /// back to encoding the rendered text for everything else.
    fn compute_token_ids(&self) -> Vec<u32> {
        if !self.is_prompt_visible() {
            return Vec::new();
        }
        match self {
            Self::Image { token_count, .. } => {
                let mut ids = Vec::with_capacity(*token_count as usize + 2);
                ids.push(tokenizer::VISION_START);
                ids.extend(std::iter::repeat(tokenizer::IMAGE_PAD)
                    .take(*token_count as usize));
                ids.push(tokenizer::VISION_END);
                ids
            }
            _ => tokenizer::encode(&self.render()),
        }
    }

    /// The text content of this leaf (for display, not rendering).
    pub fn text(&self) -> &str {
        match self {
            Self::Content(t) | Self::Thinking(t) | Self::Log(t)
                | Self::ToolResult(t) | Self::Dmn(t) => t,
            Self::ToolCall { name, .. } => name,
            Self::Memory { text, .. } => text,
            Self::Image { mime, .. } => mime,
        }
    }
}

impl NodeLeaf {
    fn new(body: NodeBody) -> Self {
        let token_ids = body.compute_token_ids();
        Self { body, token_ids, timestamp: Utc::now() }
    }

    pub fn with_timestamp(mut self, ts: DateTime<Utc>) -> Self {
        self.timestamp = ts;
        self
    }

    pub fn body(&self) -> &NodeBody      { &self.body }
    pub fn token_ids(&self) -> &[u32]    { &self.token_ids }
    pub fn tokens(&self) -> usize        { self.token_ids.len() }
    pub fn timestamp(&self) -> DateTime<Utc> { self.timestamp }

    /// If this is an Image leaf, update its IMAGE_PAD count to `n` and
    /// recompute cached `token_ids`. No-op on non-Image leaves —
    /// callers know the body shape via `body()`.
    pub fn set_image_token_count(&mut self, n: u32) {
        if let NodeBody::Image { token_count, .. } = &mut self.body {
            *token_count = n;
            self.token_ids = self.body.compute_token_ids();
        }
    }
}

impl AstNode {
    // -- Leaf constructors ----------------------------------------------------

    pub fn content(text: impl Into<String>) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::Content(text.into())))
    }

    pub fn thinking(text: impl Into<String>) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::Thinking(text.into())))
    }

    pub fn tool_call(name: impl Into<String>, arguments: impl Into<String>) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::ToolCall {
            name: name.into(),
            arguments: arguments.into(),
        }))
    }

    pub fn tool_result(text: impl Into<String>) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::ToolResult(text.into())))
    }

    pub fn memory(key: impl Into<String>, text: impl Into<String>) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::Memory {
            key: key.into(),
            text: text.into(),
            score: None,
        }))
    }

    pub fn dmn(text: impl Into<String>) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::Dmn(text.into())))
    }

    pub fn log(text: impl Into<String>) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::Log(text.into())))
    }

    /// Build an Image leaf. `token_count` is computed from the image
    /// dimensions using Qwen3-VL's resizing rules.
    pub fn image(
        bytes: Vec<u8>,
        mime: impl Into<String>,
        orig_height: u32,
        orig_width: u32,
        token_count: u32,
    ) -> Self {
        Self::Leaf(NodeLeaf::new(NodeBody::Image {
            bytes,
            mime: mime.into(),
            orig_height,
            orig_width,
            token_count,
        }))
    }

    // -- Branch constructors --------------------------------------------------

    pub fn branch(role: Role, children: Vec<AstNode>) -> Self {
        Self::Branch { role, children, timestamp: Utc::now(), memory_scores: Default::default() }
    }

    pub fn system_msg(text: impl Into<String>) -> Self {
        Self::Branch {
            role: Role::System,
            children: vec![Self::content(text)],
            timestamp: Utc::now(),
            memory_scores: Default::default(),
        }
    }

    pub fn user_msg(text: impl Into<String>) -> Self {
        Self::Branch {
            role: Role::User,
            children: vec![Self::content(text)],
            timestamp: Utc::now(),
            memory_scores: Default::default(),
        }
    }

    // -- Builder --------------------------------------------------------------

    pub fn retokenize(self) -> Self {
        match self {
            Self::Leaf(leaf) => {
                let token_ids = leaf.body.compute_token_ids();
                Self::Leaf(NodeLeaf { token_ids, ..leaf })
            }
            Self::Branch { role, children, timestamp, memory_scores } => Self::Branch {
                role,
                children: children.into_iter().map(|c| c.retokenize()).collect(),
                timestamp,
                memory_scores,
            },
        }
    }

    pub fn with_timestamp(mut self, ts: DateTime<Utc>) -> Self {
        match &mut self {
            Self::Leaf(leaf) => leaf.timestamp = ts,
            Self::Branch { timestamp, .. } => *timestamp = ts,
        }
        self
    }

    pub fn children(&self) -> &[AstNode] {
        match self {
            Self::Branch { children, .. } => children,
            Self::Leaf(_) => &[],
        }
    }

    pub fn leaf(&self) -> Option<&NodeLeaf> {
        match self {
            Self::Leaf(l) => Some(l),
            _ => None,
        }
    }

    /// Short label for the UI.
    pub fn label(&self) -> String {
        let app = crate::config::app();
        match self {
            Self::Branch { role, children, .. } => {
                let preview = children.first()
                    .and_then(|c| c.leaf())
                    .map(|l| truncate_preview(l.body.text(), 60))
                    .unwrap_or_default();
                match role {
                    Role::System => "system".into(),
                    Role::User => format!("{}: {}", app.user_name, preview),
                    Role::Assistant => format!("{}: {}", app.assistant_name, preview),
                }
            }
            Self::Leaf(leaf) => match &leaf.body {
                NodeBody::Content(t) => truncate_preview(t, 60),
                NodeBody::Thinking(t) => format!("thinking: {}", truncate_preview(t, 60)),
                NodeBody::ToolCall { name, arguments } => format!("tool: {}({})", name, truncate_preview(arguments, 80)),
                NodeBody::ToolResult(_) => "tool_result".into(),
                NodeBody::Memory { key, score, .. } => match score {
                    Some(s) => format!("mem: {} score:{:.1}", key, s),
                    None => format!("mem: {}", key),
                },
                NodeBody::Dmn(_) => "dmn".into(),
                NodeBody::Image { orig_height, orig_width, token_count, .. } =>
                    format!("image: {}x{} ({} tokens)", orig_width, orig_height, token_count),
                NodeBody::Log(t) => format!("log: {}", truncate_preview(t, 60)),
            },
        }
    }
}

impl AstNode {
    fn render_into(&self, out: &mut String) {
        match self {
            Self::Leaf(leaf) => leaf.body.render_into(out),
            Self::Branch { role, children, .. } => {
                out.push_str(&format!("<|im_start|>{}\n", role.as_str()));
                for child in children {
                    child.render_into(out);
                }
                out.push_str("<|im_end|>\n");
            }
        }
    }

    fn token_ids_into(&self, out: &mut Vec<u32>) {
        match self {
            Self::Leaf(leaf) => out.extend_from_slice(&leaf.token_ids),
            Self::Branch { role, children, .. } => {
                out.push(tokenizer::IM_START);
                out.extend(tokenizer::encode(&format!("{}\n", role.as_str())));
                for child in children {
                    child.token_ids_into(out);
                }
                out.push(tokenizer::IM_END);
                out.extend(tokenizer::encode("\n"));
            }
        }
    }
}

impl Ast for AstNode {
    fn render(&self) -> String {
        let mut s = String::new();
        self.render_into(&mut s);
        s
    }

    fn token_ids(&self) -> Vec<u32> {
        let mut ids = Vec::new();
        self.token_ids_into(&mut ids);
        ids
    }

    fn tokens(&self) -> usize {
        match self {
            Self::Leaf(leaf) => leaf.tokens(),
            Self::Branch { role, children, .. } => {
                1 + role_header_tokens(*role)
                    + children.iter().map(|c| c.tokens()).sum::<usize>()
                    + 1 + newline_tokens()
            }
        }
    }
}

fn truncate_preview(s: &str, max: usize) -> String {
    let preview: String = s.chars().take(max).collect();
    let preview = preview.replace('\n', " ");
    if s.len() > max { format!("{}...", preview) } else { preview }
}

fn format_tool_call_xml(name: &str, args_json: &str) -> String {
    let args: serde_json::Value = serde_json::from_str(args_json)
        .unwrap_or(serde_json::Value::Object(Default::default()));
    let mut xml = format!("<function={}>\n", name);
    if let Some(obj) = args.as_object() {
        for (key, value) in obj {
            let val_str = match value {
                serde_json::Value::String(s) => s.clone(),
                other => other.to_string(),
            };
            xml.push_str(&format!("<parameter={}>\n{}\n</parameter>\n", key, val_str));
        }
    }
    xml.push_str("</function>");
    xml
}

/// Search for a sequence of literal parts separated by optional ASCII whitespace.
/// Returns (start, end) byte positions of the overall match.
///
/// Handles the case where streaming tokenization inserts whitespace inside
/// XML tag structure, e.g. `< function = bash >` instead of `<function=bash>`.
fn find_ws_seq(s: &str, parts: &[&str]) -> Option<(usize, usize)> {
    let bytes = s.as_bytes();
    let mut search_from = 0;
    'outer: loop {
        let start = s[search_from..].find(parts[0])? + search_from;
        let mut pos = start + parts[0].len();
        for &part in &parts[1..] {
            while pos < bytes.len() && bytes[pos].is_ascii_whitespace() {
                pos += 1;
            }
            if !s[pos..].starts_with(part) {
                search_from = start + 1;
                continue 'outer;
            }
            pos += part.len();
        }
        return Some((start, pos));
    }
}

/// Parse a Qwen-style XML tag: `<tag=name>body</tag>`.
/// Tolerates whitespace inside tag delimiters (streaming artifact).
/// Body content is returned verbatim except for a single leading/trailing
/// newline (XML formatting convention).
fn parse_qwen_tag<'a>(s: &'a str, tag: &str) -> Option<(&'a str, &'a str, &'a str)> {
    // Open tag: tolerate whitespace from streaming tokenization
    let (_, after_eq) = find_ws_seq(s, &["<", tag, "="])?;
    let gt_offset = s[after_eq..].find('>')?;
    let name = s[after_eq..after_eq + gt_offset].trim();
    let body_start = after_eq + gt_offset + 1;

    // Close tag: exact match — model doesn't insert whitespace in close tags
    let close = format!("</{}>", tag);
    let close_offset = s[body_start..].find(&close)?;
    let body = &s[body_start..body_start + close_offset];
    // Strip the single leading/trailing newline from XML formatting,
    // but preserve all other whitespace (indentation matters for code).
    let body = body.strip_prefix('\n').unwrap_or(body);
    let body = body.strip_suffix('\n').unwrap_or(body);
    let rest = &s[body_start + close_offset + close.len()..];

    Some((name, body, rest))
}

fn parse_tool_call_body(body: &str) -> Option<(String, String)> {
    let body = body.trim();
    parse_xml_tool_call(body)
        .or_else(|| parse_json_tool_call(body))
}

fn parse_xml_tool_call(body: &str) -> Option<(String, String)> {
    let (func_name, func_body, _) = parse_qwen_tag(body, "function")?;
    let mut args = serde_json::Map::new();
    let mut rest = func_body;
    while let Some((key, val, remainder)) = parse_qwen_tag(rest, "parameter") {
        let value = serde_json::from_str(val)
            .unwrap_or(serde_json::Value::String(val.to_string()));
        args.insert(key.to_string(), value);
        rest = remainder;
    }
    Some((func_name.to_string(), serde_json::to_string(&args).unwrap_or_default()))
}

fn parse_json_tool_call(body: &str) -> Option<(String, String)> {
    let v: serde_json::Value = serde_json::from_str(body).ok()?;
    let name = v["name"].as_str()?;
    let arguments = &v["arguments"];
    Some((name.to_string(), serde_json::to_string(arguments).unwrap_or_default()))
}

/// Search `buf` for `close_tag`. If found, append everything before it to
/// `accum`, advance `buf` past the tag, and return the accumulated content.
/// If not found, drain the safe prefix (preserving any partial tag match at
/// the end of buf) into `accum`.
fn scan_close_tag(buf: &mut String, close_tag: &str, accum: &mut String) -> Option<String> {
    if let Some(pos) = buf.find(close_tag) {
        accum.push_str(&buf[..pos]);
        *buf = buf[pos + close_tag.len()..].to_string();
        Some(std::mem::take(accum))
    } else {
        let drained = drain_safe(buf, close_tag.len());
        if !drained.is_empty() {
            accum.push_str(&drained);
        }
        None
    }
}

/// Remove everything from `buf` except the last `tag_len` bytes, which might
/// be a partial tag. Returns the removed prefix.
fn drain_safe(buf: &mut String, tag_len: usize) -> String {
    let safe = buf.len().saturating_sub(tag_len);
    if safe > 0 {
        let safe = buf.floor_char_boundary(safe);
        let drained = buf[..safe].to_string();
        *buf = buf[safe..].to_string();
        drained
    } else {
        String::new()
    }
}

impl ResponseParser {
    /// @in_think: whether the model's output begins inside a <think> block.
    /// Set when the prompt was prefilled with "<think>\n" (native thinking
    /// mode) so the parser captures reasoning tokens as Thinking until the
    /// model emits </think>.
    pub fn new(branch_idx: usize, in_think: bool) -> Self {
        Self {
            branch_idx,
            call_counter: 0,
            buf: String::new(),
            content_parts: Vec::new(),
            in_think,
            think_buf: String::new(),
            in_tool_call: false,
            tool_call_buf: String::new(),
        }
    }

    /// Consume a token stream, parse into the AST, yield tool calls.
    /// Spawns a background task. Returns a tool call receiver and a
    /// join handle that resolves to Ok(()) or the stream error.
    pub fn run(
        self,
        mut stream: tokio::sync::mpsc::UnboundedReceiver<super::api::StreamToken>,
        agent: std::sync::Arc<super::Agent>,
    ) -> (
        tokio::sync::mpsc::UnboundedReceiver<PendingToolCall>,
        tokio::task::JoinHandle<anyhow::Result<()>>,
    ) {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let handle = tokio::spawn(async move {
            let mut parser = self;
            let agent_name = agent.state.lock().await.provenance.clone();
            let log_path = format!("/tmp/poc-{}.log", agent_name);
            let mut log_file = std::fs::OpenOptions::new()
                .create(true).append(true).open(&log_path).ok();
            let mut full_text = String::new();
            while let Some(event) = stream.recv().await {
                match event {
                    super::api::StreamToken::Token { id, readout } => {
                        if let Some(r) = readout {
                            if let Ok(mut buf) = agent.readout.lock() {
                                buf.push(id, r);
                            }
                        }
                        let text = super::tokenizer::decode(&[id]);
                        full_text.push_str(&text);
                        let mut ctx = agent.context.lock().await;
                        let calls = parser.feed_token(&text, &mut ctx);
                        if !calls.is_empty() {
                            if let Some(ref mut f) = log_file {
                                use std::io::Write;
                                for c in &calls {
                                    let end = c.arguments.floor_char_boundary(c.arguments.len().min(200));
                                    let _ = writeln!(f, "tool_call: {} args={}", c.name, &c.arguments[..end]);
                                }
                            }
                        }
                        for call in calls {
                            let _ = tx.send(call);
                        }
                    }
                    super::api::StreamToken::Done { usage } => {
                        if let Some(ref mut f) = log_file {
                            use std::io::Write;
                            let ctx = agent.context.lock().await;
                            let children = ctx.conversation().get(parser.branch_idx)
                                .map(|n| n.children()).unwrap_or(&[]);
                            let n_think = children.iter().filter(|c| matches!(c.leaf().map(|l| l.body()), Some(NodeBody::Thinking(_)))).count();
                            let n_content = children.iter().filter(|c| matches!(c.leaf().map(|l| l.body()), Some(NodeBody::Content(_)))).count();
                            let n_tool = children.iter().filter(|c| matches!(c.leaf().map(|l| l.body()), Some(NodeBody::ToolCall { .. }))).count();
                            let _ = writeln!(f, "done: {} chars, {} content + {} think + {} tool_call, ctx: {} tokens",
                                full_text.len(), n_content, n_think, n_tool, ctx.tokens());
                            drop(ctx);
                            if full_text.len() > 0 && n_content == 0 && n_tool == 0 {
                                let end = full_text.floor_char_boundary(full_text.len().min(2000));
                                let _ = writeln!(f, "  unparsed text: {}", &full_text[..end]);
                            }
                        }
                        if let Some(u) = usage {
                            agent.state.lock().await.last_prompt_tokens = u.prompt_tokens;
                        }
                        let mut ctx = agent.context.lock().await;
                        parser.finish(&mut ctx);
                        return Ok(());
                    }
                    super::api::StreamToken::ImageAppended { placeholder_count } => {
                        // Commit the server-authoritative IMAGE_PAD
                        // count into the first zero-count image leaf
                        // in wire order. AppendImage always runs
                        // before the final Generate, so this fires
                        // before any Token events for this stream.
                        let mut ctx = agent.context.lock().await;
                        ctx.commit_image_token_counts(&[placeholder_count]);
                    }
                    super::api::StreamToken::Error(e) => {
                        return Err(anyhow::anyhow!("{}", e));
                    }
                }
            }
            Ok(())
        });
        (rx, handle)
    }

    pub fn feed_token(&mut self, text: &str, ctx: &mut ContextState) -> Vec<PendingToolCall> {
        const THINK_OPEN: &str = "<think>";
        const THINK_CLOSE: &str = "</think>";
        const TOOL_CALL_OPEN: &str = "<tool_call>";
        const TOOL_CALL_CLOSE: &str = "</tool_call>";
        const OPEN_TAGS: &[&str] = &[THINK_OPEN, TOOL_CALL_OPEN];

        let mut pending = Vec::new();
        self.buf.push_str(text);

        loop {
            if self.in_think {
                if let Some(content) = scan_close_tag(&mut self.buf, THINK_CLOSE, &mut self.think_buf) {
                    self.in_think = false;
                    let text = content.trim().to_string();
                    if !text.is_empty() {
                        self.push_child(ctx, AstNode::thinking(text));
                    }
                    continue;
                }
                break;
            }

            if self.in_tool_call {
                if let Some(content) = scan_close_tag(&mut self.buf, TOOL_CALL_CLOSE, &mut self.tool_call_buf) {
                    self.in_tool_call = false;
                    if let Some((name, args)) = parse_tool_call_body(&content) {
                        self.flush_content(ctx);
                        self.push_child(ctx, AstNode::tool_call(&name, &args));
                        self.call_counter += 1;
                        pending.push(PendingToolCall {
                            name,
                            arguments: args,
                            id: format!("call_{}", self.call_counter),
                        });
                    }
                    continue;
                }
                break;
            }

            // Not inside a tag — find the earliest opening tag
            let next = OPEN_TAGS.iter()
                .filter_map(|tag| self.buf.find(tag).map(|pos| (pos, *tag)))
                .min_by_key(|(pos, _)| *pos);

            match next {
                Some((pos, tag)) => {
                    if pos > 0 {
                        self.content_parts.push(self.buf[..pos].to_string());
                    }
                    self.buf = self.buf[pos + tag.len()..].to_string();
                    self.flush_content(ctx);
                    match tag {
                        THINK_OPEN     => self.in_think = true,
                        TOOL_CALL_OPEN => self.in_tool_call = true,
                        _ => unreachable!(),
                    }
                    continue;
                }
                None => {
                    // Keep a tail that might be a partial opening tag
                    let max_tag = OPEN_TAGS.iter().map(|t| t.len()).max().unwrap();
                    let drained = drain_safe(&mut self.buf, max_tag);
                    if !drained.is_empty() {
                        self.content_parts.push(drained);
                    }
                    break;
                }
            }
        }

        pending
    }

    fn push_child(&self, ctx: &mut ContextState, child: AstNode) {
        ctx.push_child(Section::Conversation, self.branch_idx, child);
    }

    fn flush_content(&mut self, ctx: &mut ContextState) {
        if !self.content_parts.is_empty() {
            let text: String = self.content_parts.drain(..).collect();
            let text = text.trim().to_string();
            if !text.is_empty() {
                self.push_child(ctx, AstNode::content(text));
            }
        }
    }

    pub fn finish(mut self, ctx: &mut ContextState) {
        if !self.buf.is_empty() {
            self.content_parts.push(std::mem::take(&mut self.buf));
        }
        self.flush_content(ctx);
    }
}

impl ContextState {
    pub fn new() -> Self {
        Self {
            system: Vec::new(),
            identity: Vec::new(),
            journal: Vec::new(),
            conversation: Vec::new(),
            conversation_log: None,
        }
    }

    // -- Read access ----------------------------------------------------------

    pub fn system(&self) -> &[AstNode]       { &self.system }
    pub fn identity(&self) -> &[AstNode]     { &self.identity }
    pub fn journal(&self) -> &[AstNode]      { &self.journal }
    pub fn conversation(&self) -> &[AstNode] { &self.conversation }
    pub fn conversation_mut(&mut self) -> &mut Vec<AstNode> { &mut self.conversation }

    pub fn sections(&self) -> [&Vec<AstNode>; 4] {
        [&self.system, &self.identity, &self.journal, &self.conversation]
    }

    /// Walk image leaves across all sections in wire order and fill in
    /// the first N leaves that have `token_count == 0` with successive
    /// values from `counts`. Used after a gRPC session's stream of
    /// AppendImage responses to commit the server's IMAGE_PAD counts
    /// back into the AST so the next wire walk doesn't see zero-count
    /// images in the already-committed prefix.
    pub fn commit_image_token_counts(&mut self, counts: &[u32]) {
        fn visit(node: &mut AstNode, counts: &[u32], idx: &mut usize) {
            if *idx >= counts.len() { return; }
            match node {
                AstNode::Leaf(leaf) => {
                    if let NodeBody::Image { token_count, .. } = leaf.body() {
                        if *token_count == 0 {
                            leaf.set_image_token_count(counts[*idx]);
                            *idx += 1;
                        }
                    }
                }
                AstNode::Branch { children, .. } => {
                    for c in children { visit(c, counts, idx); }
                }
            }
        }
        let mut idx = 0usize;
        for node in &mut self.system      { visit(node, counts, &mut idx); }
        for node in &mut self.identity    { visit(node, counts, &mut idx); }
        for node in &mut self.journal     { visit(node, counts, &mut idx); }
        for node in &mut self.conversation { visit(node, counts, &mut idx); }
    }
}

impl Ast for ContextState {
    fn render(&self) -> String {
        let mut s = String::new();
        for section in self.sections() {
            for node in section {
                s.push_str(&node.render());
            }
        }
        s
    }

    fn token_ids(&self) -> Vec<u32> {
        let mut ids = Vec::new();
        for section in self.sections() {
            for node in section {
                ids.extend(node.token_ids());
            }
        }
        ids
    }

    fn tokens(&self) -> usize {
        self.sections().iter()
            .flat_map(|s| s.iter())
            .map(|n| n.tokens())
            .sum()
    }
}

/// An image collected from the AST for a request body. The AST stores
/// the pre-expanded token form (`<|vision_start|> + <|image_pad|>×N +
/// <|vision_end|>`), and the wire form mirrors that exactly so the
/// server's `session.tokens` length matches what vLLM's engine will
/// process. The authoritative N is obtained from the server via the
/// CountImageTokens RPC before the Image leaf is constructed.
#[derive(Clone)]
pub struct WireImage {
    pub bytes: Vec<u8>,
    pub mime: String,
}

/// One piece of the wire stream for the gRPC session path. Runs of
/// text/tool/thinking tokens are batched into `Tokens`; each Image
/// leaf becomes its own `Image` chunk because the server writes the
/// full vision block on AppendImage — the client never sends vision
/// tokens inline. Order matches the AST's depth-first wire order.
#[derive(Clone)]
pub enum WireChunk {
    Tokens(Vec<u32>),
    Image {
        bytes: Vec<u8>,
        mime: String,
        /// Client's current best guess at how many tokens the server
        /// will expand this image to, including bookends. `0` means
        /// the count is unknown (view_image just loaded the image and
        /// AppendImage hasn't run yet). Callers use this only to know
        /// this chunk's contribution to the server-visible length for
        /// offset bookkeeping on chunks that were already appended on
        /// a prior turn.
        known_expanded_len: u32,
    },
}

fn wire_into(node: &AstNode, tokens: &mut Vec<u32>, images: &mut Vec<WireImage>) {
    match node {
        AstNode::Leaf(leaf) => match leaf.body() {
            NodeBody::Image { bytes, mime, .. } => {
                // Send the pre-expanded token form (includes N
                // <|image_pad|> tokens); engine's multi_modal
                // pipeline pairs them with the binary data below.
                tokens.extend_from_slice(leaf.token_ids());
                images.push(WireImage {
                    bytes: bytes.clone(),
                    mime: mime.clone(),
                });
            }
            _ => tokens.extend_from_slice(leaf.token_ids()),
        },
        AstNode::Branch { role, children, .. } => {
            tokens.push(tokenizer::IM_START);
            tokens.extend(tokenizer::encode(&format!("{}\n", role.as_str())));
            for c in children {
                wire_into(c, tokens, images);
            }
            tokens.push(tokenizer::IM_END);
            tokens.extend(tokenizer::encode("\n"));
        }
    }
}

pub fn memory_key(node: &AstNode) -> Option<&str> {
    match node {
        AstNode::Leaf(leaf) => match leaf.body() {
            NodeBody::Memory { key, .. } => Some(key),
            _ => None,
        },
        _ => None,
    }
}

pub fn is_memory_node(node: &AstNode) -> bool {
    matches!(node, AstNode::Leaf(leaf) if matches!(leaf.body(), NodeBody::Memory { .. }))
}

pub fn is_assistant(node: &AstNode) -> bool {
    matches!(node, AstNode::Branch { role: Role::Assistant, .. })
}

/// Concatenate the text of a Branch's Leaf children — what the model
/// actually produced on that turn (Content + Thinking + ToolCall name).
pub fn render_branch_text(children: &[AstNode]) -> String {
    children.iter()
        .filter_map(|c| match c {
            AstNode::Leaf(leaf) => Some(leaf.body().text().to_string()),
            _ => None,
        })
        .collect::<Vec<_>>()
        .join("")
}

/// Render the last `max_msgs` user/assistant branches before `idx` as a
/// review-friendly string with `[user]` / `[assistant]` markers.
pub fn render_prior_context(entries: &[AstNode], idx: usize, max_msgs: usize) -> String {
    let mut picked: Vec<&AstNode> = Vec::with_capacity(max_msgs);
    for i in (0..idx).rev() {
        if picked.len() >= max_msgs { break; }
        if let AstNode::Branch { role, .. } = &entries[i] {
            if matches!(role, Role::User | Role::Assistant) {
                picked.push(&entries[i]);
            }
        }
    }
    picked.reverse();

    let mut out = String::new();
    for node in picked {
        if let AstNode::Branch { role, children, .. } = node {
            let marker = match role {
                Role::User => "[user]",
                Role::Assistant => "[assistant]",
                _ => continue,
            };
            out.push_str(marker);
            out.push('\n');
            out.push_str(render_branch_text(children).trim());
            out.push_str("\n\n");
        }
    }
    out.trim_end().to_string()
}

impl ContextState {
    /// Assemble the prompt in wire form: token stream with a single
    /// `<|image_pad|>` per image (vLLM expands back to N), plus the list
    /// of images to send as multi_modal_data, plus the (start, end) token
    /// positions of each assistant message branch emitted (used by the
    /// scoring path as `score_ranges`).
    ///
    /// `conv_range` selects a prefix (or any sub-range) of conversation
    /// entries to include — the agent path passes `0..conversation().len()`;
    /// scoring / candidate generation pass a prefix up to the entry of
    /// interest.
    ///
    /// `skip` is a predicate applied to identity and conversation entries;
    /// returning true drops the node from the prompt. The agent path passes
    /// `|_| false`; memory-ablation scoring passes e.g. `is_memory_node` or
    /// `|n| memory_key(n) == Some(key)`.
    pub fn wire_prompt<F>(
        &self,
        conv_range: std::ops::Range<usize>,
        mut skip: F,
    ) -> (Vec<u32>, Vec<WireImage>, Vec<(usize, usize)>)
    where F: FnMut(&AstNode) -> bool,
    {
        let mut tokens = Vec::new();
        let mut images = Vec::new();
        let mut assistant_ranges = Vec::new();

        for node in self.system() {
            wire_into(node, &mut tokens, &mut images);
        }
        for node in self.identity() {
            if skip(node) { continue; }
            wire_into(node, &mut tokens, &mut images);
        }
        for node in self.journal() {
            wire_into(node, &mut tokens, &mut images);
        }
        for node in &self.conversation()[conv_range] {
            if skip(node) { continue; }
            let start = tokens.len();
            let is_asst = matches!(node, AstNode::Branch { role: Role::Assistant, .. });
            wire_into(node, &mut tokens, &mut images);
            if is_asst {
                assistant_ranges.push((start, tokens.len()));
            }
        }
        (tokens, images, assistant_ranges)
    }

    /// Build the wire stream as interleaved `WireChunk`s for the gRPC
    /// session path. Unlike `wire_prompt`, this preserves the order
    /// of text runs vs image blocks so the caller can drive the
    /// append flow (AppendImage for each Image, Generate append for
    /// contiguous text runs).
    ///
    /// `conv_range` and `skip` mirror `wire_prompt` — select a
    /// conversation slice and drop identity / conversation nodes by
    /// predicate.
    pub fn wire_chunks<F>(
        &self,
        conv_range: std::ops::Range<usize>,
        mut skip: F,
    ) -> Vec<WireChunk>
    where F: FnMut(&AstNode) -> bool,
    {
        let mut out: Vec<WireChunk> = Vec::new();
        let mut buf: Vec<u32> = Vec::new();

        fn flush(buf: &mut Vec<u32>, out: &mut Vec<WireChunk>) {
            if !buf.is_empty() {
                out.push(WireChunk::Tokens(std::mem::take(buf)));
            }
        }

        fn visit(node: &AstNode, buf: &mut Vec<u32>, out: &mut Vec<WireChunk>) {
            match node {
                AstNode::Leaf(leaf) => match leaf.body() {
                    NodeBody::Image { bytes, mime, token_count, .. } => {
                        flush(buf, out);
                        // Bookends (VISION_START + VISION_END) add 2
                        // to the expanded length; token_count is the
                        // IMAGE_PAD run. 0 means count is still
                        // unknown (no AppendImage yet) — don't claim
                        // a length the server will disagree with.
                        let expanded = if *token_count == 0 {
                            0
                        } else {
                            *token_count + 2
                        };
                        out.push(WireChunk::Image {
                            bytes: bytes.clone(),
                            mime: mime.clone(),
                            known_expanded_len: expanded,
                        });
                    }
                    _ => buf.extend_from_slice(leaf.token_ids()),
                },
                AstNode::Branch { role, children, .. } => {
                    buf.push(tokenizer::IM_START);
                    buf.extend(tokenizer::encode(&format!("{}\n", role.as_str())));
                    for c in children {
                        visit(c, buf, out);
                    }
                    buf.push(tokenizer::IM_END);
                    buf.extend(tokenizer::encode("\n"));
                }
            }
        }

        for node in self.system()   { visit(node, &mut buf, &mut out); }
        for node in self.identity() {
            if skip(node) { continue; }
            visit(node, &mut buf, &mut out);
        }
        for node in self.journal()  { visit(node, &mut buf, &mut out); }
        for node in &self.conversation()[conv_range] {
            if skip(node) { continue; }
            visit(node, &mut buf, &mut out);
        }
        flush(&mut buf, &mut out);
        out
    }
}

impl ContextState {
    fn section_mut(&mut self, section: Section) -> &mut Vec<AstNode> {
        match section {
            Section::System       => &mut self.system,
            Section::Identity     => &mut self.identity,
            Section::Journal      => &mut self.journal,
            Section::Conversation => &mut self.conversation,
        }
    }

    /// Push and log to conversation log.
    pub fn push_log(&mut self, section: Section, node: AstNode) {
        if let Some(ref log) = self.conversation_log {
            if let Err(e) = log.append_node(&node) {
                dbglog!("warning: log: {:#}", e);
            }
        }
        self.section_mut(section).push(node);
    }

    /// Push without logging.
    pub fn push_no_log(&mut self, section: Section, node: AstNode) {
        self.section_mut(section).push(node);
    }

    /// Replace the body of a leaf at `index` in `section`.
    /// Re-tokenizes to maintain the invariant.
    pub fn set_message(&mut self, section: Section, index: usize, body: NodeBody) {
        let nodes = self.section_mut(section);
        let node = &mut nodes[index];
        match node {
            AstNode::Leaf(leaf) => {
                let token_ids = body.compute_token_ids();
                leaf.body = body;
                leaf.token_ids = token_ids;
            }
            AstNode::Branch { .. } => panic!("set_message on branch node"),
        }
    }

    /// Set the memory score on a Memory leaf at `index` in `section`.
    pub fn set_score(&mut self, section: Section, index: usize, score: Option<f64>) {
        let node = &mut self.section_mut(section)[index];
        match node {
            AstNode::Leaf(leaf) => match &mut leaf.body {
                NodeBody::Memory { score: s, .. } => *s = score,
                _ => panic!("set_score on non-memory node"),
            },
            _ => panic!("set_score on branch node"),
        }
    }

    pub fn del(&mut self, section: Section, index: usize) -> AstNode {
        self.section_mut(section).remove(index)
    }

    pub fn clear(&mut self, section: Section) {
        self.section_mut(section).clear();
    }

    /// Total tokens across every section that gets serialized into the prompt.
    /// Cheap sum over cached `node.tokens()`; call this before assembling to
    /// decide whether to trim.
    pub fn total_tokens(&self) -> usize {
        self.system().iter().map(|n| n.tokens()).sum::<usize>()
            + self.identity().iter().map(|n| n.tokens()).sum::<usize>()
            + self.journal().iter().map(|n| n.tokens()).sum::<usize>()
            + self.conversation().iter().map(|n| n.tokens()).sum::<usize>()
    }

    /// Dedup and trim conversation entries to fit within the context budget.
    ///
    /// Phase 1: Drop duplicate memories (keep last) and DMN entries.
    /// Phase 2: While over budget, drop lowest-scored memory (if memories
    ///          are > 50% of conversation tokens) or oldest conversation entry.
    /// Phase 3: Snap to user message boundary at start.
    pub fn trim_conversation(&mut self) {
        let max_tokens = context_budget_tokens();
        let fixed = self.system.iter().map(|n| n.tokens()).sum::<usize>()
            + self.identity.iter().map(|n| n.tokens()).sum::<usize>()
            + self.journal.iter().map(|n| n.tokens()).sum::<usize>();

        // Phase 1: dedup memories by key (keep last), drop DMN
        let mut seen_keys: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
        let mut drop = std::collections::HashSet::new();

        for (i, node) in self.conversation.iter().enumerate() {
            if let AstNode::Leaf(leaf) = node {
                match leaf.body() {
                    NodeBody::Dmn(_) => { drop.insert(i); }
                    NodeBody::Memory { key, .. } => {
                        if let Some(prev) = seen_keys.insert(key.clone(), i) {
                            drop.insert(prev);
                        }
                    }
                    _ => {}
                }
            }
        }

        if !drop.is_empty() {
            let mut i = 0;
            self.conversation.retain(|_| { let keep = !drop.contains(&i); i += 1; keep });
        }

        // Phase 2: while over budget, evict
        loop {
            let total: usize = self.conversation.iter().map(|n| n.tokens()).sum();
            if fixed + total <= max_tokens { break; }
            let mt: usize = self.conversation.iter()
                .filter(|n| matches!(n, AstNode::Leaf(l) if matches!(l.body(), NodeBody::Memory { .. })))
                .map(|n| n.tokens()).sum();
            let ct = total - mt;

            if mt > ct {
                // Memories > 50% — drop lowest-scored
                if let Some(i) = self.lowest_scored_memory() {
                    self.conversation.remove(i);
                    continue;
                }
            }
            // Drop oldest non-memory entry
            if let Some(i) = self.conversation.iter().position(|n|
                !matches!(n, AstNode::Leaf(l) if matches!(l.body(), NodeBody::Memory { .. })))
            {
                self.conversation.remove(i);
            } else {
                break;
            }
        }

        // Phase 3: snap to user message boundary
        while let Some(first) = self.conversation.first() {
            if matches!(first, AstNode::Branch { role: Role::User, .. }) { break; }
            self.conversation.remove(0);
        }
    }

    fn lowest_scored_memory(&self) -> Option<usize> {
        self.conversation.iter().enumerate()
            .filter_map(|(i, n)| {
                if let AstNode::Leaf(l) = n {
                    if let NodeBody::Memory { score: Some(s), .. } = l.body() {
                        return Some((i, *s));
                    }
                }
                None
            })
            .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
    }

    /// Push a child node into a branch at `index` in `section`.
    pub fn push_child(&mut self, section: Section, index: usize, child: AstNode) {
        let node = &mut self.section_mut(section)[index];
        match node {
            AstNode::Branch { children, .. } => children.push(child),
            AstNode::Leaf(_) => panic!("push_child on leaf node"),
        }
    }

    /// Number of nodes in a section.
    pub fn len(&self, section: Section) -> usize {
        match section {
            Section::System       => self.system.len(),
            Section::Identity     => self.identity.len(),
            Section::Journal      => self.journal.len(),
            Section::Conversation => self.conversation.len(),
        }
    }
}

// ---------------------------------------------------------------------------
// Qwen3-VL image token count
//
// Port of Qwen2VLImageProcessor.smart_resize + image_token_count. We need the
// exact same answer that vLLM's Qwen3VL processor will produce, because the
// token stream in our context must match what vLLM expands `<|image_pad|>`
// to at request time. Constants come from Qwen3.5-27B's preprocessor_config.
// ---------------------------------------------------------------------------

// Test-only client-side estimate of image token expansion. Production
// callers obtain the authoritative count from the server via
// CountImageTokens; these constants and helpers stay around only to
// keep the context-shape unit tests self-contained.
#[cfg(test)]
const QWEN3_PATCH_SIZE: u32 = 16;
#[cfg(test)]
const QWEN3_MERGE_SIZE: u32 = 2;
#[cfg(test)]
const QWEN3_MIN_PIXELS: u64 = 65_536;
#[cfg(test)]
const QWEN3_MAX_PIXELS: u64 = 16_777_216;

#[cfg(test)]
fn smart_resize(h: u32, w: u32, factor: u32, min_pixels: u64, max_pixels: u64) -> (u32, u32) {
    let max_s = h.max(w) as f64;
    let min_s = h.min(w) as f64;
    assert!(max_s / min_s <= 200.0, "aspect ratio too extreme: {}x{}", h, w);

    let fh = h as f64;
    let fw = w as f64;
    let ff = factor as f64;

    let h_bar = ((fh / ff).round() as u32) * factor;
    let w_bar = ((fw / ff).round() as u32) * factor;
    let total = (h_bar as u64) * (w_bar as u64);

    if total > max_pixels {
        let beta = ((fh * fw) / max_pixels as f64).sqrt();
        let hf = ((fh / beta / ff).floor() as u32) * factor;
        let wf = ((fw / beta / ff).floor() as u32) * factor;
        (hf.max(factor), wf.max(factor))
    } else if total < min_pixels {
        let beta = (min_pixels as f64 / (fh * fw)).sqrt();
        let hc = ((fh * beta / ff).ceil() as u32) * factor;
        let wc = ((fw * beta / ff).ceil() as u32) * factor;
        (hc, wc)
    } else {
        (h_bar, w_bar)
    }
}

/// Test-only: client-side estimate of how many `<|image_pad|>` tokens
/// vLLM will emit for an image of the given dimensions. Production
/// callers use `salience::count_image_tokens` (server-authoritative).
#[cfg(test)]
fn qwen3_image_token_count(orig_h: u32, orig_w: u32) -> u32 {
    let factor = QWEN3_PATCH_SIZE * QWEN3_MERGE_SIZE;
    let (rh, rw) = smart_resize(orig_h, orig_w, factor, QWEN3_MIN_PIXELS, QWEN3_MAX_PIXELS);
    (rh / QWEN3_PATCH_SIZE) * (rw / QWEN3_PATCH_SIZE) / (QWEN3_MERGE_SIZE * QWEN3_MERGE_SIZE)
}

pub fn context_window() -> usize {
    let app = crate::config::app();
    app.backends.get(&app.default_backend)
        .and_then(|b| b.context_window)
        .unwrap_or(128_000)
}

pub fn context_budget_tokens() -> usize {
    context_window() * 80 / 100
}

pub fn is_context_overflow(err: &anyhow::Error) -> bool {
    let msg = err.to_string().to_lowercase();
    msg.contains("context length")
        || msg.contains("token limit")
        || msg.contains("too many tokens")
        || msg.contains("maximum context")
        || msg.contains("prompt is too long")
        || msg.contains("request too large")
        || msg.contains("input validation error")
        || msg.contains("content length limit")
        || (msg.contains("400") && msg.contains("tokens"))
}

pub fn is_stream_error(err: &anyhow::Error) -> bool {
    err.to_string().contains("model stream error")
}

#[cfg(test)]
mod tests {
    use super::*;

    // -- Helpers for inspecting parse results ----------------------------------

    fn bodies(nodes: &[AstNode]) -> Vec<&NodeBody> {
        nodes.iter().filter_map(|c| c.leaf()).map(|l| l.body()).collect()
    }

    fn assert_content(body: &NodeBody, expected: &str) {
        match body {
            NodeBody::Content(t) => assert_eq!(t, expected),
            other => panic!("expected Content, got {:?}", other),
        }
    }

    fn assert_thinking(body: &NodeBody, expected: &str) {
        match body {
            NodeBody::Thinking(t) => assert_eq!(t, expected),
            other => panic!("expected Thinking, got {:?}", other),
        }
    }

    fn assert_tool_call<'a>(body: &'a NodeBody, expected_name: &str) -> &'a str {
        match body {
            NodeBody::ToolCall { name, arguments } => {
                assert_eq!(name, expected_name);
                arguments
            }
            other => panic!("expected ToolCall, got {:?}", other),
        }
    }

    // -- XML parsing tests ----------------------------------------------------

    #[test]
    fn test_tool_call_xml_parse_clean() {
        let body = "<function=bash>\n<parameter=command>poc-memory used core-personality</parameter>\n</function>";
        let (name, args) = parse_tool_call_body(body).unwrap();
        assert_eq!(name, "bash");
        let args: serde_json::Value = serde_json::from_str(&args).unwrap();
        assert_eq!(args["command"], "poc-memory used core-personality");
    }

    #[test]
    fn test_tool_call_xml_parse_streamed_whitespace() {
        // Streaming tokenization can insert whitespace in opening tags,
        // but close tags are always emitted verbatim.
        let body = "<\nfunction\n=\nbash\n>\n<\nparameter\n=\ncommand\n>pwd</parameter>\n</function>";
        let (name, args) = parse_tool_call_body(body).unwrap();
        assert_eq!(name, "bash");
        let args: serde_json::Value = serde_json::from_str(&args).unwrap();
        assert_eq!(args["command"], "pwd");
    }

    #[test]
    fn test_tool_call_json_parse() {
        let body = r#"{"name": "bash", "arguments": {"command": "ls"}}"#;
        let (name, args) = parse_tool_call_body(body).unwrap();
        assert_eq!(name, "bash");
        let args: serde_json::Value = serde_json::from_str(&args).unwrap();
        assert_eq!(args["command"], "ls");
    }

    #[test]
    fn test_tool_call_preserves_code_with_angle_brackets() {
        let body = "<function=edit>\n<parameter=code>if x < y {\n    std::mem::swap(&mut a, &mut b);\n}</parameter>\n</function>";
        let (name, args) = parse_tool_call_body(body).unwrap();
        assert_eq!(name, "edit");
        let args: serde_json::Value = serde_json::from_str(&args).unwrap();
        assert_eq!(args["code"], "if x < y {\n    std::mem::swap(&mut a, &mut b);\n}");
    }

    // -- ResponseParser tests -------------------------------------------------

    /// Set up a ContextState with an assistant branch, run the parser,
    /// return the children that were pushed into the branch.
    fn parse_into_ctx(chunks: &[&str]) -> (ContextState, Vec<PendingToolCall>) {
        let mut ctx = ContextState::new();
        ctx.push_no_log(Section::Conversation, AstNode::branch(Role::Assistant, vec![]));
        let mut p = ResponseParser::new(0, false);
        let mut calls = Vec::new();
        for chunk in chunks {
            // Feed each chunk as a single token (id=0 for tests)
            calls.extend(p.feed_token(chunk, &mut ctx));
        }
        p.finish(&mut ctx);
        (ctx, calls)
    }

    fn assistant_children(ctx: &ContextState) -> &[AstNode] {
        ctx.conversation()[0].children()
    }

    #[test]
    fn test_parser_plain_text() {
        let (ctx, _) = parse_into_ctx(&["hello world"]);
        let b = bodies(assistant_children(&ctx));
        assert_eq!(b.len(), 1);
        assert_content(b[0], "hello world");
    }

    #[test]
    fn test_parser_thinking_then_content() {
        let (ctx, _) = parse_into_ctx(&["<think>reasoning</think>answer"]);
        let b = bodies(assistant_children(&ctx));
        assert_eq!(b.len(), 2);
        assert_thinking(b[0], "reasoning");
        assert_content(b[1], "answer");
    }

    #[test]
    fn test_parser_tool_call() {
        let (ctx, calls) = parse_into_ctx(&[
            "<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>"
        ]);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].name, "bash");
        let b = bodies(assistant_children(&ctx));
        assert_eq!(b.len(), 1);
        let args = assert_tool_call(b[0], "bash");
        let args: serde_json::Value = serde_json::from_str(args).unwrap();
        assert_eq!(args["command"], "ls");
    }

    #[test]
    fn test_parser_content_then_tool_call_then_content() {
        let (ctx, _) = parse_into_ctx(&[
            "before",
            "<tool_call>\n<function=bash>\n<parameter=command>pwd</parameter>\n</function>\n</tool_call>",
            "after",
        ]);
        let b = bodies(assistant_children(&ctx));
        assert_eq!(b.len(), 3);
        assert_content(b[0], "before");
        assert_tool_call(b[1], "bash");
        assert_content(b[2], "after");
    }

    #[test]
    fn test_parser_incremental_feed() {
        let text = "<think>thought</think>response";
        let mut ctx = ContextState::new();
        ctx.push_no_log(Section::Conversation, AstNode::branch(Role::Assistant, vec![]));
        let mut p = ResponseParser::new(0, false);
        for ch in text.chars() {
            p.feed_token(&ch.to_string(), &mut ctx);
        }
        p.finish(&mut ctx);
        let b = bodies(assistant_children(&ctx));
        assert_eq!(b.len(), 2);
        assert_thinking(b[0], "thought");
        assert_content(b[1], "response");
    }

    #[test]
    fn test_parser_incremental_tool_call() {
        let text = "text<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>more";
        let mut ctx = ContextState::new();
        ctx.push_no_log(Section::Conversation, AstNode::branch(Role::Assistant, vec![]));
        let mut p = ResponseParser::new(0, false);
        let mut tool_calls = 0;
        for ch in text.chars() {
            tool_calls += p.feed_token(&ch.to_string(), &mut ctx).len();
        }
        p.finish(&mut ctx);
        assert_eq!(tool_calls, 1);
        let b = bodies(assistant_children(&ctx));
        assert_eq!(b.len(), 3);
        assert_content(b[0], "text");
        assert_tool_call(b[1], "bash");
        assert_content(b[2], "more");
    }

    #[test]
    fn test_parser_thinking_tool_call_content() {
        let (ctx, _) = parse_into_ctx(&[
            "<think>let me think</think>",
            "<tool_call>\n<function=read>\n<parameter=path>/etc/hosts</parameter>\n</function>\n</tool_call>",
            "here's what I found",
        ]);
        let b = bodies(assistant_children(&ctx));
        assert_eq!(b.len(), 3);
        assert_thinking(b[0], "let me think");
        assert_tool_call(b[1], "read");
        assert_content(b[2], "here's what I found");
    }

    // -- Round-trip rendering tests -------------------------------------------

    #[test]
    fn test_render_system_msg() {
        let node = AstNode::system_msg("you are helpful");
        assert_eq!(node.render(), "<|im_start|>system\nyou are helpful<|im_end|>\n");
    }

    #[test]
    fn test_render_user_msg() {
        let node = AstNode::user_msg("hello");
        assert_eq!(node.render(), "<|im_start|>user\nhello<|im_end|>\n");
    }

    #[test]
    fn test_render_assistant_with_thinking_and_content() {
        let node = AstNode::branch(Role::Assistant, vec![
            AstNode::thinking("hmm"),
            AstNode::content("answer"),
        ]);
        // Thinking renders wrapped in <think>...</think> so the model sees
        // previous turns' reasoning (Qwen 3.6 style: CoT stays in the
        // conversation across turns).
        assert_eq!(node.render(), "<|im_start|>assistant\n<think>\nhmm\n</think>\nanswer<|im_end|>\n");
    }

    #[test]
    fn test_render_tool_result() {
        let node = AstNode::tool_result("output here");
        assert_eq!(node.render(), "<|im_start|>user\n<tool_response>\noutput here\n</tool_response><|im_end|>\n");
    }

    #[test]
    fn test_render_memory() {
        let node = AstNode::memory("identity", "I am Proof of Concept");
        assert_eq!(node.render(), "<|im_start|>memory\nI am Proof of Concept<|im_end|>\n");
    }

    #[test]
    fn test_render_dmn() {
        let node = AstNode::dmn("subconscious prompt");
        assert_eq!(node.render(), "<|im_start|>dmn\nsubconscious prompt<|im_end|>\n");
    }

    #[test]
    fn test_render_tool_call() {
        let node = AstNode::tool_call("bash", r#"{"command":"ls"}"#);
        let rendered = node.render();
        assert!(rendered.contains("<tool_call>"));
        assert!(rendered.contains("<function=bash>"));
        assert!(rendered.contains("<parameter=command>"));
        assert!(rendered.contains("ls"));
        assert!(rendered.contains("</tool_call>"));
    }

    // -- Tokenizer round-trip tests -------------------------------------------
    // These require the tokenizer file; skipped if not present.

    fn init_tokenizer() -> bool {
        let path = format!("{}/.consciousness/tokenizer-qwen35.json",
            std::env::var("HOME").unwrap_or_default());
        if std::path::Path::new(&path).exists() {
            tokenizer::init(&path);
            true
        } else {
            false
        }
    }

    fn assert_token_invariants(node: &AstNode) {
        assert_eq!(node.tokens(), node.token_ids().len(),
            "tokens() != token_ids().len()");
    }

    #[test]
    fn test_tokenize_roundtrip_leaf_types() {
        if !init_tokenizer() { return; }

        assert_token_invariants(&AstNode::system_msg("you are a helpful assistant"));
        assert_token_invariants(&AstNode::user_msg("what is 2+2?"));
        assert_token_invariants(&AstNode::tool_result("4"));
        assert_token_invariants(&AstNode::memory("identity", "I am Proof of Concept"));
        assert_token_invariants(&AstNode::dmn("check the memory store"));
        assert_token_invariants(&AstNode::tool_call("bash", r#"{"command":"ls -la"}"#));
    }

    #[test]
    fn test_tokenize_roundtrip_assistant_branch() {
        if !init_tokenizer() { return; }

        let node = AstNode::branch(Role::Assistant, vec![
            AstNode::content("here's what I found:\n"),
            AstNode::tool_call("bash", r#"{"command":"pwd"}"#),
            AstNode::content("\nthat's the current directory"),
        ]);
        assert_token_invariants(&node);
    }

    #[test]
    fn test_tokenize_invisible_nodes_are_zero() {
        if !init_tokenizer() { return; }

        assert_eq!(AstNode::log("debug info").tokens(), 0);
    }

    #[test]
    fn test_tokenize_thinking_matches_rendered_tags() {
        if !init_tokenizer() { return; }

        // Thinking is now prompt-visible (wrapped in <think>...</think>);
        // token count must match the rendered wrapping.
        let node = AstNode::thinking("deep thoughts");
        assert_eq!(node.tokens(), tokenizer::encode(&node.render()).len());
    }

    #[test]
    fn test_tokenize_decode_roundtrip() {
        if !init_tokenizer() { return; }

        // Content without special tokens round-trips through decode
        let text = "hello world, this is a test";
        let ids = tokenizer::encode(text);
        let decoded = tokenizer::decode(&ids);
        assert_eq!(decoded, text);
    }

    #[test]
    fn test_tokenize_context_state_matches_concatenation() {
        if !init_tokenizer() { return; }

        let mut ctx = ContextState::new();
        ctx.push_no_log(Section::System, AstNode::system_msg("you are helpful"));
        ctx.push_no_log(Section::Identity, AstNode::memory("name", "Proof of Concept"));
        ctx.push_no_log(Section::Conversation, AstNode::user_msg("hi"));

        assert_eq!(ctx.tokens(), ctx.token_ids().len());
    }

    #[test]
    fn test_parser_roundtrip_through_tokenizer() {
        if !init_tokenizer() { return; }

        let (ctx, _) = parse_into_ctx(&[
            "I'll check that for you",
            "<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>",
        ]);
        let node = &ctx.conversation()[0];
        assert_token_invariants(node);
        assert!(node.tokens() > 0);
    }

    // -- Timestamp deserialization tests ------------------------------------------

    #[test]
    fn test_timestamp_null_rejected() {
        // Missing/null timestamps used to be accepted via a lenient
        // deserialize fallback. Post-migration the schema is strict.
        let json = r#"{"Leaf":{"body":{"Content":"hello"},"timestamp":null}}"#;
        assert!(serde_json::from_str::<AstNode>(json).is_err());
    }

    #[test]
    fn test_timestamp_missing_rejected() {
        let json = r#"{"Leaf":{"body":{"Content":"hello"}}}"#;
        assert!(serde_json::from_str::<AstNode>(json).is_err());
    }

    #[test]
    fn test_branch_timestamp_missing_rejected() {
        let json = r#"{"Branch":{"role":"User","children":[]}}"#;
        assert!(serde_json::from_str::<AstNode>(json).is_err());
    }

    // -- Image leaf tests ---------------------------------------------------------

    #[test]
    fn test_smart_resize_within_bounds() {
        // Typical case: 1024x768 → rounded to multiples of 32, under max.
        let (h, w) = smart_resize(768, 1024, 32, 65_536, 16_777_216);
        assert_eq!(h, 768);
        assert_eq!(w, 1024);
    }

    #[test]
    fn test_smart_resize_upscales_tiny() {
        // 32x32 = 1024 pixels, below min_pixels=65536. Should scale up.
        let (h, w) = smart_resize(32, 32, 32, 65_536, 16_777_216);
        assert!((h as u64) * (w as u64) >= 65_536,
            "resized {}x{} is under min_pixels", h, w);
        assert_eq!(h % 32, 0);
        assert_eq!(w % 32, 0);
    }

    #[test]
    fn test_smart_resize_downscales_huge() {
        // 8000x6000 = 48M pixels, above max_pixels=16M. Should scale down.
        let (h, w) = smart_resize(8000, 6000, 32, 65_536, 16_777_216);
        assert!((h as u64) * (w as u64) <= 16_777_216,
            "resized {}x{} exceeds max_pixels", h, w);
        assert_eq!(h % 32, 0);
        assert_eq!(w % 32, 0);
    }

    #[test]
    fn test_qwen3_token_count_matches_formula() {
        // 512x512 → resized to 512x512 (already multiple of 32, within bounds).
        // grid = 32x32, tokens = 32*32/4 = 256.
        assert_eq!(qwen3_image_token_count(512, 512), 256);
    }

    #[test]
    fn test_image_render_and_token_ids() {
        let node = AstNode::image(vec![0u8, 1, 2, 3], "image/png", 512, 512, qwen3_image_token_count(512, 512));
        let leaf = node.leaf().unwrap();
        // 3 tokens of bookend + 256 image_pad tokens
        assert_eq!(leaf.token_ids().len(), 258);
        assert_eq!(leaf.token_ids()[0], tokenizer::VISION_START);
        assert_eq!(leaf.token_ids()[257], tokenizer::VISION_END);
        for pad in &leaf.token_ids()[1..257] {
            assert_eq!(*pad, tokenizer::IMAGE_PAD);
        }
        // Rendered text has the expected bookends.
        let rendered = leaf.body().render();
        assert!(rendered.starts_with("<|vision_start|>"));
        assert!(rendered.ends_with("<|vision_end|>"));
    }

    #[test]
    fn test_wire_prompt_preserves_expanded_image_pads() {
        let mut ctx = ContextState::new();
        ctx.push_no_log(Section::Conversation, AstNode::branch(Role::User, vec![
            AstNode::content("look:"),
            AstNode::image(vec![0xDE, 0xAD], "image/png", 512, 512, qwen3_image_token_count(512, 512)),
        ]));

        // AST side and wire side should both carry N image_pads + bookends —
        // server's session.tokens length must match what vLLM's engine will
        // actually process. Binary image bytes are shipped separately in
        // multi_modal_data via the WireImage list.
        let n_expected = qwen3_image_token_count(512, 512) as usize;

        let full = ctx.token_ids();
        let n_image_pads_full = full.iter()
            .filter(|&&t| t == tokenizer::IMAGE_PAD).count();
        assert_eq!(n_image_pads_full, n_expected);

        let (wire, images, _) = ctx.wire_prompt(0..ctx.conversation().len(), |_| false);
        let n_image_pads_wire = wire.iter()
            .filter(|&&t| t == tokenizer::IMAGE_PAD).count();
        assert_eq!(n_image_pads_wire, n_expected);

        assert_eq!(images.len(), 1);
        assert_eq!(images[0].bytes, vec![0xDE, 0xAD]);
        assert_eq!(images[0].mime, "image/png");

        // One pair of vision_start/vision_end bookends around the N pads.
        assert_eq!(wire.iter().filter(|&&t| t == tokenizer::VISION_START).count(), 1);
        assert_eq!(wire.iter().filter(|&&t| t == tokenizer::VISION_END).count(), 1);
    }

    #[test]
    fn test_image_serde_roundtrip() {
        let node = AstNode::image(vec![0xDE, 0xAD, 0xBE, 0xEF], "image/png", 64, 64, qwen3_image_token_count(64, 64));
        let json = serde_json::to_string(&node).unwrap();
        // bytes must be base64-encoded in the JSON form
        assert!(json.contains("3q2+7w=="));
        let back: AstNode = serde_json::from_str(&json).unwrap();
        let leaf = back.leaf().unwrap();
        match leaf.body() {
            NodeBody::Image { bytes, mime, orig_height, orig_width, token_count } => {
                assert_eq!(bytes, &[0xDE, 0xAD, 0xBE, 0xEF]);
                assert_eq!(mime, "image/png");
                assert_eq!(*orig_height, 64);
                assert_eq!(*orig_width, 64);
                assert_eq!(*token_count, qwen3_image_token_count(64, 64));
            }
            other => panic!("expected Image, got {:?}", other),
        }
        // token_ids are recomputed on deserialization
        assert_eq!(leaf.token_ids().len(), leaf.tokens());
    }

    #[test]
    fn test_timestamp_present_accepted() {
        let json = r#"{"Leaf":{"body":{"Content":"hi"},"timestamp":"2026-04-16T12:00:00Z"}}"#;
        let node: AstNode = serde_json::from_str(json).unwrap();
        let leaf = node.leaf().unwrap();
        assert_eq!(leaf.timestamp().to_rfc3339(),
            "2026-04-16T12:00:00+00:00");
    }
}