bamboo_engine/runtime/stream/
handler.rs

1use tokio::sync::mpsc;
2use tokio_util::sync::CancellationToken;
3
4use bamboo_agent_core::tools::ToolCall;
5use bamboo_agent_core::{AgentError, AgentEvent};
6use bamboo_config::StreamTimeoutConfig;
7use bamboo_llm::LLMStream;
8
9mod chunk_handling;
10mod consume;
11mod stream_state;
12
13/// Resolved identifiers and deadlines for one stream. Identifiers are sanitized
14/// before they reach timeout diagnostics; prompts and provider payloads are
15/// never retained here.
16#[derive(Debug, Clone)]
17pub struct StreamTimeoutContext {
18    pub(crate) policy: StreamTimeoutConfig,
19    pub(crate) provider: Option<String>,
20    pub(crate) model: Option<String>,
21}
22
23impl StreamTimeoutContext {
24    pub fn new(policy: StreamTimeoutConfig, provider: Option<&str>, model: Option<&str>) -> Self {
25        let policy = match policy.validate() {
26            Ok(()) => policy,
27            Err(error) => {
28                tracing::warn!(
29                    "invalid programmatic stream timeout policy ({error}); using safe defaults"
30                );
31                StreamTimeoutConfig::default()
32            }
33        };
34        Self {
35            policy,
36            provider: provider.and_then(sanitize_identifier),
37            model: model.and_then(sanitize_identifier),
38        }
39    }
40}
41
42impl Default for StreamTimeoutContext {
43    fn default() -> Self {
44        Self::new(StreamTimeoutConfig::default(), None, None)
45    }
46}
47
48fn sanitize_identifier(value: &str) -> Option<String> {
49    let value = value.trim();
50    if value.is_empty() {
51        return None;
52    }
53    let sanitized: String = value
54        .chars()
55        .take(120)
56        .map(|character| match character {
57            character if character.is_ascii_alphanumeric() => character,
58            '-' | '_' | '.' | ':' | '/' | '@' => character,
59            _ => '_',
60        })
61        .collect();
62    (!sanitized.is_empty()).then_some(sanitized)
63}
64
65pub struct StreamHandlingOutput {
66    pub response_id: Option<String>,
67    pub content: String,
68    pub reasoning_content: String,
69    /// Provider-minted signature covering `reasoning_content`, present only
70    /// when the turn's thinking arrived as exactly one signed Anthropic
71    /// `thinking` block — see [`bamboo_llm::LLMChunk::ReasoningSignature`] (#520).
72    pub reasoning_signature: Option<String>,
73    pub token_count: usize,
74    pub tool_calls: Vec<ToolCall>,
75    pub output_tokens: u64,
76    pub thinking_tokens: u64,
77    pub cache_creation_input_tokens: u64,
78    pub cache_read_input_tokens: u64,
79    pub input_tokens: u64,
80}
81
82pub async fn consume_llm_stream(
83    stream: LLMStream,
84    event_tx: &mpsc::Sender<AgentEvent>,
85    cancel_token: &CancellationToken,
86    session_id: &str,
87) -> Result<StreamHandlingOutput, AgentError> {
88    consume_llm_stream_with_context(
89        stream,
90        event_tx,
91        cancel_token,
92        session_id,
93        &StreamTimeoutContext::default(),
94    )
95    .await
96}
97
98pub async fn consume_llm_stream_with_context(
99    stream: LLMStream,
100    event_tx: &mpsc::Sender<AgentEvent>,
101    cancel_token: &CancellationToken,
102    session_id: &str,
103    timeout_context: &StreamTimeoutContext,
104) -> Result<StreamHandlingOutput, AgentError> {
105    consume::consume_llm_stream_internal(
106        stream,
107        Some(event_tx),
108        cancel_token,
109        session_id,
110        timeout_context,
111    )
112    .await
113}
114
115pub async fn consume_llm_stream_silent(
116    stream: LLMStream,
117    cancel_token: &CancellationToken,
118    session_id: &str,
119) -> Result<StreamHandlingOutput, AgentError> {
120    consume_llm_stream_silent_with_context(
121        stream,
122        cancel_token,
123        session_id,
124        &StreamTimeoutContext::default(),
125    )
126    .await
127}
128
129pub async fn consume_llm_stream_silent_with_context(
130    stream: LLMStream,
131    cancel_token: &CancellationToken,
132    session_id: &str,
133    timeout_context: &StreamTimeoutContext,
134) -> Result<StreamHandlingOutput, AgentError> {
135    consume::consume_llm_stream_internal(stream, None, cancel_token, session_id, timeout_context)
136        .await
137}
138
139#[cfg(test)]
140mod tests;
bamboo_engine/runtime/stream/handler.rs

bamboo_engine/runtime/stream/
handler.rs
