CodeReview/backend/app/services/agent/event_manager.py

"""
Agent 事件管理器
负责事件的创建、存储和推送
"""

import asyncio
import json
import logging
from typing import Optional, Dict, Any, List, AsyncGenerator, Callable
from app.core.timezone import get_now, get_now_iso
from dataclasses import dataclass
import uuid

logger = logging.getLogger(__name__)


@dataclass
class AgentEventData:
    """Agent 事件数据"""
    event_type: str
    phase: Optional[str] = None
    message: Optional[str] = None
    tool_name: Optional[str] = None
    tool_input: Optional[Dict[str, Any]] = None
    tool_output: Optional[Dict[str, Any]] = None
    tool_duration_ms: Optional[int] = None
    finding_id: Optional[str] = None
    tokens_used: int = 0
    metadata: Optional[Dict[str, Any]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "event_type": self.event_type,
            "phase": self.phase,
            "message": self.message,
            "tool_name": self.tool_name,
            "tool_input": self.tool_input,
            "tool_output": self.tool_output,
            "tool_duration_ms": self.tool_duration_ms,
            "finding_id": self.finding_id,
            "tokens_used": self.tokens_used,
            "metadata": self.metadata,
        }


class AgentEventEmitter:
    """
    Agent 事件发射器
    用于在 Agent 执行过程中发射事件
    """
    
    def __init__(self, task_id: str, event_manager: 'EventManager'):
        self.task_id = task_id
        self.event_manager = event_manager
        self._sequence = 0
        self._current_phase = None
    
    async def emit(self, event_data: AgentEventData):
        """发射事件"""
        self._sequence += 1
        event_data.phase = event_data.phase or self._current_phase
        
        await self.event_manager.add_event(
            task_id=self.task_id,
            sequence=self._sequence,
            **event_data.to_dict()
        )
    
    async def emit_phase_start(self, phase: str, message: Optional[str] = None):
        """发射阶段开始事件"""
        self._current_phase = phase
        await self.emit(AgentEventData(
            event_type="phase_start",
            phase=phase,
            message=message or f"开始 {phase} 阶段",
        ))
    
    async def emit_phase_complete(self, phase: str, message: Optional[str] = None):
        """发射阶段完成事件"""
        await self.emit(AgentEventData(
            event_type="phase_complete",
            phase=phase,
            message=message or f"{phase} 阶段完成",
        ))
    
    async def emit_thinking(self, message: str, metadata: Optional[Dict] = None):
        """发射思考事件"""
        await self.emit(AgentEventData(
            event_type="thinking",
            message=message,
            metadata=metadata,
        ))
    
    async def emit_llm_thought(self, thought: str, iteration: int = 0):
        """发射 LLM 思考内容事件 - 核心！展示 LLM 在想什么"""
        display = thought[:500] + "..." if len(thought) > 500 else thought
        await self.emit(AgentEventData(
            event_type="llm_thought",
            message=f"💭 LLM 思考:\n{display}",
            metadata={"thought": thought, "iteration": iteration},
        ))
    
    async def emit_llm_decision(self, decision: str, reason: str = ""):
        """发射 LLM 决策事件"""
        await self.emit(AgentEventData(
            event_type="llm_decision",
            message=f"💡 LLM 决策: {decision}" + (f" ({reason})" if reason else ""),
            metadata={"decision": decision, "reason": reason},
        ))
    
    async def emit_llm_action(self, action: str, action_input: Dict):
        """发射 LLM 动作事件"""
        import json
        input_str = json.dumps(action_input, ensure_ascii=False)[:200]
        await self.emit(AgentEventData(
            event_type="llm_action",
            message=f"⚡ LLM 动作: {action}\n   参数: {input_str}",
            metadata={"action": action, "action_input": action_input},
        ))
    
    async def emit_tool_call(
        self,
        tool_name: str,
        tool_input: Dict[str, Any],
        message: Optional[str] = None,
    ):
        """发射工具调用事件"""
        await self.emit(AgentEventData(
            event_type="tool_call",
            tool_name=tool_name,
            tool_input=tool_input,
            message=message or f"调用工具: {tool_name}",
        ))
    
    async def emit_tool_result(
        self,
        tool_name: str,
        tool_output: Any,
        duration_ms: int,
        message: Optional[str] = None,
    ):
        """发射工具结果事件"""
        # 处理输出，确保可序列化
        if hasattr(tool_output, 'to_dict'):
            output_data = tool_output.to_dict()
        elif isinstance(tool_output, str):
            output_data = {"result": tool_output[:2000]}  # 截断长输出
        else:
            output_data = {"result": str(tool_output)[:2000]}
        
        await self.emit(AgentEventData(
            event_type="tool_result",
            tool_name=tool_name,
            tool_output=output_data,
            tool_duration_ms=duration_ms,
            message=message or f"工具 {tool_name} 执行完成 ({duration_ms}ms)",
        ))
    
    async def emit_finding(
        self,
        finding_id: str,
        title: str,
        severity: str,
        vulnerability_type: str,
        is_verified: bool = False,
    ):
        """发射漏洞发现事件"""
        event_type = "finding_verified" if is_verified else "finding_new"
        await self.emit(AgentEventData(
            event_type=event_type,
            finding_id=finding_id,
            message=f"{'✅ 已验证' if is_verified else '🔍 新发现'}: [{severity.upper()}] {title}",
            metadata={
                "id": finding_id,  # 🔥 添加 id 字段供前端使用
                "title": title,
                "severity": severity,
                "vulnerability_type": vulnerability_type,
                "is_verified": is_verified,
            },
        ))
    
    async def emit_info(self, message: str, metadata: Optional[Dict] = None):
        """发射信息事件"""
        await self.emit(AgentEventData(
            event_type="info",
            message=message,
            metadata=metadata,
        ))
    
    async def emit_warning(self, message: str, metadata: Optional[Dict] = None):
        """发射警告事件"""
        await self.emit(AgentEventData(
            event_type="warning",
            message=message,
            metadata=metadata,
        ))
    
    async def emit_error(self, message: str, metadata: Optional[Dict] = None):
        """发射错误事件"""
        await self.emit(AgentEventData(
            event_type="error",
            message=message,
            metadata=metadata,
        ))
    
    async def emit_progress(
        self,
        current: int,
        total: int,
        message: Optional[str] = None,
    ):
        """发射进度事件"""
        percentage = (current / total * 100) if total > 0 else 0
        await self.emit(AgentEventData(
            event_type="progress",
            message=message or f"进度: {current}/{total} ({percentage:.1f}%)",
            metadata={
                "current": current,
                "total": total,
                "percentage": percentage,
            },
        ))
    
    async def emit_task_complete(
        self,
        findings_count: int,
        duration_ms: int,
        message: Optional[str] = None,
    ):
        """发射任务完成事件"""
        await self.emit(AgentEventData(
            event_type="task_complete",
            message=message or f"✅ 审计完成！发现 {findings_count} 个漏洞，耗时 {duration_ms/1000:.1f}秒",
            metadata={
                "findings_count": findings_count,
                "duration_ms": duration_ms,
            },
        ))
    
    async def emit_task_error(self, error: str, message: Optional[str] = None):
        """发射任务错误事件"""
        await self.emit(AgentEventData(
            event_type="task_error",
            message=message or f"❌ 任务失败: {error}",
            metadata={"error": error},
        ))
    
    async def emit_task_cancelled(self, message: Optional[str] = None):
        """发射任务取消事件"""
        await self.emit(AgentEventData(
            event_type="task_cancel",
            message=message or "⚠️ 任务已取消",
        ))


class EventManager:
    """
    事件管理器
    负责事件的存储和检索
    """
    
    def __init__(self, db_session_factory=None):
        self.db_session_factory = db_session_factory
        self._event_queues: Dict[str, asyncio.Queue] = {}
        self._event_callbacks: Dict[str, List[Callable]] = {}
    
    async def add_event(
        self,
        task_id: str,
        event_type: str,
        sequence: int = 0,
        phase: Optional[str] = None,
        message: Optional[str] = None,
        tool_name: Optional[str] = None,
        tool_input: Optional[Dict] = None,
        tool_output: Optional[Dict] = None,
        tool_duration_ms: Optional[int] = None,
        finding_id: Optional[str] = None,
        tokens_used: int = 0,
        metadata: Optional[Dict] = None,
    ):
        """添加事件"""
        event_id = str(uuid.uuid4())
        timestamp = get_now()
        
        event_data = {
            "id": event_id,
            "task_id": task_id,
            "event_type": event_type,
            "sequence": sequence,
            "phase": phase,
            "message": message,
            "tool_name": tool_name,
            "tool_input": tool_input,
            "tool_output": tool_output,
            "tool_duration_ms": tool_duration_ms,
            "finding_id": finding_id,
            "tokens_used": tokens_used,
            "metadata": metadata,
            "timestamp": timestamp.isoformat(),
        }
        
        # 保存到数据库（跳过高频事件如 thinking_token）
        skip_db_events = {"thinking_token"}
        if self.db_session_factory and event_type not in skip_db_events:
            try:
                await self._save_event_to_db(event_data)
            except Exception as e:
                logger.error(f"Failed to save event to database: {e}")
        
        # 推送到队列（非阻塞）
        if task_id in self._event_queues:
            try:
                self._event_queues[task_id].put_nowait(event_data)
                # 🔥 DEBUG: 记录重要事件被添加到队列
                if event_type in ["thinking_start", "thinking_end", "dispatch", "task_complete", "task_error", "tool_call", "tool_result", "llm_action"]:
                    logger.info(f"[EventQueue] Added {event_type} to queue for task {task_id}, queue size: {self._event_queues[task_id].qsize()}")
                elif event_type == "thinking_token":
                    # 每10个token记录一次
                    if sequence % 10 == 0:
                        logger.debug(f"[EventQueue] Added thinking_token #{sequence} to queue, size: {self._event_queues[task_id].qsize()}")
            except asyncio.QueueFull:
                logger.warning(f"Event queue full for task {task_id}, dropping event: {event_type}")
        
        # 调用回调
        if task_id in self._event_callbacks:
            for callback in self._event_callbacks[task_id]:
                try:
                    if asyncio.iscoroutinefunction(callback):
                        await callback(event_data)
                    else:
                        callback(event_data)
                except Exception as e:
                    logger.error(f"Event callback error: {e}")
        
        return event_id
    
    async def _save_event_to_db(self, event_data: Dict):
        """保存事件到数据库"""
        from app.models.agent_task import AgentEvent

        # 🔥 清理无效的 UTF-8 字符（如二进制内容）
        def sanitize_string(s):
            """清理字符串中的无效 UTF-8 字符"""
            if s is None:
                return None
            if not isinstance(s, str):
                s = str(s)
            # 移除 NULL 字节和其他不可打印的控制字符（保留换行和制表符）
            return ''.join(
                char for char in s
                if char in '\n\r\t' or (ord(char) >= 32 and ord(char) != 127)
            )

        def sanitize_dict(d):
            """递归清理字典中的字符串值"""
            if d is None:
                return None
            if isinstance(d, dict):
                return {k: sanitize_dict(v) for k, v in d.items()}
            elif isinstance(d, list):
                return [sanitize_dict(item) for item in d]
            elif isinstance(d, str):
                return sanitize_string(d)
            return d

        async with self.db_session_factory() as db:
            event = AgentEvent(
                id=event_data["id"],
                task_id=event_data["task_id"],
                event_type=event_data["event_type"],
                sequence=event_data["sequence"],
                phase=event_data["phase"],
                message=sanitize_string(event_data["message"]),  # 🔥 清理消息
                tool_name=event_data["tool_name"],
                tool_input=sanitize_dict(event_data["tool_input"]),  # 🔥 清理工具输入
                tool_output=sanitize_dict(event_data["tool_output"]),  # 🔥 清理工具输出
                tool_duration_ms=event_data["tool_duration_ms"],
                finding_id=event_data["finding_id"],
                tokens_used=event_data["tokens_used"],
                event_metadata=sanitize_dict(event_data["metadata"]),  # 🔥 清理元数据
                created_at=get_now(),
            )
            db.add(event)
            await db.commit()
    
    def create_queue(self, task_id: str) -> asyncio.Queue:
        """创建或获取事件队列"""
        if task_id not in self._event_queues:
            # 🔥 使用较大的队列容量，缓存更多 token 事件
            self._event_queues[task_id] = asyncio.Queue(maxsize=5000)
        return self._event_queues[task_id]
    
    def remove_queue(self, task_id: str):
        """移除事件队列"""
        if task_id in self._event_queues:
            del self._event_queues[task_id]
    
    def add_callback(self, task_id: str, callback: Callable):
        """添加事件回调"""
        if task_id not in self._event_callbacks:
            self._event_callbacks[task_id] = []
        self._event_callbacks[task_id].append(callback)
    
    def remove_callback(self, task_id: str, callback: Callable):
        """移除事件回调"""
        if task_id in self._event_callbacks:
            self._event_callbacks[task_id].remove(callback)
    
    async def get_events(
        self,
        task_id: str,
        after_sequence: int = 0,
        limit: int = 100,
    ) -> List[Dict]:
        """获取事件列表"""
        if not self.db_session_factory:
            return []
        
        from sqlalchemy.future import select
        from app.models.agent_task import AgentEvent
        
        async with self.db_session_factory() as db:
            result = await db.execute(
                select(AgentEvent)
                .where(AgentEvent.task_id == task_id)
                .where(AgentEvent.sequence > after_sequence)
                .order_by(AgentEvent.sequence)
                .limit(limit)
            )
            events = result.scalars().all()
            return [event.to_sse_dict() for event in events]
    
    async def stream_events(
        self,
        task_id: str,
        after_sequence: int = 0,
    ) -> AsyncGenerator[Dict, None]:
        """流式获取事件

        🔥 重要: 此方法会先排空队列中已缓存的事件（在 SSE 连接前产生的），
        然后继续实时推送新事件。
        只返回序列号 > after_sequence 的事件。
        """
        logger.info(f"[StreamEvents] Task {task_id}: Starting stream with after_sequence={after_sequence}")

        # 获取现有队列（由 AgentRunner 在初始化时创建）
        queue = self._event_queues.get(task_id)
        if not queue:
            # 如果队列不存在，创建一个新的（回退逻辑）
            queue = self.create_queue(task_id)
            logger.warning(f"Queue not found for task {task_id}, created new one")

        # 🔥 CRITICAL FIX: 记录当前队列大小，只消耗这些已存在的事件
        # 之前的 bug: while not queue.empty() 会永远循环，因为 LLM 持续添加事件
        initial_queue_size = queue.qsize()
        logger.info(f"[StreamEvents] Task {task_id}: Draining {initial_queue_size} buffered events...")

        # 🔥 先排空队列中已缓存的事件（只消耗连接时已存在的事件数量）
        buffered_count = 0
        skipped_count = 0
        max_drain = initial_queue_size  # 只消耗这么多事件，避免无限循环
        
        for _ in range(max_drain):
            try:
                buffered_event = queue.get_nowait()

                # 🔥 过滤掉序列号 <= after_sequence 的事件
                event_sequence = buffered_event.get("sequence", 0)
                if event_sequence <= after_sequence:
                    skipped_count += 1
                    continue

                buffered_count += 1
                yield buffered_event

                # 🔥 取消人为延迟，防止队列堆积
                event_type = buffered_event.get("event_type")
                # if event_type == "thinking_token":
                #     await asyncio.sleep(0.005)
                # 其他事件不加延迟，快速发送

                # 检查是否是结束事件
                if event_type in ["task_complete", "task_error", "task_cancel"]:
                    logger.info(f"[StreamEvents] Task {task_id} already completed, sent {buffered_count} buffered events (skipped {skipped_count})")
                    return
            except asyncio.QueueEmpty:
                break

        if buffered_count > 0 or skipped_count > 0:
            logger.info(f"[StreamEvents] Task {task_id}: Drained {buffered_count} buffered events, skipped {skipped_count}")

        # 🔥 DEBUG: 记录进入实时循环
        logger.info(f"[StreamEvents] Task {task_id}: Entering real-time loop, queue size: {queue.qsize()}")

        # 然后实时推送新事件
        try:
            while True:
                try:
                    logger.debug(f"[StreamEvents] Task {task_id}: Waiting for next event from queue...")
                    event = await asyncio.wait_for(queue.get(), timeout=30)
                    logger.debug(f"[StreamEvents] Task {task_id}: Got event from queue: {event.get('event_type')}")

                    # 🔥 过滤掉序列号 <= after_sequence 的事件
                    event_sequence = event.get("sequence", 0)
                    if event_sequence <= after_sequence:
                        logger.debug(f"[StreamEvents] Task {task_id}: Skipping event seq={event_sequence} (after_sequence={after_sequence})")
                        continue

                    # 🔥 DEBUG: 记录重要事件被发送
                    event_type = event.get("event_type")
                    if event_type in ["thinking_start", "thinking_end", "dispatch", "task_complete", "task_error", "tool_call", "tool_result", "llm_action"]:
                        logger.info(f"[StreamEvents] Yielding {event_type} (seq={event_sequence}) for task {task_id}")

                    yield event

                    # 🔥 取消人为延迟，防止队列堆积
                    # if event_type == "thinking_token":
                    #     await asyncio.sleep(0.01)

                    # 检查是否是结束事件
                    if event.get("event_type") in ["task_complete", "task_error", "task_cancel"]:
                        break

                except asyncio.TimeoutError:
                    # 发送心跳
                    yield {"event_type": "heartbeat", "timestamp": get_now_iso()}

        except GeneratorExit:
            # SSE 连接断开
            logger.debug(f"SSE stream closed for task {task_id}")
        # 🔥 不要移除队列，让 AgentRunner 管理队列的生命周期
    
    def create_emitter(self, task_id: str) -> AgentEventEmitter:
        """创建事件发射器"""
        return AgentEventEmitter(task_id, self)
    
    async def close(self):
        """关闭事件管理器，清理资源"""
        # 清理所有队列
        for task_id in list(self._event_queues.keys()):
            self.remove_queue(task_id)
        
        # 清理所有回调
        self._event_callbacks.clear()
        
        logger.debug("EventManager closed")