From 918641d3e2dead9b9e4705388747aa8692cbad4d Mon Sep 17 00:00:00 2001 From: Johannes du Plessis Date: Tue, 30 Jun 2026 17:59:41 -0700 Subject: [PATCH] fix: purge expired thread_wakeup crons (#1656) * fix: purge expired thread_wakeup crons One-shot wakeup crons set an end_time that stops re-firing but the cron row is never deleted, so dead rows accumulate (86 in prod). Add a purge that deletes thread_wakeup crons past their end_time, called opportunistically before scheduling a new wakeup, plus a one-time backfill script. Conservative: matches only kind=thread_wakeup with a past end_time. * chore: retrigger Open SWE review --------- Co-authored-by: open-swe[bot] (cherry picked from commit 9e5a1924ef306269322c31342a1831e57831cfee) --- agent/tools/schedule_thread_wakeup.py | 70 +++++++++++++++ scripts/purge_wakeup_crons.py | 88 ++++++++++++++++++ tests/test_schedule_thread_wakeup.py | 124 +++++++++++++++++++++++++- 3 files changed, 281 insertions(+), 1 deletion(-) create mode 100644 scripts/purge_wakeup_crons.py diff --git a/agent/tools/schedule_thread_wakeup.py b/agent/tools/schedule_thread_wakeup.py index 88b558c5..78d42e06 100644 --- a/agent/tools/schedule_thread_wakeup.py +++ b/agent/tools/schedule_thread_wakeup.py @@ -18,6 +18,9 @@ _MIN_DELAY_SECONDS = 60 _MAX_DELAY_SECONDS = 86_400 _END_TIME_PADDING_SECONDS = 90 +_WAKEUP_KIND = "thread_wakeup" +_PURGE_PAGE_SIZE = 100 + _DEFAULT_WAKEUP_PROMPT = ( "This is an automated re-trigger of this thread. The agent scheduled this " "wakeup to poll for updates. Check the current state of whatever you were " @@ -46,6 +49,71 @@ def _build_one_shot_cron(fire_time: datetime) -> str: ) +def _parse_iso(value: Any) -> datetime | None: + if not isinstance(value, str) or not value: + return None + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +async def find_expired_wakeup_cron_ids(client: Any, *, now: datetime) -> list[str]: + """Return the ids of ``thread_wakeup`` crons whose ``end_time`` has passed. + + Conservative: matches solely on ``metadata.kind == "thread_wakeup"`` AND a + past ``end_time``, so analyzer/dashboard crons are never selected. Paginates + fully before returning so the result is stable to delete afterwards. + """ + expired_ids: list[str] = [] + offset = 0 + while True: + page = await client.crons.search( + metadata={"kind": _WAKEUP_KIND}, + limit=_PURGE_PAGE_SIZE, + offset=offset, + ) + if not page: + break + for cron in page: + if not isinstance(cron, dict): + continue + end_time = _parse_iso(cron.get("end_time")) + cron_id = cron.get("cron_id") + if end_time is not None and end_time < now and isinstance(cron_id, str) and cron_id: + expired_ids.append(cron_id) + if len(page) < _PURGE_PAGE_SIZE: + break + offset += len(page) + return expired_ids + + +async def purge_expired_wakeup_crons(client: Any, *, now: datetime) -> int: + """Delete ``thread_wakeup`` crons whose ``end_time`` has already passed. + + Each wakeup is a thread-bound cron with an ``end_time`` (~90s past its fire) + that stops it re-firing, but the cron row itself is never removed, so dead + rows accumulate. This deletes only those dead rows. Returns the count deleted. + """ + expired_ids = await find_expired_wakeup_cron_ids(client, now=now) + deleted = 0 + for cron_id in expired_ids: + await client.crons.delete(cron_id) + deleted += 1 + return deleted + + +async def _purge_expired_wakeups_best_effort() -> None: + """Opportunistically purge expired wakeup crons; never raises.""" + try: + client = get_client(url=langgraph_url()) + deleted = await purge_expired_wakeup_crons(client, now=datetime.now(UTC)) + if deleted: + logger.info("Purged %d expired thread_wakeup cron(s)", deleted) + except Exception: + logger.warning("Failed to purge expired thread_wakeup crons", exc_info=True) + + async def _create_wakeup_cron( *, thread_id: str, @@ -130,6 +198,8 @@ async def schedule_thread_wakeup(delay_minutes: int, prompt: str | None = None) if value is not None: wakeup_configurable[key] = value + await _purge_expired_wakeups_best_effort() + try: return await _create_wakeup_cron( thread_id=thread_id, diff --git a/scripts/purge_wakeup_crons.py b/scripts/purge_wakeup_crons.py new file mode 100644 index 00000000..ad556593 --- /dev/null +++ b/scripts/purge_wakeup_crons.py @@ -0,0 +1,88 @@ +"""One-time backfill: delete expired ``thread_wakeup`` crons from a deployment. + +One-shot wakeup crons set an ``end_time`` that stops them re-firing, but the +cron row is never removed, so dead rows accumulate. The ``schedule_thread_wakeup`` +tool now purges these opportunistically; this script clears the backlog. + +Usage: + uv run python scripts/purge_wakeup_crons.py --dry-run + uv run python scripts/purge_wakeup_crons.py + +Resolves the deployment URL from ``--url`` or ``LANGGRAPH_URL`` / ``LANGGRAPH_URL_PROD``, +and the API key from ``LANGGRAPH_API_KEY`` / ``LANGSMITH_API_KEY`` / ``LANGSMITH_API_KEY_PROD``. +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +from datetime import UTC, datetime + +from langgraph_sdk import get_client + +from agent.tools.schedule_thread_wakeup import ( + find_expired_wakeup_cron_ids, + purge_expired_wakeup_crons, +) + +logger = logging.getLogger(__name__) + + +def _load_dotenv_if_available() -> None: + try: + from dotenv import load_dotenv + except ImportError: + return + load_dotenv() + + +def _resolve_url(arg_url: str | None) -> str: + url = arg_url or os.environ.get("LANGGRAPH_URL") or os.environ.get("LANGGRAPH_URL_PROD") + if not url: + raise RuntimeError("Set --url or LANGGRAPH_URL / LANGGRAPH_URL_PROD") + return url + + +def _resolve_api_key() -> str | None: + return ( + os.environ.get("LANGGRAPH_API_KEY") + or os.environ.get("LANGSMITH_API_KEY") + or os.environ.get("LANGSMITH_API_KEY_PROD") + ) + + +async def _run(url: str, api_key: str | None, dry_run: bool) -> None: + client = get_client(url=url, api_key=api_key) + now = datetime.now(UTC) + if dry_run: + expired = await find_expired_wakeup_cron_ids(client, now=now) + logger.info("[dry-run] %d expired thread_wakeup cron(s) would be deleted", len(expired)) + for cron_id in expired: + logger.info(" %s", cron_id) + return + deleted = await purge_expired_wakeup_crons(client, now=now) + logger.info("Deleted %d expired thread_wakeup cron(s)", deleted) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Purge expired thread_wakeup crons.") + parser.add_argument("--url", default=None, help="Deployment URL (defaults to env).") + parser.add_argument( + "--dry-run", + action="store_true", + help="List the crons that would be deleted without deleting them.", + ) + return parser.parse_args() + + +def main() -> None: + _load_dotenv_if_available() + logging.basicConfig(level=logging.INFO, format="%(message)s") + args = parse_args() + asyncio.run(_run(_resolve_url(args.url), _resolve_api_key(), args.dry_run)) + + +if __name__ == "__main__": + main() diff --git a/tests/test_schedule_thread_wakeup.py b/tests/test_schedule_thread_wakeup.py index 55b0c05e..50c4b1ad 100644 --- a/tests/test_schedule_thread_wakeup.py +++ b/tests/test_schedule_thread_wakeup.py @@ -1,13 +1,67 @@ from __future__ import annotations import importlib -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from typing import Any import pytest wakeup_tool = importlib.import_module("agent.tools.schedule_thread_wakeup") +# Captured before the autouse stub replaces it, for the one test that needs the real wrapper. +_real_purge_best_effort = wakeup_tool._purge_expired_wakeups_best_effort + + +@pytest.fixture(autouse=True) +def _stub_purge(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep the opportunistic purge from touching the network in every test.""" + + async def _noop() -> None: + return None + + monkeypatch.setattr(wakeup_tool, "_purge_expired_wakeups_best_effort", _noop) + + +class _FakeCrons: + def __init__(self, crons: list[dict[str, Any]]) -> None: + self._crons = list(crons) + self.deleted: list[str] = [] + self.search_calls: list[dict[str, Any]] = [] + + async def search( + self, + *, + metadata: dict[str, Any] | None = None, + limit: int = 10, + offset: int = 0, + **_: Any, + ) -> list[dict[str, Any]]: + self.search_calls.append({"metadata": metadata, "limit": limit, "offset": offset}) + items = [ + c + for c in self._crons + if not metadata + or all((c.get("metadata") or {}).get(k) == v for k, v in metadata.items()) + ] + return items[offset : offset + limit] + + async def delete(self, cron_id: str) -> None: + self.deleted.append(cron_id) + self._crons = [c for c in self._crons if c.get("cron_id") != cron_id] + + +class _FakeClient: + def __init__(self, crons: list[dict[str, Any]]) -> None: + self.crons = _FakeCrons(crons) + + +def _wakeup_cron(cron_id: str, end_time: datetime | None) -> dict[str, Any]: + return { + "cron_id": cron_id, + "end_time": end_time.isoformat() if end_time else None, + "metadata": {"kind": "thread_wakeup"}, + } + def _config(**overrides: Any) -> dict[str, Any]: base: dict[str, Any] = { @@ -241,3 +295,71 @@ def test_build_one_shot_cron_handles_month_boundary() -> None: assert parts[1] == "23" assert parts[2] == "31" assert parts[3] == "12" + + +async def test_purge_deletes_only_expired_wakeups() -> None: + now = datetime(2026, 6, 30, 22, 0, tzinfo=UTC) + client = _FakeClient( + [ + _wakeup_cron("expired-1", now - timedelta(hours=1)), + _wakeup_cron("expired-2", now - timedelta(days=1)), + _wakeup_cron("future-1", now + timedelta(hours=1)), + _wakeup_cron("no-end", None), + ] + ) + + deleted = await wakeup_tool.purge_expired_wakeup_crons(client, now=now) + + assert deleted == 2 + assert client.crons.deleted == ["expired-1", "expired-2"] + # Search is scoped to the thread_wakeup kind so other crons are never seen. + assert client.crons.search_calls[0]["metadata"] == {"kind": "thread_wakeup"} + + +async def test_purge_paginates(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(wakeup_tool, "_PURGE_PAGE_SIZE", 2) + now = datetime(2026, 6, 30, 22, 0, tzinfo=UTC) + client = _FakeClient([_wakeup_cron(f"expired-{i}", now - timedelta(hours=1)) for i in range(3)]) + + deleted = await wakeup_tool.purge_expired_wakeup_crons(client, now=now) + + assert deleted == 3 + assert sorted(client.crons.deleted) == ["expired-0", "expired-1", "expired-2"] + # Two pages fetched (offset 0 and 2), then a short final page ends the loop. + assert [c["offset"] for c in client.crons.search_calls] == [0, 2] + + +async def test_best_effort_purge_swallows_errors(monkeypatch: pytest.MonkeyPatch) -> None: + async def boom(*_: Any, **__: Any) -> int: + raise RuntimeError("search failed") + + monkeypatch.setattr(wakeup_tool, "purge_expired_wakeup_crons", boom) + monkeypatch.setattr(wakeup_tool, "get_client", lambda url: object()) + + # The real wrapper must never propagate — a purge failure can't block wakeups. + await _real_purge_best_effort() + + +async def test_schedule_purges_before_creating(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + + async def spy_purge() -> None: + calls.append("purge") + + async def fake_create_wakeup_cron(**kwargs: Any) -> dict[str, Any]: + calls.append("create") + return { + "success": True, + "cron_id": "cron-1", + "scheduled_for": "", + "thread_id": kwargs["thread_id"], + } + + monkeypatch.setattr(wakeup_tool, "get_config", _config) + monkeypatch.setattr(wakeup_tool, "_purge_expired_wakeups_best_effort", spy_purge) + monkeypatch.setattr(wakeup_tool, "_create_wakeup_cron", fake_create_wakeup_cron) + + result = await wakeup_tool.schedule_thread_wakeup(5) + + assert result["success"] is True + assert calls == ["purge", "create"]