codebuddy_client.py

"""
CodeBuddy 上游 HTTP 客户端 - 封装对 CodeBuddy CLI --serve API 的异步调用。

提供:
- post_run()     发起 Agent 执行，返回 runId
- stream_run()   SSE 流式订阅执行结果
- health_check()  检查 CodeBuddy 服务可用性
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from functools import lru_cache
from typing import AsyncGenerator, Optional

import httpx

from config import Settings, get_settings

logger = logging.getLogger(__name__)


class CodeBuddyClientError(Exception):
    """CodeBuddy API 调用异常"""

    def __init__(self, message: str, status_code: Optional[int] = None, detail: Optional[dict] = None):
        super().__init__(message)
        self.status_code = status_code
        self.detail = detail or {}


class CodeBuddyClient:
    """CodeBuddy CLI --serve HTTP API 异步客户端

    封装 run 创建、SSE 流订阅和健康检查，统一管理 httpx 连接池。
    通过 FastAPI 依赖注入以单例形式复用以避免频繁创建连接。

    用法:
        client = CodeBuddyClient(settings)
        run_id = await client.post_run("帮我写一个函数")
        async for line in client.stream_run(run_id):
            print(line)
    """

    def __init__(self, settings: Settings):
        self._base = settings.codebuddy_base_url.rstrip("/")
        self._header_val = settings.codebuddy_request_header
        self._timeout = settings.request_timeout
        self._sse_idle_timeout = settings.sse_idle_timeout
        self._sender_id = settings.codebuddy_sender_id
        self._sender_name = settings.codebuddy_sender_name
        self._settings = settings
        self._client: Optional[httpx.AsyncClient] = None

    @property
    def client(self) -> httpx.AsyncClient:
        """延迟初始化 httpx 客户端（首次调用时创建，复用连接池）"""
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self._base,
                timeout=httpx.Timeout(
                    connect=self._settings.httpx_connect_timeout,
                    read=self._timeout,
                    write=self._settings.httpx_write_timeout,
                    pool=self._settings.httpx_pool_timeout,
                ),
                headers={
                    "X-CodeBuddy-Request": self._header_val,
                    "Accept": "application/json, text/event-stream",
                },
            )
        return self._client

    async def close(self) -> None:
        """关闭底层 HTTP 客户端，释放连接"""
        if self._client is not None and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    # ---- 公共 API ----

    async def health_check(self) -> dict:
        """GET /api/v1/health - 检查 CodeBuddy 服务是否就绪

        Returns:
            health 端点返回的完整 data 对象

        Raises:
            CodeBuddyClientError: 连接失败或返回非 2xx
        """
        try:
            resp = await self.client.get("/api/v1/health")
            resp.raise_for_status()
            body = resp.json()
            return body.get("data", body)
        except httpx.ConnectError as e:
            raise CodeBuddyClientError(
                f"无法连接到 CodeBuddy ({self._base}): {e}",
                status_code=503,
            ) from e
        except httpx.TimeoutException as e:
            raise CodeBuddyClientError(
                f"CodeBuddy 请求超时: {e}",
                status_code=504,
            ) from e
        except httpx.HTTPStatusError as e:
            detail = self._try_extract_error(e.response)
            raise CodeBuddyClientError(
                f"CodeBuddy 返回错误 ({e.response.status_code}): {detail.get('message', str(e))}",
                status_code=e.response.status_code,
                detail=detail,
            ) from e

    async def post_run(
        self,
        text: str,
        sender_id: Optional[str] = None,
        sender_name: Optional[str] = None,
    ) -> str:
        """POST /api/v1/runs - 发起 Agent 执行

        Args:
            text: 发送给 CodeBuddy 的 prompt 文本
            sender_id: 发送者标识（默认使用配置值）
            sender_name: 发送者显示名称（默认使用配置值）

        Returns:
            runId 字符串，用于后续 SSE 流订阅

        Raises:
            CodeBuddyClientError: 请求失败
        """
        payload = {
            "id": f"msg_{uuid.uuid4().hex[:8]}",
            "type": "text",
            "text": text,
            "sender": {
                "id": sender_id or self._sender_id,
                "name": sender_name or self._sender_name,
            },
        }

        try:
            resp = await self.client.post("/api/v1/runs", json=payload)
            resp.raise_for_status()
            body = resp.json()
            data = body.get("data", body)
            run_id = data.get("runId")
            if not run_id:
                raise CodeBuddyClientError(
                    f"CodeBuddy 返回的数据中缺少 runId: {body}"
                )
            logger.info("CodeBuddy run created: runId=%s", run_id)
            return run_id
        except httpx.HTTPStatusError as e:
            detail = self._try_extract_error(e.response)
            raise CodeBuddyClientError(
                f"创建 CodeBuddy run 失败 ({e.response.status_code}): {detail.get('message', str(e))}",
                status_code=e.response.status_code,
                detail=detail,
            ) from e
        except (httpx.ConnectError, httpx.TimeoutException) as e:
            raise CodeBuddyClientError(
                f"无法连接到 CodeBuddy ({self._base}): {e}",
                status_code=503,
            ) from e

    async def stream_run(
        self,
        run_id: str,
    ) -> AsyncGenerator[str, None]:
        """GET /api/v1/runs/{run_id}/stream - SSE 流式订阅执行结果"""
        url = f"/api/v1/runs/{run_id}/stream"
        last_data_time = asyncio.get_event_loop().time()

        try:
            async with self.client.stream("GET", url) as resp:
                resp.raise_for_status()

                async for raw_line in resp.aiter_lines():
                    now = asyncio.get_event_loop().time()

                    if now - last_data_time > self._sse_idle_timeout:
                        logger.warning("SSE stream idle timeout exceeded for run %s", run_id)
                        break

                    line = raw_line.strip()

                    # SSE 注释行 - 当作心跳
                    if line.startswith(":"):
                        last_data_time = now
                        continue

                    # SSE event 行 - 检测 done 事件
                    if line.startswith("event:"):
                        event_type = line[6:].strip()
                        if event_type == "done":
                            logger.debug("SSE stream done event received for run %s", run_id)
                        last_data_time = now
                        continue

                    # SSE data 行
                    if line.startswith("data:"):
                        last_data_time = now
                        data = line[5:].strip()
                        if data and data != "{}":
                            yield data

        except httpx.HTTPStatusError as e:
            detail = self._try_extract_error(e.response)
            raise CodeBuddyClientError(
                f"订阅 CodeBuddy 流失败 ({e.response.status_code}): {detail.get('message', str(e))}",
                status_code=e.response.status_code,
                detail=detail,
            ) from e
        except (httpx.ConnectError, httpx.TimeoutException) as e:
            raise CodeBuddyClientError(
                f"CodeBuddy SSE 连接失败: {e}",
                status_code=503,
            ) from e
        except (httpx.ReadError, httpx.RemoteProtocolError) as e:
            # 服务器正常关闭连接（如发送 done 后关连接），当作流结束
            logger.debug("SSE stream ended (connection closed): run=%s reason=%s", run_id, e)
        except Exception as e:
            # 其他异常也尝试容错，让上游流自然结束
            logger.debug("SSE stream ended with exception: run=%s reason=%s", run_id, e)

    async def cancel_run(self, run_id: str) -> None:
        """POST /api/v1/runs/{run_id}/cancel - 取消正在执行的 run"""
        try:
            resp = await self.client.post(f"/api/v1/runs/{run_id}/cancel")
            resp.raise_for_status()
            logger.info("CodeBuddy run cancelled: runId=%s", run_id)
        except httpx.HTTPStatusError as e:
            logger.warning("取消 run 返回非 2xx: runId=%s status=%s", run_id, e.response.status_code)

    # ---- 内部方法 ----

    @staticmethod
    def _try_extract_error(resp: Optional[httpx.Response]) -> dict:
        """尝试从响应体中提取错误信息"""
        if resp is None:
            return {}
        try:
            body = resp.json()
            return body.get("error", {})
        except Exception:
            return {"message": resp.text[:500] if resp.text else str(resp.status_code)}


# ---- 单例工厂（用于 FastAPI 依赖注入） ----


@lru_cache()
def get_client() -> CodeBuddyClient:
    """获取全局 CodeBuddyClient 单例（使用 lru_cache 替代模块级全局变量）

    FastAPI 依赖注入用法:
        from fastapi import Depends

        @app.get("/something")
        async def handler(client: CodeBuddyClient = Depends(get_client)):
            ...
    """
    return CodeBuddyClient(get_settings())


def reset_client() -> None:
    """重置客户端单例缓存（测试用）"""
    get_client.cache_clear()