Agent/im-client/client.py.bak

import asyncio
import json
import logging
import os
import tempfile
import uuid
from datetime import datetime
from pathlib import Path

import websockets
from filelock import FileLock

from protocol import IMMessage, IMResponse
from notifier import AgentNotifier, ConsoleNotifier

logging.basicConfig(level=logging.INFO, format="%(asctime)s [CLIENT:%(name)s] %(message)s")


class IMClient:
    """IM Client 长驻服务。

    通过 WebSocket 连接 Server，通过文件与 Agent 交互。

    文件约定 (data/{contact_id}/):
        chatbox.jsonl     — 所有消息历史（收发都记录）
        in_pending.json   — 收到的待处理消息 (JSON 数组)
        out_pending.jsonl — 发送失败的消息

    窗口模式 (window_mode=True):
        每次运行生成新的 chat_id，消息按 chat_id 隔离
        文件结构变为: data/{contact_id}/{chat_id}/...
    """

    def __init__(
        self,
        contact_id: str,
        server_url: str = "ws://localhost:8000",
        data_dir: str | None = None,
        notifier: AgentNotifier | None = None,
        notify_interval: float = 30.0,
        window_mode: bool = False,
        chat_id: str | None = None,
    ):
        self.contact_id = contact_id
        self.server_url = server_url
        self.notifier = notifier or ConsoleNotifier()
        self.notify_interval = notify_interval
        self.window_mode = window_mode

        # 窗口模式：生成或使用指定的 chat_id
        base_dir = Path(data_dir) if data_dir else Path("data") / contact_id
        if window_mode:
            self.chat_id = chat_id or datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:6]
            self.data_dir = base_dir / "windows" / self.chat_id
        else:
            self.chat_id = None
            self.data_dir = base_dir

        self.data_dir.mkdir(parents=True, exist_ok=True)

        self.chatbox_path = self.data_dir / "chatbox.jsonl"
        self.in_pending_path = self.data_dir / "in_pending.json"
        self.out_pending_path = self.data_dir / "out_pending.jsonl"

        # 文件锁
        self._in_pending_lock = FileLock(str(self.data_dir / ".in_pending.lock"))
        self._out_pending_lock = FileLock(str(self.data_dir / ".out_pending.lock"))
        self._chatbox_lock = FileLock(str(self.data_dir / ".chatbox.lock"))

        self.ws = None
        self.log = logging.getLogger(f"{contact_id}:{self.chat_id}" if self.chat_id else contact_id)
        self._send_queue = asyncio.Queue()

        # 初始化文件
        if not self.chatbox_path.exists():
            self.chatbox_path.write_text("")
        if not self.in_pending_path.exists():
            self.in_pending_path.write_text("[]")
        if not self.out_pending_path.exists():
            self.out_pending_path.write_text("")

    async def run(self):
        """启动 Client 服务，自动重连。"""
        while True:
            try:
                # 构造 WebSocket URL，带上 chat_id 参数
                chat_id_param = self.chat_id or "default"
                ws_url = f"{self.server_url}/ws?contact_id={self.contact_id}&chat_id={chat_id_param}"
                self.log.info(f"连接 {ws_url} ...")
                async with websockets.connect(ws_url) as ws:
                    self.ws = ws
                    self.log.info("已连接")
                    await asyncio.gather(
                        self._ws_listener(),
                        self._send_worker(),
                        self._pending_notifier(),
                    )
            except (websockets.ConnectionClosed, ConnectionRefusedError, OSError) as e:
                self.log.warning(f"连接断开: {e}, 5 秒后重连...")
                self.ws = None
                await asyncio.sleep(5)
            except asyncio.CancelledError:
                self.log.info("服务停止")
                break

    # ── 协程 1: WebSocket 收消息 ──

    async def _ws_listener(self):
        """监听 WebSocket，聊天消息写 in_pending 和 chatbox，回执打日志。"""
        async for raw in self.ws:
            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                self.log.warning(f"收到无效 JSON: {raw}")
                continue

            if "sender" in data and "receiver" in data:
                # 聊天消息
                self.log.info(f"收到消息: {data['sender']} -> {data['content'][:50]}")
                self._append_to_in_pending(data)
                self._append_to_chatbox(data)
            elif "status" in data:
                # 发送回执
                resp = IMResponse(**data)
                if resp.status == "success":
                    self.log.info(f"消息 {resp.msg_id} 发送成功")
                else:
                    self.log.warning(f"消息 {resp.msg_id} 发送失败: {resp.error}")

    # ── 协程 2: 发送队列处理 ──

    async def _send_worker(self):
        """从队列取消息并发送，失败则写入 out_pending。"""
        while True:
            msg_data = await self._send_queue.get()
            # 填充 sender_chat_id
            msg = IMMessage(
                sender=self.contact_id,
                sender_chat_id=self.chat_id or "default",
                **msg_data
            )
            try:
                await self.ws.send(msg.model_dump_json())
                self.log.info(f"发送消息: -> {msg.receiver}:{msg.receiver_chat_id or '*'}")
                # 记录到 chatbox
                self._append_to_chatbox(msg.model_dump())
            except Exception as e:
                self.log.error(f"发送失败: {e}")
                # 写入 out_pending
                self._append_to_out_pending(msg.model_dump())

    # ── 协程 3: 轮询 in_pending 通知 Agent ──

    async def _pending_notifier(self):
        """轮询 in_pending.json，有新消息就调通知回调。"""
        while True:
            pending = self._read_in_pending()
            if pending:
                senders = list(set(m.get("sender", "unknown") for m in pending))
                count = len(pending)
                try:
                    await self.notifier.notify(count=count, from_contacts=senders)
                except Exception as e:
                    self.log.error(f"通知回调异常: {e}")
            await asyncio.sleep(self.notify_interval)

    # ── 文件操作 (原子性) ──

    def _append_to_in_pending(self, msg: dict):
        """将收到的消息追加到 in_pending.json。"""
        with self._in_pending_lock:
            pending = self._load_json_array(self.in_pending_path)
            pending.append(msg)
            self._atomic_write_json(self.in_pending_path, pending)

    def _read_in_pending(self) -> list[dict]:
        """读取 in_pending.json (不清空)。"""
        with self._in_pending_lock:
            return self._load_json_array(self.in_pending_path)

    def _append_to_chatbox(self, msg: dict):
        """追加消息到 chatbox.jsonl。"""
        with self._chatbox_lock:
            with open(self.chatbox_path, "a", encoding="utf-8") as f:
                f.write(json.dumps(msg, ensure_ascii=False) + "\n")

    def _append_to_out_pending(self, msg: dict):
        """追加发送失败的消息到 out_pending.jsonl。"""
        with self._out_pending_lock:
            with open(self.out_pending_path, "a", encoding="utf-8") as f:
                f.write(json.dumps(msg, ensure_ascii=False) + "\n")

    # ── Agent 调用的工具方法 ──

    def read_pending(self) -> list[dict]:
        """Agent 读取 in_pending 中的消息，并清空。"""
        with self._in_pending_lock:
            pending = self._load_json_array(self.in_pending_path)
            if not pending:
                return []
            # 清空 in_pending
            self._atomic_write_json(self.in_pending_path, [])
            return pending

    def send_message(self, receiver: str, content: str, msg_type: str = "chat", receiver_chat_id: str | None = None):
        """Agent 调用：将消息放入发送队列。

        Args:
            receiver: 接收方 contact_id
            content: 消息内容
            msg_type: 消息类型
            receiver_chat_id: 接收方窗口 ID（指定则定向发送，否则广播给该 contact_id 的所有窗口）
        """
        msg_data = {
            "receiver": receiver,
            "content": content,
            "msg_type": msg_type,
            "receiver_chat_id": receiver_chat_id
        }
        self._send_queue.put_nowait(msg_data)

    # ── 工具方法 ──

    @staticmethod
    def _load_json_array(path: Path) -> list:
        if not path.exists():
            return []
        text = path.read_text(encoding="utf-8").strip()
        if not text:
            return []
        try:
            data = json.loads(text)
            return data if isinstance(data, list) else []
        except json.JSONDecodeError:
            return []

    @staticmethod
    def _atomic_write_json(path: Path, data):
        """原子写入：先写临时文件再 rename。"""
        tmp_fd, tmp_path = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp")
        try:
            with os.fdopen(tmp_fd, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            os.replace(tmp_path, str(path))
        except Exception:
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)
            raise