tkmind_go/crates/goose/tests/tetrate_streaming.rs

use anyhow::Result;
use futures::StreamExt;
use goose::conversation::message::{Message, MessageContent};
use goose::model::ModelConfig;
use goose::providers::base::Provider;
use goose::providers::tetrate::TetrateProvider;
use rmcp::model::Tool;
use rmcp::object;
use serial_test::serial;

/// Test module for Tetrate Agent Router Service streaming functionality
#[cfg(test)]
mod tetrate_streaming_tests {
    use super::*;

    fn create_test_provider() -> Result<TetrateProvider> {
        // Create a test provider with the default model
        let model_config = ModelConfig::new("claude-3-5-sonnet-latest")?;
        TetrateProvider::from_env(model_config)
    }

    #[tokio::test]
    #[serial]
    #[ignore] // Ignore by default, run with --ignored flag when API key is available
    async fn test_tetrate_streaming_basic() -> Result<()> {
        let provider = create_test_provider()?;

        let messages = vec![Message::user().with_text("Count from 1 to 5, one number at a time.")];

        let mut stream = provider
            .stream(
                "You are a helpful assistant that counts numbers.",
                &messages,
                &[],
            )
            .await?;

        let mut chunk_count = 0;
        let mut content_chunks = Vec::new();

        while let Some(result) = stream.next().await {
            let (message, usage) = result?;
            chunk_count += 1;

            if let Some(msg) = message {
                let text = msg.as_concat_text();
                if !text.is_empty() {
                    content_chunks.push(text);
                }
            }

            // Check if we have usage information in the final chunk
            if usage.is_some() {
                println!("Received usage information in chunk {}", chunk_count);
            }
        }

        assert!(chunk_count > 0, "Should receive at least one chunk");
        assert!(!content_chunks.is_empty(), "Should receive some content");

        let full_content = content_chunks.join("");
        println!("Full streamed content: {}", full_content);

        // Verify the response contains numbers
        assert!(
            full_content.contains('1'),
            "Response should contain number 1"
        );
        assert!(
            full_content.contains('5'),
            "Response should contain number 5"
        );

        Ok(())
    }

    #[tokio::test]
    #[serial]
    #[ignore]
    async fn test_tetrate_streaming_with_tools() -> Result<()> {
        let provider = create_test_provider()?;

        // Define a simple tool
        let weather_tool = Tool::new(
            "get_weather",
            "Get the current weather for a location",
            object!({
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA"
                    }
                },
                "required": ["location"]
            }),
        );

        let messages = vec![Message::user().with_text("What's the weather in San Francisco?")];

        let mut stream = provider
            .stream(
                "You are a helpful assistant with access to weather information.",
                &messages,
                &[weather_tool],
            )
            .await?;

        let mut received_tool_call = false;
        let mut chunk_count = 0;

        while let Some(result) = stream.next().await {
            let (message, _usage) = result?;
            chunk_count += 1;

            if let Some(msg) = message {
                // Check if message contains tool requests
                for content in &msg.content {
                    if matches!(content, MessageContent::ToolRequest(_)) {
                        received_tool_call = true;
                        println!("Received tool call in chunk {}", chunk_count);
                    }
                }
            }
        }

        assert!(chunk_count > 0, "Should receive at least one chunk");
        // Note: Tool calls might not be supported in streaming for all models
        // This is more of a capability test than a requirement
        if received_tool_call {
            println!("✓ Streaming with tools is supported");
        } else {
            println!("⚠ Streaming with tools may not be fully supported");
        }

        Ok(())
    }

    #[tokio::test]
    #[serial]
    #[ignore]
    async fn test_tetrate_streaming_empty_response() -> Result<()> {
        let provider = create_test_provider()?;

        // This might result in a very short or empty response
        let messages = vec![Message::user().with_text("")];

        let mut stream = provider
            .stream("You are a helpful assistant.", &messages, &[])
            .await?;

        let mut chunk_count = 0;

        while let Some(result) = stream.next().await {
            let (_message, _usage) = result?;
            chunk_count += 1;
        }

        // Even with empty input, we should get at least one chunk (possibly with finish_reason)
        assert!(
            chunk_count > 0,
            "Should receive at least one chunk even with empty input"
        );

        Ok(())
    }

    #[tokio::test]
    #[serial]
    #[ignore]
    async fn test_tetrate_streaming_long_response() -> Result<()> {
        let provider = create_test_provider()?;

        let messages = vec![Message::user().with_text(
            "Write a detailed 3-paragraph essay about the importance of streaming in modern APIs.",
        )];

        let mut stream = provider
            .stream(
                "You are a helpful assistant that writes detailed essays.",
                &messages,
                &[],
            )
            .await?;

        let mut chunk_count = 0;
        let mut total_content_length = 0;

        while let Some(result) = stream.next().await {
            let (message, usage) = result?;
            chunk_count += 1;

            if let Some(msg) = message {
                let text = msg.as_concat_text();
                total_content_length += text.len();
            }

            // Final chunk should have usage information
            if let Some(usage_info) = usage {
                println!("Final usage: {:?}", usage_info.usage);
                assert!(
                    usage_info.usage.output_tokens.unwrap_or(0) > 0,
                    "Should have output tokens"
                );
            }
        }

        println!(
            "Received {} chunks with total content length: {}",
            chunk_count, total_content_length
        );

        // For a detailed essay, we expect multiple chunks and substantial content
        assert!(
            chunk_count > 5,
            "Long response should be streamed in multiple chunks"
        );
        assert!(
            total_content_length > 100,
            "Essay should have substantial content"
        );

        Ok(())
    }

    #[tokio::test]
    #[serial]
    async fn test_tetrate_streaming_error_handling() -> Result<()> {
        // Test with invalid API key to ensure error handling works
        std::env::set_var("TETRATE_API_KEY", "invalid-key-for-testing");

        let model_config = ModelConfig::new("claude-3-5-sonnet-latest")?;
        let provider = TetrateProvider::from_env(model_config)?;

        let messages = vec![Message::user().with_text("Hello")];

        let result = provider
            .stream("You are a helpful assistant.", &messages, &[])
            .await;

        // We expect this to fail with an authentication error
        assert!(result.is_err(), "Should fail with invalid API key");

        // Clean up
        std::env::remove_var("TETRATE_API_KEY");

        Ok(())
    }

    #[tokio::test]
    #[serial]
    #[ignore]
    async fn test_tetrate_streaming_concurrent_streams() -> Result<()> {
        let provider = create_test_provider()?;

        // Create multiple concurrent streams
        let messages1 = vec![Message::user().with_text("Say 'Stream 1'")];
        let messages2 = vec![Message::user().with_text("Say 'Stream 2'")];

        let stream1 = provider
            .stream("You are a helpful assistant.", &messages1, &[])
            .await?;

        let stream2 = provider
            .stream("You are a helpful assistant.", &messages2, &[])
            .await?;

        // Process both streams concurrently
        let (result1, result2) = tokio::join!(
            process_stream(stream1, "Stream 1"),
            process_stream(stream2, "Stream 2")
        );

        let content1 = result1?;
        let content2 = result2?;

        println!("Stream 1 content: {}", content1);
        println!("Stream 2 content: {}", content2);

        assert!(
            content1.contains("Stream 1") || content1.contains("1"),
            "First stream should mention Stream 1"
        );
        assert!(
            content2.contains("Stream 2") || content2.contains("2"),
            "Second stream should mention Stream 2"
        );

        Ok(())
    }

    // Helper function to process a stream and collect content
    async fn process_stream(
        mut stream: goose::providers::base::MessageStream,
        label: &str,
    ) -> Result<String> {
        let mut content = String::new();
        let mut chunk_count = 0;

        while let Some(result) = stream.next().await {
            let (message, _usage) = result?;
            chunk_count += 1;

            if let Some(msg) = message {
                let text = msg.as_concat_text();
                if !text.is_empty() {
                    content.push_str(&text);
                }
            }
        }

        println!("{}: Received {} chunks", label, chunk_count);
        Ok(content)
    }
}