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
65/// Raw authoritative usage fields preserved from provider terminal events.
66///
67/// Each field remains optional so callers can distinguish an explicit
68/// provider-reported zero from an omitted value. Flat counters on
69/// [`StreamHandlingOutput`] remain the normalized compatibility view.
70#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
71pub struct ProviderUsageSnapshot {
72    pub input_tokens: Option<u64>,
73    pub output_tokens: Option<u64>,
74    pub total_tokens: Option<u64>,
75    pub reasoning_tokens: Option<u64>,
76    pub cache_creation_input_tokens: Option<u64>,
77    pub cache_read_input_tokens: Option<u64>,
78    /// OpenAI Responses cache-write volume. This is raw provider metadata and
79    /// is not folded into the disjoint legacy prompt-cache counters.
80    pub cache_write_input_tokens: Option<u64>,
81}
82
83pub struct StreamHandlingOutput {
84    pub response_id: Option<String>,
85    pub content: String,
86    pub reasoning_content: String,
87    /// Provider-minted signature covering `reasoning_content`, present only
88    /// when the turn's thinking arrived as exactly one signed Anthropic
89    /// `thinking` block — see [`bamboo_llm::LLMChunk::ReasoningSignature`] (#520).
90    pub reasoning_signature: Option<String>,
91    pub token_count: usize,
92    pub tool_calls: Vec<ToolCall>,
93    pub output_tokens: u64,
94    pub thinking_tokens: u64,
95    pub cache_creation_input_tokens: u64,
96    pub cache_read_input_tokens: u64,
97    /// Merged authoritative provider snapshot, when at least one provider-usage
98    /// chunk was observed. Repeated cumulative snapshots are idempotent, absent
99    /// fields do not erase known values, and explicit zeros remain `Some(0)`.
100    pub provider_usage: Option<ProviderUsageSnapshot>,
101    /// Normalized non-cached ("fresh") input, disjoint from the adjacent cache
102    /// counters. When a provider total is available this is derived from that
103    /// total with a saturating, cache-subset policy; the raw total remains in
104    /// [`Self::provider_usage`].
105    pub input_tokens: u64,
106}
107
108impl StreamHandlingOutput {
109    /// Prompt usage for runtime budget enforcement.
110    ///
111    /// Provider-reported totals are authoritative when present. Legacy streams
112    /// retain their historical fallback to the flat fresh-input counter.
113    pub(crate) fn prompt_tokens_for_runtime_budget(&self) -> u64 {
114        self.provider_usage
115            .and_then(|usage| usage.input_tokens)
116            .unwrap_or(self.input_tokens)
117    }
118
119    /// Completion usage for runtime budget enforcement.
120    ///
121    /// A provider-reported output total (including explicit zero) wins over
122    /// legacy summaries. Reasoning is a subset breakdown and is never added.
123    pub(crate) fn completion_tokens_for_runtime_budget(&self) -> u64 {
124        self.provider_usage
125            .and_then(|usage| usage.output_tokens)
126            .unwrap_or(self.output_tokens)
127    }
128}
129
130#[derive(Debug, Clone, serde::Serialize)]
131pub(crate) struct PartialToolCallSnapshot {
132    pub id: String,
133    pub tool_type: String,
134    pub name: String,
135    pub arguments: String,
136    pub index: Option<u32>,
137}
138
139/// Crate-private interrupted-stream payload.  Keeping this separate from the
140/// public successful [`StreamHandlingOutput`] avoids a source-breaking public
141/// field while retaining fragments that finalization intentionally drops or
142/// normalizes.
143pub(crate) struct InterruptedStreamOutput {
144    pub content: String,
145    pub reasoning_content: String,
146    pub partial_tool_calls: Vec<PartialToolCallSnapshot>,
147}
148
149impl From<&bamboo_agent_core::tools::PartialToolCall> for PartialToolCallSnapshot {
150    fn from(value: &bamboo_agent_core::tools::PartialToolCall) -> Self {
151        Self {
152            id: value.id.clone(),
153            tool_type: value.tool_type.clone(),
154            name: value.name.clone(),
155            arguments: value.arguments.clone(),
156            index: value.index,
157        }
158    }
159}
160
161/// A stream failure together with every semantic fragment accumulated before
162/// the failure.  The agent round uses this to create a durable, explicitly
163/// interrupted assistant record instead of losing already-visible output.
164pub(crate) struct StreamHandlingFailure {
165    pub error: AgentError,
166    pub partial_output: InterruptedStreamOutput,
167}
168
169pub async fn consume_llm_stream(
170    stream: LLMStream,
171    event_tx: &mpsc::Sender<AgentEvent>,
172    cancel_token: &CancellationToken,
173    session_id: &str,
174) -> Result<StreamHandlingOutput, AgentError> {
175    consume_llm_stream_with_context(
176        stream,
177        event_tx,
178        cancel_token,
179        session_id,
180        &StreamTimeoutContext::default(),
181    )
182    .await
183}
184
185pub async fn consume_llm_stream_with_context(
186    stream: LLMStream,
187    event_tx: &mpsc::Sender<AgentEvent>,
188    cancel_token: &CancellationToken,
189    session_id: &str,
190    timeout_context: &StreamTimeoutContext,
191) -> Result<StreamHandlingOutput, AgentError> {
192    consume::consume_llm_stream_internal(
193        stream,
194        Some(event_tx),
195        cancel_token,
196        session_id,
197        timeout_context,
198    )
199    .await
200}
201
202pub(crate) async fn consume_llm_stream_with_context_and_partial(
203    stream: LLMStream,
204    event_tx: &mpsc::Sender<AgentEvent>,
205    cancel_token: &CancellationToken,
206    session_id: &str,
207    timeout_context: &StreamTimeoutContext,
208) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
209    consume::consume_llm_stream_internal_with_partial(
210        stream,
211        Some(event_tx),
212        cancel_token,
213        session_id,
214        timeout_context,
215    )
216    .await
217}
218
219pub async fn consume_llm_stream_silent(
220    stream: LLMStream,
221    cancel_token: &CancellationToken,
222    session_id: &str,
223) -> Result<StreamHandlingOutput, AgentError> {
224    consume_llm_stream_silent_with_context(
225        stream,
226        cancel_token,
227        session_id,
228        &StreamTimeoutContext::default(),
229    )
230    .await
231}
232
233pub async fn consume_llm_stream_silent_with_context(
234    stream: LLMStream,
235    cancel_token: &CancellationToken,
236    session_id: &str,
237    timeout_context: &StreamTimeoutContext,
238) -> Result<StreamHandlingOutput, AgentError> {
239    consume::consume_llm_stream_internal(stream, None, cancel_token, session_id, timeout_context)
240        .await
241}
242
243pub(crate) async fn consume_llm_stream_silent_with_context_and_partial(
244    stream: LLMStream,
245    cancel_token: &CancellationToken,
246    session_id: &str,
247    timeout_context: &StreamTimeoutContext,
248) -> Result<StreamHandlingOutput, StreamHandlingFailure> {
249    consume::consume_llm_stream_internal_with_partial(
250        stream,
251        None,
252        cancel_token,
253        session_id,
254        timeout_context,
255    )
256    .await
257}
258
259#[cfg(test)]
260mod tests;
bamboo_engine/runtime/stream/handler.rs

bamboo_engine/runtime/stream/
handler.rs
