Backend/core/flow_tracker.py

import asyncio
import json
from datetime import datetime
from typing import Optional, Callable, Any
from uuid import UUID
from dataclasses import dataclass, field, asdict

from Backend.core.logging import get_logger

logger = get_logger(__name__)


@dataclass
class AgentStep:
    agent_name: str
    status: str
    started_at: str
    completed_at: Optional[str] = None
    duration_ms: Optional[float] = None
    decision: Optional[str] = None
    reasoning: Optional[str] = None
    result: Optional[dict] = None
    error: Optional[str] = None


@dataclass
class PipelineFlow:
    issue_id: UUID
    started_at: str
    status: str = "running"
    completed_at: Optional[str] = None
    total_duration_ms: Optional[float] = None
    steps: list[AgentStep] = field(default_factory=list)
    final_result: Optional[dict] = None
    
    def to_dict(self) -> dict:
        return {
            "issue_id": str(self.issue_id),
            "started_at": self.started_at,
            "status": self.status,
            "completed_at": self.completed_at,
            "total_duration_ms": self.total_duration_ms,
            "steps": [asdict(s) for s in self.steps],
            "final_result": self.final_result,
        }


class FlowTracker:
    def __init__(self, issue_id: UUID):
        self.flow = PipelineFlow(
            issue_id=issue_id,
            started_at=datetime.utcnow().isoformat(),
        )
        self._start_time = datetime.utcnow()
        self._subscribers: list[asyncio.Queue] = []
    
    def subscribe(self) -> asyncio.Queue:
        queue = asyncio.Queue()
        
        
        for step in self.flow.steps:
            if step.started_at:
                queue.put_nowait({
                    "type": "step_started",
                    "timestamp": step.started_at,
                    "data": {
                        "agent_name": step.agent_name, 
                        "step_index": self.flow.steps.index(step)
                    }
                })
            
            
            if step.status in ("completed", "error"):
                 queue.put_nowait({
                    "type": "step_completed" if step.status == "completed" else "step_error",
                    "timestamp": step.completed_at,
                    "data": {
                        "agent_name": step.agent_name,
                        "status": step.status,
                        "decision": step.decision,
                        "reasoning": step.reasoning,
                        "result": step.result,
                        "error": step.error
                    }
                })

        self._subscribers.append(queue)
        return queue
    
    def unsubscribe(self, queue: asyncio.Queue):
        if queue in self._subscribers:
            self._subscribers.remove(queue)
    
    async def _broadcast(self, event_type: str, data: dict):
        message = {
            "type": event_type,
            "timestamp": datetime.utcnow().isoformat(),
            "data": data,
        }
        for queue in self._subscribers:
            await queue.put(message)
    
    async def start_step(self, agent_name: str):
        step = AgentStep(
            agent_name=agent_name,
            status="running",
            started_at=datetime.utcnow().isoformat(),
        )
        self.flow.steps.append(step)
        
        await self._broadcast("step_started", {
            "agent_name": agent_name,
            "step_index": len(self.flow.steps) - 1,
        })
        
        return step
    
    async def complete_step(
        self,
        agent_name: str,
        decision: str,
        reasoning: str,
        result: Optional[dict] = None,
        error: Optional[str] = None
    ):
        step = next((s for s in self.flow.steps if s.agent_name == agent_name and s.status == "running"), None)
        if step:
            now = datetime.utcnow()
            step.completed_at = now.isoformat()
            step.status = "error" if error else "completed"
            step.decision = decision
            step.reasoning = reasoning
            step.result = result
            step.error = error
            
            started = datetime.fromisoformat(step.started_at)
            step.duration_ms = (now - started).total_seconds() * 1000
        
        await self._broadcast("step_completed", {
            "agent_name": agent_name,
            "status": step.status if step else "unknown",
            "decision": decision,
            "reasoning": reasoning,
            "duration_ms": step.duration_ms if step else 0,
            "result": result,
            "error": error,
        })
    
    async def complete_flow(self, final_result: dict):
        now = datetime.utcnow()
        self.flow.completed_at = now.isoformat()
        self.flow.status = "completed"
        self.flow.total_duration_ms = (now - self._start_time).total_seconds() * 1000
        self.flow.final_result = final_result
        
        await self._broadcast("flow_completed", self.flow.to_dict())
    
    async def error_flow(self, error: str):
        now = datetime.utcnow()
        self.flow.completed_at = now.isoformat()
        self.flow.status = "error"
        self.flow.total_duration_ms = (now - self._start_time).total_seconds() * 1000
        
        await self._broadcast("flow_error", {
            "error": error,
            "flow": self.flow.to_dict(),
        })


_active_flows: dict[UUID, FlowTracker] = {}


def get_flow_tracker(issue_id: UUID) -> Optional[FlowTracker]:
    return _active_flows.get(issue_id)


def create_flow_tracker(issue_id: UUID) -> FlowTracker:
    if issue_id in _active_flows:
        return _active_flows[issue_id]
    
    tracker = FlowTracker(issue_id)
    _active_flows[issue_id] = tracker
    return tracker


def remove_flow_tracker(issue_id: UUID):
    if issue_id in _active_flows:
        del _active_flows[issue_id]