cai/tests/integration/test_streaming_agent_execution.py

#!/usr/bin/env python3
"""
Integration tests for streaming during actual agent execution.

This test suite verifies streaming works correctly when:
- Agents execute real tools and generate responses
- Multiple tool calls happen in sequence
- Agents hand off to other agents
- Long-running operations are interrupted
- Memory and context are preserved across turns
"""

import asyncio
import os
import sys
import time
import unittest
from typing import Dict, List, Optional, Any, AsyncIterator
from unittest.mock import MagicMock, patch, AsyncMock, Mock
import tempfile
import json
from dataclasses import dataclass

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))

from cai.sdk.agents.agent import Agent
from cai.sdk.agents.run import Runner, RunResult, StreamChunk
from cai.sdk.agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
from cai.sdk.agents.tool import FunctionTool
from cai.sdk.agents.items import MessageOutputItem, ToolCallItem, ToolCallOutputItem, ReasoningItem
from cai.sdk.agents.run_context import RunContext
from cai.tui.display.streaming_display import StreamingDisplay
from cai.tui.display.tool_display import ToolDisplay
from cai.tui.display.agent_display import AgentDisplay
from cai.tui.display.manager import DisplayManager
from cai.tui.display.base import DisplayContext


@dataclass
class MockStreamChunk:
    """Mock stream chunk for testing."""
    chunk_type: str  # 'thinking', 'content', 'tool_call', 'tool_output'
    content: Optional[str] = None
    tool_name: Optional[str] = None
    tool_args: Optional[Dict] = None
    tool_output: Optional[str] = None
    finished: bool = False
    final_stats: Optional[Dict] = None


class MockRunner:
    """Mock runner that simulates agent execution with streaming."""
    
    def __init__(self, agent: Agent):
        self.agent = agent
        self.messages = []
        self.stream_chunks = []
        
    async def run_async(self, messages: List[Dict], stream: bool = True) -> AsyncIterator[StreamChunk]:
        """Simulate async agent execution with streaming."""
        self.messages = messages
        
        # Simulate thinking
        if self.agent.name == "o1-preview" or "think" in str(messages).lower():
            thinking_content = "I need to analyze this request and determine the best approach..."
            for i in range(0, len(thinking_content), 10):
                chunk = MockStreamChunk(
                    chunk_type='thinking',
                    content=thinking_content[:i+10]
                )
                self.stream_chunks.append(chunk)
                yield chunk
                await asyncio.sleep(0.01)
                
        # Simulate tool calls based on agent type
        if "command" in str(messages).lower() or self.agent.name in ["Bug Bounter", "Red Teamer"]:
            # Tool call
            tool_chunk = MockStreamChunk(
                chunk_type='tool_call',
                tool_name='generic_linux_command',
                tool_args={'command': 'ls', 'args': '-la'}
            )
            self.stream_chunks.append(tool_chunk)
            yield tool_chunk
            
            # Tool output
            tool_output = "total 48\ndrwxr-xr-x  5 user user 4096 Jan 15 10:00 .\ndrwxr-xr-x 10 user user 4096 Jan 15 09:00 .."
            output_chunk = MockStreamChunk(
                chunk_type='tool_output',
                tool_output=tool_output
            )
            self.stream_chunks.append(output_chunk)
            yield output_chunk
            
        # Simulate content streaming
        response_content = f"{self.agent.name} responding to: {messages[-1].get('content', '')[:50]}..."
        for i in range(0, len(response_content), 5):
            chunk = MockStreamChunk(
                chunk_type='content',
                content=response_content[:i+5]
            )
            self.stream_chunks.append(chunk)
            yield chunk
            await asyncio.sleep(0.005)
            
        # Final chunk with stats
        final_chunk = MockStreamChunk(
            chunk_type='content',
            content=response_content,
            finished=True,
            final_stats={
                'input_tokens': 150,
                'output_tokens': 50,
                'interaction_cost': 0.0025,
                'session_total_cost': 0.0075
            }
        )
        self.stream_chunks.append(final_chunk)
        yield final_chunk
        
    def run_sync(self, messages: List[Dict]) -> RunResult:
        """Simulate sync agent execution."""
        self.messages = messages
        
        # Create mock result
        items = []
        
        # Add reasoning if applicable
        if self.agent.name == "o1-preview":
            items.append(ReasoningItem(
                content="Analyzing the request and planning the approach..."
            ))
            
        # Add tool calls
        if "command" in str(messages).lower():
            items.extend([
                ToolCallItem(
                    call_id="call_123",
                    tool_name="generic_linux_command",
                    tool_args={'command': 'ls', 'args': '-la'}
                ),
                ToolCallOutputItem(
                    call_id="call_123",
                    tool_name="generic_linux_command",
                    output="file1.txt\nfile2.py\n"
                )
            ])
            
        # Add message
        items.append(MessageOutputItem(
            role="assistant",
            content=f"{self.agent.name} response to query"
        ))
        
        return RunResult(
            items=items,
            usage={
                'input_tokens': 100,
                'output_tokens': 50,
                'total_tokens': 150
            }
        )


class TestStreamingAgentExecution(unittest.TestCase):
    """Test streaming during real agent execution flows."""
    
    def setUp(self):
        """Set up test environment."""
        os.environ["CAI_TELEMETRY"] = "false"
        os.environ["CAI_TRACING"] = "false"
        os.environ["CAI_STREAM"] = "true"
        
        # Create display manager
        self.display_manager = DisplayManager()
        
        # Mock terminal outputs
        self.terminal_outputs = {}
        
        # Patch get_terminal_output
        self.patcher = patch('cai.tui.core.terminal_console.get_terminal_output')
        self.mock_get_terminal = self.patcher.start()
        self.mock_get_terminal.side_effect = self._get_mock_terminal
        
        # Create test agents
        self.test_agents = self._create_test_agents()
        
    def tearDown(self):
        """Clean up."""
        self.patcher.stop()
        
    def _get_mock_terminal(self, terminal_id):
        """Get or create mock terminal."""
        if terminal_id not in self.terminal_outputs:
            terminal = MagicMock()
            terminal.write = MagicMock()
            terminal.start_streaming_line = MagicMock()
            terminal.update_streaming_line = MagicMock()
            terminal.finish_streaming_line = MagicMock()
            terminal._streaming_lines = {}
            self.terminal_outputs[terminal_id] = terminal
        return self.terminal_outputs[terminal_id]
        
    def _create_test_agents(self) -> Dict[str, Agent]:
        """Create test agents with different configurations."""
        agents = {}
        
        # Bug Bounter - security testing agent
        bug_tools = [
            FunctionTool(
                name="generic_linux_command",
                description="Execute Linux commands",
                parameters={
                    "type": "object",
                    "properties": {
                        "command": {"type": "string"},
                        "args": {"type": "string"}
                    }
                }
            ),
            FunctionTool(
                name="web_search",
                description="Search the web",
                parameters={
                    "type": "object",
                    "properties": {
                        "query": {"type": "string"}
                    }
                }
            )
        ]
        
        agents["bug_bounter"] = Agent(
            name="Bug Bounter",
            instructions="You are a security testing agent.",
            tools=bug_tools
        )
        
        # Red Teamer - penetration testing agent
        agents["red_teamer"] = Agent(
            name="Red Teamer",
            instructions="You are a penetration testing agent.",
            tools=bug_tools
        )
        
        # Code Agent - code analysis
        code_tools = [
            FunctionTool(
                name="read_file",
                description="Read file contents",
                parameters={
                    "type": "object",
                    "properties": {
                        "path": {"type": "string"}
                    }
                }
            ),
            FunctionTool(
                name="execute_code",
                description="Execute code",
                parameters={
                    "type": "object",
                    "properties": {
                        "code": {"type": "string"},
                        "language": {"type": "string"}
                    }
                }
            )
        ]
        
        agents["code_agent"] = Agent(
            name="Code Agent",
            instructions="You are a code analysis agent.",
            tools=code_tools
        )
        
        return agents
        
    async def test_single_agent_streaming_execution(self):
        """Test streaming during single agent execution."""
        agent = self.test_agents["bug_bounter"]
        runner = MockRunner(agent)
        
        # Create display context
        context = DisplayContext(
            terminal_id="test-1",
            terminal_number=1,
            agent_name=agent.name,
            agent_id="bug-1",
            interaction_counter=1
        )
        
        # Track display calls
        display_calls = []
        
        # Mock display manager dispatch
        original_dispatch = self.display_manager.dispatch
        def mock_dispatch(event_type, context, data):
            display_calls.append((event_type, context, data))
            return original_dispatch(event_type, context, data)
        self.display_manager.dispatch = mock_dispatch
        
        # Execute agent with streaming
        messages = [{"role": "user", "content": "Scan the application for vulnerabilities"}]
        
        streamed_content = ""
        async for chunk in runner.run_async(messages, stream=True):
            # Simulate display manager handling chunks
            if chunk.chunk_type == 'thinking':
                self.display_manager.dispatch('stream_start', context, {
                    'stream_id': 'thinking-1',
                    'content_type': 'thinking',
                    'model': agent.name
                })
                self.display_manager.dispatch('stream_update', context, {
                    'stream_id': 'thinking-1',
                    'content': chunk.content
                })
                if chunk.finished:
                    self.display_manager.dispatch('stream_finish', context, {
                        'stream_id': 'thinking-1'
                    })
                    
            elif chunk.chunk_type == 'content':
                if not streamed_content:  # First content chunk
                    self.display_manager.dispatch('stream_start', context, {
                        'stream_id': 'content-1',
                        'content_type': 'text',
                        'model': agent.name
                    })
                streamed_content = chunk.content
                self.display_manager.dispatch('stream_update', context, {
                    'stream_id': 'content-1',
                    'content': chunk.content
                })
                if chunk.finished:
                    self.display_manager.dispatch('stream_finish', context, {
                        'stream_id': 'content-1',
                        'final_stats': chunk.final_stats
                    })
                    
            elif chunk.chunk_type == 'tool_call':
                self.display_manager.dispatch('tool_start', context, {
                    'stream_id': 'tool-1',
                    'tool_name': chunk.tool_name,
                    'args': chunk.tool_args
                })
                
            elif chunk.chunk_type == 'tool_output':
                self.display_manager.dispatch('tool_update', context, {
                    'stream_id': 'tool-1',
                    'output': chunk.tool_output
                })
                self.display_manager.dispatch('tool_finish', context, {
                    'stream_id': 'tool-1',
                    'output': chunk.tool_output,
                    'execution_info': {'status': 'completed'}
                })
                
        # Verify display calls were made
        self.assertTrue(len(display_calls) > 0)
        
        # Check for expected display events
        event_types = [call[0] for call in display_calls]
        self.assertIn('stream_start', event_types)
        self.assertIn('stream_update', event_types)
        self.assertIn('stream_finish', event_types)
        
        # Verify terminal output
        terminal = self.terminal_outputs[context.terminal_id]
        self.assertTrue(terminal.start_streaming_line.called or terminal.write.called)
        
    async def test_multi_tool_execution_streaming(self):
        """Test streaming with multiple tool executions."""
        agent = self.test_agents["red_teamer"]
        runner = MockRunner(agent)
        
        context = DisplayContext(
            terminal_id="test-2",
            terminal_number=1,
            agent_name=agent.name,
            agent_id="red-1",
            interaction_counter=1
        )
        
        # Simulate agent making multiple tool calls
        class MultiToolRunner(MockRunner):
            async def run_async(self, messages, stream=True):
                # First tool call - nmap scan
                yield MockStreamChunk(
                    chunk_type='tool_call',
                    tool_name='generic_linux_command',
                    tool_args={'command': 'nmap', 'args': '-sn 192.168.1.0/24'}
                )
                yield MockStreamChunk(
                    chunk_type='tool_output',
                    tool_output='2 hosts up'
                )
                
                # Second tool call - vulnerability scan
                yield MockStreamChunk(
                    chunk_type='tool_call',
                    tool_name='generic_linux_command',
                    tool_args={'command': 'nikto', 'args': '-h target.com'}
                )
                yield MockStreamChunk(
                    chunk_type='tool_output',
                    tool_output='Found 5 potential issues'
                )
                
                # Final response
                response = "Completed network scan. Found 2 active hosts and 5 potential vulnerabilities."
                for i in range(0, len(response), 10):
                    yield MockStreamChunk(
                        chunk_type='content',
                        content=response[:i+10]
                    )
                    await asyncio.sleep(0.01)
                    
                yield MockStreamChunk(
                    chunk_type='content',
                    content=response,
                    finished=True,
                    final_stats={'output_tokens': 25}
                )
                
        runner = MultiToolRunner(agent)
        
        # Track tool executions
        tool_executions = []
        
        async for chunk in runner.run_async([{"role": "user", "content": "Scan the network"}]):
            if chunk.chunk_type == 'tool_call':
                tool_executions.append(chunk.tool_name)
                self.display_manager.dispatch('tool_start', context, {
                    'stream_id': f'tool-{len(tool_executions)}',
                    'tool_name': chunk.tool_name,
                    'args': chunk.tool_args
                })
            elif chunk.chunk_type == 'tool_output':
                self.display_manager.dispatch('tool_finish', context, {
                    'stream_id': f'tool-{len(tool_executions)}',
                    'output': chunk.tool_output,
                    'execution_info': {'status': 'completed'}
                })
                
        # Verify multiple tools were executed
        self.assertEqual(len(tool_executions), 2)
        self.assertIn('generic_linux_command', tool_executions)
        
    async def test_streaming_with_interruption(self):
        """Test handling interruption during streaming."""
        agent = self.test_agents["code_agent"]
        
        context = DisplayContext(
            terminal_id="test-3",
            terminal_number=1,
            agent_name=agent.name,
            agent_id="code-1",
            interaction_counter=1
        )
        
        # Simulate interrupted execution
        class InterruptedRunner(MockRunner):
            async def run_async(self, messages, stream=True):
                # Start streaming content
                content = "Analyzing code for security vulnerabilities..."
                for i in range(0, len(content)//2, 5):  # Only stream half
                    yield MockStreamChunk(
                        chunk_type='content',
                        content=content[:i+5]
                    )
                    await asyncio.sleep(0.01)
                    
                # Simulate interruption
                raise KeyboardInterrupt("User interrupted")
                
        runner = InterruptedRunner(agent)
        
        # Handle streaming with interruption
        stream_id = None
        partial_content = ""
        
        try:
            async for chunk in runner.run_async([{"role": "user", "content": "Analyze this code"}]):
                if chunk.chunk_type == 'content':
                    if not stream_id:
                        stream_id = 'interrupted-stream'
                        self.display_manager.dispatch('stream_start', context, {
                            'stream_id': stream_id,
                            'content_type': 'text'
                        })
                    partial_content = chunk.content
                    self.display_manager.dispatch('stream_update', context, {
                        'stream_id': stream_id,
                        'content': chunk.content
                    })
        except KeyboardInterrupt:
            # Clean up on interruption
            if stream_id:
                self.display_manager.dispatch('stream_finish', context, {
                    'stream_id': stream_id,
                    'final_stats': {'interrupted': True}
                })
                
        # Verify partial content was displayed
        self.assertTrue(len(partial_content) > 0)
        self.assertLess(len(partial_content), 40)  # Should be partial
        
    async def test_parallel_agent_streaming(self):
        """Test multiple agents streaming in parallel."""
        agents = [
            self.test_agents["bug_bounter"],
            self.test_agents["red_teamer"],
            self.test_agents["code_agent"]
        ]
        
        # Create contexts for each agent
        contexts = []
        for i, agent in enumerate(agents):
            contexts.append(DisplayContext(
                terminal_id=f"parallel-{i}",
                terminal_number=i+1,
                agent_name=agent.name,
                agent_id=f"agent-{i}",
                interaction_counter=1
            ))
            
        # Run agents in parallel
        async def run_agent(agent, context, query):
            runner = MockRunner(agent)
            
            async for chunk in runner.run_async([{"role": "user", "content": query}]):
                if chunk.chunk_type == 'content':
                    if chunk.finished:
                        self.display_manager.dispatch('stream_finish', context, {
                            'stream_id': f'{context.agent_id}-stream',
                            'final_stats': chunk.final_stats
                        })
                    else:
                        self.display_manager.dispatch('stream_update', context, {
                            'stream_id': f'{context.agent_id}-stream',
                            'content': chunk.content
                        })
                        
        # Execute all agents in parallel
        tasks = []
        queries = [
            "Find security vulnerabilities",
            "Test authentication bypass",
            "Review code for bugs"
        ]
        
        for agent, context, query in zip(agents, contexts, queries):
            task = asyncio.create_task(run_agent(agent, context, query))
            tasks.append(task)
            
        await asyncio.gather(*tasks)
        
        # Verify all agents completed
        for i, context in enumerate(contexts):
            terminal = self.terminal_outputs[context.terminal_id]
            self.assertTrue(
                terminal.write.called or 
                terminal.start_streaming_line.called or
                terminal.update_streaming_line.called
            )
            
    async def test_streaming_with_context_preservation(self):
        """Test that streaming preserves context across turns."""
        agent = self.test_agents["bug_bounter"]
        runner = MockRunner(agent)
        
        context = DisplayContext(
            terminal_id="test-4",
            terminal_number=1,
            agent_name=agent.name,
            agent_id="bug-2",
            interaction_counter=1
        )
        
        # First turn
        messages = [{"role": "user", "content": "What vulnerabilities should I look for?"}]
        
        first_response = ""
        async for chunk in runner.run_async(messages):
            if chunk.chunk_type == 'content':
                first_response = chunk.content
                
        # Second turn - context should be preserved
        context.interaction_counter = 2
        messages.append({"role": "assistant", "content": first_response})
        messages.append({"role": "user", "content": "Can you scan for SQL injection?"})
        
        second_response = ""
        async for chunk in runner.run_async(messages):
            if chunk.chunk_type == 'content':
                second_response = chunk.content
                
        # Verify both responses were generated
        self.assertTrue(len(first_response) > 0)
        self.assertTrue(len(second_response) > 0)
        self.assertNotEqual(first_response, second_response)
        
        # Verify interaction counter was incremented
        self.assertEqual(context.interaction_counter, 2)


def run_async_test(coro):
    """Helper to run async test."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(coro)
    finally:
        loop.close()


class TestStreamingAgentExecutionSync(TestStreamingAgentExecution):
    """Synchronous wrapper for async tests."""
    
    def test_single_agent_streaming_execution_sync(self):
        run_async_test(self.test_single_agent_streaming_execution())
        
    def test_multi_tool_execution_streaming_sync(self):
        run_async_test(self.test_multi_tool_execution_streaming())
        
    def test_streaming_with_interruption_sync(self):
        run_async_test(self.test_streaming_with_interruption())
        
    def test_parallel_agent_streaming_sync(self):
        run_async_test(self.test_parallel_agent_streaming())
        
    def test_streaming_with_context_preservation_sync(self):
        run_async_test(self.test_streaming_with_context_preservation())


if __name__ == "__main__":
    unittest.main(verbosity=2)