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    /// Validated native items awaiting atomic commit with this round's ordinary
190    /// assistant message. Never exposed as streaming UI events.
191    pub provider_transcript_items: Vec<bamboo_domain::ProviderTranscriptItem>,
192}
193
194#[derive(Debug, Clone, serde::Serialize)]
195pub(crate) struct PartialToolCallSnapshot {
196    pub id: String,
197    pub tool_type: String,
198    pub name: String,
199    pub arguments: String,
200    pub index: Option<u32>,
201}
202
203/// Crate-private interrupted-stream payload.  Keeping this separate from the
204/// public successful [`StreamHandlingOutput`] avoids a source-breaking public
205/// field while retaining fragments that finalization intentionally drops or
206/// normalizes.
207pub(crate) struct InterruptedStreamOutput {
208    pub content: String,
209    pub reasoning_content: String,
210    pub partial_tool_calls: Vec<PartialToolCallSnapshot>,
211}
212
213impl From<&bamboo_agent_core::tools::PartialToolCall> for PartialToolCallSnapshot {
214    fn from(value: &bamboo_agent_core::tools::PartialToolCall) -> Self {
215        Self {
216            id: value.id.clone(),
217            tool_type: value.tool_type.clone(),
218            name: value.name.clone(),
219            arguments: value.arguments.clone(),
220            index: value.index,
221        }
222    }
223}
224
225/// A stream failure together with every semantic fragment accumulated before
226/// the failure.  The agent round uses this to create a durable, explicitly
227/// interrupted assistant record instead of losing already-visible output.
228pub(crate) struct StreamHandlingFailure {
229    pub error: AgentError,
230    /// Failures are already the cold path; boxing the three-buffer snapshot
231    /// keeps this error small without changing the preserved fragment data.
232    pub partial_output: Box<InterruptedStreamOutput>,
233}
234
235pub async fn consume_llm_stream(
236    stream: LLMStream,
237    event_tx: &mpsc::Sender<AgentEvent>,
238    cancel_token: &CancellationToken,
239    session_id: &str,
240) -> Result<StreamHandlingOutput, AgentError> {
241    consume_llm_stream_with_context(
242        stream,
243        event_tx,
244        cancel_token,
245        session_id,
246        &StreamTimeoutContext::default(),
247    )
248    .await
249}
250
251pub async fn consume_llm_stream_with_context(
252    stream: LLMStream,
253    event_tx: &mpsc::Sender<AgentEvent>,
254    cancel_token: &CancellationToken,
255    session_id: &str,
256    timeout_context: &StreamTimeoutContext,
257) -> Result<StreamHandlingOutput, AgentError> {
258    consume::consume_llm_stream_internal(
259        stream,
260        Some(event_tx),
261        cancel_token,
262        session_id,
263        timeout_context,
264    )
265    .await
266}
267
268pub(crate) async fn consume_llm_stream_with_context_and_partial(
269    stream: LLMStream,
270    event_tx: &mpsc::Sender<AgentEvent>,
271    cancel_token: &CancellationToken,
272    session_id: &str,
273    timeout_context: &StreamTimeoutContext,
274) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
275    consume::consume_llm_stream_internal_with_partial(
276        stream,
277        Some(event_tx),
278        cancel_token,
279        session_id,
280        timeout_context,
281    )
282    .await
283}
284
285pub async fn consume_llm_stream_silent(
286    stream: LLMStream,
287    cancel_token: &CancellationToken,
288    session_id: &str,
289) -> Result<StreamHandlingOutput, AgentError> {
290    consume_llm_stream_silent_with_context(
291        stream,
292        cancel_token,
293        session_id,
294        &StreamTimeoutContext::default(),
295    )
296    .await
297}
298
299pub async fn consume_llm_stream_silent_with_context(
300    stream: LLMStream,
301    cancel_token: &CancellationToken,
302    session_id: &str,
303    timeout_context: &StreamTimeoutContext,
304) -> Result<StreamHandlingOutput, AgentError> {
305    consume::consume_llm_stream_internal(stream, None, cancel_token, session_id, timeout_context)
306        .await
307}
308
309pub(crate) async fn consume_llm_stream_silent_with_context_and_partial(
310    stream: LLMStream,
311    cancel_token: &CancellationToken,
312    session_id: &str,
313    timeout_context: &StreamTimeoutContext,
314) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
315    consume::consume_llm_stream_internal_with_partial(
316        stream,
317        None,
318        cancel_token,
319        session_id,
320        timeout_context,
321    )
322    .await
323}
324
325#[cfg(test)]
326mod tests;
bamboo_engine/runtime/stream/handler.rs

bamboo_engine/runtime/stream/
handler.rs
