codex-app-server/cloud/pr_exec.py

"""跨主机 PR 执行（HA 阶段 1）。

背景：`/pr` 原来只在 server 进程内同步跑，worktree 从 server 本机
文件系统解析——server 与 worker 同主机是隐含前提。多机下任务的
worktree 可能在另一台主机上，server 本机根本没有那个目录。

本模块提供两条路（与 server 的 `/pr` 端点、worker 的 pr 任务循环
共用，逻辑只写一次）：

- 同步路（worktree 在本机）：`execute_task_pr` 直接跑，与旧端点
  行为一致（写回 pr_url/pr_number、记审计 pr_created、成功发
  review-requested、落盘失败发 merge-failed、跳过不发）。
- 队列路（worktree 在别机）：`enqueue_pr_job` 把任务推进 `pr_jobs`
  表；持有 worktree 的主机上的 local worker 每轮用 CAS 认领并跑
  同一个 `execute_task_pr`，结果同样写回任务行。server 与 worker
  之间不走新网络通道（不引 SSH 密钥分发），只复用已有的 DB 队列。

只有 local worker 跑 pr 任务：host-worktrees/ 与 tenants/ 在主机上
都可见（server 今天就是这么给 Docker 任务做 /pr 的），而容器内
worker 看不到 host-worktrees/，也不一定有 gh。Docker worker 跳过
pr 任务循环。
"""

import logging
import time
import uuid

from cloud.auth import AuthStore
from cloud.db import connect as _db_connect
from cloud.github_app import resolve_app_token
from cloud.pr import (PrSkipped, PrStepError, create_task_pr,
                      parse_git_remote_owner_repo)
from cloud.push import dispatch_event as _push_dispatch

_log = logging.getLogger("cloud.pr_exec")


def ensure_pr_jobs_table(conn):
    """幂等建 pr_jobs 表（workers/webhook_deliveries 同款懒建表模式）。

    PG 适配层把 `?` 译成 `%s`、把 AUTOINCREMENT 译成 IDENTITY，
    因此同一条 DDL 在 SQLite 与 Postgres 下都能跑；调用方在
    enqueue/claim 前各调一次即可。
    """
    conn.execute(
        "CREATE TABLE IF NOT EXISTS pr_jobs("
        "id TEXT PRIMARY KEY, task_id TEXT NOT NULL, "
        "project_id TEXT NOT NULL, host TEXT NOT NULL, "
        "status TEXT NOT NULL, pr_url TEXT, pr_number INTEGER, "
        "error TEXT, created_at REAL, updated_at REAL, "
        "started_at REAL, finished_at REAL)"
    )
    conn.commit()


def find_open_pr_job(conn, task_id):
    """找该任务未完成的 pr 任务（去重：同一任务同时只排一个）。"""
    ensure_pr_jobs_table(conn)
    row = conn.execute(
        "SELECT * FROM pr_jobs WHERE task_id = ?"
        " AND status IN ('queued', 'running')"
        " ORDER BY created_at ASC LIMIT 1",
        (task_id,)).fetchone()
    if row is None:
        return None
    try:
        return dict(row)
    except TypeError:
        return {"id": row[0], "task_id": row[1], "project_id": row[2],
                "host": row[3], "status": row[4]}


def enqueue_pr_job(conn, task_id, project_id, host):
    """排一个 pr 任务（调用方已做去重）。返回 job id。"""
    ensure_pr_jobs_table(conn)
    job_id = uuid.uuid4().hex
    now = time.time()
    conn.execute(
        "INSERT INTO pr_jobs(id, task_id, project_id, host, status,"
        " created_at, updated_at) VALUES(?, ?, ?, ?, 'queued', ?, ?)",
        (job_id, task_id, project_id, host, now, now))
    conn.commit()
    return job_id


def claim_pr_job(conn, host):
    """CAS 认领本主机的一个 queued pr 任务（最老优先）。

    两个同主机 worker 抢同一任务时只有一个能拿到 rowcount == 1。
    返回任务 dict 或 None。
    """
    ensure_pr_jobs_table(conn)
    row = conn.execute(
        "SELECT * FROM pr_jobs WHERE host = ? AND status = 'queued'"
        " ORDER BY created_at ASC LIMIT 1",
        (host,)).fetchone()
    if row is None:
        return None
    job_id = row["id"] if hasattr(row, "__getitem__") else row[0]
    now = time.time()
    cur = conn.execute(
        "UPDATE pr_jobs SET status = 'running', started_at = ?,"
        " updated_at = ? WHERE id = ? AND status = 'queued'",
        (now, now, job_id))
    conn.commit()
    if cur.rowcount != 1:
        return None
    try:
        return dict(row)
    except TypeError:
        return {"id": row[0], "task_id": row[1], "project_id": row[2],
                "host": row[3], "status": "running"}


def finish_pr_job(conn, job_id, status, pr_url=None, pr_number=None,
                  error=None):
    """写回 pr 任务终态（completed/failed）。"""
    ensure_pr_jobs_table(conn)
    now = time.time()
    conn.execute(
        "UPDATE pr_jobs SET status = ?, pr_url = ?, pr_number = ?,"
        " error = ?, finished_at = ?, updated_at = ? WHERE id = ?",
        (status, pr_url, pr_number, error, now, now, job_id))
    conn.commit()


def _notify(project_id, event_type, payload, db_path):
    """PR 事件推送（best-effort，与 server._notify_pr_event 同语义）。"""
    try:
        result = _push_dispatch(db_path, event_type, project_id, payload)
    except Exception as exc:  # noqa: BLE001 -- 通知失败绝不炸 PR
        _log.warning("pr notify %s failed: %s", event_type, exc)
        return
    if result.get("errors"):
        _log.warning("pr notify %s errors: %s", event_type, result["errors"])


def execute_task_pr(db_path, task_id, user_id, worktrees_root=None,
                    git_runner=None, gh_runner=None, gh_which=None,
                    app_token_client=None, github_api=None):
    """给已 merged 的任务建 PR（同步执行体，server 端点与 worker 共用）。

    步骤与旧 `_task_pr` 端点内联的逻辑一致：查项目 auto_pr/repo_url、
    有 App 安装记录时现换 installation token 走 App 路径、否则走
    gh + PAT；成功写回任务 pr_url/pr_number、记审计 pr_created、
    发 review-requested；PrStepError 时发 merge-failed 后原样抛出；
    PrSkipped 时不发任何事件、原样抛出。调用方按旧端点语义映射
    响应（PrSkipped -> 200 pr_skipped，PrStepError -> 502）。

    user_id：审计 pr_created 的操作人；server 端点传调用方，
    worker 的 pr 任务循环传 "system"（与 sweep 同款约定）。
    token 只在本次调用内存里出现，不进库、不进日志。
    """
    conn = _db_connect(str(db_path))
    try:
        row = conn.execute(
            "SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
        if row is None:
            raise PrSkipped("task_not_found")
        task = dict(row)
        prow = conn.execute(
            "SELECT auto_pr, repo_url FROM projects WHERE id = ?",
            (task["project_id"],)).fetchone()
    finally:
        conn.close()
    if task["status"] != "merged":
        raise PrSkipped("not_merged")
    auto_pr = 0
    repo_url = None
    if prow is not None:
        try:
            auto_pr = prow["auto_pr"] or 0
        except (KeyError, IndexError, TypeError):
            auto_pr = 0
        try:
            repo_url = prow["repo_url"]
        except (KeyError, IndexError, TypeError):
            repo_url = None
    app_token = None
    auth_via = "pat"
    repo_full_name = None
    if repo_url:
        parsed = parse_git_remote_owner_repo(repo_url)
        if parsed:
            repo_full_name = "%s/%s" % parsed
    if repo_full_name:
        conn = _db_connect(str(db_path))
        try:
            app_token = resolve_app_token(
                conn, repo_full_name, http_client=app_token_client)
        except Exception:
            app_token = None
        finally:
            conn.close()
        if app_token:
            auth_via = "app"
    try:
        result = create_task_pr(
            task_id, auto_pr,
            worktrees_root=worktrees_root,
            git_runner=git_runner, gh_runner=gh_runner, gh_which=gh_which,
            app_token=app_token, github_api=github_api)
    except PrSkipped:
        raise
    except PrStepError as exc:
        _notify(task["project_id"], "merge-failed",
                {"task_id": task_id, "step": exc.step,
                 "error": exc.message}, db_path)
        raise
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        conn.execute(
            "UPDATE tasks SET pr_url = ?, pr_number = ?,"
            " updated_at = ? WHERE id = ?",
            (result["pr_url"], result["pr_number"], now, task_id))
        conn.commit()
    finally:
        conn.close()
    store = AuthStore(db_path)
    try:
        store.audit(user_id, task["project_id"], "pr_created",
                    result["pr_url"])
    finally:
        try:
            store.conn.close()
        except Exception:
            pass
    _notify(task["project_id"], "review-requested",
            {"task_id": task_id, "pr_url": result["pr_url"],
             "pr_number": result["pr_number"]}, db_path)
    return {"pr_url": result["pr_url"], "pr_number": result["pr_number"],
            "branch": result["branch"], "auth_via": auth_via}