m4
112 lines
5.3 KiB
Python
112 lines
5.3 KiB
Python
"""Local durable initialization gate. All graph access must use this gate.
|
|
|
|
The file lock covers graph writes as well as owner changes; a row CAS alone
|
|
cannot fence an old writer while it is suspended in a graph await.
|
|
"""
|
|
import asyncio
|
|
import fcntl
|
|
import sqlite3
|
|
from contextlib import asynccontextmanager, closing
|
|
from pathlib import Path
|
|
|
|
from langgraph.graph import START, END
|
|
|
|
|
|
class SqliteInitializationStore:
|
|
def __init__(self, path):
|
|
if str(path) == ":memory:":
|
|
raise ValueError("initialization requires a durable file")
|
|
self.path = str(Path(path).resolve())
|
|
with closing(self.connect()) as db, db:
|
|
db.execute("""CREATE TABLE IF NOT EXISTS history_initialization (
|
|
key TEXT PRIMARY KEY, attempt TEXT NOT NULL, digest TEXT NOT NULL,
|
|
owner TEXT NOT NULL, fence INTEGER NOT NULL, status TEXT NOT NULL)""")
|
|
|
|
def connect(self):
|
|
db = sqlite3.connect(self.path)
|
|
db.row_factory = sqlite3.Row
|
|
return db
|
|
|
|
def get(self, key):
|
|
with closing(self.connect()) as db:
|
|
row = db.execute("SELECT * FROM history_initialization WHERE key=?", (key,)).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
@asynccontextmanager
|
|
async def locked(self, key):
|
|
with open(self.path + "." + key.split(":")[-1] + ".lock", "a") as lock:
|
|
while True:
|
|
try:
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
break
|
|
except BlockingIOError:
|
|
await asyncio.sleep(0.01)
|
|
try:
|
|
yield
|
|
finally:
|
|
fcntl.flock(lock, fcntl.LOCK_UN)
|
|
|
|
def claim(self, key, attempt, digest, owner, fence):
|
|
if not isinstance(attempt, str) or not attempt.strip() or not isinstance(owner, str) or not owner.strip():
|
|
raise ValueError("attempt and owner required")
|
|
if type(fence) is not int or fence < 1:
|
|
raise ValueError("positive owner fence required")
|
|
with closing(self.connect()) as db, db:
|
|
db.execute("BEGIN IMMEDIATE")
|
|
row = db.execute("SELECT * FROM history_initialization WHERE key=?", (key,)).fetchone()
|
|
if row:
|
|
if row["attempt"] != attempt or row["digest"] != digest:
|
|
raise ValueError("history attempt/content conflict")
|
|
if fence < row["fence"] or (fence == row["fence"] and owner != row["owner"]):
|
|
raise ValueError("stale history owner fence")
|
|
if row["status"] not in ("INITIALIZING", "READY"):
|
|
raise ValueError("history checkpoint already consumed")
|
|
db.execute("UPDATE history_initialization SET owner=?, fence=? WHERE key=?", (owner, fence, key))
|
|
else:
|
|
db.execute("INSERT INTO history_initialization VALUES (?,?,?,?,?,?)",
|
|
(key, attempt, digest, owner, fence, "INITIALIZING"))
|
|
|
|
def status(self, key, attempt, owner, fence, before, after):
|
|
with closing(self.connect()) as db, db:
|
|
changed = db.execute("""UPDATE history_initialization SET status=?
|
|
WHERE key=? AND attempt=? AND owner=? AND fence=? AND status=?""",
|
|
(after, key, attempt, owner, fence, before)).rowcount
|
|
if changed != 1:
|
|
raise ValueError("history requires READY and current owner fence")
|
|
|
|
|
|
async def initialize(graph, key, messages, digest, store, attempt, owner, fence):
|
|
config = {"configurable": {"thread_id": key}}
|
|
marker = {"history_attempt": attempt, "history_digest": digest}
|
|
async with store.locked(key):
|
|
current = await graph.aget_state(config)
|
|
if store.get(key) is None and current.created_at is not None:
|
|
raise ValueError("unowned history checkpoint already exists")
|
|
store.claim(key, attempt, digest, owner, fence)
|
|
if current.created_at is None:
|
|
await graph.aupdate_state(dict(config, metadata=dict(marker, history_stage="START")),
|
|
{"messages": list(messages)}, as_node=START)
|
|
current = await graph.aget_state(config)
|
|
metadata = current.metadata or {}
|
|
if any(metadata.get(k) != v for k, v in marker.items()):
|
|
raise ValueError("unowned history checkpoint")
|
|
if current.values.get("messages", []) != list(messages):
|
|
raise ValueError("history checkpoint content changed")
|
|
stage = metadata.get("history_stage")
|
|
if stage == "START":
|
|
await graph.aupdate_state(dict(config, metadata=dict(marker, history_stage="END")),
|
|
None, as_node=END)
|
|
current = await graph.aget_state(config)
|
|
elif stage != "END":
|
|
raise ValueError("unknown history initialization stage")
|
|
if current.next or current.tasks:
|
|
raise ValueError("history checkpoint is not READY")
|
|
store.status(key, attempt, owner, fence, store.get(key)["status"], "READY")
|
|
return config
|
|
|
|
|
|
async def invoke(graph, key, store, attempt, owner, fence, input):
|
|
async with store.locked(key):
|
|
# Consume before execution: uncertain invocation must not be replayed.
|
|
store.status(key, attempt, owner, fence, "READY", "CONSUMED")
|
|
return await graph.ainvoke(input, {"configurable": {"thread_id": key}}) |