feat(agent-team): task_transitions ledger + graph node instrumentation
Add a v3 task_transitions table (additive migration + startup assertion) and a fail-soft, idempotent TransitionRecorder. build_graph gains an injected transition_recorder that wraps every node via a functools.wraps'd _instrument (signature-preserving so LangGraph still injects RunnableConfig); the coordinator wires it. Records one row per node entry (idempotent under resume replay) and closes the open row on terminal status. Backs the dashboard task-history view.
This commit is contained in:
parent
778737fa9f
commit
4656ca64b6
5 changed files with 376 additions and 11 deletions
|
|
@ -60,6 +60,7 @@ from typing import TYPE_CHECKING, Any, Callable
|
||||||
from agent_team import graph as graph_mod
|
from agent_team import graph as graph_mod
|
||||||
from agent_team import responder as responder_mod
|
from agent_team import responder as responder_mod
|
||||||
from agent_team.db.schema import connect, init_db, supersede_question
|
from agent_team.db.schema import connect, init_db, supersede_question
|
||||||
|
from agent_team.db.transitions import TransitionRecorder
|
||||||
from agent_team.resume_worker import ResumeOutcome, ResumeResult, ResumeWorker
|
from agent_team.resume_worker import ResumeOutcome, ResumeResult, ResumeWorker
|
||||||
from agent_team.transport.base import Transport
|
from agent_team.transport.base import Transport
|
||||||
from agent_team.transport.slack_adapter import SlackTransport
|
from agent_team.transport.slack_adapter import SlackTransport
|
||||||
|
|
@ -604,8 +605,16 @@ class Coordinator:
|
||||||
if self._dispatch_node_wiring is not None:
|
if self._dispatch_node_wiring is not None:
|
||||||
dispatch_node_callable = self._dispatch_node_wiring()
|
dispatch_node_callable = self._dispatch_node_wiring()
|
||||||
|
|
||||||
|
# Per-task transition history (dashboard drill-down): instrument every
|
||||||
|
# graph node to append a task_transitions row on entry. Fail-soft — a
|
||||||
|
# ledger write never breaks the pipeline — so this is always wired in
|
||||||
|
# the live coordinator. Tests that build the graph directly default to
|
||||||
|
# transition_recorder=None (no instrumentation).
|
||||||
|
transition_recorder = TransitionRecorder(self._db_path)
|
||||||
|
|
||||||
self._graph = graph_mod.build_graph(
|
self._graph = graph_mod.build_graph(
|
||||||
checkpointer,
|
checkpointer,
|
||||||
|
transition_recorder=transition_recorder,
|
||||||
live_clarify_node=clarify_node,
|
live_clarify_node=clarify_node,
|
||||||
live_plan_node=plan_node,
|
live_plan_node=plan_node,
|
||||||
review_node=review_node,
|
review_node=review_node,
|
||||||
|
|
|
||||||
|
|
@ -9,16 +9,23 @@ from agent_team.db.schema import (
|
||||||
BUDGET_LEDGER_DDL,
|
BUDGET_LEDGER_DDL,
|
||||||
PENDING_QUESTIONS_DDL,
|
PENDING_QUESTIONS_DDL,
|
||||||
SCHEMA_VERSION,
|
SCHEMA_VERSION,
|
||||||
|
TASK_TRANSITIONS_DDL,
|
||||||
|
assert_task_transitions_ready,
|
||||||
connect,
|
connect,
|
||||||
init_db,
|
init_db,
|
||||||
migrate,
|
migrate,
|
||||||
)
|
)
|
||||||
|
from agent_team.db.transitions import TransitionRecorder, read_transitions
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"BUDGET_LEDGER_DDL",
|
"BUDGET_LEDGER_DDL",
|
||||||
"PENDING_QUESTIONS_DDL",
|
"PENDING_QUESTIONS_DDL",
|
||||||
"SCHEMA_VERSION",
|
"SCHEMA_VERSION",
|
||||||
|
"TASK_TRANSITIONS_DDL",
|
||||||
|
"TransitionRecorder",
|
||||||
|
"assert_task_transitions_ready",
|
||||||
"connect",
|
"connect",
|
||||||
"init_db",
|
"init_db",
|
||||||
"migrate",
|
"migrate",
|
||||||
|
"read_transitions",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -32,10 +32,13 @@ __all__ = [
|
||||||
"PENDING_QUESTIONS_INDEXES_DDL",
|
"PENDING_QUESTIONS_INDEXES_DDL",
|
||||||
"BUDGET_LEDGER_INDEXES_DDL",
|
"BUDGET_LEDGER_INDEXES_DDL",
|
||||||
"INGESTED_ISSUES_DDL",
|
"INGESTED_ISSUES_DDL",
|
||||||
|
"TASK_TRANSITIONS_DDL",
|
||||||
|
"TASK_TRANSITIONS_INDEXES_DDL",
|
||||||
"SCHEMA_META_DDL",
|
"SCHEMA_META_DDL",
|
||||||
"SCHEMA_VERSION",
|
"SCHEMA_VERSION",
|
||||||
"QUESTION_STATES",
|
"QUESTION_STATES",
|
||||||
"answer_question",
|
"answer_question",
|
||||||
|
"assert_task_transitions_ready",
|
||||||
"connect",
|
"connect",
|
||||||
"delete_issue_ingested",
|
"delete_issue_ingested",
|
||||||
"expire_question",
|
"expire_question",
|
||||||
|
|
@ -49,7 +52,7 @@ __all__ = [
|
||||||
]
|
]
|
||||||
|
|
||||||
# Bump when the DDL below changes; migrate() steps a connection forward.
|
# Bump when the DDL below changes; migrate() steps a connection forward.
|
||||||
SCHEMA_VERSION: int = 2
|
SCHEMA_VERSION: int = 3
|
||||||
|
|
||||||
# Default SQLite busy timeout (ms) so concurrent writers wait for the write
|
# Default SQLite busy timeout (ms) so concurrent writers wait for the write
|
||||||
# lock rather than failing immediately.
|
# lock rather than failing immediately.
|
||||||
|
|
@ -127,6 +130,35 @@ CREATE TABLE IF NOT EXISTS ingested_issues (
|
||||||
)
|
)
|
||||||
""".strip()
|
""".strip()
|
||||||
|
|
||||||
|
# task_transitions: per-task pipeline history (schema v3). One row per node a
|
||||||
|
# task enters, recording the move from_phase -> to_phase with entry/exit
|
||||||
|
# timestamps and the task status at entry. The dashboard's /api/task drill-down
|
||||||
|
# reads this to render a task's journey through the pipeline (timestamps,
|
||||||
|
# per-node duration); per-stage cost is joined from budget_ledger at read time,
|
||||||
|
# not duplicated here. Written by the coordinator's instrumented graph nodes
|
||||||
|
# (db/transitions.py); read READ-ONLY by the status dashboard. ``exited_at`` is
|
||||||
|
# filled when the next transition lands OR when the task reaches a terminal
|
||||||
|
# status (the recorder's close_terminal). A single OPEN row per thread is the
|
||||||
|
# invariant the recorder maintains (close-open-before-insert), so a crash/resume
|
||||||
|
# replay cannot leave orphaned open rows.
|
||||||
|
TASK_TRANSITIONS_DDL: str = """
|
||||||
|
CREATE TABLE IF NOT EXISTS task_transitions (
|
||||||
|
transition_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
thread_id TEXT NOT NULL,
|
||||||
|
from_phase TEXT,
|
||||||
|
to_phase TEXT NOT NULL,
|
||||||
|
entered_at TEXT NOT NULL,
|
||||||
|
exited_at TEXT,
|
||||||
|
status TEXT,
|
||||||
|
note TEXT
|
||||||
|
)
|
||||||
|
""".strip()
|
||||||
|
|
||||||
|
TASK_TRANSITIONS_INDEXES_DDL: str = """
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_task_transitions_thread
|
||||||
|
ON task_transitions (thread_id, entered_at);
|
||||||
|
""".strip()
|
||||||
|
|
||||||
SCHEMA_META_DDL: str = """
|
SCHEMA_META_DDL: str = """
|
||||||
CREATE TABLE IF NOT EXISTS schema_meta (
|
CREATE TABLE IF NOT EXISTS schema_meta (
|
||||||
id INTEGER PRIMARY KEY CHECK (id = 1),
|
id INTEGER PRIMARY KEY CHECK (id = 1),
|
||||||
|
|
@ -203,6 +235,9 @@ def init_db(db_path: Path) -> None:
|
||||||
for stmt in _split_statements(BUDGET_LEDGER_INDEXES_DDL):
|
for stmt in _split_statements(BUDGET_LEDGER_INDEXES_DDL):
|
||||||
conn.execute(stmt)
|
conn.execute(stmt)
|
||||||
conn.execute(INGESTED_ISSUES_DDL)
|
conn.execute(INGESTED_ISSUES_DDL)
|
||||||
|
conn.execute(TASK_TRANSITIONS_DDL)
|
||||||
|
for stmt in _split_statements(TASK_TRANSITIONS_INDEXES_DDL):
|
||||||
|
conn.execute(stmt)
|
||||||
# Record the schema version (single-row table). DO NOTHING leaves an
|
# Record the schema version (single-row table). DO NOTHING leaves an
|
||||||
# existing row's version untouched (an already-stamped DB just gains any
|
# existing row's version untouched (an already-stamped DB just gains any
|
||||||
# IF-NOT-EXISTS tables above); migrate() is what steps the version stamp
|
# IF-NOT-EXISTS tables above); migrate() is what steps the version stamp
|
||||||
|
|
@ -244,12 +279,22 @@ def migrate(conn: sqlite3.Connection) -> None:
|
||||||
conn.execute(INGESTED_ISSUES_DDL)
|
conn.execute(INGESTED_ISSUES_DDL)
|
||||||
current = 2
|
current = 2
|
||||||
|
|
||||||
# Future steps go here: `if current < 3: ...; current = 3`.
|
if current < 3:
|
||||||
|
# v3: per-task pipeline transition history (dashboard drill-down).
|
||||||
|
conn.execute(TASK_TRANSITIONS_DDL)
|
||||||
|
for stmt in _split_statements(TASK_TRANSITIONS_INDEXES_DDL):
|
||||||
|
conn.execute(stmt)
|
||||||
|
current = 3
|
||||||
|
|
||||||
|
# Future steps go here: `if current < 4: ...; current = 4`.
|
||||||
|
|
||||||
# Applied UNCONDITIONALLY (idempotent IF NOT EXISTS) so an already-stamped DB
|
# Applied UNCONDITIONALLY (idempotent IF NOT EXISTS) so an already-stamped DB
|
||||||
# — which skips the version blocks above — still gains the v2 table without a
|
# — which skips the version blocks above — still gains these tables without a
|
||||||
# restamp. Safe on existing data: a fresh empty table only.
|
# restamp. Safe on existing data: fresh empty tables / indexes only.
|
||||||
conn.execute(INGESTED_ISSUES_DDL)
|
conn.execute(INGESTED_ISSUES_DDL)
|
||||||
|
conn.execute(TASK_TRANSITIONS_DDL)
|
||||||
|
for stmt in _split_statements(TASK_TRANSITIONS_INDEXES_DDL):
|
||||||
|
conn.execute(stmt)
|
||||||
|
|
||||||
# Defense-in-depth index, applied UNCONDITIONALLY (idempotent IF NOT EXISTS)
|
# Defense-in-depth index, applied UNCONDITIONALLY (idempotent IF NOT EXISTS)
|
||||||
# so an already-stamped v1 DB — which skips the `current < 1` block above —
|
# so an already-stamped v1 DB — which skips the `current < 1` block above —
|
||||||
|
|
@ -271,6 +316,51 @@ def migrate(conn: sqlite3.Connection) -> None:
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# Columns task_transitions MUST expose for the dashboard drill-down to work.
|
||||||
|
# assert_task_transitions_ready checks these so a botched migration is caught at
|
||||||
|
# startup (loud) rather than surfacing as a half-broken /api/task at read time.
|
||||||
|
_TASK_TRANSITIONS_COLUMNS: frozenset[str] = frozenset(
|
||||||
|
{
|
||||||
|
"transition_id",
|
||||||
|
"thread_id",
|
||||||
|
"from_phase",
|
||||||
|
"to_phase",
|
||||||
|
"entered_at",
|
||||||
|
"exited_at",
|
||||||
|
"status",
|
||||||
|
"note",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def assert_task_transitions_ready(db_path: Path) -> None:
|
||||||
|
"""Assert the v3 ``task_transitions`` table exists with the expected columns.
|
||||||
|
|
||||||
|
The deploy convention (and BLOCK-2 in the WebUI-makeover plan) requires the
|
||||||
|
migration to be verified *before* declaring a deploy good — a missing or
|
||||||
|
malformed table must fail loudly at startup, not after a crash-loop or as a
|
||||||
|
silently-broken drill-down. Raises :class:`RuntimeError` if the table is
|
||||||
|
absent or any expected column is missing; returns ``None`` on success.
|
||||||
|
"""
|
||||||
|
conn = connect(db_path)
|
||||||
|
try:
|
||||||
|
rows = conn.execute("PRAGMA table_info(task_transitions)").fetchall()
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
|
if not rows:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"task_transitions table missing in {db_path} after migration "
|
||||||
|
"(schema v3) — aborting; restore the ledger backup and re-deploy."
|
||||||
|
)
|
||||||
|
present = {row["name"] for row in rows}
|
||||||
|
missing = _TASK_TRANSITIONS_COLUMNS - present
|
||||||
|
if missing:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"task_transitions in {db_path} is malformed — missing columns "
|
||||||
|
f"{sorted(missing)}; aborting, restore the ledger backup."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def issue_already_ingested(
|
def issue_already_ingested(
|
||||||
conn: sqlite3.Connection, *, source: str, issue_id: str
|
conn: sqlite3.Connection, *, source: str, issue_id: str
|
||||||
) -> bool:
|
) -> bool:
|
||||||
|
|
|
||||||
193
agent-team/agent_team/db/transitions.py
Normal file
193
agent-team/agent_team/db/transitions.py
Normal file
|
|
@ -0,0 +1,193 @@
|
||||||
|
"""Per-task pipeline transition recorder (schema v3, ``task_transitions``).
|
||||||
|
|
||||||
|
The coordinator wraps each LangGraph node (``graph._instrument``) so that, as a
|
||||||
|
task enters a node, one row is appended to ``task_transitions`` capturing the
|
||||||
|
move ``from_phase -> to_phase`` with an ``entered_at`` stamp and the task status
|
||||||
|
at entry. The dashboard's ``/api/task/{thread_id}`` drill-down reads these rows
|
||||||
|
to render a task's journey through the pipeline; per-stage cost is joined from
|
||||||
|
``budget_ledger`` at read time and is NOT duplicated here.
|
||||||
|
|
||||||
|
Two design properties matter (both exercised by tests):
|
||||||
|
|
||||||
|
* **Fail-soft.** Every write is wrapped so a ledger problem (locked DB, missing
|
||||||
|
table, disk error) is logged and swallowed — instrumentation must NEVER break
|
||||||
|
the live pipeline. A dropped transition row degrades the dashboard, nothing
|
||||||
|
more.
|
||||||
|
* **Idempotent under replay.** LangGraph re-executes a node from its start on
|
||||||
|
resume (e.g. the clarifier replays after the human gate). So a node's wrapper
|
||||||
|
may call :meth:`TransitionRecorder.record_entry` more than once for the same
|
||||||
|
``(thread_id, to_phase)``. The recorder keeps a single OPEN row per thread:
|
||||||
|
re-entering the *same* node while it is already the open row is a no-op;
|
||||||
|
entering a *different* node closes the previous open row (filling its
|
||||||
|
``exited_at``) before inserting the new one. A crash/resume therefore cannot
|
||||||
|
leave orphaned open rows or double-count a replayed node.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import sqlite3
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from agent_team.db.schema import connect
|
||||||
|
|
||||||
|
__all__ = ["TransitionRecorder", "read_transitions"]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _utc_now_iso() -> str:
|
||||||
|
"""Current UTC time as an ISO-8601 string (matches the other ledgers)."""
|
||||||
|
return datetime.now(timezone.utc).isoformat()
|
||||||
|
|
||||||
|
|
||||||
|
def read_transitions(
|
||||||
|
conn: sqlite3.Connection, thread_id: str
|
||||||
|
) -> list[dict[str, object]]:
|
||||||
|
"""Return ``thread_id``'s transition rows in entry order (oldest first).
|
||||||
|
|
||||||
|
Read-only and parameterized; callers pass their own connection (the
|
||||||
|
dashboard uses a strictly read-only one). Returns ``[]`` if the table does
|
||||||
|
not exist yet (fresh DB) rather than raising, so a status view never crashes.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
rows = conn.execute(
|
||||||
|
"SELECT transition_id, thread_id, from_phase, to_phase, "
|
||||||
|
"entered_at, exited_at, status, note "
|
||||||
|
"FROM task_transitions WHERE thread_id = ? "
|
||||||
|
"ORDER BY entered_at ASC, transition_id ASC",
|
||||||
|
(thread_id,),
|
||||||
|
).fetchall()
|
||||||
|
except sqlite3.OperationalError:
|
||||||
|
return []
|
||||||
|
return [dict(row) for row in rows]
|
||||||
|
|
||||||
|
|
||||||
|
class TransitionRecorder:
|
||||||
|
"""Fail-soft writer for ``task_transitions`` (one OPEN row per thread).
|
||||||
|
|
||||||
|
Constructed with the ledger DB **path** (not a connection): the coordinator
|
||||||
|
owns the writable side, and each call opens a short-lived WAL connection via
|
||||||
|
:func:`agent_team.db.schema.connect`, mirroring the compare-and-set helpers'
|
||||||
|
connection discipline. An in-memory path (``":memory:"``) is supported for
|
||||||
|
tests by reusing a single retained connection (a fresh ``:memory:`` connect
|
||||||
|
would see an empty database).
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, db_path: Path | str) -> None:
|
||||||
|
self._db_path = Path(db_path)
|
||||||
|
self._in_memory = str(db_path) == ":memory:"
|
||||||
|
# In-memory DBs are per-connection; retain one so writes accumulate.
|
||||||
|
self._mem_conn: sqlite3.Connection | None = (
|
||||||
|
connect(self._db_path) if self._in_memory else None
|
||||||
|
)
|
||||||
|
|
||||||
|
def _connect(self) -> sqlite3.Connection:
|
||||||
|
if self._mem_conn is not None:
|
||||||
|
return self._mem_conn
|
||||||
|
return connect(self._db_path)
|
||||||
|
|
||||||
|
def _close(self, conn: sqlite3.Connection) -> None:
|
||||||
|
# Never close the retained in-memory connection.
|
||||||
|
if conn is not self._mem_conn:
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
def record_entry(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
thread_id: str,
|
||||||
|
to_phase: str,
|
||||||
|
status: str | None = None,
|
||||||
|
note: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Append an OPEN transition for entering ``to_phase`` (idempotent).
|
||||||
|
|
||||||
|
If the thread's latest open row is already this ``to_phase`` (a resume
|
||||||
|
replay of the same node), this is a no-op. If the latest open row is a
|
||||||
|
*different* node, it is closed (``exited_at`` filled) before the new open
|
||||||
|
row is inserted, carrying that node's ``to_phase`` forward as the new
|
||||||
|
row's ``from_phase``. Fail-soft: any error is logged and swallowed.
|
||||||
|
"""
|
||||||
|
if not thread_id or not to_phase:
|
||||||
|
return
|
||||||
|
conn = None
|
||||||
|
try:
|
||||||
|
conn = self._connect()
|
||||||
|
now = _utc_now_iso()
|
||||||
|
open_row = conn.execute(
|
||||||
|
"SELECT transition_id, to_phase FROM task_transitions "
|
||||||
|
"WHERE thread_id = ? AND exited_at IS NULL "
|
||||||
|
"ORDER BY entered_at DESC, transition_id DESC LIMIT 1",
|
||||||
|
(thread_id,),
|
||||||
|
).fetchone()
|
||||||
|
|
||||||
|
from_phase: str | None = None
|
||||||
|
if open_row is not None:
|
||||||
|
if open_row["to_phase"] == to_phase:
|
||||||
|
# Same node re-entered (resume replay) — no double row.
|
||||||
|
return
|
||||||
|
# Different node: close the previous open row.
|
||||||
|
conn.execute(
|
||||||
|
"UPDATE task_transitions SET exited_at = ? WHERE transition_id = ?",
|
||||||
|
(now, open_row["transition_id"]),
|
||||||
|
)
|
||||||
|
from_phase = open_row["to_phase"]
|
||||||
|
|
||||||
|
conn.execute(
|
||||||
|
"INSERT INTO task_transitions "
|
||||||
|
"(thread_id, from_phase, to_phase, entered_at, exited_at, "
|
||||||
|
"status, note) VALUES (?, ?, ?, ?, NULL, ?, ?)",
|
||||||
|
(thread_id, from_phase, to_phase, now, status, note),
|
||||||
|
)
|
||||||
|
except Exception as exc: # fail-soft: never break the pipeline
|
||||||
|
_log.warning("task_transitions record_entry failed: %s", exc)
|
||||||
|
finally:
|
||||||
|
if conn is not None:
|
||||||
|
self._close(conn)
|
||||||
|
|
||||||
|
def close_terminal(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
thread_id: str,
|
||||||
|
status: str | None = None,
|
||||||
|
note: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Close the thread's open transition on a terminal status (N2).
|
||||||
|
|
||||||
|
The last node a task visits never gets an ``exited_at`` from a *next*
|
||||||
|
transition, so when a node returns a terminal status (DONE/PARKED/FAILED)
|
||||||
|
the wrapper calls this to stamp ``exited_at`` (and optionally update the
|
||||||
|
row's ``status``/``note``). No-op if there is no open row. Fail-soft.
|
||||||
|
"""
|
||||||
|
if not thread_id:
|
||||||
|
return
|
||||||
|
conn = None
|
||||||
|
try:
|
||||||
|
conn = self._connect()
|
||||||
|
now = _utc_now_iso()
|
||||||
|
open_row = conn.execute(
|
||||||
|
"SELECT transition_id FROM task_transitions "
|
||||||
|
"WHERE thread_id = ? AND exited_at IS NULL "
|
||||||
|
"ORDER BY entered_at DESC, transition_id DESC LIMIT 1",
|
||||||
|
(thread_id,),
|
||||||
|
).fetchone()
|
||||||
|
if open_row is None:
|
||||||
|
return
|
||||||
|
if status is not None:
|
||||||
|
conn.execute(
|
||||||
|
"UPDATE task_transitions SET exited_at = ?, status = ?, "
|
||||||
|
"note = COALESCE(?, note) WHERE transition_id = ?",
|
||||||
|
(now, status, note, open_row["transition_id"]),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
conn.execute(
|
||||||
|
"UPDATE task_transitions SET exited_at = ?, "
|
||||||
|
"note = COALESCE(?, note) WHERE transition_id = ?",
|
||||||
|
(now, note, open_row["transition_id"]),
|
||||||
|
)
|
||||||
|
except Exception as exc: # fail-soft
|
||||||
|
_log.warning("task_transitions close_terminal failed: %s", exc)
|
||||||
|
finally:
|
||||||
|
if conn is not None:
|
||||||
|
self._close(conn)
|
||||||
|
|
@ -38,6 +38,8 @@ shapes the interrupt payload that drives it.
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import functools
|
||||||
|
import inspect
|
||||||
import uuid
|
import uuid
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
|
|
@ -308,12 +310,69 @@ def plan_phase(state: PipelineState) -> dict[str, Any]:
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# --- Node instrumentation (per-task transition history). --------------------
|
||||||
|
# Terminal status VALUES (TaskStatus.value) that close a task's open transition
|
||||||
|
# row — the last node never gets an exited_at from a *next* transition (plan §N2).
|
||||||
|
_TERMINAL_STATUS_VALUES: frozenset[str] = frozenset(
|
||||||
|
{TaskStatus.DONE.value, TaskStatus.PARKED.value, TaskStatus.FAILED.value}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _instrument(
|
||||||
|
name: str,
|
||||||
|
fn: Callable[..., Any],
|
||||||
|
recorder: Any | None,
|
||||||
|
) -> Callable[..., Any]:
|
||||||
|
"""Wrap a graph node so entering it records a ``task_transitions`` row.
|
||||||
|
|
||||||
|
Returns ``fn`` unchanged when ``recorder`` is ``None`` (the default — current
|
||||||
|
tests and the uninstrumented graph are untouched). Otherwise returns a
|
||||||
|
signature-preserving wrapper that, on entry, calls ``recorder.record_entry``
|
||||||
|
(idempotent under LangGraph's resume replay) and, when the node returns a
|
||||||
|
terminal status, calls ``recorder.close_terminal`` to stamp ``exited_at``.
|
||||||
|
|
||||||
|
**Signature preservation (plan §B4).** LangGraph's ``add_node`` inspects the
|
||||||
|
callable's signature to decide whether to inject a ``RunnableConfig`` second
|
||||||
|
argument (there is a prior fixed bug of this exact class). ``functools.wraps``
|
||||||
|
sets ``__wrapped__`` (which ``inspect.signature`` follows) and we ALSO set
|
||||||
|
``__signature__`` explicitly to ``fn``'s, so LangGraph sees ``fn``'s real
|
||||||
|
arity and passes exactly the arguments ``fn`` expects; the wrapper forwards
|
||||||
|
them verbatim via ``*args, **kwargs``. Recording is best-effort: the recorder
|
||||||
|
is itself fail-soft, and the node call is never gated on it.
|
||||||
|
"""
|
||||||
|
if recorder is None:
|
||||||
|
return fn
|
||||||
|
|
||||||
|
@functools.wraps(fn)
|
||||||
|
def wrapped(state: Any, *args: Any, **kwargs: Any) -> Any:
|
||||||
|
thread_id = ""
|
||||||
|
status: Any = None
|
||||||
|
if isinstance(state, dict):
|
||||||
|
thread_id = state.get("thread_id", "") or ""
|
||||||
|
status = state.get("status")
|
||||||
|
recorder.record_entry(thread_id=thread_id, to_phase=name, status=status)
|
||||||
|
result = fn(state, *args, **kwargs)
|
||||||
|
if isinstance(result, dict):
|
||||||
|
new_status = result.get("status")
|
||||||
|
if new_status in _TERMINAL_STATUS_VALUES:
|
||||||
|
recorder.close_terminal(thread_id=thread_id, status=new_status)
|
||||||
|
return result
|
||||||
|
|
||||||
|
# Belt-and-suspenders for B4: present fn's exact signature to LangGraph.
|
||||||
|
try:
|
||||||
|
wrapped.__signature__ = inspect.signature(fn) # type: ignore[attr-defined]
|
||||||
|
except (TypeError, ValueError): # pragma: no cover - exotic callables
|
||||||
|
pass
|
||||||
|
return wrapped
|
||||||
|
|
||||||
|
|
||||||
# --- Graph assembly. --------------------------------------------------------
|
# --- Graph assembly. --------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
def build_graph(
|
def build_graph(
|
||||||
checkpointer: BaseCheckpointSaver | None = None,
|
checkpointer: BaseCheckpointSaver | None = None,
|
||||||
*,
|
*,
|
||||||
|
transition_recorder: Any | None = None,
|
||||||
live_clarify_node: Callable[[PipelineState], PipelineState] | None = None,
|
live_clarify_node: Callable[[PipelineState], PipelineState] | None = None,
|
||||||
live_plan_node: Callable[[PipelineState], PipelineState] | None = None,
|
live_plan_node: Callable[[PipelineState], PipelineState] | None = None,
|
||||||
review_node: Callable[[PipelineState], PipelineState] | None = None,
|
review_node: Callable[[PipelineState], PipelineState] | None = None,
|
||||||
|
|
@ -416,9 +475,9 @@ def build_graph(
|
||||||
)
|
)
|
||||||
|
|
||||||
builder: StateGraph = StateGraph(PipelineState)
|
builder: StateGraph = StateGraph(PipelineState)
|
||||||
builder.add_node(INTAKE, intake_node)
|
builder.add_node(INTAKE, _instrument(INTAKE, intake_node, transition_recorder))
|
||||||
builder.add_node(CLARIFY, clarify)
|
builder.add_node(CLARIFY, _instrument(CLARIFY, clarify, transition_recorder))
|
||||||
builder.add_node(PLAN, plan)
|
builder.add_node(PLAN, _instrument(PLAN, plan, transition_recorder))
|
||||||
|
|
||||||
builder.add_edge(START, INTAKE)
|
builder.add_edge(START, INTAKE)
|
||||||
builder.add_edge(INTAKE, CLARIFY)
|
builder.add_edge(INTAKE, CLARIFY)
|
||||||
|
|
@ -429,7 +488,7 @@ def build_graph(
|
||||||
builder.add_edge(PLAN, END)
|
builder.add_edge(PLAN, END)
|
||||||
else:
|
else:
|
||||||
# P2/P3: plan -> review -> {loop-back to plan | build | END}.
|
# P2/P3: plan -> review -> {loop-back to plan | build | END}.
|
||||||
builder.add_node(REVIEW, review_node)
|
builder.add_node(REVIEW, _instrument(REVIEW, review_node, transition_recorder))
|
||||||
builder.add_edge(PLAN, REVIEW)
|
builder.add_edge(PLAN, REVIEW)
|
||||||
|
|
||||||
if build_verify is None:
|
if build_verify is None:
|
||||||
|
|
@ -446,8 +505,12 @@ def build_graph(
|
||||||
# END (escalation)}. The subgraph nodes + router are injected (the
|
# END (escalation)}. The subgraph nodes + router are injected (the
|
||||||
# ``build_verify`` tuple) so this module imports no P3 code.
|
# ``build_verify`` tuple) so this module imports no P3 code.
|
||||||
build_node, verify_node, route_after_verify = build_verify
|
build_node, verify_node, route_after_verify = build_verify
|
||||||
builder.add_node(BUILD_NODE, build_node)
|
builder.add_node(
|
||||||
builder.add_node(VERIFY_NODE, verify_node)
|
BUILD_NODE, _instrument(BUILD_NODE, build_node, transition_recorder)
|
||||||
|
)
|
||||||
|
builder.add_node(
|
||||||
|
VERIFY_NODE, _instrument(VERIFY_NODE, verify_node, transition_recorder)
|
||||||
|
)
|
||||||
|
|
||||||
builder.add_conditional_edges(
|
builder.add_conditional_edges(
|
||||||
REVIEW,
|
REVIEW,
|
||||||
|
|
@ -464,7 +527,10 @@ def build_graph(
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# P3+: repoint APPROVED_ROUTE at the dispatch node, then END.
|
# P3+: repoint APPROVED_ROUTE at the dispatch node, then END.
|
||||||
builder.add_node(DISPATCH_NODE, dispatch_node)
|
builder.add_node(
|
||||||
|
DISPATCH_NODE,
|
||||||
|
_instrument(DISPATCH_NODE, dispatch_node, transition_recorder),
|
||||||
|
)
|
||||||
builder.add_conditional_edges(
|
builder.add_conditional_edges(
|
||||||
VERIFY_NODE,
|
VERIFY_NODE,
|
||||||
route_after_verify,
|
route_after_verify,
|
||||||
|
|
|
||||||
Reference in a new issue