宜博 / lemongraph

"""c_work_node_run (03-C §3.6): work node execution + git worktree isolation.

Internal interface (not exposed over HTTP -- called by `c_work_graph_execute`
for each node needing isolated execution, per 03-C §3.6's own framing).
Validates input/output against a minimal inline JSON-Schema-lite check (no
third-party `jsonschema` dependency, per 05 §5.7 -- covers object/required,
which is all this codebase's schemas use so far), transitions
`work_node_execution.status` through
`pending -> running -> (reflecting -> retrying, up to a caller-specified cap)
-> done|blocked`, and handles `USAGE_LIMIT_REACHED` using the *exact*
detection regex 03-C §3.6 documents (`hit your (session|weekly|opus) limit`,
case-insensitive) -- more precise than the approximate phrase-matching
U-010/U-012 used before this Unit existed to give the behavior a real,
persisted home (`work_node_execution.quota_wait_started_at`/`quota_limit_type`).

U-013 scope matches REQ-03's boundary case, which is what this Unit's Done
criterion actually describes: two independent nodes running in separate git
worktrees, merged back without conflict. REQ-03's happy case (real
code_diff/test_report artifacts, unit+integration tests genuinely passing)
needs `c_artifact_write` (still unowned by any Unit, see U-011's note) and a
real Coder session actually writing code -- not attempted here, same
overreach pattern already corrected for U-009/U-010/U-011/U-012.

The "actual work" a node does is inherently node-specific (that's the whole
point of "give each node a contract") -- `run_fn` is an explicit, required
dependency (no hidden default calling out to `claude`), so this module
stays about the *contract* (schema, retries, worktree, logging, quota-wait),
not about any one node's content.
"""

import json
import os
import re
import subprocess
import time
import uuid
from datetime import datetime, timezone

from src import paths

_USAGE_LIMIT_RE = re.compile(r"hit your (session|weekly|opus) limit", re.IGNORECASE)
_BACKOFF_SCHEDULE_MINUTES = (5, 10, 20, 30)  # then repeats at 30 until the type's timeout
_TIMEOUT_SECONDS = {"session": 8 * 3600, "weekly": 9 * 24 * 3600}
_GIT_IDENTITY = ["-c", "user.name=graphlemon", "-c", "user.email=graphlemon@localhost"]
_VALID_WORK_GRAPH_RUN_STATUSES = {
    "planning", "running", "merging", "verifying", "reflecting", "blocked", "done", "failed",
}
_VALID_NODE_STATUSES = {
    "pending", "running", "reflecting", "retrying", "waiting_for_quota", "blocked", "done", "failed",
}


class SchemaValidationFailedError(Exception):
    """SCHEMA_VALIDATION_FAILED -- 03-C §3.6 (output validation; input
    validation failure raises InvalidInputError instead, see below -- 03-C
    describes input rejection as immediate/non-retryable, distinct from the
    retryable output-validation-failure case)."""


class InvalidInputError(Exception):
    """input_payload doesn't match input_schema_ref -- 03-C §3.6: "校验失败
    直接拒绝，不进入执行" (rejected outright, execution never starts)."""


class UsageLimitWaitTimeoutError(Exception):
    """p_usage_limit_wait exceeded its timeout (04-P §4.4.1). work_node_execution
    has no 'blocked_pending_human' status of its own (01-M 1.4's enum doesn't
    include it) -- this function sets status='blocked' and raises; bubbling to
    task.status='blocked_pending_human' is the caller's job, same layering as
    c_work_node_run itself not knowing about `task` at all (03-C §3.6)."""


class CorruptPersistedStateError(Exception):
    """REQ-06 failure case: persisted work_graph_run/work_node_execution state
    is internally inconsistent after a restart -- must error clearly, not
    silently continue on bogus state. No 03-C error code is named for this
    (it's not a normal state-machine transition), but REQ-06 requires exactly
    this behavior; callers should route it to task.status='blocked_pending_human'
    same as c_work_node_run itself doesn't touch `task` directly (03-C §3.6)."""


def _parse_sqlite_timestamp(text):
    """SQLite's CURRENT_TIMESTAMP is UTC 'YYYY-MM-DD HH:MM:SS' with no
    timezone marker -- parsed as UTC explicitly so it's comparable against
    time.time() (also UTC epoch seconds) across a process restart. A plain
    monotonic clock can't do this: it has no fixed relationship to wall-clock
    time and resets meaning across restarts, which is exactly REQ-06's
    scenario -- so `now`'s contract here is wall-clock, not monotonic."""
    return datetime.strptime(text, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc).timestamp()


def _validate_schema_lite(payload, schema):
    """object/required only -- everything this codebase's schemas use so far."""
    if schema.get("type") == "object":
        if not isinstance(payload, dict):
            return False
        for key in schema.get("required", []):
            if key not in payload:
                return False
    return True


def _quota_limit_type(text):
    match = _USAGE_LIMIT_RE.search(text or "")
    return match.group(1).lower() if match else None


def create_worktree(repo_path, node_exec_id):
    """`git worktree add` for an isolated code-writing node (01-M
    work_node_execution.git_worktree_ref)."""
    worktree_path = os.path.join(repo_path, ".graphlemon-worktrees", node_exec_id)
    branch = "graphlemon/" + node_exec_id
    subprocess.run(
        ["git", "worktree", "add", "-b", branch, worktree_path],
        cwd=repo_path, check=True, capture_output=True, text=True,
    )
    return worktree_path


def merge_worktree(repo_path, node_exec_id):
    """Merges the node's branch back into whichever branch is currently
    checked out in the main worktree at repo_path."""
    branch = "graphlemon/" + node_exec_id
    subprocess.run(
        _prefix_git(["merge", "--no-ff", branch, "-m", "graphlemon: merge " + node_exec_id]),
        cwd=repo_path, check=True, capture_output=True, text=True,
    )


def remove_worktree(repo_path, worktree_path):
    subprocess.run(
        ["git", "worktree", "remove", "--force", worktree_path],
        cwd=repo_path, check=True, capture_output=True, text=True,
    )


def _prefix_git(args):
    return ["git"] + _GIT_IDENTITY + args


def run_work_node(
    conn,
    node_exec_id,
    input_payload,
    input_schema,
    output_schema,
    run_fn,
    max_retries=2,
    needs_worktree=False,
    repo_path=None,
    sleep=time.sleep,
    now=time.time,
    opus_downgrade_model="claude-sonnet-5",
):
    """Runs one Work Node to contract. `run_fn(cwd, model) -> output_payload`
    is the caller-supplied actual work (a real `claude -p` call in
    production, a fake in tests); `model` is None unless this call is an
    opus-quota-triggered retry, in which case it's `opus_downgrade_model`.
    Returns the validated output_payload.

    REQ-06 (breakpoint recovery): safe to call again after a process restart
    for a node_exec_id that was already mid-flight. A node already 'done'
    is never re-run (its stored output_payload is returned as-is); a node
    interrupted mid-'retrying' or mid-'waiting_for_quota' resumes from its
    persisted retry_count / quota_wait_started_at rather than restarting
    those counters at zero -- `now` must be wall-clock (default time.time)
    for the quota-wait resume math to mean anything across a restart.

    Raises InvalidInputError (input schema), UsageLimitWaitTimeoutError
    (p_usage_limit_wait exhausted), or leaves work_node_execution.status
    ='blocked' with a SchemaValidationFailedError raised once max_retries is
    exceeded on output validation.
    """
    if not _validate_schema_lite(input_payload, input_schema):
        raise InvalidInputError("input_payload does not match input_schema_ref for " + node_exec_id)

    existing = conn.execute(
        "SELECT p.project_id, wgr.task_id, wne.status, wne.retry_count, wne.output_payload, "
        "wne.quota_wait_started_at, wne.quota_limit_type "
        "FROM work_node_execution wne "
        "JOIN work_graph_run wgr ON wgr.run_id = wne.run_id "
        "JOIN task t ON t.task_id = wgr.task_id "
        "JOIN project p ON p.project_id = t.project_id "
        "WHERE wne.node_exec_id = ?",
        (node_exec_id,),
    ).fetchone()
    if existing is None:
        raise CorruptPersistedStateError("no such work_node_execution: " + node_exec_id)
    (
        project_id, task_id, existing_status, existing_retry_count, existing_output_payload,
        existing_quota_wait_started_at, existing_quota_limit_type,
    ) = existing

    if existing_status == "done":
        # REQ-06 happy: never re-run an already-completed node.
        if existing_output_payload is None:
            raise CorruptPersistedStateError(
                "work_node_execution {} is status=done but has no output_payload".format(node_exec_id)
            )
        return json.loads(existing_output_payload)
    if existing_status not in _VALID_NODE_STATUSES:
        raise CorruptPersistedStateError(
            "work_node_execution {} has invalid status {!r}".format(node_exec_id, existing_status)
        )
    if existing_status == "waiting_for_quota" and not existing_quota_wait_started_at:
        raise CorruptPersistedStateError(
            "work_node_execution {} is status=waiting_for_quota but has no quota_wait_started_at".format(node_exec_id)
        )

    log_file_id = str(uuid.uuid4())
    log_path = paths.log_file_path(project_id, task_id, node_exec_id, node_exec_id)
    conn.execute(
        "INSERT INTO execution_log_file (log_file_id, node_exec_id, project_id, file_path, format, status) "
        "VALUES (?, ?, ?, ?, 'jsonl', 'active')",
        (log_file_id, node_exec_id, project_id, log_path),
    )
    conn.execute("UPDATE work_node_execution SET status = 'running' WHERE node_exec_id = ?", (node_exec_id,))
    conn.commit()

    worktree_path = None
    try:
        if needs_worktree:
            existing_worktree_ref = conn.execute(
                "SELECT git_worktree_ref FROM work_node_execution WHERE node_exec_id = ?", (node_exec_id,)
            ).fetchone()[0]
            # REQ-06: resuming a node that already had a worktree from before
            # the restart must reuse it, not `git worktree add` the same
            # branch/path again (which would just fail).
            if existing_worktree_ref and os.path.isdir(existing_worktree_ref):
                worktree_path = existing_worktree_ref
            else:
                worktree_path = create_worktree(repo_path, node_exec_id)
                conn.execute(
                    "UPDATE work_node_execution SET git_worktree_ref = ? WHERE node_exec_id = ?",
                    (worktree_path, node_exec_id),
                )
                conn.commit()
        cwd = worktree_path or repo_path

        retry_count = existing_retry_count or 0
        quota_wait_start = (
            _parse_sqlite_timestamp(existing_quota_wait_started_at)
            if existing_status == "waiting_for_quota"
            else None
        )
        current_model = None
        opus_already_downgraded = False
        while True:
            try:
                output_payload = run_fn(cwd, model=current_model)
            except _UsageLimitSignal as signal:
                quota_type = signal.quota_limit_type
                if quota_wait_start is None:
                    quota_wait_start = now()
                    conn.execute(
                        "UPDATE work_node_execution SET status = 'waiting_for_quota', "
                        "quota_wait_started_at = CURRENT_TIMESTAMP, quota_limit_type = ? WHERE node_exec_id = ?",
                        (quota_type, node_exec_id),
                    )
                    conn.execute(
                        "INSERT INTO execution_event (event_id, project_id, task_id, node_exec_id, event_type, payload) "
                        "VALUES (?, ?, ?, ?, 'quota_wait_started', ?)",
                        (str(uuid.uuid4()), project_id, task_id, node_exec_id, json.dumps({"quota_limit_type": quota_type})),
                    )
                    conn.commit()
                if quota_type == "opus":
                    # 04-P §4.4.1: "不用等待退避——直接把该节点这次调用降级到非
                    # Opus 模型重试", not just detect-and-fail. Only one
                    # downgrade attempt -- if the downgraded model *also*
                    # reports an opus-quota signal, something's wrong (a
                    # non-opus model shouldn't hit an opus-specific quota),
                    # so this doesn't loop forever chasing it.
                    if opus_already_downgraded:
                        raise UsageLimitWaitTimeoutError(
                            "opus quota signal persisted for {} even after downgrading to {}".format(
                                node_exec_id, opus_downgrade_model
                            )
                        ) from signal
                    current_model = opus_downgrade_model
                    opus_already_downgraded = True
                    continue
                elapsed = now() - quota_wait_start
                if elapsed >= _TIMEOUT_SECONDS[quota_type]:
                    conn.execute(
                        "UPDATE work_node_execution SET status = 'blocked' WHERE node_exec_id = ?", (node_exec_id,)
                    )
                    conn.commit()
                    raise UsageLimitWaitTimeoutError(
                        "{} quota wait exceeded timeout for {}".format(quota_type, node_exec_id)
                    ) from signal
                backoff_index = min(int(elapsed // (30 * 60)), len(_BACKOFF_SCHEDULE_MINUTES) - 1)
                sleep(_BACKOFF_SCHEDULE_MINUTES[backoff_index] * 60)
                continue

            if _validate_schema_lite(output_payload, output_schema):
                conn.execute(
                    "UPDATE work_node_execution SET status = 'done', output_payload = ?, retry_count = ? "
                    "WHERE node_exec_id = ?",
                    (json.dumps(output_payload), retry_count, node_exec_id),
                )
                conn.commit()
                if worktree_path:
                    merge_worktree(repo_path, node_exec_id)
                    remove_worktree(repo_path, worktree_path)  # 03-C §3.6: "执行完成后清理 worktree"
                return output_payload

            retry_count += 1
            if retry_count > max_retries:
                conn.execute(
                    "UPDATE work_node_execution SET status = 'blocked', retry_count = ? WHERE node_exec_id = ?",
                    (retry_count, node_exec_id),
                )
                conn.commit()
                raise SchemaValidationFailedError(
                    "output_payload failed schema validation for {} after {} retries".format(
                        node_exec_id, max_retries
                    )
                )
            conn.execute(
                "UPDATE work_node_execution SET status = 'retrying', retry_count = ? WHERE node_exec_id = ?",
                (retry_count, node_exec_id),
            )
            conn.commit()
    finally:
        conn.execute(
            "UPDATE execution_log_file SET status = 'closed' WHERE log_file_id = ?", (log_file_id,)
        )
        conn.commit()


class _UsageLimitSignal(Exception):
    """Internal: run_fn raises this (via raise_if_usage_limited helper) to
    signal a quota hit without run_work_node having to know how run_fn talks
    to `claude`."""

    def __init__(self, quota_limit_type):
        super().__init__(quota_limit_type)
        self.quota_limit_type = quota_limit_type


def raise_if_usage_limited(text):
    """Call from inside a `run_fn` after a failed claude invocation: raises
    _UsageLimitSignal if `text` matches 03-C §3.6's exact detection regex,
    otherwise does nothing (let the caller re-raise its own error)."""
    quota_type = _quota_limit_type(text)
    if quota_type:
        raise _UsageLimitSignal(quota_type)


def resume_work_graph_run(conn, run_id):
    """REQ-06 (04-P §4.2 p_work_graph_execution, breakpoint recovery):
    after a process restart, determine which nodes in `run_id` still need
    work. Nodes already 'done' are excluded (run_work_node itself also
    refuses to re-run them; this is what a caller uses to decide what to
    re-dispatch at all). Raises CorruptPersistedStateError instead of
    returning a bogus resume plan if the persisted state doesn't hold
    together -- an invalid work_graph_run.status, an invalid node status, a
    'done' node with no output, or a 'waiting_for_quota' node with no
    quota_wait_started_at (the same per-node checks run_work_node applies,
    checked here for every node up front rather than one at a time).
    """
    run_row = conn.execute("SELECT status FROM work_graph_run WHERE run_id = ?", (run_id,)).fetchone()
    if run_row is None:
        raise CorruptPersistedStateError("no such work_graph_run: " + run_id)
    if run_row[0] not in _VALID_WORK_GRAPH_RUN_STATUSES:
        raise CorruptPersistedStateError(
            "work_graph_run {} has invalid status {!r}".format(run_id, run_row[0])
        )

    nodes = conn.execute(
        "SELECT node_exec_id, status, output_payload, quota_wait_started_at FROM work_node_execution "
        "WHERE run_id = ?",
        (run_id,),
    ).fetchall()

    resumable = []
    for node_exec_id, status, output_payload, quota_wait_started_at in nodes:
        if status not in _VALID_NODE_STATUSES:
            raise CorruptPersistedStateError(
                "work_node_execution {} has invalid status {!r}".format(node_exec_id, status)
            )
        if status == "done":
            if output_payload is None:
                raise CorruptPersistedStateError(
                    "work_node_execution {} is status=done but has no output_payload".format(node_exec_id)
                )
            continue
        if status == "waiting_for_quota" and not quota_wait_started_at:
            raise CorruptPersistedStateError(
                "work_node_execution {} is status=waiting_for_quota but has no quota_wait_started_at".format(
                    node_exec_id
                )
            )
        resumable.append(node_exec_id)
    return resumable