"""Tool that schedules a one-shot re-trigger of the current agent thread.""" from __future__ import annotations import logging from datetime import UTC, datetime, timedelta from typing import Any from langgraph.config import get_config from langgraph_sdk import get_client from ..utils.thread_ops import langgraph_url logger = logging.getLogger(__name__) _AGENT_ASSISTANT_ID = "agent" _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 " "waiting on and continue from there." ) def _ceil_to_next_minute(value: datetime) -> datetime: """Round a datetime up to the next whole minute.""" rounded = value.replace(second=0, microsecond=0) if rounded == value: return rounded return rounded + timedelta(minutes=1) def _build_one_shot_cron(fire_time: datetime) -> str: """Build a 5-field cron expression that fires at ``fire_time`` (UTC).""" return " ".join( [ str(fire_time.minute), str(fire_time.hour), str(fire_time.day), str(fire_time.month), "*", ] ) 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, fire_time: datetime, prompt: str, configurable: dict[str, Any], ) -> dict[str, Any]: client = get_client(url=langgraph_url()) schedule = _build_one_shot_cron(fire_time) end_time = fire_time + timedelta(seconds=_END_TIME_PADDING_SECONDS) run_config: dict[str, Any] = {"configurable": configurable} cron = await client.crons.create_for_thread( thread_id, _AGENT_ASSISTANT_ID, schedule=schedule, input={"messages": [{"role": "user", "content": prompt}]}, config=run_config, end_time=end_time, timezone="UTC", metadata={ "kind": "thread_wakeup", "thread_id": thread_id, }, ) cron_id = cron.get("cron_id") if isinstance(cron, dict) else getattr(cron, "cron_id", None) return { "success": True, "cron_id": cron_id, "scheduled_for": fire_time.isoformat(), "thread_id": thread_id, } async def schedule_thread_wakeup(delay_minutes: int, prompt: str | None = None) -> dict[str, Any]: """Schedule a one-shot re-trigger of the current thread after a delay. Use this when you need to poll or check back on something later — e.g. waiting for CI to finish, a deploy to complete, or an external process to settle. The current thread will be re-invoked with the given prompt (or a default wakeup message) after the specified delay. Args: delay_minutes: How many minutes from now to wait before re-triggering. Minimum 1 minute, maximum 1440 (24 hours). prompt: Optional message to send to the thread when it wakes up. If omitted, a default polling prompt is used. Returns a dict with ``success``, ``cron_id``, ``scheduled_for`` (ISO UTC), and ``thread_id``. """ if not isinstance(delay_minutes, int) or delay_minutes < 1: return {"success": False, "error": "delay_minutes must be a positive integer (>= 1)"} delay_seconds = delay_minutes * 60 if delay_seconds < _MIN_DELAY_SECONDS: return {"success": False, "error": "delay must be at least 1 minute"} if delay_seconds > _MAX_DELAY_SECONDS: return {"success": False, "error": "delay must be at most 1440 minutes (24 hours)"} config = get_config() configurable = config.get("configurable", {}) if isinstance(config, dict) else {} thread_id = configurable.get("thread_id") if not isinstance(thread_id, str) or not thread_id: return {"success": False, "error": "No thread_id in current run config"} fire_time = _ceil_to_next_minute(datetime.now(UTC) + timedelta(seconds=delay_seconds)) wakeup_prompt = ( prompt.strip() if isinstance(prompt, str) and prompt.strip() else _DEFAULT_WAKEUP_PROMPT ) passthrough_keys = ( "repo", "source", "slack_thread", "linear_issue", "github_login", "user_email", "schedule_id", ) wakeup_configurable: dict[str, Any] = {"thread_id": thread_id} for key in passthrough_keys: value = configurable.get(key) 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, fire_time=fire_time, prompt=wakeup_prompt, configurable=wakeup_configurable, ) except Exception as exc: logger.exception("Failed to schedule thread wakeup for %s", thread_id) return {"success": False, "error": str(exc)}