Lab 08 · Durable refund workflow

langgraph==1.2.11
langgraph-checkpoint==4.2.0
langgraph-checkpoint-sqlite==3.1.1

"""Lab 08: a refund workflow that pauses for a manager and survives a restart.

Plain-Python nodes on a SQLite checkpointer: no model, no API key.
Run:  python refund_workflow.py
"""

import operator
from collections import Counter
from pathlib import Path
from typing import Annotated, Any, Literal

from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, START, StateGraph
from langgraph.graph.state import CompiledStateGraph
from langgraph.types import Command, interrupt
from typing_extensions import TypedDict

DB_PATH = Path(__file__).resolve().parent / "refund_lab.db"
APPROVAL_THRESHOLD = 100.0
HOLDS: dict[str, str] = {}  # stands in for the payments API
LEDGER: dict[str, str] = {}  # stands in for the ledger
NODE_CALLS: Counter = Counter()  # how often each node BODY ran

class RefundState(TypedDict):
    """One refund request moving through the workflow."""

    request_id: str
    customer: str
    amount: float
    decision: str  # "" | "auto" | "approved" | "rejected"
    receipt_id: str
    audit: Annotated[list[str], operator.add]

def validate(state: RefundState) -> dict[str, Any]:
    """Pure, so replaying it costs nothing. Small refunds skip the manager."""
    auto = state["amount"] <= APPROVAL_THRESHOLD
    return {"decision": "auto" if auto else "",
            "audit": [f"validated {state['request_id']} for {state['amount']:.2f}"]}

def await_approval(state: RefundState) -> dict[str, Any]:
    """Pause for a manager. The hold is keyed because this body re-runs from line 1."""
    NODE_CALLS["await_approval"] += 1
    HOLDS.setdefault(state["request_id"], f"HOLD-{state['request_id']}")
    answer = interrupt({"type": "refund_approval", "request_id": state["request_id"],
                        "customer": state["customer"], "amount": state["amount"]})
    decision = "approved" if answer["approved"] else "rejected"
    return {"decision": decision, "audit": [f"{answer['approver']} {decision} the refund"]}

def apply_refund(state: RefundState) -> dict[str, Any]:
    """Move the money. Keyed by request id, so a resume or a replay pays out once."""
    NODE_CALLS["apply_refund"] += 1
    receipt = LEDGER.setdefault(state["request_id"], f"RCPT-{state['request_id']}")
    return {"receipt_id": receipt, "audit": [f"paid {state['amount']:.2f} as {receipt}"]}


def notify_customer(state: RefundState) -> dict[str, Any]:
    """Tell the customer the outcome."""
    return {"audit": [f"emailed {state['customer']}: {state['decision']}"]}

def route_after_validation(state: RefundState) -> Literal["await_approval", "apply_refund"]:
    """Small refunds are paid straight away; the rest wait for a manager."""
    return "apply_refund" if state["decision"] == "auto" else "await_approval"


def route_after_approval(state: RefundState) -> Literal["apply_refund", "notify_customer"]:
    """Approved refunds are paid; rejected ones only notify the customer."""
    return "apply_refund" if state["decision"] == "approved" else "notify_customer"

def build_graph(checkpointer: BaseCheckpointSaver | None = None) -> CompiledStateGraph:
    """Compile the workflow. Without a checkpointer, interrupt() cannot pause."""
    builder = StateGraph(RefundState)
    builder.add_node("validate", validate)
    builder.add_node("await_approval", await_approval)
    builder.add_node("apply_refund", apply_refund)
    builder.add_node("notify_customer", notify_customer)
    builder.add_edge(START, "validate")
    builder.add_conditional_edges("validate", route_after_validation,
                                  ["await_approval", "apply_refund"])
    builder.add_conditional_edges("await_approval", route_after_approval,
                                  ["apply_refund", "notify_customer"])
    builder.add_edge("apply_refund", "notify_customer")
    builder.add_edge("notify_customer", END)
    return builder.compile(checkpointer=checkpointer)

def first_process(request: RefundState, config: dict[str, Any]) -> None:
    """Start the refund and stop at the manager gate, then close the database."""
    with SqliteSaver.from_conn_string(str(DB_PATH)) as saver:
        graph = build_graph(saver)
        paused = graph.invoke(request, config)
        interrupts = paused.get("__interrupt__", [])
        if not interrupts:
            print("The workflow did not pause. Is the amount above APPROVAL_THRESHOLD?")
            return
        print("paused with:", interrupts[0].value)
        print("hold placed before the pause:", HOLDS)

def second_process(config: dict[str, Any]) -> None:
    """Open the same database with a new saver, resume with approval, then fork a rejection."""
    with SqliteSaver.from_conn_string(str(DB_PATH)) as saver:
        graph = build_graph(saver)
        print("recovered, next:", graph.get_state(config).next)
        final = graph.invoke(Command(resume={"approved": True, "approver": "m.chen"}), config)
        print("decision:", final["decision"], "| receipt:", final["receipt_id"])
        print("await_approval ran", NODE_CALLS["await_approval"], "times | HOLDS:", HOLDS)
        print("audit:", final["audit"])

        # Time travel: fork the checkpoint taken just before the manager decided.
        before = next(s for s in graph.get_state_history(config) if s.next == ("await_approval",))
        fork_config = graph.update_state(before.config, None)
        other = graph.invoke(Command(resume={"approved": False, "approver": "director"}), fork_config)
        print("forked ->", other["decision"], "| ledger still:", LEDGER)
        print("history now holds", len(list(graph.get_state_history(config))), "checkpoints")

if __name__ == "__main__":
    DB_PATH.unlink(missing_ok=True)  # start from an empty checkpoint store
    request: RefundState = {"request_id": "R-2", "customer": "grace@example.com", "amount": 640.0,
                            "decision": "", "receipt_id": "", "audit": []}
    config = {"configurable": {"thread_id": "R-2"}}

    first_process(request, config)
    print("--- the process exits; the SQLite file keeps the paused thread ---")
    second_process(config)

paused with: {'type': 'refund_approval', 'request_id': 'R-2', 'customer': 'grace@example.com', 'amount': 640.0}
hold placed before the pause: {'R-2': 'HOLD-R-2'}
--- the process exits; the SQLite file keeps the paused thread ---
recovered, next: ('await_approval',)
decision: approved | receipt: RCPT-R-2
await_approval ran 2 times | HOLDS: {'R-2': 'HOLD-R-2'}
audit: ['validated R-2 for 640.00', 'm.chen approved the refund', 'paid 640.00 as RCPT-R-2', 'emailed grace@example.com: approved']
forked -> rejected | ledger still: {'R-2': 'RCPT-R-2'}
history now holds 9 checkpoints