bamboo-engine 2026.7.25

use std::time::Duration;

use futures::{stream, StreamExt};
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TryRecvError;
use tokio_util::sync::CancellationToken;

use bamboo_agent_core::tools::{FunctionCall, ToolCall};
use bamboo_agent_core::{AgentError, AgentEvent};
use bamboo_config::StreamTimeoutConfig;
use bamboo_llm::provider::LLMError;
use bamboo_llm::{LLMChunk, LLMStream};

use super::consume::consume_llm_stream_internal;
use super::{consume_llm_stream, consume_llm_stream_silent, StreamTimeoutContext};

fn build_stream(items: Vec<bamboo_llm::provider::Result<LLMChunk>>) -> LLMStream {
    Box::pin(stream::iter(items))
}

fn timeout_context(
    transport_secs: u64,
    first_semantic_secs: u64,
    semantic_secs: u64,
) -> StreamTimeoutContext {
    StreamTimeoutContext::new(
        StreamTimeoutConfig {
            transport_idle_timeout_secs: transport_secs,
            first_semantic_timeout_secs: first_semantic_secs,
            semantic_idle_timeout_secs: semantic_secs,
        },
        Some("test-provider"),
        Some("test-model"),
    )
}

#[tokio::test]
async fn consume_llm_stream_accumulates_tokens_and_tool_calls() {
    let stream = build_stream(vec![
        Ok(LLMChunk::ResponseId("resp_123".to_string())),
        Ok(LLMChunk::ReasoningToken("thinking".to_string())),
        Ok(LLMChunk::Token("hi".to_string())),
        Ok(LLMChunk::ToolCalls(vec![ToolCall {
            id: "call_1".to_string(),
            tool_type: "function".to_string(),
            function: FunctionCall {
                name: "test_tool".to_string(),
                arguments: "{".to_string(),
            },
        }])),
        Ok(LLMChunk::ToolCalls(vec![ToolCall {
            id: "call_1".to_string(),
            tool_type: "function".to_string(),
            function: FunctionCall {
                name: String::new(),
                arguments: "}".to_string(),
            },
        }])),
        Ok(LLMChunk::Done),
    ]);

    let (event_tx, mut event_rx) = mpsc::channel::<AgentEvent>(8);
    let output = consume_llm_stream(stream, &event_tx, &CancellationToken::new(), "session-1")
        .await
        .expect("stream should succeed");

    assert_eq!(output.response_id.as_deref(), Some("resp_123"));
    assert_eq!(output.content, "hi");
    assert_eq!(output.reasoning_content, "thinking");
    assert_eq!(output.token_count, 2);
    assert_eq!(output.tool_calls.len(), 1);
    assert_eq!(output.tool_calls[0].function.name, "test_tool");
    assert_eq!(output.tool_calls[0].function.arguments, "{}");

    let reasoning_event = event_rx.recv().await.expect("missing reasoning event");
    assert!(matches!(reasoning_event, AgentEvent::ReasoningToken { .. }));

    let token_event = event_rx.recv().await.expect("missing token event");
    assert!(matches!(token_event, AgentEvent::Token { .. }));
}

/// #520: a provider-minted reasoning signature is surfaced on the output so
/// the persisted assistant message can replay a SIGNED thinking block.
#[tokio::test]
async fn consume_llm_stream_captures_reasoning_signature() {
    let stream = build_stream(vec![
        Ok(LLMChunk::ReasoningToken("thinking".to_string())),
        Ok(LLMChunk::ReasoningSignature("sig_abc".to_string())),
        Ok(LLMChunk::Token("hi".to_string())),
        Ok(LLMChunk::Done),
    ]);

    let (event_tx, _event_rx) = mpsc::channel::<AgentEvent>(8);
    let output = consume_llm_stream(stream, &event_tx, &CancellationToken::new(), "session-sig")
        .await
        .expect("stream should succeed");

    assert_eq!(output.reasoning_content, "thinking");
    assert_eq!(output.reasoning_signature.as_deref(), Some("sig_abc"));
}

/// #520: the empty-string marker permanently invalidates the signature for
/// the stream (multi-block/redacted turns), even if another one follows.
#[tokio::test]
async fn consume_llm_stream_honors_signature_invalidation_marker() {
    let stream = build_stream(vec![
        Ok(LLMChunk::ReasoningSignature("sig_first".to_string())),
        Ok(LLMChunk::ReasoningSignature(String::new())),
        Ok(LLMChunk::ReasoningSignature("sig_late".to_string())),
        Ok(LLMChunk::Done),
    ]);

    let (event_tx, _event_rx) = mpsc::channel::<AgentEvent>(8);
    let output = consume_llm_stream(
        stream,
        &event_tx,
        &CancellationToken::new(),
        "session-sig-invalid",
    )
    .await
    .expect("stream should succeed");

    assert_eq!(
        output.reasoning_signature, None,
        "invalidation is permanent for the stream"
    );
}

#[tokio::test]
async fn consume_llm_stream_silent_does_not_emit_events() {
    let stream = build_stream(vec![
        Ok(LLMChunk::Token("hello".to_string())),
        Ok(LLMChunk::Done),
    ]);

    let output = consume_llm_stream_silent(stream, &CancellationToken::new(), "session-2")
        .await
        .expect("silent stream should succeed");

    assert!(output.response_id.is_none());
    assert_eq!(output.content, "hello");
    assert!(output.reasoning_content.is_empty());
    assert_eq!(output.token_count, 5);
    assert!(output.tool_calls.is_empty());
}

#[tokio::test]
async fn consume_llm_stream_returns_single_prefix_stream_error_message() {
    let stream = build_stream(vec![Err(LLMError::Stream(
        "Transport error: error decoding response body".to_string(),
    ))]);

    let (event_tx, mut event_rx) = mpsc::channel::<AgentEvent>(4);
    let err =
        match consume_llm_stream(stream, &event_tx, &CancellationToken::new(), "session-3").await {
            Ok(_) => panic!("stream should fail"),
            Err(err) => err,
        };

    match err {
        AgentError::LLM(message) => {
            assert_eq!(
                message,
                "Stream error: Transport error: error decoding response body"
            );
            assert!(!message.starts_with("Stream error: Stream error:"));
        }
        other => panic!("expected AgentError::LLM, got {other:?}"),
    }

    assert!(matches!(event_rx.try_recv(), Err(TryRecvError::Empty)));
}

#[tokio::test]
async fn consume_llm_stream_aborts_already_cancelled_stalled_stream() {
    // A provider stream that never yields and never ends. Before the `select!`
    // fix, `.next().await` would block forever; now a cancelled token must
    // return promptly instead of hanging.
    let stream: LLMStream = Box::pin(stream::pending());
    let cancel = CancellationToken::new();
    cancel.cancel();

    let result = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        consume_llm_stream_silent(stream, &cancel, "session-cancelled"),
    )
    .await
    .expect("must not hang: cancellation should interrupt the stalled stream");

    assert!(matches!(result, Err(AgentError::Cancelled)));
}

#[tokio::test]
async fn consume_llm_stream_interrupts_blocked_next_on_mid_stream_cancel() {
    // Proves a *blocked* `stream.next().await` (not just between chunks) is
    // interrupted when the token is cancelled while the consume call is running.
    let stream: LLMStream = Box::pin(stream::pending());
    let cancel = CancellationToken::new();
    let canceller = cancel.clone();
    tokio::spawn(async move {
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        canceller.cancel();
    });

    let result = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        consume_llm_stream_silent(stream, &cancel, "session-cancel-mid"),
    )
    .await
    .expect("must not hang: mid-stream cancellation should interrupt the blocked next()");

    assert!(matches!(result, Err(AgentError::Cancelled)));
}

#[tokio::test(start_paused = true)]
async fn truly_silent_transport_times_out_with_actionable_diagnostic() {
    let stream: LLMStream = Box::pin(stream::pending());
    let context = timeout_context(2, 20, 20);

    let result = consume_llm_stream_internal(
        stream,
        None,
        &CancellationToken::new(),
        "session-transport-timeout",
        &context,
    )
    .await;

    let message = match result {
        Err(AgentError::StreamTimeout(message)) => message,
        Err(other) => panic!("expected transport StreamTimeout, got {other:?}"),
        Ok(_) => panic!("expected transport StreamTimeout, got success"),
    };
    assert!(message.contains("phase=transport_idle"));
    assert!(message.contains("deadline_ms=2000"));
    assert!(message.contains("provider=test-provider"));
    assert!(message.contains("model=test-model"));
    assert!(message.contains("last_transport_ms_ago=2000"));
    assert!(message.contains("last_semantic_ms_ago=never"));
    assert!(message.contains("semantic_output_started=false"));
    assert!(message.contains("retry_safe=true"));
    assert!(!message.contains("prompt"));
}

#[tokio::test(start_paused = true)]
async fn stream_stall_after_semantic_output_is_not_retry_safe() {
    let stream: LLMStream = Box::pin(
        stream::once(async { Ok::<_, LLMError>(LLMChunk::Token("first".to_string())) })
            .chain(stream::pending()),
    );
    let context = timeout_context(2, 20, 20);

    let result = consume_llm_stream_internal(
        stream,
        None,
        &CancellationToken::new(),
        "session-timeout",
        &context,
    )
    .await;

    let message = match result {
        Err(AgentError::StreamTimeout(message)) => message,
        Err(other) => panic!("expected transport StreamTimeout, got {other:?}"),
        Ok(_) => panic!("expected transport StreamTimeout, got success"),
    };
    assert!(message.contains("phase=transport_idle"));
    assert!(message.contains("semantic_output_started=true"));
    assert!(message.contains("retry_safe=false"));
}

#[tokio::test(start_paused = true)]
async fn transport_keepalives_allow_first_semantic_output_after_120_seconds() {
    let stream: LLMStream = Box::pin(stream::unfold(0u8, |step| async move {
        match step {
            0..=4 => {
                tokio::time::sleep(Duration::from_secs(30)).await;
                Some((Ok::<_, LLMError>(LLMChunk::TransportActivity), step + 1))
            }
            5 => {
                tokio::time::sleep(Duration::from_secs(30)).await;
                Some((Ok::<_, LLMError>(LLMChunk::Token("late".to_string())), 6))
            }
            6 => Some((Ok::<_, LLMError>(LLMChunk::Done), 7)),
            _ => None,
        }
    }));
    let context = timeout_context(60, 240, 60);

    let output = consume_llm_stream_internal(
        stream,
        None,
        &CancellationToken::new(),
        "session-keepalive",
        &context,
    )
    .await
    .expect("stream should succeed");

    assert_eq!(output.content, "late");
}

#[tokio::test(start_paused = true)]
async fn transport_keepalives_allow_midstream_semantic_gap_after_120_seconds() {
    let stream: LLMStream = Box::pin(
        stream::once(async { Ok::<_, LLMError>(LLMChunk::Token("first".to_string())) }).chain(
            stream::unfold(0u8, |step| async move {
                match step {
                    0..=4 => {
                        tokio::time::sleep(Duration::from_secs(30)).await;
                        Some((Ok::<_, LLMError>(LLMChunk::TransportActivity), step + 1))
                    }
                    5 => {
                        tokio::time::sleep(Duration::from_secs(30)).await;
                        Some((Ok::<_, LLMError>(LLMChunk::Token("second".to_string())), 6))
                    }
                    6 => Some((Ok::<_, LLMError>(LLMChunk::Done), 7)),
                    _ => None,
                }
            }),
        ),
    );
    let context = timeout_context(60, 240, 240);

    let output = consume_llm_stream_internal(
        stream,
        None,
        &CancellationToken::new(),
        "session-midstream-keepalive",
        &context,
    )
    .await
    .expect("live stream should survive a 180-second semantic gap");

    assert_eq!(output.content, "firstsecond");
}

#[tokio::test(start_paused = true)]
async fn keepalives_do_not_make_first_semantic_deadline_unbounded() {
    let stream: LLMStream = Box::pin(stream::unfold((), |_| async {
        tokio::time::sleep(Duration::from_secs(20)).await;
        Some((Ok::<_, LLMError>(LLMChunk::TransportActivity), ()))
    }));
    let context = timeout_context(60, 120, 120);

    let result = consume_llm_stream_internal(
        stream,
        None,
        &CancellationToken::new(),
        "session-first-semantic",
        &context,
    )
    .await;

    let message = match result {
        Err(AgentError::StreamTimeout(message)) => message,
        Err(other) => panic!("expected first-semantic StreamTimeout, got {other:?}"),
        Ok(_) => panic!("expected first-semantic StreamTimeout, got success"),
    };
    assert!(message.contains("phase=first_semantic"));
    assert!(message.contains("semantic_output_started=false"));
}

#[tokio::test(start_paused = true)]
async fn keepalives_do_not_hide_midstream_semantic_stall() {
    let stream: LLMStream = Box::pin(
        stream::once(async { Ok::<_, LLMError>(LLMChunk::Token("first".to_string())) }).chain(
            stream::unfold((), |_| async {
                tokio::time::sleep(Duration::from_secs(20)).await;
                Some((Ok::<_, LLMError>(LLMChunk::TransportActivity), ()))
            }),
        ),
    );
    let context = timeout_context(60, 120, 90);

    let result = consume_llm_stream_internal(
        stream,
        None,
        &CancellationToken::new(),
        "session-semantic-stall",
        &context,
    )
    .await;

    let message = match result {
        Err(AgentError::StreamTimeout(message)) => message,
        Err(other) => panic!("expected semantic-idle StreamTimeout, got {other:?}"),
        Ok(_) => panic!("expected semantic-idle StreamTimeout, got success"),
    };
    assert!(message.contains("phase=semantic_idle"));
    assert!(message.contains("semantic_output_started=true"));
    assert!(message.contains("retry_safe=false"));
}

#[tokio::test]
async fn consume_llm_stream_continues_when_subscriber_disconnects() {
    // Issue #23: when the subscriber drops its receiver, every token send
    // fails. Previously this was `let _ = event_tx.send(...).await;`, so the
    // failure was invisible (no log). The send must not panic and — because
    // accumulation happens *before* the forward — the stream must still
    // complete and return its full content. The failure path now emits a warn
    // instead of silently swallowing; the await (backpressure) semantics are
    // preserved (no try_send / timeout drop is introduced).
    let stream = build_stream(vec![
        Ok(LLMChunk::ReasoningToken("think".to_string())),
        Ok(LLMChunk::Token("hello".to_string())),
        Ok(LLMChunk::Token("world".to_string())),
        Ok(LLMChunk::Done),
    ]);

    // Create the channel and immediately drop the receiver, modelling a
    // subscriber that disconnected before streaming began. A capacity of 1 is
    // deliberate: it forces every send straight onto the disconnected path.
    let (event_tx, event_rx) = mpsc::channel::<AgentEvent>(1);
    drop(event_rx);

    // Must not panic, hang, or error despite every event send failing.
    let output = tokio::time::timeout(
        Duration::from_secs(5),
        consume_llm_stream(
            stream,
            &event_tx,
            &CancellationToken::new(),
            "session-disconnect",
        ),
    )
    .await
    .expect("stream must complete when the subscriber is disconnected, not hang")
    .expect("stream should succeed even though event sends fail");

    // Content is accumulated into `state` before the forward attempt, so it is
    // fully preserved regardless of the event channel state. This is the key
    // behavioral guarantee: a dropped subscriber must never corrupt the run.
    assert_eq!(output.reasoning_content, "think");
    assert_eq!(output.content, "helloworld");
    assert_eq!(output.token_count, 10);
}