fix(openai): emit prompt-guided stream scrubber's flushed suffix before Completed

crates/tinyinference-llm/src/providers/openai/test.rs

@@ -1883,6 +1883,66 @@ async fn sse_stream_recovers_tool_args_with_leaked_template_marker() {
     assert_eq!(calls[0].name, "composio_execute");
     assert_eq!(calls[0].arguments, json!({ "q": 1 }));
 }
+#[test]
+fn prompt_guided_streaming_emits_the_scrubbers_flushed_suffix_before_completed() {
+    // Regression for a dropped-stream-output bug: the terminal `Completed`
+    // item used to discard `TextScrubber::flush()`'s return value outright,
+    // so a narrative suffix withheld pending marker disambiguation (here, a
+    // trailing `<tool_` that never resolves into a full `<tool_call>`) never
+    // reached a streaming consumer, even though the terminal response's own
+    // text still carried it.
+    let tools = vec![ToolSchema::new("x", "x", json!({"type": "object"}))];
+    let mut scrubber = crate::prompt_tools::TextScrubber::new(&tools);
+
+    let delta_items = transport::scrub_prompt_guided_item(
+        ModelStreamItem::MessageDelta(crate::message::MessageDelta::text("before <tool_")),
+        &mut scrubber,
+        &tools,
+    );
+    assert_eq!(delta_items.len(), 1);
+    match &delta_items[0] {
+        ModelStreamItem::MessageDelta(delta) => assert_eq!(delta.text, "before "),
+        other => panic!("expected a MessageDelta, got {other:?}"),
+    }
+
+    let response = crate::model::ModelResponse::assistant("before <tool_");
+    let completed_items = transport::scrub_prompt_guided_item(
+        ModelStreamItem::Completed(response),
+        &mut scrubber,
+        &tools,
+    );
+    assert_eq!(completed_items.len(), 2);
+    match &completed_items[0] {
+        ModelStreamItem::MessageDelta(delta) => assert_eq!(delta.text, "<tool_"),
+        other => panic!("expected the flushed suffix as a MessageDelta, got {other:?}"),
+    }
+    assert!(matches!(completed_items[1], ModelStreamItem::Completed(_)));
+}
+
+#[test]
+fn prompt_guided_streaming_completed_without_buffered_text_emits_one_item() {
+    // The common case: nothing was withheld, so `Completed` must not gain a
+    // spurious empty `MessageDelta` ahead of it.
+    let tools = vec![ToolSchema::new("x", "x", json!({"type": "object"}))];
+    let mut scrubber = crate::prompt_tools::TextScrubber::new(&tools);
+
+    let delta_items = transport::scrub_prompt_guided_item(
+        ModelStreamItem::MessageDelta(crate::message::MessageDelta::text("plain text")),
+        &mut scrubber,
+        &tools,
+    );
+    assert_eq!(delta_items.len(), 1);
+
+    let response = crate::model::ModelResponse::assistant("plain text");
+    let completed_items = transport::scrub_prompt_guided_item(
+        ModelStreamItem::Completed(response),
+        &mut scrubber,
+        &tools,
+    );
+    assert_eq!(completed_items.len(), 1);
+    assert!(matches!(completed_items[0], ModelStreamItem::Completed(_)));
+}
+
 // `ChatModel::profile` is generic over `State`; pin `State = ()` so the concrete
 // `OpenAiModel` handle disambiguates without a turbofish at every call site.
 fn profile_of(model: &OpenAiModel) -> &crate::model::ModelProfile {

crates/tinyinference-llm/src/providers/openai/transport.rs

@@ -1864,34 +1864,16 @@ impl<State: Send + Sync> ChatModel<State> for OpenAiModel {
         // deltas as they arrive, and recover the calls from the terminal
         // `Completed` response into `message.tool_calls`. Calls are dispatched
         // from the terminal response only, so a consumer sees each exactly
-        // once; the scrubber's own releases are dropped.
+        // once; the scrubber's own releases are dropped. Any narrative text
+        // the scrubber was withholding pending disambiguation (an in-progress
+        // marker prefix that never completed) is flushed as one final delta
+        // ahead of `Completed`, so a streaming consumer sees the same visible
+        // suffix the terminal response carries.
         if self.prompt_guided_for(&request) {
             let tools = request.tools.clone();
             let mut scrubber = crate::prompt_tools::TextScrubber::new(&tools);
-            let stream = ModelStream::new(Box::pin(stream.filter_map(move |item| {
-                let mapped = match item {
-                    ModelStreamItem::MessageDelta(mut delta) if !delta.text.is_empty() => {
-                        let (text, _released) = scrubber.feed(&delta.text);
-                        delta.text = text;
-                        // A delta the scrubber emptied carries nothing worth
-                        // waking a consumer for.
-                        if delta.text.is_empty()
-                            && delta.reasoning.is_empty()
-                            && delta.tool_call.is_none()
-                        {
-                            return futures::future::ready(None);
-                        }
-                        ModelStreamItem::MessageDelta(delta)
-                    }
-                    ModelStreamItem::Completed(response) => {
-                        let _ = scrubber.flush();
-                        ModelStreamItem::Completed(crate::prompt_tools::recover_tool_calls(
-                            response, &tools,
-                        ))
-                    }
-                    other => other,
-                };
-                futures::future::ready(Some(mapped))
+            let stream = ModelStream::new(Box::pin(stream.flat_map(move |item| {
+                futures::stream::iter(scrub_prompt_guided_item(item, &mut scrubber, &tools))
             })));
             return Ok(match correlation {
                 Some(correlation) => stream.with_correlation(correlation),
@@ -1906,6 +1888,53 @@ impl<State: Send + Sync> ChatModel<State> for OpenAiModel {
     }
 }
 
+/// Applies prompt-guided scrubbing/recovery to one streamed item, returning
+/// zero or more items to forward.
+///
+/// * A `MessageDelta` is fed through the scrubber; a delta the scrubber
+///   emptied entirely (and that carries no reasoning or tool-call fragment
+///   either) is dropped rather than waking a consumer for nothing.
+/// * The terminal `Completed` response runs tool-call recovery. Any
+///   narrative text the scrubber was withholding pending disambiguation — an
+///   in-progress marker prefix that never resolved — is flushed and emitted
+///   as one final `MessageDelta` immediately before `Completed`, so a
+///   streaming consumer sees the same visible suffix the terminal response
+///   carries. The scrubber's own recovered calls are discarded here: calls
+///   are dispatched from the terminal response only, so a consumer sees each
+///   exactly once.
+/// * Every other item passes through unchanged.
+pub(super) fn scrub_prompt_guided_item(
+    item: ModelStreamItem,
+    scrubber: &mut crate::prompt_tools::TextScrubber,
+    tools: &[crate::tool::ToolSchema],
+) -> Vec<ModelStreamItem> {
+    match item {
+        ModelStreamItem::MessageDelta(mut delta) if !delta.text.is_empty() => {
+            let (text, _released) = scrubber.feed(&delta.text);
+            delta.text = text;
+            if delta.text.is_empty() && delta.reasoning.is_empty() && delta.tool_call.is_none() {
+                Vec::new()
+            } else {
+                vec![ModelStreamItem::MessageDelta(delta)]
+            }
+        }
+        ModelStreamItem::Completed(response) => {
+            let (flushed_text, _released) = scrubber.flush();
+            let mut items = Vec::with_capacity(2);
+            if !flushed_text.is_empty() {
+                items.push(ModelStreamItem::MessageDelta(MessageDelta::text(
+                    flushed_text,
+                )));
+            }
+            items.push(ModelStreamItem::Completed(
+                crate::prompt_tools::recover_tool_calls(response, tools),
+            ));
+            items
+        }
+        other => vec![other],
+    }
+}
+
 fn responses_sse_failure(body: &str, model: &OpenAiModel) -> Option<ProviderError> {
     for line in body.lines() {
         let Some(payload) = line.strip_prefix("data:").map(str::trim) else {