ResponseParser returns children incrementally, add push_child/PendingToolCall

feed() now returns all completed children (not just tool calls) so the
caller can push them into the AST as they arrive. finish() returns
remaining buffered children. The caller manages the assistant branch.

Added ContextState::push_child() for appending to an existing branch,
PendingToolCall for ephemeral dispatch handles, and len() for section
size queries.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>

src/agent/context_new.rs

@@ -96,6 +96,14 @@ pub enum Section {
     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>;
@@ -464,10 +472,9 @@ impl ResponseParser {
         }
     }
 
-    /// Feed a text chunk. Returns newly completed tool call nodes
-    /// (for immediate dispatch).
+    /// Feed a text chunk. Returns completed child nodes — the caller
+    /// pushes them into the assistant branch and dispatches any tool calls.
     pub fn feed(&mut self, text: &str) -> Vec<AstNode> {
-        let mut new_calls = vec![];
         self.buf.push_str(text);
 
         loop {
@@ -482,7 +489,6 @@ impl ResponseParser {
                         continue;
                     }
                     None => {
-                        // Keep last 8 chars ("</think>".len()) as lookahead
                         let safe = self.buf.len().saturating_sub(8);
                         if safe > 0 {
                             let safe = self.buf.floor_char_boundary(safe);
@@ -501,16 +507,13 @@ impl ResponseParser {
                         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) {
-                            let node = AstNode::tool_call(name, args);
-                            new_calls.push(node.clone());
                             self.flush_content();
-                            self.children.push(node);
+                            self.children.push(AstNode::tool_call(name, args));
                         }
                         self.tool_call_buf.clear();
                         continue;
                     }
                     None => {
-                        // Keep last 12 chars ("</tool_call>".len()) as lookahead
                         let safe = self.buf.len().saturating_sub(12);
                         if safe > 0 {
                             let safe = self.buf.floor_char_boundary(safe);
@@ -559,7 +562,7 @@ impl ResponseParser {
             }
         }
 
-        new_calls
+        self.children.drain(..).collect()
     }
 
     fn flush_content(&mut self) {
@@ -571,16 +574,16 @@ impl ResponseParser {
         }
     }
 
-    /// Finalize the parse. Returns the completed assistant AstNode.
-    pub fn finish(mut self) -> AstNode {
+    /// Flush remaining buffer and return any final children.
+    pub fn finish(mut self) -> Vec<AstNode> {
         if !self.buf.is_empty() {
             self.content_parts.push(std::mem::take(&mut self.buf));
         }
         self.flush_content();
-        AstNode::branch(Role::Assistant, self.children)
+        self.children
     }
 
-    /// Get the current display text (for streaming to UI).
+    /// Current display text (content accumulated since last drain).
     pub fn display_content(&self) -> String {
         self.content_parts.join("")
     }
@@ -686,6 +689,25 @@ impl ContextState {
     pub fn del(&mut self, section: Section, index: usize) -> AstNode {
         self.section_mut(section).remove(index)
     }
+
+    /// 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 {
@@ -719,13 +741,8 @@ mod tests {
 
     // -- Helpers for inspecting parse results ----------------------------------
 
-    /// Extract child bodies from an Assistant branch node.
-    fn child_bodies(node: &AstNode) -> Vec<&NodeBody> {
-        match node {
-            AstNode::Branch { children, .. } =>
-                children.iter().filter_map(|c| c.leaf()).map(|l| l.body()).collect(),
-            _ => panic!("expected branch"),
-        }
+    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) {
@@ -795,36 +812,39 @@ mod tests {
 
     // -- ResponseParser tests -------------------------------------------------
 
+    /// Collect all children from feed + finish.
+    fn parse_all(text: &str) -> Vec<AstNode> {
+        let mut p = ResponseParser::new();
+        let mut all = p.feed(text);
+        all.extend(p.finish());
+        all
+    }
+
     #[test]
     fn test_parser_plain_text() {
-        let mut p = ResponseParser::new();
-        p.feed("hello world");
-        let node = p.finish();
-        let bodies = child_bodies(&node);
-        assert_eq!(bodies.len(), 1);
-        assert_content(bodies[0], "hello world");
+        let nodes = parse_all("hello world");
+        let b = bodies(&nodes);
+        assert_eq!(b.len(), 1);
+        assert_content(b[0], "hello world");
     }
 
     #[test]
     fn test_parser_thinking_then_content() {
-        let mut p = ResponseParser::new();
-        p.feed("<think>reasoning</think>answer");
-        let node = p.finish();
-        let bodies = child_bodies(&node);
-        assert_eq!(bodies.len(), 2);
-        assert_thinking(bodies[0], "reasoning");
-        assert_content(bodies[1], "answer");
+        let nodes = parse_all("<think>reasoning</think>answer");
+        let b = bodies(&nodes);
+        assert_eq!(b.len(), 2);
+        assert_thinking(b[0], "reasoning");
+        assert_content(b[1], "answer");
     }
 
     #[test]
     fn test_parser_tool_call() {
         let mut p = ResponseParser::new();
-        let calls = p.feed("<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>");
-        assert_eq!(calls.len(), 1);  // returned for immediate dispatch
-        let node = p.finish();
-        let bodies = child_bodies(&node);
-        assert_eq!(bodies.len(), 1);
-        let args = assert_tool_call(bodies[0], "bash");
+        let children = p.feed("<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>");
+        // Tool call returned immediately from feed
+        assert_eq!(children.len(), 1);
+        let b = bodies(&children);
+        let args = assert_tool_call(b[0], "bash");
         let args: serde_json::Value = serde_json::from_str(args).unwrap();
         assert_eq!(args["command"], "ls");
     }
@@ -832,72 +852,70 @@ mod tests {
     #[test]
     fn test_parser_content_then_tool_call_then_content() {
         let mut p = ResponseParser::new();
-        p.feed("before");
-        p.feed("<tool_call>\n<function=bash>\n<parameter=command>pwd</parameter>\n</function>\n</tool_call>");
-        p.feed("after");
-        let node = p.finish();
-        let bodies = child_bodies(&node);
-        assert_eq!(bodies.len(), 3);
-        assert_content(bodies[0], "before");
-        assert_tool_call(bodies[1], "bash");
-        assert_content(bodies[2], "after");
+        let mut all = p.feed("before");
+        all.extend(p.feed("<tool_call>\n<function=bash>\n<parameter=command>pwd</parameter>\n</function>\n</tool_call>"));
+        all.extend(p.feed("after"));
+        all.extend(p.finish());
+        let b = bodies(&all);
+        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() {
-        // Feed the response one character at a time
         let text = "<think>thought</think>response";
         let mut p = ResponseParser::new();
+        let mut all = Vec::new();
         for ch in text.chars() {
-            p.feed(&ch.to_string());
+            all.extend(p.feed(&ch.to_string()));
         }
-        let node = p.finish();
-        let bodies = child_bodies(&node);
-        assert_eq!(bodies.len(), 2);
-        assert_thinking(bodies[0], "thought");
-        assert_content(bodies[1], "response");
+        all.extend(p.finish());
+        let b = bodies(&all);
+        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 p = ResponseParser::new();
-        let mut total_calls = 0;
+        let mut all = Vec::new();
+        let mut tool_calls = 0;
         for ch in text.chars() {
-            total_calls += p.feed(&ch.to_string()).len();
+            let children = p.feed(&ch.to_string());
+            for c in &children {
+                if let AstNode::Leaf(l) = c {
+                    if matches!(l.body(), NodeBody::ToolCall { .. }) {
+                        tool_calls += 1;
+                    }
+                }
+            }
+            all.extend(children);
         }
-        assert_eq!(total_calls, 1);  // exactly one tool call dispatched
-        let node = p.finish();
-        let bodies = child_bodies(&node);
-        assert_eq!(bodies.len(), 3);
-        assert_content(bodies[0], "text");
-        assert_tool_call(bodies[1], "bash");
-        assert_content(bodies[2], "more");
+        all.extend(p.finish());
+        assert_eq!(tool_calls, 1);
+        let b = bodies(&all);
+        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 mut p = ResponseParser::new();
-        p.feed("<think>let me think</think>");
-        p.feed("<tool_call>\n<function=read>\n<parameter=path>/etc/hosts</parameter>\n</function>\n</tool_call>");
-        p.feed("here's what I found");
-        let node = p.finish();
-        let bodies = child_bodies(&node);
-        assert_eq!(bodies.len(), 3);
-        assert_thinking(bodies[0], "let me think");
-        assert_tool_call(bodies[1], "read");
-        assert_content(bodies[2], "here's what I found");
-    }
-
-    #[test]
-    fn test_parser_finish_produces_assistant_branch() {
-        let mut p = ResponseParser::new();
-        p.feed("hello");
-        let node = p.finish();
-        match &node {
-            AstNode::Branch { role, .. } => assert_eq!(*role, Role::Assistant),
-            _ => panic!("expected branch"),
-        }
+        let mut all = p.feed("<think>let me think</think>");
+        all.extend(p.feed("<tool_call>\n<function=read>\n<parameter=path>/etc/hosts</parameter>\n</function>\n</tool_call>"));
+        all.extend(p.feed("here's what I found"));
+        all.extend(p.finish());
+        let b = bodies(&all);
+        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 -------------------------------------------
@@ -1032,10 +1050,12 @@ mod tests {
         if !init_tokenizer() { return; }
 
         let mut p = ResponseParser::new();
-        p.feed("I'll check that for you");
-        p.feed("<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>");
-        let node = p.finish();
+        let mut all = p.feed("I'll check that for you");
+        all.extend(p.feed("<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>"));
+        all.extend(p.finish());
 
+        // Wrap in assistant branch to test full tokenization
+        let node = AstNode::branch(Role::Assistant, all);
         assert_token_invariants(&node);
         assert!(node.tokens() > 0);
     }