diff --git a/crates/tinyinference-llm/src/providers/openai/test.rs b/crates/tinyinference-llm/src/providers/openai/test.rs index 5fdc6be..cf3fd51 100644 --- a/crates/tinyinference-llm/src/providers/openai/test.rs +++ b/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 ``) 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 assert_eq!(delta.text, "before "), + other => panic!("expected a MessageDelta, got {other:?}"), + } + + let response = crate::model::ModelResponse::assistant("before assert_eq!(delta.text, " 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 { diff --git a/crates/tinyinference-llm/src/providers/openai/transport.rs b/crates/tinyinference-llm/src/providers/openai/transport.rs index 56f1cff..a114f70 100644 --- a/crates/tinyinference-llm/src/providers/openai/transport.rs +++ b/crates/tinyinference-llm/src/providers/openai/transport.rs @@ -1864,34 +1864,16 @@ impl ChatModel 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 ChatModel 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 { + 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 { for line in body.lines() { let Some(payload) = line.strip_prefix("data:").map(str::trim) else {