atlas-iac/dockerfiles/patch-hermes-execution-safety.py

"""Patch upstream Hermes exact-run completion and decomposition safety."""

import os
from pathlib import Path


def replace_once(source: str, before: str, after: str, label: str) -> str:
    """Replace one anchored upstream fragment and fail closed on source drift."""
    count = source.count(before)
    if count != 1:
        raise SystemExit(
            f"Hermes {label} patch context changed: expected 1, found {count}"
        )
    return source.replace(before, after, 1)


source_root = Path(os.environ.get("HERMES_SOURCE_ROOT", "/opt/hermes"))
db_path = source_root / "hermes_cli/kanban_db.py"
db = db_path.read_text(encoding="utf-8")

complete_signature_before = '''def complete_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    result: Optional[str] = None,
    summary: Optional[str] = None,
    metadata: Optional[dict] = None,
    created_cards: Optional[Iterable[str]] = None,
    expected_run_id: Optional[int] = None,
) -> bool:
'''
complete_signature_after = '''def complete_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    result: Optional[str] = None,
    summary: Optional[str] = None,
    metadata: Optional[dict] = None,
    created_cards: Optional[Iterable[str]] = None,
    expected_run_id: Optional[int] = None,
    replay_ended_run_id: Optional[int] = None,
) -> bool:
'''
db = replace_once(
    db,
    complete_signature_before,
    complete_signature_after,
    "exact ended-run completion signature",
)

complete_branch_before = '''    with write_txn(conn):
        if expected_run_id is None:
            cur = conn.execute(
                """
                UPDATE tasks
                   SET status       = 'done',
                       result       = ?,
                       completed_at = ?,
                       claim_lock   = NULL,
                       claim_expires= NULL,
                       worker_pid   = NULL,
                       block_kind   = NULL,
                       block_recurrences = 0
                 WHERE id = ?
                   AND status IN ('running', 'ready', 'blocked', 'scheduled')
                """,
                (result, now, task_id),
            )
        else:
'''
complete_branch_after = '''    # Journal recovery may arrive after a legacy error path ended the
    # accepted run. Verify that it is still the latest run in this same
    # transaction; a replacement attempt must make the replay fail closed.
    with write_txn(conn):
        replayed_run_id = None
        if expected_run_id is not None and replay_ended_run_id is not None:
            raise ValueError("expected_run_id and replay_ended_run_id are mutually exclusive")
        if replay_ended_run_id is not None:
            replayed_run_id = int(replay_ended_run_id)
            cur = conn.execute(
                """
                UPDATE tasks
                   SET status       = 'done',
                       result       = ?,
                       completed_at = ?,
                       claim_lock   = NULL,
                       claim_expires= NULL,
                       worker_pid   = NULL,
                       block_kind   = NULL,
                       block_recurrences = 0
                 WHERE id = ?
                   AND status IN ('ready', 'blocked')
                   AND current_run_id IS NULL
                   AND ? = (
                       SELECT id FROM task_runs
                        WHERE task_id = ?
                        ORDER BY id DESC
                        LIMIT 1
                   )
                   AND EXISTS (
                       SELECT 1 FROM task_runs
                        WHERE id = ?
                          AND task_id = ?
                          AND ended_at IS NOT NULL
                   )
                """,
                (
                    result,
                    now,
                    task_id,
                    replayed_run_id,
                    task_id,
                    replayed_run_id,
                    task_id,
                ),
            )
        elif expected_run_id is None:
            cur = conn.execute(
                """
                UPDATE tasks
                   SET status       = 'done',
                       result       = ?,
                       completed_at = ?,
                       claim_lock   = NULL,
                       claim_expires= NULL,
                       worker_pid   = NULL,
                       block_kind   = NULL,
                       block_recurrences = 0
                 WHERE id = ?
                   AND status IN ('running', 'ready', 'blocked', 'scheduled')
                """,
                (result, now, task_id),
            )
        else:
'''
db = replace_once(
    db,
    complete_branch_before,
    complete_branch_after,
    "transactional exact ended-run completion",
)

run_completion_before = '''        run_id = _end_run(
            conn, task_id,
            outcome="completed", status="done",
            summary=summary if summary is not None else result,
            metadata=metadata,
        )
        # If complete_task was called on a never-claimed task (ready or
        # blocked → done with no run in flight), synthesize a
        # zero-duration run so the handoff fields are persisted in
        # attempt history instead of silently lost.
        if run_id is None and (summary or metadata or result):
            run_id = _synthesize_ended_run(
                conn, task_id,
                outcome="completed",
                summary=summary if summary is not None else result,
                metadata=metadata,
            )
'''
run_completion_after = '''        if replayed_run_id is None:
            run_id = _end_run(
                conn, task_id,
                outcome="completed", status="done",
                summary=summary if summary is not None else result,
                metadata=metadata,
            )
            # If complete_task was called on a never-claimed task (ready or
            # blocked → done with no run in flight), synthesize a
            # zero-duration run so the handoff fields are persisted in
            # attempt history instead of silently lost.
            if run_id is None and (summary or metadata or result):
                run_id = _synthesize_ended_run(
                    conn, task_id,
                    outcome="completed",
                    summary=summary if summary is not None else result,
                    metadata=metadata,
                )
        else:
            run_id = replayed_run_id
            replayed = conn.execute(
                """
                UPDATE task_runs
                   SET status = 'done',
                       outcome = 'completed',
                       summary = ?,
                       error = NULL,
                       metadata = ?,
                       claim_lock = NULL,
                       claim_expires = NULL,
                       worker_pid = NULL
                 WHERE id = ?
                   AND task_id = ?
                   AND ended_at IS NOT NULL
                """,
                (
                    summary if summary is not None else result,
                    json.dumps(metadata, ensure_ascii=False) if metadata else None,
                    run_id,
                    task_id,
                ),
            )
            if replayed.rowcount != 1:
                raise RuntimeError("exact ended-run completion lost its run identity")
        conn.execute(
            "UPDATE tasks SET completed_run_id = ? WHERE id = ?",
            (run_id, task_id),
        )
'''
db = replace_once(
    db,
    run_completion_before,
    run_completion_after,
    "exact completion run provenance",
)

reclaim_before = '''def reclaim_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    reason: Optional[str] = None,
    signal_fn=None,
) -> bool:
    """Operator-driven reclaim: release the claim and reset to ``ready``.

    Unlike :func:`release_stale_claims` which only acts on tasks whose
    ``claim_expires`` has passed, this function reclaims immediately
    regardless of TTL. Intended for the dashboard/CLI recovery flow
    when an operator wants to abort a running worker without waiting
    for the TTL to expire (e.g. after seeing a hallucination warning).

    Returns True if a reclaim happened, False if the task isn't in a
    reclaimable state (not running, or doesn't exist).
    """
    row = conn.execute(
        "SELECT status, claim_lock, worker_pid FROM tasks WHERE id = ?",
        (task_id,),
    ).fetchone()
    if not row:
        return False
    if row["status"] != "running" and row["claim_lock"] is None:
        # Nothing to reclaim — already ready / blocked / done.
        return False
    prev_lock = row["claim_lock"]
    termination = _terminate_reclaimed_worker(
        row["worker_pid"], prev_lock, signal_fn=signal_fn,
    )
    with write_txn(conn):
        cur = conn.execute(
            "UPDATE tasks SET status = 'ready', claim_lock = NULL, "
            "claim_expires = NULL, worker_pid = NULL "
            "WHERE id = ? AND status IN ('running', 'ready', 'blocked') "
            "AND claim_lock IS ?",
            (task_id, prev_lock),
        )
        if cur.rowcount != 1:
            return False
        run_id = _end_run(
            conn, task_id,
            outcome="reclaimed", status="reclaimed",
            error=(
                f"manual_reclaim: {reason}" if reason
                else f"manual_reclaim lock={prev_lock}"
            ),
            metadata=termination,
        )
        payload = {
            "manual": True,
            "reason": reason,
            "prev_lock": prev_lock,
        }
        payload.update(termination)
        _append_event(
            conn, task_id, "reclaimed",
            payload,
            run_id=run_id,
        )
    # Operator intervention — they've looked at the task, so the
    # consecutive-failures counter is now stale. Give the next retry
    # a fresh budget. (_clear_failure_counter opens its own write_txn,
    # so it runs after the enclosing one commits.)
    _clear_failure_counter(conn, task_id)
    return True
'''
reclaim_after = '''def reclaim_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    reason: Optional[str] = None,
    signal_fn=None,
    expected_run_id: Optional[int] = None,
) -> bool:
    """Operator-driven reclaim: release the claim and reset to ``ready``.

    Unlike :func:`release_stale_claims` which only acts on tasks whose
    ``claim_expires`` has passed, this function reclaims immediately
    regardless of TTL. Intended for the dashboard/CLI recovery flow
    when an operator wants to abort a running worker without waiting
    for the TTL to expire (e.g. after seeing a hallucination warning).

    When ``expected_run_id`` is supplied, the run identity is checked after
    ``BEGIN IMMEDIATE`` and included in the guarded update. A replacement run
    is therefore neither signalled nor reclaimed by stale recovery evidence.

    Returns True if a reclaim happened, False if the task isn't in a
    reclaimable state (not running, or doesn't exist), or its run changed.
    """
    with write_txn(conn):
        row = conn.execute(
            "SELECT status, claim_lock, worker_pid, current_run_id "
            "FROM tasks WHERE id = ?",
            (task_id,),
        ).fetchone()
        if not row:
            return False
        if expected_run_id is not None and row["current_run_id"] != expected_run_id:
            return False
        if row["status"] != "running" and row["claim_lock"] is None:
            # Nothing to reclaim — already ready / blocked / done.
            return False
        prev_lock = row["claim_lock"]
        termination = _terminate_reclaimed_worker(
            row["worker_pid"], prev_lock, signal_fn=signal_fn,
        )
        if expected_run_id is None:
            cur = conn.execute(
                "UPDATE tasks SET status = 'ready', claim_lock = NULL, "
                "claim_expires = NULL, worker_pid = NULL "
                "WHERE id = ? AND status IN ('running', 'ready', 'blocked') "
                "AND claim_lock IS ?",
                (task_id, prev_lock),
            )
        else:
            cur = conn.execute(
                "UPDATE tasks SET status = 'ready', claim_lock = NULL, "
                "claim_expires = NULL, worker_pid = NULL "
                "WHERE id = ? AND status IN ('running', 'ready', 'blocked') "
                "AND claim_lock IS ? AND current_run_id = ?",
                (task_id, prev_lock, int(expected_run_id)),
            )
        if cur.rowcount != 1:
            return False
        run_id = _end_run(
            conn, task_id,
            outcome="reclaimed", status="reclaimed",
            error=(
                f"manual_reclaim: {reason}" if reason
                else f"manual_reclaim lock={prev_lock}"
            ),
            metadata=termination,
        )
        payload = {
            "manual": True,
            "reason": reason,
            "prev_lock": prev_lock,
        }
        payload.update(termination)
        _append_event(
            conn, task_id, "reclaimed",
            payload,
            run_id=run_id,
        )
    # Operator intervention — they've looked at the task, so the
    # consecutive-failures counter is now stale. Give the next retry
    # a fresh budget. (_clear_failure_counter opens its own write_txn,
    # so it runs after the enclosing one commits.)
    _clear_failure_counter(conn, task_id)
    return True
'''
db = replace_once(
    db,
    reclaim_before,
    reclaim_after,
    "transactional exact-run reclaim",
)

task_field_before = '''    current_run_id: Optional[int] = None
    workflow_template_id: Optional[str] = None
'''
task_field_after = '''    current_run_id: Optional[int] = None
    # Durable provenance for the run that authoritatively completed this task.
    completed_run_id: Optional[int] = None
    workflow_template_id: Optional[str] = None
'''
db = replace_once(
    db,
    task_field_before,
    task_field_after,
    "completed-run task field",
)

task_row_before = '''            current_run_id=(
                row["current_run_id"] if "current_run_id" in keys else None
            ),
            workflow_template_id=(
'''
task_row_after = '''            current_run_id=(
                row["current_run_id"] if "current_run_id" in keys else None
            ),
            completed_run_id=(
                row["completed_run_id"] if "completed_run_id" in keys else None
            ),
            workflow_template_id=(
'''
db = replace_once(
    db,
    task_row_before,
    task_row_after,
    "completed-run row mapping",
)

task_schema_before = '''    -- run is in-flight). Denormalised for cheap reads.
    current_run_id       INTEGER,
    -- Forward-compat for v2 workflow routing. In v1 the kernel writes
'''
task_schema_after = '''    -- run is in-flight). Denormalised for cheap reads.
    current_run_id       INTEGER,
    -- Immutable provenance for the task result's authoritative run.
    completed_run_id     INTEGER,
    -- Forward-compat for v2 workflow routing. In v1 the kernel writes
'''
db = replace_once(
    db,
    task_schema_before,
    task_schema_after,
    "completed-run schema",
)

task_migration_before = '''    if "current_run_id" not in cols:
        _add_column_if_missing(
            conn, "tasks", "current_run_id", "current_run_id INTEGER"
        )
    if "workflow_template_id" not in cols:
'''
task_migration_after = '''    if "current_run_id" not in cols:
        _add_column_if_missing(
            conn, "tasks", "current_run_id", "current_run_id INTEGER"
        )
    if "completed_run_id" not in cols:
        _add_column_if_missing(
            conn, "tasks", "completed_run_id", "completed_run_id INTEGER"
        )
    if "workflow_template_id" not in cols:
'''
db = replace_once(
    db,
    task_migration_before,
    task_migration_after,
    "completed-run migration",
)

specify_signature_before = '''def specify_triage_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    title: Optional[str] = None,
    body: Optional[str] = None,
    assignee: Optional[str] = None,
    author: Optional[str] = None,
) -> bool:
'''
specify_signature_after = '''def specify_triage_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    title: Optional[str] = None,
    body: Optional[str] = None,
    assignee: Optional[str] = None,
    author: Optional[str] = None,
    require_no_runs: bool = False,
) -> bool:
'''
db = replace_once(
    db,
    specify_signature_before,
    specify_signature_after,
    "transactional specify signature",
)

specify_guard_before = '''        if existing is None:
            return False
        sets: list[str] = ["status = 'todo'"]
'''
specify_guard_after = '''        if existing is None:
            return False
        if require_no_runs and conn.execute(
            "SELECT 1 FROM task_runs WHERE task_id = ? LIMIT 1",
            (task_id,),
        ).fetchone() is not None:
            return False
        sets: list[str] = ["status = 'todo'"]
'''
db = replace_once(
    db,
    specify_guard_before,
    specify_guard_after,
    "transactional specify execution-history guard",
)

decompose_signature_before = '''def decompose_triage_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    root_assignee: Optional[str],
    children: list[dict],
    author: Optional[str] = None,
    auto_promote: bool = True,
) -> Optional[list[str]]:
'''
decompose_signature_after = '''def decompose_triage_task(
    conn: sqlite3.Connection,
    task_id: str,
    *,
    root_assignee: Optional[str],
    children: list[dict],
    author: Optional[str] = None,
    auto_promote: bool = True,
    require_no_runs: bool = False,
) -> Optional[list[str]]:
'''
db = replace_once(
    db,
    decompose_signature_before,
    decompose_signature_after,
    "transactional decomposition signature",
)

decompose_guard_before = '''        if root_row["status"] != "triage":
            return None
        tenant = root_row["tenant"]
'''
decompose_guard_after = '''        if root_row["status"] != "triage":
            return None
        if require_no_runs and conn.execute(
            "SELECT 1 FROM task_runs WHERE task_id = ? LIMIT 1",
            (task_id,),
        ).fetchone() is not None:
            return None
        tenant = root_row["tenant"]
'''
db = replace_once(
    db,
    decompose_guard_before,
    decompose_guard_after,
    "transactional decomposition execution-history guard",
)
db_path.write_text(db, encoding="utf-8")

decompose_path = source_root / "hermes_cli/kanban_decompose.py"
decompose = decompose_path.read_text(encoding="utf-8")

helper_anchor = '''    if chosen not in valid_names:
        return default_assignee
    return chosen


'''
helper_replacement = helper_anchor + '''def _has_execution_history(task_id: str) -> bool:
    """Return whether a task has ever entered a worker run."""
    with kb.connect_closing() as conn:
        return bool(kb.list_runs(conn, task_id))


'''
decompose = replace_once(
    decompose,
    helper_anchor,
    helper_replacement,
    "auto-decompose execution-history helper",
)

signature_before = '''def decompose_task(
    task_id: str,
    *,
    author: Optional[str] = None,
    timeout: Optional[int] = None,
) -> DecomposeOutcome:
'''
signature_after = '''def decompose_task(
    task_id: str,
    *,
    author: Optional[str] = None,
    timeout: Optional[int] = None,
    automatic: bool = False,
) -> DecomposeOutcome:
'''
decompose = replace_once(
    decompose,
    signature_before,
    signature_after,
    "auto-decompose function signature",
)

status_before = '''    if task.status != "triage":
        return DecomposeOutcome(
            task_id, False, f"task is not in triage (status={task.status!r})"
        )

    cfg = _load_config()
'''
status_after = '''    if task.status != "triage":
        return DecomposeOutcome(
            task_id, False, f"task is not in triage (status={task.status!r})"
        )
    if automatic and _has_execution_history(task_id):
        return DecomposeOutcome(
            task_id,
            False,
            "task has execution history and requires deliberate manual triage",
        )

    cfg = _load_config()
'''
decompose = replace_once(
    decompose,
    status_before,
    status_after,
    "auto-decompose preflight guard",
)

single_before = '''                author=audit_author,
            )
        if not ok:
'''
single_after = '''                author=audit_author,
                require_no_runs=automatic,
            )
        if not ok:
            if automatic and _has_execution_history(task_id):
                return DecomposeOutcome(
                    task_id,
                    False,
                    "task gained execution history and requires deliberate manual triage",
                )
'''
decompose = replace_once(
    decompose,
    single_before,
    single_after,
    "auto-decompose single-task commit guard",
)

fanout_before = '''                author=audit_author,
                auto_promote=auto_promote,
            )
'''
fanout_after = '''                author=audit_author,
                auto_promote=auto_promote,
                require_no_runs=automatic,
            )
'''
decompose = replace_once(
    decompose,
    fanout_before,
    fanout_after,
    "auto-decompose transactional fanout guard",
)

fanout_outcome_before = '''    if child_ids is None:
        return DecomposeOutcome(
            task_id, False, "task moved out of triage before decomposition",
        )
'''
fanout_outcome_after = '''    if child_ids is None:
        if automatic and _has_execution_history(task_id):
            return DecomposeOutcome(
                task_id,
                False,
                "task gained execution history and requires deliberate manual triage",
            )
        return DecomposeOutcome(
            task_id, False, "task moved out of triage before decomposition",
        )
'''
decompose = replace_once(
    decompose,
    fanout_outcome_before,
    fanout_outcome_after,
    "auto-decompose fanout rejection reason",
)
decompose_path.write_text(decompose, encoding="utf-8")


watcher_path = source_root / "gateway/kanban_watchers.py"
watcher = watcher_path.read_text(encoding="utf-8")
call_before = '''                            outcome = _decomp.decompose_task(
                                tid, author="auto-decomposer",
                            )
'''
call_after = '''                            outcome = _decomp.decompose_task(
                                tid,
                                author="auto-decomposer",
                                automatic=True,
                            )
'''
watcher = replace_once(
    watcher,
    call_before,
    call_after,
    "gateway automatic-decomposition call",
)
watcher_path.write_text(watcher, encoding="utf-8")