| 1 | import re |
| 2 | import json |
| 3 | import time |
| 4 | import asyncio |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | from helpers import plugins |
| 8 | from helpers import history as history_helpers |
| 9 | from helpers.errors import HandledException |
| 10 | from langchain_core.messages import AIMessage, HumanMessage, SystemMessage |
| 11 | |
| 12 | if TYPE_CHECKING: |
| 13 | from agent import Agent |
| 14 | from helpers.log import LogItem |
| 15 | |
| 16 | PLUGIN_NAME = "_infection_check" |
| 17 | DATA_KEY = f"_plugin.{PLUGIN_NAME}" |
| 18 | DATA_KEY_PASSED = f"{DATA_KEY}.passed" |
| 19 | DATA_KEY_MONO = f"{DATA_KEY}.mono" |
| 20 | |
| 21 | _RE_OK = re.compile(r"<ok\s*/>") |
| 22 | _RE_TERMINATE = re.compile(r"<terminate\s*/>") |
| 23 | _RE_CLARIFY = re.compile(r"<clarify>(.*?)</clarify>", re.DOTALL) |
| 24 | |
| 25 | |
| 26 | def get_config(agent: "Agent") -> dict: |
| 27 | return plugins.get_plugin_config(PLUGIN_NAME, agent=agent) or {} |
| 28 | |
| 29 | |
| 30 | def get_checker(agent: "Agent") -> "InfectionChecker": |
| 31 | """Get or create a checker for the current iteration. |
| 32 | |
| 33 | A new InfectionChecker is created when: |
| 34 | - The monologue changes (detected via id(loop_data)). |
| 35 | - The iteration within the same monologue changes. |
| 36 | """ |
| 37 | loop = getattr(agent, "loop_data", None) |
| 38 | iteration = loop.iteration if loop else -1 |
| 39 | mono_id = id(loop) if loop else None |
| 40 | |
| 41 | # Reset passed flag on new monologue |
| 42 | if agent.get_data(DATA_KEY_MONO) != mono_id: |
| 43 | agent.set_data(DATA_KEY_MONO, mono_id) |
| 44 | agent.set_data(DATA_KEY_PASSED, False) |
| 45 | agent.set_data(DATA_KEY, None) # discard stale checker from previous monologue |
| 46 | |
| 47 | checker: "InfectionChecker | None" = agent.get_data(DATA_KEY) |
| 48 | if checker is None or checker.iteration != iteration: |
| 49 | checker = InfectionChecker(config=get_config(agent), iteration=iteration) |
| 50 | agent.set_data(DATA_KEY, checker) |
| 51 | agent.set_data(DATA_KEY_PASSED, False) # re-check each iteration |
| 52 | return checker |
| 53 | |
| 54 | |
| 55 | def parse_result(text: str) -> tuple[str, str]: |
| 56 | """Find the *last* occurrence of any verdict tag in *text*.""" |
| 57 | pos = -1 |
| 58 | action, detail = "ok", "" |
| 59 | for m in _RE_OK.finditer(text): |
| 60 | if m.start() > pos: |
| 61 | pos, action, detail = m.start(), "ok", "" |
| 62 | for m in _RE_TERMINATE.finditer(text): |
| 63 | if m.start() > pos: |
| 64 | pos, action, detail = m.start(), "terminate", "" |
| 65 | for m in _RE_CLARIFY.finditer(text): |
| 66 | if m.start() > pos: |
| 67 | pos, action, detail = m.start(), "clarify", m.group(1).strip() |
| 68 | return action, detail |
| 69 | |
| 70 | |
| 71 | class InfectionChecker: |
| 72 | |
| 73 | def __init__(self, config: dict, iteration: int): |
| 74 | self.mode: str = config.get("mode", "thoughts") |
| 75 | self.model_choice: str = config.get("model", "utility") |
| 76 | self.prompt: str = config.get("prompt", "") |
| 77 | self.history_size: int = int(config.get("history_size", 10)) |
| 78 | self.max_clarifications: int = int(config.get("max_clarifications", 3)) |
| 79 | self.iteration = iteration |
| 80 | |
| 81 | # Accumulated text from stream callbacks |
| 82 | self.reasoning_log = "" |
| 83 | self.response_log = "" |
| 84 | |
| 85 | # Background analysis task |
| 86 | self._task: asyncio.Task | None = None |
| 87 | self._check_msgs: list = [] |
| 88 | |
| 89 | # -- collection ---------------------------------------------------------- |
| 90 | |
| 91 | def collect_reasoning(self, full_text: str): |
| 92 | self.reasoning_log = full_text |
| 93 | |
| 94 | def collect_response(self, full_text: str): |
| 95 | # Stop collecting once background analysis has started |
| 96 | if self._task is None: |
| 97 | self.response_log = full_text |
| 98 | |
| 99 | # -- analysis trigger ---------------------------------------------------- |
| 100 | |
| 101 | def start_analysis(self, agent: "Agent"): |
| 102 | """Fire-and-forget background check (called from stream extensions).""" |
| 103 | if self._task is not None: |
| 104 | return |
| 105 | snapshot = self._build_log() |
| 106 | if not snapshot.strip(): |
| 107 | return |
| 108 | self._task = asyncio.create_task(self._run_check(agent, snapshot)) |
| 109 | |
| 110 | # -- gate (called before every tool execution) --------------------------- |
| 111 | |
| 112 | async def gate(self, agent: "Agent", tool_name: str = "", tool_args: dict | None = None): |
| 113 | """Block until the safety check passes or terminate the agent.""" |
| 114 | try: |
| 115 | await self._gate_inner(agent, tool_name, tool_args) |
| 116 | except HandledException: |
| 117 | raise |
| 118 | except Exception as e: |
| 119 | from helpers.print_style import PrintStyle |
| 120 | PrintStyle(font_color="red", padding=True).print( |
| 121 | f"Infection check error (non-fatal): {e}" |
| 122 | ) |
| 123 | try: |
| 124 | agent.context.log.set_progress("Infection check: error (non-fatal)") |
| 125 | except Exception: |
| 126 | pass |
| 127 | |
| 128 | async def _gate_inner(self, agent: "Agent", tool_name: str, tool_args: dict | None): |
| 129 | if agent.get_data(DATA_KEY_PASSED): |
| 130 | return |
| 131 | if not self.reasoning_log and not self.response_log: |
| 132 | return |
| 133 | |
| 134 | _log = agent.context.log |
| 135 | _log.set_progress("Infection check: analyzing...") |
| 136 | |
| 137 | # Attach tool context for _build_log() |
| 138 | if tool_name: |
| 139 | self._tool_name = tool_name |
| 140 | try: |
| 141 | self._tool_args = dict(tool_args) if tool_args else {} |
| 142 | except Exception: |
| 143 | self._tool_args = {} |
| 144 | else: |
| 145 | self._tool_name = "" |
| 146 | self._tool_args = {} |
| 147 | |
| 148 | action, detail, cot = None, "", "" |
| 149 | |
| 150 | # Fast path: reuse result if background task already finished. |
| 151 | if self._task is not None and self._task.done(): |
| 152 | try: |
| 153 | action, detail, cot = self._task.result() |
| 154 | _log.set_progress("Infection check: evaluating result...") |
| 155 | except Exception: |
| 156 | pass |
| 157 | |
| 158 | # Slow path: rebuild with full tool context. |
| 159 | if action is None: |
| 160 | _log.set_progress("Infection check: analyzing with tool context...") |
| 161 | if self._task is not None: |
| 162 | self._task.cancel() |
| 163 | try: |
| 164 | await self._task |
| 165 | except (asyncio.CancelledError, Exception): |
| 166 | pass |
| 167 | self._task = None |
| 168 | |
| 169 | snapshot = self._build_log() |
| 170 | if not snapshot.strip(): |
| 171 | return |
| 172 | |
| 173 | self._task = asyncio.create_task(self._run_check(agent, snapshot)) |
| 174 | try: |
| 175 | action, detail, cot = await self._task |
| 176 | except asyncio.CancelledError: |
| 177 | return |
| 178 | except Exception: |
| 179 | return |
| 180 | |
| 181 | if action == "ok": |
| 182 | _log.set_progress("Infection check: passed") |
| 183 | agent.set_data(DATA_KEY_PASSED, True) |
| 184 | return |
| 185 | |
| 186 | if action == "clarify": |
| 187 | warn = agent.context.log.log( |
| 188 | type="warning", |
| 189 | heading="Infection check: requesting clarification", |
| 190 | content=f"Safety concern:\n{cot}" if cot else "", |
| 191 | ) |
| 192 | action, detail, cot = await self._clarify_loop(agent, detail, warn) |
| 193 | if action == "ok": |
| 194 | warn.update(heading="Infection check: clarification passed") |
| 195 | agent.set_data(DATA_KEY_PASSED, True) |
| 196 | return |
| 197 | |
| 198 | # terminate |
| 199 | self._do_terminate(agent, detail, cot) |
| 200 | |
| 201 | # -- internals ----------------------------------------------------------- |
| 202 | |
| 203 | def _build_log(self) -> str: |
| 204 | parts: list[str] = [] |
| 205 | if self.reasoning_log: |
| 206 | parts.append(f"## Agent Reasoning\n{self.reasoning_log}") |
| 207 | if self.response_log: |
| 208 | parts.append(f"## Agent Response\n{self.response_log}") |
| 209 | if getattr(self, "_tool_name", ""): |
| 210 | try: |
| 211 | args_str = json.dumps(self._tool_args, ensure_ascii=False, default=str, indent=2) |
| 212 | except Exception: |
| 213 | args_str = str(getattr(self, "_tool_args", {})) |
| 214 | parts.append(f"## Tool About to Execute\nTool: {self._tool_name}\nArguments:\n{args_str}") |
| 215 | return "\n\n".join(parts) |
| 216 | |
| 217 | def _get_model(self, agent: "Agent"): |
| 218 | if self.model_choice == "main": |
| 219 | return agent.get_chat_model() |
| 220 | return agent.get_utility_model() |
| 221 | |
| 222 | async def _run_check(self, agent: "Agent", log_text: str) -> tuple[str, str, str]: |
| 223 | # Build context from recent history |
| 224 | hist = agent.history.output() |
| 225 | if self.history_size > 0: |
| 226 | hist = hist[-self.history_size :] |
| 227 | |
| 228 | # Filter out previously blocked entries |
| 229 | filtered: list = [] |
| 230 | for entry in hist: |
| 231 | content = str(entry.get("content", "")) if isinstance(entry, dict) else "" |
| 232 | if "[BLOCKED]" in content: |
| 233 | if filtered: |
| 234 | filtered.pop() # also remove the user message before it |
| 235 | continue |
| 236 | filtered.append(entry) |
| 237 | |
| 238 | hist_text = history_helpers.output_text(filtered, ai_label="assistant", human_label="user") |
| 239 | user_msg = ( |
| 240 | f"## Recent Conversation History\n{hist_text}\n\n" |
| 241 | f"## Current Agent Output to Analyze\n{log_text}" |
| 242 | ) |
| 243 | self._check_msgs = [ |
| 244 | SystemMessage(content=self.prompt), |
| 245 | HumanMessage(content=user_msg), |
| 246 | ] |
| 247 | |
| 248 | model = self._get_model(agent) |
| 249 | response, _ = await model.unified_call( |
| 250 | messages=list(self._check_msgs), |
| 251 | ) |
| 252 | self._check_msgs.append(AIMessage(content=response)) |
| 253 | |
| 254 | action, detail = parse_result(response) |
| 255 | return action, detail, response |
| 256 | |
| 257 | async def _clarify_loop( |
| 258 | self, agent: "Agent", clarify_text: str, log_item: "LogItem" |
| 259 | ) -> tuple[str, str, str]: |
| 260 | cot_parts: list[str] = [] |
| 261 | |
| 262 | # Throttled log display — avoids per-token stream() and O(n²) masking. |
| 263 | _buf = "" |
| 264 | _last_flush = 0.0 |
| 265 | _INTERVAL = 0.25 # seconds between UI pushes |
| 266 | |
| 267 | def _flush(): |
| 268 | nonlocal _last_flush |
| 269 | log_item.update(content=_buf) |
| 270 | _last_flush = time.monotonic() |
| 271 | |
| 272 | def _append(text: str): |
| 273 | nonlocal _buf |
| 274 | _buf += text |
| 275 | if time.monotonic() - _last_flush >= _INTERVAL: |
| 276 | _flush() |
| 277 | |
| 278 | for i in range(self.max_clarifications): |
| 279 | log_item.update( |
| 280 | heading=f"Infection check: clarification {i + 1}/{self.max_clarifications}", |
| 281 | ) |
| 282 | _buf = f"Safety model question:\n{clarify_text}\n\nAgent response:\n" |
| 283 | _flush() |
| 284 | |
| 285 | # Clone conversation and ask the agent to explain |
| 286 | chat_msgs = agent.history.output_langchain() |
| 287 | chat_msgs.append(HumanMessage(content=clarify_text)) |
| 288 | |
| 289 | async def _agent_cb(chunk: str, full: str): |
| 290 | if chunk: |
| 291 | _append(chunk) |
| 292 | |
| 293 | agent_resp, _ = await agent.get_chat_model().unified_call( |
| 294 | messages=chat_msgs, |
| 295 | response_callback=_agent_cb, |
| 296 | ) |
| 297 | cot_parts.append(f"Q: {clarify_text}\nA: {agent_resp}") |
| 298 | _append("\n\nSafety model verdict:\n") |
| 299 | _flush() |
| 300 | |
| 301 | # Feed agent's response back to the check model |
| 302 | self._check_msgs.append( |
| 303 | HumanMessage( |
| 304 | content=( |
| 305 | f"The agent responded:\n\n{agent_resp}\n\n" |
| 306 | "Re-evaluate and provide your verdict." |
| 307 | ) |
| 308 | ) |
| 309 | ) |
| 310 | |
| 311 | async def _check_cb(chunk: str, full: str): |
| 312 | if chunk: |
| 313 | _append(chunk) |
| 314 | |
| 315 | check_resp, _ = await self._get_model(agent).unified_call( |
| 316 | messages=list(self._check_msgs), |
| 317 | response_callback=_check_cb, |
| 318 | ) |
| 319 | self._check_msgs.append(AIMessage(content=check_resp)) |
| 320 | cot_parts.append(f"Safety: {check_resp}") |
| 321 | _flush() |
| 322 | |
| 323 | action, detail = parse_result(check_resp) |
| 324 | if action != "clarify": |
| 325 | return action, detail, "\n\n".join(cot_parts) |
| 326 | clarify_text = detail |
| 327 | |
| 328 | return "terminate", "Max clarifications exceeded.", "\n\n".join(cot_parts) |
| 329 | |
| 330 | def _do_terminate(self, agent: "Agent", detail: str, cot: str): |
| 331 | import uuid as _uuid |
| 332 | content = cot or detail or "Malicious behavior detected." |
| 333 | msg_id = str(_uuid.uuid4()) |
| 334 | agent.context.log.log( |
| 335 | type="warning", |
| 336 | heading="Infection check: TERMINATED", |
| 337 | content=content, |
| 338 | id=msg_id, |
| 339 | ) |
| 340 | |
| 341 | # Replace last AI message with a blocked marker |
| 342 | try: |
| 343 | msgs = agent.history.current.messages |
| 344 | if msgs and msgs[-1].ai: |
| 345 | msgs.pop() |
| 346 | agent.history.add_message( |
| 347 | ai=True, content="[BLOCKED] Response terminated by security policy.", |
| 348 | id=msg_id, |
| 349 | ) |
| 350 | except Exception: |
| 351 | pass |
| 352 | |
| 353 | # Desktop notification |
| 354 | from helpers.notification import ( |
| 355 | NotificationManager, |
| 356 | NotificationType, |
| 357 | NotificationPriority, |
| 358 | ) |
| 359 | |
| 360 | NotificationManager.send_notification( |
| 361 | type=NotificationType.ERROR, |
| 362 | priority=NotificationPriority.HIGH, |
| 363 | title="Infection Check", |
| 364 | message="Threat detected — agent execution terminated.", |
| 365 | detail=detail or "Malicious behavior detected.", |
| 366 | display_time=8, |
| 367 | ) |
| 368 | |
| 369 | # process_chain_end won't fire after HandledException, |
| 370 | # so schedule queue resumption before raising. |
| 371 | self._schedule_queue_resume(agent) |
| 372 | |
| 373 | raise HandledException( |
| 374 | Exception("Infection check terminated: " + (detail or "threat detected")) |
| 375 | ) |
| 376 | |
| 377 | @staticmethod |
| 378 | def _schedule_queue_resume(agent: "Agent"): |
| 379 | from helpers import message_queue as mq |
| 380 | |
| 381 | if agent.number != 0: |
| 382 | return |
| 383 | context = agent.context |
| 384 | if not mq.has_queue(context): |
| 385 | return |
| 386 | |
| 387 | async def _resume(): |
| 388 | total_wait = 0.0 |
| 389 | while context.is_running() and total_wait < 60: |
| 390 | await asyncio.sleep(0.1) |
| 391 | total_wait += 0.1 |
| 392 | if not context.is_running(): |
| 393 | mq.send_next(context) |
| 394 | |
| 395 | asyncio.create_task(_resume()) |