* fix: session lost after evoscientist restart * feat: Implement memory worker thread deletion on completion - Added synchronous and asynchronous functions to delete memory worker threads after they finish execution, ensuring no residual checkpoints are left in the database. - Enhanced `_watch_memory_worker_run_sync` and `_watch_memory_worker_run_async` to invoke deletion functions upon confirming worker completion. - Introduced tests to verify that worker threads are deleted correctly upon successful completion and that failures in deletion do not affect the overall worker status. - Updated session management to ensure that only relevant threads are restored from the database, preventing exposure of internal or unrelated workspace threads. - Implemented a purge function to clean up leftover worker checkpoints during server startup. * feat: Implement short thread ID display for CLI and session hints * fix: ensure proper accounting and deletion order for memory worker threads --------- Co-authored-by: z00827015 <zhoulun1@huawei.com>
This commit is contained in:
Xi Zhang
@@ -39,6 +39,7 @@ from ..sessions import (
|
||||
get_thread_messages,
|
||||
get_thread_metadata,
|
||||
resolve_thread_id_prefix,
|
||||
short_thread_id,
|
||||
thread_exists,
|
||||
)
|
||||
from ..stream.console import console
|
||||
@@ -1423,7 +1424,7 @@ def cmd_run(
|
||||
console.print(sep)
|
||||
console.print(Text(f"> {prompt}"))
|
||||
console.print(sep)
|
||||
console.print(f"[dim]Thread: {thread_id}[/dim]")
|
||||
console.print(f"[dim]Thread: {short_thread_id(thread_id)}[/dim]")
|
||||
if workspace_dir:
|
||||
console.print(f"[dim]Workspace: {_shorten_path(workspace_dir)}[/dim]")
|
||||
console.print()
|
||||
|
||||
@@ -14,6 +14,8 @@ def print_resume_hint(
|
||||
out = console or Console()
|
||||
out.print("[dim]Goodbye![/dim]")
|
||||
if thread_id:
|
||||
from ..sessions import short_thread_id
|
||||
|
||||
out.print()
|
||||
out.print("[dim]Resume this session with:[/dim]")
|
||||
out.print(f"[cyan]EvoSci --resume {escape(thread_id)}[/cyan]")
|
||||
out.print(f"[cyan]EvoSci --resume {escape(short_thread_id(thread_id))}[/cyan]")
|
||||
|
||||
@@ -202,9 +202,9 @@ def build_row_text(
|
||||
indented: bool = False,
|
||||
) -> Text:
|
||||
"""Thread row. *indented* adds extra leading space for L2-grouped rows."""
|
||||
from ...sessions import _format_relative_time
|
||||
from ...sessions import _format_relative_time, short_thread_id
|
||||
|
||||
tid = thread["thread_id"]
|
||||
tid = short_thread_id(thread["thread_id"])
|
||||
preview = thread.get("preview", "") or ""
|
||||
msgs = thread.get("message_count", 0)
|
||||
model = thread.get("model", "") or ""
|
||||
|
||||
@@ -98,12 +98,14 @@ class ThreadsCommand(Command):
|
||||
table.add_column("Model", style="dim")
|
||||
table.add_column("Last Used", style="dim")
|
||||
|
||||
from ...sessions import short_thread_id
|
||||
|
||||
for thread in threads:
|
||||
thread_id_value = thread["thread_id"]
|
||||
marker = " *" if thread_id_value == ctx.thread_id else ""
|
||||
|
||||
row = [
|
||||
f"{thread_id_value}{marker}",
|
||||
f"{short_thread_id(thread_id_value)}{marker}",
|
||||
thread.get("preview", "") or "",
|
||||
str(thread.get("message_count", 0)),
|
||||
]
|
||||
|
||||
@@ -7,6 +7,10 @@
|
||||
"evomemory-subagent-worker": "EvoScientist.langgraph_dev.graphs:evomemory_subagent_worker",
|
||||
"evomemory-turn-worker": "EvoScientist.langgraph_dev.graphs:evomemory_turn_worker"
|
||||
},
|
||||
"checkpointer": {
|
||||
"backend": "custom",
|
||||
"path": "EvoScientist.sessions.create_checkpointer_for_langgraph_api"
|
||||
},
|
||||
"config": {
|
||||
"recursion_limit": 1000000
|
||||
}
|
||||
|
||||
@@ -958,6 +958,32 @@ def _status_from_run_response(run: object) -> str:
|
||||
return str(value or "").strip().lower()
|
||||
|
||||
|
||||
def _delete_memory_worker_thread(client: Any, thread_id: str) -> None:
|
||||
"""Best-effort delete of a finished worker thread.
|
||||
|
||||
Worker conversations have no value after the run: the durable artifact
|
||||
is the memory files they write, and worker threads are never resumed.
|
||||
Deleting the thread drops its checkpoints from the shared sessions.db
|
||||
so short-lived workers leave no per-turn residue behind.
|
||||
"""
|
||||
try:
|
||||
client.threads.delete(thread_id)
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Failed to delete EvoMemory worker thread %s", thread_id, exc_info=True
|
||||
)
|
||||
|
||||
|
||||
async def _adelete_memory_worker_thread(client: Any, thread_id: str) -> None:
|
||||
"""Async variant of :func:`_delete_memory_worker_thread`."""
|
||||
try:
|
||||
await client.threads.delete(thread_id)
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Failed to delete EvoMemory worker thread %s", thread_id, exc_info=True
|
||||
)
|
||||
|
||||
|
||||
def _spawn_memory_worker_status_thread(
|
||||
*,
|
||||
url: str,
|
||||
@@ -984,6 +1010,7 @@ def _watch_memory_worker_run_sync(
|
||||
|
||||
failures = 0
|
||||
worker_confirmed_finished = False
|
||||
client = None
|
||||
try:
|
||||
client = get_sync_client(url=url, headers={"x-auth-scheme": "langsmith"})
|
||||
while True:
|
||||
@@ -1010,7 +1037,14 @@ def _watch_memory_worker_run_sync(
|
||||
time.sleep(_MEMORY_WORKER_POLL_INTERVAL_SECONDS)
|
||||
finally:
|
||||
if worker_confirmed_finished:
|
||||
# Accounting first, then best-effort deletion (mirrors the
|
||||
# async watcher's cancellation-safe ordering). Only delete
|
||||
# once the run is terminal — deleting a thread with a live
|
||||
# run would break it. Crash residue is handled by the
|
||||
# restore whitelist + startup purge in sessions.py.
|
||||
mark_memory_worker_finished(thread_id, run_id)
|
||||
if client is not None:
|
||||
_delete_memory_worker_thread(client, thread_id)
|
||||
else:
|
||||
forget_memory_worker(thread_id, run_id)
|
||||
|
||||
@@ -1064,7 +1098,13 @@ async def _watch_memory_worker_run_async(
|
||||
await asyncio.sleep(_MEMORY_WORKER_POLL_INTERVAL_SECONDS)
|
||||
finally:
|
||||
if worker_confirmed_finished:
|
||||
# Accounting BEFORE the best-effort deletion: if this task is
|
||||
# cancelled mid-finally, only the deletion await is lost
|
||||
# (startup purge covers the residue). The to_thread side
|
||||
# effect completes even if its await is cancelled, so the
|
||||
# worker is never stuck "running".
|
||||
await asyncio.to_thread(mark_memory_worker_finished, thread_id, run_id)
|
||||
await _adelete_memory_worker_thread(client, thread_id)
|
||||
else:
|
||||
forget_memory_worker(thread_id, run_id)
|
||||
|
||||
|
||||
+378
-2
@@ -13,6 +13,18 @@ Per-step pruning:
|
||||
``get_checkpointer()`` yields a ``PruningCheckpointer`` that prunes
|
||||
older rows for the same ``(thread_id, checkpoint_ns)`` after every
|
||||
``aput()``. The first-run migration sweep cleans up legacy bloat.
|
||||
|
||||
WebUI / langgraph-dev checkpointer:
|
||||
``create_checkpointer_for_langgraph_api()`` — the ``checkpointer.path``
|
||||
target in ``langgraph_dev/langgraph.json`` — backs every ``langgraph
|
||||
dev`` subprocess (deploy / WebUI / CLI-spawned) with this same SQLite
|
||||
file instead of the default pickle-based ``InMemorySaver``, whose flush
|
||||
window and pickle-compatibility failures lose session history on
|
||||
restart (issue #277). On startup it purges leftover evomemory-worker
|
||||
rows and rebuilds the in-memory thread registry from SQLite. See
|
||||
``_restore_webui_threads_to_global_store`` for the restore scope and
|
||||
``_ApiPruningCheckpointer`` for the metadata stamping that makes WebUI
|
||||
threads first-class CLI sessions.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -99,9 +111,25 @@ def get_db_path() -> Path:
|
||||
return Path(_to_short_path(str(db_dir))) / "sessions.db"
|
||||
|
||||
|
||||
def short_thread_id(thread_id: str) -> str:
|
||||
"""First 8 chars for display (git-style); legacy 8-hex ids pass through.
|
||||
|
||||
All lookup commands (``/resume``, ``/delete``) accept prefixes, so the
|
||||
shortened form is always a usable handle.
|
||||
"""
|
||||
return thread_id[:8]
|
||||
|
||||
|
||||
def generate_thread_id() -> str:
|
||||
"""Generate an 8-char hex thread ID."""
|
||||
return uuid.uuid4().hex[:8]
|
||||
"""Generate a full-UUID thread ID.
|
||||
|
||||
UUID format (not the legacy 8-char hex) so CLI threads are addressable
|
||||
by langgraph-api — its thread endpoints reject non-UUID ids — which is
|
||||
what lets the WebUI list and resume CLI sessions. UIs display the
|
||||
first 8 chars; ``/resume`` prefix matching is unaffected. Pre-existing
|
||||
8-char hex threads keep working in the CLI but stay CLI-only.
|
||||
"""
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1369,3 +1397,351 @@ async def db_stats(top_n: int = 5) -> dict[str, Any]:
|
||||
# Read-only — corrupt/locked DB → return zeroed stats rather than crash.
|
||||
return out
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# langgraph-api / WebUI checkpointer factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _api_workspace_dir() -> str:
|
||||
"""Resolve the langgraph dev subprocess's workspace directory.
|
||||
|
||||
``start_langgraph_dev`` injects ``EVOSCIENTIST_WORKSPACE_DIR`` and sets
|
||||
the subprocess cwd to the workspace, so either source identifies the
|
||||
workspace this server instance is serving.
|
||||
"""
|
||||
import os
|
||||
|
||||
ws = os.environ.get("EVOSCIENTIST_WORKSPACE_DIR", "").strip()
|
||||
if ws:
|
||||
return str(Path(ws).expanduser().resolve())
|
||||
return str(Path.cwd().resolve())
|
||||
|
||||
|
||||
class _ApiPruningCheckpointer(PruningCheckpointer):
|
||||
"""``PruningCheckpointer`` that stamps CLI-compatible ownership metadata.
|
||||
|
||||
langgraph-api run metadata carries ``graph_id``/``assistant_id`` but not
|
||||
the ``agent_name`` / ``workspace_dir`` / ``updated_at`` keys that the CLI
|
||||
session surface (``list_threads``, ``/resume``, ``/delete``,
|
||||
``_prune_after_put``) filters and sorts on. Stamping them at write time
|
||||
— for main-graph runs only — makes WebUI threads first-class CLI
|
||||
sessions in the same workspace, and brings them under the existing
|
||||
pruning/retention machinery. Worker and async-subagent graphs are left
|
||||
unstamped on purpose: they must not surface in CLI listings.
|
||||
"""
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
config: Any,
|
||||
checkpoint: Any,
|
||||
metadata: Any,
|
||||
new_versions: Any,
|
||||
) -> Any:
|
||||
if isinstance(metadata, dict) and metadata.get("graph_id") == AGENT_NAME:
|
||||
metadata = dict(metadata)
|
||||
metadata.setdefault("agent_name", AGENT_NAME)
|
||||
metadata.setdefault("workspace_dir", _api_workspace_dir())
|
||||
metadata["updated_at"] = datetime.now(UTC).isoformat()
|
||||
return await super().aput(config, checkpoint, metadata, new_versions)
|
||||
|
||||
|
||||
async def _purge_internal_worker_threads() -> None:
|
||||
"""Best-effort removal of evomemory-worker checkpoint residue.
|
||||
|
||||
Finished workers delete their own thread (see
|
||||
``middleware/memory_lifecycle.py``), but a crash between run completion
|
||||
and deletion leaves rows behind — and rows written before that cleanup
|
||||
existed are still in the DB. Idempotent, runs on every server start,
|
||||
and never blocks startup on failure.
|
||||
"""
|
||||
try:
|
||||
db_path = str(get_db_path())
|
||||
async with aiosqlite.connect(db_path, timeout=30.0) as conn:
|
||||
if not await _table_exists(conn, "checkpoints"):
|
||||
return
|
||||
if await _table_exists(conn, "writes"):
|
||||
await conn.execute(
|
||||
"""
|
||||
DELETE FROM writes WHERE thread_id IN (
|
||||
SELECT DISTINCT thread_id FROM checkpoints
|
||||
WHERE json_extract(metadata, '$.graph_id') LIKE 'evomemory-%'
|
||||
)
|
||||
"""
|
||||
)
|
||||
cur = await conn.execute(
|
||||
"""
|
||||
DELETE FROM checkpoints
|
||||
WHERE json_extract(metadata, '$.graph_id') LIKE 'evomemory-%'
|
||||
"""
|
||||
)
|
||||
await conn.commit()
|
||||
if cur.rowcount:
|
||||
_logger.info(
|
||||
"Purged %d leftover evomemory-worker checkpoint row(s).",
|
||||
cur.rowcount,
|
||||
)
|
||||
except Exception:
|
||||
_logger.warning(
|
||||
"evomemory-worker residue purge failed (non-fatal).", exc_info=True
|
||||
)
|
||||
|
||||
|
||||
async def _restore_webui_threads_to_global_store() -> None:
|
||||
"""Re-populate ``GlobalStore["threads"]`` from SQLite on server startup.
|
||||
|
||||
The inmem runtime's thread registry lives in memory (pickled to
|
||||
``.langgraph_ops.pckl``) and is cleared on every start — if the pickle
|
||||
is absent or corrupt, the WebUI sidebar is empty even though all
|
||||
checkpoint data sits safely in SQLite. This rebuilds it: ghost entries
|
||||
whose threads have no checkpoint rows are dropped, surviving entries
|
||||
are normalized in place, and missing threads are appended as stub
|
||||
dicts that satisfy ``POST /threads/search``.
|
||||
|
||||
Restore scope — only threads that are BOTH main-graph
|
||||
(``metadata.graph_id == AGENT_NAME``) and owned by this server's
|
||||
workspace (``metadata.workspace_dir`` matches): sessions.db is
|
||||
machine-global, and an unscoped restore would expose every workspace's
|
||||
history (and internal worker threads) on the unauthenticated API —
|
||||
worst case ``--tunnel``. CLI/TUI threads (8-char hex IDs, managed by
|
||||
``list_threads()``) and pre-stamping rows without ``workspace_dir``
|
||||
are excluded.
|
||||
|
||||
Best-effort: any exception is logged and swallowed so a broken restore
|
||||
never prevents the ``langgraph dev`` server from starting.
|
||||
"""
|
||||
try:
|
||||
from langgraph_runtime_inmem.database import ( # type: ignore[import-untyped]
|
||||
GLOBAL_STORE,
|
||||
)
|
||||
except ImportError:
|
||||
# langgraph_runtime_inmem not available (unit tests, plain CLI mode).
|
||||
return
|
||||
|
||||
def _to_uuid_safe(v: Any) -> uuid.UUID | None:
|
||||
try:
|
||||
return uuid.UUID(str(v))
|
||||
except (ValueError, AttributeError):
|
||||
return None
|
||||
|
||||
try:
|
||||
rows: list[Any] = []
|
||||
# All UUID threads that have ANY checkpoint rows — the existence
|
||||
# check for ghost removal (deliberately unscoped: a thread whose
|
||||
# checkpoints exist but fall outside the restore scope is not a
|
||||
# ghost, its state still loads when opened).
|
||||
uuid_threads_in_db: set[uuid.UUID] = set()
|
||||
# Restore scope: ONLY main-graph threads belonging to THIS server's
|
||||
# workspace. sessions.db is machine-global, so an unscoped restore
|
||||
# would resurrect every workspace's history (and internal
|
||||
# worker/subagent threads) into this server's thread registry — and
|
||||
# expose it over the unauthenticated API / --tunnel. Main-graph =
|
||||
# metadata.graph_id == AGENT_NAME (langgraph-api rows, stamped by
|
||||
# _ApiPruningCheckpointer) OR no graph_id but agent_name ==
|
||||
# AGENT_NAME (CLI rows via build_metadata). Worker residue carries
|
||||
# graph_id='evomemory-*' and is excluded by the first clause even
|
||||
# though it also stamps agent_name. Rows predating stamping have no
|
||||
# workspace_dir and are deliberately excluded.
|
||||
current_workspace = _api_workspace_dir()
|
||||
sqlite_data: dict[uuid.UUID, tuple[str | None, str | None, str]] = {}
|
||||
titles: dict[uuid.UUID, str] = {}
|
||||
db_path = str(get_db_path())
|
||||
async with aiosqlite.connect(db_path, timeout=30.0) as conn:
|
||||
# No checkpoints table (fresh DB) → rows stays empty, but the
|
||||
# ghost-removal pass below must still run: every UUID entry the
|
||||
# .pckl registry loaded is then stale by definition.
|
||||
if await _table_exists(conn, "checkpoints"):
|
||||
# UUID regex: 8-4-4-4-12 hex groups. assistant_id/graph_id
|
||||
# are needed by the WebUI's POST /threads/search filters.
|
||||
# Every metadata field is MAX-aggregated: an interop thread
|
||||
# mixes CLI rows (no assistant_id/graph_id) with WebUI rows,
|
||||
# and a bare column under GROUP BY would let SQLite pick an
|
||||
# arbitrary row's NULL.
|
||||
query = """
|
||||
SELECT thread_id,
|
||||
MAX(json_extract(metadata, '$.updated_at')) as updated_at,
|
||||
MAX(json_extract(metadata, '$.assistant_id')) as assistant_id,
|
||||
MAX(json_extract(metadata, '$.graph_id')) as graph_id,
|
||||
MAX(json_extract(metadata, '$.workspace_dir')) as workspace_dir,
|
||||
MAX(json_extract(metadata, '$.agent_name')) as agent_name
|
||||
FROM checkpoints
|
||||
WHERE thread_id LIKE '________-____-____-____-____________'
|
||||
GROUP BY thread_id
|
||||
ORDER BY updated_at DESC
|
||||
"""
|
||||
async with conn.execute(query) as cur:
|
||||
rows = await cur.fetchall()
|
||||
|
||||
for row in rows:
|
||||
(
|
||||
thread_id_str,
|
||||
updated_at,
|
||||
assistant_id,
|
||||
graph_id,
|
||||
workspace_dir,
|
||||
agent_name,
|
||||
) = row
|
||||
thread_uuid = _to_uuid_safe(thread_id_str)
|
||||
if thread_uuid is None:
|
||||
continue
|
||||
uuid_threads_in_db.add(thread_uuid)
|
||||
is_main_graph = graph_id == AGENT_NAME or (
|
||||
graph_id is None and agent_name == AGENT_NAME
|
||||
)
|
||||
if not is_main_graph:
|
||||
continue
|
||||
if not workspace_dir or workspace_dir != current_workspace:
|
||||
continue
|
||||
sqlite_data[thread_uuid] = (updated_at, assistant_id, AGENT_NAME)
|
||||
|
||||
# Derive a sidebar title from each scoped thread's first human
|
||||
# message (stubs carry values=None, so the WebUI would otherwise
|
||||
# render every restored thread as "Untitled Thread").
|
||||
if sqlite_data:
|
||||
saver = AsyncSqliteSaver(conn, serde=JsonPlusSerializer())
|
||||
for thread_uuid in sqlite_data:
|
||||
try:
|
||||
msgs = await _load_checkpoint_messages(saver, str(thread_uuid))
|
||||
preview = _extract_preview(msgs)
|
||||
if preview:
|
||||
titles[thread_uuid] = preview
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
def _parse_dt(s: str | None) -> datetime:
|
||||
"""Parse an ISO timestamp string to datetime, falling back to now."""
|
||||
if s:
|
||||
try:
|
||||
return datetime.fromisoformat(s)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
return datetime.now(UTC)
|
||||
|
||||
# Drop ghost entries: a .pckl-loaded UUID entry with no checkpoint
|
||||
# rows opens as an empty session (the #277 symptom). Slice
|
||||
# assignment mutates the live registry list.
|
||||
store_threads: list[dict[str, Any]] = GLOBAL_STORE.get("threads", [])
|
||||
before = len(store_threads)
|
||||
store_threads[:] = [
|
||||
entry
|
||||
for entry in store_threads
|
||||
if (tid := _to_uuid_safe(entry.get("thread_id"))) is None
|
||||
or tid in uuid_threads_in_db
|
||||
]
|
||||
removed = before - len(store_threads)
|
||||
|
||||
# Normalize surviving .pckl-loaded entries in place.
|
||||
fixed = 0
|
||||
existing_uuids: set[uuid.UUID] = set()
|
||||
for entry in store_threads:
|
||||
tid_uuid = _to_uuid_safe(entry.get("thread_id"))
|
||||
if tid_uuid is None:
|
||||
continue
|
||||
existing_uuids.add(tid_uuid)
|
||||
changed = False
|
||||
# Threads.get() compares against _ensure_uuid() — str never matches.
|
||||
if not isinstance(entry.get("thread_id"), uuid.UUID):
|
||||
entry["thread_id"] = tid_uuid
|
||||
changed = True
|
||||
if tid_uuid in sqlite_data:
|
||||
_updated_at, asst_id_str, gid = sqlite_data[tid_uuid]
|
||||
meta: dict[str, Any] = entry.setdefault("metadata", {})
|
||||
if asst_id_str and "assistant_id" not in meta:
|
||||
# str, not uuid.UUID: the runtime stores str and search
|
||||
# filters compare with raw == against JSON strings.
|
||||
meta["assistant_id"] = str(asst_id_str)
|
||||
changed = True
|
||||
if gid and "graph_id" not in meta:
|
||||
meta["graph_id"] = gid
|
||||
changed = True
|
||||
if "title" not in meta and tid_uuid in titles:
|
||||
meta["title"] = titles[tid_uuid]
|
||||
changed = True
|
||||
# State.get() KeyErrors without "config".
|
||||
if "config" not in entry:
|
||||
entry["config"] = {}
|
||||
changed = True
|
||||
# Threads.search() sorted() raises on datetime-vs-str mixes.
|
||||
for ts_key in ("created_at", "updated_at", "state_updated_at"):
|
||||
if isinstance(entry.get(ts_key), str):
|
||||
entry[ts_key] = _parse_dt(entry[ts_key])
|
||||
changed = True
|
||||
if changed:
|
||||
fixed += 1
|
||||
|
||||
# Append threads present in SQLite but absent from the registry.
|
||||
restored = 0
|
||||
for thread_uuid, (updated_at, assistant_id, graph_id) in sqlite_data.items():
|
||||
if thread_uuid in existing_uuids:
|
||||
continue
|
||||
stub_metadata: dict[str, Any] = {"graph_id": graph_id}
|
||||
if assistant_id:
|
||||
# str, not uuid.UUID — same convention as above.
|
||||
stub_metadata["assistant_id"] = str(assistant_id)
|
||||
if thread_uuid in titles:
|
||||
stub_metadata["title"] = titles[thread_uuid]
|
||||
ts = _parse_dt(updated_at)
|
||||
stub: dict[str, Any] = {
|
||||
"thread_id": thread_uuid,
|
||||
"created_at": ts,
|
||||
"updated_at": ts,
|
||||
"state_updated_at": ts,
|
||||
"metadata": stub_metadata,
|
||||
"status": "idle",
|
||||
"config": {},
|
||||
"values": None,
|
||||
}
|
||||
GLOBAL_STORE["threads"].append(stub)
|
||||
existing_uuids.add(thread_uuid)
|
||||
restored += 1
|
||||
|
||||
if fixed or restored or removed:
|
||||
_logger.info(
|
||||
"WebUI thread restore: fixed %d existing + appended %d new + "
|
||||
"removed %d ghost thread(s) in GlobalStore "
|
||||
"(langgraph_runtime_inmem).",
|
||||
fixed,
|
||||
restored,
|
||||
removed,
|
||||
)
|
||||
except Exception:
|
||||
_logger.warning(
|
||||
"WebUI thread restore failed (non-fatal); WebUI session list may be "
|
||||
"empty until new threads are created.",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def create_checkpointer_for_langgraph_api() -> AsyncIterator[PruningCheckpointer]:
|
||||
"""SQLite-backed checkpointer for the ``langgraph dev`` subprocess.
|
||||
|
||||
``checkpointer.path`` target in ``langgraph_dev/langgraph.json``
|
||||
(applies to every ``langgraph dev`` launch: deploy, WebUI, and the
|
||||
CLI-spawned subprocess). Replaces the default pickle-based
|
||||
``InMemorySaver``, whose 10s flush window drops recent checkpoints on
|
||||
SIGKILL and whose pickle-incompatible upgrades wipe the whole store
|
||||
(issue #277); here every ``aput()`` commits a WAL transaction and a
|
||||
bad row only loses that row. The langgraph-api adapter detects async
|
||||
context managers and enters them automatically.
|
||||
|
||||
The yielded ``_ApiPruningCheckpointer`` stamps main-graph rows with
|
||||
``agent_name`` / ``workspace_dir`` / ``updated_at`` so WebUI threads
|
||||
surface in the CLI session commands and participate in
|
||||
``_prune_after_put`` retention.
|
||||
|
||||
Capability note: ``adelete_thread`` is real, but ``aprune`` /
|
||||
``adelete_for_runs`` / ``acopy_thread`` remain ``BaseCheckpointSaver``
|
||||
raising stubs — langgraph-api's probe reports them missing and
|
||||
degrades (``multitask_strategy='rollback'`` cleanup raises; thread
|
||||
copy uses the slow generic fallback).
|
||||
"""
|
||||
keep = _resolve_keep_per_ns()
|
||||
async with _ApiPruningCheckpointer.from_conn_string_with_keep(
|
||||
str(get_db_path()), keep_per_ns=keep
|
||||
) as saver:
|
||||
await saver.setup()
|
||||
await _purge_internal_worker_threads()
|
||||
await _restore_webui_threads_to_global_store()
|
||||
yield saver
|
||||
|
||||
@@ -754,6 +754,151 @@ def test_memory_worker_watcher_finishes_on_terminal_status(tmp_path, monkeypatch
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
|
||||
|
||||
def test_sync_watcher_deletes_worker_thread_on_terminal_status(tmp_path, monkeypatch):
|
||||
"""Finished workers leave no checkpoint residue: thread is deleted."""
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
worker_activity.mark_memory_worker_started(
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
memory_dir=tmp_path / "memories",
|
||||
)
|
||||
deleted: list[str] = []
|
||||
|
||||
class _Runs:
|
||||
def get(self, **_kwargs):
|
||||
return {"status": "success"}
|
||||
|
||||
class _Threads:
|
||||
def delete(self, thread_id):
|
||||
deleted.append(thread_id)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"langgraph_sdk.get_sync_client",
|
||||
lambda **_kwargs: SimpleNamespace(runs=_Runs(), threads=_Threads()),
|
||||
)
|
||||
monkeypatch.setattr(memory_lifecycle, "_MEMORY_WORKER_POLL_INTERVAL_SECONDS", 0)
|
||||
|
||||
try:
|
||||
memory_lifecycle._watch_memory_worker_run_sync(
|
||||
url="http://x",
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
)
|
||||
assert deleted == ["worker-thread"]
|
||||
assert worker_activity.memory_worker_status().is_running is False
|
||||
finally:
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
|
||||
|
||||
def test_sync_watcher_delete_failure_still_marks_finished(tmp_path, monkeypatch):
|
||||
"""Thread deletion is best-effort: a failure must not break accounting."""
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
worker_activity.mark_memory_worker_started(
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
memory_dir=tmp_path / "memories",
|
||||
)
|
||||
|
||||
class _Runs:
|
||||
def get(self, **_kwargs):
|
||||
return {"status": "success"}
|
||||
|
||||
class _Threads:
|
||||
def delete(self, thread_id):
|
||||
raise RuntimeError("delete failed")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"langgraph_sdk.get_sync_client",
|
||||
lambda **_kwargs: SimpleNamespace(runs=_Runs(), threads=_Threads()),
|
||||
)
|
||||
monkeypatch.setattr(memory_lifecycle, "_MEMORY_WORKER_POLL_INTERVAL_SECONDS", 0)
|
||||
|
||||
try:
|
||||
memory_lifecycle._watch_memory_worker_run_sync(
|
||||
url="http://x",
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
)
|
||||
assert worker_activity.memory_worker_status().is_running is False
|
||||
finally:
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
|
||||
|
||||
def test_sync_watcher_does_not_delete_thread_on_poll_abort(tmp_path, monkeypatch):
|
||||
"""A run we lost track of may still be live — never delete its thread."""
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
worker_activity.mark_memory_worker_started(
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
memory_dir=tmp_path / "memories",
|
||||
)
|
||||
deleted: list[str] = []
|
||||
|
||||
class _Runs:
|
||||
def get(self, **_kwargs):
|
||||
raise RuntimeError("poll failed")
|
||||
|
||||
class _Threads:
|
||||
def delete(self, thread_id):
|
||||
deleted.append(thread_id)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"langgraph_sdk.get_sync_client",
|
||||
lambda **_kwargs: SimpleNamespace(runs=_Runs(), threads=_Threads()),
|
||||
)
|
||||
monkeypatch.setattr(memory_lifecycle, "_MEMORY_WORKER_POLL_INTERVAL_SECONDS", 0)
|
||||
monkeypatch.setattr(memory_lifecycle, "_MEMORY_WORKER_MAX_POLL_FAILURES", 1)
|
||||
|
||||
try:
|
||||
memory_lifecycle._watch_memory_worker_run_sync(
|
||||
url="http://x",
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
)
|
||||
assert deleted == []
|
||||
finally:
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
|
||||
|
||||
def test_async_watcher_deletes_worker_thread_on_terminal_status(
|
||||
tmp_path, monkeypatch, run_async
|
||||
):
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
worker_activity.mark_memory_worker_started(
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
memory_dir=tmp_path / "memories",
|
||||
)
|
||||
deleted: list[str] = []
|
||||
|
||||
class _Runs:
|
||||
async def get(self, **_kwargs):
|
||||
return {"status": "success"}
|
||||
|
||||
class _Threads:
|
||||
async def delete(self, thread_id):
|
||||
# Accounting must complete BEFORE the best-effort deletion —
|
||||
# cancellation mid-deletion must never leave the worker
|
||||
# stuck as "running" (CodeRabbit on #279).
|
||||
assert worker_activity.memory_worker_status().is_running is False
|
||||
deleted.append(thread_id)
|
||||
|
||||
monkeypatch.setattr(memory_lifecycle, "_MEMORY_WORKER_POLL_INTERVAL_SECONDS", 0)
|
||||
|
||||
try:
|
||||
run_async(
|
||||
memory_lifecycle._watch_memory_worker_run_async(
|
||||
SimpleNamespace(runs=_Runs(), threads=_Threads()),
|
||||
thread_id="worker-thread",
|
||||
run_id="run-1",
|
||||
)
|
||||
)
|
||||
assert deleted == ["worker-thread"]
|
||||
assert worker_activity.memory_worker_status().is_running is False
|
||||
finally:
|
||||
worker_activity.reset_memory_worker_status_for_tests()
|
||||
|
||||
|
||||
def test_memory_worker_skips_when_langgraph_dev_unavailable(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(memory_lifecycle, "_memory_worker_url", lambda: "http://x")
|
||||
monkeypatch.setattr(
|
||||
|
||||
+810
-6
@@ -5,6 +5,7 @@ import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
import uuid
|
||||
from datetime import UTC
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -50,13 +51,10 @@ def _mock_path(db_path: str):
|
||||
|
||||
|
||||
class TestGenerateThreadId(unittest.TestCase):
|
||||
def test_length(self):
|
||||
def test_uuid_format(self):
|
||||
# Full UUID so langgraph-api can address CLI threads (WebUI interop).
|
||||
tid = generate_thread_id()
|
||||
assert len(tid) == 8
|
||||
|
||||
def test_hex(self):
|
||||
tid = generate_thread_id()
|
||||
int(tid, 16) # Should not raise
|
||||
assert tid == str(uuid.UUID(tid))
|
||||
|
||||
def test_uniqueness(self):
|
||||
ids = {generate_thread_id() for _ in range(100)}
|
||||
@@ -2004,5 +2002,811 @@ class TestReduceMessagesDeltaUpstreamParity:
|
||||
assert _signature(out) == [("HumanMessage", "d1", "x")]
|
||||
|
||||
|
||||
class TestCreateCheckpointerForLanggraphApi(unittest.TestCase):
|
||||
"""Tests for ``create_checkpointer_for_langgraph_api`` — the WebUI/deploy
|
||||
SQLite checkpointer factory that replaces the default ``InMemorySaver``."""
|
||||
|
||||
def test_yields_pruning_checkpointer(self):
|
||||
"""Factory yields a ``PruningCheckpointer`` instance."""
|
||||
from EvoScientist.sessions import (
|
||||
PruningCheckpointer,
|
||||
create_checkpointer_for_langgraph_api,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
with patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
):
|
||||
|
||||
async def _run_inner():
|
||||
async with create_checkpointer_for_langgraph_api() as cp:
|
||||
assert isinstance(cp, PruningCheckpointer)
|
||||
|
||||
_run(_run_inner())
|
||||
|
||||
def test_checkpointer_is_set_up(self):
|
||||
"""Factory calls ``setup()`` so tables exist before yielding."""
|
||||
import aiosqlite
|
||||
|
||||
from EvoScientist.sessions import create_checkpointer_for_langgraph_api
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
with patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
):
|
||||
|
||||
async def _run_inner():
|
||||
async with create_checkpointer_for_langgraph_api():
|
||||
async with aiosqlite.connect(db) as conn:
|
||||
async with conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='checkpoints'"
|
||||
) as cur:
|
||||
row = await cur.fetchone()
|
||||
assert row is not None, (
|
||||
"checkpoints table must exist after setup()"
|
||||
)
|
||||
|
||||
_run(_run_inner())
|
||||
|
||||
def test_checkpointer_persists_across_contexts(self):
|
||||
"""Data written in one context manager is readable in a new one.
|
||||
|
||||
This is the core regression test: verifies that session data
|
||||
survives process restarts (simulated as two separate ``async with``
|
||||
blocks sharing the same DB file).
|
||||
"""
|
||||
from EvoScientist.sessions import create_checkpointer_for_langgraph_api
|
||||
|
||||
thread_id = "testthread1"
|
||||
|
||||
async def _run_inner():
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
|
||||
def _patch():
|
||||
return patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
)
|
||||
|
||||
# --- First "process": write a checkpoint ---
|
||||
with _patch():
|
||||
async with create_checkpointer_for_langgraph_api() as cp:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
checkpoint = {
|
||||
"v": 1,
|
||||
"id": "ckpt-001",
|
||||
"ts": "2024-01-01T00:00:00+00:00",
|
||||
"channel_values": {"messages": []},
|
||||
"channel_versions": {},
|
||||
"versions_seen": {},
|
||||
"pending_sends": [],
|
||||
}
|
||||
metadata = {
|
||||
"source": "input",
|
||||
"step": 0,
|
||||
"writes": {},
|
||||
"parents": {},
|
||||
"agent_name": "EvoScientist",
|
||||
}
|
||||
await cp.aput(config, checkpoint, metadata, {})
|
||||
|
||||
# --- Second "process": read back the checkpoint ---
|
||||
with _patch():
|
||||
async with create_checkpointer_for_langgraph_api() as cp2:
|
||||
result = await cp2.aget_tuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
)
|
||||
assert result is not None, (
|
||||
"Checkpoint written in first context must be readable in second context "
|
||||
"(simulates data survival across process restarts)"
|
||||
)
|
||||
assert result.config["configurable"]["thread_id"] == thread_id
|
||||
|
||||
_run(_run_inner())
|
||||
|
||||
def test_capability_surface_matches_langgraph_api_probe(self):
|
||||
"""Document the REAL capability surface langgraph-api will detect.
|
||||
|
||||
The adapter's ``_is_overridden`` probe compares each method against
|
||||
``BaseCheckpointSaver``: inherited raising stubs do NOT count as
|
||||
capability. ``callable()`` checks would pass on the stubs and give
|
||||
false confidence (the original PR test did exactly that).
|
||||
"""
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
|
||||
from EvoScientist.sessions import _ApiPruningCheckpointer
|
||||
|
||||
def overridden(name: str) -> bool:
|
||||
base = getattr(BaseCheckpointSaver, name, None)
|
||||
sub = getattr(_ApiPruningCheckpointer, name, None)
|
||||
return base is not None and sub is not None and sub is not base
|
||||
|
||||
# Real implementations the adapter will detect and use.
|
||||
for method in ("adelete_thread", "aget_tuple", "aput", "aput_writes"):
|
||||
assert overridden(method), f"'{method}' must be a real implementation"
|
||||
# Known degradations: still BaseCheckpointSaver raising stubs.
|
||||
# rollback cleanup raises at runtime; thread copy uses the adapter's
|
||||
# slow generic fallback. If these start passing, the docstring in
|
||||
# create_checkpointer_for_langgraph_api should be updated.
|
||||
for method in ("aprune", "adelete_for_runs", "acopy_thread"):
|
||||
assert not overridden(method), (
|
||||
f"'{method}' is now overridden — update the capability "
|
||||
"docstring in create_checkpointer_for_langgraph_api"
|
||||
)
|
||||
|
||||
def test_aput_stamps_cli_metadata_for_main_graph_rows(self):
|
||||
"""Main-graph (graph_id == AGENT_NAME) rows get agent_name /
|
||||
workspace_dir / updated_at so they surface in CLI listings and
|
||||
participate in pruning."""
|
||||
import json
|
||||
|
||||
import aiosqlite
|
||||
|
||||
from EvoScientist.sessions import (
|
||||
AGENT_NAME,
|
||||
create_checkpointer_for_langgraph_api,
|
||||
)
|
||||
|
||||
def _checkpoint(cid: str) -> dict:
|
||||
return {
|
||||
"v": 1,
|
||||
"id": cid,
|
||||
"ts": "2024-01-01T00:00:00+00:00",
|
||||
"channel_values": {"messages": []},
|
||||
"channel_versions": {},
|
||||
"versions_seen": {},
|
||||
"pending_sends": [],
|
||||
}
|
||||
|
||||
async def _run_inner(db: str):
|
||||
async with create_checkpointer_for_langgraph_api() as cp:
|
||||
# Simulates a WebUI main-graph run (langgraph-api metadata).
|
||||
await cp.aput(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": "11111111-1111-1111-1111-111111111111",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
},
|
||||
_checkpoint("ckpt-main"),
|
||||
{"source": "loop", "step": 1, "graph_id": AGENT_NAME},
|
||||
{},
|
||||
)
|
||||
# Simulates a memory-worker run: must stay unstamped.
|
||||
await cp.aput(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": "22222222-2222-2222-2222-222222222222",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
},
|
||||
_checkpoint("ckpt-worker"),
|
||||
{"source": "loop", "step": 1, "graph_id": "evomemory-turn-worker"},
|
||||
{},
|
||||
)
|
||||
async with aiosqlite.connect(db) as conn:
|
||||
rows = {}
|
||||
async with conn.execute(
|
||||
"SELECT thread_id, metadata FROM checkpoints"
|
||||
) as cur:
|
||||
async for tid, meta in cur:
|
||||
rows[tid] = json.loads(meta)
|
||||
main = rows["11111111-1111-1111-1111-111111111111"]
|
||||
worker = rows["22222222-2222-2222-2222-222222222222"]
|
||||
assert main.get("agent_name") == AGENT_NAME
|
||||
assert main.get("workspace_dir") == "/tmp/test-workspace"
|
||||
assert main.get("updated_at"), "updated_at drives /threads ordering"
|
||||
assert "agent_name" not in worker
|
||||
assert "workspace_dir" not in worker
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
os.environ,
|
||||
{"EVOSCIENTIST_WORKSPACE_DIR": "/tmp/test-workspace"},
|
||||
),
|
||||
patch(
|
||||
"EvoScientist.sessions._api_workspace_dir",
|
||||
return_value="/tmp/test-workspace",
|
||||
),
|
||||
):
|
||||
_run(_run_inner(db))
|
||||
|
||||
|
||||
class TestRestoreWebuiThreadsToGlobalStore(unittest.TestCase):
|
||||
"""Tests for ``_restore_webui_threads_to_global_store``.
|
||||
|
||||
Verifies that UUID-format threads written to SQLite by ``langgraph dev``
|
||||
runs are re-populated into ``GlobalStore["threads"]`` on server restart,
|
||||
so the WebUI sidebar is not empty after a package upgrade or clean restart.
|
||||
"""
|
||||
|
||||
_WS = "/tmp/restore-test-workspace"
|
||||
|
||||
def _make_db_with_threads(
|
||||
self,
|
||||
db_path: str,
|
||||
thread_ids: list[str],
|
||||
assistant_id: str | None = "aaaa-bbbb",
|
||||
graph_id: str | None = "EvoScientist",
|
||||
workspace_dir: str | None = _WS,
|
||||
agent_name: str | None = "EvoScientist",
|
||||
ckpt_prefix: str = "ckpt",
|
||||
) -> None:
|
||||
"""Insert minimal checkpoint rows for the given thread_ids into a fresh DB."""
|
||||
import json
|
||||
import sqlite3
|
||||
|
||||
con = sqlite3.connect(db_path)
|
||||
con.execute(
|
||||
"CREATE TABLE IF NOT EXISTS checkpoints "
|
||||
"(thread_id TEXT, checkpoint_ns TEXT, checkpoint_id TEXT PRIMARY KEY, "
|
||||
" parent_checkpoint_id TEXT, type TEXT, checkpoint BLOB, metadata TEXT)"
|
||||
)
|
||||
for tid in thread_ids:
|
||||
meta_dict: dict = {"updated_at": "2025-01-01T00:00:00+00:00"}
|
||||
if agent_name is not None:
|
||||
meta_dict["agent_name"] = agent_name
|
||||
if assistant_id is not None:
|
||||
meta_dict["assistant_id"] = assistant_id
|
||||
if graph_id is not None:
|
||||
meta_dict["graph_id"] = graph_id
|
||||
if workspace_dir is not None:
|
||||
meta_dict["workspace_dir"] = workspace_dir
|
||||
meta = json.dumps(meta_dict)
|
||||
con.execute(
|
||||
"INSERT INTO checkpoints VALUES (?,?,?,?,?,?,?)",
|
||||
(tid, "", f"{ckpt_prefix}-{tid[:8]}", None, "empty", b"", meta),
|
||||
)
|
||||
con.commit()
|
||||
con.close()
|
||||
|
||||
def _patch_workspace(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
return patch("EvoScientist.sessions._api_workspace_dir", return_value=self._WS)
|
||||
|
||||
def test_restores_uuid_threads_into_global_store(self):
|
||||
"""UUID-format thread IDs from SQLite are injected into GlobalStore."""
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
uuid_id = "12345678-1234-1234-1234-123456789abc"
|
||||
asst_uuid_id = "a2b49500-c49b-5560-b664-d42ee8b66d3c"
|
||||
short_id = "abcd1234" # CLI-style, must be excluded
|
||||
|
||||
mock_store: dict = {"threads": []}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
mock_global_store.__setitem__ = lambda self, k, v: mock_store.__setitem__(k, v)
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
self._make_db_with_threads(
|
||||
db, [uuid_id, short_id], assistant_id=asst_uuid_id
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
self._patch_workspace(),
|
||||
):
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
# Only the UUID thread should have been added; the short-hex CLI thread
|
||||
# should not appear because it doesn't match the UUID LIKE pattern.
|
||||
added = mock_store["threads"]
|
||||
assert len(added) == 1, f"Expected 1 restored thread, got {len(added)}: {added}"
|
||||
import uuid as _uuid_mod
|
||||
|
||||
# thread_id MUST be stored as a uuid.UUID object, not a plain string.
|
||||
# langgraph_runtime_inmem._get_with_filters compares with == against
|
||||
# _ensure_uuid(thread_id), which returns a UUID object. A string never
|
||||
# equals a UUID object, causing every Threads.get() call to 404.
|
||||
assert isinstance(added[0]["thread_id"], _uuid_mod.UUID), (
|
||||
f"thread_id must be uuid.UUID, got {type(added[0]['thread_id'])}"
|
||||
)
|
||||
assert added[0]["thread_id"] == _uuid_mod.UUID(uuid_id)
|
||||
assert added[0]["status"] == "idle"
|
||||
# metadata.assistant_id must stay a STRING: the runtime stores
|
||||
# str(assistant_id) (ops.py Threads.create) and search filters compare
|
||||
# with raw == against JSON strings — a uuid.UUID here would silently
|
||||
# exclude restored threads from assistant_id-filtered searches.
|
||||
assert added[0]["metadata"].get("assistant_id") == asst_uuid_id
|
||||
assert isinstance(added[0]["metadata"].get("assistant_id"), str)
|
||||
assert added[0]["metadata"].get("graph_id") == "EvoScientist"
|
||||
# created_at / updated_at must be datetime objects, not ISO strings.
|
||||
# Threads.search() sorts by these fields using sorted(); mixing
|
||||
# datetime and str raises TypeError: '<' not supported.
|
||||
from datetime import datetime as _dt
|
||||
|
||||
assert isinstance(added[0]["created_at"], _dt), (
|
||||
f"created_at must be datetime, got {type(added[0]['created_at'])}"
|
||||
)
|
||||
assert isinstance(added[0]["updated_at"], _dt)
|
||||
|
||||
def test_fixes_existing_string_thread_ids_in_place(self):
|
||||
"""Threads already in GlobalStore with string thread_id get fixed in-place.
|
||||
|
||||
When .pckl loads successfully, threads are already in the store but
|
||||
thread_id is a plain string (as pickled). The restore must:
|
||||
1. Convert thread_id to uuid.UUID so Threads.get() comparison works.
|
||||
2. Backfill missing metadata.assistant_id / graph_id from SQLite.
|
||||
3. Not create duplicate entries.
|
||||
"""
|
||||
import sys
|
||||
import uuid as _uuid_mod
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
uuid_id = "aaaabbbb-aaaa-bbbb-cccc-ddddeeeeffff"
|
||||
asst_uuid_id = "a2b49500-c49b-5560-b664-d42ee8b66d3c"
|
||||
|
||||
# Simulate .pckl-restored thread: thread_id is a string, metadata empty.
|
||||
existing_stub: dict = {"thread_id": uuid_id, "status": "idle", "metadata": {}}
|
||||
mock_store: dict = {"threads": [existing_stub]}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
mock_global_store.__setitem__ = lambda self, k, v: mock_store.__setitem__(k, v)
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
self._make_db_with_threads(db, [uuid_id], assistant_id=asst_uuid_id)
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
self._patch_workspace(),
|
||||
):
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
# No duplicate: still exactly one entry.
|
||||
assert len(mock_store["threads"]) == 1, (
|
||||
f"Expected 1 thread, got {len(mock_store['threads'])}"
|
||||
)
|
||||
t = mock_store["threads"][0]
|
||||
# thread_id must now be a UUID object, not a string.
|
||||
assert isinstance(t["thread_id"], _uuid_mod.UUID), (
|
||||
f"thread_id must be uuid.UUID after fix, got {type(t['thread_id'])}"
|
||||
)
|
||||
assert t["thread_id"] == _uuid_mod.UUID(uuid_id)
|
||||
# metadata must be backfilled — assistant_id as str (runtime convention).
|
||||
assert t["metadata"].get("assistant_id") == asst_uuid_id
|
||||
assert isinstance(t["metadata"].get("assistant_id"), str)
|
||||
assert t["metadata"].get("graph_id") == "EvoScientist"
|
||||
|
||||
def test_restore_excludes_other_workspaces_and_internal_graphs(self):
|
||||
"""The restore scope is graph_id==AGENT_NAME AND current workspace.
|
||||
|
||||
Threads from other workspaces, internal worker/subagent graphs, and
|
||||
pre-stamping rows without workspace_dir must NOT be resurrected —
|
||||
sessions.db is machine-global and an unscoped restore would expose
|
||||
them on the unauthenticated API (worst case --tunnel).
|
||||
"""
|
||||
import sys
|
||||
import uuid as _uuid_mod
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
mine = "11111111-1111-1111-1111-111111111111"
|
||||
other_ws = "22222222-2222-2222-2222-222222222222"
|
||||
worker = "33333333-3333-3333-3333-333333333333"
|
||||
subagent = "44444444-4444-4444-4444-444444444444"
|
||||
legacy_no_ws = "55555555-5555-5555-5555-555555555555"
|
||||
|
||||
mock_store: dict = {"threads": []}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
mock_global_store.__setitem__ = lambda self, k, v: mock_store.__setitem__(k, v)
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
self._make_db_with_threads(db, [mine])
|
||||
self._make_db_with_threads(db, [other_ws], workspace_dir="/elsewhere")
|
||||
self._make_db_with_threads(db, [worker], graph_id="evomemory-turn-worker")
|
||||
self._make_db_with_threads(db, [subagent], graph_id="writing-agent")
|
||||
self._make_db_with_threads(db, [legacy_no_ws], workspace_dir=None)
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
self._patch_workspace(),
|
||||
):
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
added = mock_store["threads"]
|
||||
assert len(added) == 1, (
|
||||
f"Only the current-workspace main-graph thread may be restored, "
|
||||
f"got {len(added)}: {[t['thread_id'] for t in added]}"
|
||||
)
|
||||
assert added[0]["thread_id"] == _uuid_mod.UUID(mine)
|
||||
|
||||
def test_purge_removes_only_evomemory_rows(self):
|
||||
"""Startup purge drops evomemory-* residue, leaves everything else."""
|
||||
import sqlite3
|
||||
from unittest.mock import patch
|
||||
|
||||
from EvoScientist.sessions import _purge_internal_worker_threads
|
||||
|
||||
keep_main = "11111111-1111-1111-1111-111111111111"
|
||||
keep_cli = "abcd1234"
|
||||
drop_worker = "33333333-3333-3333-3333-333333333333"
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
self._make_db_with_threads(db, [keep_main, keep_cli])
|
||||
self._make_db_with_threads(
|
||||
db, [drop_worker], graph_id="evomemory-turn-worker"
|
||||
)
|
||||
with patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
):
|
||||
_run(_purge_internal_worker_threads())
|
||||
# Idempotent: second run is a no-op, not an error.
|
||||
_run(_purge_internal_worker_threads())
|
||||
|
||||
con = sqlite3.connect(db)
|
||||
remaining = {
|
||||
r[0] for r in con.execute("SELECT DISTINCT thread_id FROM checkpoints")
|
||||
}
|
||||
con.close()
|
||||
assert remaining == {keep_main, keep_cli}
|
||||
|
||||
def test_restores_cli_rows_and_excludes_worker_residue(self):
|
||||
"""CLI rows (agent_name, no graph_id) are restored with graph_id
|
||||
backfilled; crashed-worker residue (agent_name AND graph_id=
|
||||
evomemory-*) stays excluded — graph_id wins over agent_name."""
|
||||
import sys
|
||||
import uuid as _uuid_mod
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
cli_thread = "11111111-1111-1111-1111-111111111111"
|
||||
worker_residue = "22222222-2222-2222-2222-222222222222"
|
||||
|
||||
mock_store: dict = {"threads": []}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
mock_global_store.__setitem__ = lambda self, k, v: mock_store.__setitem__(k, v)
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
# CLI row: build_metadata stamps agent_name/workspace_dir but
|
||||
# never graph_id or assistant_id.
|
||||
self._make_db_with_threads(
|
||||
db, [cli_thread], assistant_id=None, graph_id=None
|
||||
)
|
||||
# Crashed memory-worker residue: stamps BOTH.
|
||||
self._make_db_with_threads(
|
||||
db, [worker_residue], graph_id="evomemory-turn-worker"
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
self._patch_workspace(),
|
||||
):
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
added = mock_store["threads"]
|
||||
assert len(added) == 1, f"expected only the CLI thread, got {added}"
|
||||
assert added[0]["thread_id"] == _uuid_mod.UUID(cli_thread)
|
||||
# graph_id backfilled so Threads.State.get works on the stub.
|
||||
assert added[0]["metadata"].get("graph_id") == "EvoScientist"
|
||||
|
||||
def test_mixed_cli_webui_rows_keep_assistant_and_graph_id(self):
|
||||
"""Interop thread (CLI rows + WebUI rows under one UUID): bare
|
||||
columns under GROUP BY let SQLite pick an arbitrary row's NULL —
|
||||
all metadata fields must be MAX-aggregated (Codex F2)."""
|
||||
import sys
|
||||
import uuid as _uuid_mod
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
tid = "11111111-1111-1111-1111-111111111111"
|
||||
asst = "a2b49500-c49b-5560-b664-d42ee8b66d3c"
|
||||
|
||||
mock_store: dict = {"threads": []}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
mock_global_store.__setitem__ = lambda self, k, v: mock_store.__setitem__(k, v)
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
# Older CLI row: no assistant_id / graph_id. "a-..." checkpoint
|
||||
# id sorts BEFORE the WebUI row's so a bare-column GROUP BY
|
||||
# would tend to surface this row's NULLs.
|
||||
self._make_db_with_threads(
|
||||
db, [tid], assistant_id=None, graph_id=None, ckpt_prefix="a"
|
||||
)
|
||||
# Newer WebUI row on the SAME thread: carries both.
|
||||
self._make_db_with_threads(db, [tid], assistant_id=asst, ckpt_prefix="b")
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
self._patch_workspace(),
|
||||
):
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
added = mock_store["threads"]
|
||||
assert len(added) == 1, f"expected 1 thread, got {added}"
|
||||
assert added[0]["thread_id"] == _uuid_mod.UUID(tid)
|
||||
assert added[0]["metadata"].get("assistant_id") == asst
|
||||
assert added[0]["metadata"].get("graph_id") == "EvoScientist"
|
||||
|
||||
def test_restored_stub_gets_title_from_first_human_message(self):
|
||||
"""Stubs carry metadata.title derived from the thread's first human
|
||||
message, so the WebUI sidebar doesn't show "Untitled Thread"."""
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from EvoScientist.sessions import (
|
||||
AGENT_NAME,
|
||||
_restore_webui_threads_to_global_store,
|
||||
create_checkpointer_for_langgraph_api,
|
||||
)
|
||||
|
||||
thread_id = "33333333-3333-3333-3333-333333333333"
|
||||
|
||||
mock_store: dict = {"threads": []}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
mock_global_store.__setitem__ = lambda self, k, v: mock_store.__setitem__(k, v)
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
async def _write_then_restore():
|
||||
async with create_checkpointer_for_langgraph_api() as cp:
|
||||
await cp.aput(
|
||||
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
||||
{
|
||||
"v": 1,
|
||||
"id": "ckpt-title",
|
||||
"ts": "2024-01-01T00:00:00+00:00",
|
||||
"channel_values": {
|
||||
"messages": [HumanMessage(content="hello title test")]
|
||||
},
|
||||
"channel_versions": {},
|
||||
"versions_seen": {},
|
||||
"pending_sends": [],
|
||||
},
|
||||
{"source": "loop", "step": 1, "graph_id": AGENT_NAME},
|
||||
{},
|
||||
)
|
||||
await _restore_webui_threads_to_global_store()
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
patch(
|
||||
"EvoScientist.sessions._api_workspace_dir",
|
||||
return_value=self._WS,
|
||||
),
|
||||
):
|
||||
_run(_write_then_restore())
|
||||
|
||||
added = mock_store["threads"]
|
||||
assert len(added) == 1, f"expected 1 restored thread, got {added}"
|
||||
assert added[0]["metadata"].get("title") == "hello title test"
|
||||
|
||||
def test_removes_ghost_entries_absent_from_sqlite(self):
|
||||
"""Stale .pckl UUID entries with no checkpoint rows are dropped.
|
||||
|
||||
Ghost entries point at deleted/lost state and render as empty
|
||||
sessions (the #277 symptom). Existence is checked against ALL UUID
|
||||
threads in the DB, not the scoped restore set: a thread whose
|
||||
checkpoints exist but fall outside the restore scope still opens
|
||||
fine, so it must NOT be treated as a ghost. CLI-style non-UUID
|
||||
entries are never touched.
|
||||
"""
|
||||
import sys
|
||||
import uuid as _uuid_mod
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
in_scope = "11111111-1111-1111-1111-111111111111"
|
||||
out_of_scope = "22222222-2222-2222-2222-222222222222"
|
||||
ghost = "99999999-9999-9999-9999-999999999999"
|
||||
|
||||
ghost_entry: dict = {"thread_id": ghost, "status": "idle", "metadata": {}}
|
||||
out_of_scope_entry: dict = {
|
||||
"thread_id": out_of_scope,
|
||||
"status": "idle",
|
||||
"metadata": {},
|
||||
}
|
||||
cli_entry: dict = {"thread_id": "notauuid", "status": "idle"}
|
||||
mock_store: dict = {"threads": [ghost_entry, out_of_scope_entry, cli_entry]}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
mock_global_store.__setitem__ = lambda self, k, v: mock_store.__setitem__(k, v)
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
self._make_db_with_threads(db, [in_scope])
|
||||
self._make_db_with_threads(db, [out_of_scope], workspace_dir="/elsewhere")
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
self._patch_workspace(),
|
||||
):
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
ids = [t["thread_id"] for t in mock_store["threads"]]
|
||||
assert _uuid_mod.UUID(ghost) not in ids, f"ghost must be removed, got {ids}"
|
||||
assert ghost not in ids, f"ghost must be removed (str form), got {ids}"
|
||||
assert "notauuid" in ids, "CLI-style entries must never be touched"
|
||||
# Out-of-scope but existing in DB: kept (state still loads when opened).
|
||||
assert _uuid_mod.UUID(out_of_scope) in ids
|
||||
# In-scope thread restored as usual.
|
||||
assert _uuid_mod.UUID(in_scope) in ids
|
||||
|
||||
def test_no_op_when_langgraph_runtime_inmem_absent(self):
|
||||
"""ImportError for langgraph_runtime_inmem is silently swallowed."""
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
with patch.dict(sys.modules, {"langgraph_runtime_inmem.database": None}):
|
||||
# Must not raise.
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
def test_no_op_when_db_has_no_checkpoints_table(self):
|
||||
"""Missing checkpoints table is handled gracefully."""
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from EvoScientist.sessions import _restore_webui_threads_to_global_store
|
||||
|
||||
mock_store: dict = {"threads": []}
|
||||
mock_global_store = MagicMock()
|
||||
mock_global_store.get.side_effect = mock_store.get
|
||||
mock_global_store.__getitem__ = lambda self, k: mock_store[k]
|
||||
|
||||
fake_module = MagicMock()
|
||||
fake_module.GLOBAL_STORE = mock_global_store
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "empty.db")
|
||||
# Create a valid but empty SQLite DB (no checkpoints table).
|
||||
import sqlite3
|
||||
|
||||
sqlite3.connect(db).close()
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"langgraph_runtime_inmem.database": fake_module}
|
||||
),
|
||||
):
|
||||
_run(_restore_webui_threads_to_global_store())
|
||||
|
||||
# threads list untouched.
|
||||
assert mock_store["threads"] == []
|
||||
|
||||
def test_create_checkpointer_calls_restore(self):
|
||||
"""create_checkpointer_for_langgraph_api calls _restore_webui_threads_to_global_store."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from EvoScientist.sessions import create_checkpointer_for_langgraph_api
|
||||
|
||||
restore_called = []
|
||||
|
||||
async def fake_restore():
|
||||
restore_called.append(True)
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
db = os.path.join(td, "sessions.db")
|
||||
with (
|
||||
patch(
|
||||
"EvoScientist.sessions.get_db_path",
|
||||
return_value=_mock_path(db),
|
||||
),
|
||||
patch(
|
||||
"EvoScientist.sessions._restore_webui_threads_to_global_store",
|
||||
side_effect=fake_restore,
|
||||
),
|
||||
):
|
||||
|
||||
async def _run_inner():
|
||||
async with create_checkpointer_for_langgraph_api():
|
||||
pass
|
||||
|
||||
_run(_run_inner())
|
||||
|
||||
assert restore_called, "_restore_webui_threads_to_global_store must be called"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user