Skip to main content

rig_core/operation/
completion.rs

1//! Generating an assistant turn. A completion decoder writes its reply one
2//! provider item at a time, keyed by the item's wire index: the item opens
3//! when its first event arrives and takes the next position in the reply,
4//! its fragments grow it, and closing it finalizes its block with the
5//! provider's item as the block's native. Nothing regroups: blocks keep the
6//! order their items opened in. The provider's end of the reply is a
7//! [`Finish`], which the fold needs to produce the response.
8//!
9//! ```
10//! use rig_core::operation::Finish;
11//! use rig_core::completion::FinishReason;
12//!
13//! let finish = Finish {
14//!     reason: Some(FinishReason::Stop),
15//!     ..Finish::default()
16//! };
17//! assert_eq!(finish.reason, Some(FinishReason::Stop));
18//! ```
19
20use std::collections::{BTreeMap, HashMap, HashSet};
21
22use crate::completion::{CompletionRequest, CompletionResponse, FinishReason, Usage};
23use crate::error::ProviderError;
24use crate::message::citation;
25use crate::message::{
26    Api, AssistantContent, CallId, Image, LocalCallId, Opaque, Origin, Reasoning, Text, ToolCall,
27    ToolFunction, ToolName,
28};
29use crate::streaming::{Item, Part, PartKind, StreamEvent};
30use crate::telemetry::{GenAiOperation, SpanBuilder, SpanCombinator};
31use crate::wire::{
32    Assembled, Call, Descriptor, Emit, Fold, Mode, Operation, Out, Reply, Shared, WireCitation,
33};
34
35/// Generating an assistant turn, unary or streamed.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct Completion;
38
39impl Operation for Completion {
40    type Request = CompletionRequest;
41    type Event = StreamEvent;
42    type End = Finish;
43    type Response = CompletionResponse;
44    type Fold = Turn;
45    type Emit = Assembled;
46
47    /// The call's span names the model the request overrides to, when it
48    /// names one: every wire honours the override on encode.
49    fn fold(request: &Self::Request, call: &mut Call<'_>) -> Self::Fold {
50        let telemetry = call
51            .wire
52            .telemetry
53            .map_or(GenAiOperation::Chat, |telemetry| telemetry(call.mode));
54        debug_assert!(telemetry.is_completion());
55        // A wire that picks its API per request records the one it used.
56        let replay = call
57            .wire
58            .replay
59            .map(|target| target.route(request).unwrap_or(target));
60        let model = request
61            .model
62            .as_deref()
63            .or(call.wire.model)
64            .or_else(|| replay.map(|target| target.model()))
65            .unwrap_or_default();
66        let span = SpanBuilder::new(call.wire.name, model, telemetry)
67            .streaming(matches!(call.mode, Mode::Streaming))
68            .system_instructions(
69                request.system_instructions(),
70                request.record_telemetry_content,
71            )
72            .build();
73        call.instrument(span.clone());
74        let api = replay.map_or_else(
75            || Api::from(call.wire.name.to_owned()),
76            |target| target.api(),
77        );
78        let mut origin = Origin::new(api, call.wire.name, model);
79        // Only a wire that binds items to the context needs it recorded.
80        if replay.is_some_and(|target| target.binds_context(model)) {
81            origin.context =
82                replay.map(|target| crate::completion::history::context_of(request, target, model));
83        }
84        Turn {
85            span,
86            wire: replay.is_some_and(|target| target.states_finish_reason()),
87            call_id_slot: replay.and_then(|target| target.call_id_slot()),
88            accept_unknown_finish: request.accept_unknown_finish_reasons,
89            ..Turn::new(origin)
90        }
91    }
92
93    /// Resolve the model the request addresses once, as the request's
94    /// `model`: the one it names, else the wire's. The history is checked
95    /// with [`CompletionRequest::validate_message_content`], then
96    /// [`adapt`](crate::completion::adapt)ed for that model on the wire's
97    /// replay target and checked again. Encoders, the fold and replay all
98    /// read the one resolved model. Each option the routed target refuses
99    /// is reported first, through
100    /// [`options::check`](crate::completion::options::check).
101    fn prepare(
102        mut request: Self::Request,
103        wire: &Descriptor<'_>,
104    ) -> Result<Self::Request, ProviderError> {
105        let Some(target) = wire.replay else {
106            return Err(ProviderError::request(format!(
107                "completion wire `{}` names no replay target",
108                wire.name
109            )));
110        };
111        // The caller's history is checked as written, so a rejection names
112        // its own messages; adapting never empties a message it keeps.
113        request.validate_message_content()?;
114        // A wire that addresses no model (an interaction read back) leaves
115        // it to the reply, which names the model the turn is from.
116        request.model = request
117            .model
118            .take()
119            .filter(|model| !model.is_empty())
120            .or_else(|| Some(target.model().to_owned()).filter(|model| !model.is_empty()));
121        let target = target.route(&request).unwrap_or(target);
122        // Every option the routed wire cannot honour is reported here, before
123        // any wire encodes, whatever its encoder does.
124        crate::completion::options::check(target, &mut request)?;
125        // Documents join the history before it is adapted, so the adapter's
126        // rules apply to them, unless the encoder sends them itself.
127        if !target.takes_documents() {
128            request.chat_history = request.chat_history_with_documents();
129            request.documents.clear();
130        }
131        let stored = target.continues_stored(&request);
132        let shape = crate::completion::history::Request {
133            model: request.model.as_deref(),
134            stored,
135            // A conversation the provider stores holds its own tools, so a
136            // continuation that declares none still calls them.
137            tools: stored || target.declares_tools(&request),
138            context: (!target.drops_unbound_items(&request)).then(|| {
139                crate::completion::history::context_of(
140                    &request,
141                    target,
142                    request.model.as_deref().unwrap_or(target.model()),
143                )
144            }),
145        };
146        request.chat_history =
147            crate::completion::history::adapt_for(&request.chat_history, target, &shape);
148        request.validate_message_content()?;
149        Ok(request)
150    }
151}
152
153impl crate::wire::reply::Closing<Completion> for Assembled {
154    fn close(shared: &mut Shared<Completion>) {
155        let Shared { fold, items, .. } = shared;
156        fold.close_open(items);
157    }
158}
159
160impl Emit<Completion> for Assembled {}
161
162/// What the provider sends when it ends a completion reply.
163#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
164pub struct Finish {
165    /// Token usage the provider reported. A counter it did not report is
166    /// `None`. A cost the provider reported goes in `usage.cost` and is
167    /// kept; without one the fold prices the counters at the catalog's
168    /// [`Pricing`](crate::catalog::Pricing) for the request's model, and
169    /// leaves the cost `None` when the catalog has no price for it.
170    pub usage: Usage,
171    /// Why the model stopped, when the provider said.
172    #[serde(default, skip_serializing_if = "Option::is_none")]
173    pub reason: Option<FinishReason>,
174    /// The response id.
175    pub response_id: Option<String>,
176    /// The model the provider reports.
177    pub model: Option<String>,
178    /// The provider's report that the turn failed, such as a refusal's
179    /// explanation. The turn is then never replayed.
180    #[serde(default, skip_serializing_if = "Option::is_none")]
181    pub error: Option<String>,
182}
183
184/// What a provider item becomes when it opens.
185#[derive(Debug, Clone, PartialEq)]
186pub enum Block {
187    /// Answer text.
188    Text,
189    /// Reasoning, possibly withheld by the provider.
190    Reasoning {
191        /// Whether the provider withheld the text.
192        redacted: bool,
193    },
194    /// An image, its canonical fields stated at open; fragments append
195    /// base64 data.
196    Image(Image),
197    /// A tool call with its id and name.
198    Call {
199        /// The call's id.
200        id: CallId,
201        /// The tool's name.
202        name: ToolName,
203    },
204    /// An item with no canonical meaning.
205    Opaque {
206        /// Whether the item goes back to the model that produced it.
207        replay: bool,
208    },
209}
210
211/// The completion fold. The driver and the bus writer build it; no other
212/// code can feed one.
213///
214/// It is both sides of one reply: the writer state a decoder writes through
215/// (the open items by wire index) and the blocks the consumer has taken, in
216/// their position.
217pub struct Turn {
218    span: tracing::Span,
219    /// Who the reply is from; the end adds the provider's model and id.
220    origin: Origin,
221    // The writer.
222    open: BTreeMap<usize, Draft>,
223    /// The position of the block each wire index last closed, whose item
224    /// [`Out::edit`] may still amend until the reply ends.
225    ended: HashMap<usize, usize>,
226    next_part: u32,
227    call_ids: HashSet<CallId>,
228    /// Whether the reply's end must name a finish reason: a wire's fold
229    /// whose target states one (a relayed or written reply states its own).
230    wire: bool,
231    /// Where the wire's call items hold their id ([`ReplayTarget::call_id_slot`]).
232    ///
233    /// [`ReplayTarget::call_id_slot`]: crate::completion::ReplayTarget::call_id_slot
234    call_id_slot: Option<&'static str>,
235    /// Whether the request accepts an unknown finish reason as a normal stop
236    /// ([`CompletionRequest::accept_unknown_finish_reasons`]).
237    ///
238    /// [`CompletionRequest::accept_unknown_finish_reasons`]: field@CompletionRequest::accept_unknown_finish_reasons
239    accept_unknown_finish: bool,
240    /// The first position whose item closed without the provider stating it
241    /// complete. An item there may be the partner a later one needs, so
242    /// blocks from it on replay from their canonical fields.
243    first_incomplete: Option<usize>,
244    /// Whether a call was still open when the provider ended the reply.
245    unfinished_call: bool,
246    /// The call the last index-less fragment went to.
247    last_call: Option<usize>,
248    /// The index of the block a boundary-less wire is streaming.
249    run: Option<usize>,
250    next_auto: usize,
251    /// The wire index of each position, for a wire whose choice follows
252    /// its indices ([`Out::order_by_index`]).
253    by_index: Option<BTreeMap<usize, usize>>,
254    /// The position of the block the choice holds first ([`Out::lead`]).
255    lead: Option<usize>,
256    // The fold.
257    choice: Vec<Option<AssistantContent>>,
258    /// The text the consumer took of text and reasoning parts still open,
259    /// by position.
260    open_text: BTreeMap<usize, AssistantContent>,
261}
262
263/// One open provider item.
264struct Draft {
265    part: Part,
266    /// Whether its start was emitted.
267    started: bool,
268    /// The provider's item as assembled so far; `Null` for none.
269    item: serde_json::Value,
270    body: Body,
271    /// The citations of a text item, resolved when it closes.
272    citations: Vec<WireCitation>,
273}
274
275/// Whether a closing item keeps its provider item as the block's native.
276#[derive(Clone, Copy, PartialEq, Eq)]
277enum Closing {
278    /// The provider stated the item complete: it becomes the native.
279    Complete,
280    /// The item never completed: the block has no native, and an opaque
281    /// item does not replay.
282    Incomplete,
283}
284
285enum Body {
286    Text(String),
287    Image(Image),
288    Reasoning {
289        text: String,
290        redacted: bool,
291    },
292    Call {
293        id: Option<CallId>,
294        name: String,
295        arguments: Arguments,
296    },
297    Opaque {
298        replay: bool,
299    },
300}
301
302impl Body {
303    fn kind(&self) -> PartKind {
304        match self {
305            Self::Text(_) => PartKind::Text,
306            Self::Reasoning { .. } => PartKind::Reasoning,
307            Self::Image(_) => PartKind::Image,
308            Self::Call { .. } => PartKind::ToolCall,
309            Self::Opaque { .. } => PartKind::Opaque,
310        }
311    }
312}
313
314/// A tool call's argument text as it arrives.
315#[derive(Default)]
316struct Arguments {
317    text: String,
318    /// How many bytes of `text` the stream has carried.
319    streamed: usize,
320    overflowed: bool,
321    /// Whether any fragment carried a non-blank byte.
322    substantive: bool,
323    /// Arguments the provider announced when the call opened, used only
324    /// when no fragment arrives.
325    announced: Option<serde_json::Value>,
326}
327
328/// The most argument bytes one call accumulates.
329const MAX_TOOL_INPUT_BYTES: usize = 32 * 1024 * 1024;
330
331impl Arguments {
332    fn push(&mut self, fragment: &str, name: &str) {
333        self.substantive |= !fragment.trim().is_empty();
334        // Some OpenAI-compatible gateways send a literal `null` before the
335        // real fragments; a non-blank fragment supersedes it.
336        if self.text.trim() == "null" && !fragment.trim().is_empty() {
337            self.text.clear();
338        }
339        if self.text.len().saturating_add(fragment.len()) > MAX_TOOL_INPUT_BYTES {
340            if !self.overflowed {
341                self.overflowed = true;
342                tracing::warn!(
343                    tool = name,
344                    "streamed tool-call input exceeded the accumulation bound; truncating"
345                );
346            }
347        } else {
348            self.text.push_str(fragment);
349        }
350    }
351
352    /// The call to `name` these arguments make: the announced ones when no
353    /// fragment arrived, else the text, read by [`ToolFunction::parse`].
354    fn function(&self, name: ToolName) -> ToolFunction {
355        if self.text.is_empty()
356            && let Some(announced) = &self.announced
357        {
358            return ToolFunction::new(name, announced.clone());
359        }
360        let mut function = ToolFunction::parse(name, &self.text);
361        if self.overflowed {
362            function.invalid_arguments = Some(self.text.clone());
363        }
364        function
365    }
366
367    /// The text the stream has not carried yet, now marked carried. A text
368    /// that may still be a gateway's placeholder `null` waits, since the
369    /// first real fragment replaces it.
370    fn unstreamed(&mut self) -> Option<String> {
371        if "null".starts_with(self.text.trim()) {
372            return None;
373        }
374        let fragment = self.text.get(self.streamed..)?;
375        if fragment.is_empty() {
376            return None;
377        }
378        let fragment = fragment.to_owned();
379        self.streamed = self.text.len();
380        Some(fragment)
381    }
382
383    /// Whether the text is a complete JSON object.
384    fn complete(&self) -> bool {
385        !self.overflowed
386            && matches!(
387                crate::json_utils::parse_tool_arguments(&self.text),
388                Ok(serde_json::Value::Object(_))
389            )
390    }
391}
392
393/// Whether `id` names a call: some providers send `""` or `"null"` for a
394/// call they gave no id.
395fn stated(id: &str) -> bool {
396    !id.is_empty() && id != "null"
397}
398
399/// What one fragment of a buffered tool call carries.
400#[derive(Debug, Clone, Copy, Default)]
401pub struct CallFragment<'a> {
402    /// The provider's id for the call.
403    pub id: Option<&'a str>,
404    /// The tool's name.
405    pub name: Option<&'a str>,
406    /// A fragment of the argument JSON.
407    pub arguments: Option<&'a str>,
408}
409
410/// The first index the writer hands out itself ([`Out::fresh_index`]); a
411/// provider's wire indices stay below it.
412pub const AUTO_INDEX: usize = 1 << 30;
413
414fn not_open(index: usize) -> ProviderError {
415    ProviderError::Response(format!(
416        "the reply wrote to item {index}, which is not open"
417    ))
418}
419
420/// Carry what the open call `draft` has not streamed yet: its start once it
421/// names a tool, then the argument text since the last fragment. A call
422/// that never names a tool stays unseen, as its close drops it.
423fn stream_call(items: &mut Items, draft: &mut Draft) {
424    let Body::Call {
425        name, arguments, ..
426    } = &mut draft.body
427    else {
428        return;
429    };
430    if !draft.started {
431        let Ok(name) = ToolName::new(name.clone()) else {
432            return;
433        };
434        draft.started = true;
435        emit(
436            items,
437            StreamEvent::Start {
438                part: draft.part,
439                kind: PartKind::ToolCall,
440                name: Some(name),
441            },
442        );
443    }
444    if let Some(json) = arguments.unstreamed() {
445        emit(
446            items,
447            StreamEvent::Arguments {
448                part: draft.part,
449                json,
450            },
451        );
452    }
453}
454
455impl Turn {
456    /// The writer of a reply from `origin`.
457    pub(crate) fn new(origin: Origin) -> Self {
458        Self {
459            span: tracing::Span::none(),
460            origin,
461            open: BTreeMap::new(),
462            ended: HashMap::new(),
463            next_part: 0,
464            call_ids: HashSet::new(),
465            wire: false,
466            call_id_slot: None,
467            accept_unknown_finish: false,
468            first_incomplete: None,
469            unfinished_call: false,
470            last_call: None,
471            run: None,
472            next_auto: AUTO_INDEX,
473            by_index: None,
474            lead: None,
475            choice: Vec::new(),
476            open_text: BTreeMap::new(),
477        }
478    }
479
480    /// The fold of a stream relayed from another fold, which built its
481    /// events: it only collects them, and takes its origin from the relay
482    /// ([`Self::set_origin`]). Until then it names the relay `label`.
483    pub(crate) fn relayed(label: impl Into<String>) -> Self {
484        let label = label.into();
485        Self::new(Origin::new(label.clone(), label.clone(), label))
486    }
487
488    /// Who the reply is from, as the relay or the writer's owner states it.
489    pub(crate) fn set_origin(&mut self, origin: Origin) {
490        self.origin = origin;
491    }
492
493    /// The next position in the choice.
494    fn next(&mut self) -> Part {
495        let part = Part::new(self.next_part);
496        self.next_part += 1;
497        part
498    }
499
500    fn fresh_index(&mut self) -> usize {
501        let index = self.next_auto;
502        self.next_auto += 1;
503        index
504    }
505
506    fn draft(&mut self, index: usize) -> Result<&mut Draft, ProviderError> {
507        self.open.get_mut(&index).ok_or_else(|| not_open(index))
508    }
509
510    pub(crate) fn open_item(
511        &mut self,
512        items: &mut Items,
513        index: usize,
514        block: Block,
515        item: serde_json::Value,
516    ) -> Result<(), ProviderError> {
517        if self.open.contains_key(&index) {
518            return Err(ProviderError::Response(format!(
519                "the reply opened item {index} twice"
520            )));
521        }
522        let body = match block {
523            Block::Text => Body::Text(String::new()),
524            Block::Image(image) => Body::Image(image),
525            Block::Reasoning { redacted } => Body::Reasoning {
526                text: String::new(),
527                redacted,
528            },
529            Block::Call { id, name } => Body::Call {
530                id: Some(id),
531                name: name.into(),
532                arguments: Arguments::default(),
533            },
534            Block::Opaque { replay } => Body::Opaque { replay },
535        };
536        self.insert(index, body, item);
537        stream_call(items, self.draft(index)?);
538        Ok(())
539    }
540
541    fn insert(&mut self, index: usize, body: Body, item: serde_json::Value) {
542        self.ended.remove(&index);
543        let part = self.next();
544        if let Some(by_index) = &mut self.by_index {
545            by_index.insert(part.index(), index);
546        }
547        let started = false;
548        self.open.insert(
549            index,
550            Draft {
551                part,
552                started,
553                item,
554                body,
555                citations: Vec::new(),
556            },
557        );
558    }
559
560    pub(crate) fn push_item(
561        &mut self,
562        items: &mut Items,
563        index: usize,
564        fragment: &str,
565    ) -> Result<(), ProviderError> {
566        if fragment.is_empty() {
567            return Ok(());
568        }
569        let draft = self.draft(index)?;
570        let event = match &mut draft.body {
571            Body::Text(text) => {
572                text.push_str(fragment);
573                StreamEvent::Text {
574                    part: draft.part,
575                    text: fragment.to_owned(),
576                }
577            }
578            Body::Reasoning { text, .. } => {
579                text.push_str(fragment);
580                StreamEvent::Reasoning {
581                    part: draft.part,
582                    text: fragment.to_owned(),
583                }
584            }
585            Body::Call {
586                name, arguments, ..
587            } => {
588                arguments.push(fragment, name);
589                stream_call(items, draft);
590                return Ok(());
591            }
592            Body::Image(image) => {
593                if let crate::message::DocumentSourceKind::Base64(data) = &mut image.data {
594                    data.push_str(fragment);
595                    return Ok(());
596                }
597                return Err(ProviderError::Response(format!(
598                    "the reply wrote data to the image item {index}, which holds no base64 data"
599                )));
600            }
601            Body::Opaque { .. } => {
602                return Err(ProviderError::Response(format!(
603                    "the reply wrote text to the opaque item {index}"
604                )));
605            }
606        };
607        if !draft.started {
608            draft.started = true;
609            emit(
610                items,
611                StreamEvent::Start {
612                    part: draft.part,
613                    kind: draft.body.kind(),
614                    name: None,
615                },
616            );
617        }
618        emit(items, event);
619        Ok(())
620    }
621
622    /// The id a call closing now takes, and the provider item it keeps:
623    /// `id`, or, when the provider sent none or one an earlier call of the
624    /// reply holds, a fresh rig-issued id (as pi issues). A renamed call
625    /// keeps no item, since the item names the id it lost; a call sent with
626    /// no id keeps its item only on a wire whose items have an id slot,
627    /// which replay then fills with the id its result gets.
628    fn distinct_call_id(
629        &mut self,
630        id: Option<CallId>,
631        item: serde_json::Value,
632    ) -> (CallId, serde_json::Value) {
633        let (id, item) = match id {
634            Some(id) if !self.call_ids.contains(&id) => (id, item),
635            Some(id) => {
636                tracing::warn!(%id, "the provider named two tool calls with one id; renaming the second");
637                (CallId::Local(LocalCallId::new()), serde_json::Value::Null)
638            }
639            None if self.call_id_slot.is_some() => (CallId::Local(LocalCallId::new()), item),
640            None => (CallId::Local(LocalCallId::new()), serde_json::Value::Null),
641        };
642        self.call_ids.insert(id.clone());
643        (id, item)
644    }
645
646    /// Close the item at `index`: its block becomes visible, holding the
647    /// provider's item as its native when `closing` says it is complete.
648    fn close_item(
649        &mut self,
650        items: &mut Items,
651        index: usize,
652        closing: Closing,
653    ) -> Result<(), ProviderError> {
654        let mut draft = self.open.remove(&index).ok_or_else(|| not_open(index))?;
655        // A call streams what it has not yet: its start, if it names a tool,
656        // and the argument text since its last fragment.
657        stream_call(items, &mut draft);
658        if self.run == Some(index) {
659            self.run = None;
660        }
661        self.ended.insert(index, draft.part.index());
662        let Draft {
663            part,
664            started,
665            item,
666            body,
667            citations,
668        } = draft;
669        let item = match (closing, &body) {
670            (Closing::Complete, _) | (Closing::Incomplete, Body::Opaque { .. }) => item,
671            (Closing::Incomplete, _) => {
672                self.cut_at(part.index());
673                serde_json::Value::Null
674            }
675        };
676        let content = match body {
677            Body::Text(text) => {
678                if text.is_empty() && item.is_null() {
679                    return Ok(());
680                }
681                // Resolved at close: a provider may cite a text before its
682                // last fragment arrives.
683                let mut text = Text::new(text);
684                citation::attach(
685                    &mut text,
686                    Vec::new(),
687                    citations,
688                    &self.origin.provider,
689                    index,
690                );
691                with_item(AssistantContent::Text(text), item)
692            }
693            Body::Reasoning { text, redacted } => {
694                if text.is_empty() && !redacted && item.is_null() {
695                    return Ok(());
696                }
697                let reasoning = Reasoning {
698                    text,
699                    redacted,
700                    native: None,
701                };
702                with_item(AssistantContent::Reasoning(reasoning), item)
703            }
704            Body::Image(image) => with_item(AssistantContent::Image(image), item),
705            Body::Opaque { replay } => AssistantContent::Opaque(Opaque {
706                item,
707                replay: replay && closing == Closing::Complete,
708            }),
709            Body::Call {
710                id,
711                name,
712                arguments,
713            } => {
714                let Ok(name) = ToolName::new(name) else {
715                    tracing::warn!(
716                        index,
717                        "the provider closed a tool call without a name; nothing can answer it"
718                    );
719                    return Ok(());
720                };
721                let function = arguments.function(name);
722                let (id, item) = self.distinct_call_id(id, item);
723                // A call that streamed no text (none arrived, or only a
724                // placeholder `null`) states the arguments it ends with.
725                if arguments.streamed == 0 {
726                    let json = serde_json::Value::Object(function.arguments.clone()).to_string();
727                    emit(items, StreamEvent::Arguments { part, json });
728                }
729                let call = ToolCall::new(id, function);
730                let content = with_item(AssistantContent::ToolCall(call), item);
731                emit(items, StreamEvent::End { part, content });
732                return Ok(());
733            }
734        };
735        if !started {
736            emit(
737                items,
738                StreamEvent::Start {
739                    part,
740                    kind: kind_of(&content),
741                    name: None,
742                },
743            );
744        }
745        emit(items, StreamEvent::End { part, content });
746        Ok(())
747    }
748
749    /// Edit the item of the block the wire index `index` last closed, while
750    /// its end event waits to be taken or once the fold holds it.
751    fn edit_ended(
752        &mut self,
753        items: &mut Items,
754        index: usize,
755        edit: impl FnOnce(&mut serde_json::Value),
756    ) -> Result<(), ProviderError> {
757        let position = *self.ended.get(&index).ok_or_else(|| not_open(index))?;
758        let queued = items.iter_mut().find_map(|item| match item {
759            Ok(Item::Event(StreamEvent::End { part, content })) if part.index() == position => {
760                Some(content)
761            }
762            _ => None,
763        });
764        let content = match queued {
765            Some(content) => Some(content),
766            None => self.choice.get_mut(position).and_then(Option::as_mut),
767        };
768        if let Some(item) = content.and_then(item_of) {
769            edit(item);
770        }
771        Ok(())
772    }
773
774    /// Add `citations` to the text item at `index`, or with `replace` put
775    /// them in place of its own. An open item resolves them when it closes;
776    /// a closed one, against its final text now, while its end event waits
777    /// to be taken or once the fold holds it. Citations for any other item
778    /// are dropped with a warning.
779    fn cite_item(
780        &mut self,
781        items: &mut Items,
782        index: usize,
783        citations: Vec<WireCitation>,
784        replace: bool,
785    ) {
786        let provider = self.origin.provider.as_str();
787        if let Some(draft) = self.open.get_mut(&index) {
788            if !matches!(draft.body, Body::Text(_)) {
789                tracing::warn!(
790                    provider,
791                    index,
792                    "dropped citations of an item that is not text"
793                );
794            } else if replace {
795                draft.citations = citations;
796            } else {
797                draft.citations.extend(citations);
798            }
799            return;
800        }
801        let Some(position) = self.ended.get(&index).copied() else {
802            tracing::warn!(
803                provider,
804                index,
805                "dropped citations of an item the reply never opened"
806            );
807            return;
808        };
809        let queued = items.iter_mut().find_map(|item| match item {
810            Ok(Item::Event(StreamEvent::End { part, content })) if part.index() == position => {
811                Some(content)
812            }
813            _ => None,
814        });
815        let content = match queued {
816            Some(content) => Some(content),
817            None => self.choice.get_mut(position).and_then(Option::as_mut),
818        };
819        match content {
820            Some(AssistantContent::Text(text)) => {
821                let kept = if replace {
822                    Vec::new()
823                } else {
824                    text.citations().to_vec()
825                };
826                citation::attach(text, kept, citations, provider, index);
827            }
828            _ => tracing::warn!(
829                provider,
830                index,
831                "dropped citations of an item that is not text"
832            ),
833        }
834    }
835
836    pub(crate) fn run_item(
837        &mut self,
838        items: &mut Items,
839        block: Block,
840        fragment: &str,
841    ) -> Result<usize, ProviderError> {
842        let current = self.run.filter(|index| {
843            self.open.get(index).is_some_and(|draft| {
844                matches!(
845                    (&draft.body, &block),
846                    (Body::Text(_), Block::Text)
847                        | (
848                            Body::Reasoning {
849                                redacted: false,
850                                ..
851                            },
852                            Block::Reasoning { redacted: false }
853                        )
854                )
855            })
856        });
857        let index = match current {
858            Some(index) => index,
859            None => {
860                self.end_run(items)?;
861                let index = self.fresh_index();
862                self.open_item(items, index, block, serde_json::Value::Null)?;
863                self.run = Some(index);
864                index
865            }
866        };
867        self.push_item(items, index, fragment)?;
868        Ok(index)
869    }
870
871    pub(crate) fn end_run(&mut self, items: &mut Items) -> Result<(), ProviderError> {
872        match self.run.take() {
873            Some(index) => self.close_item(items, index, Closing::Complete),
874            None => Ok(()),
875        }
876    }
877
878    /// One whole block at the next position, its native kept as given.
879    pub(crate) fn write_content(
880        &mut self,
881        items: &mut Items,
882        content: AssistantContent,
883    ) -> Result<(), ProviderError> {
884        let content = match content {
885            AssistantContent::ToolCall(mut call) => {
886                let item = call
887                    .native
888                    .take()
889                    .map_or(serde_json::Value::Null, |native| native.item);
890                let (id, item) = self.distinct_call_id(Some(call.id), item);
891                call.id = id;
892                with_item(AssistantContent::ToolCall(call), item)
893            }
894            content => content,
895        };
896        let part = self.next();
897        let name = match &content {
898            AssistantContent::ToolCall(call) => Some(call.function.name.clone()),
899            _ => None,
900        };
901        emit(
902            items,
903            StreamEvent::Start {
904                part,
905                kind: kind_of(&content),
906                name,
907            },
908        );
909        match &content {
910            AssistantContent::Text(text) if !text.text.is_empty() => emit(
911                items,
912                StreamEvent::Text {
913                    part,
914                    text: text.text.clone(),
915                },
916            ),
917            AssistantContent::Reasoning(reasoning) if !reasoning.text.is_empty() => emit(
918                items,
919                StreamEvent::Reasoning {
920                    part,
921                    text: reasoning.text.clone(),
922                },
923            ),
924            AssistantContent::ToolCall(call) => emit(
925                items,
926                StreamEvent::Arguments {
927                    part,
928                    json: call.function.arguments_value().to_string(),
929                },
930            ),
931            _ => {}
932        }
933        emit(items, StreamEvent::End { part, content });
934        Ok(())
935    }
936
937    /// Close every item still open when the provider ends the reply, in the
938    /// order they opened. The block a boundary-less wire was streaming is
939    /// complete; any other item was never stated complete, so its block has
940    /// no native. Calls keep what their arguments state.
941    pub(crate) fn close_open(&mut self, items: &mut Items) {
942        let mut open: Vec<(Part, usize)> = self
943            .open
944            .iter()
945            .map(|(index, draft)| (draft.part, *index))
946            .collect();
947        open.sort();
948        for (_, index) in open {
949            let closing = if self.run == Some(index) {
950                Closing::Complete
951            } else {
952                Closing::Incomplete
953            };
954            if self
955                .open
956                .get(&index)
957                .is_some_and(|draft| matches!(draft.body, Body::Call { .. }))
958            {
959                self.unfinished_call = true;
960            }
961            if let Err(error) = self.close_item(items, index, closing) {
962                items.push_back(Err(error));
963            }
964        }
965    }
966
967    /// Record that the item at `position` never completed.
968    fn cut_at(&mut self, position: usize) {
969        self.first_incomplete = Some(
970            self.first_incomplete
971                .map_or(position, |first| first.min(position)),
972        );
973    }
974
975    /// The parts taken so far, in their position (in wire-index order on a
976    /// wire that asked for it, after the lead block); a part that has not
977    /// ended is not among them, and parts from the first incomplete one on
978    /// keep no provider item.
979    pub fn snapshot(&self) -> Vec<AssistantContent> {
980        self.ordered(
981            self.choice.iter().cloned().enumerate().collect(),
982            self.first_incomplete,
983        )
984    }
985
986    /// Who the reply is from.
987    pub fn origin(&self) -> &Origin {
988        &self.origin
989    }
990
991    /// What arrived so far as a response: every part that ended, the text
992    /// the consumer already took of a text or reasoning part still open,
993    /// and the provider's end when it arrived. A reply the provider did not
994    /// end is a failed turn: `failure` is the error that ended it, and
995    /// without one the caller stopped reading. Unless the provider ended the
996    /// reply and the consumer took every part, no block keeps its provider
997    /// item.
998    pub(crate) fn partial(
999        &self,
1000        end: Option<&Finish>,
1001        reply: &Reply,
1002        failure: Option<&ProviderError>,
1003    ) -> CompletionResponse {
1004        let mut response = self.response(end.cloned().unwrap_or_default(), reply.clone());
1005        if end.is_none() {
1006            match failure {
1007                Some(error) => {
1008                    response.error.get_or_insert_with(|| error.to_string());
1009                }
1010                None => {
1011                    response.aborted = Some(
1012                        "the caller stopped reading before the provider ended the reply".to_owned(),
1013                    );
1014                }
1015            }
1016        }
1017        // A turn the consumer has not wholly taken replays only its canonical
1018        // fields, as pi drops an unfinished turn's provider data: no item of
1019        // it can then need a partner that never arrived.
1020        let unfinished =
1021            end.is_none() || !self.open.is_empty() || self.choice.iter().any(Option::is_none);
1022        let canonical_from = if unfinished {
1023            Some(0)
1024        } else {
1025            self.first_incomplete
1026        };
1027        let parts = self
1028            .choice
1029            .iter()
1030            .enumerate()
1031            .map(|(index, part)| {
1032                let part = part.clone().or_else(|| self.open_text.get(&index).cloned());
1033                (index, part)
1034            })
1035            .collect();
1036        response.choice = self.ordered(parts, canonical_from);
1037        response
1038    }
1039
1040    /// `parts` by position, in wire-index order on a wire that asked for
1041    /// it and after the lead block, each from position `canonical_from` on
1042    /// with no provider item.
1043    fn ordered(
1044        &self,
1045        parts: Vec<(usize, Option<AssistantContent>)>,
1046        canonical_from: Option<usize>,
1047    ) -> Vec<AssistantContent> {
1048        let mut parts: Vec<(usize, AssistantContent)> = parts
1049            .into_iter()
1050            .filter_map(|(position, part)| {
1051                let part = part?;
1052                Some(match canonical_from {
1053                    Some(first) if position >= first => (position, canonical(part)),
1054                    _ => (position, part),
1055                })
1056            })
1057            .collect();
1058        if let Some(by_index) = &self.by_index {
1059            parts.sort_by_key(|(position, _)| {
1060                (
1061                    by_index.get(position).copied().unwrap_or(usize::MAX),
1062                    *position,
1063                )
1064            });
1065        }
1066        if let Some(lead) = self.lead {
1067            parts.sort_by_key(|(position, _)| *position != lead);
1068        }
1069        parts.into_iter().map(|(_, part)| part).collect()
1070    }
1071
1072    fn response(&self, end: Finish, reply: Reply) -> CompletionResponse {
1073        let Finish {
1074            usage,
1075            reason,
1076            response_id,
1077            model,
1078            error,
1079        } = end;
1080        use crate::provider_response::reported;
1081        let mut origin = self.origin.clone();
1082        origin.response_model = reported(model);
1083        // A request that named no model is from the model the provider
1084        // reports.
1085        if origin.model.is_empty()
1086            && let Some(model) = &origin.response_model
1087        {
1088            origin.model.clone_from(model);
1089        }
1090        origin.response_id = reported(response_id);
1091        let usage = priced(usage, &origin);
1092        let error = error.or_else(|| {
1093            if !self.wire {
1094                return None;
1095            }
1096            match &reason {
1097                None => Some("the provider ended the reply without a finish reason".to_owned()),
1098                Some(FinishReason::Length) => None,
1099                Some(_) if self.unfinished_call => Some(
1100                    "the provider ended the reply with a tool call it never finished".to_owned(),
1101                ),
1102                Some(_) => None,
1103            }
1104        });
1105        let mut response = CompletionResponse::new(self.snapshot(), usage, origin, reply.raw)
1106            .with_optional_finish_reason(reason)
1107            .accept_unknown_finish_reasons(self.accept_unknown_finish);
1108        response.error = error;
1109        response.provider_request_id = reported(reply.provider_request_id);
1110        response
1111    }
1112}
1113
1114/// `usage` with its cost: the one the provider reported, else the counters
1115/// priced at the catalog's pricing for the requested model, else none.
1116fn priced(usage: Usage, origin: &Origin) -> Usage {
1117    if usage.cost.is_some() {
1118        return usage;
1119    }
1120    let cost = catalog_cost(origin, &usage);
1121    usage.cost(cost)
1122}
1123
1124/// The providers whose replies the catalog does not price: ChatGPT and
1125/// Copilot bill a subscription, not the per-token API rates the catalog
1126/// lists, and an Ollama server is usually local, while the catalog's
1127/// `ollama` prices are Ollama Cloud's and the reply does not say which
1128/// served it.
1129const UNPRICED: [&str; 3] = [
1130    crate::providers::chatgpt::PROVIDER_NAME,
1131    crate::providers::copilot::PROVIDER_NAME,
1132    crate::providers::ollama::PROVIDER_NAME,
1133];
1134
1135/// What `usage` costs at the built-in catalog's pricing for `origin`'s
1136/// model, or `None` when the catalog has no price for it or the provider
1137/// is one it does not price ([`UNPRICED`]).
1138fn catalog_cost(origin: &Origin, usage: &Usage) -> Option<crate::completion::Cost> {
1139    if UNPRICED.contains(&origin.provider.as_str()) {
1140        return None;
1141    }
1142    crate::catalog::lookup(&origin.provider, &origin.model)
1143        .or_else(|| crate::catalog::lookup_snapshot(&origin.provider, &origin.model))
1144        .and_then(|spec| spec.pricing.as_ref())
1145        .and_then(|pricing| pricing.cost(usage))
1146}
1147
1148/// `block` with no provider item: an opaque item no longer replays.
1149pub(crate) fn canonical(block: AssistantContent) -> AssistantContent {
1150    match block {
1151        AssistantContent::Opaque(opaque) => AssistantContent::Opaque(Opaque {
1152            replay: false,
1153            ..opaque
1154        }),
1155        block => block.canonical(),
1156    }
1157}
1158
1159/// The provider item `content` holds, if any.
1160fn item_of(content: &mut AssistantContent) -> Option<&mut serde_json::Value> {
1161    let native = match content {
1162        AssistantContent::Text(text) => text.native.as_mut(),
1163        AssistantContent::ToolCall(call) => call.native.as_mut(),
1164        AssistantContent::Reasoning(reasoning) => reasoning.native.as_mut(),
1165        AssistantContent::Image(image) => image.native.as_mut(),
1166        AssistantContent::Opaque(opaque) => return Some(&mut opaque.item),
1167    };
1168    native.map(|native| &mut native.item)
1169}
1170
1171/// `block` holding `item` as its native, unless there is no item.
1172fn with_item(block: AssistantContent, item: serde_json::Value) -> AssistantContent {
1173    if item.is_null() {
1174        block
1175    } else {
1176        block.with_native(item)
1177    }
1178}
1179
1180fn kind_of(content: &AssistantContent) -> PartKind {
1181    match content {
1182        AssistantContent::Text(_) => PartKind::Text,
1183        AssistantContent::Reasoning(_) => PartKind::Reasoning,
1184        AssistantContent::ToolCall(_) => PartKind::ToolCall,
1185        AssistantContent::Image(_) => PartKind::Image,
1186        AssistantContent::Opaque(_) => PartKind::Opaque,
1187    }
1188}
1189
1190pub(crate) type Items = std::collections::VecDeque<Result<Item<StreamEvent>, ProviderError>>;
1191
1192fn emit(items: &mut Items, event: StreamEvent) {
1193    items.push_back(Ok(Item::Event(event)));
1194}
1195
1196/// Merge `delta` into `item`: each string field but `type` appends to the
1197/// item's string of that key, an array extends the item's array, and any
1198/// other value replaces the key. A delta kind rig has never seen still
1199/// lands in the item.
1200pub fn merge(item: &mut serde_json::Value, delta: &serde_json::Map<String, serde_json::Value>) {
1201    use serde_json::Value;
1202    if !item.is_object() {
1203        *item = Value::Object(serde_json::Map::new());
1204    }
1205    let Value::Object(item) = item else {
1206        return;
1207    };
1208    for (key, value) in delta {
1209        if key == "type" {
1210            continue;
1211        }
1212        match (item.get_mut(key), value) {
1213            (Some(Value::String(existing)), Value::String(fragment)) => {
1214                existing.push_str(fragment);
1215            }
1216            (Some(Value::Array(existing)), Value::Array(more)) => {
1217                existing.extend(more.iter().cloned());
1218            }
1219            _ => {
1220                item.insert(key.clone(), value.clone());
1221            }
1222        }
1223    }
1224}
1225
1226impl Fold<Completion> for Turn {
1227    fn absorb(&mut self, event: &StreamEvent) -> Result<(), ProviderError> {
1228        match event {
1229            StreamEvent::Start { part, .. } => {
1230                if self.choice.len() <= part.index() {
1231                    self.choice.resize(part.index() + 1, None);
1232                }
1233            }
1234            StreamEvent::End { part, content } => {
1235                if self.choice.len() <= part.index() {
1236                    self.choice.resize(part.index() + 1, None);
1237                }
1238                if let Some(slot) = self.choice.get_mut(part.index()) {
1239                    *slot = Some(content.clone());
1240                }
1241                self.open_text.remove(&part.index());
1242            }
1243            StreamEvent::Text { part, text } => {
1244                if let AssistantContent::Text(open) = self
1245                    .open_text
1246                    .entry(part.index())
1247                    .or_insert_with(|| AssistantContent::text(""))
1248                {
1249                    open.text.push_str(text);
1250                }
1251            }
1252            StreamEvent::Reasoning { part, text } => {
1253                if let AssistantContent::Reasoning(open) = self
1254                    .open_text
1255                    .entry(part.index())
1256                    .or_insert_with(|| AssistantContent::reasoning(""))
1257                {
1258                    open.text.push_str(text);
1259                }
1260            }
1261            StreamEvent::Arguments { .. } => {}
1262        }
1263        Ok(())
1264    }
1265
1266    fn finish(self, end: Finish, reply: Reply) -> Result<CompletionResponse, ProviderError> {
1267        if let Some(FinishReason::Other(reason)) = &end.reason {
1268            tracing::warn!(
1269                provider = %self.origin.provider,
1270                reason = %reason,
1271                accepted = self.accept_unknown_finish,
1272                "the provider ended the reply with an unknown finish reason"
1273            );
1274        }
1275        let response = self.response(end, reply);
1276        self.span
1277            .record_response(response.response_id(), response.model(), &response.usage);
1278        Ok(response)
1279    }
1280}
1281
1282impl<'id> Out<'id, Completion> {
1283    /// What `usage` costs at the built-in catalog's pricing for the
1284    /// requested model, for a decoder whose provider bills other counters
1285    /// than the ones it reports as usage.
1286    pub(crate) fn catalog_cost(&self, usage: &Usage) -> Option<crate::completion::Cost> {
1287        catalog_cost(&self.lock().fold.origin, usage)
1288    }
1289
1290    /// Open the provider item at wire `index` as `block`: it takes the next
1291    /// position in the reply. `item` is the item as its first event states
1292    /// it, the base its deltas merge into; `Null` when the block keeps no
1293    /// provider item. An index already open is an error.
1294    pub fn open(
1295        &mut self,
1296        index: usize,
1297        block: Block,
1298        item: serde_json::Value,
1299    ) -> Result<(), ProviderError> {
1300        let mut shared = self.lock();
1301        let Shared { fold, items, .. } = &mut *shared;
1302        fold.open_item(items, index, block, item)
1303    }
1304
1305    /// Append a fragment to the open item at `index`: text, reasoning, or a
1306    /// call's argument JSON, by the item's block.
1307    pub fn push(&mut self, index: usize, fragment: &str) -> Result<(), ProviderError> {
1308        let mut shared = self.lock();
1309        let Shared { fold, items, .. } = &mut *shared;
1310        fold.push_item(items, index, fragment)
1311    }
1312
1313    /// Edit the item at `index` in place: apply a delta with [`merge`], or
1314    /// replace the whole item a provider restates when it finishes.
1315    /// Until the reply ends, an index whose block already closed edits that
1316    /// block's item: the terminal backfill of a field the provider states
1317    /// only at its end.
1318    pub fn edit(
1319        &mut self,
1320        index: usize,
1321        edit: impl FnOnce(&mut serde_json::Value),
1322    ) -> Result<(), ProviderError> {
1323        let mut shared = self.lock();
1324        let Shared { fold, items, .. } = &mut *shared;
1325        match fold.open.get_mut(&index) {
1326            Some(draft) => {
1327                edit(&mut draft.item);
1328                Ok(())
1329            }
1330            None => fold.edit_ended(items, index, edit),
1331        }
1332    }
1333
1334    /// Close the item at `index` the provider never stated complete: its
1335    /// block becomes visible with no native, and an opaque item is kept but
1336    /// does not replay. Empty text and reasoning are dropped, and so is a
1337    /// call with no name. A call's arguments are read by
1338    /// [`ToolFunction::parse`], so malformed ones never fail the reply; a
1339    /// call that never got an id, or reuses one an earlier call took, gets
1340    /// one rig issues.
1341    pub fn close(&mut self, index: usize) -> Result<(), ProviderError> {
1342        let mut shared = self.lock();
1343        let Shared { fold, items, .. } = &mut *shared;
1344        fold.close_item(items, index, Closing::Incomplete)
1345    }
1346
1347    /// Close the item at `index` the provider stated complete: [`Self::close`],
1348    /// with the item as assembled becoming the block's native.
1349    pub fn finish(&mut self, index: usize) -> Result<(), ProviderError> {
1350        let mut shared = self.lock();
1351        let Shared { fold, items, .. } = &mut *shared;
1352        fold.close_item(items, index, Closing::Complete)
1353    }
1354
1355    /// End every item still open as stated complete, for a wire whose end
1356    /// of reply is the provider's statement that its items are done.
1357    pub fn finish_open(&mut self) -> Result<(), ProviderError> {
1358        let mut shared = self.lock();
1359        let Shared { fold, items, .. } = &mut *shared;
1360        let mut open: Vec<(Part, usize)> = fold
1361            .open
1362            .iter()
1363            .map(|(index, draft)| (draft.part, *index))
1364            .collect();
1365        open.sort();
1366        for (_, index) in open {
1367            fold.close_item(items, index, Closing::Complete)?;
1368        }
1369        Ok(())
1370    }
1371
1372    /// Open and finish the item at `index` in one step: a whole block a
1373    /// provider states in one piece, `item` its provider item and `text`
1374    /// its text, reasoning, argument JSON or base64 image data.
1375    pub fn whole(
1376        &mut self,
1377        index: usize,
1378        block: Block,
1379        item: serde_json::Value,
1380        text: &str,
1381    ) -> Result<(), ProviderError> {
1382        self.open(index, block, item)?;
1383        self.push(index, text)?;
1384        self.finish(index)
1385    }
1386
1387    /// Order the response's blocks by wire index rather than by when they
1388    /// opened, for a wire whose indices are the provider's item order and
1389    /// whose end may state items it never streamed. Call it before the
1390    /// first item opens; blocks with no wire index go last.
1391    pub fn order_by_index(&mut self) {
1392        let mut shared = self.lock();
1393        if shared.fold.by_index.is_none() {
1394            shared.fold.by_index = Some(BTreeMap::new());
1395        }
1396    }
1397
1398    /// Hold the open item at `index` first in the response whenever it
1399    /// opened, for a wire that states one block of the turn apart from its
1400    /// order. Its events keep their arrival order. A turn leads with the
1401    /// first item named.
1402    pub fn lead(&mut self, index: usize) -> Result<(), ProviderError> {
1403        let mut shared = self.lock();
1404        let position = shared.fold.draft(index)?.part.index();
1405        shared.fold.lead.get_or_insert(position);
1406        Ok(())
1407    }
1408
1409    /// Replace the text, reasoning or argument JSON of the open item at
1410    /// `index` with `text`, the whole of it as the provider restates it at
1411    /// its end. The fragments already streamed stand and nothing more
1412    /// streams for them; the block ends holding `text`.
1413    pub fn restate(&mut self, index: usize, text: &str) -> Result<(), ProviderError> {
1414        let mut shared = self.lock();
1415        let draft = shared.fold.draft(index)?;
1416        match &mut draft.body {
1417            Body::Text(body) | Body::Reasoning { text: body, .. } => text.clone_into(body),
1418            Body::Call { arguments, .. } => {
1419                text.clone_into(&mut arguments.text);
1420                // A call not yet started streams the restated text whole.
1421                arguments.streamed = if draft.started { text.len() } else { 0 };
1422            }
1423            Body::Image(_) | Body::Opaque { .. } => {}
1424        }
1425        Ok(())
1426    }
1427
1428    /// Write one fragment of a tool call the provider streams, opening the
1429    /// call at its first fragment. Its id and name may arrive in any
1430    /// fragment, and the first name it is given stays: the call starts
1431    /// streaming once it names a tool, carrying
1432    /// the argument text held until then, and its end states its id. Calls
1433    /// are told apart by `index` when the wire gives one. A new id under an
1434    /// index starts a new call once the held call's arguments are a complete
1435    /// object or when it names a tool: some providers send a fresh id, and
1436    /// no name, with every chunk of one call. Without an index (`None`, or
1437    /// the wire's `null`), a fragment with an unseen id opens a call, and one
1438    /// with no id continues the latest call while its arguments are
1439    /// incomplete.
1440    pub fn fragment(
1441        &mut self,
1442        index: Option<usize>,
1443        fragment: CallFragment<'_>,
1444    ) -> Result<(), ProviderError> {
1445        let mut shared = self.lock();
1446        let Shared {
1447            fold: turn, items, ..
1448        } = &mut *shared;
1449        let new_id = fragment.id.filter(|id| stated(id));
1450        let index = match index {
1451            Some(index) => {
1452                let names = fragment.name.is_some_and(|name| !name.is_empty());
1453                let held = turn.open.get(&index).and_then(|draft| match &draft.body {
1454                    Body::Call {
1455                        id: Some(id),
1456                        arguments,
1457                        ..
1458                    } if new_id.is_some_and(|new| id.wire() != new) => Some(arguments.complete()),
1459                    _ => None,
1460                });
1461                if let Some(complete) = held.filter(|complete| *complete || names) {
1462                    // Another call took over the index: the one it held ends.
1463                    let moved = turn.fresh_index();
1464                    if let Some(draft) = turn.open.remove(&index) {
1465                        turn.open.insert(moved, draft);
1466                    }
1467                    let closing = if complete {
1468                        Closing::Complete
1469                    } else {
1470                        Closing::Incomplete
1471                    };
1472                    turn.close_item(items, moved, closing)?;
1473                }
1474                index
1475            }
1476            None => {
1477                let owner = new_id.and_then(|new| {
1478                    turn.open
1479                        .iter()
1480                        .find_map(|(index, draft)| match &draft.body {
1481                            Body::Call { id: Some(id), .. } if id.wire() == new => Some(*index),
1482                            _ => None,
1483                        })
1484                });
1485                let continues = |last: &usize| {
1486                    turn.open.get(last).is_some_and(|draft| {
1487                        matches!(&draft.body, Body::Call { arguments, .. } if !arguments.complete())
1488                    })
1489                };
1490                match (owner, new_id, turn.last_call) {
1491                    (Some(index), _, _) => index,
1492                    (None, None, Some(last)) if continues(&last) => last,
1493                    _ => turn.fresh_index(),
1494                }
1495            }
1496        };
1497        turn.last_call = Some(index);
1498        if !turn.open.contains_key(&index) {
1499            let body = Body::Call {
1500                id: None,
1501                name: String::new(),
1502                arguments: Arguments::default(),
1503            };
1504            turn.insert(index, body, serde_json::Value::Null);
1505        }
1506        let Body::Call {
1507            id,
1508            name,
1509            arguments,
1510        } = &mut turn.draft(index)?.body
1511        else {
1512            return Err(ProviderError::Response(format!(
1513                "the reply sent a call fragment for item {index}, which is not a call"
1514            )));
1515        };
1516        if let Some(call_id) = fragment.id.filter(|id| stated(id)) {
1517            *id = Some(CallId::from_wire(call_id));
1518        }
1519        // A call keeps the first name it is given, so the tool its start
1520        // names is the tool its end names.
1521        if name.is_empty()
1522            && let Some(fragment) = fragment.name.filter(|name| !name.is_empty())
1523        {
1524            fragment.clone_into(name);
1525        }
1526        if let Some(fragment) = fragment.arguments {
1527            arguments.push(fragment, name);
1528        }
1529        stream_call(items, turn.draft(index)?);
1530        Ok(())
1531    }
1532
1533    /// Arguments the provider announced for the call at `index`, used only
1534    /// if no fragment arrives.
1535    pub fn announce(
1536        &mut self,
1537        index: usize,
1538        announced: serde_json::Value,
1539    ) -> Result<(), ProviderError> {
1540        if let Body::Call { arguments, .. } = &mut self.lock().fold.draft(index)?.body {
1541            arguments.announced = Some(announced);
1542        }
1543        Ok(())
1544    }
1545
1546    /// An index no provider item uses, for a wire that indexes nothing.
1547    pub fn fresh_index(&mut self) -> usize {
1548        self.lock().fold.fresh_index()
1549    }
1550
1551    /// A fragment on a wire that marks no item boundaries: it continues the
1552    /// block the last fragment went to while `block` is the same kind, and
1553    /// otherwise closes that block and opens a new one. Returns the block's
1554    /// index, for merging its provider item.
1555    pub fn run(&mut self, block: Block, fragment: &str) -> Result<usize, ProviderError> {
1556        let mut shared = self.lock();
1557        let Shared { fold, items, .. } = &mut *shared;
1558        fold.run_item(items, block, fragment)
1559    }
1560
1561    /// Close the block [`Self::run`] is extending, if any: output of
1562    /// another kind interleaved it.
1563    pub fn end_run(&mut self) -> Result<(), ProviderError> {
1564        let mut shared = self.lock();
1565        let Shared { fold, items, .. } = &mut *shared;
1566        fold.end_run(items)
1567    }
1568
1569    /// Add `citation` to the text item at wire `index`, open or ended. Its
1570    /// span is resolved to bytes of the text when the item closes, or at
1571    /// once for an item already closed. A span that does not resolve, or
1572    /// covers other text than its `quoted`, drops the citation with a
1573    /// warning, as do citations of an item that is not text: a citation
1574    /// never fails a reply.
1575    pub fn cite(&mut self, index: usize, citation: WireCitation) {
1576        let mut shared = self.lock();
1577        let Shared { fold, items, .. } = &mut *shared;
1578        fold.cite_item(items, index, vec![citation], false);
1579    }
1580
1581    /// Put `citations` in place of the text item's own, for a provider that
1582    /// restates them; [`Self::cite`] otherwise.
1583    pub fn set_citations(&mut self, index: usize, citations: Vec<WireCitation>) {
1584        let mut shared = self.lock();
1585        let Shared { fold, items, .. } = &mut *shared;
1586        fold.cite_item(items, index, citations, true);
1587    }
1588
1589    /// A whole block of an already assembled response, at the next
1590    /// position, its native kept as given.
1591    #[cfg(any(test, feature = "test-utils"))]
1592    pub(crate) fn content(&mut self, content: AssistantContent) -> Result<(), ProviderError> {
1593        let mut shared = self.lock();
1594        let Shared { fold, items, .. } = &mut *shared;
1595        fold.write_content(items, content)
1596    }
1597}
1598
1599/// The events a stream of `response` would have carried: each part whole,
1600/// in order.
1601pub(crate) fn events_of(
1602    response: &CompletionResponse,
1603) -> Result<Vec<Item<StreamEvent>>, ProviderError> {
1604    let mut turn = Turn::new(response.origin.clone());
1605    let mut items = Items::new();
1606    for content in &response.choice {
1607        turn.write_content(&mut items, content.clone())?;
1608    }
1609    items.into_iter().collect()
1610}
1611
1612#[cfg(test)]
1613mod tests;