宜博 / 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 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"]


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)."""


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.monotonic,
):
    """Runs one Work Node to contract. `run_fn(cwd) -> output_payload` is the
    caller-supplied actual work (a real `claude -p` call in production, a
    fake in tests). Returns the validated output_payload.

    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)

    project_id, task_id = conn.execute(
        "SELECT p.project_id, wgr.task_id 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()

    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:
            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 = 0
        quota_wait_start = None
        while True:
            try:
                output_payload = run_fn(cwd)
            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":
                    raise UsageLimitWaitTimeoutError(
                        "opus quota exhausted for {} -- caller should retry on a non-opus model".format(node_exec_id)
                    ) from signal
                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)