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
82#[derive(Debug, Clone, serde::Serialize)]
83pub(crate) struct PartialToolCallSnapshot {
84    pub id: String,
85    pub tool_type: String,
86    pub name: String,
87    pub arguments: String,
88    pub index: Option<u32>,
89}
90
91/// Crate-private interrupted-stream payload.  Keeping this separate from the
92/// public successful [`StreamHandlingOutput`] avoids a source-breaking public
93/// field while retaining fragments that finalization intentionally drops or
94/// normalizes.
95pub(crate) struct InterruptedStreamOutput {
96    pub content: String,
97    pub reasoning_content: String,
98    pub partial_tool_calls: Vec<PartialToolCallSnapshot>,
99}
100
101impl From<&bamboo_agent_core::tools::PartialToolCall> for PartialToolCallSnapshot {
102    fn from(value: &bamboo_agent_core::tools::PartialToolCall) -> Self {
103        Self {
104            id: value.id.clone(),
105            tool_type: value.tool_type.clone(),
106            name: value.name.clone(),
107            arguments: value.arguments.clone(),
108            index: value.index,
109        }
110    }
111}
112
113/// A stream failure together with every semantic fragment accumulated before
114/// the failure.  The agent round uses this to create a durable, explicitly
115/// interrupted assistant record instead of losing already-visible output.
116pub(crate) struct StreamHandlingFailure {
117    pub error: AgentError,
118    pub partial_output: InterruptedStreamOutput,
119}
120
121pub async fn consume_llm_stream(
122    stream: LLMStream,
123    event_tx: &mpsc::Sender<AgentEvent>,
124    cancel_token: &CancellationToken,
125    session_id: &str,
126) -> Result<StreamHandlingOutput, AgentError> {
127    consume_llm_stream_with_context(
128        stream,
129        event_tx,
130        cancel_token,
131        session_id,
132        &StreamTimeoutContext::default(),
133    )
134    .await
135}
136
137pub async fn consume_llm_stream_with_context(
138    stream: LLMStream,
139    event_tx: &mpsc::Sender<AgentEvent>,
140    cancel_token: &CancellationToken,
141    session_id: &str,
142    timeout_context: &StreamTimeoutContext,
143) -> Result<StreamHandlingOutput, AgentError> {
144    consume::consume_llm_stream_internal(
145        stream,
146        Some(event_tx),
147        cancel_token,
148        session_id,
149        timeout_context,
150    )
151    .await
152}
153
154pub(crate) async fn consume_llm_stream_with_context_and_partial(
155    stream: LLMStream,
156    event_tx: &mpsc::Sender<AgentEvent>,
157    cancel_token: &CancellationToken,
158    session_id: &str,
159    timeout_context: &StreamTimeoutContext,
160) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
161    consume::consume_llm_stream_internal_with_partial(
162        stream,
163        Some(event_tx),
164        cancel_token,
165        session_id,
166        timeout_context,
167    )
168    .await
169}
170
171pub async fn consume_llm_stream_silent(
172    stream: LLMStream,
173    cancel_token: &CancellationToken,
174    session_id: &str,
175) -> Result<StreamHandlingOutput, AgentError> {
176    consume_llm_stream_silent_with_context(
177        stream,
178        cancel_token,
179        session_id,
180        &StreamTimeoutContext::default(),
181    )
182    .await
183}
184
185pub async fn consume_llm_stream_silent_with_context(
186    stream: LLMStream,
187    cancel_token: &CancellationToken,
188    session_id: &str,
189    timeout_context: &StreamTimeoutContext,
190) -> Result<StreamHandlingOutput, AgentError> {
191    consume::consume_llm_stream_internal(stream, None, cancel_token, session_id, timeout_context)
192        .await
193}
194
195pub(crate) async fn consume_llm_stream_silent_with_context_and_partial(
196    stream: LLMStream,
197    cancel_token: &CancellationToken,
198    session_id: &str,
199    timeout_context: &StreamTimeoutContext,
200) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
201    consume::consume_llm_stream_internal_with_partial(
202        stream,
203        None,
204        cancel_token,
205        session_id,
206        timeout_context,
207    )
208    .await
209}
210
211#[cfg(test)]
212mod tests;
bamboo_engine/runtime/stream/handler.rs

bamboo_engine/runtime/stream/
handler.rs
