Skip to main content

rig_core/
driver.rs

1//! Calls a model. A [`Model`] pairs a [`Wire`] (what to send and how to read
2//! the reply) with a [`Transport`] (how the payload travels). Its `stream`
3//! yields any operation's events as a [`Streamed`], and its `call` is that
4//! stream, finished: one decoder, one fold and one response constructor for
5//! both. The `_observed` twins take the observation context a bus records
6//! under.
7//!
8//! ```no_run
9//! use rig_core::completion::CompletionRequest;
10//! use rig_core::driver::Model;
11//! use rig_core::providers::openai::{self, OpenAI};
12//!
13//! # async fn example(http: rig_core::http_client::DynHttpClient) -> Result<(), Box<dyn std::error::Error>> {
14//! let model = OpenAI::from_env()?.with_http(http).responses(openai::GPT_5_2);
15//! let response = model.call(CompletionRequest::new("Hello")).await?;
16//! # let _ = response;
17//! # Ok(())
18//! # }
19//! ```
20
21use std::future::Future;
22use std::pin::Pin;
23use std::sync::{Arc, Mutex};
24use std::task::{Context, Poll};
25
26use futures::StreamExt;
27use tracing::Instrument;
28
29use crate::error::ProviderError;
30use crate::observe::{AdapterContext, AdapterEnding, AdapterSlot};
31use crate::streaming::Streamed;
32use crate::wasm_compat::{WasmBoxedFuture, WasmBoxedStream, WasmCompatSend, WasmCompatSync};
33use crate::wire::{
34    Call, Capabilities, Decoder, Flow, Mode, Operation, Out, Request, Response, Shared, Wire,
35    WireEvent,
36};
37
38mod dyn_model;
39mod http_transport;
40mod local;
41
42pub use dyn_model::DynModel;
43pub use local::{Local, Step};
44
45/// An endpoint of one provider: a wire bound to a transport.
46///
47/// The pair holds no invariant, so both halves are public. To share one
48/// transport across models, clone it; to send a model's requests another
49/// way, replace its transport. The transport defaults to the erased HTTP
50/// client, so a model on the default transport is `Model<W>`.
51#[derive(Clone, Debug, Default, PartialEq)]
52pub struct Model<W, T = crate::http_client::DynHttpClient> {
53    /// What to send and how to read the reply.
54    pub wire: W,
55    /// How the payload travels.
56    pub transport: T,
57}
58
59impl<W, T> Model<W, T> {
60    /// Pair `wire` with the `transport` that sends it.
61    pub fn new(wire: W, transport: T) -> Self {
62        Self { wire, transport }
63    }
64}
65
66/// Sends a wire's payloads and delivers its replies' frames.
67///
68/// HTTP clients ([`HttpClientExt`](crate::http_client::HttpClientExt)) are
69/// transports for every wire that encodes [`Encoded`](crate::wire::Encoded)
70/// requests and reads [`WireFrame`](crate::wire::WireFrame)s. A local
71/// runtime, a proxy or a mock is a transport for the wires it serves.
72pub trait Transport<W: Wire>: Clone + WasmCompatSend + WasmCompatSync + 'static {
73    /// Open the reply to one payload. Nothing is sent until the opening is
74    /// first polled. A payload the transport cannot send in the exchange's
75    /// mode fails the opening; a failure after the reply opened is its last
76    /// frame.
77    fn send(&self, payload: W::Payload, exchange: Exchange) -> Opening<W::Frame>;
78}
79
80/// What the driver tells a transport about one send.
81pub struct Exchange {
82    /// How the reply is read.
83    pub mode: Mode,
84    /// The observation an observing transport records the send under.
85    pub(crate) observation: Option<AdapterContext>,
86}
87
88/// One reply, opened when first polled.
89pub struct Opening<F>(WasmBoxedFuture<'static, Result<Opened<F>, ProviderError>>);
90
91impl<F: WasmCompatSend + 'static> Opening<F> {
92    /// The reply `open` opens.
93    pub fn new(
94        open: impl Future<Output = Result<Opened<F>, ProviderError>> + WasmCompatSend + 'static,
95    ) -> Self {
96        Self(Box::pin(open))
97    }
98
99    /// A reply that is already open.
100    pub fn ready(opened: Opened<F>) -> Self {
101        Self::new(std::future::ready(Ok(opened)))
102    }
103
104    /// A reply that failed to open with `error`.
105    pub fn failed(error: ProviderError) -> Self {
106        Self::new(std::future::ready(Err(error)))
107    }
108}
109
110impl<F> Future for Opening<F> {
111    type Output = Result<Opened<F>, ProviderError>;
112
113    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
114        self.0.as_mut().poll(cx)
115    }
116}
117
118/// One reply a transport opened: its frames, and the facts the transport
119/// owns.
120pub struct Opened<F> {
121    pub(crate) frames: WasmBoxedStream<'static, Result<F, ProviderError>>,
122    pub(crate) request_id: Option<String>,
123    pub(crate) status: Option<http::StatusCode>,
124    pub(crate) headers: Option<http::HeaderMap>,
125    pub(crate) route: Option<String>,
126    pub(crate) document: Option<serde_json::Value>,
127    pub(crate) slot: Option<AdapterSlot>,
128    pub(crate) analysis_only: Option<fn(&F) -> bool>,
129}
130
131impl<F: WasmCompatSend + 'static> Opened<F> {
132    /// A reply of `frames`, in order. A transport failure is the last item.
133    pub fn new(
134        frames: impl futures::Stream<Item = Result<F, ProviderError>> + WasmCompatSend + 'static,
135    ) -> Self {
136        Self {
137            frames: Box::pin(frames),
138            request_id: None,
139            status: None,
140            headers: None,
141            route: None,
142            document: None,
143            slot: None,
144            analysis_only: None,
145        }
146    }
147
148    /// A reply that failed with `error` before any frame.
149    pub fn failed(error: ProviderError) -> Self {
150        Self::new(futures::stream::once(async move { Err(error) }))
151    }
152
153    /// The provider's transport request id, when the reply carried one. An
154    /// empty id is no id.
155    pub fn with_request_id(mut self, request_id: Option<String>) -> Self {
156        self.request_id = crate::provider_response::reported(request_id);
157        self
158    }
159
160    /// The reply's HTTP status and headers, which enrich its unary failure.
161    pub fn with_http(mut self, status: http::StatusCode, headers: http::HeaderMap) -> Self {
162        self.status = Some(status);
163        self.headers = Some(headers);
164        self
165    }
166
167    /// The concrete request path, which a listing's failure names.
168    pub fn with_route(mut self, route: impl Into<String>) -> Self {
169        self.route = Some(route.into());
170        self
171    }
172
173    /// The whole reply as one document: a response's `raw` when the decoder
174    /// records none.
175    pub fn with_document(mut self, document: serde_json::Value) -> Self {
176        self.document = Some(document);
177        self
178    }
179
180    /// The same reply, its frames passed through `frames`: what a transport
181    /// that wraps another uses to watch or rewrite a reply in flight.
182    pub fn map_frames<S>(
183        mut self,
184        frames: impl FnOnce(WasmBoxedStream<'static, Result<F, ProviderError>>) -> S,
185    ) -> Self
186    where
187        S: futures::Stream<Item = Result<F, ProviderError>> + WasmCompatSend + 'static,
188    {
189        self.frames = Box::pin(frames(self.frames));
190        self
191    }
192}
193
194impl<W, T> Model<W, T>
195where
196    W: Wire,
197    T: Transport<W>,
198{
199    /// The wire's provider descriptor name (`"anthropic"`).
200    pub fn name(&self) -> &str {
201        self.wire.describe().name
202    }
203
204    /// The model id the wire addresses, when the operation addresses one.
205    pub fn id(&self) -> Option<&str> {
206        self.wire.describe().model
207    }
208
209    /// What a runtime accounts for about this model, such as an embedding
210    /// model's width.
211    pub fn capabilities(&self) -> Capabilities {
212        self.wire.describe().capabilities
213    }
214
215    /// Send `request` and fold the whole reply into the operation's
216    /// response: [`Self::stream`] of a unary reply, finished. A completion
217    /// takes a prompt, a conversation or a
218    /// [`CompletionRequest`](crate::completion::CompletionRequest).
219    pub fn call(
220        &self,
221        request: impl Into<Request<W>>,
222    ) -> impl Future<Output = Result<Response<W>, ProviderError>> + WasmCompatSend + 'static {
223        self.finished(request.into(), None)
224    }
225
226    /// [`Self::call`], with the attempt observed under `observation`.
227    pub fn call_observed(
228        &self,
229        request: impl Into<Request<W>>,
230        observation: AdapterContext,
231    ) -> impl Future<Output = Result<Response<W>, ProviderError>> + WasmCompatSend + 'static {
232        self.finished(request.into(), Some(observation))
233    }
234
235    /// The call opens when first polled, so its span is created under the
236    /// caller's instrumented context.
237    fn finished(
238        &self,
239        request: Request<W>,
240        observation: Option<AdapterContext>,
241    ) -> impl Future<Output = Result<Response<W>, ProviderError>> + WasmCompatSend + 'static {
242        let model = self.clone();
243        async move {
244            model
245                .open(request, Mode::Unary, observation)?
246                .finish()
247                .await
248        }
249    }
250
251    /// Open a streamed reply. Encoding errors return here; every later
252    /// failure arrives in-band, as the stream's last item. Nothing is sent
253    /// until the stream is first polled.
254    pub fn stream(&self, request: impl Into<Request<W>>) -> Result<Streamed<W::Op>, ProviderError> {
255        self.open(request.into(), Mode::Streaming, None)
256    }
257
258    /// [`Self::stream`], with the attempt observed under `observation`.
259    pub fn stream_observed(
260        &self,
261        request: impl Into<Request<W>>,
262        observation: AdapterContext,
263    ) -> Result<Streamed<W::Op>, ProviderError> {
264        self.open(request.into(), Mode::Streaming, Some(observation))
265    }
266
267    /// [`Self::call`], with a failure paired with the request path of the
268    /// reply that failed, for an operation whose errors name their route.
269    pub(crate) async fn call_routed(
270        &self,
271        request: Request<W>,
272    ) -> Result<Response<W>, (ProviderError, String)> {
273        self.open(request, Mode::Unary, None)
274            .map_err(|error| (error, String::new()))?
275            .finish_routed()
276            .await
277    }
278
279    /// The one entry to the driver: the operation's fold for the reply,
280    /// the encoded payload, and the transport's reply, in `mode`. A request
281    /// the operation rejects fails here, before it is encoded.
282    pub(crate) fn open(
283        &self,
284        request: Request<W>,
285        mode: Mode,
286        observation: Option<AdapterContext>,
287    ) -> Result<Streamed<W::Op>, ProviderError> {
288        <W::Op as Operation>::validate(&request)?;
289        let describe = self.wire.describe();
290        let provider = describe.name.to_owned();
291        let mut call = Call::new(&describe, mode);
292        let fold = <W::Op as Operation>::fold(&request, &mut call);
293        let span = call.span;
294        let payload = self.wire.encode(request, mode)?;
295        let opening = self.transport.send(payload, Exchange { mode, observation });
296        let shared = Arc::new(Mutex::new(Shared::new(fold)));
297        let reading = read(
298            self.wire.clone(),
299            opening,
300            Arc::clone(&shared),
301            span.clone(),
302            mode,
303        );
304        Ok(Streamed::new(reading, shared, span, provider))
305    }
306}
307
308/// Read one reply: open it, classify each frame, decode it into the reply's
309/// writer, and stop at the provider's end. The stream yields once per frame
310/// read, so its items reach the consumer frame by frame.
311///
312/// The reply's `'id` brand is the borrow of its shared state inside this
313/// stream: the decoder and every part handle it holds are tied to it, and
314/// none can leave.
315fn read<W: Wire>(
316    wire: W,
317    opening: Opening<W::Frame>,
318    shared: Arc<Mutex<Shared<W::Op>>>,
319    span: tracing::Span,
320    mode: Mode,
321) -> WasmBoxedStream<'static, ()> {
322    let decoding = span.clone();
323    let reading = async_stream::stream! {
324        let reply: &Mutex<Shared<W::Op>> = &shared;
325        // A unary call sends under the call's span; a stream reads under it.
326        let opened = match mode {
327            Mode::Unary => opening.instrument(span.clone()).await,
328            Mode::Streaming => opening.await,
329        };
330        let Opened {
331            mut frames,
332            request_id,
333            status,
334            headers,
335            route,
336            document,
337            slot,
338            analysis_only,
339        } = match opened {
340            Ok(opened) => opened,
341            Err(error) => {
342                fail(reply, slot_none(), error);
343                return;
344            }
345        };
346        record_request_id(&span, request_id.as_deref());
347        {
348            let mut state = lock(reply);
349            state.request_id.clone_from(&request_id);
350            state.document = document;
351            state.route = route.unwrap_or_default();
352        }
353        let enrich = |error: ProviderError| match mode {
354            // A unary reply's failure carries what the transport reported.
355            Mode::Unary => error
356                .with_provider_status(status)
357                .with_provider_request_id(request_id.clone())
358                .with_response_headers(headers.clone()),
359            Mode::Streaming => error,
360        };
361        let mut decoder = wire.decoder();
362        // Frames counted for observation's EOF and corruption positions.
363        let mut counted = 0usize;
364        loop {
365            let step = match frames.next().await {
366                Some(Ok(frame)) => {
367                    let analysis = slot.is_some()
368                        && analysis_only.is_some_and(|analysis_only| analysis_only(&frame));
369                    let classified = decoder.classify(frame);
370                    // A corrupt frame is never exempt, whatever its shape.
371                    let corrupt = matches!(classified, WireEvent::Corrupt(_));
372                    if slot.is_some() && (corrupt || !analysis) {
373                        counted += 1;
374                    }
375                    match classified {
376                        WireEvent::Known(event) => decoder.decode(event, Out::new(reply)),
377                        // Unmodeled, but always delivered: the consumer sees
378                        // it, and aggregation never folds it into the answer.
379                        WireEvent::Unknown { event_type, value } => {
380                            warn_unmodeled(&event_type, &value);
381                            Out::new(reply).unknown(value);
382                            Ok(Flow::More)
383                        }
384                        WireEvent::Corrupt(error) => {
385                            if let Some(slot) = &slot {
386                                slot.corrupt(counted);
387                            }
388                            Err(ProviderError::from(error))
389                        }
390                    }
391                }
392                Some(Err(error)) => {
393                    fail(reply, slot.as_ref(), error);
394                    return;
395                }
396                None => {
397                    if let Some(slot) = &slot {
398                        slot.transport_eof(counted);
399                    }
400                    // A decoder that saw the provider's end ends the reply
401                    // here; otherwise the frames ran out on it.
402                    let step = decoder.eof(Out::new(reply));
403                    if !matches!(step, Ok(Flow::Ended(_)))
404                        && let Some(slot) = &slot
405                    {
406                        slot.eof(counted);
407                    }
408                    match step {
409                        Ok(Flow::More) => Err(ProviderError::Truncated),
410                        step => step,
411                    }
412                }
413            };
414            match step {
415                Ok(Flow::More) => yield (),
416                Ok(Flow::Ended(_)) => break,
417                Err(error) => {
418                    fail(reply, slot.as_ref(), enrich(error));
419                    return;
420                }
421            }
422        }
423        if let Some(slot) = &slot {
424            slot.finish(AdapterEnding::Terminal);
425        }
426        {
427            let state = lock(reply);
428            if let Some(document) = state.raw.as_ref().or(state.document.as_ref()) {
429                crate::providers::internal::trace_json(
430                    crate::providers::internal::LogTarget::Completions,
431                    "reply",
432                    document,
433                );
434            }
435        }
436        yield ();
437    };
438    let mut reading: WasmBoxedStream<'static, ()> = Box::pin(reading);
439    match mode {
440        // A stream decodes under the call's span.
441        Mode::Streaming => Box::pin(futures::stream::poll_fn(move |cx| {
442            let _decoding = decoding.enter();
443            reading.as_mut().poll_next(cx)
444        })),
445        Mode::Unary => reading,
446    }
447}
448
449/// A reply decoded from frames already in hand: its items, then its
450/// response or the error that ended it.
451#[cfg(any(test, feature = "websocket", feature = "test-utils"))]
452pub(crate) struct Decoded<Op: Operation> {
453    /// Read only by relays and tests; a websocket session needs the outcome.
454    #[cfg(any(test, feature = "test-utils"))]
455    pub(crate) items: Vec<Result<crate::streaming::Item<Op::Event>, ProviderError>>,
456    pub(crate) outcome: Result<Op::Response, ProviderError>,
457}
458
459/// What the driver does with a classified frame: a known event is decoded,
460/// an unmodeled payload is warned about and delivered, and a corrupt frame
461/// is the error that ends the reply.
462#[cfg(any(test, feature = "websocket", feature = "test-utils"))]
463pub(crate) fn triage<E>(event: WireEvent<E>) -> Result<crate::streaming::Item<E>, ProviderError> {
464    match event {
465        WireEvent::Known(event) => Ok(crate::streaming::Item::Event(event)),
466        WireEvent::Unknown { event_type, value } => {
467            warn_unmodeled(&event_type, &value);
468            Ok(crate::streaming::Item::Unknown(value))
469        }
470        WireEvent::Corrupt(error) => Err(ProviderError::from(error)),
471    }
472}
473
474/// Decode one frame already in hand into `reply`, as [`read`] does.
475#[cfg(any(test, feature = "websocket", feature = "test-utils"))]
476pub(crate) fn step<'id, Op, F, D>(
477    decoder: &mut D,
478    reply: &'id Mutex<Shared<Op>>,
479    frame: F,
480) -> Result<Flow, ProviderError>
481where
482    Op: Operation,
483    D: crate::wire::Decoder<'id, Op, F>,
484{
485    match triage(decoder.classify(frame))? {
486        crate::streaming::Item::Event(event) => decoder.decode(event, Out::new(reply)),
487        crate::streaming::Item::Unknown(value) => {
488            Out::new(reply).unknown(value);
489            Ok(Flow::More)
490        }
491    }
492}
493
494/// Feed frames already in hand through `decoder` into `reply` as [`read`]
495/// does: the reply ends at the provider's end, or the decoder decides at EOF.
496#[cfg(any(test, feature = "websocket", feature = "test-utils"))]
497fn feed<'id, Op, F, D>(
498    decoder: &mut D,
499    reply: &'id Mutex<Shared<Op>>,
500    frames: impl IntoIterator<Item = F>,
501) -> Result<(), ProviderError>
502where
503    Op: Operation,
504    D: crate::wire::Decoder<'id, Op, F>,
505{
506    for frame in frames {
507        if let Flow::Ended(_) = step(decoder, reply, frame)? {
508            return Ok(());
509        }
510    }
511    match decoder.eof(Out::new(reply))? {
512        Flow::Ended(_) => Ok(()),
513        Flow::More => Err(ProviderError::Truncated),
514    }
515}
516
517/// Fold a fed reply: its items, then its response, or the error that ended
518/// it. `reply.raw` stands unless it is null, when the decoder's record does.
519#[cfg(any(test, feature = "websocket", feature = "test-utils"))]
520pub(crate) fn settle<Op: Operation>(
521    shared: Mutex<Shared<Op>>,
522    fed: Result<(), ProviderError>,
523    reply: crate::wire::Reply,
524) -> Decoded<Op> {
525    let Shared {
526        mut fold,
527        items,
528        end,
529        raw,
530        ..
531    } = shared
532        .into_inner()
533        .unwrap_or_else(std::sync::PoisonError::into_inner);
534    let mut absorbed = Ok(());
535    for item in &items {
536        if let (Ok(crate::streaming::Item::Event(event)), Ok(())) = (item, &absorbed) {
537            absorbed = crate::wire::Fold::absorb(&mut fold, event);
538        }
539    }
540    let outcome = fed.and(absorbed).and_then(|()| {
541        let reply = crate::wire::Reply {
542            raw: if reply.raw.is_null() {
543                raw.unwrap_or(serde_json::Value::Null)
544            } else {
545                reply.raw
546            },
547            ..reply
548        };
549        crate::wire::Fold::finish(fold, end.ok_or(ProviderError::Truncated)?, reply)
550    });
551    #[cfg(any(test, feature = "test-utils"))]
552    let items = {
553        let mut items: Vec<_> = items.into_iter().collect();
554        if let Err(error) = &outcome
555            && !items.iter().any(Result::is_err)
556        {
557            items.push(Err(error.clone()));
558        }
559        items
560    };
561    Decoded {
562        #[cfg(any(test, feature = "test-utils"))]
563        items,
564        outcome,
565    }
566}
567
568/// Decode a reply whose frames are already in hand through `wire`'s
569/// decoder and `fold`, without a transport: what a caller that reads a
570/// provider's frames itself (a websocket session, a whole body) finishes a
571/// reply with.
572#[cfg(any(test, feature = "websocket"))]
573pub(crate) fn decode_frames<W: Wire>(
574    wire: &W,
575    fold: <W::Op as Operation>::Fold,
576    frames: impl IntoIterator<Item = W::Frame>,
577    reply: crate::wire::Reply,
578) -> Result<Response<W>, ProviderError> {
579    let shared = Mutex::new(Shared::new(fold));
580    let fed = feed(&mut wire.decoder(), &shared, frames);
581    settle(shared, fed, reply).outcome
582}
583
584/// A whole reply body of an HTTP wire, decoded as its one frame: what a
585/// caller holding the body finishes a reply with.
586#[cfg(any(test, feature = "websocket"))]
587pub(crate) fn decode_body<W: Wire<Frame = crate::wire::WireFrame>>(
588    wire: &W,
589    fold: <W::Op as Operation>::Fold,
590    body: String,
591    reply: crate::wire::Reply,
592) -> Result<Response<W>, ProviderError> {
593    decode_frames(wire, fold, [crate::wire::WireFrame::Text(body)], reply)
594}
595
596/// A completion reply decoded from frames already in hand, as the bus
597/// relays it: its items, then the response, or the error that ended it.
598#[cfg(any(test, feature = "test-utils"))]
599pub(crate) fn relay_frames<W>(
600    wire: &W,
601    frames: impl IntoIterator<Item = W::Frame>,
602) -> crate::streaming::StreamEvents
603where
604    W: Wire<Op = crate::operation::Completion>,
605{
606    use crate::error::ErrorReport;
607    use crate::streaming::Relayed;
608
609    let provider = wire.describe().name.to_owned();
610    let shared = Mutex::new(Shared::new(crate::operation::Turn::new(provider.clone())));
611    let fed = feed(&mut wire.decoder(), &shared, frames);
612    let decoded = settle(
613        shared,
614        fed,
615        crate::wire::Reply {
616            provider,
617            raw: serde_json::Value::Null,
618            provider_request_id: None,
619        },
620    );
621    let relayed: Vec<Result<Relayed, ErrorReport>> = decoded
622        .items
623        .into_iter()
624        .map(|item| match item {
625            Ok(item) => Ok(Relayed::Item(item)),
626            Err(error) => Err(ErrorReport::from(&error)),
627        })
628        // An error that ended the reply is already its last item.
629        .chain(
630            decoded
631                .outcome
632                .ok()
633                .map(|response| Ok(Relayed::Done(Box::new(response)))),
634        )
635        .collect();
636    Box::pin(futures::stream::iter(relayed))
637}
638
639#[cfg(test)]
640impl Decoded<crate::operation::Completion> {
641    /// The events the reply yielded, without unmodeled payloads.
642    pub(crate) fn events(&self) -> Vec<&crate::streaming::StreamEvent> {
643        self.items
644            .iter()
645            .filter_map(|item| match item {
646                Ok(crate::streaming::Item::Event(event)) => Some(event),
647                _ => None,
648            })
649            .collect()
650    }
651
652    /// What each part ended with, in the order the parts ended.
653    pub(crate) fn ended(&self) -> Vec<crate::message::AssistantContent> {
654        self.events()
655            .into_iter()
656            .filter_map(|event| match event {
657                crate::streaming::StreamEvent::End { content, .. } => Some(content.clone()),
658                _ => None,
659            })
660            .collect()
661    }
662}
663
664/// Decode typed `events` through `decoder` as one completion reply from
665/// `provider`, then EOF: a decoder's test harness, without frames.
666#[cfg(test)]
667macro_rules! decode_events {
668    ($decoder:expr, $provider:expr, $events:expr) => {
669        $crate::driver::decode_with(
670            $crate::operation::Turn::new($provider),
671            $provider,
672            |reply| {
673                let mut decoder = $decoder;
674                for event in $events {
675                    if let $crate::wire::Flow::Ended(_) =
676                        $crate::wire::Decoder::decode(&mut decoder, event, reply.out())?
677                    {
678                        return Ok(());
679                    }
680                }
681                match $crate::wire::Decoder::eof(&mut decoder, reply.out())? {
682                    $crate::wire::Flow::Ended(_) => Ok(()),
683                    $crate::wire::Flow::More => Err($crate::error::ProviderError::Truncated),
684                }
685            },
686        )
687    };
688}
689#[cfg(test)]
690pub(crate) use decode_events;
691
692/// Decode `frames` through `decoder`, classifier included, as one
693/// completion reply from `provider`: what the driver does with frames it
694/// read.
695#[cfg(test)]
696macro_rules! feed_frames {
697    ($decoder:expr, $provider:expr, $frames:expr) => {
698        $crate::driver::decode_with(
699            $crate::operation::Turn::new($provider),
700            $provider,
701            |reply| {
702                let mut decoder = $decoder;
703                reply.feed(&mut decoder, $frames)
704            },
705        )
706    };
707}
708#[cfg(test)]
709pub(crate) use feed_frames;
710
711/// One reply's writer, for a test that drives a decoder by hand.
712#[cfg(test)]
713pub(crate) struct Replying<'id, Op: Operation>(&'id Mutex<Shared<Op>>);
714
715#[cfg(test)]
716impl<'id, Op: Operation> Replying<'id, Op> {
717    /// A writer for the next event.
718    pub(crate) fn out(&self) -> Out<'id, Op> {
719        Out::new(self.0)
720    }
721
722    /// Feed frames through `decoder`, as the driver does.
723    pub(crate) fn feed<F, D: crate::wire::Decoder<'id, Op, F>>(
724        &self,
725        decoder: &mut D,
726        frames: impl IntoIterator<Item = F>,
727    ) -> Result<(), ProviderError> {
728        feed(decoder, self.0, frames)
729    }
730}
731
732/// Drive one reply by hand: `run` decodes into it until the reply ends, or
733/// fails with what ended it, and the reply folds into `fold` as a stream
734/// would.
735#[cfg(test)]
736pub(crate) fn decode_with<Op: Operation>(
737    fold: Op::Fold,
738    provider: &str,
739    run: impl for<'id> FnOnce(Replying<'id, Op>) -> Result<(), ProviderError>,
740) -> Decoded<Op> {
741    let shared = Mutex::new(Shared::new(fold));
742    let fed = run(Replying(&shared));
743    settle(
744        shared,
745        fed,
746        crate::wire::Reply {
747            provider: provider.to_owned(),
748            raw: serde_json::Value::Null,
749            provider_request_id: None,
750        },
751    )
752}
753
754/// No observation slot, for a reply that failed before it opened.
755fn slot_none() -> Option<&'static AdapterSlot> {
756    None
757}
758
759/// The reply failed: the error is its last item.
760fn fail<Op: Operation>(
761    reply: &Mutex<Shared<Op>>,
762    slot: Option<&AdapterSlot>,
763    error: ProviderError,
764) {
765    if let Some(slot) = slot {
766        slot.fail(&error);
767    }
768    lock(reply).items.push_back(Err(error));
769}
770
771pub(crate) fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
772    mutex
773        .lock()
774        .unwrap_or_else(std::sync::PoisonError::into_inner)
775}
776
777/// Read pages until one names no cursor, names the cursor it was asked
778/// with, or [`MAX_CONTINUATION_PAGES`] were read. `page` reads the page at a
779/// cursor (`None` for the first) and returns it with the next cursor.
780pub(crate) async fn follow_cursors<P, Fut>(
781    provider: &str,
782    operation: &str,
783    mut page: impl FnMut(Option<String>) -> Fut,
784) -> Result<Vec<P>, ProviderError>
785where
786    Fut: Future<Output = Result<(P, Option<String>), ProviderError>>,
787{
788    let mut pages = Vec::new();
789    let mut cursor = None;
790    loop {
791        let (read, next) = page(cursor.clone()).await?;
792        pages.push(read);
793        // Warn only when an offered page is refused: normal exhaustion is
794        // not truncation, while repeated or cycling cursors need a bound.
795        let Some(next) = next else { break };
796        if cursor.as_deref() == Some(next.as_str()) {
797            // The next request would be identical to the one just answered,
798            // so the page would repeat forever.
799            tracing::warn!(
800                provider,
801                operation,
802                pages = pages.len(),
803                "listing repeated its pagination cursor; returning the pages fetched so far"
804            );
805            break;
806        }
807        if pages.len() >= MAX_CONTINUATION_PAGES {
808            tracing::warn!(
809                provider,
810                operation,
811                pages = pages.len(),
812                "listing hit its page ceiling with a cursor still advancing; returning the pages \
813                 fetched so far"
814            );
815            break;
816        }
817        cursor = Some(next);
818    }
819    Ok(pages)
820}
821
822/// Logs an unmodeled payload's kind and serialized size, never its content.
823/// Callers must supply a structural kind label without sensitive data.
824pub fn warn_unmodeled(kind: &str, payload: &impl serde::Serialize) {
825    tracing::warn!(
826        kind,
827        payload_bytes = unknown_payload_bytes(payload),
828        "skipping unmodeled wire payload"
829    );
830}
831
832/// Serialized byte size of an unknown frame's payload, for the structural
833/// warn log (the log never carries the payload itself).
834fn unknown_payload_bytes(value: &impl serde::Serialize) -> u64 {
835    /// Counter sink: measures how many bytes serialization would write
836    /// without buffering them.
837    struct CountingWriter(u64);
838
839    impl std::io::Write for CountingWriter {
840        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
841            self.0 += buf.len() as u64;
842            Ok(buf.len())
843        }
844
845        fn flush(&mut self) -> std::io::Result<()> {
846            Ok(())
847        }
848    }
849
850    let mut counter = CountingWriter(0);
851    // A `Value` cannot fail to serialize; degrade to 0 rather than panic.
852    let _ = serde_json::to_writer(&mut counter, value);
853    counter.0
854}
855
856/// Page count after which a listing's cursors are ignored, preventing
857/// infinite cursor cycles.
858const MAX_CONTINUATION_PAGES: usize = 1000;
859
860/// Record the transport request id on the call's span, success or failure.
861pub(crate) fn record_request_id(span: &tracing::Span, request_id: Option<&str>) {
862    if let Some(request_id) = request_id
863        && !span.is_disabled()
864    {
865        span.record(crate::telemetry::PROVIDER_REQUEST_ID_FIELD, request_id);
866    }
867}
868
869#[cfg(test)]
870pub(crate) mod tests;