src/swf_common_lib/episodes.py

0001 """Episode building: the generic engine behind workflow episode records.
0002
0003 An episode is a durable Snapper record of one bounded activity — a
0004 workflow execution — whose record is the event sequence at native
0005 resolution (snapper-ai docs/EPISODES.md). This module carries the
0006 workflow-agnostic machinery:
0007
0008 - ``EpisodeDefinition`` — the contract a workflow-specific module
0009   implements: which bus messages are events, how participants are
0010   recognized, and the completion pass that joins late records.
0011 - ``EpisodeBuilder`` — the engine an agent drives: it watches bus
0012   traffic, opens the episode on first sight of an execution, appends
0013   events as they arrive, and runs the definition's completion pass
0014   before closing.
0015 - ``MonitorEpisodeIngest`` — the REST client for the swf-monitor
0016   episode ingest endpoints.
0017
0018 A workflow gains an episode record by implementing one definition and
0019 naming it in the episode builder agent's configuration; nothing here
0020 changes.
0021 """
0022
0023 import logging
0024 from datetime import datetime, timezone
0025 from typing import Dict, List, Optional
0026
0027 import requests
0028
0029 logger = logging.getLogger(__name__)
0030
0031
0032 def utc_now_iso() -> str:
0033     return datetime.now(timezone.utc).isoformat()
0034
0035
0036 class EpisodeIngestError(Exception):
0037     """A rejected or failed episode ingest call."""
0038
0039
0040 class MonitorEpisodeIngest:
0041     """REST client for the swf-monitor episode ingest endpoints.
0042
0043     Endpoints (token-authenticated):
0044       POST {base}/api/snapper/episodes/open/
0045       POST {base}/api/snapper/episodes/append/
0046       POST {base}/api/snapper/episodes/close/
0047     """
0048
0049     def __init__(self, base_url: str, token: str,
0050                  builder_identity: str, session=None):
0051         self.base_url = (base_url or "").rstrip("/")
0052         self.builder_identity = builder_identity
0053         self.session = session or requests.Session()
0054         if token:
0055             self.session.headers.update({"Authorization": f"Token {token}"})
0056
0057     def _post(self, path: str, payload: Dict) -> Dict:
0058         url = f"{self.base_url}/api/snapper/episodes/{path}/"
0059         payload = dict(payload, builder_identity=self.builder_identity)
0060         try:
0061             response = self.session.post(url, json=payload, timeout=30)
0062         except requests.RequestException as exc:
0063             raise EpisodeIngestError(f"POST {url} failed: {exc}") from exc
0064         if response.status_code >= 400:
0065             raise EpisodeIngestError(
0066                 f"POST {url} returned {response.status_code}: "
0067                 f"{response.text[:500]}"
0068             )
0069         return response.json()
0070
0071     def open(self, scope: str, episode_id: str, started_at: str,
0072              label: str = "", kind: str = "",
0073              summary: Optional[Dict] = None) -> Dict:
0074         return self._post("open", {
0075             "scope": scope, "episode_id": episode_id,
0076             "started_at": started_at, "label": label, "kind": kind,
0077             "summary": summary or {},
0078         })
0079
0080     def append(self, scope: str, episode_id: str,
0081                events: Optional[List[Dict]] = None,
0082                participants: Optional[List[Dict]] = None) -> Dict:
0083         return self._post("append", {
0084             "scope": scope, "episode_id": episode_id,
0085             "events": events or [], "participants": participants or [],
0086         })
0087
0088     def close(self, scope: str, episode_id: str, ended_at: str,
0089               summary: Optional[Dict] = None) -> Dict:
0090         return self._post("close", {
0091             "scope": scope, "episode_id": episode_id,
0092             "ended_at": ended_at, "summary": summary or {},
0093         })
0094
0095     def _get(self, path: str) -> Dict:
0096         url = f"{self.base_url}/api/snapper/{path}"
0097         try:
0098             response = self.session.get(url, timeout=30)
0099         except requests.RequestException as exc:
0100             raise EpisodeIngestError(f"GET {url} failed: {exc}") from exc
0101         if response.status_code >= 400:
0102             raise EpisodeIngestError(
0103                 f"GET {url} returned {response.status_code}: "
0104                 f"{response.text[:500]}")
0105         return response.json()
0106
0107     def open_episodes(self, scope: str) -> List[Dict]:
0108         """The scope's episodes without a recorded end."""
0109         listing = self._get(f"{scope}/episodes/")
0110         return [e for e in listing.get("episodes", [])
0111                 if not e.get("ended_at")]
0112
0113     def episode(self, scope: str, episode_id: str) -> Dict:
0114         """One episode's full record."""
0115         return self._get(f"{scope}/episodes/{episode_id}/")
0116
0117
0118 class EpisodeDefinition:
0119     """The workflow-specific contract. Subclass per workflow.
0120
0121     The builder consults ``matches`` to route bus messages, converts
0122     them through ``event_from_message`` / ``participants_from_message``,
0123     and treats ``is_end`` as the execution's end signal. After the end
0124     signal the builder calls ``completion_poll`` on every tick until it
0125     returns True or ``completion_deadline_seconds`` passes, then closes
0126     the episode.
0127     """
0128
0129     #: Snapper scope the episodes belong to (e.g. 'testbed').
0130     scope = ""
0131     #: Workflow name this definition covers; used by the default
0132     #: ``matches`` against the execution id prefix.
0133     workflow_name = ""
0134     #: Seconds after the end signal within which completion_poll must
0135     #: finish; the episode closes regardless when the deadline passes.
0136     completion_deadline_seconds = 1800
0137
0138     def matches(self, message: Dict) -> bool:
0139         execution_id = message.get("execution_id") or ""
0140         return bool(self.workflow_name) and execution_id.startswith(
0141             f"{self.workflow_name}-"
0142         )
0143
0144     def label(self, message: Dict) -> str:
0145         run_id = message.get("run_id")
0146         return f"run {run_id}" if run_id else ""
0147
0148     def started_at(self, message: Dict) -> str:
0149         """Timezone-aware ISO start for the episode. The default is the
0150         arrival time; definitions that trust their messages' stamps
0151         override this with a normalized message time."""
0152         return utc_now_iso()
0153
0154     def ended_at(self, message: Dict) -> str:
0155         """Timezone-aware ISO end for the episode, from the end-signal
0156         message. Same default and override contract as started_at —
0157         essential for backfilled episodes, whose close must carry the
0158         recorded end rather than the replay time."""
0159         return utc_now_iso()
0160
0161     def event_from_message(self, message: Dict) -> Optional[Dict]:
0162         """Bus message -> event dict ``{time, kind, participant,
0163         counterpart?, payload?}``, or None to ignore the message."""
0164         raise NotImplementedError
0165
0166     def participants_from_message(self, message: Dict) -> List[Dict]:
0167         """Bus message -> participant upserts ``[{id, label?, kind?,
0168         born_at?, died_at?, detail?}]``."""
0169         return []
0170
0171     def is_end(self, message: Dict) -> bool:
0172         return message.get("msg_type") == "end_run"
0173
0174     def completion_poll(self, context: "EpisodeContext",
0175                         ingest: MonitorEpisodeIngest) -> bool:
0176         """Join late records (workload states, registries) by appending
0177         further events and participants; return True when complete.
0178         Called repeatedly after the end signal until True or deadline."""
0179         return True
0180
0181     def summary(self, context: "EpisodeContext") -> Dict:
0182         """The episode's closing summary document."""
0183         return {}
0184
0185
0186 class EpisodeContext:
0187     """Mutable per-execution state the builder shares with a definition."""
0188
0189     def __init__(self, definition: EpisodeDefinition, episode_id: str):
0190         self.definition = definition
0191         self.episode_id = episode_id
0192         self.opened_at = utc_now_iso()
0193         self.end_seen_at: Optional[str] = None
0194         self.ended_at: Optional[str] = None
0195         self.first_message: Optional[Dict] = None
0196         self.last_message: Optional[Dict] = None
0197         #: Scratch space for the definition (run ids, task ids, ...).
0198         self.notes: Dict = {}
0199         #: Participant ids already reported, so steady message traffic
0200         #: does not re-upsert its sender on every message.
0201         self.seen_participants: set = set()
0202
0203
0204 class EpisodeBuilder:
0205     """The engine: routes bus messages to armed definitions and drives
0206     each execution's episode through open, append, completion, close.
0207
0208     The driving agent calls ``handle_message`` for every bus message it
0209     receives and ``tick`` periodically (heartbeat cadence is enough);
0210     ticks run completion polls and close episodes past their deadline.
0211     Ingest failures are logged and surfaced through the return values;
0212     the builder never raises out of ``handle_message`` or ``tick`` so a
0213     broken ingest cannot take the listening agent down with it.
0214     """
0215
0216     def __init__(self, definitions: List[EpisodeDefinition],
0217                  ingest: MonitorEpisodeIngest):
0218         self.definitions = list(definitions)
0219         self.ingest = ingest
0220         self.active: Dict[str, EpisodeContext] = {}
0221
0222     def adopt_open_episodes(self) -> int:
0223         """Adopt the builder identity's open episodes, so a restarted
0224         builder resumes what its predecessor left live: an episode
0225         whose end signal already passed is driven to completion and
0226         close by the next ticks; one still mid-flight keeps appending
0227         as its messages arrive. Returns the number adopted."""
0228         adopted = 0
0229         scopes = {d.scope for d in self.definitions if d.scope}
0230         for scope in scopes:
0231             try:
0232                 open_records = self.ingest.open_episodes(scope)
0233             except EpisodeIngestError as exc:
0234                 logger.error("open-episode listing failed for %s: %s",
0235                              scope, exc)
0236                 continue
0237             for entry in open_records:
0238                 episode_id = entry.get("episode_id")
0239                 if not episode_id or episode_id in self.active:
0240                     continue
0241                 definition = next(
0242                     (d for d in self.definitions
0243                      if d.scope == scope
0244                      and d.workflow_name == entry.get("kind")), None)
0245                 if definition is None:
0246                     continue
0247                 try:
0248                     record = self.ingest.episode(scope, episode_id)
0249                 except EpisodeIngestError as exc:
0250                     logger.error("episode fetch failed for %s: %s",
0251                                  episode_id, exc)
0252                     continue
0253                 context = EpisodeContext(definition, episode_id)
0254                 for event in record.get("events", []):
0255                     context.seen_participants.add(event.get("participant"))
0256                     if definition.is_end({"msg_type": event.get("kind")}):
0257                         context.end_seen_at = utc_now_iso()
0258                         context.ended_at = event.get("time")
0259                 self.active[episode_id] = context
0260                 adopted += 1
0261                 logger.info("adopted open episode %s (%s)", episode_id,
0262                             "end seen" if context.end_seen_at
0263                             else "still live")
0264         return adopted
0265
0266     def handle_message(self, message: Dict) -> bool:
0267         """Route one bus message; returns True if it joined an episode."""
0268         execution_id = message.get("execution_id")
0269         if not execution_id:
0270             return False
0271         # Arrival stamp: the fallback event time for messages whose
0272         # writer stamps no timestamp of its own.
0273         message.setdefault("_received_at", utc_now_iso())
0274         for definition in self.definitions:
0275             if definition.matches(message):
0276                 break
0277         else:
0278             return False
0279         try:
0280             context = self.active.get(execution_id)
0281             if context is None:
0282                 context = EpisodeContext(definition, execution_id)
0283                 context.first_message = message
0284                 self.active[execution_id] = context
0285                 self.ingest.open(
0286                     scope=definition.scope,
0287                     episode_id=execution_id,
0288                     started_at=definition.started_at(message),
0289                     label=definition.label(message),
0290                     kind=definition.workflow_name,
0291                 )
0292             context.last_message = message
0293             event = definition.event_from_message(message)
0294             participants = [
0295                 entry for entry in definition.participants_from_message(message)
0296                 if not (entry.get("id") in context.seen_participants
0297                         and "died_at" not in entry)
0298             ]
0299             for entry in participants:
0300                 context.seen_participants.add(entry.get("id"))
0301             if event or participants:
0302                 self.ingest.append(
0303                     scope=definition.scope,
0304                     episode_id=execution_id,
0305                     events=[event] if event else [],
0306                     participants=participants,
0307                 )
0308             if definition.is_end(message):
0309                 context.end_seen_at = utc_now_iso()
0310                 context.ended_at = definition.ended_at(message)
0311             return True
0312         except EpisodeIngestError as exc:
0313             logger.error("episode ingest failed for %s: %s",
0314                          execution_id, exc)
0315             return False
0316
0317     def tick(self) -> None:
0318         """Drive pending completions; safe to call at any cadence."""
0319         for execution_id in list(self.active):
0320             context = self.active[execution_id]
0321             if context.end_seen_at is None:
0322                 continue
0323             definition = context.definition
0324             try:
0325                 done = definition.completion_poll(context, self.ingest)
0326             except Exception as exc:
0327                 logger.error("completion poll failed for %s: %s",
0328                              execution_id, exc)
0329                 done = False
0330             deadline_passed = self._deadline_passed(context)
0331             if not done and not deadline_passed:
0332                 continue
0333             if deadline_passed and not done:
0334                 logger.warning(
0335                     "episode %s closed at completion deadline with the "
0336                     "completion pass unfinished", execution_id)
0337             try:
0338                 self.ingest.close(
0339                     scope=definition.scope,
0340                     episode_id=execution_id,
0341                     ended_at=context.ended_at or context.end_seen_at,
0342                     summary=definition.summary(context),
0343                 )
0344             except EpisodeIngestError as exc:
0345                 logger.error("episode close failed for %s: %s",
0346                              execution_id, exc)
0347             del self.active[execution_id]
0348
0349     def _deadline_passed(self, context: EpisodeContext) -> bool:
0350         if context.end_seen_at is None:
0351             return False
0352         seen = datetime.fromisoformat(context.end_seen_at)
0353         elapsed = datetime.now(timezone.utc) - seen
0354         return elapsed.total_seconds() > (
0355             context.definition.completion_deadline_seconds
0356         )