terminal_streamer.py

import asyncio
import sys
import json
import threading
import queue
from datetime import datetime
from fastapi import WebSocket
from contextlib import redirect_stdout, redirect_stderr
import io
import logging

class TerminalStreamer:
    def __init__(self):
        self.connected_clients = set()
        self.output_queue = queue.Queue()
        self.original_stdout = sys.stdout
        self.original_stderr = sys.stderr
        self.capture_enabled = False
        
    def start_capture(self):
        """Start capturing terminal output"""
        if not self.capture_enabled:
            sys.stdout = self._create_capture_wrapper(sys.stdout, "stdout")
            sys.stderr = self._create_capture_wrapper(sys.stderr, "stderr")
            self.capture_enabled = True
            
    def stop_capture(self):
        """Stop capturing terminal output"""
        if self.capture_enabled:
            sys.stdout = self.original_stdout
            sys.stderr = self.original_stderr
            self.capture_enabled = False
            
    def _create_capture_wrapper(self, original_stream, stream_type):
        """Create a wrapper that captures output and forwards to clients"""
        class StreamWrapper:
            def __init__(self, original, streamer, stream_type):
                self.original = original
                self.streamer = streamer
                self.stream_type = stream_type
                
            def write(self, text):
                # Write to original stream
                self.original.write(text)
                self.original.flush()
                
                # Send to connected clients
                if text.strip():  # Only send non-empty messages
                    message = {
                        "type": "terminal_output",
                        "stream": self.stream_type,
                        "content": text,
                        "timestamp": datetime.now().isoformat()
                    }
                    self.streamer._broadcast_message(message)
                return len(text)
                
            def flush(self):
                self.original.flush()
                
            def __getattr__(self, name):
                return getattr(self.original, name)
                
        return StreamWrapper(original_stream, self, stream_type)
    
    def _broadcast_message(self, message):
        """Broadcast message to all connected clients"""
        if self.connected_clients:
            # Use asyncio to send to all clients
            asyncio.create_task(self._send_to_all_clients(message))
    
    async def _send_to_all_clients(self, message):
        """Send message to all connected WebSocket clients"""
        if not self.connected_clients:
            return
            
        disconnected_clients = set()
        for client in self.connected_clients.copy():
            try:
                await client.send_json(message)
            except Exception as e:
                print(f"Error sending to client: {e}")
                disconnected_clients.add(client)
        
        # Remove disconnected clients
        self.connected_clients -= disconnected_clients
    
    async def add_client(self, websocket: WebSocket):
        """Add a new WebSocket client"""
        await websocket.accept()
        self.connected_clients.add(websocket)
        
        # Send welcome message
        welcome_message = {
            "type": "connection_established",
            "message": "Terminal output streaming started",
            "timestamp": datetime.now().isoformat()
        }
        await websocket.send_json(welcome_message)
        
        try:
            # Keep connection alive and handle incoming messages
            while True:
                try:
                    # Wait for messages (client can send ping/pong)
                    data = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
                    message = json.loads(data)
                    
                    if message.get("type") == "ping":
                        await websocket.send_json({"type": "pong", "timestamp": datetime.now().isoformat()})
                        
                except asyncio.TimeoutError:
                    # Send periodic heartbeat
                    await websocket.send_json({
                        "type": "heartbeat", 
                        "timestamp": datetime.now().isoformat()
                    })
                    
        except Exception as e:
            print(f"Client disconnected: {e}")
        finally:
            self.connected_clients.discard(websocket)
    
    def remove_client(self, websocket: WebSocket):
        """Remove a WebSocket client"""
        self.connected_clients.discard(websocket)

# Global instance
terminal_streamer = TerminalStreamer()