chore: JailStream Optimizations

[ayushag-nv]

Overview:

Details:

Where should the reviewer start?

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: #xxx

Summary by CodeRabbit

• New Features

  • Unified end-position detection for tool calls across parsers, improving streaming accuracy.
  • Automatic detection of jail/start/end markers from configuration for more reliable streaming.

• Bug Fixes

  • More robust start/end handling of tool-call sections, including single-chunk completions.
  • Consistent emission of pass-through and trailing content.
  • Reduced stray content by clearing buffers reliably after matches.

• Refactor

  • Tool-call parsing APIs are now synchronous, simplifying integration and reducing overhead.
  • Centralized and standardized stream emission logic for clearer, more predictable behavior.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

Refactors OpenAI chat_completions jail streaming to centralize choice emission and tool-call handling, introduces parser end-position utilities, and converts multiple tool-calling parser APIs from async to sync. Adds dispatch for end-position detection across parsers and updates builders to derive jail markers from parser config.

Changes

Cohort / File(s)	Summary of changes
OpenAI jail streaming refactor lib/llm/src/protocols/openai/chat_completions/jail.rs	Adds helper create_choice_stream and ChoiceJailState emission helpers; refactors content processing to use helpers; standardizes tool-call vs content emission; clears partial buffers consistently; removes ad-hoc end logic; imports ChatChoiceLogprobs and unified end-position helper; builder auto-populates jail start/end sequences from parser config and common markers.
Harmony parser sync conversion lib/parsers/src/tool_calling/harmony/harmony_parser.rs	Converts async functions to sync; replaces tokio::sync::OnceCell with std::sync::OnceLock; updates encoding initialization path; adjusts tests to non-async.
End-position helpers (Harmony) lib/parsers/src/tool_calling/harmony/mod.rs	Adds find_tool_call_end_position_harmony(chunk: &str) -> usize returning chunk.len().
End-position helpers (JSON) lib/parsers/src/tool_calling/json/mod.rs	Adds find_tool_call_end_position_json(chunk, parser, config) with parser-specific logic for end detection (hermes/nemotron_deci via end tokens; mistral/phi4 via last ]; default to chunk.len()).
End-position helpers (Pythonic) lib/parsers/src/tool_calling/pythonic/mod.rs	Adds find_tool_call_end_position_pythonic(chunk: &str) -> usize returning chunk.len().
Parser dispatch: sync + end-position lib/parsers/src/tool_calling/parsers.rs	Adds find_tool_call_end_position(chunk, parser_str) dispatcher; converts try_tool_call_parse and detect_and_parse_tool_call to sync; wires Harmony/JSON/Pythonic end-position helpers; maintains unsupported branches for Typescript/XML.
Top-level parser re-exports lib/parsers/src/tool_calling/mod.rs	Re-exports find_tool_call_end_position publicly.
Tool-call utilities sync conversion lib/parsers/src/tool_calling/tools.rs	Converts try_tool_call_parse_aggregate and try_tool_call_parse_stream to sync and removes awaits on detection/parse.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor Client
  participant Stream as OpenAI Stream
  participant Jail as JailedStream
  participant Parser as ToolCall Parsers
  participant Emit as Choice Emission

  Client->>Stream: receive chunk
  Stream->>Jail: process_content(chunk)
  Jail->>Parser: find_tool_call_end_position(chunk, parser_str)
  Parser-->>Jail: end_index
  alt Tool-call markers detected
    Jail->>Jail: accumulate tool_calls / trailing
    Jail->>Emit: emit_tool_or_content(...)
  else No markers
    Jail->>Emit: emit_pass_through or content
  end
  Emit-->>Client: ChatChoiceStream chunk(s)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• feat: parse normal text along with tool calls #2709 — Introduces/adjusts tool-call end-position APIs and re-exports that align with the new find_tool_call_end_position dispatch used in jail processing.
• feat: Tool calling support stream=True #2932 — Modifies jail/tool-call streaming integration; closely related to this refactor that standardizes emissions and end detection.
• feat: added support for pythonic tool parser #2788 — Expands tool-calling parser modules and wiring, overlapping with added parser end-position helpers and public re-exports.

Poem

In streams I nibble, byte by byte,
Markers blink like stars at night.
Tools or tales? I sniff the end—
Then hop along the flow I send.
Sync now swift, my whiskers preen—
A tidy warren, fast and clean. 🐇✨

Pre-merge checks

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description Check	⚠️ Warning	The pull request description contains only the repository template headings with placeholder text and lacks any substantive summary, list of changes, files to review, rationale, or testing information, so it does not meet the repository's required description content.	Populate the Overview with a brief summary of intent and high-level changes, fill Details with the specific file changes and behavioral impact, list Where should the reviewer start with key files and functions to inspect, and add Related Issues plus testing/integration notes and any migration or compatibility considerations.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The PR title "chore: JailStream Optimizations" is concise and directly describes the primary change (optimizations and refactors to the JailStream logic) and is specific enough for a teammate scanning history to understand the main intent.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 7

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

lib/parsers/src/tool_calling/tools.rs (1)

10-16: Remove .await from sync API callsites; lower parser log level to debug.

try_tool_call_parse_aggregate and try_tool_call_parse_stream are synchronous; awaiting them will fail to compile.

• Remove .await at: lib/llm/src/protocols/openai/chat_completions/jail.rs:641, 664-665, 723.
• Change tracing::info! → tracing::debug! in lib/parsers/src/tool_calling/tools.rs (around lines 19–21) to avoid noisy logs on hot paths.

lib/llm/src/protocols/openai/chat_completions/jail.rs (2)

639-645: Call the free function find_tool_call_end_position correctly

You're calling a non‑existent method on self and passing &String where Option<&str> is expected.

Apply this diff:

-                if let Ok((_, _)) =
-                    try_tool_call_parse_aggregate(accumulated_content, Some(parser)).await
+                if let Ok((_, _)) = try_tool_call_parse_aggregate(
+                    accumulated_content,
+                    Some(parser.as_str()),
+                )
+                .await
                 {
-                    let split_pos = self.find_tool_call_end_position(accumulated_content, parser);
+                    let split_pos = find_tool_call_end_position(
+                        accumulated_content,
+                        Some(parser.as_str()),
+                    );
                     (true, split_pos)
                 } else {

---

663-699: Avoid unstable let‑chains in if let; restructure for stable Rust

if let PATTERN = EXPR && guard is not stable. Use nested if.

Apply this diff:

-        if let Ok((tool_calls, normal_text)) =
-            try_tool_call_parse_aggregate(accumulated_content, self.tool_call_parser.as_deref())
-                .await
-            && !tool_calls.is_empty()
-        {
-            // Convert to streaming format
-            let tool_call_chunks: Vec<ChatCompletionMessageToolCallChunk> = tool_calls
-                .into_iter()
-                .enumerate()
-                .map(|(idx, tool_call)| ChatCompletionMessageToolCallChunk {
-                    index: idx as u32,
-                    id: Some(tool_call.id),
-                    r#type: Some(tool_call.r#type),
-                    function: Some(FunctionCallStream {
-                        name: Some(tool_call.function.name),
-                        arguments: Some(tool_call.function.arguments),
-                    }),
-                })
-                .collect();
-
-            // Create choice with tool calls
-            #[allow(deprecated)]
-            return ChatChoiceStream {
-                index: choice_index,
-                delta: ChatCompletionStreamResponseDelta {
-                    role: Some(Role::Assistant),
-                    content: normal_text.filter(|t| !t.is_empty()),
-                    tool_calls: Some(tool_call_chunks),
-                    function_call: None,
-                    refusal: None,
-                    reasoning_content: None,
-                },
-                finish_reason: Some(FinishReason::ToolCalls),
-                logprobs: None,
-            };
-        }
+        if let Ok((tool_calls, normal_text)) =
+            try_tool_call_parse_aggregate(accumulated_content, self.tool_call_parser.as_deref())
+                .await
+        {
+            if !tool_calls.is_empty() {
+                // Convert to streaming format
+                let tool_call_chunks: Vec<ChatCompletionMessageToolCallChunk> = tool_calls
+                    .into_iter()
+                    .enumerate()
+                    .map(|(idx, tool_call)| ChatCompletionMessageToolCallChunk {
+                        index: idx as u32,
+                        id: Some(tool_call.id),
+                        r#type: Some(tool_call.r#type),
+                        function: Some(FunctionCallStream {
+                            name: Some(tool_call.function.name),
+                            arguments: Some(tool_call.function.arguments),
+                        }),
+                    })
+                    .collect();
+
+                // Create choice with tool calls
+                #[allow(deprecated)]
+                return ChatChoiceStream {
+                    index: choice_index,
+                    delta: ChatCompletionStreamResponseDelta {
+                        role: Some(Role::Assistant),
+                        content: normal_text.filter(|t| !t.is_empty()),
+                        tool_calls: Some(tool_call_chunks),
+                        function_call: None,
+                        refusal: None,
+                        reasoning_content: None,
+                    },
+                    finish_reason: Some(FinishReason::ToolCalls),
+                    logprobs: None,
+                };
+            }
+        }

🧹 Nitpick comments (16)

lib/parsers/src/tool_calling/json/mod.rs (1)

45-71: Guard against empty end tokens; consider scanning all configured end tokens.

If the first configured end token is an empty string, chunk.find(end_token) returns 0, causing a premature end position. Also, when multiple non-empty end tokens exist, picking the earliest match is safer.

Apply:

-        "hermes" | "nemotron_deci" => {
-            if let Some(end_token) = config.tool_call_end_tokens.first() {
-                if let Some(pos) = chunk.find(end_token) {
-                    pos + end_token.len()
-                } else {
-                    chunk.len()
-                }
-            } else {
-                chunk.len()
-            }
-        }
+        "hermes" | "nemotron_deci" => {
+            let end_pos = config
+                .tool_call_end_tokens
+                .iter()
+                .filter(|t| !t.is_empty())
+                .filter_map(|t| chunk.find(t).map(|pos| pos + t.len()))
+                .min()
+                .unwrap_or(chunk.len());
+            end_pos
+        }

Please confirm ToolCallConfig::hermes() and ToolCallConfig::nemotron_deci() set a non-empty end token first (e.g., "</tool_call>", "</TOOLCALL>"), or accept the above defensive change.

lib/parsers/src/tool_calling/pythonic/mod.rs (1)

9-11: Return end of closing bracket if present.

Mirrors Mistral/Phi4 handling and avoids over-reading beyond a completed call list.

 pub fn find_tool_call_end_position_pythonic(chunk: &str) -> usize {
-    chunk.len()
+    if let Some(pos) = chunk.rfind(']') {
+        pos + 1
+    } else {
+        chunk.len()
+    }
 }

lib/parsers/src/tool_calling/harmony/mod.rs (1)

11-13: LGTM; parity with dispatcher achieved.

Returning chunk.len() is acceptable for Harmony given complete-parse behavior. Consider future extension to accept config if Harmony ever exposes end tokens.

lib/parsers/src/tool_calling/parsers.rs (2)

127-155: End-position dispatcher: sensible defaults; minor hardening optional.

• Unknown or unimplemented parsers return chunk.len() — sane.
• JSON path forwards parser_key to JSON-specific logic — good.

Add targeted tests for find_tool_call_end_position across: hermes, nemotron_deci, mistral, phi4 (including edge with trailing normal text), to lock behavior.

---

189-206: Tests now call sync APIs — fine; consider dropping tokio where unnecessary.

Async test harness remains but calls are sync; optional cleanup to #[test] for purely sync cases.

Also applies to: 202-219, 230-234, 245-250, 283-293, 299-306, 315-319, 336-337, 367-368, 386-387, 401-413, 431-433, 484-485, 511-512, 525-526, 539-540, 558-559, 572-573, 590-591, 616-617, 634-635, 668-669, 682-683, 701-702, 729-730, 746-747, 759-760, 775-776, 788-789, 801-802, 817-818, 835-836, 853-854, 864-871, 899-904, 939-947, 970-982, 987-995, 1001-1009, 1022-1030, 1040-1048, 1053-1061, 1081-1087, 1092-1098, 1106-1117, 1125-1136, 1143-1151, 1158-1166, 1172-1179, 1185-1192, 1199-1206, 1213-1218, 1224-1228, 1236-1266, 1274-1286, 1293-1309, 1325-1338, 1353-1366, 1372-1383, 1389-1401, 1407-1444

lib/parsers/src/tool_calling/harmony/harmony_parser.rs (10)

11-19: OnceLock switch LGTM; confirm encoding choice/name.

Good move to std::sync::OnceLock. Please confirm HarmonyEncodingName::HarmonyGptOss is correct for this generic harmony tool-calling parser. If this parser is used beyond GPT‑OSS, consider parameterizing the encoding or renaming the static to a model‑agnostic identifier to avoid confusion.

---

137-144: Avoid needless clones and unwrap when building ToolCallResponse.

Use the owned fname and Value::to_string(); serde_json::to_string on Value shouldn’t fail, and this trims one clone and an unwrap.

Apply this diff:

-                res.push(ToolCallResponse {
-                    id: format!("call-{}", call_idx),
-                    tp: ToolCallType::Function,
-                    function: CalledFunction {
-                        name: fname.to_string(),
-                        // Safety: `Value::Object` is always valid JSON, so serialization cannot fail
-                        arguments: serde_json::to_string(&args).unwrap(),
-                    },
-                });
+                res.push(ToolCallResponse {
+                    id: format!("call-{}", call_idx),
+                    tp: ToolCallType::Function,
+                    function: CalledFunction {
+                        name: fname,
+                        arguments: args.to_string(),
+                    },
+                });

---

147-153: Potential panic: indexing message.content[0].

Consistency with prior parsers is noted, but this can panic if content is empty. Safer pattern keeps behavior identical when present.

Apply these diffs:

-        if message.author.role == Role::Assistant && message.channel.as_deref() == Some("analysis")
-        {
-            normal_text.push_str(match &message.content[0] {
-                Text(t) => &t.text,
-                _ => "",
-            });
-        }
+        if message.author.role == Role::Assistant && message.channel.as_deref() == Some("analysis") {
+            if let Some(Text(t)) = message.content.first() {
+                normal_text.push_str(&t.text);
+            }
+        }

-        } else if channel == Some("analysis") {
-            normal_text.push_str(match &message.content[0] {
-                Text(t) => &t.text,
-                _ => "",
-            });
-        }
+        } else if channel == Some("analysis") {
+            if let Some(Text(t)) = message.content.first() {
+                normal_text.push_str(&t.text);
+            }
+        }

Also applies to: 251-255

---

155-155: Avoid unnecessary String clone on return.

Return the owned String directly.

Apply these diffs:

-    Ok((res, Some(normal_text.to_string())))
+    Ok((res, Some(normal_text)))

-    Ok((res, Some(normal_text.to_string())))
+    Ok((res, Some(normal_text)))

Also applies to: 257-257

---

225-248: Duplicate clone/unwrap in complete parser.

Same cleanup as streaming variant.

Apply this diff:

-                res.push(ToolCallResponse {
-                    id: format!("call-{}", call_idx),
-                    tp: ToolCallType::Function,
-                    function: CalledFunction {
-                        name: fname.to_string(),
-                        // Safety: `Value::Object` is always valid JSON, so serialization cannot fail
-                        arguments: serde_json::to_string(&args).unwrap(),
-                    },
-                });
+                res.push(ToolCallResponse {
+                    id: format!("call-{}", call_idx),
+                    tp: ToolCallType::Function,
+                    function: CalledFunction {
+                        name: fname,
+                        arguments: args.to_string(),
+                    },
+                });

---

283-299: Tighten partial-token detection loop (remove redundant get/len slicing).

Current code allocates and then calls .get(..) only to slice the entire string again. Simplify without changing logic.

Apply these diffs:

-        config.tool_call_start_tokens.iter().any(|token| {
+        config.tool_call_start_tokens.iter().any(|token| {
             if token.is_empty() {
                 return false;
             }
             // Check if the chunk could be a prefix of this start token
-            // Handle Unicode character boundaries properly
-            for i in 1..=token.chars().count() {
-                if let Some(prefix) = token.chars().take(i).collect::<String>().get(..) {
-                    let prefix_str = &prefix[..prefix.len()];
-                    if trimmed == prefix_str || trimmed.ends_with(prefix_str) {
-                        return true;
-                    }
-                }
-            }
+            for i in 1..=token.chars().count() {
+                let prefix: String = token.chars().take(i).collect();
+                if trimmed == prefix || trimmed.ends_with(&prefix) {
+                    return true;
+                }
+            }
             false
         })

-        let has_partial_token = config.tool_call_start_tokens.iter().any(|token| {
+        let has_partial_token = config.tool_call_start_tokens.iter().any(|token| {
             if token.is_empty() {
                 return false;
             }
             // Check if the chunk could be a prefix of this start token
-            // Handle Unicode character boundaries properly
-            for i in 1..=token.chars().count() {
-                if let Some(prefix) = token.chars().take(i).collect::<String>().get(..) {
-                    let prefix_str = &prefix[..prefix.len()];
-                    if trimmed == prefix_str || trimmed.ends_with(prefix_str) {
-                        return true;
-                    }
-                }
-            }
+            for i in 1..=token.chars().count() {
+                let prefix: String = token.chars().take(i).collect();
+                if trimmed == prefix || trimmed.ends_with(&prefix) {
+                    return true;
+                }
+            }
             false
         });

Also applies to: 311-329

---

188-190: Nit: stray comment markers.

Remove the duplicate slashes.

Apply this diff:

-    // // Encode the text into tokens using harmony encoding
+    // Encode the text into tokens using harmony encoding

---

86-86: Nit: typo in comment.

“Iteratate” → “Iterate”.

Apply this diff:

-    // Iteratate through messages and extract tool calls if there
+    // Iterate through messages and extract tool calls if present

---

364-375: Tests: make sync to match new API (drop tokio/async).

These tests don’t await anything. Convert to plain #[test] for consistency.

Apply these diffs:

-    #[tokio::test]
-    async fn test_parse_tool_calls_harmony_complete_basic() {
+    #[test]
+    fn test_parse_tool_calls_harmony_complete_basic() {
         let text = r#"<|channel|>commentary to=functions.get_current_weather <|constrain|>json<|message|>{"format":"celsius","location":"San Francisco"}"#;
-        let (tool_calls, normal_content) =
-            parse_tool_calls_harmony_complete(text, &Default::default()).unwrap();
+        let (tool_calls, normal_content) =
+            parse_tool_calls_harmony_complete(text, &Default::default()).unwrap();

-    #[tokio::test]
-    async fn test_parse_tools_harmony_without_start_token() {
+    #[test]
+    fn test_parse_tools_harmony_without_start_token() {
         let text = r#"
 <|channel|>analysis<|message|>Need to use function get_current_weather.<|end|>
 <|message|>{"location":"San Francisco"}<|call|>
 "#;

Also applies to: 376-390

---

475-477: Nit: typo in test comment.

“warkaround” → “workaround”.

Apply this diff:

-        // This is a warkaround for now. Right now everything is treated as tool call start token.
+        // This is a workaround for now. Right now everything is treated as a tool call start token.

lib/llm/src/protocols/openai/chat_completions/jail.rs (1)

831-877: Dedup marker patterns before constructing the matcher

Prevents redundant work and ambiguous matches if the same marker is added from multiple sources.

Apply this diff:

         // Create the marker matcher (fallback to empty patterns if none configured)
-        let marker_matcher = if all_patterns.is_empty() {
+        // Dedup to avoid redundant comparisons
+        all_patterns.sort();
+        all_patterns.dedup();
+        let marker_matcher = if all_patterns.is_empty() {

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c63ccea and b804181.

📒 Files selected for processing (8)

• lib/llm/src/protocols/openai/chat_completions/jail.rs (8 hunks)
• lib/parsers/src/tool_calling/harmony/harmony_parser.rs (11 hunks)
• lib/parsers/src/tool_calling/harmony/mod.rs (1 hunks)
• lib/parsers/src/tool_calling/json/mod.rs (1 hunks)
• lib/parsers/src/tool_calling/mod.rs (1 hunks)
• lib/parsers/src/tool_calling/parsers.rs (77 hunks)
• lib/parsers/src/tool_calling/pythonic/mod.rs (1 hunks)
• lib/parsers/src/tool_calling/tools.rs (3 hunks)

🧰 Additional context used

🧠 Learnings (5)

📚 Learning: 2025-09-10T22:32:12.978Z

Learnt from: zhongdaor-nv
PR: ai-dynamo/dynamo#2999
File: lib/parsers/src/tool_calling/harmony/harmony_parser.rs:250-256
Timestamp: 2025-09-10T22:32:12.978Z
Learning: In lib/parsers/src/tool_calling/harmony/harmony_parser.rs, the team prefers to maintain identical code patterns between parse_tool_calls_harmony and parse_tool_calls_harmony_complete functions, including message.content[0] indexing, to ensure consistency between streaming and complete parser implementations.

Applied to files:

• lib/parsers/src/tool_calling/mod.rs
• lib/parsers/src/tool_calling/harmony/mod.rs
• lib/parsers/src/tool_calling/tools.rs
• lib/llm/src/protocols/openai/chat_completions/jail.rs
• lib/parsers/src/tool_calling/parsers.rs
• lib/parsers/src/tool_calling/harmony/harmony_parser.rs

📚 Learning: 2025-09-10T05:04:58.417Z

Learnt from: ayushag-nv
PR: ai-dynamo/dynamo#2932
File: lib/llm/src/protocols/openai/chat_completions/aggregator.rs:66-86
Timestamp: 2025-09-10T05:04:58.417Z
Learning: In the dynamo codebase, tool call chunks from streaming responses always contain complete tool calls (one chunk = one tool call), unlike standard OpenAI streaming where tool calls can be fragmented across multiple chunks. The convert_tool_chunk_to_message_tool_call function correctly assumes complete tool call data in each chunk.

Applied to files:

• lib/parsers/src/tool_calling/tools.rs
• lib/llm/src/protocols/openai/chat_completions/jail.rs

📚 Learning: 2025-09-10T15:27:42.511Z

Learnt from: ayushag-nv
PR: ai-dynamo/dynamo#2932
File: lib/llm/src/preprocessor.rs:768-844
Timestamp: 2025-09-10T15:27:42.511Z
Learning: In the tool calling jail implementation in lib/llm/src/preprocessor.rs, the design intentionally emits only the first accumulated choice that contains tool calls during unjailing, dropping other accumulated choices. This is a deliberate design decision, not a bug.

Applied to files:

• lib/llm/src/protocols/openai/chat_completions/jail.rs

📚 Learning: 2025-08-22T19:55:41.608Z

Learnt from: nachiketb-nvidia
PR: ai-dynamo/dynamo#2656
File: lib/llm/src/protocols/openai/chat_completions/delta.rs:320-327
Timestamp: 2025-08-22T19:55:41.608Z
Learning: The create_choice method exists on multiple different objects in the codebase. The DeltaGenerator::create_choice in lib/llm/src/protocols/openai/chat_completions/delta.rs has its own signature that was updated to include reasoning_content, but other objects in lib/llm/src/engines.rs have their own separate create_choice methods with different signatures that are not related to chat completions.

Applied to files:

• lib/llm/src/protocols/openai/chat_completions/jail.rs

📚 Learning: 2025-08-22T19:55:41.608Z

Learnt from: nachiketb-nvidia
PR: ai-dynamo/dynamo#2656
File: lib/llm/src/protocols/openai/chat_completions/delta.rs:320-327
Timestamp: 2025-08-22T19:55:41.608Z
Learning: There are two separate DeltaGenerator classes in the codebase: one for chat completions (lib/llm/src/protocols/openai/chat_completions/delta.rs with object "chat.completion.chunk") and one for text completions (lib/llm/src/protocols/openai/completions/delta.rs with object "text_completion"). They have different create_choice method signatures and serve different OpenAI API endpoints. The reasoning parsing functionality is only relevant to the chat completions DeltaGenerator.

Applied to files:

• lib/llm/src/protocols/openai/chat_completions/jail.rs

🧬 Code graph analysis (5)

lib/parsers/src/tool_calling/mod.rs (1)

lib/parsers/src/tool_calling/parsers.rs (4)

• detect_and_parse_tool_call (72-96)
• detect_tool_call_start (98-125)
• find_tool_call_end_position (127-155)
• try_tool_call_parse (43-69)

lib/parsers/src/tool_calling/tools.rs (1)

lib/parsers/src/tool_calling/parsers.rs (1)

• detect_and_parse_tool_call (72-96)

lib/llm/src/protocols/openai/chat_completions/jail.rs (2)

lib/parsers/src/tool_calling/parsers.rs (3)

• get_tool_parser_map (23-37)
• detect_tool_call_start (98-125)
• find_tool_call_end_position (127-155)

lib/parsers/src/tool_calling/tools.rs (1)

• try_tool_call_parse_aggregate (10-42)

lib/parsers/src/tool_calling/parsers.rs (6)

lib/parsers/src/tool_calling/config.rs (5)

• harmony (145-154)
• pythonic (138-143)
• default (45-55)
• default (68-73)
• mistral (116-125)

lib/parsers/src/tool_calling/harmony/harmony_parser.rs (2)

• detect_tool_call_start_harmony (260-330)
• parse_tool_calls_harmony_complete (176-258)

lib/parsers/src/tool_calling/harmony/mod.rs (1)

• find_tool_call_end_position_harmony (11-13)

lib/parsers/src/tool_calling/json/mod.rs (4)

• detect_tool_call_start_json (38-43)
• find_tool_call_end_position_json (45-71)
• try_tool_call_parse_json (28-36)
• default (23-25)

lib/parsers/src/tool_calling/pythonic/pythonic_parser.rs (2)

• detect_tool_call_start_pythonic (190-198)
• try_tool_call_parse_pythonic (162-188)

lib/parsers/src/tool_calling/pythonic/mod.rs (1)

• find_tool_call_end_position_pythonic (9-11)

lib/parsers/src/tool_calling/harmony/harmony_parser.rs (1)

lib/parsers/src/reasoning/gpt_oss_parser.rs (2)

• new (39-56)
• get_harmony_encoding (20-23)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: Build and Test - dynamo
• GitHub Check: pre-merge-rust (lib/bindings/python)
• GitHub Check: pre-merge-rust (lib/runtime/examples)
• GitHub Check: pre-merge-rust (.)

🔇 Additional comments (5)

lib/parsers/src/tool_calling/mod.rs (1)

16-19: Re-export of end-position dispatcher looks good.

Public surface now includes find_tool_call_end_position.

lib/parsers/src/tool_calling/parsers.rs (3)

5-15: Consolidated per-parser imports (incl. end-position helpers) — good.

Keeps dispatch localized and explicit.

---

43-69: Synchronous try_tool_call_parse conversion is consistent.

Harmony path uses parse_tool_calls_harmony_complete, JSON/Pythonic paths unchanged. Looks correct.

Confirm external crates no longer expect an async API for this function.

---

72-96: Sync detect_and_parse_tool_call — OK.

Defaulting None/empty to "default" parser is clear; error message enumerates available parsers.

lib/llm/src/protocols/openai/chat_completions/jail.rs (1)

78-99: Helper factory looks good

Centralizing ChatChoiceStream construction reduces duplication and error‑prone field copying.

[coderabbitai]

⚠️ Potential issue

Don't move out of a borrowed choice; clone role and finish_reason

This will otherwise try to move fields from an immutably borrowed choice.

Apply this diff:

-            let pass_through_choice = create_choice_stream(
-                content,
-                choice.delta.role,
-                choice.index,
-                choice.finish_reason,
-                choice.logprobs.clone(),
-            );
+            let pass_through_choice = create_choice_stream(
+                content,
+                choice.delta.role.clone(),
+                choice.index,
+                choice.finish_reason.clone(),
+                choice.logprobs.clone(),
+            );

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	let pass_through_choice = create_choice_stream(
	content,
	choice.delta.role,
	choice.index,
	choice.finish_reason,
	choice.logprobs.clone(),
	);
	let pass_through_choice = create_choice_stream(
	content,
	choice.delta.role.clone(),
	choice.index,
	choice.finish_reason.clone(),
	choice.logprobs.clone(),
	);

🤖 Prompt for AI Agents

In lib/llm/src/protocols/openai/chat_completions/jail.rs around lines 134 to
140, the code is attempting to move fields out of an immutably borrowed
`choice`; change the arguments passed into create_choice_stream so you pass
clones for non-Copy fields (specifically clone `choice.delta.role` and
`choice.finish_reason`) instead of moving them, e.g. call
create_choice_stream(content, choice.delta.role.clone(), choice.index,
choice.finish_reason.clone(), choice.logprobs.clone()) so nothing is moved out
of the borrowed `choice`.

[coderabbitai]

⚠️ Potential issue

Clone role and finish_reason here as well

Same issue as above in trailing emission.

Apply this diff:

-            let trailing_choice = create_choice_stream(
-                trailing,
-                choice.delta.role,
-                choice.index,
-                choice.finish_reason,
-                choice.logprobs.clone(),
-            );
+            let trailing_choice = create_choice_stream(
+                trailing,
+                choice.delta.role.clone(),
+                choice.index,
+                choice.finish_reason.clone(),
+                choice.logprobs.clone(),
+            );

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	let trailing_choice = create_choice_stream(
	trailing,
	choice.delta.role,
	choice.index,
	choice.finish_reason,
	choice.logprobs.clone(),
	);
	let trailing_choice = create_choice_stream(
	trailing,
	choice.delta.role.clone(),
	choice.index,
	choice.finish_reason.clone(),
	choice.logprobs.clone(),
	);

🤖 Prompt for AI Agents

In lib/llm/src/protocols/openai/chat_completions/jail.rs around lines 154 to
160, the trailing_choice creation passes choice.delta.role and
choice.finish_reason by reference (or moves) and does not clone them like
earlier choices do; update the call to clone the role and finish_reason (e.g.,
use choice.delta.role.clone() and choice.finish_reason.clone()) so the trailing
emitted choice owns its data and matches the handling used for other choice
emissions.

[coderabbitai]

⚠️ Potential issue

Await the async end‑detection (jailed path)

Missing .await here too.

Apply this diff:

-            let (should_end, split_pos) = jail_stream.should_end_jail(&self.accumulated_content);
+            let (should_end, split_pos) = jail_stream
+                .should_end_jail(&self.accumulated_content)
+                .await;

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	let (should_end, split_pos) = jail_stream.should_end_jail(&self.accumulated_content);
	let (should_end, split_pos) = jail_stream
	.should_end_jail(&self.accumulated_content)
	.await;

🤖 Prompt for AI Agents

In lib/llm/src/protocols/openai/chat_completions/jail.rs around line 282, the
call to jail_stream.should_end_jail(&self.accumulated_content) is missing an
await; change it to await the async function (i.e., call .await on the returned
future) and update surrounding code if necessary to be inside an async context
so the result (should_end, split_pos) is obtained correctly.

[coderabbitai]

⚠️ Potential issue

Await the async tool‑choice creation (jailed path)

Add .await.

Apply this diff:

-                let unjailed_choice =
-                    jail_stream.create_tool_call_choice(choice.index, jailed_part, choice);
+                let unjailed_choice = jail_stream
+                    .create_tool_call_choice(choice.index, jailed_part, choice)
+                    .await;

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	let unjailed_choice =
	jail_stream.create_tool_call_choice(choice.index, jailed_part, choice);
	Self::emit_tool_or_content(&mut emissions, unjailed_choice);
	let unjailed_choice = jail_stream
	.create_tool_call_choice(choice.index, jailed_part, choice)
	.await;
	Self::emit_tool_or_content(&mut emissions, unjailed_choice);

🤖 Prompt for AI Agents

In lib/llm/src/protocols/openai/chat_completions/jail.rs around lines 293 to
296, the jailed path calls jail_stream.create_tool_call_choice(...) which is
async but is not awaited; change the call to await the future (add .await) so
you assign the resolved value to unjailed_choice before passing it to
Self::emit_tool_or_content; ensure the enclosing function is async or otherwise
able to await (and adjust error/Result handling if needed).

[coderabbitai]

⚠️ Potential issue

Await during finalize

create_tool_call_choice is async; missing .await.

Apply this diff:

-            let final_choice = jail_stream.create_tool_call_choice(
-                self.index,
-                &self.accumulated_content,
-                &dummy_choice,
-            );
+            let final_choice = jail_stream
+                .create_tool_call_choice(self.index, &self.accumulated_content, &dummy_choice)
+                .await;

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	let final_choice = jail_stream.create_tool_call_choice(
	self.index,
	&self.accumulated_content,
	&dummy_choice,
	);
	let final_choice = jail_stream
	.create_tool_call_choice(self.index, &self.accumulated_content, &dummy_choice)
	.await;

🤖 Prompt for AI Agents

In lib/llm/src/protocols/openai/chat_completions/jail.rs around lines 327 to
331, the call to create_tool_call_choice is asynchronous but is invoked without
awaiting; update the call to await it (add .await) and ensure the enclosing
function is async (or otherwise able to await) and that any returned
Result/Value is handled or propagated accordingly so compilation and semantics
remain correct.