kent/consciousness
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

ResponseParser mutates AST directly, returns PendingToolCalls

Browse source 
The parser takes &mut ContextState on feed()/finish() and pushes
completed children (content, thinking, tool calls) directly into
the assistant branch. Only PendingToolCall handles are returned
to the caller for dispatch — the caller no longer manages AST
mutation.

Tests verify by reading back from ContextState after parsing.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-08 14:33:57 -04:00
parent 6139d43942
commit 648356ae40
1 changed files with 89 additions and 71 deletions
Download patch file Download diff file Expand all files Collapse all files

160
src/agent/context_new.rs
View file

@ -111,9 +111,10 @@ pub trait Ast {
}
pub struct ResponseParser {
branch_idx: usize,
call_counter: u32,
buf: String,
content_parts: Vec<String>,
children: Vec<AstNode>,
in_think: bool,
think_buf: String,
in_tool_call: bool,
@ -460,11 +461,14 @@ fn parse_json_tool_call(body: &str) -> Option<(String, String)> {
}
impl ResponseParser {
pub fn new() -> Self {
/// 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(),
children: Vec::new(),
in_think: false,
think_buf: String::new(),
in_tool_call: false,
@ -472,9 +476,10 @@ impl ResponseParser {
}
}
/// 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> {
/// 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<PendingToolCall> {
let mut pending = Vec::new();
self.buf.push_str(text);
loop {
@ -484,7 +489,7 @@ impl ResponseParser {
self.think_buf.push_str(&self.buf[..end]);
self.buf = self.buf[end + 8..].to_string();
self.in_think = false;
self.children.push(AstNode::thinking(&self.think_buf));
self.push_child(ctx, AstNode::thinking(&self.think_buf));
self.think_buf.clear();
continue;
}
@ -507,8 +512,14 @@ 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) {
self.flush_content();
self.children.push(AstNode::tool_call(name, args));
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;
@ -541,11 +552,11 @@ impl ResponseParser {
}
if self.buf[pos..].starts_with("<think>") {
self.buf = self.buf[pos + 7..].to_string();
self.flush_content();
self.flush_content(ctx);
self.in_think = true;
} else {
self.buf = self.buf[pos + 11..].to_string();
self.flush_content();
self.flush_content(ctx);
self.in_tool_call = true;
}
continue;
@ -562,25 +573,28 @@ impl ResponseParser {
}
}
self.children.drain(..).collect()
pending
}
fn flush_content(&mut self) {
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.children.push(AstNode::content(text));
self.push_child(ctx, AstNode::content(text));
}
}
}
/// Flush remaining buffer and return any final children.
pub fn finish(mut self) -> Vec<AstNode> {
/// 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();
self.children
self.flush_content(ctx);
}
/// Current display text (content accumulated since last drain).
@ -812,26 +826,36 @@ 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
/// Set up a ContextState with an assistant branch, run the parser,
/// return the children that were pushed into the branch.
fn parse_into_ctx(chunks: &[&str]) -> (ContextState, Vec<PendingToolCall>) {
let mut ctx = ContextState::new();
ctx.push(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 nodes = parse_all("hello world");
let b = bodies(&nodes);
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 nodes = parse_all("<think>reasoning</think>answer");
let b = bodies(&nodes);
let (ctx, _) = parse_into_ctx(&["<think>reasoning</think>answer"]);
let b = bodies(assistant_children(&ctx));
assert_eq!(b.len(), 2);
assert_thinking(b[0], "reasoning");
assert_content(b[1], "answer");
@ -839,11 +863,13 @@ mod tests {
#[test]
fn test_parser_tool_call() {
let mut p = ResponseParser::new();
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 (ctx, calls) = parse_into_ctx(&[
"<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>"
]);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "bash");
let b = bodies(assistant_children(&ctx));
assert_eq!(b.len(), 1);
let args = assert_tool_call(b[0], "bash");
let args: serde_json::Value = serde_json::from_str(args).unwrap();
assert_eq!(args["command"], "ls");
@ -851,12 +877,12 @@ mod tests {
#[test]
fn test_parser_content_then_tool_call_then_content() {
let mut p = ResponseParser::new();
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);
let (ctx, _) = parse_into_ctx(&[
"before",
"<tool_call>\n<function=bash>\n<parameter=command>pwd</parameter>\n</function>\n</tool_call>",
"after",
]);
let b = bodies(assistant_children(&ctx));
assert_eq!(b.len(), 3);
assert_content(b[0], "before");
assert_tool_call(b[1], "bash");
@ -866,13 +892,14 @@ mod tests {
#[test]
fn test_parser_incremental_feed() {
let text = "<think>thought</think>response";
let mut p = ResponseParser::new();
let mut all = Vec::new();
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() {
all.extend(p.feed(&ch.to_string()));
p.feed(&ch.to_string(), &mut ctx);
}
all.extend(p.finish());
let b = bodies(&all);
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");
@ -881,23 +908,16 @@ mod tests {
#[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 all = Vec::new();
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() {
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;
tool_calls += p.feed(&ch.to_string(), &mut ctx).len();
}
}
}
all.extend(children);
}
all.extend(p.finish());
p.finish(&mut ctx);
assert_eq!(tool_calls, 1);
let b = bodies(&all);
let b = bodies(assistant_children(&ctx));
assert_eq!(b.len(), 3);
assert_content(b[0], "text");
assert_tool_call(b[1], "bash");
@ -906,12 +926,12 @@ mod tests {
#[test]
fn test_parser_thinking_tool_call_content() {
let mut p = ResponseParser::new();
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);
let (ctx, _) = parse_into_ctx(&[
"<think>let me think</think>",
"<tool_call>\n<function=read>\n<parameter=path>/etc/hosts</parameter>\n</function>\n</tool_call>",
"here's what I found",
]);
let b = bodies(assistant_children(&ctx));
assert_eq!(b.len(), 3);
assert_thinking(b[0], "let me think");
assert_tool_call(b[1], "read");
@ -1049,14 +1069,12 @@ mod tests {
fn test_parser_roundtrip_through_tokenizer() {
if !init_tokenizer() { return; }
let mut p = ResponseParser::new();
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);
let (ctx, _) = parse_into_ctx(&[
"I'll check that for you",
"<tool_call>\n<function=bash>\n<parameter=command>ls</parameter>\n</function>\n</tool_call>",
]);
let node = &ctx.conversation()[0];
assert_token_invariants(node);
assert!(node.tokens() > 0);
}
}