宜博 / lemongraph

"""D (2026-08-06): Work Graph time-travel debugging -- fork a new
work_graph_run from an already-'done' work_node_execution, substituting a
human-provided output for that node and re-executing everything strictly
downstream of it for real, while every other node reuses the original
run's persisted output_payload untouched. See
docs/todo/work-graph-time-travel-debugging-design.md for the full design.

**Output substitution, not re-prompting** (design doc §2, resolving the
requirements doc's FR-1 ambiguity): the fork point's own output_payload is
replaced directly with what the human supplies -- it is NOT re-executed
via a new claude -p call. This sidesteps having to reconcile the different
build_prompt/input_payload conventions different callers use (Planner's
fixed artifact_type lookup, E's generic task_prompt field, F's jury) --
this module only needs the persisted graph_manifest and each node's
output_payload, nothing caller-specific, for the SUBSTITUTED node. The
nodes that DO get genuinely re-executed (the strict downstream subgraph)
still need a real build_prompt, which is reconstructed generically -- see
_reconstruct_build_prompt.

**v1 restricts forking to a parent run that has already reached a
terminal status** ('done'/'failed'/'blocked') -- design doc §1: this keeps
the new fork completely outside ux_work_graph_run_active's partial unique
index (that index only applies to non-terminal statuses), so no schema/
constraint change is needed to avoid a collision. Forking from a
still-in-flight parent run is left for a future iteration.

**Zero changes to the execution engine** (design doc §3): pre-seeding
every non-downstream node as a 'done' work_node_execution row before
calling run_graph() on the UNCHANGED original graph_manifest is enough --
run_work_node()'s existing REQ-06 breakpoint-recovery short-circuit
(already-'done' nodes return their persisted output without calling
run_fn) does all the "reuse without re-executing" work for free.
"""

import json
import threading
import time
import uuid

from src import e2e_pipeline, schema, work_graph_orchestrate, work_graph_plan

_FORK_INPUT_SCHEMA = {"type": "object", "required": ["node_id"]}
_FORK_OUTPUT_SCHEMA = {"type": "object", "required": ["content"]}
_TERMINAL_RUN_STATUSES = ("done", "failed", "blocked")


class ForkSourceNotDoneError(Exception):
    """FR-1: can only fork from a work_node_execution with status='done' --
    its output_payload has to exist to be a valid substitution point."""


class ForkParentRunNotTerminalError(Exception):
    """v1 scope cut (design doc §1): forking from a node inside a
    still-in-flight parent run would collide with ux_work_graph_run_active
    (same task_id+org_node, both non-terminal) -- not supported yet."""


class ForkGraphShapeUnsupportedError(Exception):
    """The original graph_manifest doesn't carry enough self-contained
    information to reconstruct a build_prompt generically for the
    re-executed downstream nodes -- e.g. F's jury nodes or Coder's single
    ad hoc node, which rely on caller-side closures over state that isn't
    persisted in graph_manifest itself (design doc §2)."""


def _downstream_node_ids(graph_manifest, start_node_id):
    """Every node_id reachable from start_node_id via depends_on/edges,
    INCLUDING start_node_id itself. Same dependents-graph construction
    topological_layers() uses, kept local here rather than imported since
    it's a small, self-contained piece not worth threading a shared helper
    for across modules."""
    nodes = {n["node_id"]: n for n in graph_manifest["nodes"]}
    dependents = {node_id: [] for node_id in nodes}
    for node in graph_manifest["nodes"]:
        for dep in (node.get("depends_on") or []):
            if dep in dependents:
                dependents[dep].append(node["node_id"])
    for edge in graph_manifest.get("edges", []):
        from_id, to_id = edge.get("from"), edge.get("to")
        if from_id in dependents and to_id in nodes:
            dependents[from_id].append(to_id)

    seen = set()
    frontier = [start_node_id]
    while frontier:
        node_id = frontier.pop()
        if node_id in seen:
            continue
        seen.add(node_id)
        frontier.extend(dependents[node_id])
    return seen


def _reconstruct_build_prompt(conn, task_id, org_node, graph_manifest):
    """Rebuilds a build_prompt(node, outputs) usable for the re-executed
    downstream subgraph, purely from persisted state -- no access to the
    original call site's Python closures. Two supported shapes (design doc
    §2): Planner's canonical doc chain (every node has 'artifact_type',
    prompt reconstructed via e2e_pipeline's own builder + task.goal_text)
    and E's generic dynamic path (every node carries its own
    self-contained 'task_prompt'). Anything else raises rather than
    silently building broken prompts."""
    nodes = graph_manifest["nodes"]
    if org_node == "planner" and nodes and all("artifact_type" in n for n in nodes):
        goal_text = conn.execute("SELECT goal_text FROM task WHERE task_id = ?", (task_id,)).fetchone()[0]
        node_artifact_types = {n["node_id"]: n["artifact_type"] for n in nodes}
        return e2e_pipeline._planner_prompt_builder(goal_text, node_artifact_types)
    if nodes and all("task_prompt" in n for n in nodes):
        return work_graph_plan.generic_node_prompt_builder
    raise ForkGraphShapeUnsupportedError(
        "graph_manifest nodes must either all have 'artifact_type' (Planner's canonical chain) "
        "or all have 'task_prompt' (the generic dynamic-planning path) to reconstruct prompts for a fork"
    )


def create_fork(conn, source_node_exec_id, new_output_payload):
    """Synchronous, DB-only half of a fork: validates the source node and
    parent run, pre-seeds every non-downstream node as 'done' under a
    brand-new work_graph_run, and returns everything execute_fork() needs
    to actually run the downstream subgraph. Split from execute_fork() so
    an HTTP caller can hand the (fast) run_id back immediately and let the
    (slow, real claude -p calls) execution happen in the background --
    same split task_execution.start_task_execution_in_background already
    uses for the manual Start button."""
    row = conn.execute(
        "SELECT wne.run_id, wne.node_id, wne.status, wne.subagent_type, wne.skill_ref, "
        "wgr.task_id, wgr.org_node, wgr.graph_manifest, wgr.status "
        "FROM work_node_execution wne JOIN work_graph_run wgr ON wgr.run_id = wne.run_id "
        "WHERE wne.node_exec_id = ?",
        (source_node_exec_id,),
    ).fetchone()
    if row is None:
        raise ValueError("no such work_node_execution: " + source_node_exec_id)
    (
        source_run_id, source_node_id, source_status, subagent_type, skill_ref,
        task_id, org_node, graph_manifest_json, parent_run_status,
    ) = row
    if source_status != "done":
        raise ForkSourceNotDoneError(
            "work_node_execution {} is status={}, must be 'done' to fork from".format(
                source_node_exec_id, source_status
            )
        )
    if parent_run_status not in _TERMINAL_RUN_STATUSES:
        raise ForkParentRunNotTerminalError(
            "parent work_graph_run {} is status={}, must be terminal ({}) to fork from -- forking from "
            "a still-in-flight run isn't supported yet".format(
                source_run_id, parent_run_status, _TERMINAL_RUN_STATUSES
            )
        )

    graph_manifest = json.loads(graph_manifest_json)
    build_prompt = _reconstruct_build_prompt(conn, task_id, org_node, graph_manifest)

    downstream = _downstream_node_ids(graph_manifest, source_node_id) - {source_node_id}
    original_outputs = {
        r[0]: r[1] for r in conn.execute(
            "SELECT node_id, output_payload FROM work_node_execution WHERE run_id = ? AND status = 'done'",
            (source_run_id,),
        ).fetchall()
    }

    new_run_id = str(uuid.uuid4())
    conn.execute(
        "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status, started_at, "
        "parent_run_id, fork_of_node_exec_id) VALUES (?, ?, ?, ?, 'running', CURRENT_TIMESTAMP, ?, ?)",
        (new_run_id, task_id, org_node, graph_manifest_json, source_run_id, source_node_exec_id),
    )
    conn.commit()

    for node in graph_manifest["nodes"]:
        node_id = node["node_id"]
        if node_id in downstream:
            continue
        payload = new_output_payload if node_id == source_node_id else json.loads(original_outputs[node_id])
        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, output_payload, status) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'done')",
            (
                str(uuid.uuid4()), new_run_id, node_id, subagent_type, skill_ref,
                "work_graph_node_input.schema.json", "work_graph_node_output.schema.json",
                json.dumps({"node_id": node_id}), json.dumps(payload),
            ),
        )
    conn.commit()

    return {
        "run_id": new_run_id, "org_node": org_node, "skill_ref": skill_ref,
        "graph_manifest": graph_manifest, "build_prompt": build_prompt,
    }


def execute_fork(conn, fork_info, target_folder_path, timeout=600, sleep=time.sleep, now=time.time):
    """The real (slow) half: runs the downstream subgraph via the
    unmodified run_graph() -- every pre-seeded node short-circuits, every
    downstream node executes for real (see module docstring). Sets the new
    run's status to 'done'/'failed' on completion -- unlike every other
    real graph_manifest run in this codebase, which never transitions
    work_graph_run.status past its initial value (a pre-existing gap, out
    of D's scope to fix generically -- see design doc §3's own note); the
    forked run gets this narrowly because FR-4's console visibility needs
    SOME real completion signal to show, not because this module is fixing
    the wider gap."""
    run_id = fork_info["run_id"]
    try:
        outputs = work_graph_orchestrate.run_graph(
            conn, run_id, fork_info["org_node"], fork_info["skill_ref"], target_folder_path,
            fork_info["graph_manifest"], fork_info["build_prompt"], _FORK_INPUT_SCHEMA, _FORK_OUTPUT_SCHEMA,
            timeout=timeout, sleep=sleep, now=now,
        )
    except Exception:
        conn.execute(
            "UPDATE work_graph_run SET status = 'failed', completed_at = CURRENT_TIMESTAMP WHERE run_id = ?",
            (run_id,),
        )
        conn.commit()
        raise
    conn.execute(
        "UPDATE work_graph_run SET status = 'done', completed_at = CURRENT_TIMESTAMP WHERE run_id = ?", (run_id,)
    )
    conn.commit()
    return outputs


def fork_work_graph_run(conn, source_node_exec_id, new_output_payload, target_folder_path, timeout=600):
    """create_fork() + execute_fork() in one synchronous call -- what
    tests and any fully-synchronous caller should use. HTTP handlers use
    the two halves separately (see start_fork_in_background)."""
    fork_info = create_fork(conn, source_node_exec_id, new_output_payload)
    outputs = execute_fork(conn, fork_info, target_folder_path, timeout=timeout)
    return {"run_id": fork_info["run_id"], "outputs": outputs}


def start_fork_in_background(db_path, source_node_exec_id, new_output_payload, target_folder_path, timeout=600):
    """HTTP entrypoint: create_fork() runs synchronously (fast, DB-only)
    so the caller gets run_id back immediately and can navigate the
    console straight to it; execute_fork()'s real claude -p calls then run
    in a background daemon thread, own connection (check_same_thread=False
    -- same reasoning as scheduler.run_goal_to_completion's own conn, see
    that module's comment: this can reach run_graph()'s multi-node
    concurrent layers). Returns (run_id, thread) -- thread is already
    started, callers/tests can .join() it if needed (mirrors
    task_execution.start_task_execution_in_background's own Thread
    return)."""
    conn = schema.connect(db_path, check_same_thread=False)
    fork_info = create_fork(conn, source_node_exec_id, new_output_payload)

    def run():
        try:
            execute_fork(conn, fork_info, target_folder_path, timeout=timeout)
        except Exception:
            pass  # already reflected in work_graph_run/work_node_execution status
        finally:
            conn.close()

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    return fork_info["run_id"], thread