cfc/运营文案/opencode_client.py

# -*- coding: utf-8 -*-
"""
opencode_client.py — 通过本机常驻 opencode 服务的 HTTP API 进行会话交互

负责：
    1. 认证（basic auth）
    2. chat_id → session_id 映射的持久化（JSON 文件，每用户/群一个 opencode 会话）
    3. 确保某 chat 的会话存在（build agent）
    4. 发消息 + 轮询取 assistant 最终文本回复

权限策略（最安全）：
    - 只为会话启用「只读」类工具（读文件、浏览、文本生成）
    - 禁用运行命令 / 写文件等需要授权的操作
    - 因此 opencode 只能做无副作用的分析、生成回复，无法改动本机文件
"""
import io
import json
import os
import sys
import time
import urllib.parse
import urllib.request
import base64

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
OPERATION_DIR = SCRIPT_DIR
CFC_ROOT = os.path.dirname(OPERATION_DIR)

# opencode 常驻服务
OPCODE_BASE = "http://127.0.0.1:4090"
OPCODE_USER = "opencode"
OPCODE_PASSWORD = os.environ.get("OPENCODE_SERVER_PASSWORD", "IwinTrue@123")

# 会话映射文件（chat_id -> session_id）
SESSION_MAP_PATH = os.path.join(OPERATION_DIR, "opencode_sessions.json")

# opencode 会话工作目录（项目根，含素材库/脚本，open code 可直接只读分析）
WORK_DIR = CFC_ROOT

# build agent 默认工作于该 project 的模型取服务端默认
AGENT = "build"

# 只读工具白名单：其余工具（bash 执行、写文件等）一律禁用
# 此处列出 opencode 内置常用只读工具；未列出的默认不可用（服务端按 tools 参数过滤）
READONLY_TOOLS = []

# 危险的需授权工具，显式禁用（防御性，即便服务端宽松也不给）
FORBIDDEN_TOOLS = [
    "bash", "shell", "write", "edit", "apply_patch",
    "task", "dispatch", "webfetch_post", "chrome_launch",
]


def _auth_header() -> str:
    token = base64.b64encode(
        f"{OPCODE_USER}:{OPCODE_PASSWORD}".encode("utf-8")
    ).decode("ascii")
    return f"Basic {token}"


def _api(method: str, path: str, body=None, timeout: float = 60):
    """请求 opencode HTTP API，返回 (status, decoded_body)。"""
    url = OPCODE_BASE + path
    data = None
    if body is not None:
        data = json.dumps(body, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("Authorization", _auth_header())
    if body is not None:
        req.add_header("Content-Type", "application/json; charset=utf-8")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read().decode("utf-8", errors="replace")
        return resp.status, raw


def _load_session_map() -> dict:
    if not os.path.exists(SESSION_MAP_PATH):
        return {}
    try:
        with open(SESSION_MAP_PATH, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return {}


def _save_session_map(mapping: dict):
    with open(SESSION_MAP_PATH, "w", encoding="utf-8") as f:
        json.dump(mapping, f, ensure_ascii=False, indent=2)


def _create_session(chat_id: str) -> str:
    """为某 chat 新建一个 build agent 会话，返回 session_id。"""
    dir_q = urllib.parse.quote(WORK_DIR)
    body = {
        "title": f"xhs-bot-{chat_id[:24]}",
        "agent": AGENT,
    }
    st, raw = _api("POST", f"/session?directory={dir_q}", body=body, timeout=30)
    if st not in (200, 201):
        raise RuntimeError(f"创建会话失败 status={st}: {raw[:300]}")
    data = json.loads(raw)
    sid = data.get("id")
    if not sid:
        raise RuntimeError(f"创建会话未返回 session_id: {raw[:300]}")
    return sid


def ensure_session(chat_id: str) -> str:
    """确保某 chat 有对应 opencode 会话，返回 session_id。"""
    mapping = _load_session_map()
    sid = mapping.get(chat_id)
    if sid:
        return sid
    sid = _create_session(chat_id)
    mapping[chat_id] = sid
    _save_session_map(mapping)
    return sid


def send_prompt(session_id: str, text: str, timeout: float = 30):
    """向指定会话异步发送一条用户消息。返回 HTTP 状态。"""
    parts = [{"type": "text", "text": text}]
    body = {
        "parts": parts,
        # 只读约束：明确禁用需授权的执行/写入类工具
        "tools": {t: False for t in FORBIDDEN_TOOLS},
    }
    st, raw = _api("POST", f"/session/{session_id}/prompt_async",
                   body=body, timeout=timeout)
    return st, raw


def _extract_assistant_text(msgs) -> str:
    """从 v1 message 列表里提取最后一条 assistant 的最终可见文本。"""
    last = msgs[-1] if msgs else None
    if not last:
        return ""
    if last.get("info", {}).get("role") != "assistant":
        return ""
    texts = [p.get("text", "") for p in last.get("parts", [])
             if p.get("type") == "text" and p.get("text")]
    return texts[-1] if texts else ""


def wait_for_reply(session_id: str, timeout: float = 180,
                   poll_interval: float = 5.0) -> str:
    """等待会话最新 assistant 回复完成，返回最终文本。

    处理两种情形：
      1) assistant 回复极快，进入本函数时已生成 → 直接看最后一条 assistant 文本
      2) 正常流式：等待「发消息后新增」的 assistant 带文本出现
    """
    # 进入时先记录当前消息条数，作为"已有"基线
    try:
        _, raw0 = _api("GET", f"/session/{session_id}/message", timeout=30)
        before = len(json.loads(raw0)) if raw0 else 0
    except Exception:
        before = 0

    deadline = time.time() + timeout
    while time.time() < deadline:
        time.sleep(poll_interval)
        try:
            st, raw = _api("GET", f"/session/{session_id}/message", timeout=30)
            if st != 200 or not raw:
                continue
            msgs = json.loads(raw)
        except Exception:
            continue
        if not msgs:
            continue

        # 情况1：最后一条是 assistant 且有文本（覆盖超快回复：before 已含该条）
        last = msgs[-1]
        if last.get("info", {}).get("role") == "assistant":
            lt = [p.get("text", "") for p in last.get("parts", [])
                  if p.get("type") == "text" and (p.get("text") or "").strip()]
            if lt:
                return lt[-1]

        # 情况2：出现新增的 assistant 消息（条数 > 基线）
        if len(msgs) > before:
            for m in reversed(msgs):
                role = m.get("info", {}).get("role")
                if role == "assistant":
                    txts = [p.get("text", "") for p in m.get("parts", [])
                            if p.get("type") == "text" and (p.get("text") or "").strip()]
                    if txts:
                        return txts[-1]
    raise TimeoutError(f"等待 opencode 回复超时（{timeout}s）")


def ask(chat_id: str, text: str, timeout: float = 180) -> str:
    """高层封装：确保会话 → 发消息 → 等回复 → 返回最终文本。"""
    sid = ensure_session(chat_id)
    st, raw = send_prompt(sid, text)
    if st != 204:
        raise RuntimeError(f"发送到 opencode 失败 status={st}: {raw[:300]}")
    return wait_for_reply(sid, timeout=timeout)


if __name__ == "__main__":

    if not sys.stdout.isatty():
        try:
            sys.stdout.reconfigure(encoding='utf-8', errors='replace')
        except Exception:
            pass
    test_chat = "test-opencode-client"
    reply = ask(test_chat, "请用一句话回复：收到消息了吗？")
    print(f"REPLY: {reply}")