m4
120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
"""Instance-local stop adapter for LangGraph 1.2.6, SQLite saver 3.0.3.
|
|
|
|
No detached writers or foreign update_state/raw SQL writers are supported.
|
|
aiosqlite 0.22.1's FIFO commit is the connection barrier, not its saver lock.
|
|
"""
|
|
import asyncio
|
|
from importlib.metadata import version
|
|
from typing import Any
|
|
|
|
|
|
class StreamStop:
|
|
def __init__(self):
|
|
self.graph: Any = None
|
|
self.stream = None
|
|
self.pulls = set()
|
|
self.exits = set()
|
|
self.errors = []
|
|
self.task = None
|
|
self.unsupported = False
|
|
self.exit_tasks_observed = 0
|
|
self.unconfirmed_exit = asyncio.Event()
|
|
|
|
def attach(self, stream):
|
|
from langgraph.stream.run_stream import AsyncGraphRunStream
|
|
if version("langgraph") != "1.2.6" or type(stream) is not AsyncGraphRunStream:
|
|
self.unsupported = True
|
|
raise RuntimeError("CHECKPOINT_STOP_ADAPTER_UNSUPPORTED")
|
|
self.stream = stream
|
|
self.graph = stream._graph_aiter
|
|
stream._graph_aiter = self
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
def _record(self, exc):
|
|
self.exits.update(arg for arg in exc.args if isinstance(arg, asyncio.Future))
|
|
|
|
async def __anext__(self):
|
|
task = asyncio.current_task()
|
|
self.pulls.add(task)
|
|
try:
|
|
return await self.graph.__anext__()
|
|
except BaseException as exc:
|
|
self._record(exc)
|
|
raise
|
|
finally:
|
|
self.pulls.discard(task)
|
|
|
|
async def _retry_close(self, close):
|
|
while True:
|
|
try:
|
|
await close()
|
|
return
|
|
except BaseException as exc:
|
|
self._record(exc)
|
|
self.errors.append(exc)
|
|
await self._settle_exits()
|
|
await asyncio.sleep(0.05)
|
|
|
|
async def _settle_exits(self):
|
|
while self.exits:
|
|
tasks = tuple(self.exits)
|
|
self.exit_tasks_observed += len(tasks)
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
self.exits.difference_update(tasks)
|
|
for result in results:
|
|
if isinstance(result, BaseException):
|
|
self._record(result)
|
|
self.errors.append(result)
|
|
# A failed/cancelled stack exit is not executor settlement.
|
|
# There is no supported recovery handle for that case.
|
|
await self.unconfirmed_exit.wait()
|
|
|
|
def start(self, producers):
|
|
if self.task is None:
|
|
self.task = asyncio.create_task(self._close(producers), name="evo-v3-stop")
|
|
return self.task
|
|
|
|
async def _close(self, producers):
|
|
if self.unsupported:
|
|
await self.unconfirmed_exit.wait()
|
|
stream = self.stream
|
|
if stream is None:
|
|
return
|
|
# Freeze pumping before cancellation can discard the dependency's pull.
|
|
async with stream._pump_cond:
|
|
stream._exhausted = True
|
|
stream._aborting = True
|
|
pulls = set(self.pulls)
|
|
if stream._anext_task is not None:
|
|
pulls.add(stream._anext_task)
|
|
stream._pump_cond.notify_all()
|
|
for task in pulls:
|
|
if not task.done():
|
|
task.cancel()
|
|
results = await asyncio.gather(*pulls, return_exceptions=True)
|
|
for result in results:
|
|
if isinstance(result, BaseException):
|
|
self._record(result)
|
|
await self._settle_exits()
|
|
await self._retry_close(self.graph.aclose)
|
|
await self._settle_exits()
|
|
await self._retry_close(stream._mux.aclose)
|
|
for task in producers:
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(*producers, return_exceptions=True)
|
|
|
|
|
|
async def sqlite_barrier(saver):
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
import aiosqlite
|
|
if (type(saver) is not AsyncSqliteSaver or type(saver.conn) is not aiosqlite.Connection
|
|
or version("langgraph-checkpoint-sqlite") != "3.0.3"
|
|
or version("aiosqlite") != "0.22.1"):
|
|
raise RuntimeError("CHECKPOINT_STOP_ADAPTER_UNSUPPORTED")
|
|
async with saver.lock:
|
|
# Worker executes even cancelled queued futures. Await a later commit
|
|
# after graph/executor/repair writers settle, covering those operations.
|
|
await saver.conn.commit() |