From 7cec411cba7cdefef47650820b5ef57a63a68a57 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Mon, 3 Aug 2026 11:03:26 +0700 Subject: [PATCH] fix(stream): robust reasoning/content split + flush un-tagged output The incremental search missed tags split across tokens, inverting reasoning/content classification. Detect the first on the full buffer and track the content boundary as a byte offset. If the model never closes , flush the buffered text as content so clients always receive the response. Chat UI now renders reasoning_content too. Co-Authored-By: Claude Opus 5 (1M context) --- src/application/chat/mod.rs | 4 +- src/application/chat/use_cases.rs | 93 +++++++++++++-------- src/presentation/handler/chat-ui/index.html | 46 ++++++++-- src/presentation/handler/chat.rs | 41 +++++++-- 4 files changed, 134 insertions(+), 50 deletions(-) diff --git a/src/application/chat/mod.rs b/src/application/chat/mod.rs index 184c6a7..78be815 100644 --- a/src/application/chat/mod.rs +++ b/src/application/chat/mod.rs @@ -1,6 +1,6 @@ pub mod use_cases; pub use use_cases::{ - build_prompt, build_sampler, clean_text, parse_tool_calls, split_stream_chunk, validate_model, - SamplerParams, + build_prompt, build_sampler, clean_text, parse_tool_calls, split_stream_chunk, strip_markup, + validate_model, SamplerParams, }; diff --git a/src/application/chat/use_cases.rs b/src/application/chat/use_cases.rs index d0aa200..fe44e6a 100644 --- a/src/application/chat/use_cases.rs +++ b/src/application/chat/use_cases.rs @@ -335,7 +335,7 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec) { /// Handles both fixed tags (`<|im_end|>`, ``, ``, …) and the /// attribute-bearing openers used by this model's tool format (``, /// ``), even when a tag straddles a token boundary. -fn strip_markup(text: &str) -> String { +pub fn strip_markup(text: &str) -> String { let mut out = String::with_capacity(text.len()); let mut rest = text; @@ -383,22 +383,42 @@ fn strip_markup(text: &str) -> String { out } -/// Split an incremental streamed text fragment into `(reasoning, content)` -/// deltas for SSE. +/// Compute `(reasoning, content)` deltas for an incremental streamed fragment. /// -/// * `think_done` means the `` boundary was already crossed **before** -/// this fragment (i.e. it is not the chunk containing the first ``). -/// * Whitespace **inside** a fragment is preserved — only the whitespace -/// sitting immediately around the `` boundary is trimmed, so -/// reasoning does not end with a dangling newline and content does not begin -/// with one. (Trimming every fragment corrupted inter-word spaces.) -/// * Before the first ``, everything is emitted as `reasoning_content`; -/// after it, as `content`. A stray second `` is stripped, not split. -pub fn split_stream_chunk(new_text: &str, think_done: bool) -> (Option, String) { - if !think_done { - if let Some(pos) = new_text.find("") { - let mut before = strip_markup(&new_text[..pos]); - let mut after = strip_markup(&new_text[pos + 8..]); +/// * `new_text` — the not-yet-emitted fragment (`text_buf[sent_len..]`). +/// * `sent_len` — byte offset in the full buffer where `new_text` begins. +/// * `content_start` — byte offset in the full buffer where the content phase +/// begins (immediately after the first ``); `None` while still +/// thinking. +/// +/// The boundary is a position in the *full* buffer, not a string search in the +/// fragment — this stays correct even when `` is split across tokens. +/// Whitespace inside a fragment is preserved; only the edges around the +/// boundary are trimmed. A stray second `` is stripped, not re-split. +pub fn split_stream_chunk( + new_text: &str, + sent_len: usize, + content_start: Option, +) -> (Option, String) { + match content_start { + // Still thinking — everything is reasoning. + None => { + let cleaned = strip_markup(new_text); + let reasoning = if cleaned.is_empty() { + None + } else { + Some(cleaned) + }; + (reasoning, String::new()) + } + // Boundary already emitted — everything is content. + Some(cs) if cs <= sent_len => (None, strip_markup(new_text)), + // Boundary falls inside this fragment (or beyond it, defensively). + Some(cs) => { + let rel = (cs - sent_len).min(new_text.len()); + let (before, after) = new_text.split_at(rel); + let mut before = strip_markup(before); + let mut after = strip_markup(after); while before.ends_with(['\n', ' ', '\t']) { before.pop(); @@ -413,19 +433,7 @@ pub fn split_stream_chunk(new_text: &str, think_done: bool) -> (Option, Some(before) }; (reasoning, after) - } else { - // Still thinking — everything is reasoning. - let cleaned = strip_markup(new_text); - let reasoning = if cleaned.is_empty() { - None - } else { - Some(cleaned) - }; - (reasoning, String::new()) } - } else { - // Think phase already ended — everything is content; strip stray markup. - (None, strip_markup(new_text)) } } @@ -471,7 +479,7 @@ mod tests { #[test] fn split_chunk_before_think_is_reasoning() { - let (reasoning, content) = split_stream_chunk("Hello ", false); + let (reasoning, content) = split_stream_chunk("Hello ", 0, None); assert_eq!(reasoning.as_deref(), Some("Hello ")); assert_eq!(content, ""); } @@ -479,21 +487,23 @@ mod tests { #[test] fn split_chunk_preserves_internal_spaces() { // Regression: trimming every fragment used to eat inter-word spaces. - let (r1, _) = split_stream_chunk("Hello", false); - let (r2, _) = split_stream_chunk(" world", false); + let (r1, _) = split_stream_chunk("Hello", 0, None); + let (r2, _) = split_stream_chunk(" world", 5, None); assert_eq!(format!("{}{}", r1.unwrap(), r2.unwrap()), "Hello world"); } #[test] fn split_chunk_boundary_trims_only_edges() { - let (reasoning, content) = split_stream_chunk("question\n\n\nAnswer ", false); + // text_buf = "question\n\n\nAnswer "; boundary right after the + // tag at byte 17. + let (reasoning, content) = split_stream_chunk("question\n\n\nAnswer ", 0, Some(17)); assert_eq!(reasoning.as_deref(), Some("question")); assert_eq!(content, "Answer "); } #[test] fn split_chunk_after_think_is_content() { - let (reasoning, content) = split_stream_chunk(" answer", true); + let (reasoning, content) = split_stream_chunk(" answer", 0, Some(0)); assert_eq!(reasoning, None); assert_eq!(content, " answer"); } @@ -502,16 +512,29 @@ mod tests { fn split_chunk_stray_think_tag_is_stripped_not_split() { // A second (already past the boundary) must not restart // reasoning classification. - let (reasoning, content) = split_stream_chunk("...more", true); + let (reasoning, content) = split_stream_chunk("...more", 0, Some(0)); assert_eq!(reasoning, None); assert_eq!(content, "...more"); } + #[test] + fn split_chunk_boundary_straddling_tokens() { + // `` split as "" across two fragments: the boundary + // is detected on the full buffer, so the answer still becomes content. + let (r1, _) = split_stream_chunk("reasoning..."). let (reasoning, content) = split_stream_chunk( "<|im_end|>Hello\n", - false, + 0, + Some(30), ); assert_eq!(reasoning.as_deref(), Some("Hello")); assert_eq!(content, ""); diff --git a/src/presentation/handler/chat-ui/index.html b/src/presentation/handler/chat-ui/index.html index ad440c1..4e4545c 100644 --- a/src/presentation/handler/chat-ui/index.html +++ b/src/presentation/handler/chat-ui/index.html @@ -72,6 +72,13 @@ display: flex; align-items: center; gap: 6px; } .msg .tool-call::before { content: '\1F527'; } + .msg .reasoning { + font-size: 12px; font-style: italic; + color: var(--text2); + white-space: pre-wrap; word-break: break-word; + border-bottom: 1px solid var(--border); + padding-bottom: 8px; margin-bottom: 8px; + } .msg.error { background: #2a1818; border-color: #4a2828; color: #f08080; } @@ -186,8 +193,33 @@ async function send() { const decoder = new TextDecoder(); let buffer = ''; let full = ''; + let reasoning = ''; let toolCalls = null; + // Re-render the assistant bubble: reasoning (muted) above the answer. + function render() { + el.innerHTML = ''; + if (reasoning) { + const r = document.createElement('div'); + r.className = 'reasoning'; + r.textContent = reasoning; + el.appendChild(r); + } + if (full) { + const c = document.createElement('div'); + c.textContent = full; + el.appendChild(c); + } + if (toolCalls?.length) { + for (const tc of toolCalls) { + const t = document.createElement('div'); + t.className = 'tool-call'; + t.textContent = 'Calling tool: ' + (tc.function?.name || 'tool'); + el.appendChild(t); + } + } + } + while (true) { const { done, value } = await reader.read(); if (done) break; @@ -206,18 +238,20 @@ async function send() { const delta = chunk.choices?.[0]?.delta; const finish = chunk.choices?.[0]?.finish_reason; + if (delta?.reasoning_content) { + reasoning += delta.reasoning_content; + render(); + } if (delta?.content) { full += delta.content; - el.textContent = full; + render(); } if (delta?.tool_calls) { toolCalls = delta.tool_calls; + render(); } - if (finish === 'tool_calls' && toolCalls) { - const t = document.createElement('div'); - t.className = 'tool-call'; - t.textContent = 'Calling tool: ' + toolCalls.map(tc => tc.function?.name).join(', '); - el.appendChild(t); + if (finish === 'tool_calls') { + render(); } } catch (e) { /* skip malformed chunk */ } } diff --git a/src/presentation/handler/chat.rs b/src/presentation/handler/chat.rs index d5e5d51..92f7caa 100644 --- a/src/presentation/handler/chat.rs +++ b/src/presentation/handler/chat.rs @@ -176,7 +176,9 @@ async fn handle_streaming( let mut sampler = SendSampler(chat::build_sampler(¶ms)); let mut text_buf = String::new(); let mut sent_len: usize = 0; - let mut think_done = false; + // Byte offset in text_buf where the content phase begins (right after + // the first ``); None while still thinking. + let mut content_start: Option = None; let outcome = engine.generate( &input_tokens, @@ -187,11 +189,13 @@ async fn handle_streaming( &mut |_token, piece| { text_buf.push_str(piece); - // `was_thinking` is passed to split_stream_chunk so the chunk - // containing the first is treated as the boundary. - let was_thinking = !think_done; - if !think_done && text_buf.contains("") { - think_done = true; + // Robust boundary detection on the *full* buffer — a `` + // tag may be split across tokens, which would defeat a search + // over the incremental fragment only. + if content_start.is_none() { + if let Some(pos) = text_buf.find("") { + content_start = Some(pos + 8); + } } let new_text = &text_buf[sent_len..]; @@ -199,7 +203,8 @@ async fn handle_streaming( return true; } - let (reasoning, content) = chat::split_stream_chunk(new_text, was_thinking); + let (reasoning, content) = + chat::split_stream_chunk(new_text, sent_len, content_start); if let Some(reasoning) = reasoning { let chunk = SseChunk::delta( @@ -251,6 +256,28 @@ async fn handle_streaming( total_tokens: prompt_tokens + completion_tokens, }; + // If the model never emitted ``, everything was streamed + // as reasoning_content. Flush it as content so the client always + // receives the response text. + if content_start.is_none() && !text_buf.is_empty() { + let cleaned = chat::strip_markup(&text_buf); + if !cleaned.is_empty() { + let chunk = SseChunk::delta( + chat_id.clone(), + created, + model_name.clone(), + SseDelta { + role: None, + content: Some(cleaned), + tool_calls: None, + reasoning_content: None, + }, + ); + let event = serde_json::to_string(&chunk).unwrap(); + let _ = tx.blocking_send(Ok(Event::default().data(event))); + } + } + // Single-shot tool-calls delta (this model emits whole blocks). let mut sent_tool_calls = false; if outcome.finish == FinishReason::ToolCalls {