Skip to main content

gproxy_protocol/adapt/generate/stream/fanout/
run.rs

1use super::*;
2use crate::{capability::Upstream, transform::TransformErrorKind};
3
4impl<B: FanoutBridge> FanoutStream<B>
5where
6    B::ClientEvent: FanoutEvent,
7{
8    /// Starts each unstarted child once and returns incremental client bytes.
9    /// Errors retain native receipts in memory; cancellation during a POST
10    /// never triggers a second POST.
11    pub async fn next<U: Upstream, S: StateStore>(
12        &mut self,
13        upstream: &U,
14        target: &U::Target,
15        state: &GenerationStateAccess<'_, S>,
16    ) -> Result<Option<StreamChunk<B::ClientEvent>>, TransformError> {
17        if self
18            .children
19            .iter()
20            .any(|child| child.image_resources_required)
21        {
22            return Err(super::super::missing(
23                "fanout URI delivery requires next_with_image_resources",
24            ));
25        }
26        self.next_using(upstream, target, state, &mut super::driver::Plain)
27            .await
28    }
29    pub(super) async fn next_using<
30        U: Upstream,
31        S: StateStore,
32        D: super::driver::ChildNext<B, S>,
33    >(
34        &mut self,
35        upstream: &U,
36        target: &U::Target,
37        state: &GenerationStateAccess<'_, S>,
38        driver: &mut D,
39    ) -> Result<Option<StreamChunk<B::ClientEvent>>, TransformError> {
40        // Every child was prepared against one state; a step against any
41        // other is refused before a child reads or writes anything.
42        for child in &self.children {
43            child.binding.check(state)?;
44        }
45        if self.failed {
46            return Err(invalid(
47                "fanout stream failed; retained children require reconciliation",
48            ));
49        }
50        if self.finished {
51            return Ok(None);
52        }
53        let result = self.next_inner(upstream, target, state, driver).await;
54        // StateStore/transport interruption may have an uncertain receipt and is
55        // retained for explicit recovery. Proven mapping/limit failure is final.
56        if let Err(error) = &result
57            && matches!(
58                error.kind(),
59                TransformErrorKind::InvalidInput
60                    | TransformErrorKind::InvalidResult
61                    | TransformErrorKind::Limit
62                    | TransformErrorKind::Unsupported
63            )
64        {
65            self.fail();
66        }
67        result
68    }
69    async fn next_inner<U: Upstream, S: StateStore, D: super::driver::ChildNext<B, S>>(
70        &mut self,
71        upstream: &U,
72        target: &U::Target,
73        state: &GenerationStateAccess<'_, S>,
74        driver: &mut D,
75    ) -> Result<Option<StreamChunk<B::ClientEvent>>, TransformError> {
76        while self.index < self.children.len() {
77            if !self.seeded {
78                let reserved: BTreeSet<String> = self.seen.keys().cloned().collect();
79                let child = &mut self.children[self.index];
80                child
81                    .bridge
82                    .as_mut()
83                    .ok_or_else(|| invalid("child bridge consumed"))?
84                    .reserve(IdentityRole::ToolCall, &reserved)?;
85                child
86                    .bridge
87                    .as_mut()
88                    .unwrap()
89                    .reserve(IdentityRole::Response, &self.response_ids)?;
90                child.flow = child.bridge.as_ref().unwrap().identities().clone();
91                self.seeded = true;
92            }
93            if !self.children[self.index].sent
94                && let StreamStart::Rejected(_) = self.children[self.index]
95                    .start(upstream, target, state)
96                    .await?
97            {
98                return Err(invalid(format!(
99                    "fanout child {} rejected request; native receipt retained",
100                    self.index
101                )));
102            }
103            if self.children[self.index].finished {
104                let child = &self.children[self.index];
105                let native = child
106                    .native_result()
107                    .ok_or_else(|| invalid("finished child lacks native receipt"))?;
108                let client = child
109                    .client_result()
110                    .ok_or_else(|| invalid("finished child lacks acknowledged result"))?;
111                let retained = crate::codec::encode_json(&(native, client), self.settings.codec)
112                    .map_err(codec_error)?
113                    .len();
114                self.retained_bytes = self
115                    .retained_bytes
116                    .checked_add(retained)
117                    .ok_or_else(|| limit("fanout retained byte overflow"))?;
118                if self.retained_bytes > self.settings.events.max_bytes
119                    || self.retained_bytes as u64 > self.settings.codec.max_body_bytes
120                {
121                    return Err(limit("fanout retained results exceed aggregate budget"));
122                }
123                for diagnostic in &child.report.diagnostics {
124                    if !self.report.diagnostics.contains(diagnostic) {
125                        self.report.diagnostics.push(diagnostic.clone());
126                    }
127                }
128
129                self.response_ids.extend(
130                    child
131                        .flow
132                        .handles()
133                        .filter(|handle| handle.role == IdentityRole::Response)
134                        .map(|handle| handle.emitted_id),
135                );
136                self.completed.push(client.clone());
137                self.index += 1;
138                self.seeded = false;
139                continue;
140            }
141            let chunk = match driver
142                .next(&mut self.children[self.index], self.index, state)
143                .await
144            {
145                Ok(chunk) => chunk,
146                Err(error) => {
147                    // A prior child already exposed this immutable identity.
148                    // This remains a proven collision even if child persistence
149                    // rejected the alias before its pending event could return.
150                    let duplicate = self.children[self.index]
151                        .ready
152                        .as_ref()
153                        .and_then(|ready| ready.chunk.event.as_ref())
154                        .is_some_and(|event| {
155                            event.tool_declarations(true).iter().any(|(id, _, _)| {
156                                self.seen.get(id).is_some_and(|index| *index != self.index)
157                            })
158                        });
159                    if duplicate {
160                        return Err(invalid("independent children repeat an immutable tool ID"));
161                    }
162                    return Err(error);
163                }
164            };
165            let Some(mut event) = chunk.and_then(|v| v.event) else {
166                continue;
167            };
168            for (id, _, _) in event.tool_declarations(true) {
169                if self.seen.get(&id).is_some_and(|index| *index != self.index) {
170                    return Err(invalid("independent children repeat an immutable tool ID"));
171                }
172                self.seen.insert(id, self.index);
173            }
174
175            event.project(self.index, &self.id, &mut self.created)?;
176            if !event.visible() {
177                continue;
178            }
179            B::ClientEvent::collect(
180                self.collector
181                    .as_mut()
182                    .ok_or_else(|| invalid("aggregate collector consumed"))?,
183                event.clone(),
184            )?;
185            return self.emit(event).map(Some);
186        }
187        if self.aggregate.is_none() {
188            let aggregate = B::ClientEvent::aggregate(
189                self.completed.clone(),
190                self.id.clone(),
191                &mut self.report,
192            )?;
193            // Collect actual projected chunks plus actual summed usage before
194            // accepting the final aggregate. include_usage controls wire only.
195            if let Some(tail) = B::ClientEvent::tail(&aggregate, true)? {
196                B::ClientEvent::collect(
197                    self.collector
198                        .as_mut()
199                        .ok_or_else(|| invalid("aggregate collector consumed"))?,
200                    tail,
201                )?;
202            }
203            if B::ClientEvent::DONE {
204                B::ClientEvent::collect_done(self.collector.as_mut().unwrap())?;
205            }
206            let actual = B::ClientEvent::collected(self.collector.take().unwrap())?;
207            if !B::ClientEvent::equivalent(&actual.value.value, &aggregate) {
208                return Err(invalid(
209                    "fanout actual incremental output differs from completed aggregate",
210                ));
211            }
212            self.tail = B::ClientEvent::tail(&aggregate, self.emit_usage)?;
213            self.aggregate = Some(aggregate);
214        }
215        if !self.tail_emitted {
216            self.tail_emitted = true;
217            if let Some(tail) = self.tail.take() {
218                return self.emit(tail).map(Some);
219            }
220        }
221        let bytes = self.encoder.finish::<B::ClientEvent>()?;
222        self.finished = true;
223        Ok(Some(StreamChunk {
224            event: None,
225            bytes,
226            finished: true,
227        }))
228    }
229    pub async fn collect<U: Upstream, S: StateStore>(
230        &mut self,
231        upstream: &U,
232        target: &U::Target,
233        state: &GenerationStateAccess<'_, S>,
234    ) -> Result<crate::transform::Converted<ClientFull<B>>, TransformError> {
235        while self.next(upstream, target, state).await?.is_some() {}
236        Ok(crate::transform::Converted {
237            value: self
238                .client_result()
239                .ok_or_else(|| invalid("aggregate not acknowledged"))?
240                .clone(),
241            report: self.report.clone(),
242        })
243    }
244}