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
214/// Publish an already consumed response through the existing visible-message
215/// protocol. Final-answer checks use this after deciding which text to keep.
216pub(crate) async fn publish_buffered_response(
217    event_tx: &mpsc::Sender<AgentEvent>,
218    identity: &VisibleMessageIdentity,
219    content: &str,
220    reasoning: Option<&str>,
221) {
222    let _ = event_tx
223        .send(AgentEvent::VisibleMessageStart {
224            message_id: identity.message_id.clone(),
225            created_at: identity.created_at,
226        })
227        .await;
228    if let Some(reasoning) = reasoning.filter(|text| !text.is_empty()) {
229        let _ = event_tx
230            .send(AgentEvent::ReasoningToken {
231                content: reasoning.to_string(),
232            })
233            .await;
234    }
235    if !content.is_empty() {
236        let _ = event_tx
237            .send(AgentEvent::Token {
238                content: content.to_string(),
239            })
240            .await;
241    }
242}
243
244pub struct StreamHandlingOutput {
245    pub response_id: Option<String>,
246    /// Stable identity shared by the safe realtime text and persisted message.
247    pub visible_message: Option<VisibleMessageIdentity>,
248    pub content: String,
249    pub reasoning_content: String,
250    /// Provider-minted signature covering `reasoning_content`, present only
251    /// when the turn's thinking arrived as exactly one signed Anthropic
252    /// `thinking` block — see [`bamboo_llm::LLMChunk::ReasoningSignature`] (#520).
253    pub reasoning_signature: Option<String>,
254    pub token_count: usize,
255    pub tool_calls: Vec<ToolCall>,
256    pub output_tokens: u64,
257    pub thinking_tokens: u64,
258    pub cache_creation_input_tokens: u64,
259    pub cache_read_input_tokens: u64,
260    /// Merged authoritative provider snapshot, when at least one provider-usage
261    /// chunk was observed. Repeated cumulative snapshots are idempotent, absent
262    /// fields do not erase known values, and explicit zeros remain `Some(0)`.
263    pub provider_usage: Option<ProviderUsageSnapshot>,
264    /// Normalized non-cached ("fresh") input, disjoint from the adjacent cache
265    /// counters. When a provider total is available this is derived from that
266    /// total with a saturating, cache-subset policy; the raw total remains in
267    /// [`Self::provider_usage`].
268    pub input_tokens: u64,
269    /// Validated native items awaiting atomic commit with this round's ordinary
270    /// assistant message. Never exposed as streaming UI events.
271    pub provider_transcript_items: Vec<bamboo_domain::ProviderTranscriptItem>,
272}
273
274#[derive(Debug, Clone, serde::Serialize)]
275pub(crate) struct PartialToolCallSnapshot {
276    pub id: String,
277    pub tool_type: String,
278    pub name: String,
279    pub arguments: String,
280    pub index: Option<u32>,
281}
282
283/// Crate-private interrupted-stream payload.  Keeping this separate from the
284/// public successful [`StreamHandlingOutput`] avoids a source-breaking public
285/// field while retaining fragments that finalization intentionally drops or
286/// normalizes.
287pub(crate) struct InterruptedStreamOutput {
288    pub visible_message: Option<VisibleMessageIdentity>,
289    pub content: String,
290    pub reasoning_content: String,
291    pub partial_tool_calls: Vec<PartialToolCallSnapshot>,
292    pub provider_usage: Option<ProviderUsageSnapshot>,
293    pub input_tokens: u64,
294    pub output_tokens: u64,
295}
296
297impl From<&bamboo_agent_core::tools::PartialToolCall> for PartialToolCallSnapshot {
298    fn from(value: &bamboo_agent_core::tools::PartialToolCall) -> Self {
299        Self {
300            id: value.id.clone(),
301            tool_type: value.tool_type.clone(),
302            name: value.name.clone(),
303            arguments: value.arguments.clone(),
304            index: value.index,
305        }
306    }
307}
308
309/// A stream failure together with every semantic fragment accumulated before
310/// the failure.  The agent round uses this to create a durable, explicitly
311/// interrupted assistant record instead of losing already-visible output.
312pub(crate) struct StreamHandlingFailure {
313    pub error: AgentError,
314    /// Failures are already the cold path; boxing the three-buffer snapshot
315    /// keeps this error small without changing the preserved fragment data.
316    pub partial_output: Box<InterruptedStreamOutput>,
317}
318
319pub async fn consume_llm_stream(
320    stream: LLMStream,
321    event_tx: &mpsc::Sender<AgentEvent>,
322    cancel_token: &CancellationToken,
323    session_id: &str,
324) -> Result<StreamHandlingOutput, AgentError> {
325    consume_llm_stream_with_context(
326        stream,
327        event_tx,
328        cancel_token,
329        session_id,
330        &StreamTimeoutContext::default(),
331    )
332    .await
333}
334
335pub async fn consume_llm_stream_with_context(
336    stream: LLMStream,
337    event_tx: &mpsc::Sender<AgentEvent>,
338    cancel_token: &CancellationToken,
339    session_id: &str,
340    timeout_context: &StreamTimeoutContext,
341) -> Result<StreamHandlingOutput, AgentError> {
342    consume::consume_llm_stream_internal(
343        stream,
344        Some(event_tx),
345        cancel_token,
346        session_id,
347        timeout_context,
348    )
349    .await
350}
351
352pub(crate) async fn consume_llm_stream_with_context_and_partial(
353    stream: LLMStream,
354    event_tx: &mpsc::Sender<AgentEvent>,
355    cancel_token: &CancellationToken,
356    session_id: &str,
357    timeout_context: &StreamTimeoutContext,
358) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
359    consume::consume_llm_stream_internal_with_partial(
360        stream,
361        Some(event_tx),
362        cancel_token,
363        session_id,
364        timeout_context,
365    )
366    .await
367}
368
369pub async fn consume_llm_stream_silent(
370    stream: LLMStream,
371    cancel_token: &CancellationToken,
372    session_id: &str,
373) -> Result<StreamHandlingOutput, AgentError> {
374    consume_llm_stream_silent_with_context(
375        stream,
376        cancel_token,
377        session_id,
378        &StreamTimeoutContext::default(),
379    )
380    .await
381}
382
383pub async fn consume_llm_stream_silent_with_context(
384    stream: LLMStream,
385    cancel_token: &CancellationToken,
386    session_id: &str,
387    timeout_context: &StreamTimeoutContext,
388) -> Result<StreamHandlingOutput, AgentError> {
389    consume::consume_llm_stream_internal(stream, None, cancel_token, session_id, timeout_context)
390        .await
391}
392
393pub(crate) async fn consume_llm_stream_silent_with_context_and_partial(
394    stream: LLMStream,
395    cancel_token: &CancellationToken,
396    session_id: &str,
397    timeout_context: &StreamTimeoutContext,
398) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
399    consume::consume_llm_stream_internal_with_partial(
400        stream,
401        None,
402        cancel_token,
403        session_id,
404        timeout_context,
405    )
406    .await
407}
408
409#[cfg(test)]
410mod tests;
bamboo_engine/runtime/stream/handler.rs

bamboo_engine/runtime/stream/
handler.rs
