| 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 | |
| 9 | #[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| 10 | pub(super) enum ContentBlockKind { |
| 11 | Text, |
| 12 | Thinking, |
| 13 | ToolUse, |
| 14 | } |
| 15 | |
| 16 | #[derive(Debug, Clone)] |
| 17 | pub(super) struct ToolUseState { |
| 18 | pub(super) id: String, |
| 19 | pub(super) name: String, |
| 20 | pub(super) input: serde_json::Value, |
| 21 | pub(super) caller: Option<ToolCaller>, |
| 22 | pub(super) input_buffer: String, |
| 23 | } |
| 24 | |
| 25 | /// Maximum time to wait for a single stream chunk before assuming a stall. |
| 26 | /// **This is the idle timeout** — it resets on every SSE chunk, so long |
| 27 | /// thinking turns that ARE producing reasoning_content stay alive. Only a |
| 28 | /// genuine `chunk_timeout` window of silence kills the stream. |
| 29 | pub(super) const STREAM_CHUNK_TIMEOUT_SECS: u64 = 90; |
| 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 — `STREAM_CHUNK_TIMEOUT_SECS` (idle) 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 to address `TODO_FIXES.md` #1. 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, so DeepSeek hasn't billed us and |
| 51 | /// the user hasn't seen anything. Two attempts is enough to 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 — otherwise DeepSeek |
| 59 | /// has already billed us for output tokens and the user has seen partial |
| 60 | /// deltas; resending would double-bill and desync the UI. |
| 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 | pub(crate) const TOOL_CALL_START_MARKERS: [&str; 5] = [ |
| 75 | "[TOOL_CALL]", |
| 76 | "<deepseek:tool_call", |
| 77 | "<tool_call", |
| 78 | "<invoke ", |
| 79 | "<function_calls>", |
| 80 | ]; |
| 81 | |
| 82 | pub(crate) const TOOL_CALL_END_MARKERS: [&str; 5] = [ |
| 83 | "[/TOOL_CALL]", |
| 84 | "</deepseek:tool_call>", |
| 85 | "</tool_call>", |
| 86 | "</invoke>", |
| 87 | "</function_calls>", |
| 88 | ]; |
| 89 | |
| 90 | /// Compact one-shot notice emitted when a model attempts to forge a tool-call |
| 91 | /// wrapper in plain text instead of using the API tool channel. The visible |
| 92 | /// content is still scrubbed; this exists so the user can see why their text |
| 93 | /// shrank. |
| 94 | pub(crate) const FAKE_WRAPPER_NOTICE: &str = |
| 95 | "Stripped non-API tool-call wrapper from model output (use the API tool channel)"; |
| 96 | |
| 97 | /// True if `text` contains any of the known fake-wrapper start markers. Used by |
| 98 | /// the streaming loop to decide whether to emit `FAKE_WRAPPER_NOTICE`. |
| 99 | pub(crate) fn contains_fake_tool_wrapper(text: &str) -> bool { |
| 100 | TOOL_CALL_START_MARKERS.iter().any(|m| text.contains(m)) |
| 101 | } |
| 102 | |
| 103 | fn find_first_marker(text: &str, markers: &[&str]) -> Option<(usize, usize)> { |
| 104 | markers |
| 105 | .iter() |
| 106 | .filter_map(|marker| text.find(marker).map(|idx| (idx, marker.len()))) |
| 107 | .min_by_key(|(idx, _)| *idx) |
| 108 | } |
| 109 | |
| 110 | pub(crate) fn filter_tool_call_delta(delta: &str, in_tool_call: &mut bool) -> String { |
| 111 | if delta.is_empty() { |
| 112 | return String::new(); |
| 113 | } |
| 114 | |
| 115 | let mut output = String::new(); |
| 116 | let mut rest = delta; |
| 117 | |
| 118 | loop { |
| 119 | if *in_tool_call { |
| 120 | let Some((idx, len)) = find_first_marker(rest, &TOOL_CALL_END_MARKERS) else { |
| 121 | break; |
| 122 | }; |
| 123 | rest = &rest[idx + len..]; |
| 124 | *in_tool_call = false; |
| 125 | } else { |
| 126 | let Some((idx, len)) = find_first_marker(rest, &TOOL_CALL_START_MARKERS) else { |
| 127 | output.push_str(rest); |
| 128 | break; |
| 129 | }; |
| 130 | output.push_str(&rest[..idx]); |
| 131 | rest = &rest[idx + len..]; |
| 132 | *in_tool_call = true; |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | output |
| 137 | } |
| 138 |