| 1 | from __future__ import annotations |
| 2 | |
| 3 | import types |
| 4 | from typing import Any, Mapping, TypedDict, Union, get_args, get_origin, get_type_hints |
| 5 | |
| 6 | from dataclasses import dataclass, replace |
| 7 | |
| 8 | import pytz # type: ignore[import-untyped] |
| 9 | |
| 10 | from agent import AgentContext, AgentContextType |
| 11 | |
| 12 | from helpers.dotenv import get_dotenv_value |
| 13 | from helpers.localization import Localization |
| 14 | from helpers.task_scheduler import TaskScheduler |
| 15 | |
| 16 | |
| 17 | class SnapshotV1(TypedDict): |
| 18 | deselect_chat: bool |
| 19 | context: str |
| 20 | contexts: list[dict[str, Any]] | None |
| 21 | tasks: list[dict[str, Any]] | None |
| 22 | logs: list[dict[str, Any]] |
| 23 | log_guid: str |
| 24 | log_version: int |
| 25 | # Historical behavior: when no context is selected, log_progress is 0 (falsy). |
| 26 | # When a context is active, it is usually a string. |
| 27 | log_progress: str | int |
| 28 | log_progress_active: bool |
| 29 | paused: bool |
| 30 | notifications: list[dict[str, Any]] |
| 31 | notifications_guid: str |
| 32 | notifications_version: int |
| 33 | |
| 34 | @dataclass(frozen=True) |
| 35 | class StateRequestV1: |
| 36 | context: str | None |
| 37 | log_from: int |
| 38 | notifications_from: int |
| 39 | timezone: str |
| 40 | collections_delta: bool = False |
| 41 | |
| 42 | |
| 43 | class StateRequestValidationError(ValueError): |
| 44 | def __init__( |
| 45 | self, |
| 46 | *, |
| 47 | reason: str, |
| 48 | message: str, |
| 49 | details: dict[str, Any] | None = None, |
| 50 | ) -> None: |
| 51 | super().__init__(message) |
| 52 | self.reason = reason |
| 53 | self.details = details or {} |
| 54 | |
| 55 | |
| 56 | def _annotation_to_isinstance_types(annotation: Any) -> tuple[type, ...]: |
| 57 | """Convert type annotation to tuple suitable for isinstance().""" |
| 58 | origin = get_origin(annotation) |
| 59 | |
| 60 | # Handle Union (typing.Union or types.UnionType from X | Y) |
| 61 | _union_type = getattr(types, "UnionType", None) |
| 62 | if origin is Union or origin is _union_type: |
| 63 | result: list[type] = [] |
| 64 | for arg in get_args(annotation): |
| 65 | result.extend(_annotation_to_isinstance_types(arg)) |
| 66 | return tuple(result) |
| 67 | |
| 68 | # Generic aliases: list[X] -> list, dict[K,V] -> dict |
| 69 | if origin is not None: |
| 70 | return (origin,) |
| 71 | |
| 72 | if isinstance(annotation, type): |
| 73 | return (annotation,) |
| 74 | |
| 75 | return () |
| 76 | |
| 77 | |
| 78 | def _build_schema_from_typeddict(td: type) -> dict[str, tuple[type, ...]]: |
| 79 | """Extract field names and isinstance-compatible types from TypedDict.""" |
| 80 | return {k: _annotation_to_isinstance_types(v) for k, v in get_type_hints(td).items()} |
| 81 | |
| 82 | |
| 83 | _SNAPSHOT_V1_SCHEMA = _build_schema_from_typeddict(SnapshotV1) |
| 84 | SNAPSHOT_SCHEMA_V1_KEYS: tuple[str, ...] = tuple(_SNAPSHOT_V1_SCHEMA.keys()) |
| 85 | |
| 86 | |
| 87 | def validate_snapshot_schema_v1(snapshot: Mapping[str, Any]) -> None: |
| 88 | if not isinstance(snapshot, dict): |
| 89 | raise TypeError("snapshot must be a dict") |
| 90 | expected = set(SNAPSHOT_SCHEMA_V1_KEYS) |
| 91 | actual = set(snapshot.keys()) |
| 92 | missing = sorted(expected - actual) |
| 93 | extra = sorted(actual - expected) |
| 94 | if missing or extra: |
| 95 | message = "snapshot schema mismatch" |
| 96 | if missing: |
| 97 | message += f"; missing={missing}" |
| 98 | if extra: |
| 99 | message += f"; unexpected={extra}" |
| 100 | raise ValueError(message) |
| 101 | |
| 102 | for key, expected_types in _SNAPSHOT_V1_SCHEMA.items(): |
| 103 | if expected_types and not isinstance(snapshot.get(key), expected_types): |
| 104 | type_desc = " | ".join(t.__name__ for t in expected_types) |
| 105 | raise TypeError(f"snapshot.{key} must be {type_desc}") |
| 106 | |
| 107 | |
| 108 | def _coerce_non_negative_int(value: Any, default: int = 0) -> int: |
| 109 | try: |
| 110 | as_int = int(value) |
| 111 | except (TypeError, ValueError): |
| 112 | return default |
| 113 | return as_int if as_int >= 0 else default |
| 114 | |
| 115 | |
| 116 | def _get_agent_profile_labels() -> dict[str, str]: |
| 117 | try: |
| 118 | from helpers import subagents |
| 119 | |
| 120 | return { |
| 121 | str(item.get("key") or ""): str(item.get("label") or item.get("key") or "") |
| 122 | for item in subagents.get_all_agents_list() |
| 123 | if item.get("key") |
| 124 | } |
| 125 | except Exception: |
| 126 | return {} |
| 127 | |
| 128 | |
| 129 | def _apply_agent_profile_metadata( |
| 130 | context_data: dict[str, Any], |
| 131 | ctx: AgentContext, |
| 132 | labels: dict[str, str], |
| 133 | ) -> None: |
| 134 | agent_config = getattr(getattr(ctx, "agent0", None), "config", None) |
| 135 | profile = str( |
| 136 | getattr(agent_config, "profile", None) |
| 137 | or getattr(getattr(ctx, "config", None), "profile", "") |
| 138 | or "" |
| 139 | ) |
| 140 | context_data["agent_profile"] = profile |
| 141 | context_data["agent_profile_label"] = labels.get(profile, profile) if profile else "" |
| 142 | |
| 143 | |
| 144 | def _prune_missing_saved_contexts() -> None: |
| 145 | from helpers import persist_chat |
| 146 | |
| 147 | saved_ids = persist_chat.saved_chat_ids() |
| 148 | for ctx in AgentContext.all(): |
| 149 | if ctx.type == AgentContextType.BACKGROUND or ctx.is_running(): |
| 150 | continue |
| 151 | if ( |
| 152 | ctx.data.get(persist_chat.SAVED_CHAT_CONTEXT_DATA_KEY) |
| 153 | and ctx.id not in saved_ids |
| 154 | ): |
| 155 | AgentContext.remove(ctx.id) |
| 156 | |
| 157 | |
| 158 | def parse_state_request_payload(payload: Mapping[str, Any]) -> StateRequestV1: |
| 159 | context = payload.get("context") |
| 160 | log_from = payload.get("log_from") |
| 161 | notifications_from = payload.get("notifications_from") |
| 162 | timezone = payload.get("timezone") |
| 163 | collections_delta = payload.get("collections_delta", False) |
| 164 | |
| 165 | if context is not None and not isinstance(context, str): |
| 166 | raise StateRequestValidationError( |
| 167 | reason="context_type", |
| 168 | message="context must be a string or null", |
| 169 | details={"context_type": type(context).__name__}, |
| 170 | ) |
| 171 | if not isinstance(log_from, int) or log_from < 0: |
| 172 | raise StateRequestValidationError( |
| 173 | reason="log_from", |
| 174 | message="log_from must be an integer >= 0", |
| 175 | details={"log_from": log_from}, |
| 176 | ) |
| 177 | if not isinstance(notifications_from, int) or notifications_from < 0: |
| 178 | raise StateRequestValidationError( |
| 179 | reason="notifications_from", |
| 180 | message="notifications_from must be an integer >= 0", |
| 181 | details={"notifications_from": notifications_from}, |
| 182 | ) |
| 183 | if not isinstance(timezone, str) or not timezone.strip(): |
| 184 | raise StateRequestValidationError( |
| 185 | reason="timezone_empty", |
| 186 | message="timezone must be a non-empty string", |
| 187 | details={"timezone": timezone}, |
| 188 | ) |
| 189 | if not isinstance(collections_delta, bool): |
| 190 | raise StateRequestValidationError( |
| 191 | reason="collections_delta_type", |
| 192 | message="collections_delta must be a boolean", |
| 193 | details={"collections_delta_type": type(collections_delta).__name__}, |
| 194 | ) |
| 195 | |
| 196 | tz = timezone.strip() |
| 197 | try: |
| 198 | pytz.timezone(tz) |
| 199 | except pytz.exceptions.UnknownTimeZoneError as exc: |
| 200 | raise StateRequestValidationError( |
| 201 | reason="timezone_invalid", |
| 202 | message="timezone must be a valid IANA timezone name", |
| 203 | details={"timezone": tz}, |
| 204 | ) from exc |
| 205 | |
| 206 | ctxid: str | None = context.strip() if isinstance(context, str) else None |
| 207 | if ctxid == "": |
| 208 | ctxid = None |
| 209 | return StateRequestV1( |
| 210 | context=ctxid, |
| 211 | log_from=log_from, |
| 212 | notifications_from=notifications_from, |
| 213 | timezone=tz, |
| 214 | collections_delta=collections_delta, |
| 215 | ) |
| 216 | |
| 217 | |
| 218 | def _coerce_state_request_inputs( |
| 219 | *, |
| 220 | context: Any, |
| 221 | log_from: Any, |
| 222 | notifications_from: Any, |
| 223 | timezone: Any, |
| 224 | ) -> StateRequestV1: |
| 225 | tz = timezone if isinstance(timezone, str) and timezone else None |
| 226 | tz = tz or get_dotenv_value("DEFAULT_USER_TIMEZONE", Localization.get().get_timezone()) |
| 227 | |
| 228 | ctxid: str | None = context.strip() if isinstance(context, str) else None |
| 229 | if ctxid == "": |
| 230 | ctxid = None |
| 231 | |
| 232 | return StateRequestV1( |
| 233 | context=ctxid, |
| 234 | log_from=_coerce_non_negative_int(log_from, default=0), |
| 235 | notifications_from=_coerce_non_negative_int(notifications_from, default=0), |
| 236 | timezone=tz, |
| 237 | ) |
| 238 | |
| 239 | |
| 240 | def advance_state_request_after_snapshot( |
| 241 | request: StateRequestV1, |
| 242 | snapshot: Mapping[str, Any], |
| 243 | ) -> StateRequestV1: |
| 244 | log_from = request.log_from |
| 245 | notifications_from = request.notifications_from |
| 246 | |
| 247 | try: |
| 248 | log_from = int(snapshot.get("log_version", log_from)) |
| 249 | except (TypeError, ValueError): |
| 250 | pass |
| 251 | |
| 252 | try: |
| 253 | notifications_from = int(snapshot.get("notifications_version", notifications_from)) |
| 254 | except (TypeError, ValueError): |
| 255 | pass |
| 256 | |
| 257 | return replace( |
| 258 | request, |
| 259 | log_from=log_from, |
| 260 | notifications_from=notifications_from, |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | async def build_snapshot_from_request( |
| 265 | *, request: StateRequestV1, include_collections: bool = True |
| 266 | ) -> SnapshotV1: |
| 267 | """Build a poll-shaped snapshot for both /poll and state_push.""" |
| 268 | |
| 269 | localization = Localization.get() |
| 270 | previous_timezone = localization.get_timezone() |
| 271 | localization.set_timezone(request.timezone) |
| 272 | current_timezone = localization.get_timezone() |
| 273 | if current_timezone != previous_timezone: |
| 274 | _notify_timezone_changed(previous_timezone, current_timezone) |
| 275 | |
| 276 | ctxid = request.context if isinstance(request.context, str) else "" |
| 277 | ctxid = ctxid.strip() |
| 278 | |
| 279 | from_no = _coerce_non_negative_int(request.log_from, default=0) |
| 280 | notifications_from_no = _coerce_non_negative_int(request.notifications_from, default=0) |
| 281 | |
| 282 | if include_collections: |
| 283 | _prune_missing_saved_contexts() |
| 284 | |
| 285 | active_context = AgentContext.get(ctxid) if ctxid else None |
| 286 | |
| 287 | if active_context: |
| 288 | log_output = active_context.log.output(start=from_no) |
| 289 | logs = log_output.items |
| 290 | log_end = log_output.end |
| 291 | else: |
| 292 | logs = [] |
| 293 | log_end = 0 |
| 294 | |
| 295 | notification_manager = AgentContext.get_notification_manager() |
| 296 | notifications, notifications_guid, notifications_version = ( |
| 297 | notification_manager.output_with_state(start=notifications_from_no) |
| 298 | ) |
| 299 | |
| 300 | scheduler = TaskScheduler.get() |
| 301 | |
| 302 | ctxs: list[dict[str, Any]] = [] |
| 303 | tasks: list[dict[str, Any]] = [] |
| 304 | processed_contexts: set[str] = set() |
| 305 | agent_profile_labels = _get_agent_profile_labels() if include_collections else {} |
| 306 | |
| 307 | all_ctxs = AgentContext.all() if include_collections else [] |
| 308 | for ctx in all_ctxs: |
| 309 | if ctx.id in processed_contexts: |
| 310 | continue |
| 311 | |
| 312 | if ctx.type == AgentContextType.BACKGROUND: |
| 313 | processed_contexts.add(ctx.id) |
| 314 | continue |
| 315 | |
| 316 | context_data = ctx.output() |
| 317 | _apply_agent_profile_metadata(context_data, ctx, agent_profile_labels) |
| 318 | |
| 319 | context_task = scheduler.get_task_by_uuid(ctx.id) |
| 320 | is_task_context = context_task is not None and context_task.context_id == ctx.id |
| 321 | |
| 322 | if not is_task_context: |
| 323 | ctxs.append(context_data) |
| 324 | else: |
| 325 | task_details = scheduler.serialize_task(ctx.id) |
| 326 | if task_details: |
| 327 | context_data.update( |
| 328 | { |
| 329 | "task_name": task_details.get("name"), |
| 330 | "uuid": task_details.get("uuid"), |
| 331 | "state": task_details.get("state"), |
| 332 | "type": task_details.get("type"), |
| 333 | "system_prompt": task_details.get("system_prompt"), |
| 334 | "prompt": task_details.get("prompt"), |
| 335 | "last_run": task_details.get("last_run"), |
| 336 | "last_result": task_details.get("last_result"), |
| 337 | "attachments": task_details.get("attachments", []), |
| 338 | "context_id": task_details.get("context_id"), |
| 339 | } |
| 340 | ) |
| 341 | |
| 342 | if task_details.get("type") == "scheduled": |
| 343 | context_data["schedule"] = task_details.get("schedule") |
| 344 | elif task_details.get("type") == "planned": |
| 345 | context_data["plan"] = task_details.get("plan") |
| 346 | else: |
| 347 | context_data["token"] = task_details.get("token") |
| 348 | |
| 349 | tasks.append(context_data) |
| 350 | |
| 351 | processed_contexts.add(ctx.id) |
| 352 | |
| 353 | ctxs.sort(key=lambda x: x["created_at"], reverse=True) |
| 354 | tasks.sort(key=lambda x: x["created_at"], reverse=True) |
| 355 | |
| 356 | snapshot: SnapshotV1 = { |
| 357 | "deselect_chat": bool(ctxid) and active_context is None, |
| 358 | "context": active_context.id if active_context else "", |
| 359 | "contexts": ctxs if include_collections else None, |
| 360 | "tasks": tasks if include_collections else None, |
| 361 | "logs": logs, |
| 362 | "log_guid": active_context.log.guid if active_context else "", |
| 363 | "log_version": log_end, |
| 364 | "log_progress": active_context.log.progress if active_context else 0, |
| 365 | "log_progress_active": bool(active_context.log.progress_active) if active_context else False, |
| 366 | "paused": active_context.paused if active_context else False, |
| 367 | "notifications": notifications, |
| 368 | "notifications_guid": notifications_guid, |
| 369 | "notifications_version": notifications_version, |
| 370 | } |
| 371 | |
| 372 | validate_snapshot_schema_v1(snapshot) |
| 373 | return snapshot |
| 374 | |
| 375 | |
| 376 | def _notify_timezone_changed(previous_timezone: str, current_timezone: str) -> None: |
| 377 | try: |
| 378 | from helpers import plugins |
| 379 | |
| 380 | plugins.call_plugin_hook( |
| 381 | "_office", |
| 382 | "timezone_changed", |
| 383 | None, |
| 384 | previous_timezone=previous_timezone, |
| 385 | timezone=current_timezone, |
| 386 | ) |
| 387 | except Exception: |
| 388 | return |
| 389 | |
| 390 | |
| 391 | async def build_snapshot( |
| 392 | *, |
| 393 | context: str | None, |
| 394 | log_from: int, |
| 395 | notifications_from: int, |
| 396 | timezone: str | None, |
| 397 | ) -> SnapshotV1: |
| 398 | request = _coerce_state_request_inputs( |
| 399 | context=context, |
| 400 | log_from=log_from, |
| 401 | notifications_from=notifications_from, |
| 402 | timezone=timezone, |
| 403 | ) |
| 404 | return await build_snapshot_from_request(request=request) |