| 1 | //! Streaming response state and guardrails. |
| 2 | //! |
| 3 | //! This module owns the local state used while decoding one model stream: |
| 4 | //! content block kind tracking, streamed tool-use buffers, transparent retry |
| 5 | //! policy, and scrubbers for text that looks like a forged tool-call wrapper. |
| 6 | |
| 7 | use codewhale_models::ToolCaller; |
| 8 | use std::time::Duration; |
| 9 | |
| 10 | #[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| 11 | pub(super) enum ContentBlockKind { |
| 12 | Text, |
| 13 | Thinking, |
| 14 | ToolUse, |
| 15 | } |
| 16 | |
| 17 | #[derive(Debug, Clone)] |
| 18 | pub(super) struct ToolUseState { |
| 19 | pub(super) id: String, |
| 20 | pub(super) name: String, |
| 21 | pub(super) input: serde_json::Value, |
| 22 | pub(super) caller: Option<ToolCaller>, |
| 23 | /// Google thought signature captured on the tool call; replayed with the |
| 24 | /// assistant tool-call message on later turns. |
| 25 | pub(super) thought_signature: Option<String>, |
| 26 | pub(super) input_buffer: String, |
| 27 | pub(super) input_parse_error: Option<String>, |
| 28 | } |
| 29 | |
| 30 | /// Maximum total bytes of text/thinking content before aborting the stream. |
| 31 | pub(super) const STREAM_MAX_CONTENT_BYTES: usize = 10 * 1024 * 1024; // 10 MB |
| 32 | /// Sanity backstop for total stream wall-clock duration. **Not** a routine |
| 33 | /// kill switch — the stream chunk idle timeout is the primary stall |
| 34 | /// detector. The wall-clock cap is here only to bound pathological cases |
| 35 | /// (e.g. a server that keeps sending heartbeats forever without progress). |
| 36 | /// |
| 37 | /// History: this used to be 300s (5 min) which was too aggressive — V4 |
| 38 | /// thinking turns on hard prompts legitimately exceed 5 minutes wall-clock |
| 39 | /// while still emitting reasoning_content chunks the whole way. Bumped to |
| 40 | /// 30 min in v0.6.6 after long-reasoning turns hit the old cap. Codex defaults to a |
| 41 | /// per-chunk idle of 300s with no wall-clock cap; we keep both layers but |
| 42 | /// give the wall-clock a generous window so it never fires in practice. |
| 43 | pub(super) const STREAM_MAX_DURATION_SECS: u64 = 1800; // 30 minutes (was 300s; #103/#1) |
| 44 | /// Hard cap on consecutive recoverable stream errors before we surface a turn |
| 45 | /// failure. Bumped 3 → 5 in v0.6.7 along with the HTTP/2 keepalive defaults |
| 46 | /// (#103) — keepalive should make spurious decode errors rarer, so we can |
| 47 | /// tolerate a longer streak before giving up on the turn. |
| 48 | pub(super) const MAX_STREAM_ERRORS_BEFORE_FAIL: u32 = 5; |
| 49 | /// Cap on transparent stream-level retries — these only happen when the wire |
| 50 | /// dies before any content was streamed. The user has seen nothing, but |
| 51 | /// provider usage or billing may already exist. Two attempts can ride out a |
| 52 | /// flaky edge node without amplifying real outages (#103). |
| 53 | pub(super) const MAX_TRANSPARENT_STREAM_RETRIES: u32 = 2; |
| 54 | |
| 55 | /// Decide whether a stream error is eligible for a transparent retry. |
| 56 | /// |
| 57 | /// True only when ALL three conditions hold: |
| 58 | /// 1. No content has been received on the current attempt. Reissuing after |
| 59 | /// visible partial deltas needs a separate recovery policy. This content |
| 60 | /// check is not evidence that the provider consumed or billed zero tokens. |
| 61 | /// 2. We still have transparent-retry budget remaining. |
| 62 | /// 3. The turn has not been cancelled. |
| 63 | /// |
| 64 | /// Extracted as a pure function so the four #103 retry cases can be exercised |
| 65 | /// in unit tests without booting the full engine state machine. |
| 66 | pub(super) fn should_transparently_retry_stream( |
| 67 | any_content_received: bool, |
| 68 | transparent_attempts: u32, |
| 69 | cancelled: bool, |
| 70 | ) -> bool { |
| 71 | !any_content_received && transparent_attempts < MAX_TRANSPARENT_STREAM_RETRIES && !cancelled |
| 72 | } |
| 73 | |
| 74 | /// Budget for re-issuing the whole request after a dead stream. Shared by the |
| 75 | /// nothing-streamed outer retry (#103 Phase 3) and the sleep-resume retry |
| 76 | /// (#2990). |
| 77 | pub(super) const MAX_STREAM_RETRIES: u32 = 3; |
| 78 | |
| 79 | /// Typed, engine-internal state for one mid-stream drop recovery. |
| 80 | /// |
| 81 | /// This enum **is** the retry mechanism. A resumed turn used to append a |
| 82 | /// synthetic `[runtime]` *user* message to the persisted conversation, which |
| 83 | /// polluted the transcript and — when only hidden reasoning had streamed — |
| 84 | /// promised a preserved partial answer that never existed (0.9.10 |
| 85 | /// regression). The retry is now modeled as this value, carried out of the |
| 86 | /// stream decoder in [`StreamOutcome`] and consumed exactly once per drop: |
| 87 | /// |
| 88 | /// * it is never persisted to the user transcript, and |
| 89 | /// * nothing it triggers is serialized into the provider request history as |
| 90 | /// a user role — the retried request is simply the persisted conversation |
| 91 | /// re-issued, ending (when a visible fragment was preserved) with that |
| 92 | /// assistant fragment so the provider continues from it. |
| 93 | #[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| 94 | pub(super) enum StreamResume { |
| 95 | /// The stream died before anything actionable was streamed (#103 |
| 96 | /// Phase 3): discard the fragment and re-issue the identical request. |
| 97 | NoContentStreamDeath, |
| 98 | /// The host slept mid-stream (#2990): the partial output predates the |
| 99 | /// sleep and no operator watched it — discard and re-issue. |
| 100 | AfterSleep, |
| 101 | /// Mid-stream network drop on a headless host (v0.9.4 Terminal-Bench |
| 102 | /// P0): the fragment was never committed and no tool from it ran, so |
| 103 | /// discard and re-issue the identical request. |
| 104 | HeadlessNetworkDrop, |
| 105 | /// Mid-stream network drop in the interactive TUI. A fragment with |
| 106 | /// sendable content is preserved as the trailing assistant message and |
| 107 | /// the request is re-issued; a thinking-only fragment has nothing |
| 108 | /// visible to preserve, so it is discarded exactly like the headless |
| 109 | /// resume and the copy must never claim otherwise. |
| 110 | InteractiveNetworkDrop, |
| 111 | } |
| 112 | |
| 113 | /// Bounded authorization for drop-resume retries. |
| 114 | /// |
| 115 | /// Mechanism, not comment: [`StreamRetryBudget::authorize`] is the only way |
| 116 | /// to spend a resume and it returns `None` once [`MAX_STREAM_RETRIES`] |
| 117 | /// resumes have been issued, so no call site can loop past the budget even |
| 118 | /// if a guard predicate is relaxed. A healthy stream round resets it. |
| 119 | #[derive(Debug, Default)] |
| 120 | pub(super) struct StreamRetryBudget { |
| 121 | spent: u32, |
| 122 | } |
| 123 | |
| 124 | impl StreamRetryBudget { |
| 125 | /// Drop-resumes already issued without a healthy round in between. |
| 126 | pub(super) fn spent(&self) -> u32 { |
| 127 | self.spent |
| 128 | } |
| 129 | |
| 130 | /// Spend one resume and return its 1-based attempt number, or `None` |
| 131 | /// when the budget is exhausted. |
| 132 | pub(super) fn authorize(&mut self) -> Option<u32> { |
| 133 | if self.spent >= MAX_STREAM_RETRIES { |
| 134 | return None; |
| 135 | } |
| 136 | self.spent = self.spent.saturating_add(1); |
| 137 | Some(self.spent) |
| 138 | } |
| 139 | |
| 140 | /// A healthy round clears the chain: the next drop starts a fresh, |
| 141 | /// still-bounded budget. |
| 142 | pub(super) fn reset(&mut self) { |
| 143 | self.spent = 0; |
| 144 | } |
| 145 | } |
| 146 | |
| 147 | /// Wall-clock vs monotonic divergence above which we conclude the host slept |
| 148 | /// mid-stream (#2990). `Instant` pauses during system sleep (CLOCK_UPTIME_RAW |
| 149 | /// on macOS, CLOCK_MONOTONIC on Linux) while `SystemTime` keeps advancing, so |
| 150 | /// a large positive gap can only come from a suspend/resume cycle — ordinary |
| 151 | /// network flakes never produce one. Windows `Instant` may keep ticking |
| 152 | /// through sleep, in which case this simply never fires (no behavior change). |
| 153 | pub(super) const SLEEP_GAP_THRESHOLD: Duration = Duration::from_secs(10); |
| 154 | |
| 155 | /// True when the gap between wall-clock and monotonic elapsed time since the |
| 156 | /// last stream progress says the host was suspended. |
| 157 | pub(super) fn sleep_gap_detected(monotonic_elapsed: Duration, wallclock_elapsed: Duration) -> bool { |
| 158 | wallclock_elapsed.saturating_sub(monotonic_elapsed) > SLEEP_GAP_THRESHOLD |
| 159 | } |
| 160 | |
| 161 | /// Decide whether a failed stream should be silently re-issued because the |
| 162 | /// host slept mid-turn (#2990). |
| 163 | /// |
| 164 | /// Unlike the transparent retry (#103), this fires even after content has |
| 165 | /// streamed: the partial output predates the sleep, the user was not |
| 166 | /// watching, and re-running the identical request is the correct |
| 167 | /// user-visible behavior. The double-billing concern that blocks ordinary |
| 168 | /// post-content retries is accepted here because the alternative is a dead |
| 169 | /// turn the user must re-prompt (and pay for) anyway. |
| 170 | pub(super) fn should_resume_after_sleep( |
| 171 | sleep_detected: bool, |
| 172 | retry_attempts: u32, |
| 173 | cancelled: bool, |
| 174 | ) -> bool { |
| 175 | sleep_detected && retry_attempts < MAX_STREAM_RETRIES && !cancelled |
| 176 | } |
| 177 | |
| 178 | /// Decide whether a failed stream should be re-issued after a mid-stream |
| 179 | /// network drop in a headless host (`exec` / stream-json / app-server), even |
| 180 | /// though content already streamed. |
| 181 | /// |
| 182 | /// This extends the #2990 sleep-resume contract to ordinary transport drops |
| 183 | /// for hosts with no operator watching: the partial assistant fragment has |
| 184 | /// not been committed to the conversation and no tool call from the |
| 185 | /// incomplete response has executed, so discarding the fragment and |
| 186 | /// re-issuing the identical request cannot duplicate side effects. The |
| 187 | /// double-billing risk that blocks post-content retries in the interactive |
| 188 | /// TUI (#103) is accepted here because the alternative is a dead turn that |
| 189 | /// forfeits the entire headless run — the exact tradeoff #2990 already makes |
| 190 | /// for sleep-resume. Interactive sessions keep the #103 surface-the-warning |
| 191 | /// behavior: the user saw the partial deltas, and replaying would render the |
| 192 | /// same prefix twice. |
| 193 | pub(super) fn should_resume_after_network_drop( |
| 194 | headless_host: bool, |
| 195 | network_class_error: bool, |
| 196 | retry_attempts: u32, |
| 197 | cancelled: bool, |
| 198 | ) -> bool { |
| 199 | headless_host && network_class_error && retry_attempts < MAX_STREAM_RETRIES && !cancelled |
| 200 | } |
| 201 | |
| 202 | /// Decide whether an interactive TUI stream should be re-issued after a |
| 203 | /// mid-stream network drop, preserving a visible partial reply. |
| 204 | /// |
| 205 | /// Unlike the headless resume, this keeps a sendable fragment: the user has |
| 206 | /// already seen the deltas, so the assistant message is committed and the |
| 207 | /// re-issued request ends with that fragment, which is the provider-neutral |
| 208 | /// continuation contract. No synthetic user turn is appended — the retry is |
| 209 | /// typed state ([`StreamResume::InteractiveNetworkDrop`]), invisible to the |
| 210 | /// transcript and to the provider request history as a user role. A |
| 211 | /// thinking-only fragment preserves nothing and must not be described as a |
| 212 | /// preserved reply. Tool calls are never resumed because an incomplete tool |
| 213 | /// call could be re-issued and duplicate side effects. Bounded by |
| 214 | /// `MAX_STREAM_RETRIES` and gated on a network/timeout-class error so |
| 215 | /// model/parse/auth failures still surface normally. |
| 216 | pub(super) fn should_resume_interactive_after_network_drop( |
| 217 | terminal_chrome_enabled: bool, |
| 218 | network_class_error: bool, |
| 219 | any_content_received: bool, |
| 220 | tool_uses_empty: bool, |
| 221 | retry_attempts: u32, |
| 222 | cancelled: bool, |
| 223 | ) -> bool { |
| 224 | terminal_chrome_enabled |
| 225 | && network_class_error |
| 226 | && any_content_received |
| 227 | && tool_uses_empty |
| 228 | && retry_attempts < MAX_STREAM_RETRIES |
| 229 | && !cancelled |
| 230 | } |
| 231 | |
| 232 | /// Convert low-level reqwest/hyper stream read errors into an operator-facing |
| 233 | /// message. The raw provider error remains attached, but the lead sentence |
| 234 | /// explains why Codewhale may retry before any output and why it must surface |
| 235 | /// the warning once partial output has already streamed. |
| 236 | pub(super) fn stream_read_error_user_message(message: &str, any_content_received: bool) -> String { |
| 237 | let lower = message.to_ascii_lowercase(); |
| 238 | let is_stream_read = lower.contains("stream read error") |
| 239 | || lower.contains("error decoding response body") |
| 240 | || lower.contains("chunk decode error") |
| 241 | || lower.contains("body decode"); |
| 242 | if !is_stream_read { |
| 243 | return message.to_string(); |
| 244 | } |
| 245 | |
| 246 | let retry_note = if any_content_received { |
| 247 | "Some output had already streamed, so Codewhale is surfacing the warning instead of replaying the request and risking duplicated output." |
| 248 | } else { |
| 249 | "No output had streamed yet, so Codewhale will retry automatically while retry budget remains." |
| 250 | }; |
| 251 | format!( |
| 252 | "Provider stream connection dropped while reading the response body. {retry_note} Details: {message}" |
| 253 | ) |
| 254 | } |
| 255 | |
| 256 | /// Wrapper shapes a model may emit as plain text instead of using the API tool |
| 257 | /// channel. Each pair is `(start, end)`; the tables below are projections of |
| 258 | /// this one and must stay in sync with it. |
| 259 | /// |
| 260 | /// Three families are covered: |
| 261 | /// |
| 262 | /// 1. Generic/Anthropic-style (`[TOOL_CALL]`, `<invoke …>`, `<function_calls>`). |
| 263 | /// 2. DSML wrappers, in fullwidth `｜` (U+FF5C) and ASCII `|` delimiters, upper |
| 264 | /// and lower case. |
| 265 | /// 3. **DeepSeek's native tool-call tokens** (#3880). DeepSeek's chat template |
| 266 | /// separates words with `▁` (U+2581 LOWER ONE EIGHTH BLOCK), not a space or |
| 267 | /// underscore, so `<｜tool▁calls▁begin｜>` does not match any DSML entry and |
| 268 | /// leaked into visible output. Both the `▁` and `_` separators are listed |
| 269 | /// because a partially-normalizing tokenizer can emit either, and both |
| 270 | /// delimiter forms because the ASCII fallback shows up in some renderings. |
| 271 | /// |
| 272 | /// When adding a shape, add it here and to the two marker tables below. |
| 273 | /// `marker_tables_are_consistent` enforces that they agree. |
| 274 | pub(crate) const TOOL_CALL_MARKER_PAIRS: [(&str, &str); 28] = [ |
| 275 | ("[TOOL_CALL]", "[/TOOL_CALL]"), |
| 276 | ("<codewhale:tool_call", "</codewhale:tool_call>"), |
| 277 | ("<tool_call", "</tool_call>"), |
| 278 | ("<invoke ", "</invoke>"), |
| 279 | ("<function_calls>", "</function_calls>"), |
| 280 | ("<｜DSML｜tool_calls>", "</｜DSML｜tool_calls>"), |
| 281 | ("<｜DSML｜invoke ", "</｜DSML｜invoke>"), |
| 282 | ("<|DSML|tool_calls>", "</|DSML|tool_calls>"), |
| 283 | ("<|DSML|invoke ", "</|DSML|invoke>"), |
| 284 | ("<|dsml|tool_calls>", "</|dsml|tool_calls>"), |
| 285 | ("<|dsml|invoke ", "</|dsml|invoke>"), |
| 286 | ("<|tool_calls>", "</|tool_calls>"), |
| 287 | // DeepSeek native, fullwidth delimiters, U+2581 separator. |
| 288 | ("<｜tool▁calls▁begin｜>", "<｜tool▁calls▁end｜>"), |
| 289 | ("<｜tool▁call▁begin｜>", "<｜tool▁call▁end｜>"), |
| 290 | ("<｜tool▁outputs▁begin｜>", "<｜tool▁outputs▁end｜>"), |
| 291 | ("<｜tool▁output▁begin｜>", "<｜tool▁output▁end｜>"), |
| 292 | // DeepSeek native, ASCII delimiters, U+2581 separator. |
| 293 | ("<|tool▁calls▁begin|>", "<|tool▁calls▁end|>"), |
| 294 | ("<|tool▁call▁begin|>", "<|tool▁call▁end|>"), |
| 295 | ("<|tool▁outputs▁begin|>", "<|tool▁outputs▁end|>"), |
| 296 | ("<|tool▁output▁begin|>", "<|tool▁output▁end|>"), |
| 297 | // DeepSeek native, underscore separator. |
| 298 | ("<｜tool_calls_begin｜>", "<｜tool_calls_end｜>"), |
| 299 | ("<｜tool_call_begin｜>", "<｜tool_call_end｜>"), |
| 300 | ("<｜tool_outputs_begin｜>", "<｜tool_outputs_end｜>"), |
| 301 | ("<｜tool_output_begin｜>", "<｜tool_output_end｜>"), |
| 302 | ("<|tool_calls_begin|>", "<|tool_calls_end|>"), |
| 303 | ("<|tool_call_begin|>", "<|tool_call_end|>"), |
| 304 | ("<|tool_outputs_begin|>", "<|tool_outputs_end|>"), |
| 305 | ("<|tool_output_begin|>", "<|tool_output_end|>"), |
| 306 | ]; |
| 307 | |
| 308 | pub(crate) const TOOL_CALL_START_MARKERS: [&str; 28] = [ |
| 309 | "[TOOL_CALL]", |
| 310 | "<codewhale:tool_call", |
| 311 | "<tool_call", |
| 312 | "<invoke ", |
| 313 | "<function_calls>", |
| 314 | "<｜DSML｜tool_calls>", |
| 315 | "<｜DSML｜invoke ", |
| 316 | "<|DSML|tool_calls>", |
| 317 | "<|DSML|invoke ", |
| 318 | "<|dsml|tool_calls>", |
| 319 | "<|dsml|invoke ", |
| 320 | "<|tool_calls>", |
| 321 | "<｜tool▁calls▁begin｜>", |
| 322 | "<｜tool▁call▁begin｜>", |
| 323 | "<｜tool▁outputs▁begin｜>", |
| 324 | "<｜tool▁output▁begin｜>", |
| 325 | "<|tool▁calls▁begin|>", |
| 326 | "<|tool▁call▁begin|>", |
| 327 | "<|tool▁outputs▁begin|>", |
| 328 | "<|tool▁output▁begin|>", |
| 329 | "<｜tool_calls_begin｜>", |
| 330 | "<｜tool_call_begin｜>", |
| 331 | "<｜tool_outputs_begin｜>", |
| 332 | "<｜tool_output_begin｜>", |
| 333 | "<|tool_calls_begin|>", |
| 334 | "<|tool_call_begin|>", |
| 335 | "<|tool_outputs_begin|>", |
| 336 | "<|tool_output_begin|>", |
| 337 | ]; |
| 338 | |
| 339 | pub(crate) const TOOL_CALL_END_MARKERS: [&str; 28] = [ |
| 340 | "[/TOOL_CALL]", |
| 341 | "</codewhale:tool_call>", |
| 342 | "</tool_call>", |
| 343 | "</invoke>", |
| 344 | "</function_calls>", |
| 345 | "</｜DSML｜tool_calls>", |
| 346 | "</｜DSML｜invoke>", |
| 347 | "</|DSML|tool_calls>", |
| 348 | "</|DSML|invoke>", |
| 349 | "</|dsml|tool_calls>", |
| 350 | "</|dsml|invoke>", |
| 351 | "</|tool_calls>", |
| 352 | "<｜tool▁calls▁end｜>", |
| 353 | "<｜tool▁call▁end｜>", |
| 354 | "<｜tool▁outputs▁end｜>", |
| 355 | "<｜tool▁output▁end｜>", |
| 356 | "<|tool▁calls▁end|>", |
| 357 | "<|tool▁call▁end|>", |
| 358 | "<|tool▁outputs▁end|>", |
| 359 | "<|tool▁output▁end|>", |
| 360 | "<｜tool_calls_end｜>", |
| 361 | "<｜tool_call_end｜>", |
| 362 | "<｜tool_outputs_end｜>", |
| 363 | "<｜tool_output_end｜>", |
| 364 | "<|tool_calls_end|>", |
| 365 | "<|tool_call_end|>", |
| 366 | "<|tool_outputs_end|>", |
| 367 | "<|tool_output_end|>", |
| 368 | ]; |
| 369 | |
| 370 | #[derive(Debug, Default)] |
| 371 | pub(crate) struct ToolCallDeltaFilterState { |
| 372 | in_tool_call: bool, |
| 373 | marker_carry: String, |
| 374 | active_end_marker: Option<&'static str>, |
| 375 | } |
| 376 | |
| 377 | /// Compact one-shot notice emitted when a model attempts to forge a tool-call |
| 378 | /// wrapper in plain text instead of using the API tool channel. The visible |
| 379 | /// content is still scrubbed; this exists so the user can see why their text |
| 380 | /// shrank. |
| 381 | pub(crate) const FAKE_WRAPPER_NOTICE: &str = |
| 382 | "Stripped non-API tool-call wrapper from model output (use the API tool channel)"; |
| 383 | |
| 384 | /// True if `text` contains any of the known fake-wrapper start markers. Used by |
| 385 | /// the streaming loop to decide whether to emit `FAKE_WRAPPER_NOTICE`. |
| 386 | pub(crate) fn contains_fake_tool_wrapper(text: &str) -> bool { |
| 387 | TOOL_CALL_START_MARKERS.iter().any(|m| text.contains(m)) |
| 388 | } |
| 389 | |
| 390 | fn find_first_marker(text: &str, markers: &[&str]) -> Option<(usize, usize)> { |
| 391 | markers |
| 392 | .iter() |
| 393 | .filter_map(|marker| text.find(marker).map(|idx| (idx, marker.len()))) |
| 394 | .min_by_key(|(idx, _)| *idx) |
| 395 | } |
| 396 | |
| 397 | fn find_first_start_marker(text: &str) -> Option<(usize, usize, &'static str)> { |
| 398 | TOOL_CALL_MARKER_PAIRS |
| 399 | .iter() |
| 400 | .filter_map(|(start, end)| text.find(start).map(|idx| (idx, start.len(), *end))) |
| 401 | .min_by_key(|(idx, _, _)| *idx) |
| 402 | } |
| 403 | |
| 404 | /// Cheap rejection: every marker prefix ends with the marker's own first |
| 405 | /// byte, so a text without any marker's first byte cannot end with one. |
| 406 | /// Every tool-call marker starts with `<` or `[`, so this is one scan of a |
| 407 | /// usually short delta and keeps per-token `ends_with` probing off the hot |
| 408 | /// path for plain prose deltas. |
| 409 | fn fast_reject_marker_text(text: &str) -> bool { |
| 410 | !text.bytes().any(|b| b == b'<' || b == b'[') |
| 411 | } |
| 412 | |
| 413 | fn trailing_marker_prefix_len(text: &str, markers: &[&str]) -> usize { |
| 414 | if fast_reject_marker_text(text) { |
| 415 | return 0; |
| 416 | } |
| 417 | markers |
| 418 | .iter() |
| 419 | .flat_map(|marker| { |
| 420 | marker |
| 421 | .char_indices() |
| 422 | .map(|(idx, _)| idx) |
| 423 | .filter(|idx| *idx > 0) |
| 424 | .chain(std::iter::once(marker.len())) |
| 425 | .filter(|idx| *idx < marker.len()) |
| 426 | .filter(|idx| { |
| 427 | let prefix = &marker[..*idx]; |
| 428 | text.ends_with(prefix) |
| 429 | }) |
| 430 | }) |
| 431 | .max() |
| 432 | .unwrap_or(0) |
| 433 | } |
| 434 | |
| 435 | fn trailing_start_marker_prefix_len(text: &str) -> usize { |
| 436 | if fast_reject_marker_text(text) { |
| 437 | return 0; |
| 438 | } |
| 439 | TOOL_CALL_MARKER_PAIRS |
| 440 | .iter() |
| 441 | .flat_map(|(marker, _)| { |
| 442 | marker |
| 443 | .char_indices() |
| 444 | .map(|(idx, _)| idx) |
| 445 | .filter(|idx| *idx > 0) |
| 446 | .chain(std::iter::once(marker.len())) |
| 447 | .filter(|idx| *idx < marker.len()) |
| 448 | .filter(|idx| { |
| 449 | let prefix = &marker[..*idx]; |
| 450 | text.ends_with(prefix) |
| 451 | }) |
| 452 | }) |
| 453 | .max() |
| 454 | .unwrap_or(0) |
| 455 | } |
| 456 | |
| 457 | #[cfg(test)] |
| 458 | pub(crate) fn filter_tool_call_delta(delta: &str, in_tool_call: &mut bool) -> String { |
| 459 | let mut state = ToolCallDeltaFilterState { |
| 460 | in_tool_call: *in_tool_call, |
| 461 | ..ToolCallDeltaFilterState::default() |
| 462 | }; |
| 463 | let output = filter_tool_call_delta_with_state(delta, &mut state); |
| 464 | *in_tool_call = state.in_tool_call; |
| 465 | output |
| 466 | } |
| 467 | |
| 468 | pub(crate) fn filter_tool_call_delta_with_state( |
| 469 | delta: &str, |
| 470 | state: &mut ToolCallDeltaFilterState, |
| 471 | ) -> String { |
| 472 | if delta.is_empty() { |
| 473 | return String::new(); |
| 474 | } |
| 475 | |
| 476 | let chunk; |
| 477 | let mut rest = if state.marker_carry.is_empty() { |
| 478 | delta |
| 479 | } else { |
| 480 | chunk = format!("{}{delta}", state.marker_carry); |
| 481 | state.marker_carry.clear(); |
| 482 | &chunk |
| 483 | }; |
| 484 | let mut output = String::new(); |
| 485 | |
| 486 | loop { |
| 487 | if state.in_tool_call { |
| 488 | let active_end_marker = state.active_end_marker; |
| 489 | let found = active_end_marker |
| 490 | .and_then(|marker| rest.find(marker).map(|idx| (idx, marker.len()))) |
| 491 | .or_else(|| find_first_marker(rest, &TOOL_CALL_END_MARKERS)); |
| 492 | let Some((idx, len)) = found else { |
| 493 | let keep = active_end_marker.map_or_else( |
| 494 | || trailing_marker_prefix_len(rest, &TOOL_CALL_END_MARKERS), |
| 495 | |marker| trailing_marker_prefix_len(rest, &[marker]), |
| 496 | ); |
| 497 | if keep > 0 { |
| 498 | state.marker_carry.push_str(&rest[rest.len() - keep..]); |
| 499 | } |
| 500 | break; |
| 501 | }; |
| 502 | rest = &rest[idx + len..]; |
| 503 | state.in_tool_call = false; |
| 504 | state.active_end_marker = None; |
| 505 | } else { |
| 506 | let Some((idx, len, end_marker)) = find_first_start_marker(rest) else { |
| 507 | let keep = trailing_start_marker_prefix_len(rest); |
| 508 | if keep > 0 { |
| 509 | let split = rest.len() - keep; |
| 510 | output.push_str(&rest[..split]); |
| 511 | state.marker_carry.push_str(&rest[split..]); |
| 512 | } else { |
| 513 | output.push_str(rest); |
| 514 | } |
| 515 | break; |
| 516 | }; |
| 517 | output.push_str(&rest[..idx]); |
| 518 | rest = &rest[idx + len..]; |
| 519 | state.in_tool_call = true; |
| 520 | state.active_end_marker = Some(end_marker); |
| 521 | } |
| 522 | } |
| 523 | |
| 524 | output |
| 525 | } |
| 526 | |
| 527 | pub(crate) fn flush_tool_call_delta_state(state: &mut ToolCallDeltaFilterState) -> String { |
| 528 | if state.in_tool_call { |
| 529 | state.marker_carry.clear(); |
| 530 | return String::new(); |
| 531 | } |
| 532 | std::mem::take(&mut state.marker_carry) |
| 533 | } |
| 534 |