bamboo_engine/runtime/stream/
handler.rs

1use std::future::Future;
2use std::sync::Arc;
3use std::time::Duration;
4
5use chrono::{DateTime, Utc};
6
7use tokio::sync::mpsc;
8use tokio::time::Instant;
9use tokio_util::sync::CancellationToken;
10
11use bamboo_agent_core::tools::ToolCall;
12use bamboo_agent_core::{AgentError, AgentEvent, StreamTimeoutError, StreamTimeoutPhase};
13use bamboo_config::StreamTimeoutConfig;
14use bamboo_llm::retry::{InitialHttpRetry, InitialHttpRetryObserver};
15use bamboo_llm::LLMStream;
16
17mod chunk_handling;
18mod consume;
19mod stream_state;
20
21/// Resolved identifiers and deadlines for one stream. Identifiers are sanitized
22/// before they reach timeout diagnostics; prompts and provider payloads are
23/// never retained here.
24#[derive(Debug, Clone)]
25pub struct StreamTimeoutContext {
26    pub(crate) policy: StreamTimeoutConfig,
27    pub(crate) provider: Option<String>,
28    pub(crate) model: Option<String>,
29    request_started_at: Option<Instant>,
30    turn_retry_eligible: bool,
31}
32
33impl StreamTimeoutContext {
34    pub fn new(policy: StreamTimeoutConfig, provider: Option<&str>, model: Option<&str>) -> Self {
35        let policy = match policy.validate() {
36            Ok(()) => policy,
37            Err(error) => {
38                tracing::warn!(
39                    "invalid programmatic stream timeout policy ({error}); using safe defaults"
40                );
41                StreamTimeoutConfig::default()
42            }
43        };
44        Self {
45            policy,
46            provider: provider.and_then(sanitize_identifier),
47            model: model.and_then(sanitize_identifier),
48            request_started_at: None,
49            turn_retry_eligible: false,
50        }
51    }
52
53    /// Mark this as the primary turn response. Only that stream can safely ask
54    /// the outer agent loop to replay a timeout that occurs before semantic
55    /// output; auxiliary streams may run after durable turn side effects.
56    pub(crate) fn allow_turn_retry_before_semantic_output(mut self) -> Self {
57        self.turn_retry_eligible = true;
58        self
59    }
60
61    /// Bind a fresh request-dispatch timestamp so the first-semantic deadline
62    /// includes provider bootstrap time as documented.
63    pub(crate) fn begin_request(mut self) -> Self {
64        self.request_started_at = Some(Instant::now());
65        self
66    }
67
68    fn timeout_error(
69        &self,
70        session_id: &str,
71        phase: StreamTimeoutPhase,
72        deadline: Duration,
73        last_transport: Duration,
74        last_semantic: Option<Duration>,
75    ) -> AgentError {
76        self.timeout_error_with_http_retry(
77            session_id,
78            phase,
79            deadline,
80            last_transport,
81            last_semantic,
82            None,
83        )
84    }
85
86    fn timeout_error_with_http_retry(
87        &self,
88        session_id: &str,
89        phase: StreamTimeoutPhase,
90        deadline: Duration,
91        last_transport: Duration,
92        last_semantic: Option<Duration>,
93        http_retry: Option<InitialHttpRetry>,
94    ) -> AgentError {
95        let mut timeout = StreamTimeoutError::new(
96            phase,
97            deadline,
98            self.provider.clone(),
99            self.model.clone(),
100            last_transport,
101            last_semantic,
102            self.turn_retry_eligible,
103        );
104        if let Some(retry) = http_retry {
105            timeout = timeout.with_last_http_retry(retry.status, retry.delay);
106        }
107        tracing::warn!("[{}] LLM stream watchdog expired: {}", session_id, timeout,);
108        AgentError::StreamTimeout(timeout)
109    }
110}
111
112impl Default for StreamTimeoutContext {
113    fn default() -> Self {
114        Self::new(StreamTimeoutConfig::default(), None, None)
115    }
116}
117
118/// Wait for a provider call to establish its response stream while preserving
119/// cancellation and the existing transport-idle policy.
120///
121/// Provider implementations return [`LLMStream`] only after the initial HTTP
122/// response has been established. Without this outer watchdog, a proxy that
123/// accepts the request but never returns response headers can hold the agent
124/// forever before the normal per-frame stream watchdog starts.
125pub(crate) async fn await_stream_bootstrap<F, T>(
126    future: F,
127    cancel_token: &CancellationToken,
128    session_id: &str,
129    timeout_context: &StreamTimeoutContext,
130) -> Result<T, AgentError>
131where
132    F: Future<Output = T>,
133{
134    let started_at = timeout_context
135        .request_started_at
136        .unwrap_or_else(Instant::now);
137    let deadline = Duration::from_secs(timeout_context.policy.transport_idle_timeout_secs);
138    let expires_at = started_at + deadline;
139    let http_retry_observer = Arc::new(InitialHttpRetryObserver::default());
140
141    bamboo_llm::retry::observe_initial_http_retries(http_retry_observer.clone(), async {
142        tokio::select! {
143            biased;
144            _ = cancel_token.cancelled() => Err(AgentError::Cancelled),
145            result = future => Ok(result),
146            _ = tokio::time::sleep_until(expires_at) => {
147                let now = Instant::now();
148                Err(timeout_context.timeout_error_with_http_retry(
149                    session_id,
150                    StreamTimeoutPhase::Bootstrap,
151                    deadline,
152                    now.saturating_duration_since(started_at),
153                    None,
154                    http_retry_observer.last(),
155                ))
156            }
157        }
158    })
159    .await
160}
161
162fn sanitize_identifier(value: &str) -> Option<String> {
163    let value = value.trim();
164    if value.is_empty() {
165        return None;
166    }
167    let sanitized: String = value
168        .chars()
169        .take(120)
170        .map(|character| match character {
171            character if character.is_ascii_alphanumeric() => character,
172            '-' | '_' | '.' | ':' | '/' | '@' => character,
173            _ => '_',
174        })
175        .collect();
176    (!sanitized.is_empty()).then_some(sanitized)
177}
178
179/// Raw authoritative usage fields preserved from provider terminal events.
180///
181/// Each field remains optional so callers can distinguish an explicit
182/// provider-reported zero from an omitted value. Flat counters on
183/// [`StreamHandlingOutput`] remain the normalized compatibility view.
184#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
185pub struct ProviderUsageSnapshot {
186    pub input_tokens: Option<u64>,
187    pub output_tokens: Option<u64>,
188    pub total_tokens: Option<u64>,
189    pub reasoning_tokens: Option<u64>,
190    pub cache_creation_input_tokens: Option<u64>,
191    pub cache_read_input_tokens: Option<u64>,
192    /// OpenAI Responses cache-write volume. This is raw provider metadata and
193    /// is not folded into the disjoint legacy prompt-cache counters.
194    pub cache_write_input_tokens: Option<u64>,
195}
196
197#[derive(Debug, Clone, PartialEq, Eq)]
198pub struct VisibleMessageIdentity {
199    pub message_id: String,
200    pub created_at: DateTime<Utc>,
201}
202
203impl VisibleMessageIdentity {
204    pub(crate) fn apply_to(
205        self,
206        mut message: bamboo_agent_core::Message,
207    ) -> bamboo_agent_core::Message {
208        message.id = self.message_id;
209        message.created_at = self.created_at;
210        message
211    }
212}
213
214pub struct StreamHandlingOutput {
215    pub response_id: Option<String>,
216    /// Stable identity shared by the safe realtime text and persisted message.
217    pub visible_message: Option<VisibleMessageIdentity>,
218    pub content: String,
219    pub reasoning_content: String,
220    /// Provider-minted signature covering `reasoning_content`, present only
221    /// when the turn's thinking arrived as exactly one signed Anthropic
222    /// `thinking` block — see [`bamboo_llm::LLMChunk::ReasoningSignature`] (#520).
223    pub reasoning_signature: Option<String>,
224    pub token_count: usize,
225    pub tool_calls: Vec<ToolCall>,
226    pub output_tokens: u64,
227    pub thinking_tokens: u64,
228    pub cache_creation_input_tokens: u64,
229    pub cache_read_input_tokens: u64,
230    /// Merged authoritative provider snapshot, when at least one provider-usage
231    /// chunk was observed. Repeated cumulative snapshots are idempotent, absent
232    /// fields do not erase known values, and explicit zeros remain `Some(0)`.
233    pub provider_usage: Option<ProviderUsageSnapshot>,
234    /// Normalized non-cached ("fresh") input, disjoint from the adjacent cache
235    /// counters. When a provider total is available this is derived from that
236    /// total with a saturating, cache-subset policy; the raw total remains in
237    /// [`Self::provider_usage`].
238    pub input_tokens: u64,
239    /// Validated native items awaiting atomic commit with this round's ordinary
240    /// assistant message. Never exposed as streaming UI events.
241    pub provider_transcript_items: Vec<bamboo_domain::ProviderTranscriptItem>,
242}
243
244#[derive(Debug, Clone, serde::Serialize)]
245pub(crate) struct PartialToolCallSnapshot {
246    pub id: String,
247    pub tool_type: String,
248    pub name: String,
249    pub arguments: String,
250    pub index: Option<u32>,
251}
252
253/// Crate-private interrupted-stream payload.  Keeping this separate from the
254/// public successful [`StreamHandlingOutput`] avoids a source-breaking public
255/// field while retaining fragments that finalization intentionally drops or
256/// normalizes.
257pub(crate) struct InterruptedStreamOutput {
258    pub visible_message: Option<VisibleMessageIdentity>,
259    pub content: String,
260    pub reasoning_content: String,
261    pub partial_tool_calls: Vec<PartialToolCallSnapshot>,
262}
263
264impl From<&bamboo_agent_core::tools::PartialToolCall> for PartialToolCallSnapshot {
265    fn from(value: &bamboo_agent_core::tools::PartialToolCall) -> Self {
266        Self {
267            id: value.id.clone(),
268            tool_type: value.tool_type.clone(),
269            name: value.name.clone(),
270            arguments: value.arguments.clone(),
271            index: value.index,
272        }
273    }
274}
275
276/// A stream failure together with every semantic fragment accumulated before
277/// the failure.  The agent round uses this to create a durable, explicitly
278/// interrupted assistant record instead of losing already-visible output.
279pub(crate) struct StreamHandlingFailure {
280    pub error: AgentError,
281    /// Failures are already the cold path; boxing the three-buffer snapshot
282    /// keeps this error small without changing the preserved fragment data.
283    pub partial_output: Box<InterruptedStreamOutput>,
284}
285
286pub async fn consume_llm_stream(
287    stream: LLMStream,
288    event_tx: &mpsc::Sender<AgentEvent>,
289    cancel_token: &CancellationToken,
290    session_id: &str,
291) -> Result<StreamHandlingOutput, AgentError> {
292    consume_llm_stream_with_context(
293        stream,
294        event_tx,
295        cancel_token,
296        session_id,
297        &StreamTimeoutContext::default(),
298    )
299    .await
300}
301
302pub async fn consume_llm_stream_with_context(
303    stream: LLMStream,
304    event_tx: &mpsc::Sender<AgentEvent>,
305    cancel_token: &CancellationToken,
306    session_id: &str,
307    timeout_context: &StreamTimeoutContext,
308) -> Result<StreamHandlingOutput, AgentError> {
309    consume::consume_llm_stream_internal(
310        stream,
311        Some(event_tx),
312        cancel_token,
313        session_id,
314        timeout_context,
315    )
316    .await
317}
318
319pub(crate) async fn consume_llm_stream_with_context_and_partial(
320    stream: LLMStream,
321    event_tx: &mpsc::Sender<AgentEvent>,
322    cancel_token: &CancellationToken,
323    session_id: &str,
324    timeout_context: &StreamTimeoutContext,
325) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
326    consume::consume_llm_stream_internal_with_partial(
327        stream,
328        Some(event_tx),
329        cancel_token,
330        session_id,
331        timeout_context,
332    )
333    .await
334}
335
336pub async fn consume_llm_stream_silent(
337    stream: LLMStream,
338    cancel_token: &CancellationToken,
339    session_id: &str,
340) -> Result<StreamHandlingOutput, AgentError> {
341    consume_llm_stream_silent_with_context(
342        stream,
343        cancel_token,
344        session_id,
345        &StreamTimeoutContext::default(),
346    )
347    .await
348}
349
350pub async fn consume_llm_stream_silent_with_context(
351    stream: LLMStream,
352    cancel_token: &CancellationToken,
353    session_id: &str,
354    timeout_context: &StreamTimeoutContext,
355) -> Result<StreamHandlingOutput, AgentError> {
356    consume::consume_llm_stream_internal(stream, None, cancel_token, session_id, timeout_context)
357        .await
358}
359
360pub(crate) async fn consume_llm_stream_silent_with_context_and_partial(
361    stream: LLMStream,
362    cancel_token: &CancellationToken,
363    session_id: &str,
364    timeout_context: &StreamTimeoutContext,
365) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
366    consume::consume_llm_stream_internal_with_partial(
367        stream,
368        None,
369        cancel_token,
370        session_id,
371        timeout_context,
372    )
373    .await
374}
375
376#[cfg(test)]
377mod tests;
bamboo_engine/runtime/stream/handler.rs

bamboo_engine/runtime/stream/
handler.rs
