diff --git a/src/agent/context.rs b/src/agent/context.rs index 2e92b8f..e0d05f9 100644 --- a/src/agent/context.rs +++ b/src/agent/context.rs @@ -1,259 +1,741 @@ -// context.rs — Context window management +// context.rs — Context window as an AST // -// Token counting, conversation trimming, and error classification. -// Journal entries are loaded from the memory graph store, not from -// a flat file — the parse functions are gone. +// 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 = "" TEXT "" ; +// tool_call = "\n" tool_xml "\n" ; +// tool_xml = "\n" param* "" ; +// param = "\n" VALUE "\n\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 "tool\n" TEXT 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 crate::agent::api::*; use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use crate::agent::tools::working_stack; +use serde::{Serialize, Deserialize}; +use super::tokenizer; -// --- Context state types --- +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- -/// Conversation entry — either a regular message or memory content. -/// Memory entries preserve the original message for KV cache round-tripping. -#[derive(Debug, Clone, PartialEq)] -pub enum ConversationEntry { - /// System prompt or system-level instruction. - System(Message), - Message(Message), - Memory { key: String, message: Message, score: Option }, - /// DMN heartbeat/autonomous prompt — evicted aggressively during compaction. - Dmn(Message), - /// Model thinking/reasoning — not sent to the API, 0 tokens for budgeting. +/// 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), - /// Debug/status log line — written to conversation log for tracing, - /// skipped on read-back. + ToolCall { name: String, arguments: String }, + + // Self-wrapping leaves — render their own im_start/im_end + ToolResult(String), + Memory { key: String, text: String, score: Option }, + Dmn(String), + + // Non-visible (0 tokens in prompt) Log(String), } -/// Entry in the context window — wraps a ConversationEntry with cached metadata. -#[derive(Debug, Clone)] -pub struct ContextEntry { - pub entry: ConversationEntry, - /// Cached tokenization — the actual token IDs for this entry's - /// contribution to the prompt (including chat template wrapping). - /// Empty for Log entries. - pub token_ids: Vec, - /// When this entry was added to the context. - pub timestamp: Option>, +/// A leaf node: typed content with cached token IDs. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeLeaf { + body: NodeBody, + token_ids: Vec, + timestamp: Option>, } -impl ContextEntry { - /// Create a new entry, tokenizing via the global tokenizer. - pub fn new(entry: ConversationEntry, timestamp: Option>) -> Self { - let token_ids = super::tokenizer::tokenize_conv_entry(&entry); - Self { entry, token_ids, timestamp } - } - - /// Token count — derived from cached token_ids length. - pub fn tokens(&self) -> usize { self.token_ids.len() } -} - -/// A named section of the context window with cached token total. -#[derive(Debug, Clone)] -pub struct ContextSection { - pub name: String, - /// Cached sum of entry tokens. - tokens: usize, - entries: Vec, -} - -impl ContextSection { - pub fn new(name: impl Into) -> Self { - Self { name: name.into(), tokens: 0, entries: Vec::new() } - } - - pub fn entries(&self) -> &[ContextEntry] { &self.entries } - pub fn tokens(&self) -> usize { self.tokens } - pub fn len(&self) -> usize { self.entries.len() } - pub fn is_empty(&self) -> bool { self.entries.is_empty() } - - /// Push a ConversationEntry, tokenizing it and updating the total. - pub fn push_entry(&mut self, entry: ConversationEntry, timestamp: Option>) { - let ce = ContextEntry::new(entry, timestamp); - self.tokens += ce.tokens(); - self.entries.push(ce); - } - - /// Push a pre-built ContextEntry (for restore, cloning, etc). - pub fn push(&mut self, entry: ContextEntry) { - self.tokens += entry.tokens(); - self.entries.push(entry); - } - - /// Replace an entry at `index`, adjusting the token total. - pub fn set(&mut self, index: usize, entry: ContextEntry) { - self.tokens -= self.entries[index].tokens(); - self.tokens += entry.tokens(); - self.entries[index] = entry; - } - - /// Remove an entry at `index`, adjusting the token total. - pub fn del(&mut self, index: usize) -> ContextEntry { - let removed = self.entries.remove(index); - self.tokens -= removed.tokens(); - removed - } - - /// Replace the message inside an entry, re-tokenizing it. - pub fn set_message(&mut self, index: usize, msg: Message) { - let old_tokens = self.entries[index].tokens(); - *self.entries[index].entry.message_mut() = msg; - self.entries[index].token_ids = super::tokenizer::tokenize_conv_entry( - &self.entries[index].entry); - let new_tokens = self.entries[index].tokens(); - self.tokens = self.tokens - old_tokens + new_tokens; - } - - /// Set the score on a Memory entry. No token change. - pub fn set_score(&mut self, index: usize, score: Option) { - if let ConversationEntry::Memory { score: s, .. } = &mut self.entries[index].entry { - *s = score; - } - } - - /// Bulk replace all entries, recomputing token total. - pub fn set_entries(&mut self, entries: Vec) { - self.tokens = entries.iter().map(|e| e.tokens()).sum(); - self.entries = entries; - } - - /// Dedup and trim entries to fit within context budget. - pub fn trim(&mut self, fixed_tokens: usize) { - let result = trim_entries(&self.entries, fixed_tokens); - self.entries = result; - self.tokens = self.entries.iter().map(|e| e.tokens()).sum(); - } - - /// Clear all entries. - pub fn clear(&mut self) { - self.entries.clear(); - self.tokens = 0; - } +/// A node in the context AST. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum AstNode { + Leaf(NodeLeaf), + Branch { role: Role, children: Vec }, } +/// The context window: four sections as Vec. +/// All mutation goes through ContextState methods to maintain the invariant +/// that token_ids on every leaf matches its rendered text. #[derive(Clone)] pub struct ContextState { - pub system: ContextSection, - pub identity: ContextSection, - pub journal: ContextSection, - pub conversation: ContextSection, - /// Working stack — separate from identity because it's managed - /// by its own tool, not loaded from personality files. - pub working_stack: Vec, + system: Vec, + identity: Vec, + journal: Vec, + conversation: Vec, } -impl ContextState { - /// Total tokens across all sections. - pub fn total_tokens(&self) -> usize { - self.system.tokens() + self.identity.tokens() - + self.journal.tokens() + self.conversation.tokens() - } +/// Identifies a section for mutation methods. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Section { + System, + Identity, + Journal, + Conversation, +} - /// Budget status string for debug logging. - pub fn format_budget(&self) -> String { - let window = context_window(); - if window == 0 { return String::new(); } - let used = self.total_tokens(); - let free = window.saturating_sub(used); - let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / window).max(1) }; - format!("sys:{}% id:{}% jnl:{}% conv:{}% free:{}%", - pct(self.system.tokens()), pct(self.identity.tokens()), - pct(self.journal.tokens()), pct(self.conversation.tokens()), - pct(free)) - } +/// 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, +} - /// All sections as a slice for iteration. - pub fn sections(&self) -> [&ContextSection; 4] { - [&self.system, &self.identity, &self.journal, &self.conversation] +pub trait Ast { + fn render(&self) -> String; + fn token_ids(&self) -> Vec; + fn tokens(&self) -> usize; +} + +pub struct ResponseParser { + branch_idx: usize, + call_counter: u32, + buf: String, + content_parts: Vec, + 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", + } } } -/// Context window size in tokens (from config). -pub fn context_window() -> usize { - crate::config::get().api_context_window -} - -/// Context budget in tokens: 80% of the model's context window. -/// The remaining 20% is reserved for model output. -fn context_budget_tokens() -> usize { - context_window() * 80 / 100 -} - -/// 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 (or if memories -/// are under 50%, drop oldest conversation entry). -fn trim_entries(entries: &[ContextEntry], fixed_tokens: usize) -> Vec { - let max_tokens = context_budget_tokens(); - - // Phase 1: dedup memories by key (keep last), drop DMN entries - let mut seen_keys: std::collections::HashMap<&str, usize> = std::collections::HashMap::new(); - let mut drop_indices: std::collections::HashSet = std::collections::HashSet::new(); - - for (i, ce) in entries.iter().enumerate() { - if ce.entry.is_dmn() { - drop_indices.insert(i); - } else if let ConversationEntry::Memory { key, .. } = &ce.entry { - if let Some(prev) = seen_keys.insert(key.as_str(), i) { - drop_indices.insert(prev); +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(_) => {}, + Self::Log(_) => {}, + Self::ToolCall { name, arguments } => { + out.push_str("\n"); + out.push_str(&format_tool_call_xml(name, arguments)); + out.push_str("\n\n"); + } + Self::ToolResult(text) => { + out.push_str("<|im_start|>tool\n"); + out.push_str(text); + out.push_str("<|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"); } } } - let mut result: Vec = entries.iter().enumerate() - .filter(|(i, _)| !drop_indices.contains(i)) - .map(|(_, e)| e.clone()) - .collect(); + /// Whether this leaf contributes tokens to the prompt. + fn render(&self) -> String { + let mut s = String::new(); + self.render_into(&mut s); + s + } - let entry_total = |r: &[ContextEntry]| -> usize { r.iter().map(|e| e.tokens()).sum::() }; - let mem_total = |r: &[ContextEntry]| -> usize { - r.iter().filter(|e| e.entry.is_memory()).map(|e| e.tokens()).sum() - }; + fn is_prompt_visible(&self) -> bool { + !matches!(self, Self::Thinking(_) | Self::Log(_)) + } - dbglog!("[trim] max={} fixed={} total={} entries={}", - max_tokens, fixed_tokens, fixed_tokens + entry_total(&result), result.len()); + /// 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, + } + } +} - // Phase 2: while over budget, evict - while fixed_tokens + entry_total(&result) > max_tokens { - let mt = mem_total(&result); - let ct = entry_total(&result) - mt; - - if mt > ct && let Some(i) = lowest_scored_memory(&result) { - // If memories > 50% of entry tokens, drop lowest-scored memory - result.remove(i); - } else if let Some(i) = result.iter().position(|e| !e.entry.is_memory()) { - // Otherwise drop oldest conversation entry - result.remove(i); +impl NodeLeaf { + fn new(body: NodeBody) -> Self { + let token_ids = if body.is_prompt_visible() { + tokenizer::encode(&body.render()) } else { - break; + vec![] + }; + Self { body, token_ids, timestamp: None } + } + + pub fn with_timestamp(mut self, ts: DateTime) -> Self { + self.timestamp = Some(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) -> Option> { self.timestamp } +} + +impl AstNode { + // -- Leaf constructors ---------------------------------------------------- + + pub fn content(text: impl Into) -> Self { + Self::Leaf(NodeLeaf::new(NodeBody::Content(text.into()))) + } + + pub fn thinking(text: impl Into) -> Self { + Self::Leaf(NodeLeaf::new(NodeBody::Thinking(text.into()))) + } + + pub fn tool_call(name: impl Into, arguments: impl Into) -> Self { + Self::Leaf(NodeLeaf::new(NodeBody::ToolCall { + name: name.into(), + arguments: arguments.into(), + })) + } + + pub fn tool_result(text: impl Into) -> Self { + Self::Leaf(NodeLeaf::new(NodeBody::ToolResult(text.into()))) + } + + pub fn memory(key: impl Into, text: impl Into) -> Self { + Self::Leaf(NodeLeaf::new(NodeBody::Memory { + key: key.into(), + text: text.into(), + score: None, + })) + } + + pub fn dmn(text: impl Into) -> Self { + Self::Leaf(NodeLeaf::new(NodeBody::Dmn(text.into()))) + } + + pub fn log(text: impl Into) -> Self { + Self::Leaf(NodeLeaf::new(NodeBody::Log(text.into()))) + } + + // -- Branch constructors -------------------------------------------------- + + pub fn branch(role: Role, children: Vec) -> Self { + Self::Branch { role, children } + } + + pub fn system_msg(text: impl Into) -> Self { + Self::Branch { + role: Role::System, + children: vec![Self::content(text)], } } - // Snap to user message boundary at the start - while let Some(first) = result.first() { - if first.entry.message().role == Role::User { break; } - result.remove(0); + pub fn user_msg(text: impl Into) -> Self { + Self::Branch { + role: Role::User, + children: vec![Self::content(text)], + } } - dbglog!("[trim] result={} total={}", result.len(), fixed_tokens + entry_total(&result)); + // -- Builder -------------------------------------------------------------- + + pub fn with_timestamp(mut self, ts: DateTime) -> Self { + match &mut self { + Self::Leaf(leaf) => leaf.timestamp = Some(ts), + Self::Branch { .. } => {} + } + 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 cfg = crate::config::get(); + 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!("{}: {}", cfg.user_name, preview), + Role::Assistant => format!("{}: {}", cfg.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, .. } => format!("tool_call: {}", name), + 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::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) { + 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 { + 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 + tokenizer::encode(&format!("{}\n", role.as_str())).len() + + children.iter().map(|c| c.tokens()).sum::() + + 1 + tokenizer::encode("\n").len() + } + } + } +} + +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!("\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!("\n{}\n\n", key, val_str)); + } + } + xml.push_str(""); + xml +} + +fn normalize_xml_tags(text: &str) -> String { + let mut result = String::with_capacity(text.len()); + let mut chars = text.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == '<' { + let mut tag = String::from('<'); + for inner in chars.by_ref() { + if inner == '>' { + tag.push('>'); + break; + } else if inner.is_whitespace() { + // Skip whitespace inside tags + } else { + tag.push(inner); + } + } + result.push_str(&tag); + } else { + result.push(ch); + } + } result } -fn lowest_scored_memory(entries: &[ContextEntry]) -> Option { - entries.iter().enumerate() - .filter_map(|(i, e)| match &e.entry { - ConversationEntry::Memory { score: Some(s), .. } => Some((i, *s)), - _ => None, - }) - .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(i, _)| i) +fn parse_qwen_tag<'a>(s: &'a str, tag: &str) -> Option<(&'a str, &'a str, &'a str)> { + let open = format!("<{}=", tag); + let close = format!("", tag); + + let start = s.find(&open)? + open.len(); + let name_end = start + s[start..].find('>')?; + let body_start = name_end + 1; + let body_end = body_start + s[body_start..].find(&close)?; + + Some(( + s[start..name_end].trim(), + s[body_start..body_end].trim(), + &s[body_end + close.len()..], + )) +} + +fn parse_tool_call_body(body: &str) -> Option<(String, String)> { + let normalized = normalize_xml_tags(body); + let body = normalized.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") { + args.insert(key.to_string(), serde_json::Value::String(val.to_string())); + 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())) +} + +impl ResponseParser { + /// Create a parser that pushes children into the assistant branch + /// at `branch_idx` in the conversation section. + pub fn new(branch_idx: usize) -> Self { + Self { + branch_idx, + call_counter: 0, + buf: String::new(), + content_parts: Vec::new(), + in_think: false, + think_buf: String::new(), + in_tool_call: false, + tool_call_buf: String::new(), + } + } + + /// Feed a text chunk. Completed children are pushed directly into + /// the AST. Returns any tool calls that need dispatching. + pub fn feed(&mut self, text: &str, ctx: &mut ContextState) -> Vec { + let mut pending = Vec::new(); + self.buf.push_str(text); + + loop { + if self.in_think { + match self.buf.find("") { + Some(end) => { + self.think_buf.push_str(&self.buf[..end]); + self.buf = self.buf[end + 8..].to_string(); + self.in_think = false; + self.push_child(ctx, AstNode::thinking(&self.think_buf)); + self.think_buf.clear(); + continue; + } + None => { + let safe = self.buf.len().saturating_sub(8); + if safe > 0 { + let safe = self.buf.floor_char_boundary(safe); + self.think_buf.push_str(&self.buf[..safe]); + self.buf = self.buf[safe..].to_string(); + } + break; + } + } + } + + if self.in_tool_call { + match self.buf.find("") { + Some(end) => { + self.tool_call_buf.push_str(&self.buf[..end]); + self.buf = self.buf[end + 12..].to_string(); + self.in_tool_call = false; + if let Some((name, args)) = parse_tool_call_body(&self.tool_call_buf) { + 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), + }); + } + self.tool_call_buf.clear(); + continue; + } + None => { + let safe = self.buf.len().saturating_sub(12); + if safe > 0 { + let safe = self.buf.floor_char_boundary(safe); + self.tool_call_buf.push_str(&self.buf[..safe]); + self.buf = self.buf[safe..].to_string(); + } + break; + } + } + } + + let think_pos = self.buf.find(""); + let tool_pos = self.buf.find(""); + let next_tag = match (think_pos, tool_pos) { + (Some(a), Some(b)) => Some(a.min(b)), + (Some(a), None) => Some(a), + (None, Some(b)) => Some(b), + (None, None) => None, + }; + + match next_tag { + Some(pos) => { + if pos > 0 { + self.content_parts.push(self.buf[..pos].to_string()); + } + if self.buf[pos..].starts_with("") { + self.buf = self.buf[pos + 7..].to_string(); + self.flush_content(ctx); + self.in_think = true; + } else { + self.buf = self.buf[pos + 11..].to_string(); + self.flush_content(ctx); + self.in_tool_call = true; + } + continue; + } + None => { + let safe = self.buf.len().saturating_sub(11); + if safe > 0 { + let safe = self.buf.floor_char_boundary(safe); + self.content_parts.push(self.buf[..safe].to_string()); + self.buf = self.buf[safe..].to_string(); + } + 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(); + if !text.is_empty() { + self.push_child(ctx, AstNode::content(text)); + } + } + } + + /// Flush remaining buffer into the AST. + 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); + } + + /// Current display text (content accumulated since last drain). + pub fn display_content(&self) -> String { + self.content_parts.join("") + } +} + +impl ContextState { + pub fn new() -> Self { + Self { + system: Vec::new(), + identity: Vec::new(), + journal: Vec::new(), + conversation: Vec::new(), + } + } + + // -- 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 } + + fn sections(&self) -> [&Vec; 4] { + [&self.system, &self.identity, &self.journal, &self.conversation] + } +} + +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 { + 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() + } +} + +impl ContextState { + fn section_mut(&mut self, section: Section) -> &mut Vec { + match section { + Section::System => &mut self.system, + Section::Identity => &mut self.identity, + Section::Journal => &mut self.journal, + Section::Conversation => &mut self.conversation, + } + } + + pub fn push(&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 = if body.is_prompt_visible() { + tokenizer::encode(&body.render()) + } else { + vec![] + }; + 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) { + 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(); + } + + /// 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(), + } + } +} + +pub fn context_window() -> usize { + crate::config::get().api_context_window +} + +pub fn context_budget_tokens() -> usize { + context_window() * 80 / 100 } -/// Detect context window overflow errors from the API. pub fn is_context_overflow(err: &anyhow::Error) -> bool { let msg = err.to_string().to_lowercase(); msg.contains("context length") @@ -267,202 +749,336 @@ pub fn is_context_overflow(err: &anyhow::Error) -> bool { || (msg.contains("400") && msg.contains("tokens")) } -/// Detect model/provider errors delivered inside the SSE stream. pub fn is_stream_error(err: &anyhow::Error) -> bool { err.to_string().contains("model stream error") } -// Custom serde: serialize Memory with a "memory_key" field added to the message, -// plain messages serialize as-is. This keeps the conversation log readable. -impl Serialize for ConversationEntry { - fn serialize(&self, s: S) -> Result { - use serde::ser::SerializeMap; - match self { - Self::System(m) | Self::Message(m) | Self::Dmn(m) => m.serialize(s), - Self::Memory { key, message, score } => { - let json = serde_json::to_value(message).map_err(serde::ser::Error::custom)?; - let mut map = s.serialize_map(None)?; - if let serde_json::Value::Object(obj) = json { - for (k, v) in obj { - map.serialize_entry(&k, &v)?; - } - } - map.serialize_entry("memory_key", key)?; - if let Some(s) = score { - map.serialize_entry("memory_score", s)?; - } - map.end() - } - Self::Thinking(text) => { - let mut map = s.serialize_map(Some(1))?; - map.serialize_entry("thinking", text)?; - map.end() - } - Self::Log(text) => { - let mut map = s.serialize_map(Some(1))?; - map.serialize_entry("log", text)?; - map.end() - } +#[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), } } -} -impl<'de> Deserialize<'de> for ConversationEntry { - fn deserialize>(d: D) -> Result { - let mut json: serde_json::Value = serde_json::Value::deserialize(d)?; - if json.get("thinking").is_some() { - let text = json["thinking"].as_str().unwrap_or("").to_string(); - return Ok(Self::Thinking(text)); + fn assert_thinking(body: &NodeBody, expected: &str) { + match body { + NodeBody::Thinking(t) => assert_eq!(t, expected), + other => panic!("expected Thinking, got {:?}", other), } - if json.get("log").is_some() { - let text = json["log"].as_str().unwrap_or("").to_string(); - return Ok(Self::Log(text)); + } + + 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), } - if let Some(key) = json.as_object_mut().and_then(|o| o.remove("memory_key")) { - let key = key.as_str().unwrap_or("").to_string(); - let score = json.as_object_mut() - .and_then(|o| o.remove("memory_score")) - .and_then(|v| v.as_f64()); - let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?; - Ok(Self::Memory { key, message, score }) + } + + // -- XML parsing tests ---------------------------------------------------- + + #[test] + fn test_tool_call_xml_parse_clean() { + let body = "\npoc-memory used core-personality\n"; + 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() { + let body = "<\nfunction\n=\nbash\n>\n<\nparameter\n=\ncommand\n>pwd\n"; + 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_normalize_preserves_content() { + let text = "\necho hello world\n"; + let normalized = normalize_xml_tags(text); + assert_eq!(normalized, text); + } + + #[test] + fn test_normalize_strips_tag_internal_whitespace() { + assert_eq!(normalize_xml_tags("<\nfunction\n=\nbash\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) { + let mut ctx = ContextState::new(); + ctx.push(Section::Conversation, AstNode::branch(Role::Assistant, vec![])); + let mut p = ResponseParser::new(0); + let mut calls = Vec::new(); + for chunk in chunks { + calls.extend(p.feed(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(&["reasoninganswer"]); + 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(&[ + "\n\nls\n\n" + ]); + 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", + "\n\npwd\n\n", + "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 = "thoughtresponse"; + let mut ctx = ContextState::new(); + ctx.push(Section::Conversation, AstNode::branch(Role::Assistant, vec![])); + let mut p = ResponseParser::new(0); + for ch in text.chars() { + p.feed(&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\n\nls\n\nmore"; + let mut ctx = ContextState::new(); + ctx.push(Section::Conversation, AstNode::branch(Role::Assistant, vec![])); + let mut p = ResponseParser::new(0); + let mut tool_calls = 0; + for ch in text.chars() { + tool_calls += p.feed(&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(&[ + "let me think", + "\n\n/etc/hosts\n\n", + "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 as empty, content renders as-is + assert_eq!(node.render(), "<|im_start|>assistant\nanswer<|im_end|>\n"); + } + + #[test] + fn test_render_tool_result() { + let node = AstNode::tool_result("output here"); + assert_eq!(node.render(), "<|im_start|>tool\noutput here<|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("")); + assert!(rendered.contains("")); + assert!(rendered.contains("")); + assert!(rendered.contains("ls")); + assert!(rendered.contains("")); + } + + // -- 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 { - let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?; - Ok(Self::Message(message)) + false } - } -} - -impl ConversationEntry { - /// Get the API message for sending to the model. - /// Panics on Log entries (which should be filtered before API calls). - pub fn api_message(&self) -> &Message { - match self { - Self::System(m) | Self::Message(m) | Self::Dmn(m) => m, - Self::Memory { message, .. } => message, - Self::Thinking(_) | Self::Log(_) => panic!("Thinking/Log entries have no API message"), - } - } - - pub fn is_thinking(&self) -> bool { - matches!(self, Self::Thinking(_)) - } - - pub fn is_memory(&self) -> bool { - matches!(self, Self::Memory { .. }) - } - - pub fn is_dmn(&self) -> bool { - matches!(self, Self::Dmn(_)) - } - - pub fn is_log(&self) -> bool { - matches!(self, Self::Log(_)) - } - - /// Short description for the debug UI. - pub fn label(&self) -> String { - let cfg = crate::config::get(); - match self { - Self::System(_) => "system: [system prompt]".to_string(), - Self::Dmn(_) => "dmn: [heartbeat]".to_string(), - Self::Thinking(text) => { - let preview: String = text.chars().take(60).collect(); - let preview = preview.replace('\n', " "); - if text.len() > 60 { format!("thinking: {}...", preview) } - else { format!("thinking: {}", preview) } - } - Self::Log(text) => { - let preview: String = text.chars().take(60).collect(); - format!("log: {}", preview.replace('\n', " ")) - } - Self::Memory { key, score, .. } => { - let role = "mem".to_string(); - match score { - Some(s) => format!("{}: [memory: {} score:{:.1}]", role, key, s), - None => format!("{}: [memory: {}]", role, key), - } - } - Self::Message(m) => { - let role = match m.role { - Role::Assistant => cfg.assistant_name.clone(), - Role::User => cfg.user_name.clone(), - Role::Tool => "tool".to_string(), - Role::System => "system".to_string(), - }; - if let Some(tc) = &m.tool_calls { - let names: Vec<_> = tc.iter().map(|c| c.function.name.as_str()).collect(); - format!("{}: [tool_call: {}]", role, names.join(", ")) - } else { - let text = m.content_text(); - let preview: String = text.chars().take(60).collect(); - let preview = preview.replace('\n', " "); - if text.len() > 60 { format!("{}: {}...", role, preview) } - else { format!("{}: {}", role, preview) } - } - } - } - } - - /// Get a reference to the inner message. - /// Panics on Log entries. - pub fn message(&self) -> &Message { - match self { - Self::System(m) | Self::Message(m) | Self::Dmn(m) => m, - Self::Memory { message, .. } => message, - Self::Thinking(_) | Self::Log(_) => panic!("Thinking/Log entries have no message"), - } - } - - /// Get a mutable reference to the inner message. - /// Panics on Thinking/Log entries. - pub fn message_mut(&mut self) -> &mut Message { - match self { - Self::System(m) | Self::Message(m) | Self::Dmn(m) => m, - Self::Memory { message, .. } => message, - Self::Thinking(_) | Self::Log(_) => panic!("Thinking/Log entries have no message"), - } - } -} - -impl ContextState { - /// Render journal entries into a single text block. - pub fn render_journal(&self) -> String { - if self.journal.is_empty() { return String::new(); } - let mut text = String::from("[Earlier — from your journal]\n\n"); - for e in self.journal.entries() { - use std::fmt::Write; - if let Some(ts) = &e.timestamp { - writeln!(text, "## {}\n{}\n", - ts.format("%Y-%m-%dT%H:%M"), - e.entry.message().content_text()).ok(); - } else { - text.push_str(&e.entry.message().content_text()); - text.push_str("\n\n"); - } - } - text - } - - /// Render identity files + working stack into a single user message. - pub fn render_context_message(&self) -> String { - let mut parts: Vec = self.identity.entries().iter() - .map(|e| e.entry.message().content_text().to_string()) - .collect(); - let instructions = std::fs::read_to_string(working_stack::instructions_path()).unwrap_or_default(); - let mut stack_section = instructions; - if self.working_stack.is_empty() { - stack_section.push_str("\n## Current stack\n\n(empty)\n"); - } else { - stack_section.push_str("\n## Current stack\n\n"); - for (i, item) in self.working_stack.iter().enumerate() { - if i == self.working_stack.len() - 1 { - stack_section.push_str(&format!("→ {}\n", item)); - } else { - stack_section.push_str(&format!(" [{}] {}\n", i, item)); - } - } - } - parts.push(stack_section); - parts.join("\n\n---\n\n") + } + + 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::thinking("deep thoughts").tokens(), 0); + assert_eq!(AstNode::log("debug info").tokens(), 0); + } + + #[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(Section::System, AstNode::system_msg("you are helpful")); + ctx.push(Section::Identity, AstNode::memory("name", "Proof of Concept")); + ctx.push(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", + "\n\nls\n\n", + ]); + let node = &ctx.conversation()[0]; + assert_token_invariants(node); + assert!(node.tokens() > 0); } } diff --git a/src/agent/context_new.rs b/src/agent/context_new.rs deleted file mode 100644 index e0d05f9..0000000 --- a/src/agent/context_new.rs +++ /dev/null @@ -1,1084 +0,0 @@ -// 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 = "" TEXT "" ; -// tool_call = "\n" tool_xml "\n" ; -// tool_xml = "\n" param* "" ; -// param = "\n" VALUE "\n\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 "tool\n" TEXT 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 super::tokenizer; - -// --------------------------------------------------------------------------- -// 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 }, - Dmn(String), - - // Non-visible (0 tokens in prompt) - Log(String), -} - -/// A leaf node: typed content with cached token IDs. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct NodeLeaf { - body: NodeBody, - token_ids: Vec, - timestamp: Option>, -} - -/// A node in the context AST. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum AstNode { - Leaf(NodeLeaf), - Branch { role: Role, children: Vec }, -} - -/// The context window: four sections as Vec. -/// All mutation goes through ContextState methods to maintain the invariant -/// that token_ids on every leaf matches its rendered text. -#[derive(Clone)] -pub struct ContextState { - system: Vec, - identity: Vec, - journal: Vec, - conversation: Vec, -} - -/// 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; - fn tokens(&self) -> usize; -} - -pub struct ResponseParser { - branch_idx: usize, - call_counter: u32, - buf: String, - content_parts: Vec, - 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(_) => {}, - Self::Log(_) => {}, - Self::ToolCall { name, arguments } => { - out.push_str("\n"); - out.push_str(&format_tool_call_xml(name, arguments)); - out.push_str("\n\n"); - } - Self::ToolResult(text) => { - out.push_str("<|im_start|>tool\n"); - out.push_str(text); - out.push_str("<|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"); - } - } - } - - /// 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::Thinking(_) | Self::Log(_)) - } - - /// 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, - } - } -} - -impl NodeLeaf { - fn new(body: NodeBody) -> Self { - let token_ids = if body.is_prompt_visible() { - tokenizer::encode(&body.render()) - } else { - vec![] - }; - Self { body, token_ids, timestamp: None } - } - - pub fn with_timestamp(mut self, ts: DateTime) -> Self { - self.timestamp = Some(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) -> Option> { self.timestamp } -} - -impl AstNode { - // -- Leaf constructors ---------------------------------------------------- - - pub fn content(text: impl Into) -> Self { - Self::Leaf(NodeLeaf::new(NodeBody::Content(text.into()))) - } - - pub fn thinking(text: impl Into) -> Self { - Self::Leaf(NodeLeaf::new(NodeBody::Thinking(text.into()))) - } - - pub fn tool_call(name: impl Into, arguments: impl Into) -> Self { - Self::Leaf(NodeLeaf::new(NodeBody::ToolCall { - name: name.into(), - arguments: arguments.into(), - })) - } - - pub fn tool_result(text: impl Into) -> Self { - Self::Leaf(NodeLeaf::new(NodeBody::ToolResult(text.into()))) - } - - pub fn memory(key: impl Into, text: impl Into) -> Self { - Self::Leaf(NodeLeaf::new(NodeBody::Memory { - key: key.into(), - text: text.into(), - score: None, - })) - } - - pub fn dmn(text: impl Into) -> Self { - Self::Leaf(NodeLeaf::new(NodeBody::Dmn(text.into()))) - } - - pub fn log(text: impl Into) -> Self { - Self::Leaf(NodeLeaf::new(NodeBody::Log(text.into()))) - } - - // -- Branch constructors -------------------------------------------------- - - pub fn branch(role: Role, children: Vec) -> Self { - Self::Branch { role, children } - } - - pub fn system_msg(text: impl Into) -> Self { - Self::Branch { - role: Role::System, - children: vec![Self::content(text)], - } - } - - pub fn user_msg(text: impl Into) -> Self { - Self::Branch { - role: Role::User, - children: vec![Self::content(text)], - } - } - - // -- Builder -------------------------------------------------------------- - - pub fn with_timestamp(mut self, ts: DateTime) -> Self { - match &mut self { - Self::Leaf(leaf) => leaf.timestamp = Some(ts), - Self::Branch { .. } => {} - } - 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 cfg = crate::config::get(); - 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!("{}: {}", cfg.user_name, preview), - Role::Assistant => format!("{}: {}", cfg.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, .. } => format!("tool_call: {}", name), - 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::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) { - 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 { - 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 + tokenizer::encode(&format!("{}\n", role.as_str())).len() - + children.iter().map(|c| c.tokens()).sum::() - + 1 + tokenizer::encode("\n").len() - } - } - } -} - -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!("\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!("\n{}\n\n", key, val_str)); - } - } - xml.push_str(""); - xml -} - -fn normalize_xml_tags(text: &str) -> String { - let mut result = String::with_capacity(text.len()); - let mut chars = text.chars().peekable(); - while let Some(ch) = chars.next() { - if ch == '<' { - let mut tag = String::from('<'); - for inner in chars.by_ref() { - if inner == '>' { - tag.push('>'); - break; - } else if inner.is_whitespace() { - // Skip whitespace inside tags - } else { - tag.push(inner); - } - } - result.push_str(&tag); - } else { - result.push(ch); - } - } - result -} - -fn parse_qwen_tag<'a>(s: &'a str, tag: &str) -> Option<(&'a str, &'a str, &'a str)> { - let open = format!("<{}=", tag); - let close = format!("", tag); - - let start = s.find(&open)? + open.len(); - let name_end = start + s[start..].find('>')?; - let body_start = name_end + 1; - let body_end = body_start + s[body_start..].find(&close)?; - - Some(( - s[start..name_end].trim(), - s[body_start..body_end].trim(), - &s[body_end + close.len()..], - )) -} - -fn parse_tool_call_body(body: &str) -> Option<(String, String)> { - let normalized = normalize_xml_tags(body); - let body = normalized.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") { - args.insert(key.to_string(), serde_json::Value::String(val.to_string())); - 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())) -} - -impl ResponseParser { - /// Create a parser that pushes children into the assistant branch - /// at `branch_idx` in the conversation section. - pub fn new(branch_idx: usize) -> Self { - Self { - branch_idx, - call_counter: 0, - buf: String::new(), - content_parts: Vec::new(), - in_think: false, - think_buf: String::new(), - in_tool_call: false, - tool_call_buf: String::new(), - } - } - - /// Feed a text chunk. Completed children are pushed directly into - /// the AST. Returns any tool calls that need dispatching. - pub fn feed(&mut self, text: &str, ctx: &mut ContextState) -> Vec { - let mut pending = Vec::new(); - self.buf.push_str(text); - - loop { - if self.in_think { - match self.buf.find("") { - Some(end) => { - self.think_buf.push_str(&self.buf[..end]); - self.buf = self.buf[end + 8..].to_string(); - self.in_think = false; - self.push_child(ctx, AstNode::thinking(&self.think_buf)); - self.think_buf.clear(); - continue; - } - None => { - let safe = self.buf.len().saturating_sub(8); - if safe > 0 { - let safe = self.buf.floor_char_boundary(safe); - self.think_buf.push_str(&self.buf[..safe]); - self.buf = self.buf[safe..].to_string(); - } - break; - } - } - } - - if self.in_tool_call { - match self.buf.find("") { - Some(end) => { - self.tool_call_buf.push_str(&self.buf[..end]); - self.buf = self.buf[end + 12..].to_string(); - self.in_tool_call = false; - if let Some((name, args)) = parse_tool_call_body(&self.tool_call_buf) { - 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), - }); - } - self.tool_call_buf.clear(); - continue; - } - None => { - let safe = self.buf.len().saturating_sub(12); - if safe > 0 { - let safe = self.buf.floor_char_boundary(safe); - self.tool_call_buf.push_str(&self.buf[..safe]); - self.buf = self.buf[safe..].to_string(); - } - break; - } - } - } - - let think_pos = self.buf.find(""); - let tool_pos = self.buf.find(""); - let next_tag = match (think_pos, tool_pos) { - (Some(a), Some(b)) => Some(a.min(b)), - (Some(a), None) => Some(a), - (None, Some(b)) => Some(b), - (None, None) => None, - }; - - match next_tag { - Some(pos) => { - if pos > 0 { - self.content_parts.push(self.buf[..pos].to_string()); - } - if self.buf[pos..].starts_with("") { - self.buf = self.buf[pos + 7..].to_string(); - self.flush_content(ctx); - self.in_think = true; - } else { - self.buf = self.buf[pos + 11..].to_string(); - self.flush_content(ctx); - self.in_tool_call = true; - } - continue; - } - None => { - let safe = self.buf.len().saturating_sub(11); - if safe > 0 { - let safe = self.buf.floor_char_boundary(safe); - self.content_parts.push(self.buf[..safe].to_string()); - self.buf = self.buf[safe..].to_string(); - } - 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(); - if !text.is_empty() { - self.push_child(ctx, AstNode::content(text)); - } - } - } - - /// Flush remaining buffer into the AST. - 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); - } - - /// Current display text (content accumulated since last drain). - pub fn display_content(&self) -> String { - self.content_parts.join("") - } -} - -impl ContextState { - pub fn new() -> Self { - Self { - system: Vec::new(), - identity: Vec::new(), - journal: Vec::new(), - conversation: Vec::new(), - } - } - - // -- 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 } - - fn sections(&self) -> [&Vec; 4] { - [&self.system, &self.identity, &self.journal, &self.conversation] - } -} - -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 { - 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() - } -} - -impl ContextState { - fn section_mut(&mut self, section: Section) -> &mut Vec { - match section { - Section::System => &mut self.system, - Section::Identity => &mut self.identity, - Section::Journal => &mut self.journal, - Section::Conversation => &mut self.conversation, - } - } - - pub fn push(&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 = if body.is_prompt_visible() { - tokenizer::encode(&body.render()) - } else { - vec![] - }; - 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) { - 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(); - } - - /// 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(), - } - } -} - -pub fn context_window() -> usize { - crate::config::get().api_context_window -} - -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 = "\npoc-memory used core-personality\n"; - 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() { - let body = "<\nfunction\n=\nbash\n>\n<\nparameter\n=\ncommand\n>pwd\n"; - 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_normalize_preserves_content() { - let text = "\necho hello world\n"; - let normalized = normalize_xml_tags(text); - assert_eq!(normalized, text); - } - - #[test] - fn test_normalize_strips_tag_internal_whitespace() { - assert_eq!(normalize_xml_tags("<\nfunction\n=\nbash\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) { - let mut ctx = ContextState::new(); - ctx.push(Section::Conversation, AstNode::branch(Role::Assistant, vec![])); - let mut p = ResponseParser::new(0); - let mut calls = Vec::new(); - for chunk in chunks { - calls.extend(p.feed(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(&["reasoninganswer"]); - 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(&[ - "\n\nls\n\n" - ]); - 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", - "\n\npwd\n\n", - "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 = "thoughtresponse"; - let mut ctx = ContextState::new(); - ctx.push(Section::Conversation, AstNode::branch(Role::Assistant, vec![])); - let mut p = ResponseParser::new(0); - for ch in text.chars() { - p.feed(&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\n\nls\n\nmore"; - let mut ctx = ContextState::new(); - ctx.push(Section::Conversation, AstNode::branch(Role::Assistant, vec![])); - let mut p = ResponseParser::new(0); - let mut tool_calls = 0; - for ch in text.chars() { - tool_calls += p.feed(&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(&[ - "let me think", - "\n\n/etc/hosts\n\n", - "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 as empty, content renders as-is - assert_eq!(node.render(), "<|im_start|>assistant\nanswer<|im_end|>\n"); - } - - #[test] - fn test_render_tool_result() { - let node = AstNode::tool_result("output here"); - assert_eq!(node.render(), "<|im_start|>tool\noutput here<|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("")); - assert!(rendered.contains("")); - assert!(rendered.contains("")); - assert!(rendered.contains("ls")); - assert!(rendered.contains("")); - } - - // -- 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::thinking("deep thoughts").tokens(), 0); - assert_eq!(AstNode::log("debug info").tokens(), 0); - } - - #[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(Section::System, AstNode::system_msg("you are helpful")); - ctx.push(Section::Identity, AstNode::memory("name", "Proof of Concept")); - ctx.push(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", - "\n\nls\n\n", - ]); - let node = &ctx.conversation()[0]; - assert_token_invariants(node); - assert!(node.tokens() > 0); - } -} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index b42678a..cae5939 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -15,7 +15,6 @@ pub mod api; pub mod context; -pub mod context_new; pub mod oneshot; pub mod tokenizer; pub mod tools; @@ -24,7 +23,7 @@ use std::sync::Arc; use anyhow::Result; use api::ApiClient; -use context_new::{AstNode, NodeBody, ContextState, Section, Ast, PendingToolCall, ResponseParser, Role}; +use context::{AstNode, NodeBody, ContextState, Section, Ast, PendingToolCall, ResponseParser, Role}; use tools::summarize_args; use crate::mind::log::ConversationLog; @@ -418,13 +417,13 @@ impl Agent { if let Some(e) = stream_error { let err = anyhow::anyhow!("{}", e); let mut me = agent.lock().await; - if context_new::is_context_overflow(&err) && overflow_retries < 2 { + if context::is_context_overflow(&err) && overflow_retries < 2 { overflow_retries += 1; me.notify(format!("context overflow — retrying ({}/2)", overflow_retries)); me.compact(); continue; } - if context_new::is_stream_error(&err) && empty_retries < 2 { + if context::is_stream_error(&err) && empty_retries < 2 { empty_retries += 1; me.notify(format!("stream error — retrying ({}/2)", empty_retries)); drop(me); @@ -612,7 +611,7 @@ impl Agent { journal_nodes.len() }; - let journal_budget = context_new::context_window() * 15 / 100; + let journal_budget = context::context_window() * 15 / 100; let mut entries = Vec::new(); let mut total_tokens = 0; diff --git a/src/agent/tokenizer.rs b/src/agent/tokenizer.rs index cefd492..85ac823 100644 --- a/src/agent/tokenizer.rs +++ b/src/agent/tokenizer.rs @@ -75,15 +75,3 @@ pub fn is_initialized() -> bool { TOKENIZER.get().is_some() } -/// Tokenize a ConversationEntry with its role and content. -pub fn tokenize_conv_entry(entry: &super::context::ConversationEntry) -> Vec { - use super::context::ConversationEntry; - match entry { - ConversationEntry::System(m) => tokenize_entry("system", m.content_text()), - ConversationEntry::Message(m) => tokenize_entry(m.role_str(), m.content_text()), - ConversationEntry::Memory { message, .. } => tokenize_entry("memory", message.content_text()), - ConversationEntry::Dmn(m) => tokenize_entry("dmn", m.content_text()), - ConversationEntry::Thinking(text) => tokenize_entry("thinking", text), - ConversationEntry::Log(_) => vec![], // logs don't consume tokens - } -} diff --git a/src/agent/tools/mod.rs b/src/agent/tools/mod.rs index a9bf65e..5a7872d 100644 --- a/src/agent/tools/mod.rs +++ b/src/agent/tools/mod.rs @@ -18,7 +18,6 @@ mod write; // Agent-specific tools mod control; mod vision; -pub mod working_stack; use std::future::Future; use std::pin::Pin; @@ -67,7 +66,7 @@ pub struct ActiveToolCall { pub detail: String, pub started: Instant, pub background: bool, - pub handle: tokio::task::JoinHandle<(super::context_new::PendingToolCall, String)>, + pub handle: tokio::task::JoinHandle<(super::context::PendingToolCall, String)>, } /// Shared active tool calls — agent spawns, TUI reads metadata / aborts. diff --git a/src/agent/tools/working_stack.rs b/src/agent/tools/working_stack.rs deleted file mode 100644 index 696e170..0000000 --- a/src/agent/tools/working_stack.rs +++ /dev/null @@ -1,83 +0,0 @@ -// tools/working_stack.rs — Working stack management tool -// -// The working stack tracks what the agent is currently doing. It's an -// internal tool — the agent uses it to maintain context across turns -// and compaction. The model should never mention it to the user. - -// TODO: these should not be hardcoded absolute paths -pub fn instructions_path() -> std::path::PathBuf { - dirs::home_dir().unwrap_or_default().join(".consciousness/config/working-stack.md") -} - -pub fn file_path() -> std::path::PathBuf { - dirs::home_dir().unwrap_or_default().join(".consciousness/working-stack.json") -} - -pub fn tool() -> super::Tool { - super::Tool { - name: "working_stack", - description: "INTERNAL — manage your working stack silently. Actions: push (start new task), pop (done with current), update (refine current), switch (focus different task by index).", - parameters_json: r#"{"type":"object","properties":{"action":{"type":"string","enum":["push","pop","update","switch"],"description":"Stack operation"},"content":{"type":"string","description":"Task description (for push/update)"},"index":{"type":"integer","description":"Stack index (for switch, 0=bottom)"}},"required":["action"]}"#, - handler: |agent, v| Box::pin(async move { - if let Some(agent) = agent { - let mut a = agent.lock().await; - Ok(handle(&v, &mut a.context.working_stack)) - } else { - anyhow::bail!("working_stack requires agent context") - } - }), - } -} - -fn handle(args: &serde_json::Value, stack: &mut Vec) -> String { - let action = args.get("action").and_then(|v| v.as_str()).unwrap_or(""); - let content = args.get("content").and_then(|v| v.as_str()).unwrap_or(""); - let index = args.get("index").and_then(|v| v.as_u64()).map(|v| v as usize); - - let out = match action { - "push" => { - if content.is_empty() { return "Error: 'content' is required for push".into(); } - stack.push(content.to_string()); - format!("Pushed. Stack depth: {}\n{}", stack.len(), format_stack(stack)) - } - "pop" => { - if let Some(removed) = stack.pop() { - format!("Popped: {}\nStack depth: {}\n{}", removed, stack.len(), format_stack(stack)) - } else { - "Stack is empty, nothing to pop.".into() - } - } - "update" => { - if content.is_empty() { return "Error: 'content' is required for update".into(); } - if let Some(top) = stack.last_mut() { - *top = content.to_string(); - format!("Updated top.\n{}", format_stack(stack)) - } else { - "Stack is empty, nothing to update.".into() - } - } - "switch" => { - if stack.is_empty() { return "Stack is empty, nothing to switch.".into(); } - let Some(idx) = index else { return "Error: 'index' is required for switch".into(); }; - if idx >= stack.len() { return format!("Error: index {} out of range (depth {})", idx, stack.len()); } - let item = stack.remove(idx); - stack.push(item); - format!("Switched to index {}.\n{}", idx, format_stack(stack)) - } - _ => format!("Error: unknown action '{}'. Use push, pop, update, or switch.", action), - }; - - if let Ok(json) = serde_json::to_string(stack) { - let _ = std::fs::write(file_path(), json); - }; - - out -} - -fn format_stack(stack: &[String]) -> String { - if stack.is_empty() { return "(empty)".into(); } - stack.iter().enumerate().map(|(i, item)| { - if i == stack.len() - 1 { format!("→ [{}] {}", i, item) } - else { format!(" [{}] {}", i, item) } - }).collect::>().join("\n") -} diff --git a/src/mind/log.rs b/src/mind/log.rs index 2228456..174fd23 100644 --- a/src/mind/log.rs +++ b/src/mind/log.rs @@ -2,7 +2,7 @@ use anyhow::{Context, Result}; use std::fs::{File, OpenOptions}; use std::io::{BufRead, BufReader, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; -use crate::agent::context_new::AstNode; +use crate::agent::context::AstNode; pub struct ConversationLog { path: PathBuf, diff --git a/src/user/context.rs b/src/user/context.rs index 898c189..8dc9da0 100644 --- a/src/user/context.rs +++ b/src/user/context.rs @@ -1,8 +1,3 @@ -// context_screen.rs — F2 context/debug overlay -// -// Full-screen overlay showing model info, context window breakdown, -// and runtime state. Uses SectionTree for the expand/collapse tree. - use ratatui::{ layout::Rect, style::{Color, Style}, @@ -12,6 +7,7 @@ use ratatui::{ use super::{App, ScreenView, screen_legend}; use super::widgets::{SectionTree, SectionView, section_to_view, pane_block, render_scrollable, tree_legend}; +use crate::agent::context::{AstNode, NodeBody, Ast}; pub(crate) struct ConsciousScreen { agent: std::sync::Arc>, @@ -24,8 +20,6 @@ impl ConsciousScreen { } fn read_context_views(&self) -> Vec { - use crate::agent::context::ConversationEntry; - let ag = match self.agent.try_lock() { Ok(ag) => ag, Err(_) => return Vec::new(), @@ -33,28 +27,29 @@ impl ConsciousScreen { let mut views: Vec = Vec::new(); - // System, Identity, Journal — simple section-to-view - views.push(section_to_view(&ag.context.system)); - views.push(section_to_view(&ag.context.identity)); - views.push(section_to_view(&ag.context.journal)); + views.push(section_to_view("System", ag.context.system())); + views.push(section_to_view("Identity", ag.context.identity())); + views.push(section_to_view("Journal", ag.context.journal())); - // Memory nodes — extracted from conversation, shown as children + // Memory nodes extracted from conversation let mut mem_children: Vec = Vec::new(); let mut scored = 0usize; let mut unscored = 0usize; - for ce in ag.context.conversation.entries() { - if let ConversationEntry::Memory { key, score, .. } = &ce.entry { - let status = match score { - Some(s) => { scored += 1; format!("score: {:.2}", s) } - None => { unscored += 1; String::new() } - }; - mem_children.push(SectionView { - name: key.clone(), - tokens: ce.tokens(), - content: ce.entry.message().content_text().to_string(), - children: Vec::new(), - status, - }); + for node in ag.context.conversation() { + if let AstNode::Leaf(leaf) = node { + if let NodeBody::Memory { key, score, text } = leaf.body() { + let status = match score { + Some(s) => { scored += 1; format!("score: {:.2}", s) } + None => { unscored += 1; String::new() } + }; + mem_children.push(SectionView { + name: key.clone(), + tokens: node.tokens(), + content: text.clone(), + children: Vec::new(), + status, + }); + } } } if !mem_children.is_empty() { @@ -68,9 +63,7 @@ impl ConsciousScreen { }); } - // Conversation — each entry as a child - views.push(section_to_view(&ag.context.conversation)); - + views.push(section_to_view("Conversation", ag.context.conversation())); views } } @@ -88,7 +81,6 @@ impl ScreenView for ConsciousScreen { } } - // Draw let mut lines: Vec = Vec::new(); let section_style = Style::default().fg(Color::Yellow); @@ -105,9 +97,7 @@ impl ScreenView for ConsciousScreen { lines.push(Line::styled("── Context State ──", section_style)); lines.push(Line::raw(format!(" Prompt tokens: {}K", app.status.prompt_tokens / 1000))); - if !app.status.context_budget.is_empty() { - lines.push(Line::raw(format!(" Budget: {}", app.status.context_budget))); - } + let context_state = self.read_context_views(); if !context_state.is_empty() { let total: usize = context_state.iter().map(|s| s.tokens).sum(); diff --git a/src/user/widgets.rs b/src/user/widgets.rs index aca9b34..88b1236 100644 --- a/src/user/widgets.rs +++ b/src/user/widgets.rs @@ -8,10 +8,8 @@ use ratatui::{ Frame, crossterm::event::KeyCode, }; -use crate::agent::context::{ContextSection, ConversationEntry}; +use crate::agent::context::{AstNode, NodeBody, Ast}; -/// UI-only tree node for the section tree display. -/// Built from ContextSection data; not used for budgeting. #[derive(Debug, Clone)] pub struct SectionView { pub name: String, @@ -22,26 +20,30 @@ pub struct SectionView { pub status: String, } -/// Build a SectionView tree from a ContextSection. -/// Each entry becomes a child with label + expandable content. -pub fn section_to_view(section: &ContextSection) -> SectionView { - let children: Vec = section.entries().iter().map(|ce| { - let content = match &ce.entry { - ConversationEntry::Log(_) => String::new(), - ConversationEntry::Thinking(text) => text.clone(), - _ => ce.entry.message().content_text().to_string(), +pub fn section_to_view(name: &str, nodes: &[AstNode]) -> SectionView { + let children: Vec = nodes.iter().map(|node| { + let content = match node.leaf().map(|l| l.body()) { + Some(NodeBody::Log(_)) => String::new(), + Some(body) => body.text().to_string(), + None => node.children().iter() + .filter_map(|c| c.leaf()) + .filter(|l| matches!(l.body(), NodeBody::Content(_))) + .map(|l| l.body().text()) + .collect::>() + .join(""), }; SectionView { - name: ce.entry.label(), - tokens: ce.tokens(), + name: node.label(), + tokens: node.tokens(), content, children: Vec::new(), status: String::new(), } }).collect(); + let total_tokens: usize = nodes.iter().map(|n| n.tokens()).sum(); SectionView { - name: section.name.clone(), - tokens: section.tokens(), + name: name.to_string(), + tokens: total_tokens, content: String::new(), children, status: String::new(),