mao1/Rumeng/bridge/group_chat.py

"""Workgroup chat — lightweight multi-agent messaging.
Adapted from CcCompanion group_chat.py, stripped to core."""
from __future__ import annotations

import json
import os
import re
import secrets
import threading
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

DATA_DIR = Path(__file__).resolve().parent / "data"
GROUP_FILE = DATA_DIR / "group_messages.jsonl"

MENTION_RE = re.compile(r"@([A-Za-z0-9_\-]+|[一-鿿]+)")
ALL_TOKEN = "__all__"

ROSTER: list[dict[str, Any]] = [
    {
        "id": "amian",
        "display_name": "眠眠",
        "kind": "human",
        "avatar": "眠",
        "color": "neutral",
        "model": None,
        "tmux": None,
        "can_reply": False,
    },
    {
        "id": "shenyan",
        "display_name": "沈晏",
        "kind": "agent",
        "avatar": "沈",
        "color": "orange",
        "model": "DeepSeek",
        "tmux": "shenyan",
        "can_reply": True,
        "default_responder": True,
    },
    {
        "id": "zhangsan",
        "display_name": "张三",
        "kind": "agent",
        "avatar": "张",
        "color": "green",
        "model": "Qwen 35B",
        "tmux": None,
        "can_reply": True,
    },
]

ROSTER_BY_ID = {m["id"]: m for m in ROSTER}
REPLY_AGENT_IDS = [m["id"] for m in ROSTER if m.get("can_reply")]

MENTION_ALIASES: dict[str, str] = {
    "all": ALL_TOKEN,
    "__all__": ALL_TOKEN,
    "眠眠": "amian",
    "amian": "amian",
    "沈晏": "shenyan",
    "shenyan": "shenyan",
    "阿晏": "shenyan",
    "张三": "zhangsan",
    "zhangsan": "zhangsan",
    "zs": "zhangsan",
}


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds")


def _make_id(prefix: str) -> str:
    return f"{prefix}_{int(time.time() * 1000)}_{secrets.token_hex(4)}"


def _read_jsonl(path: Path) -> list[dict]:
    if not path.exists():
        return []
    rows: list[dict] = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    rows.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    return rows


def _append_jsonl(path: Path, record: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")


def normalize_mentions(mentions_raw: Any = None, text: str = "") -> list[str]:
    """Extract and resolve @mentions from explicit list + text body."""
    raw: list[str] = []
    if isinstance(mentions_raw, str):
        raw.extend(m.strip() for m in mentions_raw.split(",") if m.strip())
    elif isinstance(mentions_raw, list):
        raw.extend(str(m).strip() for m in mentions_raw if str(m).strip())
    if text:
        raw.extend(m.group(1).strip() for m in MENTION_RE.finditer(text))

    seen: set[str] = set()
    resolved: list[str] = []
    for item in raw:
        key = item.strip().lstrip("@").lower()
        agent_id = MENTION_ALIASES.get(key)
        if agent_id and agent_id not in seen:
            resolved.append(agent_id)
            seen.add(agent_id)
    return resolved


def targets_for(sender_id: str, mentions: list[str]) -> list[str]:
    """Decide which agents should receive a message."""
    if sender_id == "amian":
        # human: no @mention → broadcast to all agents
        if not mentions or ALL_TOKEN in mentions:
            return list(REPLY_AGENT_IDS)
        return [m for m in mentions if m in REPLY_AGENT_IDS]
    else:
        # agent → agent: only explicit @mentions
        if not mentions or ALL_TOKEN in mentions:
            return []
        return [m for m in mentions if m in REPLY_AGENT_IDS and m != sender_id]


def append_group_message(
    sender_id: str,
    text: str,
    *,
    mentions: list[str] | None = None,
    reply_to: str | None = None,
    message_type: str = "chat",
) -> dict:
    """Write a message to the group JSONL and return the record."""
    member = ROSTER_BY_ID.get(sender_id)
    if not member:
        raise ValueError(f"unknown sender_id: {sender_id}")
    text = str(text or "").strip()
    if not text:
        raise ValueError("text required")

    mentions = mentions or normalize_mentions(text=text)
    targets = targets_for(sender_id, mentions)

    record = {
        "id": _make_id("grp"),
        "ts": _now_iso(),
        "sender_id": sender_id,
        "text": text,
        "mentions": mentions,
        "targets": targets,
        "reply_to": reply_to,
        "message_type": message_type,
    }
    _append_jsonl(GROUP_FILE, record)
    return record


def read_since(since_ts: str | None = None, limit: int = 100) -> list[dict]:
    """Return messages newer than since_ts, up to limit."""
    rows = _read_jsonl(GROUP_FILE)
    if since_ts:
        rows = [r for r in rows if r.get("ts", "") > since_ts]
        return rows[:limit]
    return rows[-limit:]


def roster() -> list[dict]:
    return [dict(m) for m in ROSTER]


def agent_status(session_exists: Callable[[str], bool] | None = None) -> dict:
    """Return online/offline status for each agent."""
    agents: dict[str, Any] = {}
    for m in ROSTER:
        if m.get("kind") != "agent":
            continue
        aid = m["id"]
        tmux = m.get("tmux")
        online = bool(tmux and session_exists and session_exists(tmux))
        agents[aid] = {"state": "online" if online else "offline", "tmux": tmux}
    return {"agents": agents}


def context_lines(limit: int = 20) -> list[str]:
    """Last N messages formatted for fan-out context injection."""
    records = read_since(limit=limit)
    lines: list[str] = []
    for rec in records:
        sender_id = rec.get("sender_id", "")
        member = ROSTER_BY_ID.get(sender_id) or {}
        name = member.get("display_name") or sender_id
        ts = str(rec.get("ts", ""))[11:16]
        text = str(rec.get("text", "")).replace("\n", " ")
        if len(text) > 180:
            text = text[:177] + "..."
        lines.append(f"[{ts}] {name}: {text}")
        if len(lines) >= limit:
            break
    return lines