4use std::cell::{Cell, RefCell};
5use std::collections::VecDeque;
6use std::future::Future;
7use std::pin::pin;
8use std::task::{Context, Poll, Waker};
9use std::time::{Duration, Instant, SystemTime};

11use bytes::Bytes;
12use jev_protocol::{ApiErrorKind, Json, ModelId, Noul, Questions, REQUEST_ID_HEADER};
13
14use crate::{Client, ClientError, Event, HttpRequest, Observer, HttpResponse, Runtime, Transport, TransportError, TransportErrorKind};
15
16enum Step {
17    Reply { status: u16, headers: Vec<(&'static str, &'static str)>, body: String },
18    Fail(TransportErrorKind),

20    Stall,
21}

23fn ok(p: f64, id: &'static str) -> Step {
24    Step::Reply {
25        status: 200,
26        headers: vec![(REQUEST_ID_HEADER, id)],
27        body: format!(
28            r#"{{"model":"jev-1.13.0","answers":{{"q":{{"type":"noul","noul":{p}}}}},"usage":{{"input_tokens":1,"output_tokens":1}}}}"#
29        ),
30    }
31}
32
33fn status(status: u16, headers: Vec<(&'static str, &'static str)>) -> Step {
34    Step::Reply { status, headers, body: r#"{"detail":{"error_type":"e","message":"m"}}"#.into() }
35}
36
37struct Fake {
38    script: RefCell<VecDeque<Step>>,
39    events: RefCell<Vec<String>>,
40    sent: RefCell<Vec<HttpRequest>>,
41    distrusted: Cell<u32>,
42    clock: Cell<Duration>,
43    base: Instant,
44}
45
46struct FakeTransport<'a>(&'a Fake);
47struct FakeRuntime<'a>(&'a Fake);
48struct FakeObserver<'a>(&'a Fake);
49
50impl Observer for FakeObserver<'_> {
51    fn event(&self, event: Event<'_>) {
52        let line = match event {
53            Event::Sending { retry } => format!("send {retry}"),
54            Event::Retrying { attempt, delay, cause, request_id } => {
55                format!("retry after attempt {attempt} in {delay:?} ({cause}) id {request_id:?}")
56            }
57            Event::Redialing { cause } => format!("redial ({cause})"),
58            Event::Answered { request_id, attempts, usage } => {
59                format!("answered id {request_id:?} in {attempts} tokens {}+{}", usage.input_tokens, usage.output_tokens)
60            }
61        };
62        self.0.events.borrow_mut().push(line);
63    }
64}
65
66impl Transport for FakeTransport<'_> {
67    async fn send(&self, request: HttpRequest) -> Result<HttpResponse, TransportError> {
68        self.0.sent.borrow_mut().push(request);
69        let step = self.0.script.borrow_mut().pop_front().expect("the script ran out");
70        match step {
71            Step::Reply { status, headers, body } => {
72                let mut map = http::HeaderMap::new();
73                for (k, v) in headers {
74                    map.insert(k, http::HeaderValue::from_static(v));
75                }
76                Ok(HttpResponse { status, headers: map, body: Bytes::from(body) })
77            }
78            Step::Fail(kind) => Err(TransportError::new(kind, format!("{kind:?}"))),
79            Step::Stall => std::future::pending().await,
80        }
81    }
82
83    fn distrust_connection(&self) {
84        self.0.distrusted.set(self.0.distrusted.get() + 1);
85    }
86}
87
88impl Runtime for FakeRuntime<'_> {
89    fn now(&self) -> Instant {
90        self.0.base + self.0.clock.get()
91    }
92
93    fn wall_clock(&self) -> SystemTime {
94        SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000_000) + self.0.clock.get()
95    }
96
97    async fn sleep(&self, duration: Duration) {
98        self.0.clock.set(self.0.clock.get() + duration);
99    }
100
101    async fn timeout<F: Future>(&self, duration: Duration, future: F) -> Option<F::Output> {
102        // Fakes either answer at once or never: poll once, and a stall
103        // costs exactly the timeout.
104        let mut future = pin!(future);
105        match future.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
106            Poll::Ready(out) => Some(out),
107            Poll::Pending => {
108                self.0.clock.set(self.0.clock.get() + duration);
109                None
110            }
111        }
112    }
113
114    fn jitter(&self) -> f64 {
115        0.0
116    }
117}
118
119fn run(script: Vec<Step>) -> (Result<crate::Answered, ClientError>, Fake) {
120    let fake = Fake {
121        script: RefCell::new(script.into()),
122        events: RefCell::default(),
123        sent: RefCell::default(),
124        distrusted: Cell::default(),
125        clock: Cell::default(),
126        base: Instant::now(),
127    };
128    let result = {
129        let model = ModelId::pinned("jev-1.13.0").unwrap();
130        let client =
131            Client::new(FakeTransport(&fake), FakeRuntime(&fake), model, "k").unwrap().observed(FakeObserver(&fake));
132        let mut questions = Questions::new();
133        questions.noul("q", Noul::new(Json::text("?"))).unwrap();
134        let state = Json::text("s");
135        let future = pin!(client.ask(&state, &questions));
136        match future.poll(&mut Context::from_waker(Waker::noop())) {
137            Poll::Ready(r) => r,
138            Poll::Pending => panic!("fakes never leave the client pending"),
139        }
140    };
141    (result, fake)
142}
143
144fn retry_counts(fake: &Fake) -> Vec<Option<String>> {
145    fake.sent
146        .borrow()
147        .iter()
148        .map(|r| r.headers.get("x-typesafe-retry-count").map(|v| v.to_str().unwrap().to_owned()))
149        .collect()
150}
151
152#[test]
153fn answers_first_time() {
154    let (r, fake) = run(vec![ok(0.95, "a")]);
155    let a = r.unwrap();
156    assert_eq!((a.attempts, a.request_id.as_deref()), (1, Some("a")));
157    assert_eq!(fake.clock.get(), Duration::ZERO);
158    assert!(fake.sent.borrow()[0].headers[http::header::AUTHORIZATION].is_sensitive());
159}
160
161#[test]
162fn honours_the_servers_delay() {
163    let (r, fake) = run(vec![status(429, vec![("retry-after-ms", "300")]), ok(0.95, "b")]);
164    assert_eq!(r.unwrap().attempts, 2);
165    assert_eq!(fake.clock.get(), Duration::from_millis(300));
166    assert_eq!(retry_counts(&fake), vec![None, Some("1".into())]);
167}
168
169#[test]
170fn backs_off_then_gives_up_with_the_last_request_id() {
171    let (r, fake) = run(vec![
172        status(503, vec![(REQUEST_ID_HEADER, "r1")]),
173        status(503, vec![(REQUEST_ID_HEADER, "r2")]),
174        status(503, vec![(REQUEST_ID_HEADER, "r3")]),
175    ]);
176    match r.unwrap_err() {
177        ClientError::Api { error, attempts } => {
178            assert_eq!((error.status, attempts, error.request_id.as_deref()), (503, 3, Some("r3")));
179        }
180        other => panic!("{other:?}"),
181    }
182    assert_eq!(fake.clock.get(), Duration::from_millis(1500));
183    assert_eq!(retry_counts(&fake), vec![None, Some("1".into()), Some("2".into())]);
184}
185
186#[test]
187fn other_4xx_fail_at_once() {
188    let (r, fake) = run(vec![status(403, vec![])]);
189    assert!(matches!(r.unwrap_err(), ClientError::Api { attempts: 1, .. }));
190    assert_eq!(fake.sent.borrow().len(), 1);
191}
192
193#[test]
194fn context_overflow_is_not_resent() {
195    let (r, _) = run(vec![Step::Reply {
196        status: 400,
197        headers: vec![],
198        body: r#"{"detail":{"error_type":"max_tokens_exceeded","message":"too long"}}"#.into(),
199    }]);
200    match r.unwrap_err() {
201        ClientError::Api { error, .. } => assert_eq!(error.kind, ApiErrorKind::ContextOverflow),
202        other => panic!("{other:?}"),
203    }
204}
205
206#[test]
207fn never_sent_is_redialled_not_retried() {
208    let (r, fake) = run(vec![Step::Fail(TransportErrorKind::NotSent), ok(0.95, "c")]);
209    assert_eq!(r.unwrap().attempts, 1);
210    assert_eq!(fake.clock.get(), Duration::ZERO, "no backoff");
211    assert_eq!(retry_counts(&fake), vec![None, None], "not a retry");
212}
213
214#[test]
215fn repeated_never_sent_becomes_a_retry() {
216    use TransportErrorKind::NotSent;
217    let (r, fake) = run(vec![Step::Fail(NotSent), Step::Fail(NotSent), Step::Fail(NotSent), ok(0.95, "d")]);
218    assert_eq!(r.unwrap().attempts, 2);
219    assert_eq!(retry_counts(&fake), vec![None, None, None, Some("1".into())]);
220}
221
222#[test]
223fn unretryable_transport_failures_fail_at_once() {
224    for kind in [TransportErrorKind::Tls, TransportErrorKind::TooLarge, TransportErrorKind::Config] {
225        let (r, fake) = run(vec![Step::Fail(kind)]);
226        assert!(matches!(r.unwrap_err(), ClientError::Transport { attempts: 1, .. }), "{kind:?}");
227        assert_eq!(fake.sent.borrow().len(), 1, "{kind:?}");
228    }
229}
230
231#[test]
232fn a_stalled_attempt_times_out_distrusts_the_connection_and_retries() {
233    let (r, fake) = run(vec![Step::Stall, ok(0.95, "e")]);
234    assert_eq!(r.unwrap().attempts, 2);
235    assert_eq!(fake.distrusted.get(), 1);
236    assert_eq!(fake.clock.get(), Duration::from_millis(10_500));
237}
238
239#[test]
240fn the_budget_bounds_everything() {
241    let (r, fake) = run(vec![Step::Stall, Step::Stall, Step::Stall]);
242    // 10 s, 0.5 s backoff, 10 s, 1 s backoff, then the 8.5 s left.
243    assert!(matches!(r.unwrap_err(), ClientError::TimedOut { attempts: 3, .. }));
244    assert_eq!(fake.clock.get(), Duration::from_secs(30));
245}
246
247#[test]
248fn transport_errors_keep_an_earlier_request_id() {
249    use TransportErrorKind::Interrupted;
250    let (r, _) = run(vec![
251        status(503, vec![(REQUEST_ID_HEADER, "seen")]),
252        Step::Fail(Interrupted),
253        Step::Fail(Interrupted),
254    ]);
255    let e = r.unwrap_err();
256    assert_eq!(e.request_id(), Some("seen"));
257    assert!(e.to_string().ends_with("[request seen]"), "{e}");
258}
259
260#[test]
261fn a_wrong_answer_carries_its_request_id() {
262    let (r, _) = run(vec![Step::Reply {
263        status: 200,
264        headers: vec![(REQUEST_ID_HEADER, "m")],
265        body: r#"{"model":"jev-9.9.9","answers":{"q":{"type":"noul","noul":0.5}},"usage":{"input_tokens":1,"output_tokens":1}}"#.into(),
266    }]);
267    let e = r.unwrap_err();
268    assert!(matches!(e, ClientError::Response { .. }));
269    assert_eq!(e.request_id(), Some("m"));
270}
271
272fn events(fake: &Fake) -> Vec<String> {
273    fake.events.borrow().clone()
274}
275
276#[test]
277fn a_success_reports_its_request_id_and_usage() {
278    let (_, fake) = run(vec![ok(0.95, "a")]);
279    assert_eq!(events(&fake), ["send 0", r#"answered id Some("a") in 1 tokens 1+1"#]);
280}
281
282#[test]
283fn a_retry_reports_its_cause_delay_and_request_id() {
284    let (_, fake) = run(vec![
285        status(503, vec![(REQUEST_ID_HEADER, "r1")]),
286        Step::Stall,
287        ok(0.95, "b"),
288    ]);
289    let e = events(&fake);
290    assert_eq!(e.len(), 6, "{e:?}");
291    assert_eq!(e[0], "send 0");
292    assert!(e[1].starts_with("retry after attempt 1 in 500ms (Jev answered HTTP 503"), "{}", e[1]);
293    assert!(e[1].ends_with(r#"id Some("r1")"#), "{}", e[1]);
294    assert_eq!(e[2], "send 1");
295    assert!(e[3].starts_with("retry after attempt 2 in 1s (attempt timed out"), "{}", e[3]);
296    assert_eq!(e[4], "send 2");
297    assert_eq!(e[5], r#"answered id Some("b") in 3 tokens 1+1"#);
298}
299
300#[test]
301fn a_redial_is_reported_and_is_not_a_retry() {
302    let (_, fake) = run(vec![Step::Fail(TransportErrorKind::NotSent), ok(0.95, "c")]);
303    assert_eq!(events(&fake), ["send 0", "redial (NotSent)", "send 0", r#"answered id Some("c") in 1 tokens 1+1"#]);
304}