bamboo_engine/runtime/stream/
handler.rs

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

bamboo_engine/runtime/stream/
handler.rs
