Cancellation and cleanup

import asyncio
from typing import Any
from langgraph.graph import StateGraph, START, END
from langgraph.types import StateSnapshot
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import START as LANGGRAPH_START

# Simulating a long-running operation that can be cancelled
class ResourceManager:
    def __init__(self):
        self.resources = []
        self.is_active = True

    async def acquire_resource(self, name: str):
        if not self.is_active:
            raise RuntimeError(f"Cannot acquire {name}: manager is shut down")
        self.resources.append(name)
        print(f"Acquired resource: {name}")
        await asyncio.sleep(0.1)
        return name

    async def cleanup(self):
        self.is_active = False
        for resource in self.resources:
            print(f"Cleaning up resource: {resource}")
            await asyncio.sleep(0.05)
        self.resources.clear()

resource_manager = ResourceManager()

def node_a(state: dict[str, Any]) -> dict[str, Any]:
    """Expensive operation that acquires resources."""
    print("Node A: Starting expensive task")
    # In async context, this would be awaited
    state["step_a_complete"] = True
    state["acquired"] = "database_connection"
    print("Node A: Acquired database_connection")
    return state

def node_b(state: dict[str, Any]) -> dict[str, Any]:
    """Second operation that depends on A."""
    print("Node B: Processing with acquired resource")
    state["step_b_complete"] = True
    return state

def cleanup_handler(state: dict[str, Any]) -> dict[str, Any]:
    """Cleanup node that runs on cancellation."""
    print("Cleanup: Releasing all resources")
    if "acquired" in state:
        print(f"Cleanup: Releasing {state['acquired']}")
    state["cleaned_up"] = True
    return state

class State(dict):
    pass

# Build the graph with interrupt points
graph_builder = StateGraph(State)
graph_builder.add_node("node_a", node_a)
graph_builder.add_node("node_b", node_b)
graph_builder.add_node("cleanup", cleanup_handler)

graph_builder.add_edge(START, "node_a")
graph_builder.add_edge("node_a", "node_b")
graph_builder.add_edge("node_b", END)

# Compile with memory checkpointer to support interruption
graph = graph_builder.compile(checkpointer=MemorySaver())

# Run normally
print("=== Normal execution ===")
result = graph.invoke(
    {"step_a_complete": False, "step_b_complete": False},
    config={"configurable": {"thread_id": "test_1"}}
)
print(f"Final state: {dict(result)}")
print()

# Simulate interruption at a specific point
print("=== Execution with interrupt point ===")
graph_with_interrupt = StateGraph(State)
graph_with_interrupt.add_node("node_a", node_a)
graph_with_interrupt.add_node("node_b", node_b)
graph_with_interrupt.add_node("cleanup", cleanup_handler)

graph_with_interrupt.add_edge(START, "node_a")
# Interrupt BEFORE node_b to allow decision
graph_with_interrupt.add_edge("node_a", "node_b")
graph_with_interrupt.add_edge("node_b", "cleanup")
graph_with_interrupt.add_edge("cleanup", END)

compiled = graph_with_interrupt.compile(
    checkpointer=MemorySaver(),
    interrupt_before=["node_b"]
)

config = {"configurable": {"thread_id": "test_2"}}
state_snapshot = compiled.invoke(
    {"step_a_complete": False, "step_b_complete": False},
    config=config
)

print(f"Interrupted state: {dict(state_snapshot)}")
print(f"Graph halted before: node_b")
print()

# Demonstrate cleanup on early termination
print("=== Cleanup on cancellation ===")
print("Simulating user cancellation during execution...")
state_with_cleanup = compiled.invoke(
    {"step_a_complete": False, "step_b_complete": False, "acquired": "api_token"},
    config={"configurable": {"thread_id": "test_3"}}
)
print("Cleanup node would execute here to release api_token")