| 1 | import copy |
| 2 | import json |
| 3 | import threading |
| 4 | import time |
| 5 | import uuid |
| 6 | from collections import OrderedDict |
| 7 | from dataclasses import dataclass |
| 8 | from typing import Any, Literal, Optional, TYPE_CHECKING, TypeVar, cast |
| 9 | |
| 10 | from helpers.secrets import get_secrets_manager |
| 11 | from helpers.strings import truncate_text_by_ratio |
| 12 | |
| 13 | |
| 14 | if TYPE_CHECKING: |
| 15 | from agent import AgentContext |
| 16 | |
| 17 | |
| 18 | _MARK_DIRTY_ALL = None |
| 19 | _MARK_DIRTY_FOR_CONTEXT = None |
| 20 | |
| 21 | |
| 22 | def _lazy_mark_dirty_all(*, reason: str | None = None) -> None: |
| 23 | # Lazy import to avoid circular import at module load time (AgentContext -> Log). |
| 24 | global _MARK_DIRTY_ALL |
| 25 | if _MARK_DIRTY_ALL is None: |
| 26 | from helpers.state_monitor_integration import mark_dirty_all |
| 27 | |
| 28 | _MARK_DIRTY_ALL = mark_dirty_all |
| 29 | _MARK_DIRTY_ALL(reason=reason) |
| 30 | |
| 31 | |
| 32 | def _lazy_mark_dirty_for_context(context_id: str, *, reason: str | None = None) -> None: |
| 33 | # Lazy import to avoid circular import at module load time (AgentContext -> Log). |
| 34 | global _MARK_DIRTY_FOR_CONTEXT |
| 35 | if _MARK_DIRTY_FOR_CONTEXT is None: |
| 36 | from helpers.state_monitor_integration import mark_dirty_for_context |
| 37 | |
| 38 | _MARK_DIRTY_FOR_CONTEXT = mark_dirty_for_context |
| 39 | _MARK_DIRTY_FOR_CONTEXT(context_id, reason=reason, include_collections=False) |
| 40 | |
| 41 | |
| 42 | T = TypeVar("T") |
| 43 | |
| 44 | Type = Literal[ |
| 45 | "agent", |
| 46 | "browser", |
| 47 | "code_exe", |
| 48 | "subagent", |
| 49 | "error", |
| 50 | "hint", |
| 51 | "info", |
| 52 | "progress", |
| 53 | "response", |
| 54 | "tool", |
| 55 | "mcp", |
| 56 | "input", |
| 57 | "user", |
| 58 | "util", |
| 59 | "warning", |
| 60 | ] |
| 61 | |
| 62 | ProgressUpdate = Literal["persistent", "temporary", "none"] |
| 63 | |
| 64 | |
| 65 | HEADING_MAX_LEN: int = 120 |
| 66 | CONTENT_MAX_LEN: int = 15_000 |
| 67 | RESPONSE_CONTENT_MAX_LEN: int = 250_000 |
| 68 | KEY_MAX_LEN: int = 60 |
| 69 | VALUE_MAX_LEN: int = 5000 |
| 70 | PROGRESS_MAX_LEN: int = 120 |
| 71 | |
| 72 | |
| 73 | def _truncate_heading(text: str | None) -> str: |
| 74 | if text is None: |
| 75 | return "" |
| 76 | return truncate_text_by_ratio(str(text), HEADING_MAX_LEN, "...", ratio=1.0) |
| 77 | |
| 78 | |
| 79 | def _truncate_progress(text: str | None) -> str: |
| 80 | if text is None: |
| 81 | return "" |
| 82 | return truncate_text_by_ratio(str(text), PROGRESS_MAX_LEN, "...", ratio=1.0) |
| 83 | |
| 84 | |
| 85 | def _truncate_key(text: str) -> str: |
| 86 | return truncate_text_by_ratio(str(text), KEY_MAX_LEN, "...", ratio=1.0) |
| 87 | |
| 88 | |
| 89 | def _truncate_value(val: T) -> T: |
| 90 | # If dict, recursively truncate each value |
| 91 | if isinstance(val, dict): |
| 92 | for k in list(val.keys()): |
| 93 | v = val[k] |
| 94 | del val[k] |
| 95 | val[_truncate_key(k)] = _truncate_value(v) |
| 96 | return cast(T, val) |
| 97 | # If list or tuple, recursively truncate each item |
| 98 | if isinstance(val, list): |
| 99 | for i in range(len(val)): |
| 100 | val[i] = _truncate_value(val[i]) |
| 101 | return cast(T, val) |
| 102 | if isinstance(val, tuple): |
| 103 | return cast(T, tuple(_truncate_value(x) for x in val)) |
| 104 | |
| 105 | # Convert non-str values to json for consistent length measurement |
| 106 | if isinstance(val, str): |
| 107 | raw = val |
| 108 | else: |
| 109 | try: |
| 110 | raw = json.dumps(val, ensure_ascii=False) |
| 111 | except Exception: |
| 112 | raw = str(val) |
| 113 | |
| 114 | if len(raw) <= VALUE_MAX_LEN: |
| 115 | return val # No truncation needed, preserve original type |
| 116 | |
| 117 | # Do a single truncation calculation |
| 118 | removed = len(raw) - VALUE_MAX_LEN |
| 119 | replacement = f"\n\n<< {removed} Characters hidden >>\n\n" |
| 120 | truncated = truncate_text_by_ratio(raw, VALUE_MAX_LEN, replacement, ratio=0.3) |
| 121 | return cast(T, truncated) |
| 122 | |
| 123 | |
| 124 | def _truncate_content(text: str | None, type: Type) -> str: |
| 125 | |
| 126 | max_len = CONTENT_MAX_LEN if type != "response" else RESPONSE_CONTENT_MAX_LEN |
| 127 | |
| 128 | if text is None: |
| 129 | return "" |
| 130 | raw = str(text) |
| 131 | if len(raw) <= max_len: |
| 132 | return raw |
| 133 | |
| 134 | # Same dynamic replacement logic as value truncation |
| 135 | removed = len(raw) - max_len |
| 136 | while True: |
| 137 | replacement = f"\n\n<< {removed} Characters hidden >>\n\n" |
| 138 | truncated = truncate_text_by_ratio(raw, max_len, replacement, ratio=0.3) |
| 139 | new_removed = len(raw) - (len(truncated) - len(replacement)) |
| 140 | if new_removed == removed: |
| 141 | break |
| 142 | removed = new_removed |
| 143 | return truncated |
| 144 | |
| 145 | |
| 146 | @dataclass |
| 147 | class LogItem: |
| 148 | log: "Log" |
| 149 | no: int |
| 150 | type: Type |
| 151 | heading: str = "" |
| 152 | content: str = "" |
| 153 | update_progress: Optional[ProgressUpdate] = "persistent" |
| 154 | kvps: Optional[OrderedDict] = None # Use OrderedDict for kvps |
| 155 | id: Optional[str] = None # Add id field |
| 156 | guid: str = "" |
| 157 | timestamp: float = 0.0 |
| 158 | agentno: int = 0 |
| 159 | |
| 160 | def __post_init__(self): |
| 161 | self.guid = self.log.guid |
| 162 | self.timestamp = self.timestamp or time.time() |
| 163 | |
| 164 | def update( |
| 165 | self, |
| 166 | type: Type | None = None, |
| 167 | heading: str | None = None, |
| 168 | content: str | None = None, |
| 169 | kvps: dict | None = None, |
| 170 | update_progress: ProgressUpdate | None = None, |
| 171 | **kwargs, |
| 172 | ): |
| 173 | if self.guid == self.log.guid: |
| 174 | self.log._update_item( |
| 175 | self.no, |
| 176 | type=type, |
| 177 | heading=heading, |
| 178 | content=content, |
| 179 | kvps=kvps, |
| 180 | update_progress=update_progress, |
| 181 | **kwargs, |
| 182 | ) |
| 183 | |
| 184 | def stream( |
| 185 | self, |
| 186 | heading: str | None = None, |
| 187 | content: str | None = None, |
| 188 | **kwargs, |
| 189 | ): |
| 190 | if heading is not None: |
| 191 | self.update(heading=self.heading + heading) |
| 192 | if content is not None: |
| 193 | self.update(content=self.content + content) |
| 194 | |
| 195 | for k, v in kwargs.items(): |
| 196 | prev = self.kvps.get(k, "") if self.kvps else "" |
| 197 | self.update(**{k: prev + v}) |
| 198 | |
| 199 | def output(self): |
| 200 | return { |
| 201 | "no": self.no, |
| 202 | "id": self.id, # Include id in output |
| 203 | "type": self.type, |
| 204 | "heading": self.heading, |
| 205 | "content": self.content, |
| 206 | "kvps": self.kvps, |
| 207 | "timestamp": self.timestamp, |
| 208 | "agentno": self.agentno, |
| 209 | } |
| 210 | |
| 211 | |
| 212 | @dataclass(frozen=True) |
| 213 | class LogOutput: |
| 214 | items: list[dict[str, Any]] |
| 215 | start: int |
| 216 | end: int |
| 217 | |
| 218 | |
| 219 | class Log: |
| 220 | |
| 221 | def __init__(self): |
| 222 | self._lock = threading.RLock() |
| 223 | self.context: "AgentContext|None" = None # set from outside |
| 224 | self.guid: str = str(uuid.uuid4()) |
| 225 | self.updates: list[int] = [] |
| 226 | self.logs: list[LogItem] = [] |
| 227 | self.progress: str = "" |
| 228 | self.progress_no: int = 0 |
| 229 | self.progress_active: bool = False |
| 230 | self.set_initial_progress() |
| 231 | |
| 232 | def log( |
| 233 | self, |
| 234 | type: Type, |
| 235 | heading: str | None = None, |
| 236 | content: str | None = None, |
| 237 | kvps: dict | None = None, |
| 238 | update_progress: ProgressUpdate | None = None, |
| 239 | id: Optional[str] = None, |
| 240 | **kwargs, |
| 241 | ) -> LogItem: |
| 242 | with self._lock: |
| 243 | # add a minimal item to the log |
| 244 | # Determine agent number from streaming agent |
| 245 | agentno = 0 |
| 246 | if self.context and self.context.streaming_agent: |
| 247 | agentno = self.context.streaming_agent.number |
| 248 | |
| 249 | item = LogItem( |
| 250 | log=self, |
| 251 | no=len(self.logs), |
| 252 | type=type, |
| 253 | agentno=agentno, |
| 254 | ) |
| 255 | |
| 256 | self.logs.append(item) |
| 257 | |
| 258 | # Update outside the lock - the heavy masking/truncation work should not hold |
| 259 | # the lock; we only need locking while mutating shared arrays/fields. |
| 260 | self._update_item( |
| 261 | no=item.no, |
| 262 | type=type, |
| 263 | heading=heading, |
| 264 | content=content, |
| 265 | kvps=kvps, |
| 266 | update_progress=update_progress, |
| 267 | id=id, |
| 268 | notify_state_monitor=False, |
| 269 | **kwargs, |
| 270 | ) |
| 271 | |
| 272 | self._notify_state_monitor() |
| 273 | return item |
| 274 | |
| 275 | def _update_item( |
| 276 | self, |
| 277 | no: int, |
| 278 | type: Type | None = None, |
| 279 | heading: str | None = None, |
| 280 | content: str | None = None, |
| 281 | kvps: dict | None = None, |
| 282 | update_progress: ProgressUpdate | None = None, |
| 283 | id: Optional[str] = None, |
| 284 | notify_state_monitor: bool = True, |
| 285 | **kwargs, |
| 286 | ): |
| 287 | # Capture the effective type for truncation without holding the lock during |
| 288 | # masking/truncation work. |
| 289 | with self._lock: |
| 290 | current_type = self.logs[no].type |
| 291 | type_for_truncation = type if type is not None else current_type |
| 292 | |
| 293 | heading_out: str | None = None |
| 294 | if heading is not None: |
| 295 | heading_out = _truncate_heading(self._mask_recursive(heading)) |
| 296 | |
| 297 | content_out: str | None = None |
| 298 | if content is not None: |
| 299 | content_out = _truncate_content(self._mask_recursive(content), type_for_truncation) |
| 300 | |
| 301 | kvps_out: OrderedDict | None = None |
| 302 | if kvps is not None: |
| 303 | kvps_out_tmp = OrderedDict(copy.deepcopy(kvps)) |
| 304 | kvps_out_tmp = self._mask_recursive(kvps_out_tmp) |
| 305 | kvps_out_tmp = _truncate_value(kvps_out_tmp) |
| 306 | kvps_out = OrderedDict(kvps_out_tmp) |
| 307 | |
| 308 | kwargs_out: dict | None = None |
| 309 | if kwargs: |
| 310 | kwargs_out = copy.deepcopy(kwargs) |
| 311 | kwargs_out = self._mask_recursive(kwargs_out) |
| 312 | kwargs_out = _truncate_value(kwargs_out) |
| 313 | |
| 314 | with self._lock: |
| 315 | item = self.logs[no] |
| 316 | |
| 317 | if id is not None: |
| 318 | item.id = id |
| 319 | |
| 320 | if type is not None: |
| 321 | item.type = type |
| 322 | |
| 323 | if update_progress is not None: |
| 324 | item.update_progress = update_progress |
| 325 | |
| 326 | if heading_out is not None: |
| 327 | item.heading = heading_out |
| 328 | |
| 329 | if content_out is not None: |
| 330 | item.content = content_out |
| 331 | |
| 332 | if kvps_out is not None: |
| 333 | item.kvps = kvps_out |
| 334 | elif item.kvps is None: |
| 335 | item.kvps = OrderedDict() |
| 336 | |
| 337 | if kwargs_out: |
| 338 | if item.kvps is None: |
| 339 | item.kvps = OrderedDict() |
| 340 | item.kvps.update(kwargs_out) |
| 341 | |
| 342 | self.updates.append(item.no) |
| 343 | |
| 344 | if item.heading and item.update_progress != "none": |
| 345 | if item.no >= self.progress_no: |
| 346 | self.progress = item.heading |
| 347 | self.progress_no = ( |
| 348 | item.no if item.update_progress == "persistent" else -1 |
| 349 | ) |
| 350 | self.progress_active = True |
| 351 | if notify_state_monitor: |
| 352 | self._notify_state_monitor_for_context_update() |
| 353 | |
| 354 | def _notify_state_monitor(self) -> None: |
| 355 | ctx = self.context |
| 356 | if not ctx: |
| 357 | return |
| 358 | # Logs update both the active chat stream (sid-bound) and the global chats list |
| 359 | # (context metadata like last_message/log_version). Broadcast so all tabs refresh |
| 360 | # their chat/task lists without leaking logs (logs are still scoped per-sid). |
| 361 | _lazy_mark_dirty_all(reason="log.Log._notify_state_monitor") |
| 362 | |
| 363 | def _notify_state_monitor_for_context_update(self) -> None: |
| 364 | ctx = self.context |
| 365 | if not ctx: |
| 366 | return |
| 367 | # Log item updates only need to refresh the active chat stream for any sid |
| 368 | # currently projecting this context. Avoid global fanout at high frequency. |
| 369 | _lazy_mark_dirty_for_context(ctx.id, reason="log.Log._update_item") |
| 370 | |
| 371 | def set_progress(self, progress: str, no: int = 0, active: bool = True): |
| 372 | progress = self._mask_recursive(progress) |
| 373 | progress = _truncate_progress(progress) |
| 374 | changed = False |
| 375 | ctx = self.context |
| 376 | with self._lock: |
| 377 | prev_progress = self.progress |
| 378 | prev_active = self.progress_active |
| 379 | |
| 380 | self.progress = progress |
| 381 | if not no: |
| 382 | no = len(self.logs) |
| 383 | self.progress_no = no |
| 384 | self.progress_active = active |
| 385 | |
| 386 | changed = self.progress != prev_progress or self.progress_active != prev_active |
| 387 | |
| 388 | if changed and ctx: |
| 389 | # Progress changes are included in every snapshot, but push sync requires a |
| 390 | # dirty mark even when no log items changed. |
| 391 | _lazy_mark_dirty_for_context(ctx.id, reason="log.Log.set_progress") |
| 392 | |
| 393 | def set_initial_progress(self): |
| 394 | self.set_progress("Waiting for input", 0, False) |
| 395 | |
| 396 | def output(self, start=None, end=None): |
| 397 | with self._lock: |
| 398 | if start is None: |
| 399 | start = 0 |
| 400 | if end is None: |
| 401 | end = len(self.updates) |
| 402 | updates = self.updates[start:end] |
| 403 | logs = list(self.logs) |
| 404 | |
| 405 | out = [] |
| 406 | seen = set() |
| 407 | for update in updates: |
| 408 | if update not in seen and update < len(logs): |
| 409 | out.append(logs[update].output()) |
| 410 | seen.add(update) |
| 411 | return LogOutput(items=out, start=start, end=end) |
| 412 | |
| 413 | def reset(self): |
| 414 | with self._lock: |
| 415 | self.guid = str(uuid.uuid4()) |
| 416 | self.updates = [] |
| 417 | self.logs = [] |
| 418 | self.set_initial_progress() |
| 419 | |
| 420 | def _mask_recursive(self, obj: T) -> T: |
| 421 | """Recursively mask secrets in nested objects.""" |
| 422 | try: |
| 423 | from agent import AgentContext |
| 424 | secrets_mgr = get_secrets_manager(self.context or AgentContext.current()) |
| 425 | |
| 426 | # debug helper to identify context mismatch |
| 427 | # self_id = self.context.id if self.context else None |
| 428 | # current_ctx = AgentContext.current() |
| 429 | # current_id = current_ctx.id if current_ctx else None |
| 430 | # if self_id != current_id: |
| 431 | # print(f"Context ID mismatch: {self_id} != {current_id}") |
| 432 | |
| 433 | if isinstance(obj, str): |
| 434 | return cast(Any, secrets_mgr.mask_values(obj)) |
| 435 | elif isinstance(obj, dict): |
| 436 | return {k: self._mask_recursive(v) for k, v in obj.items()} # type: ignore |
| 437 | elif isinstance(obj, list): |
| 438 | return [self._mask_recursive(item) for item in obj] # type: ignore |
| 439 | else: |
| 440 | return obj |
| 441 | except Exception: |
| 442 | # If masking fails, return original object |
| 443 | return obj |