宜博 / lemongraph

"""Bridges `plan_work_graph()`'s `graph_manifest` to real per-node execution
via `execute_work_graph()` + `run_work_node()` -- the missing integration
between three previously isolated, individually-tested modules (U-011/
U-012/U-013), closed during U-031's regression audit (2026-08-01): each of
the three had its own unit tests (and even real-API validation) but none of
them had ever actually been called by production code, and the one
genuinely-validated closed loop (`e2e_pipeline.py`, U-014) used a much
simpler, separate mechanism (flat `org_node_trigger.trigger_org_node()`
calls) that bypassed all three.

Design: each node's actual work happens via its OWN one-node dynamic
workflow script (`write_workflow_script` + `execute_work_graph`), not one
big multi-node script for the whole graph. This is a deliberate trade-off,
not the "amortize orchestration cost across many nodes in one `claude -p`
call" pattern 03-C §3.5 originally motivated `execute_work_graph` with --
REQ-06's breakpoint-recovery requirement (resume after a process restart
without re-running already-`done` nodes) is only achievable with per-node,
Python-level persisted state (`work_node_run.py`'s whole design), which is
fundamentally incompatible with wrapping the WHOLE graph in one opaque
`claude` subprocess call Python has no visibility into mid-flight.

Validated for real (2026-08-01): a one-node workflow script with a custom
`agentType` + `schema` genuinely loads the target org_node's migrated
persona (confirmed via a real call asking the agent to self-identify) and
returns real `structured_output` -- see 05 §5.4.1.5.

**`verify_fn` (2026-08-02, PH4-closure audit before starting PH5)**: closes
REQ-02's boundary case (PRD internal-contradiction check must intercept
before `verifying`, not silently pass) -- found genuinely missing
(`artifact_write.write_artifact` always marks `validation_status='passed'`
unconditionally, and no Unit ever built a consistency-judgment call).
`work_graph_run.status` was ALSO never once written past its initial
'planning' value anywhere in this codebase before this -- `verify_fn`'s
'reflecting'/'verifying'/'blocked' writes below are the first real use of
that enum. Scoped deliberately narrow: one optional per-node verify
callback, one regeneration attempt on failure (matching 04-P's "回灌受影响
的最小子图" -- for a single content-producing node, the minimal affected
subgraph IS just itself), escalate to `work_graph_run.status='blocked'` +
raise if the regenerated content still fails. Not a general multi-node
re-planning mechanism -- that would be new, unscoped complexity 04-P itself
doesn't describe either.
"""

import json
import time
import uuid

from src import paths, work_graph_execute, work_node_run


class ConsistencyCheckFailedError(Exception):
    """`verify_fn` rejected a node's output twice in a row (original +
    one regeneration attempt) -- `work_graph_run.status` is already
    'blocked' by the time this is raised."""


def topological_order(graph_manifest):
    """Kahn's algorithm. Dependencies are read from BOTH `node['depends_on']`
    and `graph_manifest['edges']` and merged -- `plan_work_graph()`'s own
    output isn't schema-validated (it can't be, `--json-schema` doesn't
    work with `--agent`, see work_graph_plan.py), so a real LLM response
    might populate one and not the other; trusting only one would silently
    mis-order nodes rather than raise. Raises ValueError on a cycle or a
    dependency referencing an unknown node_id.
    """
    nodes = {n["node_id"]: n for n in graph_manifest["nodes"]}
    deps = {node_id: set(n.get("depends_on") or []) for node_id, n in nodes.items()}
    for edge in graph_manifest.get("edges", []):
        to_id, from_id = edge.get("to"), edge.get("from")
        if to_id in deps and from_id in nodes:
            deps[to_id].add(from_id)

    in_degree = {node_id: len(d) for node_id, d in deps.items()}
    dependents = {node_id: [] for node_id in nodes}
    for node_id, d in deps.items():
        for dep in d:
            if dep in dependents:
                dependents[dep].append(node_id)

    ready = sorted(node_id for node_id, degree in in_degree.items() if degree == 0)
    ordered = []
    while ready:
        node_id = ready.pop(0)
        ordered.append(nodes[node_id])
        for dependent in sorted(dependents[node_id]):
            in_degree[dependent] -= 1
            if in_degree[dependent] == 0:
                ready.append(dependent)
        ready.sort()
    if len(ordered) != len(nodes):
        raise ValueError(
            "graph_manifest has a dependency cycle or a depends_on/edge references an unknown node_id"
        )
    return ordered


def ensure_placeholder_skill(conn, skill_id, owner_org_node):
    """No Skill content has actually been migrated from the Hermes profiles
    yet (skill_update.py's own documented gap) -- work_node_execution.
    skill_ref is a required FK, so a real row has to exist regardless.
    Idempotent."""
    exists = conn.execute("SELECT 1 FROM skill WHERE skill_id = ?", (skill_id,)).fetchone()
    if exists is None:
        conn.execute(
            "INSERT INTO skill (skill_id, owner_org_node, file_path, current_git_commit_sha, source, status) "
            "VALUES (?, ?, ?, 'unknown', 'new', 'active')",
            (skill_id, owner_org_node, "skills/" + skill_id + "/"),
        )
        conn.commit()


def _ensure_node_row(conn, run_id, node_id, subagent_type, skill_ref, input_schema_ref, output_schema_ref, input_payload):
    existing = conn.execute(
        "SELECT node_exec_id FROM work_node_execution WHERE run_id = ? AND node_id = ?", (run_id, node_id)
    ).fetchone()
    if existing:
        return existing[0]
    node_exec_id = str(uuid.uuid4())
    conn.execute(
        "INSERT INTO work_node_execution (node_exec_id, run_id, node_id, subagent_type, skill_ref, "
        "input_schema_ref, output_schema_ref, input_payload, status) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending')",
        (node_exec_id, run_id, node_id, subagent_type, skill_ref, input_schema_ref, output_schema_ref, json.dumps(input_payload)),
    )
    conn.commit()
    return node_exec_id


def _render_node_script(script_name, org_node, prompt, output_schema, model):
    model_opt = ", model: " + json.dumps(model) if model else ""
    return (
        "export const meta = {{ name: {name}, description: 'graphlemon Work Node ({org_node})' }}\n\n"
        "const result = await agent({prompt}, {{ agentType: {org_node_json}, schema: {schema}{model_opt} }})\n\n"
        "return result\n"
    ).format(
        name=json.dumps(script_name), org_node=org_node, org_node_json=json.dumps(org_node),
        prompt=json.dumps(prompt), schema=json.dumps(output_schema), model_opt=model_opt,
    )


def _make_run_fn(target_folder_path, node_exec_id, org_node, prompt, output_schema, timeout):
    script_name = "graphlemon-node-" + node_exec_id.replace("-", "")

    def run_fn(cwd, model=None):
        script_content = _render_node_script(script_name, org_node, prompt, output_schema, model)
        work_graph_execute.write_workflow_script(cwd, script_name, script_content)
        try:
            return work_graph_execute.execute_work_graph(cwd, script_name, {}, output_schema, timeout=timeout)
        except work_graph_execute.UsageLimitReachedError as e:
            work_node_run.raise_usage_limit_signal(e.quota_limit_type)
        # ClassifierBlockedError / generic WorkGraphExecuteError propagate
        # as-is -- run_work_node's own `except Exception` clause marks the
        # node 'blocked' and re-raises, no special handling needed here.

    return run_fn


def run_graph(
    conn, run_id, org_node, skill_id, target_folder_path, graph_manifest,
    build_prompt, input_schema, output_schema, timeout=600, needs_worktree=False,
    verify_fn=None, sleep=time.sleep, now=time.time,
):
    """Runs every node in `graph_manifest` (topologically ordered) through
    `run_work_node()`, each backed by a real one-node dynamic-workflow
    `agent()` call using `org_node`'s migrated persona.

    `build_prompt(node, outputs)` is called right before running that
    node -- `outputs` contains every already-completed node's
    output_payload so far (topological order guarantees a node's
    dependencies are always already present), so prompts needing upstream
    content (e.g. "here is the approved BRD...") can be built just-in-time
    instead of requiring the whole graph's prompts precomputed before
    anything runs. This also makes resume-after-restart work correctly:
    `run_work_node` short-circuits already-`done` nodes and returns their
    persisted output, which lands in `outputs` the same as a freshly-run
    one would.

    `needs_worktree=True` (03-C §3.6: "并行写代码的节点必须先申请
    git_worktree_ref") isolates each node's actual file changes in its own
    git worktree, merged back only once the node succeeds -- callers
    writing code (Coder) should pass this; callers only producing document
    content (Planner) don't need it.

    **U-035 (05 §5.8) self-target override**: when `target_folder_path` IS
    graphlemon's own checkout (`paths.is_self_target`), `needs_worktree` is
    forced True regardless of what the caller passed -- the currently-running
    backend process is reading this exact code from disk, so an in-place edit
    to the main checkout (even a serial, non-parallel one) risks corrupting
    state out from under itself mid-execution. This is unconditional (not
    "only when parallel", 03-C §3.6's general rule) precisely because
    self-target has no safe non-worktree case at all.

    `verify_fn(node, output) -> (ok: bool, feedback)`, if given, is called
    right after each node completes, BEFORE moving on to nodes that depend
    on it -- lets a caller intercept a bad output before it propagates
    downstream (e.g. REQ-02's PRD-internal-contradiction check), not just
    after the fact. On `ok=False`: `work_graph_run.status` -> 'reflecting',
    the SAME node is regenerated once with `feedback` appended to its
    prompt, and re-verified. If it fails again: `work_graph_run.status`
    -> 'blocked' and `ConsistencyCheckFailedError` is raised. If it now
    passes: `work_graph_run.status` -> 'verifying' and the regenerated
    output replaces the original in `outputs`.

    Returns {node_id: output_payload}.
    """
    ensure_placeholder_skill(conn, skill_id, org_node)
    effective_needs_worktree = needs_worktree or paths.is_self_target(target_folder_path)
    ordered = topological_order(graph_manifest)
    outputs = {}
    for node in ordered:
        node_id = node["node_id"]
        node_exec_id = _ensure_node_row(
            conn, run_id, node_id, org_node, skill_id,
            "work_graph_node_input.schema.json", "work_graph_node_output.schema.json",
            {"node_id": node_id},
        )
        prompt = build_prompt(node, outputs)
        run_fn = _make_run_fn(target_folder_path, node_exec_id, org_node, prompt, output_schema, timeout)
        output = work_node_run.run_work_node(
            conn, node_exec_id, {"node_id": node_id}, input_schema, output_schema, run_fn,
            needs_worktree=effective_needs_worktree, repo_path=target_folder_path,
            sleep=sleep, now=now,
        )

        if verify_fn is not None:
            output = _verify_and_maybe_regenerate(
                conn, run_id, org_node, skill_id, target_folder_path, node, node_id, prompt,
                output, input_schema, output_schema, timeout, effective_needs_worktree, verify_fn, sleep, now,
            )

        outputs[node_id] = output
    return outputs


def _verify_and_maybe_regenerate(
    conn, run_id, org_node, skill_id, target_folder_path, node, node_id, prompt,
    output, input_schema, output_schema, timeout, needs_worktree, verify_fn, sleep, now,
):
    ok, feedback = verify_fn(node, output)
    if ok:
        return output

    conn.execute("UPDATE work_graph_run SET status = 'reflecting' WHERE run_id = ?", (run_id,))
    conn.commit()

    reflect_node_id = node_id + "-reflect"
    reflect_node_exec_id = _ensure_node_row(
        conn, run_id, reflect_node_id, org_node, skill_id,
        "work_graph_node_input.schema.json", "work_graph_node_output.schema.json",
        {"node_id": node_id, "feedback": feedback},
    )
    reflect_prompt = (
        prompt + "\n\nA consistency reviewer flagged issues with your previous draft: "
        + "; ".join(feedback or []) + "\nProduce a corrected version resolving these issues."
    )
    reflect_run_fn = _make_run_fn(target_folder_path, reflect_node_exec_id, org_node, reflect_prompt, output_schema, timeout)
    regenerated = work_node_run.run_work_node(
        conn, reflect_node_exec_id, {"node_id": node_id}, input_schema, output_schema, reflect_run_fn,
        needs_worktree=needs_worktree, repo_path=target_folder_path,
        sleep=sleep, now=now,
    )

    ok2, feedback2 = verify_fn(node, regenerated)
    if not ok2:
        conn.execute("UPDATE work_graph_run SET status = 'blocked' WHERE run_id = ?", (run_id,))
        conn.commit()
        raise ConsistencyCheckFailedError(
            "node {} still failed verification after one regeneration attempt: {}".format(node_id, feedback2)
        )

    conn.execute("UPDATE work_graph_run SET status = 'verifying' WHERE run_id = ?", (run_id,))
    conn.commit()
    return regenerated