codex-app-server/cloud/peer_messages.py

"""任务间互发消息（peer messages）。

Muse 式通用互发消息在 codex-app-server 里的落地。两个发送方：

1. API：任何持有用户 key 的调用方（控制台、外部编排器）调
   POST /tasks/<目标>/messages，带 from_task_id，消息以发送
   任务的名义投递。
2. Agent：任务里跑的 agent 在回合输出里写标记块，worker 在
   回合结束后解析并投递——agent 不用联网、不用凭证，本地
   与 Docker 执行器都通（两者都经 sandbox 的结果解析）。

投递语义（核心约束：互发消息不中断任务）：

- 目标排队中 / 运行中：消息排成一条后续消息（followup），
  带 sender_task_id；由 worker 排空循环在轮次边界执行——
  direct 优先、队内 FIFO，与用户后续消息同款，当前轮绝不
  被打断。执行仍受目标任务的配额门约束（排空前查配额，
  超限时该条回 pending、任务走配额中断，与后续消息一致）。
- 目标已终态（completed / failed / interrupted）：消息复活
  它——把带发送方前缀的正文写进 resume_message 并把状态
  置回 queued，下一次 worker 认领时按续跑跑这一轮。这样
  A→B→A 的往返不受双方回合结束时序影响；没有复活，B 回
  给已结束的 A 的消息会直接丢失。
- 其它情形一律拒绝，不炸任务、不改状态：发给自己、跨项目、
  目标是 goal 任务（goal 循环自有轮次）、目标已合并、
  目标已有待续跑消息（复活位被占）、目标不存在。

发送方约束：

- API 路径由 server 鉴权：调用方须对目标项目有 create_task
  权限、对发送任务项目有 read 权限（防冒名）。
- Worker 路径的发送方就是当前任务本身；跨项目在投递层再
  挡一次（agent 猜到别项目任务 id 也发不出去）。
- 每轮最多投递 5 个标记块；只有 completed 轮次的标记块被
  投递（失败 / 超时 / 取消 / 配额中断的轮次不投递）；goal
  任务的回合不解析标记块（只收不发，见 README）。
- 标记块正文里不能含围栏——围栏是非贪婪匹配，正文里的
  围栏会把块截断成非法 JSON 而被跳过。

本模块只依赖标准库，server 与 worker 共用，不 import
cloud 其它模块，避免循环 import。
"""

import json
import re
import time
import uuid

MAX_PEER_MESSAGES_PER_TURN = 5
MAX_PEER_MESSAGE_BODY = 20000

# 复活可接的目标终态。merged / merge-failed 不复活（代码已
# 合并，续跑语义不明）；其它未知状态一律拒绝。
REVIVABLE_STATUSES = ("completed", "failed", "interrupted")

_FENCE_RE = re.compile(r"```task-message\s*\n(.*?)\n```", re.DOTALL)


def peer_message_prefix(from_task_id):
    """发送方前缀：排空轮次 prompt 与复活消息都带它，接收方
    agent 凭此知道回给哪个任务 id。"""
    return "[来自任务 %s 的消息]\n" % from_task_id


def row_get(row, name):
    """行取值容错：sqlite3.Row 缺列报 IndexError、dict 缺键
    报 KeyError、空行报 TypeError，一律回 None（worker 直连
    老库、列还没迁移时排空不能因此停摆）。"""
    try:
        return row[name]
    except (KeyError, IndexError, TypeError):
        return None


def extract_peer_messages(text):
    """从回合输出文本里提取互发消息标记块。

    块格式::

        ```task-message
        {"to": "<目标任务 id>", "body": "...", "mode": "queued"}
        ```

    mode 缺省 queued，非法值也回落 queued；to / body 缺失、
    非字符串、空串、正文超长、JSON 非法时整块跳过；每轮
    最多取 5 块，多余的忽略。"""
    if not text:
        return []
    out = []
    for match in _FENCE_RE.finditer(str(text)):
        if len(out) >= MAX_PEER_MESSAGES_PER_TURN:
            break
        try:
            payload = json.loads(match.group(1))
        except (TypeError, ValueError):
            continue
        if not isinstance(payload, dict):
            continue
        to_task_id = payload.get("to") or payload.get("to_task_id")
        body = payload.get("body")
        if not isinstance(to_task_id, str) or not to_task_id.strip():
            continue
        if not isinstance(body, str) or not body.strip():
            continue
        body = body.strip()
        if len(body) > MAX_PEER_MESSAGE_BODY:
            continue
        mode = payload.get("mode", "queued")
        if mode not in ("queued", "direct"):
            mode = "queued"
        out.append({"to_task_id": to_task_id.strip(),
                    "body": body, "mode": mode})
    return out


def ensure_followups_table(conn):
    """幂等建后续消息表并补 sender_task_id 列。

    server 的 _ensure_followups_table 委托到这里，worker 投递
    路径也调它——两边同一份 DDL，不漂移。列补齐写法与
    tasks 迁移同款：PRAGMA 查现有列，缺了再 ALTER（db 适配
    层把 PRAGMA 翻译成 information_schema 查询，Postgres 下
    同样生效）。"""
    conn.execute(
        "CREATE TABLE IF NOT EXISTS task_followups("
        "id TEXT PRIMARY KEY, "
        "task_id TEXT, "
        "project_id TEXT, "
        "body TEXT, "
        "mode TEXT, "
        "status TEXT, "
        "created_by TEXT, "
        "sender_task_id TEXT, "
        "created_at DOUBLE PRECISION, "
        "updated_at DOUBLE PRECISION)"
    )
    try:
        existing = {row[1] for row in conn.execute(
            "PRAGMA table_info(task_followups)").fetchall()}
    except Exception:
        existing = set()
    if "sender_task_id" not in existing:
        try:
            conn.execute(
                "ALTER TABLE task_followups"
                " ADD COLUMN sender_task_id TEXT")
        except Exception:
            pass
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_followups_task"
        " ON task_followups(task_id, status)")
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_followups_sender"
        " ON task_followups(sender_task_id)")


def deliver_peer_message(conn, from_task_id, to_task_id, body, mode="queued"):
    """投递一条互发消息。调用方自备连接，本函数在投递成功时
    commit，拒绝时不写任何东西。

    返回 (status, detail)：delivered 时 detail 为 followup id；
    revived 时 detail 为空；rejected 时 detail 为原因（直接
    可进 409 / 审计）。mode 非 queued/direct 时回落 queued。"""
    if not isinstance(body, str) or not body.strip():
        return "rejected", "body required"
    body = body.strip()
    if len(body) > MAX_PEER_MESSAGE_BODY:
        return "rejected", "body too long"
    if mode not in ("queued", "direct"):
        mode = "queued"
    if not from_task_id or not to_task_id:
        return "rejected", "task not found"
    if from_task_id == to_task_id:
        return "rejected", "cannot message self"
    sender = conn.execute(
        "SELECT * FROM tasks WHERE id = ?", (from_task_id,)).fetchone()
    if sender is None:
        return "rejected", "sender task not found"
    target = conn.execute(
        "SELECT * FROM tasks WHERE id = ?", (to_task_id,)).fetchone()
    if target is None:
        return "rejected", "target task not found"
    if row_get(sender, "project_id") != row_get(target, "project_id"):
        return "rejected", "cross-project messages not allowed"
    if row_get(target, "verify_command") or row_get(target, "goal_mode"):
        return "rejected", "goal tasks do not accept peer messages"
    status = row_get(target, "status")
    now = time.time()
    if status in ("queued", "running"):
        ensure_followups_table(conn)
        fu_id = uuid.uuid4().hex
        conn.execute(
            "INSERT INTO task_followups(id, task_id, project_id,"
            " body, mode, status, created_by, sender_task_id,"
            " created_at, updated_at)"
            " VALUES(?, ?, ?, ?, ?, 'pending', ?, ?, ?, ?)",
            (fu_id, to_task_id, row_get(target, "project_id"),
             body, mode, row_get(sender, "created_by"),
             from_task_id, now, now),
        )
        conn.commit()
        return "delivered", fu_id
    if status in REVIVABLE_STATUSES:
        cur = conn.execute(
            "UPDATE tasks SET status = 'queued',"
            " resume_message = ?, updated_at = ?"
            " WHERE id = ? AND status IN"
            " ('completed', 'failed', 'interrupted')"
            " AND resume_message IS NULL",
            (peer_message_prefix(from_task_id) + body, now, to_task_id))
        conn.commit()
        if cur.rowcount == 1:
            return "revived", ""
        return "rejected", "target has a pending resume message"
    return "rejected", "target task is not receivable: %s" % status