Skip to main content

bamboo_engine/runtime/stream/
handler.rs

1use std::time::Duration;
2
3use tokio::sync::mpsc;
4use tokio_util::sync::CancellationToken;
5
6use bamboo_agent_core::tools::ToolCall;
7use bamboo_agent_core::{AgentError, AgentEvent};
8use bamboo_llm::LLMStream;
9
10mod chunk_handling;
11mod consume;
12mod stream_state;
13
14/// Default *inter-chunk* idle timeout for LLM streams (issue #28).
15///
16/// A provider connection that stops sending chunks for longer than this is
17/// treated as a hung stream. Because the deadline resets on every received
18/// chunk (see [`consume::consume_llm_stream_internal`]), a legitimately long
19/// stream that keeps producing data — e.g. a thinking model that streams for
20/// minutes — is never affected; only a true stall (no chunk at all) trips it.
21///
22/// This is a well-named const default rather than a threaded config field: the
23/// consume function is shared across the main stream path and several
24/// auxiliary (silent) callers, and wiring a new field through `AgentLoopConfig`
25/// plus every call site is invasive for a streaming hot path. The internal
26/// function still accepts an `idle_timeout` override so it remains configurable
27/// and testable; full config wiring is deferred.
28const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(120);
29
30pub struct StreamHandlingOutput {
31    pub response_id: Option<String>,
32    pub content: String,
33    pub reasoning_content: String,
34    /// Provider-minted signature covering `reasoning_content`, present only
35    /// when the turn's thinking arrived as exactly one signed Anthropic
36    /// `thinking` block — see [`bamboo_llm::LLMChunk::ReasoningSignature`] (#520).
37    pub reasoning_signature: Option<String>,
38    pub token_count: usize,
39    pub tool_calls: Vec<ToolCall>,
40    pub output_tokens: u64,
41    pub thinking_tokens: u64,
42    pub cache_creation_input_tokens: u64,
43    pub cache_read_input_tokens: u64,
44    pub input_tokens: u64,
45}
46
47pub async fn consume_llm_stream(
48    stream: LLMStream,
49    event_tx: &mpsc::Sender<AgentEvent>,
50    cancel_token: &CancellationToken,
51    session_id: &str,
52) -> Result<StreamHandlingOutput, AgentError> {
53    consume::consume_llm_stream_internal(
54        stream,
55        Some(event_tx),
56        cancel_token,
57        session_id,
58        DEFAULT_STREAM_IDLE_TIMEOUT,
59    )
60    .await
61}
62
63pub async fn consume_llm_stream_silent(
64    stream: LLMStream,
65    cancel_token: &CancellationToken,
66    session_id: &str,
67) -> Result<StreamHandlingOutput, AgentError> {
68    consume::consume_llm_stream_internal(
69        stream,
70        None,
71        cancel_token,
72        session_id,
73        DEFAULT_STREAM_IDLE_TIMEOUT,
74    )
75    .await
76}
77
78#[cfg(test)]
79mod tests;