| 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 crate::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 | pub(super) input_buffer: String, |
| 24 | pub(super) input_parse_error: Option<String>, |
| 25 | } |
| 26 | |
| 27 | /// Maximum total bytes of text/thinking content before aborting the stream. |
| 28 | pub(super) const STREAM_MAX_CONTENT_BYTES: usize = 10 * 1024 * 1024; // 10 MB |
| 29 | /// Sanity backstop for total stream wall-clock duration. **Not** a routine |
| 30 | /// kill switch — the stream chunk idle timeout is the primary stall |
| 31 | /// detector. The wall-clock cap is here only to bound pathological cases |
| 32 | /// (e.g. a server that keeps sending heartbeats forever without progress). |
| 33 | /// |
| 34 | /// History: this used to be 300s (5 min) which was too aggressive — V4 |
| 35 | /// thinking turns on hard prompts legitimately exceed 5 minutes wall-clock |
| 36 | /// while still emitting reasoning_content chunks the whole way. Bumped to |
| 37 | /// 30 min in v0.6.6 after long-reasoning turns hit the old cap. Codex defaults to a |
| 38 | /// per-chunk idle of 300s with no wall-clock cap; we keep both layers but |
| 39 | /// give the wall-clock a generous window so it never fires in practice. |
| 40 | pub(super) const STREAM_MAX_DURATION_SECS: u64 = 1800; // 30 minutes (was 300s; #103/#1) |
| 41 | /// Hard cap on consecutive recoverable stream errors before we surface a turn |
| 42 | /// failure. Bumped 3 → 5 in v0.6.7 along with the HTTP/2 keepalive defaults |
| 43 | /// (#103) — keepalive should make spurious decode errors rarer, so we can |
| 44 | /// tolerate a longer streak before giving up on the turn. |
| 45 | pub(super) const MAX_STREAM_ERRORS_BEFORE_FAIL: u32 = 5; |
| 46 | /// Cap on transparent stream-level retries — these only happen when the wire |
| 47 | /// dies before any content was streamed, so DeepSeek hasn't billed us and |
| 48 | /// the user hasn't seen anything. Two attempts is enough to ride out a |
| 49 | /// flaky edge node without amplifying real outages (#103). |
| 50 | pub(super) const MAX_TRANSPARENT_STREAM_RETRIES: u32 = 2; |
| 51 | |
| 52 | /// Decide whether a stream error is eligible for a transparent retry. |
| 53 | /// |
| 54 | /// True only when ALL three conditions hold: |
| 55 | /// 1. No content has been received on the current attempt — otherwise DeepSeek |
| 56 | /// has already billed us for output tokens and the user has seen partial |
| 57 | /// deltas; resending would double-bill and desync the UI. |
| 58 | /// 2. We still have transparent-retry budget remaining. |
| 59 | /// 3. The turn has not been cancelled. |
| 60 | /// |
| 61 | /// Extracted as a pure function so the four #103 retry cases can be exercised |
| 62 | /// in unit tests without booting the full engine state machine. |
| 63 | pub(super) fn should_transparently_retry_stream( |
| 64 | any_content_received: bool, |
| 65 | transparent_attempts: u32, |
| 66 | cancelled: bool, |
| 67 | ) -> bool { |
| 68 | !any_content_received && transparent_attempts < MAX_TRANSPARENT_STREAM_RETRIES && !cancelled |
| 69 | } |
| 70 | |
| 71 | /// Budget for re-issuing the whole request after a dead stream. Shared by the |
| 72 | /// nothing-streamed outer retry (#103 Phase 3) and the sleep-resume retry |
| 73 | /// (#2990). |
| 74 | pub(super) const MAX_STREAM_RETRIES: u32 = 3; |
| 75 | |
| 76 | /// Wall-clock vs monotonic divergence above which we conclude the host slept |
| 77 | /// mid-stream (#2990). `Instant` pauses during system sleep (CLOCK_UPTIME_RAW |
| 78 | /// on macOS, CLOCK_MONOTONIC on Linux) while `SystemTime` keeps advancing, so |
| 79 | /// a large positive gap can only come from a suspend/resume cycle — ordinary |
| 80 | /// network flakes never produce one. Windows `Instant` may keep ticking |
| 81 | /// through sleep, in which case this simply never fires (no behavior change). |
| 82 | pub(super) const SLEEP_GAP_THRESHOLD: Duration = Duration::from_secs(10); |
| 83 | |
| 84 | /// True when the gap between wall-clock and monotonic elapsed time since the |
| 85 | /// last stream progress says the host was suspended. |
| 86 | pub(super) fn sleep_gap_detected(monotonic_elapsed: Duration, wallclock_elapsed: Duration) -> bool { |
| 87 | wallclock_elapsed.saturating_sub(monotonic_elapsed) > SLEEP_GAP_THRESHOLD |
| 88 | } |
| 89 | |
| 90 | /// Decide whether a failed stream should be silently re-issued because the |
| 91 | /// host slept mid-turn (#2990). |
| 92 | /// |
| 93 | /// Unlike the transparent retry (#103), this fires even after content has |
| 94 | /// streamed: the partial output predates the sleep, the user was not |
| 95 | /// watching, and re-running the identical request is the correct |
| 96 | /// user-visible behavior. The double-billing concern that blocks ordinary |
| 97 | /// post-content retries is accepted here because the alternative is a dead |
| 98 | /// turn the user must re-prompt (and pay for) anyway. |
| 99 | pub(super) fn should_resume_after_sleep( |
| 100 | sleep_detected: bool, |
| 101 | retry_attempts: u32, |
| 102 | cancelled: bool, |
| 103 | ) -> bool { |
| 104 | sleep_detected && retry_attempts < MAX_STREAM_RETRIES && !cancelled |
| 105 | } |
| 106 | |
| 107 | /// Decide whether a failed stream should be re-issued after a mid-stream |
| 108 | /// network drop in a headless host (`exec` / stream-json / app-server), even |
| 109 | /// though content already streamed. |
| 110 | /// |
| 111 | /// This extends the #2990 sleep-resume contract to ordinary transport drops |
| 112 | /// for hosts with no operator watching: the partial assistant fragment has |
| 113 | /// not been committed to the conversation and no tool call from the |
| 114 | /// incomplete response has executed, so discarding the fragment and |
| 115 | /// re-issuing the identical request cannot duplicate side effects. The |
| 116 | /// double-billing risk that blocks post-content retries in the interactive |
| 117 | /// TUI (#103) is accepted here because the alternative is a dead turn that |
| 118 | /// forfeits the entire headless run — the exact tradeoff #2990 already makes |
| 119 | /// for sleep-resume. Interactive sessions keep the #103 surface-the-warning |
| 120 | /// behavior: the user saw the partial deltas, and replaying would render the |
| 121 | /// same prefix twice. |
| 122 | pub(super) fn should_resume_after_network_drop( |
| 123 | headless_host: bool, |
| 124 | network_class_error: bool, |
| 125 | retry_attempts: u32, |
| 126 | cancelled: bool, |
| 127 | ) -> bool { |
| 128 | headless_host && network_class_error && retry_attempts < MAX_STREAM_RETRIES && !cancelled |
| 129 | } |
| 130 | |
| 131 | /// Convert low-level reqwest/hyper stream read errors into an operator-facing |
| 132 | /// message. The raw provider error remains attached, but the lead sentence |
| 133 | /// explains why Codewhale may retry before any output and why it must surface |
| 134 | /// the warning once partial output has already streamed. |
| 135 | pub(super) fn stream_read_error_user_message(message: &str, any_content_received: bool) -> String { |
| 136 | let lower = message.to_ascii_lowercase(); |
| 137 | let is_stream_read = lower.contains("stream read error") |
| 138 | || lower.contains("error decoding response body") |
| 139 | || lower.contains("chunk decode error") |
| 140 | || lower.contains("body decode"); |
| 141 | if !is_stream_read { |
| 142 | return message.to_string(); |
| 143 | } |
| 144 | |
| 145 | let retry_note = if any_content_received { |
| 146 | "Some output had already streamed, so Codewhale is surfacing the warning instead of replaying the request and risking duplicated output." |
| 147 | } else { |
| 148 | "No output had streamed yet, so Codewhale will retry automatically while retry budget remains." |
| 149 | }; |
| 150 | format!( |
| 151 | "Provider stream connection dropped while reading the response body. {retry_note} Details: {message}" |
| 152 | ) |
| 153 | } |
| 154 | |
| 155 | /// Wrapper shapes a model may emit as plain text instead of using the API tool |
| 156 | /// channel. Each pair is `(start, end)`; the tables below are projections of |
| 157 | /// this one and must stay in sync with it. |
| 158 | /// |
| 159 | /// Three families are covered: |
| 160 | /// |
| 161 | /// 1. Generic/Anthropic-style (`[TOOL_CALL]`, `<invoke …>`, `<function_calls>`). |
| 162 | /// 2. DSML wrappers, in fullwidth `｜` (U+FF5C) and ASCII `|` delimiters, upper |
| 163 | /// and lower case. |
| 164 | /// 3. **DeepSeek's native tool-call tokens** (#3880). DeepSeek's chat template |
| 165 | /// separates words with `▁` (U+2581 LOWER ONE EIGHTH BLOCK), not a space or |
| 166 | /// underscore, so `<｜tool▁calls▁begin｜>` does not match any DSML entry and |
| 167 | /// leaked into visible output. Both the `▁` and `_` separators are listed |
| 168 | /// because a partially-normalizing tokenizer can emit either, and both |
| 169 | /// delimiter forms because the ASCII fallback shows up in some renderings. |
| 170 | /// |
| 171 | /// When adding a shape, add it here and to the two marker tables below. |
| 172 | /// `marker_tables_are_consistent` enforces that they agree. |
| 173 | pub(crate) const TOOL_CALL_MARKER_PAIRS: [(&str, &str); 28] = [ |
| 174 | ("[TOOL_CALL]", "[/TOOL_CALL]"), |
| 175 | ("<codewhale:tool_call", "</codewhale:tool_call>"), |
| 176 | ("<tool_call", "</tool_call>"), |
| 177 | ("<invoke ", "</invoke>"), |
| 178 | ("<function_calls>", "</function_calls>"), |
| 179 | ("<｜DSML｜tool_calls>", "</｜DSML｜tool_calls>"), |
| 180 | ("<｜DSML｜invoke ", "</｜DSML｜invoke>"), |
| 181 | ("<|DSML|tool_calls>", "</|DSML|tool_calls>"), |
| 182 | ("<|DSML|invoke ", "</|DSML|invoke>"), |
| 183 | ("<|dsml|tool_calls>", "</|dsml|tool_calls>"), |
| 184 | ("<|dsml|invoke ", "</|dsml|invoke>"), |
| 185 | ("<|tool_calls>", "</|tool_calls>"), |
| 186 | // DeepSeek native, fullwidth delimiters, U+2581 separator. |
| 187 | ("<｜tool▁calls▁begin｜>", "<｜tool▁calls▁end｜>"), |
| 188 | ("<｜tool▁call▁begin｜>", "<｜tool▁call▁end｜>"), |
| 189 | ("<｜tool▁outputs▁begin｜>", "<｜tool▁outputs▁end｜>"), |
| 190 | ("<｜tool▁output▁begin｜>", "<｜tool▁output▁end｜>"), |
| 191 | // DeepSeek native, ASCII delimiters, U+2581 separator. |
| 192 | ("<|tool▁calls▁begin|>", "<|tool▁calls▁end|>"), |
| 193 | ("<|tool▁call▁begin|>", "<|tool▁call▁end|>"), |
| 194 | ("<|tool▁outputs▁begin|>", "<|tool▁outputs▁end|>"), |
| 195 | ("<|tool▁output▁begin|>", "<|tool▁output▁end|>"), |
| 196 | // DeepSeek native, underscore separator. |
| 197 | ("<｜tool_calls_begin｜>", "<｜tool_calls_end｜>"), |
| 198 | ("<｜tool_call_begin｜>", "<｜tool_call_end｜>"), |
| 199 | ("<｜tool_outputs_begin｜>", "<｜tool_outputs_end｜>"), |
| 200 | ("<｜tool_output_begin｜>", "<｜tool_output_end｜>"), |
| 201 | ("<|tool_calls_begin|>", "<|tool_calls_end|>"), |
| 202 | ("<|tool_call_begin|>", "<|tool_call_end|>"), |
| 203 | ("<|tool_outputs_begin|>", "<|tool_outputs_end|>"), |
| 204 | ("<|tool_output_begin|>", "<|tool_output_end|>"), |
| 205 | ]; |
| 206 | |
| 207 | pub(crate) const TOOL_CALL_START_MARKERS: [&str; 28] = [ |
| 208 | "[TOOL_CALL]", |
| 209 | "<codewhale:tool_call", |
| 210 | "<tool_call", |
| 211 | "<invoke ", |
| 212 | "<function_calls>", |
| 213 | "<｜DSML｜tool_calls>", |
| 214 | "<｜DSML｜invoke ", |
| 215 | "<|DSML|tool_calls>", |
| 216 | "<|DSML|invoke ", |
| 217 | "<|dsml|tool_calls>", |
| 218 | "<|dsml|invoke ", |
| 219 | "<|tool_calls>", |
| 220 | "<｜tool▁calls▁begin｜>", |
| 221 | "<｜tool▁call▁begin｜>", |
| 222 | "<｜tool▁outputs▁begin｜>", |
| 223 | "<｜tool▁output▁begin｜>", |
| 224 | "<|tool▁calls▁begin|>", |
| 225 | "<|tool▁call▁begin|>", |
| 226 | "<|tool▁outputs▁begin|>", |
| 227 | "<|tool▁output▁begin|>", |
| 228 | "<｜tool_calls_begin｜>", |
| 229 | "<｜tool_call_begin｜>", |
| 230 | "<｜tool_outputs_begin｜>", |
| 231 | "<｜tool_output_begin｜>", |
| 232 | "<|tool_calls_begin|>", |
| 233 | "<|tool_call_begin|>", |
| 234 | "<|tool_outputs_begin|>", |
| 235 | "<|tool_output_begin|>", |
| 236 | ]; |
| 237 | |
| 238 | pub(crate) const TOOL_CALL_END_MARKERS: [&str; 28] = [ |
| 239 | "[/TOOL_CALL]", |
| 240 | "</codewhale:tool_call>", |
| 241 | "</tool_call>", |
| 242 | "</invoke>", |
| 243 | "</function_calls>", |
| 244 | "</｜DSML｜tool_calls>", |
| 245 | "</｜DSML｜invoke>", |
| 246 | "</|DSML|tool_calls>", |
| 247 | "</|DSML|invoke>", |
| 248 | "</|dsml|tool_calls>", |
| 249 | "</|dsml|invoke>", |
| 250 | "</|tool_calls>", |
| 251 | "<｜tool▁calls▁end｜>", |
| 252 | "<｜tool▁call▁end｜>", |
| 253 | "<｜tool▁outputs▁end｜>", |
| 254 | "<｜tool▁output▁end｜>", |
| 255 | "<|tool▁calls▁end|>", |
| 256 | "<|tool▁call▁end|>", |
| 257 | "<|tool▁outputs▁end|>", |
| 258 | "<|tool▁output▁end|>", |
| 259 | "<｜tool_calls_end｜>", |
| 260 | "<｜tool_call_end｜>", |
| 261 | "<｜tool_outputs_end｜>", |
| 262 | "<｜tool_output_end｜>", |
| 263 | "<|tool_calls_end|>", |
| 264 | "<|tool_call_end|>", |
| 265 | "<|tool_outputs_end|>", |
| 266 | "<|tool_output_end|>", |
| 267 | ]; |
| 268 | |
| 269 | #[derive(Debug, Default)] |
| 270 | pub(crate) struct ToolCallDeltaFilterState { |
| 271 | in_tool_call: bool, |
| 272 | marker_carry: String, |
| 273 | active_end_marker: Option<&'static str>, |
| 274 | } |
| 275 | |
| 276 | /// Compact one-shot notice emitted when a model attempts to forge a tool-call |
| 277 | /// wrapper in plain text instead of using the API tool channel. The visible |
| 278 | /// content is still scrubbed; this exists so the user can see why their text |
| 279 | /// shrank. |
| 280 | pub(crate) const FAKE_WRAPPER_NOTICE: &str = |
| 281 | "Stripped non-API tool-call wrapper from model output (use the API tool channel)"; |
| 282 | |
| 283 | /// True if `text` contains any of the known fake-wrapper start markers. Used by |
| 284 | /// the streaming loop to decide whether to emit `FAKE_WRAPPER_NOTICE`. |
| 285 | pub(crate) fn contains_fake_tool_wrapper(text: &str) -> bool { |
| 286 | TOOL_CALL_START_MARKERS.iter().any(|m| text.contains(m)) |
| 287 | } |
| 288 | |
| 289 | fn find_first_marker(text: &str, markers: &[&str]) -> Option<(usize, usize)> { |
| 290 | markers |
| 291 | .iter() |
| 292 | .filter_map(|marker| text.find(marker).map(|idx| (idx, marker.len()))) |
| 293 | .min_by_key(|(idx, _)| *idx) |
| 294 | } |
| 295 | |
| 296 | fn find_first_start_marker(text: &str) -> Option<(usize, usize, &'static str)> { |
| 297 | TOOL_CALL_MARKER_PAIRS |
| 298 | .iter() |
| 299 | .filter_map(|(start, end)| text.find(start).map(|idx| (idx, start.len(), *end))) |
| 300 | .min_by_key(|(idx, _, _)| *idx) |
| 301 | } |
| 302 | |
| 303 | fn trailing_marker_prefix_len(text: &str, markers: &[&str]) -> usize { |
| 304 | markers |
| 305 | .iter() |
| 306 | .flat_map(|marker| { |
| 307 | marker |
| 308 | .char_indices() |
| 309 | .map(|(idx, _)| idx) |
| 310 | .filter(|idx| *idx > 0) |
| 311 | .chain(std::iter::once(marker.len())) |
| 312 | .filter(|idx| *idx < marker.len()) |
| 313 | .filter(|idx| { |
| 314 | let prefix = &marker[..*idx]; |
| 315 | text.ends_with(prefix) |
| 316 | }) |
| 317 | }) |
| 318 | .max() |
| 319 | .unwrap_or(0) |
| 320 | } |
| 321 | |
| 322 | fn trailing_start_marker_prefix_len(text: &str) -> usize { |
| 323 | TOOL_CALL_MARKER_PAIRS |
| 324 | .iter() |
| 325 | .flat_map(|(marker, _)| { |
| 326 | marker |
| 327 | .char_indices() |
| 328 | .map(|(idx, _)| idx) |
| 329 | .filter(|idx| *idx > 0) |
| 330 | .chain(std::iter::once(marker.len())) |
| 331 | .filter(|idx| *idx < marker.len()) |
| 332 | .filter(|idx| { |
| 333 | let prefix = &marker[..*idx]; |
| 334 | text.ends_with(prefix) |
| 335 | }) |
| 336 | }) |
| 337 | .max() |
| 338 | .unwrap_or(0) |
| 339 | } |
| 340 | |
| 341 | #[cfg(test)] |
| 342 | pub(crate) fn filter_tool_call_delta(delta: &str, in_tool_call: &mut bool) -> String { |
| 343 | let mut state = ToolCallDeltaFilterState { |
| 344 | in_tool_call: *in_tool_call, |
| 345 | ..ToolCallDeltaFilterState::default() |
| 346 | }; |
| 347 | let output = filter_tool_call_delta_with_state(delta, &mut state); |
| 348 | *in_tool_call = state.in_tool_call; |
| 349 | output |
| 350 | } |
| 351 | |
| 352 | pub(crate) fn filter_tool_call_delta_with_state( |
| 353 | delta: &str, |
| 354 | state: &mut ToolCallDeltaFilterState, |
| 355 | ) -> String { |
| 356 | if delta.is_empty() { |
| 357 | return String::new(); |
| 358 | } |
| 359 | |
| 360 | let chunk; |
| 361 | let mut rest = if state.marker_carry.is_empty() { |
| 362 | delta |
| 363 | } else { |
| 364 | chunk = format!("{}{delta}", state.marker_carry); |
| 365 | state.marker_carry.clear(); |
| 366 | &chunk |
| 367 | }; |
| 368 | let mut output = String::new(); |
| 369 | |
| 370 | loop { |
| 371 | if state.in_tool_call { |
| 372 | let active_end_marker = state.active_end_marker; |
| 373 | let found = active_end_marker |
| 374 | .and_then(|marker| rest.find(marker).map(|idx| (idx, marker.len()))) |
| 375 | .or_else(|| find_first_marker(rest, &TOOL_CALL_END_MARKERS)); |
| 376 | let Some((idx, len)) = found else { |
| 377 | let keep = active_end_marker.map_or_else( |
| 378 | || trailing_marker_prefix_len(rest, &TOOL_CALL_END_MARKERS), |
| 379 | |marker| trailing_marker_prefix_len(rest, &[marker]), |
| 380 | ); |
| 381 | if keep > 0 { |
| 382 | state.marker_carry.push_str(&rest[rest.len() - keep..]); |
| 383 | } |
| 384 | break; |
| 385 | }; |
| 386 | rest = &rest[idx + len..]; |
| 387 | state.in_tool_call = false; |
| 388 | state.active_end_marker = None; |
| 389 | } else { |
| 390 | let Some((idx, len, end_marker)) = find_first_start_marker(rest) else { |
| 391 | let keep = trailing_start_marker_prefix_len(rest); |
| 392 | if keep > 0 { |
| 393 | let split = rest.len() - keep; |
| 394 | output.push_str(&rest[..split]); |
| 395 | state.marker_carry.push_str(&rest[split..]); |
| 396 | } else { |
| 397 | output.push_str(rest); |
| 398 | } |
| 399 | break; |
| 400 | }; |
| 401 | output.push_str(&rest[..idx]); |
| 402 | rest = &rest[idx + len..]; |
| 403 | state.in_tool_call = true; |
| 404 | state.active_end_marker = Some(end_marker); |
| 405 | } |
| 406 | } |
| 407 | |
| 408 | output |
| 409 | } |
| 410 | |
| 411 | pub(crate) fn flush_tool_call_delta_state(state: &mut ToolCallDeltaFilterState) -> String { |
| 412 | if state.in_tool_call { |
| 413 | state.marker_carry.clear(); |
| 414 | return String::new(); |
| 415 | } |
| 416 | std::mem::take(&mut state.marker_carry) |
| 417 | } |
| 418 |