返回 CodeWhale
stream_entry.rs
根目录 / crates / tui / src / client / stream_entry.rs
1 //! Shared stream entry seam for Chat Completions / Anthropic Messages / Responses.
2 //!
3 //! Scoped consolidation for v0.9.1: wire-protocol adapters stay at the edge
4 //! (`chat.rs`, `anthropic.rs`, `responses.rs`); this module owns the common
5 //! open path, HTTP/1.1 fallback policy, and idle-timeout envelope so providers
6 //! do not re-implement transport differently.
7 //!
8 //! Full piagent-style provider collapse is deferred — see
9 //! `docs/notes/post-0.9.1-thin-tui-and-stream.md`.
10
11 use std::future::Future;
12 use std::time::Duration;
13
14 use anyhow::Result;
15 use reqwest::Client;
16
17 /// Default bounded wait for SSE response headers. Intentionally shorter than
18 /// the per-chunk idle timeout: it covers connection setup and upstream header
19 /// return only, never model thinking time after streaming has started.
20 pub(crate) const DEFAULT_STREAM_OPEN_TIMEOUT: Duration = Duration::from_secs(45);
21
22 /// Env override (`CODEWHALE_STREAM_OPEN_TIMEOUT_SECS`, legacy
23 /// `DEEPSEEK_STREAM_OPEN_TIMEOUT_SECS`) for the response-header wait,
24 /// shared by every streaming adapter.
25 pub(crate) fn stream_open_timeout() -> Duration {
26 stream_open_timeout_from_env(
27 std::env::var("CODEWHALE_STREAM_OPEN_TIMEOUT_SECS")
28 .or_else(|_| std::env::var("DEEPSEEK_STREAM_OPEN_TIMEOUT_SECS"))
29 .ok()
30 .as_deref(),
31 )
32 }
33
34 pub(crate) fn stream_open_timeout_from_env(value: Option<&str>) -> Duration {
35 let secs = value
36 .and_then(|v| v.parse::<u64>().ok())
37 .unwrap_or(DEFAULT_STREAM_OPEN_TIMEOUT.as_secs())
38 .clamp(5, 300);
39 Duration::from_secs(secs)
40 }
41
42 /// How the shared stream open path should pin HTTP version.
43 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
44 pub enum StreamHttpPolicy {
45 /// Prefer the dual client (H2 primary, H1 twin for fallback).
46 DualWithH1Fallback,
47 /// Force HTTP/1.1 only (env pin or prior H2 stall).
48 Http1Only,
49 }
50
51 /// Inputs shared by every streaming provider adapter at open time.
52 #[derive(Debug, Clone)]
53 pub struct StreamOpenRequest {
54 pub policy: StreamHttpPolicy,
55 pub open_timeout: Duration,
56 pub idle_timeout: Duration,
57 }
58
59 impl StreamOpenRequest {
60 #[must_use]
61 pub fn new(open_timeout: Duration, idle_timeout: Duration) -> Self {
62 Self {
63 policy: if super::force_http1_from_env() {
64 StreamHttpPolicy::Http1Only
65 } else {
66 StreamHttpPolicy::DualWithH1Fallback
67 },
68 open_timeout,
69 idle_timeout,
70 }
71 }
72
73 /// After an H2 stall, retry on the HTTP/1.1 twin.
74 #[must_use]
75 pub fn with_h1_only(mut self) -> Self {
76 self.policy = StreamHttpPolicy::Http1Only;
77 self
78 }
79 }
80
81 /// Select the HTTP client for a stream open attempt.
82 #[must_use]
83 pub fn client_for_policy<'a>(
84 primary: &'a Client,
85 http1_fallback: &'a Client,
86 policy: StreamHttpPolicy,
87 ) -> &'a Client {
88 match policy {
89 StreamHttpPolicy::DualWithH1Fallback => primary,
90 StreamHttpPolicy::Http1Only => http1_fallback,
91 }
92 }
93
94 /// Whether a transport error should trigger H1 fallback retry.
95 #[must_use]
96 pub fn should_retry_with_h1(policy: StreamHttpPolicy, err_text: &str) -> bool {
97 if policy != StreamHttpPolicy::DualWithH1Fallback {
98 return false;
99 }
100 let lower = err_text.to_ascii_lowercase();
101 lower.contains("http2")
102 || lower.contains("h2 ")
103 || lower.contains("stream closed")
104 || lower.contains("connection reset")
105 || lower.contains("protocol error")
106 || lower.contains("frame size")
107 }
108
109 /// Open an SSE response through the shared transport policy.
110 ///
111 /// `attempt` builds and sends one wire-specific request on the client
112 /// selected for the given policy (via [`client_for_policy`]); everything
113 /// transport-shared lives here:
114 ///
115 /// - the response-header wait is bounded by `open_req.open_timeout`;
116 /// - a header stall on the dual client retries exactly once on the
117 /// HTTP/1.1 twin ([`should_retry_with_h1`] classification);
118 /// - a stall on an already H1-pinned request never retries;
119 /// - once response headers have been received the seam never retries —
120 /// body/stream errors belong to the adapter's decode loop.
121 pub(crate) async fn open_sse_response<F, Fut>(
122 open_req: &StreamOpenRequest,
123 attempt: F,
124 ) -> Result<reqwest::Response>
125 where
126 F: Fn(StreamHttpPolicy) -> Fut,
127 Fut: Future<Output = Result<reqwest::Response>>,
128 {
129 match tokio::time::timeout(open_req.open_timeout, attempt(open_req.policy)).await {
130 Ok(result) => result,
131 Err(_elapsed) => {
132 // A header stall on the dual client is eligible for one explicit
133 // retry through the prebuilt HTTP/1.1 twin.
134 if should_retry_with_h1(open_req.policy, "http2 stream closed") {
135 let h1_req = open_req.clone().with_h1_only();
136 crate::logging::warn(
137 "SSE stream headers timed out over HTTP/2; retrying once with HTTP/1.1",
138 );
139 match tokio::time::timeout(h1_req.open_timeout, attempt(h1_req.policy)).await {
140 Ok(Ok(response)) => Ok(response),
141 Ok(Err(err)) => Err(anyhow::anyhow!(
142 "SSE stream request failed after HTTP/1.1 fallback: {err}. \
143 `codewhale doctor` can still pass when non-streaming requests work; \
144 on Windows or proxy networks, try `CODEWHALE_FORCE_HTTP1=1` and rerun `codewhale`."
145 )),
146 Err(_elapsed) => Err(anyhow::anyhow!(
147 "SSE stream request did not receive response headers after {}s \
148 (HTTP/2 and HTTP/1.1). `codewhale doctor` can still pass when \
149 non-streaming requests work; try `CODEWHALE_FORCE_HTTP1=1` and \
150 rerun `codewhale`.",
151 open_req.open_timeout.as_secs()
152 )),
153 }
154 } else {
155 Err(anyhow::anyhow!(
156 "SSE stream request did not receive response headers after {}s. \
157 `codewhale doctor` can still pass when non-streaming requests work; \
158 on Windows or proxy networks, try `CODEWHALE_FORCE_HTTP1=1` and rerun `codewhale`.",
159 open_req.open_timeout.as_secs()
160 ))
161 }
162 }
163 }
164 }
165
166 /// Format a stable idle-timeout message shared across adapters.
167 #[must_use]
168 pub fn idle_timeout_message(
169 idle: Duration,
170 bytes_received: usize,
171 stream_age: Duration,
172 since_last_chunk: Duration,
173 ) -> String {
174 format!(
175 "SSE stream idle timeout after {}s — no data received \
176 (bytes_received={}, stream_age_ms={}, ms_since_last_chunk={})",
177 idle.as_secs(),
178 bytes_received,
179 stream_age.as_millis(),
180 since_last_chunk.as_millis(),
181 )
182 }
183
184 #[cfg(test)]
185 mod tests {
186 use std::sync::Arc;
187 use std::sync::atomic::{AtomicUsize, Ordering};
188
189 use wiremock::matchers::method;
190 use wiremock::{Mock, MockServer, ResponseTemplate};
191
192 use super::*;
193
194 fn open_req(policy: StreamHttpPolicy, open_timeout: Duration) -> StreamOpenRequest {
195 StreamOpenRequest {
196 policy,
197 open_timeout,
198 idle_timeout: Duration::from_secs(30),
199 }
200 }
201
202 async fn ok_server() -> MockServer {
203 let server = MockServer::start().await;
204 Mock::given(method("POST"))
205 .respond_with(ResponseTemplate::new(200))
206 .mount(&server)
207 .await;
208 server
209 }
210
211 #[tokio::test]
212 async fn open_returns_first_attempt_response_on_dual_policy() {
213 let server = ok_server().await;
214 let _ = rustls::crypto::ring::default_provider().install_default();
215 let client = reqwest::Client::new();
216 let attempts = Arc::new(AtomicUsize::new(0));
217 let response = open_sse_response(
218 &open_req(StreamHttpPolicy::DualWithH1Fallback, Duration::from_secs(5)),
219 |policy| {
220 assert_eq!(policy, StreamHttpPolicy::DualWithH1Fallback);
221 let attempts = Arc::clone(&attempts);
222 let client = client.clone();
223 let url = server.uri();
224 async move {
225 attempts.fetch_add(1, Ordering::SeqCst);
226 Ok(client.post(url).send().await?)
227 }
228 },
229 )
230 .await
231 .expect("first attempt succeeds");
232 assert_eq!(response.status(), 200);
233 assert_eq!(attempts.load(Ordering::SeqCst), 1);
234 }
235
236 #[tokio::test]
237 async fn header_stall_on_dual_policy_retries_exactly_once_on_h1() {
238 let server = ok_server().await;
239 let _ = rustls::crypto::ring::default_provider().install_default();
240 let client = reqwest::Client::new();
241 let attempts = Arc::new(AtomicUsize::new(0));
242 let response = open_sse_response(
243 &open_req(
244 StreamHttpPolicy::DualWithH1Fallback,
245 Duration::from_millis(150),
246 ),
247 |policy| {
248 let attempts = Arc::clone(&attempts);
249 let client = client.clone();
250 let url = server.uri();
251 async move {
252 let attempt = attempts.fetch_add(1, Ordering::SeqCst);
253 if attempt == 0 {
254 // First attempt stalls before response headers.
255 assert_eq!(policy, StreamHttpPolicy::DualWithH1Fallback);
256 std::future::pending::<()>().await;
257 }
258 assert_eq!(policy, StreamHttpPolicy::Http1Only);
259 Ok(client.post(url).send().await?)
260 }
261 },
262 )
263 .await
264 .expect("H1 fallback retry succeeds");
265 assert_eq!(response.status(), 200);
266 assert_eq!(
267 attempts.load(Ordering::SeqCst),
268 2,
269 "exactly one fallback retry"
270 );
271 }
272
273 #[tokio::test]
274 async fn header_stall_when_h1_pinned_never_retries_and_reports_timeout_text() {
275 let attempts = Arc::new(AtomicUsize::new(0));
276 let err = open_sse_response(
277 &open_req(StreamHttpPolicy::Http1Only, Duration::from_millis(100)),
278 |_| {
279 let attempts = Arc::clone(&attempts);
280 async move {
281 attempts.fetch_add(1, Ordering::SeqCst);
282 std::future::pending::<()>().await;
283 unreachable!("stalled attempt never resolves")
284 }
285 },
286 )
287 .await
288 .expect_err("H1-pinned stall fails without retry");
289 assert_eq!(attempts.load(Ordering::SeqCst), 1, "no retry when pinned");
290 let text = err.to_string();
291 assert!(text.contains("did not receive response headers"), "{text}");
292 assert!(text.contains("CODEWHALE_FORCE_HTTP1=1"), "{text}");
293 assert!(
294 !text.contains("HTTP/2 and HTTP/1.1"),
295 "single-protocol stall must not claim a dual-protocol attempt: {text}"
296 );
297 }
298
299 #[tokio::test]
300 async fn attempt_error_before_headers_is_not_h1_retried() {
301 let attempts = Arc::new(AtomicUsize::new(0));
302 let err = open_sse_response(
303 &open_req(StreamHttpPolicy::DualWithH1Fallback, Duration::from_secs(5)),
304 |_| {
305 let attempts = Arc::clone(&attempts);
306 async move {
307 attempts.fetch_add(1, Ordering::SeqCst);
308 Err(anyhow::anyhow!("HTTP 401: invalid api key"))
309 }
310 },
311 )
312 .await
313 .expect_err("provider error propagates");
314 assert_eq!(
315 attempts.load(Ordering::SeqCst),
316 1,
317 "non-stall errors are never H1-retried"
318 );
319 assert!(err.to_string().contains("HTTP 401"), "{err}");
320 }
321
322 #[tokio::test]
323 async fn double_stall_reports_both_protocols_in_timeout_text() {
324 let attempts = Arc::new(AtomicUsize::new(0));
325 let err = open_sse_response(
326 &open_req(
327 StreamHttpPolicy::DualWithH1Fallback,
328 Duration::from_millis(100),
329 ),
330 |_| {
331 let attempts = Arc::clone(&attempts);
332 async move {
333 attempts.fetch_add(1, Ordering::SeqCst);
334 std::future::pending::<()>().await;
335 unreachable!("stalled attempt never resolves")
336 }
337 },
338 )
339 .await
340 .expect_err("double stall fails");
341 assert_eq!(attempts.load(Ordering::SeqCst), 2, "one fallback, no more");
342 let text = err.to_string();
343 assert!(text.contains("HTTP/2 and HTTP/1.1"), "{text}");
344 }
345
346 #[test]
347 fn h1_retry_only_on_dual_policy() {
348 assert!(should_retry_with_h1(
349 StreamHttpPolicy::DualWithH1Fallback,
350 "http2 protocol error"
351 ));
352 assert!(!should_retry_with_h1(
353 StreamHttpPolicy::Http1Only,
354 "http2 protocol error"
355 ));
356 }
357
358 #[test]
359 fn stream_open_timeout_defaults_and_clamps_env_values() {
360 assert_eq!(stream_open_timeout_from_env(None), Duration::from_secs(45));
361 assert_eq!(
362 stream_open_timeout_from_env(Some("not-a-number")),
363 Duration::from_secs(45)
364 );
365 assert_eq!(
366 stream_open_timeout_from_env(Some("1")),
367 Duration::from_secs(5)
368 );
369 assert_eq!(
370 stream_open_timeout_from_env(Some("120")),
371 Duration::from_secs(120)
372 );
373 assert_eq!(
374 stream_open_timeout_from_env(Some("999")),
375 Duration::from_secs(300)
376 );
377 }
378
379 #[test]
380 fn idle_message_is_stable() {
381 let msg = idle_timeout_message(
382 Duration::from_secs(30),
383 0,
384 Duration::from_secs(30),
385 Duration::from_secs(30),
386 );
387 assert!(msg.contains("idle timeout"));
388 assert!(msg.contains("bytes_received=0"));
389 }
390 }
391
391 lines RUST