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

1use super::*;
2use crate::{
3    adapt::generate::{
4        GenerationResources,
5        image_resources::{ImageStreamProgress, ResourceSend, ResourceSync},
6    },
7    capability::ResourceAccess,
8    transform::generate::gemini_responses::stream::ResponsesToGeminiStream,
9    wire::gemini as g,
10};
11
12struct Images<'a, 'r, R: ResourceAccess> {
13    resources: &'a GenerationResources<'r, R>,
14    progress: &'a mut [ImageStreamProgress<R::PublishedHandle>],
15}
16
17impl<R: ResourceAccess + ResourceSync, S: StateStore + ResourceSync>
18    super::driver::ChildNext<ResponsesToGeminiStream, S> for Images<'_, '_, R>
19where
20    R::Scope: ResourceSync,
21    R::PublishedHandle: ResourceSend,
22    S::Scope: ResourceSync,
23{
24    async fn next(
25        &mut self,
26        child: &mut StreamInvocation<ResponsesToGeminiStream>,
27        index: usize,
28        state: &GenerationStateAccess<'_, S>,
29    ) -> Result<Option<StreamChunk<g::GenerateContentResponseBody>>, TransformError> {
30        let mut used_bytes = 0u64;
31        let mut used_count = 0usize;
32        for (other, progress) in self.progress.iter().enumerate() {
33            if other == index {
34                continue;
35            }
36            let (count, bytes) = progress.publication_budget()?;
37            used_bytes = used_bytes
38                .checked_add(bytes)
39                .ok_or_else(|| limit("fanout publication byte overflow"))?;
40            used_count = used_count
41                .checked_add(count)
42                .ok_or_else(|| limit("fanout publication count overflow"))?;
43        }
44        let mut limits = self.resources.limits;
45        limits.max_body_bytes = limits
46            .max_body_bytes
47            .checked_sub(used_bytes)
48            .ok_or_else(|| limit("fanout publication byte budget exceeded"))?;
49        let max_references = self
50            .resources
51            .max_references
52            .checked_sub(used_count)
53            .ok_or_else(|| limit("fanout publication count budget exceeded"))?;
54        let resources = GenerationResources {
55            access: self.resources.access,
56            scope: self.resources.scope,
57            limits,
58            max_references,
59            now: self.resources.now,
60            target: self.resources.target,
61        };
62        child
63            .next_with_image_resources(state, &resources, &mut self.progress[index])
64            .await
65    }
66}
67
68impl FanoutStream<ResponsesToGeminiStream> {
69    /// Keep one progress entry per child. Publication operation IDs and receipts
70    /// survive cancellation, and resource budgets apply to the whole group.
71    pub async fn next_with_image_resources<U: crate::capability::Upstream, S, R>(
72        &mut self,
73        upstream: &U,
74        target: &U::Target,
75        state: &GenerationStateAccess<'_, S>,
76        resources: &GenerationResources<'_, R>,
77        progress: &mut [ImageStreamProgress<R::PublishedHandle>],
78    ) -> Result<Option<StreamChunk<g::GenerateContentResponseBody>>, TransformError>
79    where
80        R: ResourceAccess + ResourceSync,
81        R::Scope: ResourceSync,
82        R::PublishedHandle: ResourceSend,
83        S: StateStore + ResourceSync,
84        S::Scope: ResourceSync,
85    {
86        if progress.len() != self.children.len() {
87            return Err(TransformError::shape(
88                "fanout.image_progress",
89                "one retained image progress entry per child required",
90            ));
91        }
92        for (child, image_progress) in self.children.iter().zip(progress.iter_mut()) {
93            if image_progress.revision != child.resource_revision {
94                return Err(conflict(
95                    "fanout image progress changed for an active or completed child",
96                ));
97            }
98            image_progress.bind(child.selected.identities.response.namespace())?;
99        }
100        self.next_using(
101            upstream,
102            target,
103            state,
104            &mut Images {
105                resources,
106                progress,
107            },
108        )
109        .await
110    }
111    pub async fn collect_with_image_resources<U: crate::capability::Upstream, S, R>(
112        &mut self,
113        upstream: &U,
114        target: &U::Target,
115        state: &GenerationStateAccess<'_, S>,
116        resources: &GenerationResources<'_, R>,
117        progress: &mut [ImageStreamProgress<R::PublishedHandle>],
118    ) -> Result<crate::transform::Converted<g::GenerateContentResponseBody>, TransformError>
119    where
120        R: ResourceAccess + ResourceSync,
121        R::Scope: ResourceSync,
122        R::PublishedHandle: ResourceSend,
123        S: StateStore + ResourceSync,
124        S::Scope: ResourceSync,
125    {
126        while self
127            .next_with_image_resources(upstream, target, state, resources, progress)
128            .await?
129            .is_some()
130        {}
131        Ok(crate::transform::Converted {
132            value: self
133                .client_result()
134                .ok_or_else(|| invalid("aggregate not acknowledged"))?
135                .clone(),
136            report: self.report.clone(),
137        })
138    }
139}
gproxy_protocol/adapt/generate/stream/fanout/images.rs

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