codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/handlers/websocket.py

@@ -0,0 +1,497 @@
+"""WebSocket handler for Agent communication."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any, Optional
+
+from jupyter_server.base.handlers import JupyterHandler
+from tornado.websocket import WebSocketHandler
+
+# Agent-mode dependencies (langgraph/langchain require Python >=3.10)
+try:
+    from codebuff_core.core.exceptions import AgentError
+    from langgraph.types import Command
+
+    from ..agent import AgentFactory, AgentRunner, config_from_dict
+    from ..agent.agent_factory import AgentEvent
+    from ..agent.middleware import MemoryStore
+    from ..tools import NotebookContextManager
+
+    _AGENT_AVAILABLE = True
+except ImportError:
+    _AGENT_AVAILABLE = False
+
+from ..log_paths import get_log_dir
+
+logger = logging.getLogger(__name__)
+
+_LOGS_DIR = get_log_dir()
+_MEMORY_FILE = _LOGS_DIR / "memory" / "memories.json"
+
+
+class AgentWebSocketHandler(WebSocketHandler, JupyterHandler):
+    """
+    WebSocket handler for real-time agent communication.
+
+    Uses the LangChain Agent-based AgentFactory to create and run agents.
+    HITL은 HumanInTheLoopMiddleware의 그래프 인터럽트 방식으로 처리.
+
+    Handles JSON-RPC style messages for:
+    - Starting agent runs
+    - Canceling agent runs
+    - Responding to HITL interrupts (approve/edit/reject)
+    """
+
+    def initialize(self) -> None:
+        """Initialize handler state."""
+        self._agent_runner: Optional[AgentRunner] = None
+        self._agent_task: Optional[asyncio.Task] = None
+        self._hitl_future: Optional[asyncio.Future] = None
+        self._is_closed = False
+        self._notebook_context: Optional[NotebookContextManager] = None
+        self._memory_store = MemoryStore(storage_path=_MEMORY_FILE)
+
+    def check_origin(self, origin: str) -> bool:
+        """Allow connections from same origin."""
+        return True
+
+    async def open(self) -> None:
+        """Handle WebSocket connection opened."""
+        logger.info("Agent WebSocket connected")
+        await self._send_event("connected", {"status": "ready"})
+
+    def on_close(self) -> None:
+        """Handle WebSocket connection closed."""
+        logger.info("Agent WebSocket disconnected")
+        self._is_closed = True
+        self._memory_store.flush_now()  # 보류 중인 메모리 즉시 저장
+
+        # Cancel any running agent task
+        if self._agent_task and not self._agent_task.done():
+            self._agent_task.cancel()
+
+        # Cancel any running agent
+        if self._agent_runner:
+            self._agent_runner.cancel()
+
+        # Cancel any pending HITL request
+        if self._hitl_future and not self._hitl_future.done():
+            self._hitl_future.cancel()
+
+    async def on_message(self, message: str) -> None:
+        """
+        Handle incoming WebSocket message.
+
+        Args:
+            message: JSON-RPC formatted message string
+        """
+        try:
+            data = json.loads(message)
+        except json.JSONDecodeError:
+            await self._send_error(None, -32700, "Parse error")
+            return
+
+        method = data.get("method", "")
+        params = data.get("params", {})
+        msg_id = data.get("id")
+
+        logger.debug(f"Received message: method={method}, id={msg_id}")
+
+        # Route to appropriate handler
+        if method == "agent/run":
+            await self._handle_run(msg_id, params)
+        elif method == "agent/cancel":
+            await self._handle_cancel(msg_id, params)
+        elif method == "agent/hitl_response":
+            await self._handle_hitl_response(msg_id, params)
+        else:
+            await self._send_error(msg_id, -32601, f"Method not found: {method}")
+
+    async def _handle_run(self, msg_id: Optional[str], params: dict[str, Any]) -> None:
+        """
+        Handle agent/run request.
+
+        Args:
+            msg_id: Request ID
+            params: Request parameters
+        """
+        # Check if agent is already running
+        if self._agent_task and not self._agent_task.done():
+            await self._send_error(msg_id, -32000, "Agent is already running")
+            return
+
+        prompt = params.get("prompt", "")
+        mode = params.get("mode", "default")
+        max_iterations = params.get("maxIterations", 20)
+        notebook_context = params.get("notebookContext")
+
+        if not prompt:
+            await self._send_error(msg_id, -32602, "Prompt is required")
+            return
+
+        # Create notebook context manager
+        self._notebook_context = NotebookContextManager(
+            kernel_manager=getattr(self, "kernel_manager", None),
+            contents_manager=getattr(self, "contents_manager", None),
+        )
+
+        # Set notebook context if provided
+        if notebook_context:
+            self._notebook_context.set_notebook(
+                notebook_context.get("path", ""),
+                notebook_context,
+            )
+
+        # Create agent factory (HITL은 미들웨어가 처리)
+        factory = AgentFactory(
+            notebook_context=self._notebook_context,
+            on_event=self._on_agent_event,
+            memory_store=self._memory_store,
+        )
+
+        llm_config = _turn_llm_config(params.get("llmConfig") or {})
+        llm_config["mode"] = mode
+        llm_config["maxIterations"] = max_iterations
+
+        # Create agent config from frontend format
+        config = config_from_dict(llm_config)
+
+        # Create agent runner
+        try:
+            self._agent_runner = factory.create_agent(config)
+        except Exception as e:
+            logger.error(f"Failed to create agent: {e}")
+            await self._send_error(msg_id, -32000, f"Failed to create agent: {e}")
+            return
+
+        # Run agent in background
+        self._agent_task = asyncio.create_task(self._run_agent(msg_id, prompt))
+
+    async def _run_agent(self, msg_id: Optional[str], prompt: str) -> None:
+        """
+        Run the agent and handle completion.
+
+        HITL 인터럽트 발생 시:
+        1. 에이전트가 GraphInterrupt 발생
+        2. 인터럽트 데이터를 프론트엔드에 hitl_request 이벤트로 전송
+        3. 사용자 응답을 기다린 후 Command(resume=...) 로 그래프 재개
+
+        Args:
+            msg_id: Request ID
+            prompt: User prompt
+        """
+        try:
+            result = await self._agent_runner.run(prompt)
+
+            # Send success response if we have a message ID
+            if msg_id:
+                await self._send_response(
+                    msg_id,
+                    {
+                        "status": "complete",
+                        "response": result,
+                        "iterations": self._agent_runner.iteration,
+                    },
+                )
+
+        except asyncio.CancelledError:
+            logger.info("Agent task cancelled")
+            if msg_id:
+                await self._send_error(msg_id, -32000, "Agent cancelled")
+
+        except Exception as e:
+            # GraphInterrupt는 HITL 인터럽트를 의미
+            if _is_graph_interrupt(e):
+                await self._handle_graph_interrupt(msg_id, e)
+            elif isinstance(e, AgentError):
+                logger.error(f"Agent error [{e.error_type}]: {e}", exc_info=e.cause)
+                if msg_id:
+                    await self._send_error(msg_id, -32000, str(e))
+            else:
+                logger.error(f"Unexpected agent error: {e}", exc_info=True)
+                if msg_id:
+                    await self._send_error(msg_id, -32000, str(e))
+
+    async def _handle_graph_interrupt(
+        self, msg_id: Optional[str], exc: Exception
+    ) -> None:
+        """
+        HITL 그래프 인터럽트 처리.
+
+        인터럽트 데이터를 파싱하여 프론트엔드에 전송하고,
+        사용자 응답을 기다림.
+
+        Args:
+            msg_id: Original request ID
+            exc: GraphInterrupt exception
+        """
+        # 인터럽트 데이터 추출
+        interrupt_data = _extract_interrupt_data(exc)
+
+        # 프론트엔드에 HITL 요청 전송
+        await self._send_event(
+            "hitl_request",
+            {
+                "action_requests": interrupt_data.get("action_requests", []),
+                "review_configs": interrupt_data.get("review_configs", []),
+            },
+        )
+
+        # 사용자 응답 대기 (Future 생성)
+        self._hitl_future = asyncio.Future()
+
+        try:
+            hitl_response = await asyncio.wait_for(
+                self._hitl_future, timeout=300
+            )  # 5 min timeout
+        except asyncio.TimeoutError:
+            logger.warning("HITL request timed out")
+            return
+        except asyncio.CancelledError:
+            return
+        finally:
+            self._hitl_future = None
+
+        # 사용자 응답으로 그래프 재개
+        try:
+            config = {"configurable": {"thread_id": self._agent_runner.thread_id}}
+
+            result = ""
+            async for event in self._agent_runner.agent.astream_events(
+                Command(resume=hitl_response),
+                config=config,
+                version="v2",
+            ):
+                if self._is_closed:
+                    return
+
+                event_kind = event.get("event", "")
+
+                if event_kind == "on_tool_start":
+                    tool_name = event.get("name", "unknown")
+                    tool_input = event.get("data", {}).get("input", {})
+                    await self._send_event(
+                        "tool_call",
+                        {
+                            "id": event.get("run_id", ""),
+                            "name": tool_name,
+                            "arguments": tool_input,
+                        },
+                    )
+
+                elif event_kind == "on_tool_end":
+                    tool_name = event.get("name", "unknown")
+                    tool_output = event.get("data", {}).get("output", "")
+                    await self._send_event(
+                        "tool_result",
+                        {
+                            "id": event.get("run_id", ""),
+                            "name": tool_name,
+                            "result": {"success": True, "data": tool_output},
+                        },
+                    )
+
+                elif event_kind == "on_chat_model_stream":
+                    chunk = event.get("data", {}).get("chunk", None)
+                    if chunk and hasattr(chunk, "content") and chunk.content:
+                        result += chunk.content
+
+            if msg_id:
+                await self._send_response(
+                    msg_id,
+                    {
+                        "status": "complete",
+                        "response": result,
+                        "iterations": self._agent_runner.iteration,
+                    },
+                )
+
+        except Exception as resume_exc:
+            if _is_graph_interrupt(resume_exc):
+                # 재귀적으로 또 다른 인터럽트 처리
+                await self._handle_graph_interrupt(msg_id, resume_exc)
+            else:
+                logger.error(f"Agent resume error: {resume_exc}")
+                if msg_id:
+                    await self._send_error(msg_id, -32000, str(resume_exc))
+
+    async def _handle_cancel(
+        self, msg_id: Optional[str], params: dict[str, Any]
+    ) -> None:
+        """
+        Handle agent/cancel request.
+
+        Args:
+            msg_id: Request ID
+            params: Request parameters
+        """
+        if self._agent_runner:
+            self._agent_runner.cancel()
+
+        if self._agent_task and not self._agent_task.done():
+            self._agent_task.cancel()
+            await self._send_response(msg_id, {"status": "cancelled"})
+        else:
+            await self._send_response(msg_id, {"status": "not_running"})
+
+    async def _handle_hitl_response(
+        self,
+        msg_id: Optional[str],
+        params: dict[str, Any],
+    ) -> None:
+        """
+        Handle agent/hitl_response request.
+
+        프론트엔드에서 HITL 승인/거부 응답을 수신.
+        decisions 형식: [{"type": "approve"} | {"type": "reject", "message": "..."} | ...]
+
+        Args:
+            msg_id: Request ID
+            params: Request parameters containing decisions
+        """
+        decisions = params.get("decisions", [])
+
+        if self._hitl_future and not self._hitl_future.done():
+            # HITLResponse 형식으로 전달
+            self._hitl_future.set_result({"decisions": decisions})
+            await self._send_response(msg_id, {"status": "received"})
+        else:
+            await self._send_error(msg_id, -32000, "No pending HITL request")
+
+    async def _on_agent_event(self, event: AgentEvent) -> None:
+        """
+        Callback for agent events.
+
+        Args:
+            event: Agent event
+        """
+        if self._is_closed:
+            return
+        await self._send_event(event.type, event.data)
+
+    async def _send_event(self, event_type: str, data: dict[str, Any]) -> None:
+        """
+        Send a JSON-RPC notification (no id).
+
+        Args:
+            event_type: Event type
+            data: Event data
+        """
+        if self._is_closed:
+            return
+
+        message = {
+            "jsonrpc": "2.0",
+            "method": f"agent/{event_type}",
+            "params": data,
+        }
+
+        try:
+            await self.write_message(json.dumps(message))
+        except Exception as e:
+            logger.error(f"Error sending event: {e}")
+
+    async def _send_response(
+        self,
+        msg_id: Optional[str],
+        result: Any,
+    ) -> None:
+        """
+        Send a JSON-RPC response.
+
+        Args:
+            msg_id: Request ID
+            result: Result data
+        """
+        if self._is_closed:
+            return
+
+        message = {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "result": result,
+        }
+
+        try:
+            await self.write_message(json.dumps(message))
+        except Exception as e:
+            logger.error(f"Error sending response: {e}")
+
+    async def _send_error(
+        self,
+        msg_id: Optional[str],
+        code: int,
+        message: str,
+    ) -> None:
+        """
+        Send a JSON-RPC error response.
+
+        Args:
+            msg_id: Request ID
+            code: Error code
+            message: Error message
+        """
+        if self._is_closed:
+            return
+
+        error_msg = {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "error": {
+                "code": code,
+                "message": message,
+            },
+        }
+
+        try:
+            await self.write_message(json.dumps(error_msg))
+        except Exception as e:
+            logger.error(f"Error sending error: {e}")
+
+
+# ---------------------------------------------------------------------------
+# GraphInterrupt helpers
+# ---------------------------------------------------------------------------
+
+
+# 브라우저가 보낸 llmConfig 에서 받는 키 — 개인설정뿐이다 (이슈 #113).
+# temperature 는 0.0 고정이라 받지 않는다 (이슈 #171).
+_CLIENT_LLM_PREFS = ("autoApprove",)
+
+
+def _turn_llm_config(client_prefs: dict) -> dict:
+    """한 턴의 LLM 설정. 접속·조직 값은 서버가 해석하고 브라우저 값은 개인설정만 얹는다.
+
+    접속 설정은 centri > 시크릿 env > pod 파일이다 (이슈 #113). 브라우저가 endpoint·apiKey·
+    prompts 를 보내도 반영하지 않는다.
+    """
+    from codebuff_core.managers.config_manager import get_config_manager
+
+    llm_config = get_config_manager().get_config()
+    llm_config.update({k: v for k, v in client_prefs.items() if k in _CLIENT_LLM_PREFS})
+    return llm_config
+
+
+def _is_graph_interrupt(exc: Exception) -> bool:
+    """GraphInterrupt 예외인지 확인."""
+    # langgraph의 GraphInterrupt 또는 NodeInterrupt 확인
+    exc_type = type(exc).__name__
+    return exc_type in ("GraphInterrupt", "NodeInterrupt")
+
+
+def _extract_interrupt_data(exc: Exception) -> dict[str, Any]:
+    """GraphInterrupt에서 HITL 데이터 추출."""
+    # GraphInterrupt.value에 HITLRequest 데이터가 있음
+    if hasattr(exc, "value"):
+        value = exc.value
+        if isinstance(value, dict):
+            return value
+        # interrupt() 결과가 리스트인 경우
+        if isinstance(value, (list, tuple)) and value:
+            first = value[0]
+            if isinstance(first, dict):
+                return first
+    return {"action_requests": [], "review_configs": []}