agent_engine/runtime/google_gemini/
stream.rs

1//! Broker-proxied streaming call to Code Assist's
2//! `v1internal:streamGenerateContent`. Composes:
3//!
4//!   1. `translate_generate_content_request` (pure translator)
5//!   2. `CredentialBroker::proxy_stream` (bearer + pinned host + no redirect)
6//!   3. line-buffered SSE decoding via `from_stream_line`
7//!
8//! The runtime never touches the OAuth access token, the refresh token, or
9//! auth.json — those live behind the broker boundary. Cancellation is
10//! honored between chunks so long streams terminate promptly.
11
12use std::pin::Pin;
13use std::sync::atomic::{AtomicU64, Ordering};
14
15/// Bounded producer/consumer handoff for decoded Gemini events. The producer
16/// awaits capacity (explicit backpressure); dropping the receiver makes
17/// `send` fail immediately and releases the pump.
18pub const GEMINI_EVENT_CHANNEL_CAPACITY: usize = 64;
19static GEMINI_PRODUCED: AtomicU64 = AtomicU64::new(0);
20static GEMINI_FORWARDED: AtomicU64 = AtomicU64::new(0);
21static GEMINI_DROPPED: AtomicU64 = AtomicU64::new(0);
22static GEMINI_ACTIVE_PRODUCERS: AtomicU64 = AtomicU64::new(0);
23
24#[derive(Debug, Clone, Copy)]
25pub struct GeminiEventChannelSnapshot {
26    pub produced_events: u64,
27    pub forwarded_events: u64,
28    pub dropped_events: u64,
29    pub retained_events: u64,
30    pub active_producers: u64,
31}
32
33pub fn gemini_event_channel_snapshot() -> GeminiEventChannelSnapshot {
34    let produced = GEMINI_PRODUCED.load(Ordering::Relaxed);
35    let forwarded = GEMINI_FORWARDED.load(Ordering::Relaxed);
36    let dropped = GEMINI_DROPPED.load(Ordering::Relaxed);
37    GeminiEventChannelSnapshot {
38        produced_events: produced,
39        forwarded_events: forwarded,
40        dropped_events: dropped,
41        retained_events: produced.saturating_sub(forwarded).saturating_sub(dropped),
42        active_producers: GEMINI_ACTIVE_PRODUCERS.load(Ordering::SeqCst),
43    }
44}
45
46async fn send_event(
47    tx: &tokio::sync::mpsc::Sender<Result<GeminiStreamEvent, StreamError>>,
48    event: Result<GeminiStreamEvent, StreamError>,
49) -> bool {
50    GEMINI_PRODUCED.fetch_add(1, Ordering::Relaxed);
51    if tx.send(event).await.is_ok() {
52        GEMINI_FORWARDED.fetch_add(1, Ordering::Relaxed);
53        true
54    } else {
55        GEMINI_DROPPED.fetch_add(1, Ordering::Relaxed);
56        false
57    }
58}
59
60use agent_core::auth::{BrokerError, CredentialBroker, ProxyRequest};
61use bytes::BytesMut;
62use futures::{Stream, StreamExt};
63use tokio_util::sync::CancellationToken;
64
65use super::translate::{
66    from_stream_line, translate_generate_content_request, ChatTurn, GeminiStreamEvent, ToolSpec,
67    MAX_INBOUND_LINE_BYTES,
68};
69
70/// Runtime-facing outcome of a single streamed model turn.
71#[derive(Debug, thiserror::Error)]
72pub enum StreamError {
73    #[error("gemini stream error: {0}")]
74    Broker(#[from] BrokerError),
75    #[error("gemini stream error: {0}")]
76    Decode(String),
77    #[error("gemini stream error: cancelled")]
78    Cancelled,
79    #[error("gemini stream error: line buffer exceeded {MAX_INBOUND_LINE_BYTES} bytes")]
80    LineTooLarge,
81}
82
83/// Build the exact-byte `ProxyRequest` for one streamed Code Assist turn.
84/// The JSON envelope is serialized **once** via
85/// [`ProxyRequest::post_json_exact`]; the returned [`bytes::Bytes`] handle is
86/// the very buffer `LocalBroker` sends verbatim, so a caller-side trace
87/// digest over it describes the true wire body on the local-broker path.
88pub fn build_stream_request(
89    model: impl Into<String>,
90    project: Option<String>,
91    system_prompt: Option<String>,
92    turns: &[ChatTurn],
93    tools: &[ToolSpec],
94) -> Result<(ProxyRequest, bytes::Bytes), StreamError> {
95    let body = translate_generate_content_request(model, project, system_prompt, turns, tools);
96    let body = serde_json::to_value(&body)
97        .map_err(|e| StreamError::Decode(format!("failed to serialize request: {e}")))?;
98    ProxyRequest::post_json_exact(
99        "google-gemini",
100        "/v1internal:streamGenerateContent",
101        body,
102        true,
103    )
104    .map_err(StreamError::Broker)
105}
106
107/// Start a broker-proxied streaming turn against Code Assist and yield
108/// decoded events as they arrive. The returned stream is `Send + 'static`
109/// so it can be forwarded through the runtime's event bus.
110pub async fn stream_gemini<B: CredentialBroker + ?Sized>(
111    broker: &B,
112    model: impl Into<String>,
113    project: Option<String>,
114    system_prompt: Option<String>,
115    turns: &[ChatTurn],
116    tools: &[ToolSpec],
117    cancel: CancellationToken,
118) -> Result<Pin<Box<dyn Stream<Item = Result<GeminiStreamEvent, StreamError>> + Send>>, StreamError>
119{
120    let (request, _bytes) = build_stream_request(model, project, system_prompt, turns, tools)?;
121    stream_gemini_request(broker, request, cancel).await
122}
123
124/// Start one broker stream attempt from a prebuilt request (see
125/// [`build_stream_request`]). Retry loops reuse the same request value per
126/// attempt so every attempt sends identical bytes.
127pub async fn stream_gemini_request<B: CredentialBroker + ?Sized>(
128    broker: &B,
129    request: ProxyRequest,
130    cancel: CancellationToken,
131) -> Result<Pin<Box<dyn Stream<Item = Result<GeminiStreamEvent, StreamError>> + Send>>, StreamError>
132{
133    let mut byte_stream = broker.proxy_stream(request).await?;
134
135    // Pump on a dedicated task and forward decoded events through mpsc. The
136    // pump is `'static`-bounded, so we don't need async-stream / self-refs.
137    let (tx, rx) = tokio::sync::mpsc::channel::<Result<GeminiStreamEvent, StreamError>>(
138        GEMINI_EVENT_CHANNEL_CAPACITY,
139    );
140    GEMINI_ACTIVE_PRODUCERS.fetch_add(1, Ordering::SeqCst);
141    tokio::spawn(async move {
142        struct ProducerGauge;
143        impl Drop for ProducerGauge {
144            fn drop(&mut self) {
145                GEMINI_ACTIVE_PRODUCERS.fetch_sub(1, Ordering::SeqCst);
146            }
147        }
148        let _gauge = ProducerGauge;
149        let mut buf = BytesMut::new();
150        loop {
151            tokio::select! {
152                biased;
153                _ = cancel.cancelled() => {
154                    let _ = send_event(&tx, Err(StreamError::Cancelled)).await;
155                    return;
156                }
157                chunk = byte_stream.next() => {
158                    match chunk {
159                        Some(Ok(bytes)) => {
160                            if buf.len().saturating_add(bytes.len()) > MAX_INBOUND_LINE_BYTES {
161                                let _ = send_event(&tx, Err(StreamError::LineTooLarge)).await;
162                                return;
163                            }
164                            buf.extend_from_slice(&bytes);
165                            while let Some(pos) = buf.iter().position(|b| *b == b'\n') {
166                                let line = buf.split_to(pos + 1);
167                                let line = std::str::from_utf8(&line)
168                                    .map(|s| s.trim_end_matches(['\r', '\n']).to_string());
169                                match line {
170                                    Ok(line) => match from_stream_line(&line) {
171                                        Ok(events) => {
172                                            for ev in events {
173                                                if !send_event(&tx, Ok(ev)).await {
174                                                    return;
175                                                }
176                                            }
177                                        }
178                                        Err(e) => {
179                                            let _ = send_event(&tx, Err(StreamError::Decode(e))).await;
180                                            return;
181                                        }
182                                    },
183                                    Err(_) => {
184                                        let _ = send_event(&tx, Err(StreamError::Decode(
185                                            "non-utf8 in stream".into(),
186                                        ))).await;
187                                        return;
188                                    }
189                                }
190                            }
191                        }
192                        Some(Err(e)) => {
193                            let _ = send_event(&tx, Err(StreamError::Broker(e))).await;
194                            return;
195                        }
196                        None => {
197                            if !buf.is_empty() {
198                                let tail = std::str::from_utf8(&buf)
199                                    .map(|s| s.trim_end().to_string())
200                                    .unwrap_or_default();
201                                if !tail.is_empty() {
202                                    match from_stream_line(&tail) {
203                                        Ok(events) => {
204                                            for ev in events {
205                                                let _ = send_event(&tx, Ok(ev)).await;
206                                            }
207                                        }
208                                        Err(e) => {
209                                            let _ = send_event(&tx, Err(StreamError::Decode(e))).await;
210                                        }
211                                    }
212                                }
213                            }
214                            return;
215                        }
216                    }
217                }
218            }
219        }
220    });
221
222    Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx)))
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use agent_core::auth::{AccessToken, ProxyByteStream, ProxyMethod, ProxyResponse};
229    use async_trait::async_trait;
230    use futures::stream;
231    use std::sync::{Arc, Mutex};
232
233    struct StubBroker {
234        chunks: Mutex<Option<Vec<Result<bytes::Bytes, BrokerError>>>>,
235        seen: Arc<Mutex<Option<ProxyRequest>>>,
236    }
237
238    impl StubBroker {
239        fn new(chunks: Vec<Result<bytes::Bytes, BrokerError>>) -> Self {
240            Self {
241                chunks: Mutex::new(Some(chunks)),
242                seen: Arc::default(),
243            }
244        }
245    }
246
247    #[async_trait]
248    impl CredentialBroker for StubBroker {
249        async fn access_token(
250            &self,
251            _p: agent_core::auth::OAuthProviderId,
252        ) -> Result<AccessToken, BrokerError> {
253            Err(BrokerError::NotConfigured("stub".into()))
254        }
255        async fn proxy(&self, _r: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
256            Err(BrokerError::Denied("not implemented".into()))
257        }
258        async fn proxy_stream(
259            &self,
260            request: ProxyRequest,
261        ) -> Result<ProxyByteStream, BrokerError> {
262            *self.seen.lock().unwrap() = Some(request);
263            let chunks = self.chunks.lock().unwrap().take().unwrap_or_default();
264            Ok(Box::pin(stream::iter(chunks)))
265        }
266        async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
267            Err(BrokerError::Denied("not implemented".into()))
268        }
269        async fn capabilities(&self) -> Result<Vec<agent_core::auth::ProviderStatus>, BrokerError> {
270            Ok(vec![])
271        }
272    }
273
274    fn chunk(s: &str) -> Result<bytes::Bytes, BrokerError> {
275        Ok(bytes::Bytes::copy_from_slice(s.as_bytes()))
276    }
277
278    /// A stalled consumer retains at most the explicit channel capacity;
279    /// dropping the stream releases the producer task.
280    #[tokio::test]
281    async fn slow_consumer_is_bounded_at_the_model_event_production_boundary() {
282        let mut lines = Vec::new();
283        for _ in 0..10_000 {
284            lines.push(chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\"}]}}]}}\n"));
285        }
286        let before = gemini_event_channel_snapshot();
287        let broker = StubBroker::new(lines);
288        let stream = stream_gemini(
289            &broker,
290            "gemini-2.5-pro",
291            None,
292            None,
293            &[ChatTurn::User { text: "hi".into() }],
294            &[],
295            CancellationToken::new(),
296        )
297        .await
298        .unwrap();
299        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
300        let stalled = gemini_event_channel_snapshot();
301        let retained_delta = stalled
302            .produced_events
303            .saturating_sub(before.produced_events)
304            .saturating_sub(
305                stalled
306                    .forwarded_events
307                    .saturating_sub(before.forwarded_events),
308            )
309            .saturating_sub(stalled.dropped_events.saturating_sub(before.dropped_events));
310        assert!(
311            retained_delta > 0,
312            "stalled consumer must retain queued events"
313        );
314        assert!(retained_delta <= GEMINI_EVENT_CHANNEL_CAPACITY as u64);
315        let active_before = before.active_producers;
316        drop(stream);
317        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
318        while gemini_event_channel_snapshot().active_producers > active_before
319            && tokio::time::Instant::now() < deadline
320        {
321            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
322        }
323        assert!(gemini_event_channel_snapshot().active_producers <= active_before);
324    }
325
326    #[tokio::test]
327    async fn streams_text_deltas_and_tool_calls_in_order() {
328        let broker = StubBroker::new(vec![
329            chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello \"}]}}]}}\n"),
330            chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"world\"}]}}]}}\n"),
331            chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"do\",\"args\":{}}}]},\"finishReason\":\"STOP\"}]}}\n"),
332            chunk("data: [DONE]\n"),
333        ]);
334        let mut stream = stream_gemini(
335            &broker,
336            "gemini-2.5-pro",
337            Some("proj".into()),
338            None,
339            &[ChatTurn::User { text: "hi".into() }],
340            &[],
341            CancellationToken::new(),
342        )
343        .await
344        .unwrap();
345
346        let mut got = Vec::new();
347        while let Some(ev) = stream.next().await {
348            got.push(ev.unwrap());
349        }
350        assert!(matches!(&got[0], GeminiStreamEvent::TextDelta(t) if t == "Hello "));
351        assert!(matches!(&got[1], GeminiStreamEvent::TextDelta(t) if t == "world"));
352        assert!(matches!(&got[2], GeminiStreamEvent::ToolCall(c) if c.name == "do"));
353        assert!(matches!(&got[3], GeminiStreamEvent::Finish { reason: Some(r) } if r == "STOP"));
354        assert!(matches!(
355            &got[4],
356            GeminiStreamEvent::Finish { reason: None }
357        ));
358
359        // Wire request went to the exact pinned method + provider + streamed.
360        let seen = broker.seen.lock().unwrap().clone().unwrap();
361        assert_eq!(seen.provider, "google-gemini");
362        assert_eq!(seen.path, "/v1internal:streamGenerateContent");
363        assert!(seen.stream);
364        assert_eq!(seen.method, ProxyMethod::Post);
365        let body = seen.body.unwrap();
366        assert_eq!(body["model"], "gemini-2.5-pro");
367        assert_eq!(body["project"], "proj");
368        assert_eq!(body["request"]["contents"][0]["parts"][0]["text"], "hi");
369    }
370
371    #[tokio::test]
372    async fn cancellation_terminates_stream_between_chunks() {
373        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
374        // First chunk is delivered; then the stream stays open, so the
375        // cancel token is the only path to termination.
376        tx.send(chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"one\"}]}}]}}\n"))
377            .unwrap();
378
379        struct Blocking(
380            Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<Result<bytes::Bytes, BrokerError>>>>,
381        );
382        #[async_trait]
383        impl CredentialBroker for Blocking {
384            async fn access_token(
385                &self,
386                _p: agent_core::auth::OAuthProviderId,
387            ) -> Result<AccessToken, BrokerError> {
388                Err(BrokerError::NotConfigured("stub".into()))
389            }
390            async fn proxy(&self, _r: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
391                Err(BrokerError::Denied("not implemented".into()))
392            }
393            async fn proxy_stream(&self, _r: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
394                let rx = self.0.lock().unwrap().take().unwrap();
395                Ok(Box::pin(
396                    tokio_stream::wrappers::UnboundedReceiverStream::new(rx),
397                ))
398            }
399            async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
400                Err(BrokerError::Denied("not implemented".into()))
401            }
402            async fn capabilities(
403                &self,
404            ) -> Result<Vec<agent_core::auth::ProviderStatus>, BrokerError> {
405                Ok(vec![])
406            }
407        }
408
409        let broker = Blocking(Mutex::new(Some(rx)));
410        let cancel = CancellationToken::new();
411        let mut stream = stream_gemini(
412            &broker,
413            "gemini-2.5-pro",
414            None,
415            None,
416            &[ChatTurn::User { text: "hi".into() }],
417            &[],
418            cancel.clone(),
419        )
420        .await
421        .unwrap();
422
423        // First event: TextDelta("one").
424        let first = stream.next().await.unwrap().unwrap();
425        assert!(matches!(first, GeminiStreamEvent::TextDelta(t) if t == "one"));
426
427        // Cancel and expect Cancelled next.
428        cancel.cancel();
429        let next = stream.next().await.unwrap();
430        assert!(matches!(next, Err(StreamError::Cancelled)));
431        // Sender is dropped — stream ends cleanly after cancel.
432        drop(tx);
433    }
434
435    #[tokio::test]
436    async fn malformed_chunk_yields_decode_error_and_terminates() {
437        let broker = StubBroker::new(vec![
438            chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"ok\"}]}}]}}\n"),
439            chunk("data: {not json\n"),
440            chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"never\"}]}}]}}\n"),
441        ]);
442        let mut stream = stream_gemini(
443            &broker,
444            "gemini-2.5-pro",
445            None,
446            None,
447            &[ChatTurn::User { text: "hi".into() }],
448            &[],
449            CancellationToken::new(),
450        )
451        .await
452        .unwrap();
453
454        let a = stream.next().await.unwrap().unwrap();
455        assert!(matches!(a, GeminiStreamEvent::TextDelta(t) if t == "ok"));
456        let b = stream.next().await.unwrap();
457        assert!(matches!(b, Err(StreamError::Decode(_))));
458        // No further events after a decode error.
459        assert!(stream.next().await.is_none());
460    }
461
462    #[tokio::test]
463    async fn broker_transport_error_is_surfaced() {
464        let broker = StubBroker::new(vec![
465            chunk("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"partial\"}]}}]}}\n"),
466            Err(BrokerError::Transport("upstream reset".into())),
467        ]);
468        let mut stream = stream_gemini(
469            &broker,
470            "gemini-2.5-pro",
471            None,
472            None,
473            &[ChatTurn::User { text: "hi".into() }],
474            &[],
475            CancellationToken::new(),
476        )
477        .await
478        .unwrap();
479        stream.next().await.unwrap().unwrap();
480        match stream.next().await.unwrap() {
481            Err(StreamError::Broker(BrokerError::Transport(m))) => {
482                assert!(m.contains("upstream reset"))
483            }
484            other => panic!("expected broker transport error, got {other:?}"),
485        }
486    }
487}
agent_engine/runtime/google_gemini/stream.rs

agent_engine/runtime/google_gemini/
stream.rs
