LangGraph Thread ID & Checkpoint Retrieval Failures

from fastapi import FastAPI, Cookie, Response
from langchain_core.messages import HumanMessage
import uuid

api = FastAPI()


# ❌ ANTI-PATTERN — new thread every request
@api.post("/chat-broken")
async def chat_broken(message: str):
    thread_id = str(uuid.uuid4())          # <-- fresh UUID every call
    config = {"configurable": {"thread_id": thread_id}}
    return await api.state.graph.ainvoke(
        {"messages": [HumanMessage(content=message)]},
        config,
    )   # No history ever — every turn is turn 1


# ✅ Correct — thread_id from a persistent cookie
@api.post("/chat")
async def chat(
    response: Response,
    message: str,
    session_id: str | None = Cookie(default=None),
):
    if not session_id:
        # First request — mint a session and set the cookie
        session_id = str(uuid.uuid4())
        response.set_cookie("session_id", session_id, httponly=True, max_age=3600 * 24 * 30)

    config = {"configurable": {"thread_id": session_id}}
    return await api.state.graph.ainvoke(
        {"messages": [HumanMessage(content=message)]},
        config,
    )


# ✅ Better — thread_id derives from user + chat_id, stored in your DB
@api.post("/chat/{chat_id}")
async def chat_by_id(chat_id: str, user_id: str, message: str):
    # thread_id is deterministic and portable across servers
    thread_id = f"{user_id}::{chat_id}"
    config = {"configurable": {"thread_id": thread_id}}
    return await api.state.graph.ainvoke(
        {"messages": [HumanMessage(content=message)]},
        config,
    )


# ✅ Check thread state before invoking (useful for "resume vs new" UI)
@api.get("/chat/{chat_id}/history")
async def get_history(chat_id: str, user_id: str):
    config = {"configurable": {"thread_id": f"{user_id}::{chat_id}"}}
    state = await api.state.graph.aget_state(config)
    if state is None:
        return {"exists": False, "messages": []}
    return {
        "exists": True,
        "checkpoint_id": state.config["configurable"]["checkpoint_id"],
        "messages": [
            {"role": m.type, "content": m.content} for m in state.values.get("messages", [])
        ],
    }

# ✅ Postgres — every worker and every process shares the same store
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from fastapi import FastAPI
from contextlib import asynccontextmanager
import os


@asynccontextmanager
async def lifespan(app: FastAPI):
    dsn = os.environ["DATABASE_URL"].replace("postgres://", "postgresql://", 1)
    pool = AsyncConnectionPool(dsn, min_size=2, max_size=20)
    await pool.open()
    checkpointer = AsyncPostgresSaver(pool)
    await checkpointer.setup()             # idempotent
    app.state.graph = build_graph(checkpointer)
    yield
    await pool.close()

api = FastAPI(lifespan=lifespan)


# ✅ SQLite — works across restarts (single writer, single process only)
from langgraph.checkpoint.sqlite import SqliteSaver

@asynccontextmanager
async def lifespan_sqlite(app: FastAPI):
    import sqlite3
    conn = sqlite3.connect("/var/data/langgraph.db", check_same_thread=False)
    checkpointer = SqliteSaver(conn)
    checkpointer.setup()
    app.state.graph = build_graph(checkpointer)
    yield
    conn.close()


# ✅ MemorySaver — ONLY for tests / single-process notebooks
from langgraph.checkpoint.memory import MemorySaver
if __name__ == "__main__" and os.environ.get("MODE") == "test":
    checkpointer = MemorySaver()
    graph = build_graph(checkpointer)
    # ... tests run here, data lost when process exits ...


# ✅ Inspect what's in Postgres for a thread — useful when debugging "None"
async def debug_thread(thread_id: str):
    async with AsyncPostgresSaver.from_conn_string(os.environ["DATABASE_URL"]) as saver:
        config = {"configurable": {"thread_id": thread_id}}
        count = 0
        async for tup in saver.alist(config):
            count += 1
            print(f"checkpoint {tup.checkpoint['id']}: {tup.metadata.get('step')}")
        if count == 0:
            print(f"No checkpoints for thread {thread_id!r} — check the ID.")

# ✅ Get the latest state snapshot for a thread
config = {"configurable": {"thread_id": "user-42"}}
snapshot = app.get_state(config)

if snapshot is None:
    print(f"Thread {config['configurable']['thread_id']} has no checkpoints.")
else:
    print("Latest checkpoint:", snapshot.config["configurable"]["checkpoint_id"])
    print("Values:", snapshot.values)
    print("Next nodes to run:", snapshot.next)
    print("Metadata:", snapshot.metadata)


# ✅ Walk history — newest first, one snapshot per graph step
for snap in app.get_state_history(config):
    print(
        snap.config["configurable"]["checkpoint_id"],
        "step:", snap.metadata.get("step"),
        "source:", snap.metadata.get("source"),   # "input" | "loop" | "update"
        "next:", snap.next,
    )


# ✅ Time-travel — resume from a SPECIFIC historic checkpoint
# 1. Get history and pick a checkpoint from before the buggy turn
history = list(app.get_state_history(config))
target = next(s for s in history if s.metadata.get("step") == 5)

# 2. Build a config that points at that specific checkpoint
resume_config = {
    "configurable": {
        "thread_id": "user-42",
        "checkpoint_id": target.config["configurable"]["checkpoint_id"],
    },
}

# 3. Invoke — the graph resumes from that point, creating a new branch
result = app.invoke(None, resume_config)   # None = don't add new input, just resume


# ✅ Manually update state at a specific checkpoint (correct a bad message)
from langchain_core.messages import AIMessage, RemoveMessage
new_checkpoint_id = app.update_state(
    resume_config,
    {"messages": [RemoveMessage(id="msg_bad_id")]},
)
print("Wrote correction at:", new_checkpoint_id)


# ✅ List checkpoints filtered by metadata (LangGraph 0.2+)
recent = app.get_state_history(
    config,
    limit=5,
    filter={"source": "input"},          # only user-input checkpoints
)