gbnf fix and refactor

crates/gbnf-rs/AGENTS.md

@@ -21,4 +21,9 @@ Generate GBNF grammar ASTs from `schemars::Schema` inputs.
 ## Constraints
 - Generated object rules include all properties
   - required properties appear first
-  - optional properties may follow in any order
+  - optional properties follow in fixed, deterministic order
+    - each optional property appears at most once (`?` semantics)
+    - when requireds exist, each optional is prefixed by a comma if present
+    - when no requireds exist, the first present optional has no leading comma
+    - when no requireds exist, optional properties can start at any position and may skip middle properties while preserving order and uniqueness
+      - implemented via recursive choices over downstream chains (O(n^2) size)

crates/gbnf-rs/src/lib.rs

@@ -258,35 +258,85 @@ impl Generator {
             optional_props.sort_by(|a, b| a.0.cmp(b.0));
 
             if !optional_props.is_empty() {
-                let mut choices = Vec::new();
-                for (name, subschema) in optional_props {
-                    let subschema: Schema = subschema.clone().try_into().unwrap_or_default();
-                    let expr = self.expr_from_schema(&subschema, defs, cache, grammar);
-                    let expr = self.maybe_rule(expr, grammar);
-                    choices.push(Expr::Seq(vec![
+                // Build optional property expressions in fixed order.
+                // When there are required fields already emitted (first == false),
+                // each optional field, if present, must be preceded by a comma.
+                // When no required fields exist (first == true), we construct a
+                // nested optional chain like: opt1 ("," opt2 ("," opt3)?)?
+                // and wrap the whole chain in Optional(...), allowing an empty object.
+
+                // Precompute value expressions for each optional prop so we reuse
+                // the same rule/reference across multiple chains and avoid duplicates.
+                let prepared: Vec<(String, Expr)> = optional_props
+                    .iter()
+                    .map(|(name, subschema)| {
+                        let subschema: Schema = (*subschema).clone().try_into().unwrap_or_default();
+                        let expr = self.expr_from_schema(&subschema, defs, cache, grammar);
+                        let expr = self.maybe_rule(expr, grammar);
+                        ((*name).clone(), expr)
+                    })
+                    .collect();
+
+                // Helper to build a single property expr (without any leading comma)
+                let mk_prop_idx = |idx: usize| {
+                    let (ref name, ref expr) = prepared[idx];
+                    Expr::Seq(vec![
                         Expr::Literal(format!("\"{}\"", name)),
                         Expr::Ref("ws".into()),
                         Expr::Literal(":".into()),
                         Expr::Ref("ws".into()),
-                        expr,
-                    ]));
-                }
+                        expr.clone(),
+                    ])
+                };
 
                 if first {
-                    seq.push(Expr::Optional(Box::new(Expr::Seq(vec![
-                        Expr::Choice(choices.clone()),
-                        Expr::Repeat(Box::new(Expr::Seq(vec![
+                    // No requireds: allow starting at any optional and skipping
+                    // middles while preserving order and uniqueness.
+                    // Build chain_i for each i: prop(i) followed by optional
+                    // comma + choice of any chain_j where j > i.
+                    let n = prepared.len();
+                    if n > 0 {
+                        let mut chains: Vec<Expr> = vec![Expr::Literal(String::new()); n];
+                        // Build from the end to the front so later chains exist.
+                        for i in (0..n).rev() {
+                            let prop = mk_prop_idx(i);
+                            let mut choices: Vec<Expr> = Vec::new();
+                            for j in (i + 1)..n {
+                                choices.push(chains[j].clone());
+                            }
+                            let chain_i = if choices.is_empty() {
+                                prop
+                            } else {
+                                Expr::Seq(vec![
+                                    prop,
+                                    Expr::Optional(Box::new(Expr::Seq(vec![
+                                        Expr::Ref("ws".into()),
+                                        Expr::Literal(",".into()),
+                                        Expr::Ref("ws".into()),
+                                        Expr::Choice(choices),
+                                    ]))),
+                                ])
+                            };
+                            chains[i] = chain_i;
+                        }
+                        let top_choice = if n == 1 {
+                            chains.remove(0)
+                        } else {
+                            Expr::Choice(chains)
+                        };
+                        seq.push(Expr::Optional(Box::new(top_choice)));
+                    }
+                } else {
+                    // Requireds present: add each optional as an independent ("," prop)?
+                    for (idx, _) in prepared.iter().enumerate() {
+                        let prop = mk_prop_idx(idx);
+                        seq.push(Expr::Optional(Box::new(Expr::Seq(vec![
+                            Expr::Ref("ws".into()),
                             Expr::Literal(",".into()),
                             Expr::Ref("ws".into()),
-                            Expr::Choice(choices),
-                        ]))),
-                    ]))));
-                } else {
-                    seq.push(Expr::Repeat(Box::new(Expr::Seq(vec![
-                        Expr::Literal(",".into()),
-                        Expr::Ref("ws".into()),
-                        Expr::Choice(choices),
-                    ]))));
+                            prop,
+                        ]))));
+                    }
                 }
             }
         }
@@ -346,7 +396,7 @@ mod tests {
         let generator = Generator::new();
         let grammar = generator.generate("params", &schema);
         insta::assert_snapshot!(grammar.to_string(), @r###"
-params-r0 ::= "{" ws "\"file_path\"" ws ":" ws string "," ws "\"new_string\"" ws ":" ws string "," ws "\"old_string\"" ws ":" ws string ("," ws "\"expected_replacements\"" ws ":" ws number)* ws "}"
+params-r0 ::= "{" ws "\"file_path\"" ws ":" ws string "," ws "\"new_string\"" ws ":" ws string "," ws "\"old_string\"" ws ":" ws string (ws "," ws "\"expected_replacements\"" ws ":" ws number)? ws "}"
 "###);
     }
 
@@ -364,13 +414,32 @@ params-r0 ::= "{" ws "\"file_path\"" ws ":" ws string "," ws "\"new_string\"" ws
         let generator = Generator::new();
         let grammar = generator.generate("params", &schema);
         insta::assert_snapshot!(grammar.to_string(), @r###"
-params-r0 ::= "{" ws "\"path\"" ws ":" ws string ("," ws "\"ignore\"" ws ":" ws r1 | "\"include\"" ws ":" ws r2 | "\"include_hidden\"" ws ":" ws r3)* ws "}"
+params-r0 ::= "{" ws "\"path\"" ws ":" ws string (ws "," ws "\"ignore\"" ws ":" ws r1)? (ws "," ws "\"include\"" ws ":" ws r2)? (ws "," ws "\"include_hidden\"" ws ":" ws r3)? ws "}"
 r1 ::= "[" ws (string ("," ws string)*)? ws "]"
 r2 ::= "[" ws (string ("," ws string)*)? ws "]"
 r3 ::= "true" | "false"
 "###);
     }
 
+    #[derive(JsonSchema)]
+    struct AllOptional {
+        a: Option<String>,
+        b: Option<bool>,
+        c: Option<Vec<String>>,
+    }
+
+    #[test]
+    fn all_optional_object_grammar() {
+        let schema = schemars::schema_for!(AllOptional);
+        let generator = Generator::new();
+        let grammar = generator.generate("params", &schema);
+        insta::assert_snapshot!(grammar.to_string(), @r###"
+params-r0 ::= "{" ws ("\"a\"" ws ":" ws string (ws "," ws "\"b\"" ws ":" ws r1 (ws "," ws "\"c\"" ws ":" ws r2)? | "\"c\"" ws ":" ws r2)? | "\"b\"" ws ":" ws r1 (ws "," ws "\"c\"" ws ":" ws r2)? | "\"c\"" ws ":" ws r2)? ws "}"
+r1 ::= "true" | "false"
+r2 ::= "[" ws (string ("," ws string)*)? ws "]"
+"###);
+    }
+
     #[derive(JsonSchema)]
     struct Item {
         value: String,

crates/llm/AGENTS.md

@@ -91,10 +91,10 @@ Trait-based LLM client implementations for multiple providers.
     - implements `ToolExecutor` for MCP calls
       - unrecognized or unprefixed tool names return "{name} is not a valid tool name"
       - tool execution errors with `is_error` set propagate as `Err`
-    - tool call chunks insert assistant messages immediately before execution
-      - accumulated assistant content is flushed as a separate assistant message
-      - the tool-call assistant message carries empty `content` with `tool_calls` populated
-    - accumulated streamed content is appended as an assistant message after the stream completes
+    - tool orchestration message shape
+      - streamed assistant parts (text/thinking) and tool calls from a single response are combined into one assistant message in original order
+      - the tool result is appended as a separate tool-role message with the same call id
+      - follow-up assistant responses are appended after tool execution completes
 - Test utilities
   - `TestProvider` implements `LlmClient`
     - captures `ChatMessageRequest`s for assertions

crates/llm/src/tools.rs

@@ -14,7 +14,7 @@ use tokio_stream::{Stream, StreamExt};
 
 use crate::{
     AssistantMessage, AssistantPart, ChatMessage, ChatMessageRequest, JsonResult, LlmClient,
-    ResponseChunk, ToolCall,
+    ResponseChunk,
 };
 
 #[async_trait]
@@ -69,100 +69,84 @@ pub async fn run_tool_loop(
         tx.send(ToolEvent::RequestStarted).ok();
         let mut handles: JoinSet<(String, String, Result<String, Box<dyn Error + Send + Sync>>)> =
             JoinSet::new();
-        let mut assistant_content: Option<String> = None;
-        let mut assistant_thinking: Option<String> = None;
+        let mut parts: Vec<AssistantPart> = Vec::new();
+        let mut current_part: Option<AssistantPart> = None;
         while let Some(chunk) = stream.next().await {
             let chunk = chunk?;
             let mut done = false;
-            let mut tool_calls: Vec<ToolCall> = Vec::new();
             match &chunk {
                 ResponseChunk::Content(content) => {
-                    if !content.is_empty() {
-                        assistant_content
-                            .get_or_insert_with(String::new)
-                            .push_str(content);
+                    if let Some(AssistantPart::Text { text }) = current_part.as_mut() {
+                        text.push_str(content);
+                    } else {
+                        if let Some(part) = current_part.take() {
+                            parts.push(part);
+                        }
+                        current_part = Some(AssistantPart::Text {
+                            text: content.into(),
+                        });
                     }
                 }
                 ResponseChunk::Thinking(thinking) => {
-                    if !thinking.is_empty() {
-                        assistant_thinking
-                            .get_or_insert_with(String::new)
-                            .push_str(thinking);
+                    if let Some(AssistantPart::Thinking { text }) = current_part.as_mut() {
+                        text.push_str(thinking);
+                    } else {
+                        if let Some(part) = current_part.take() {
+                            parts.push(part);
+                        }
+                        current_part = Some(AssistantPart::Thinking {
+                            text: thinking.into(),
+                        });
                     }
                 }
                 ResponseChunk::ToolCall(tc) => {
-                    tool_calls.push(tc.clone());
+                    if let Some(part) = current_part.take() {
+                        parts.push(part);
+                    }
+                    current_part = Some(AssistantPart::ToolCall(tc.clone()));
+                    tx.send(ToolEvent::ToolStarted {
+                        call_id: tc.id.clone(),
+                        name: tc.name.clone(),
+                        args: tc.arguments.clone(),
+                    })
+                    .ok();
+                    let executor = tool_executor.clone();
+                    let name = tc.name.clone();
+                    let args = tc.arguments.clone();
+                    let call_id = tc.id.clone();
+                    handles.spawn(async move {
+                        match args {
+                            JsonResult::Content { content } => {
+                                let res = executor.call(&name, content).await;
+                                (call_id, name, res)
+                            }
+                            JsonResult::Error { .. } => (
+                                call_id,
+                                name,
+                                Err::<String, Box<dyn Error + Send + Sync>>(Box::new(
+                                    std::io::Error::new(
+                                        std::io::ErrorKind::Other,
+                                        "Could not parse arguments as JSON",
+                                    ),
+                                )),
+                            ),
+                        }
+                    });
                 }
                 ResponseChunk::Usage { .. } => {}
                 ResponseChunk::Done => {
                     done = true;
                 }
             }
-            if !tool_calls.is_empty() {
-                if let Some(content) = assistant_content.take() {
-                    chat_history
-                        .lock()
-                        .unwrap()
-                        .push(ChatMessage::Assistant(AssistantMessage {
-                            content: vec![AssistantPart::Text { text: content }],
-                        }));
-                }
-                let mut parts: Vec<AssistantPart> = tool_calls
-                    .clone()
-                    .into_iter()
-                    .map(AssistantPart::ToolCall)
-                    .collect();
-                if let Some(thinking) = assistant_thinking.take() {
-                    parts.insert(0, AssistantPart::Thinking { text: thinking });
-                }
-                chat_history
-                    .lock()
-                    .unwrap()
-                    .push(ChatMessage::Assistant(AssistantMessage { content: parts }));
-            }
             tx.send(ToolEvent::Chunk(chunk)).ok();
-            for call in tool_calls {
-                tx.send(ToolEvent::ToolStarted {
-                    call_id: call.id.clone(),
-                    name: call.name.clone(),
-                    args: call.arguments.clone(),
-                })
-                .ok();
-                let executor = tool_executor.clone();
-                let name = call.name.clone();
-                let args = call.arguments.clone();
-                let call_id = call.id.clone();
-                handles.spawn(async move {
-                    match args {
-                        JsonResult::Content { content } => {
-                            let res = executor.call(&name, content).await;
-                            (call_id, name, res)
-                        }
-                        JsonResult::Error { .. } => (
-                            call_id,
-                            name,
-                            Err::<String, Box<dyn Error + Send + Sync>>(Box::new(
-                                std::io::Error::new(
-                                    std::io::ErrorKind::Other,
-                                    "Could not parse arguments as JSON",
-                                ),
-                            )),
-                        ),
-                    }
-                });
-            }
             if done {
                 break;
             }
         }
-        if assistant_content.is_some() || assistant_thinking.is_some() {
-            let mut parts: Vec<AssistantPart> = Vec::new();
-            if let Some(thinking) = assistant_thinking.take() {
-                parts.push(AssistantPart::Thinking { text: thinking });
-            }
-            if let Some(content) = assistant_content.take() {
-                parts.push(AssistantPart::Text { text: content });
-            }
+        if let Some(part) = current_part.take() {
+            parts.push(part);
+        }
+        if !parts.is_empty() {
             chat_history
                 .lock()
                 .unwrap()
@@ -280,30 +264,25 @@ mod tests {
             .await
             .unwrap();
         let updated = history.lock().unwrap().clone();
-        // Behavior: assistant preamble content, then separate tool-call message, tool result, and final assistant response
-        assert_eq!(updated.len(), 5);
-        // First assistant message should contain the preamble content with no tool calls
-        let preamble_msg = &updated[1];
-        if let ChatMessage::Assistant(a) = preamble_msg {
-            assert_eq!(a.content.len(), 1);
+        // Behavior: assistant content and tool call are combined in one assistant message,
+        // followed by the tool result message and the final assistant response
+        assert_eq!(updated.len(), 4);
+        // First assistant message should contain the preamble content followed by the tool call
+        let first_assistant = &updated[1];
+        if let ChatMessage::Assistant(a) = first_assistant {
+            assert_eq!(a.content.len(), 2);
             match &a.content[0] {
                 AssistantPart::Text { text } => assert_eq!(text, "first"),
-                _ => panic!("expected text part"),
+                _ => panic!("expected first part to be text"),
             }
-        } else {
-            panic!("expected assistant preamble message");
-        }
-        // Next assistant message should carry the tool call
-        let call_msg = &updated[2];
-        if let ChatMessage::Assistant(a) = call_msg {
-            assert_eq!(a.content.len(), 1);
-            match &a.content[0] {
+            match &a.content[1] {
                 AssistantPart::ToolCall(tc) => assert_eq!(tc.name, "test"),
-                _ => panic!("expected tool call part"),
+                _ => panic!("expected second part to be tool call"),
             }
         } else {
-            panic!("expected assistant message with tool call");
+            panic!("expected combined assistant message");
         }
+        // Final assistant message should be the follow-up content
         let final_msg = updated.last().unwrap();
         if let ChatMessage::Assistant(a) = final_msg {
             assert_eq!(a.content.len(), 1);