diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 63c010d2c..110e935f0 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -70,6 +70,53 @@ def _sanitized_env() -> dict: and "KEY" not in k and "SECRET" not in k and "TOKEN" not in k and "PASSWORD" not in k} +def _build_hooks_from_context(context: Dict) -> Dict: + """Build the runner hooks dict from ``context['drivers']``. + + Sub-runners dispatched by the ai task are constructed in-process and run + synchronously, so the framework's pickle path (``__setstate__``, which + re-registers driver hooks from ``context['drivers']``) never runs for them. + Without this, a sub-runner inherits the ai task's ``workspace_id`` / + ``drivers`` in its context but registers *no* driver hooks — so its + ``mongodb``/``api`` ``update_runner``/``update_finding`` hooks never fire and + its runner doc + findings are never persisted to the workspace. The result: + sub-runs are absent from the workspace History. + + This mirrors the normal CLI entrypoint (``cli_helper._run``): import each + driver's ``secator.hooks..HOOKS`` and ``deep_merge_dicts`` them into a + single class-keyed dict (keyed by ``Scan``/``Workflow``/``Task``). The dict is + returned raw (not flattened) because ``Task``/``Workflow`` forward + ``self._hooks.get(Task, {})`` down to their command/task signatures. + + Args: + context: Runner context dict (expects ``drivers`` list). + + Returns: + dict: Merged hooks dict suitable for ``runner_cls(..., hooks=hooks)``. + """ + from secator.loader import discover_external_drivers, get_available_drivers, order_drivers + from secator.utils import import_dynamic, deep_merge_dicts + + drivers = list(context.get('drivers', [])) + if not drivers: + return {} + discover_external_drivers() + # Order by canonical priority so authoritative backends (e.g. mongodb) register + # their hooks before relay drivers (e.g. api) — same ordering as __setstate__. + drivers = order_drivers(drivers) + supported = set(get_available_drivers()) + hooks_list = [] + for driver in drivers: + if driver not in supported: + continue + driver_hooks = import_dynamic(f'secator.hooks.{driver}', 'HOOKS') + if driver_hooks: + hooks_list.append(driver_hooks) + if not hooks_list: + return {} + return deep_merge_dicts(*hooks_list) + + def _build_action_display(action: Dict) -> str: """Build a display string for the action being checked. @@ -259,6 +306,60 @@ def dispatch_action(action: Dict, ctx: ActionContext) -> Generator: yield Warning(message=f"Unknown action: {action_type}", _context=context) +def _format_action_error(e: Exception, max_chars: int = 400) -> str: + """Build a concise, LLM-facing error string for a failed action dispatch. + + Combines the exception type + message with the last few traceback frames so + the model can see *where* it failed, then truncates to a sane length so a + deep traceback can't blow up the next prompt's token budget. + """ + import traceback + + errtype = type(e).__name__ + msg = str(e) + head = f"{errtype}: {msg}" if msg else errtype + + # Keep only the tail of the traceback (last ~3 frames) — that's where the + # actual failure is, and it keeps the feedback compact. + tb_lines = traceback.format_exc().strip().splitlines() + tb_tail = "\n".join(tb_lines[-6:]) if tb_lines else "" + + detail = f"{head}\n{tb_tail}" if tb_tail else head + if len(detail) > max_chars: + detail = detail[:max_chars] + "…(truncated)" + return ( + f"Action failed with error: {detail}\n" + "Fix the issue and try again." + ) + + +def safe_dispatch_action(action: Dict, ctx: ActionContext) -> Generator: + """Dispatch a single action, converting any raised ``Exception`` into an + ``Error`` output item instead of letting it abort the AI loop. + + A Python error during a handler (e.g. ``TypeError: 'str' object is not a + mapping`` from a malformed LLM action/opts) must NOT kill the main loop. We + wrap the per-action generator so the failure becomes an ``Error`` carrying + the action's ``tool_call_id``/``tool_call_name`` in ``_context`` — that lets + the caller group it into a tool result and feed the error back to the LLM so + it can correct itself on the next turn. + + Only ``Exception`` is caught: ``KeyboardInterrupt`` / ``SystemExit`` / + ``GeneratorExit`` (all ``BaseException`` subclasses) propagate so legitimate + control-flow and generator close are never swallowed. + """ + import traceback as _traceback + try: + yield from dispatch_action(action, ctx) + except Exception as e: # noqa: BLE001 - per-action resilience: feed error back to LLM, never abort the loop + context = _get_result_context(action, ctx) + yield Error( + message=_format_action_error(e), + traceback=_traceback.format_exc(), + _context=context, + ) + + def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: """Execute a secator task or workflow. @@ -292,9 +393,6 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator yield Info(message=f"[DRY RUN] Would run {runner_type}: {name} on {targets}", _context=context) return - if not ctx.silent: - yield Ai(content=name, ai_type=runner_type, extra_data={"targets": targets, "opts": opts}, _context=context) - run_opts = { "print_item": not ctx.silent, "print_line": ctx.verbose and not ctx.silent, @@ -315,11 +413,42 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator context["task_chunk_id"] = str(uuid.uuid4()) if ctx.subagent: context["subagent"] = ctx.context.get("subagent", True) + + # Propagate the ai task's driver hooks (mongodb/api) into the sub-runner. + # The context already carries workspace_id/workspace_name/drivers (see + # _get_result_context), but a sync sub-runner never goes through the pickle + # path that re-registers driver hooks — so without this its results would + # persist with no workspace scope and never appear in the workspace History. + hooks = _build_hooks_from_context(context) try: - runner = runner_cls(tpl, targets, run_opts=run_opts, context=context) + runner = runner_cls(tpl, targets, run_opts=run_opts, hooks=hooks, context=context) except TaskNotFoundError as e: yield Error(message=str(e), _context=context) return + + # Emit the action Ai item now that the runner exists: its on_init hook has + # stamped the runner id into context, so we can surface it on the item + # (extra_data.runner_id/runner_type) for the UI to link to a RunnerCard. + # Emit even when silent (batch mode): silent only suppresses live console + # chatter, but the action doc must still be yielded so it is persisted and + # the UI can render a RunnerCard for it. + # Prefer the context id (`{type}_id`) the on_init hook stamped — that IS the + # persisted runner doc's `_id`, which is what the UI's getRunner queries. + # `runner.id` is secator's internal id and does NOT match the persisted doc, + # so the RunnerCard showed "Runner not found". + runner_id = context.get(f"{runner_type}_id", "") or runner.id + yield Ai( + content=name, + ai_type=runner_type, + extra_data={ + "targets": targets, + "opts": opts, + "runner_id": runner_id, + "runner_type": runner_type, + }, + _context=context, + ) + yield from runner # Auto-allow reading from the spawned runner's reports folder @@ -329,15 +458,26 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator def _get_result_context(action, ctx): - """Get result context from action""" - ctx = ctx.context.copy() + """Get result context from action. + + Always stamps the ai task's ``session_id`` (the conversation id) onto the + derived context. The ai task's ``self.session_id`` may be derived (from + ``session_name`` / the runner id) and is therefore not guaranteed to already + live in ``ctx.context``. Stamping it here means every sub-runner (task / + workflow / scan) dispatched by the ai task persists a runner doc whose + ``context.session_id`` matches the conversation — so the runners spawned by a + conversation are queryable by that conversation's session_id. + """ + new_ctx = ctx.context.copy() + if ctx.session_id and not new_ctx.get("session_id"): + new_ctx["session_id"] = ctx.session_id action_context = {} tool_call_id = action.get("tool_call_id") tool_call_name = action.get("tool_call_name") if tool_call_id: action_context["tool_call_id"] = tool_call_id action_context["tool_call_name"] = tool_call_name - return {**ctx, **action_context} + return {**new_ctx, **action_context} def _handle_task(action: Dict, ctx: ActionContext) -> Generator: @@ -444,7 +584,10 @@ def _handle_follow_up(action: Dict, ctx: ActionContext) -> Generator: context = _get_result_context(action, ctx) reason = action.get("reason", "completed") choices = action.get("choices", []) - yield Ai(content=reason, ai_type="follow_up", extra_data={"choices": choices}, _context=context) + # Store choices on the top-level `choices` field (what the web UI reads) AND in + # extra_data (back-compat). Without the top-level field, the persisted follow-up + # doc has `choices: []` and the UI renders no choice buttons. + yield Ai(content=reason, ai_type="follow_up", choices=choices, extra_data={"choices": choices}, _context=context) def _handle_stop(action: Dict, ctx: ActionContext) -> Generator: @@ -454,6 +597,106 @@ def _handle_stop(action: Dict, ctx: ActionContext) -> Generator: yield Ai(content=reason, ai_type="stopped", _context=context) +def _resolve_field_type(f) -> Optional[type]: + """Resolve a dataclass field's declared type to a concrete builtin type. + + Mirrors ``OutputType.validate_fields``: ``f.type`` may be an actual type + (``bool``) or — under ``from __future__ import annotations`` — a string + annotation (``'bool'``). Returns the concrete type (``bool``/``int``/ + ``float``/``list``/``dict``/``str``) or ``None`` if it can't be resolved. + """ + t = f.type + # Actual type, e.g. bool / int / float / str + if isinstance(t, type): + return t + # Typing generic, e.g. List[str] -> list + origin = getattr(t, '__origin__', None) + if origin is not None: + return origin + # String annotation, e.g. 'bool', 'int', "List[str]" + if isinstance(t, str): + name = t.split('[', 1)[0].strip().lower() + return { + 'bool': bool, 'int': int, 'float': float, + 'str': str, 'list': list, 'dict': dict, + }.get(name) + return None + + +def _coerce_finding_fields(cls, data: Dict) -> Dict: + """Coerce AI-provided scalar values to a finding class's declared field types. + + LLMs frequently emit wrong-typed scalars (a ``bool`` field as the string + ``"true"``, an ``int`` as ``"3"``). This fixes *obvious* type mismatches + before validation so the finding isn't rejected for model type sloppiness. + + Only coerces when safe; unknown keys, already-correct values, and + unparseable values are left untouched (validation will still surface a real + error rather than silently dropping data). + """ + field_types = {f.name: _resolve_field_type(f) for f in fields(cls)} + for key, value in list(data.items()): + if key.startswith('_'): + continue + expected = field_types.get(key) + if expected is None or value is None: + continue + # Already the right type (note: bool is a subclass of int, so guard it). + if isinstance(value, expected) and not (expected is int and isinstance(value, bool)): + continue + + if expected is bool: + if isinstance(value, bool): + continue + if isinstance(value, int): + data[key] = bool(value) + elif isinstance(value, str): + s = value.strip().lower() + if s in ('true', '1', 'yes', 'on'): + data[key] = True + elif s in ('false', '0', 'no', 'off', ''): + data[key] = False + elif expected is int: + # Avoid coercing real bools into ints. + if isinstance(value, bool): + continue + if isinstance(value, float): + if value.is_integer(): + data[key] = int(value) + elif isinstance(value, str): + try: + data[key] = int(value) + except ValueError: + try: + f_val = float(value) + if f_val.is_integer(): + data[key] = int(f_val) + except ValueError: + pass + elif expected is float: + if isinstance(value, bool): + continue + if isinstance(value, int): + data[key] = float(value) + elif isinstance(value, str): + try: + data[key] = float(value) + except ValueError: + pass + elif expected is list: + if isinstance(value, str): + s = value.strip() + if s.startswith('['): + try: + parsed = json.loads(s) + if isinstance(parsed, list): + data[key] = parsed + except (json.JSONDecodeError, TypeError): + pass + # str fields: leave as-is (don't stringify); unknown types: leave untouched. + return data + + def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator: """Create a secator finding from LLM-provided data. @@ -507,6 +750,10 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator: extra.update(unknown) finding_data['extra_data'] = extra + # Coerce AI-provided scalars to declared field types (LLMs send wrong-typed + # scalars, e.g. a bool field as the string "true") before validating. + finding_data = _coerce_finding_fields(cls, finding_data) + # Validate field types before instantiation errors = cls.validate_fields(finding_data) if errors: @@ -519,6 +766,9 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator: yield Ai( content=f'{str(finding)}', ai_type="add_finding", + # Carry the created finding so the web UI can render its FindingCard + # (VulnerabilityCard/SubdomainCard/…) — it routes on `_type`. + extra_data={"finding": finding.toDict()}, _context=context ) yield finding @@ -593,8 +843,12 @@ def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator: progress_ids = {} def run_single(act: Dict, idx: int) -> Dict: + # Use safe_dispatch_action so one action raising doesn't abort the whole + # batch (the executor future.result() would otherwise re-raise into the + # main loop). The error is captured as an Error item attributed to that + # action's tool_call_id and fed back to the LLM like any other result. results = [] - for item in dispatch_action(act, batch_ctx): + for item in safe_dispatch_action(act, batch_ctx): if isinstance(item, Ai) and item.ai_type == "token_usage": if progress: extra = item.extra_data or {} diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 7744ae2e1..78e314145 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -119,7 +119,7 @@ def build_pending_prompt(self, question, choices, session_id, prompt_type="follo ) def ask_user(self, question, choices, session_id, prompt_type="follow_up", **context): - answer = self._poll_for_answer(session_id, prompt_type) + answer = self._poll_for_answer(session_id, prompt_type, prompt_uuid=context.get("prompt_uuid")) if answer is None: return None @@ -135,23 +135,100 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con # follow_up: return the answer text return {"answer": answer} - def _poll_for_answer(self, session_id, prompt_type): - """Poll DB for user answer until timeout.""" + def poll_steers(self, session_id): + """Drain pending steer docs for ``session_id`` and mark them consumed. + + A "steer" is a mid-flight user message: it's written into the channel + (``_type:"ai"``, ``ai_type:"steer"``, ``status:"pending"``) WHILE the agent + is running, and the worker picks it up at the next loop checkpoint to + redirect the next turn. This is distinct from a follow-up ``answer`` (which + the worker is *blocked* waiting on) and from a hard Stop (which revokes the + Celery task). + + Returns a list of steer content strings (oldest-first). Each returned doc is + flipped to ``status:"consumed"`` so it's injected exactly once. Robust by + design: any backend error returns ``[]`` so a steer can never crash the run. + """ + if self.query_engine is None: + return [] + base = { + "_type": "ai", + "ai_type": "steer", + # Correlate by the runner context's session_id, auto-stamped on every + # persisted item (item._context = self.context) — see _poll_for_answer. + "_context.session_id": session_id, + "status": "pending", + } + try: + results = self.query_engine.search(base, limit=50) + except Exception: # noqa: BLE001 - a steer must never crash the run + return [] + if not results: + return [] + # Oldest-first so multiple queued steers are injected in send order. + results = sorted(results, key=lambda r: r.get("_timestamp", 0)) + contents = [] + for doc in results: + content = doc.get("content") or doc.get("answer") or "" + if content: + contents.append(content) + # Mark this session's pending steers consumed so they inject exactly once. + try: + self.query_engine.update( + {**base}, + {"$set": {"status": "consumed"}}, + ) + except Exception: # noqa: BLE001 - consume failure must not crash the run + pass + return contents + + def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): + """Poll DB for the answer to the SPECIFIC pending prompt until timeout. + + The query MUST be scoped to the exact prompt the worker is currently + blocked on — identified by ``prompt_uuid`` (stamped into the pending doc's + ``extra_data.prompt_uuid`` before it was persisted). Matching only on + ``{session_id, status:"answered"}`` is a bug: a multi-turn conversation + accumulates *previously* answered follow-up docs, so an unscoped query + returns a STALE answer immediately, the worker re-injects that old answer + as a brand-new prompt, re-runs the whole turn, asks again, re-matches the + same stale doc — an infinite respawn loop that re-runs scans and burns + tokens. Scoping on ``prompt_uuid`` makes the poll resolve only THIS + prompt's own answer (and time out only THIS prompt's doc). + + A steer (mid-flight user message) breaks the wait: if a pending steer + arrives for this session while we're blocked on a follow-up, we return its + content as the "answer" so the loop redirects immediately instead of + stalling until the follow-up is explicitly answered (or times out). This + keeps follow-up semantics intact for the no-steer case. + """ + base = { + "_type": "ai", + "ai_type": prompt_type, + # Correlate by the runner context's session_id: it's auto-stamped on + # every persisted item (item._context = self.context), so it's always + # present — unlike the top-level session_id field. + "_context.session_id": session_id, + } + if prompt_uuid: + base["extra_data.prompt_uuid"] = prompt_uuid + elapsed = 0 while elapsed < self.timeout: - results = self.query_engine.search({ - "_type": "ai", - "ai_type": prompt_type, - "session_id": session_id, - "status": "answered" - }, limit=1) + results = self.query_engine.search({**base, "status": "answered"}, limit=1) if results: return results[0].get("answer") + # A steer breaks the wait: treat the steer as the user's answer so the + # blocked follow-up resolves and the next turn redirects. + steers = self.poll_steers(session_id) + if steers: + return "\n".join(steers) sleep(self.poll_interval) elapsed += self.poll_interval - # Timeout: update finding status + # Timeout: flip ONLY this prompt's still-pending doc to timed_out, so a + # concurrent/older pending doc for the same session isn't disturbed. self.query_engine.update( - {"_type": "ai", "ai_type": prompt_type, "session_id": session_id, "status": "pending"}, + {**base, "status": "pending"}, {"$set": {"status": "timed_out"}} ) return None diff --git a/secator/ai/session.py b/secator/ai/session.py index 3023b533a..a6e803404 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -177,3 +177,73 @@ def replay_session(session): except (json.JSONDecodeError, OSError) as e: console.print(Error(message=f'Failed to load history: {e}')) return None + + +def restore_history_from_db(session_id, query_engine, model=None, encryptor=None, system_prompt=None): + """Rebuild an in-memory ChatHistory from the workspace's `_type:"ai"` Mongo docs. + + Headless equivalent of ``replay_session`` for the remote (web) path: a + respawned ``ai`` task on a different worker pod has no local report files, so + the conversation is rebuilt from the channel docs themselves (queried by + ``session_id``, ordered by ``_timestamp``). + + This is a **text-only** restore. Only the user turns (``ai_type="prompt"``) + and assistant turns (``ai_type="response"``) are reconstructed as litellm + ``user``/``assistant`` messages. Intermediate tool-call / tool-result + messages are NOT persisted as ``_type:"ai"`` docs (only their human-readable + action display is), so they cannot be replayed verbatim. Fabricating + assistant ``tool_calls`` messages without their matching ``tool`` results + would produce a malformed transcript that most providers reject, so we + deliberately collapse tool activity into the surrounding text turns. This is + sufficient for ``mode="chat"`` continuation (the assistant text already + summarises what it did); for ``mode="attack"`` the intermediate tool I/O is + not replayed. See the feature spec for the richer-persistence follow-up. + + Args: + session_id: The conversation's session id (UUID generated by the UI). + query_engine: A ``QueryEngine`` (must resolve to the workspace Mongo + backend for the docs to be visible). + model: Optional LLM model name to set on the returned history. + encryptor: Optional ``SensitiveDataEncryptor``. Persisted docs hold + plaintext (response content is decrypted before it is yielded), so + when an encryptor is active we re-encrypt restored turns to keep the + in-memory convention (encrypted) consistent with a fresh run. + system_prompt: Optional system prompt to set as the first message. + + Returns: + ChatHistory: The rebuilt history (possibly with only a system prompt if + no prior docs exist). + """ + from secator.ai.history import ChatHistory + from secator.ai.encryption import maybe_encrypt + + history = ChatHistory(model=model) + if system_prompt is not None: + history.set_system(maybe_encrypt(system_prompt, encryptor)) + + try: + docs = query_engine.search({'_type': 'ai', '_context.session_id': session_id}) + except Exception as e: # noqa: BLE001 - backend errors must not crash the worker + console.print(Warning(message=f'Failed to restore session from DB: {e}')) + return history + + docs = sorted(docs or [], key=lambda d: d.get('_timestamp', 0)) + for doc in docs: + ai_type = doc.get('ai_type') + content = doc.get('content', '') + if not content: + continue + if ai_type == 'prompt': + history.add_user(maybe_encrypt(content, encryptor)) + elif ai_type == 'response': + history.add_assistant(maybe_encrypt(content, encryptor)) + elif ai_type == 'steer': + # A mid-flight steer is a real user turn (an interjection that + # redirected the run): preserve it as a user message on respawn so the + # redirect survives a history restore. Mirror the live-loop framing. + history.add_user(maybe_encrypt(f'[User interjected]: {content}', encryptor)) + # All other ai_types (action displays, follow_up/permission prompts, + # shell_output, summaries) are channel/UX artifacts, not conversation + # turns — intentionally skipped for a valid litellm transcript. + + return history diff --git a/secator/output_types/ai.py b/secator/output_types/ai.py index 192b2933a..ef5bf960f 100644 --- a/secator/output_types/ai.py +++ b/secator/output_types/ai.py @@ -67,6 +67,7 @@ def render_markdown_for_rich(text: str, title: str = '') -> str: 'query': {'label': '🟢', 'color': 'magenta'}, 'stopped': {'label': '🛑', 'color': 'orange3'}, 'follow_up': {'label': '[FOLLOW UP]', 'color': 'orange3'}, + 'steer': {'label': '[STEER]', 'color': 'cyan'}, } ACTION_TYPES = ('task', 'workflow', 'shell', 'add_finding', 'query', 'stopped') diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 0f7ca0e2d..57d04f654 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -1,6 +1,7 @@ # secator/tasks/ai.py """AI-powered penetration testing task.""" import json +import uuid from itertools import groupby from pathlib import Path from time import sleep @@ -15,7 +16,7 @@ from secator.runners import PythonRunner from secator.rich import console, maybe_status from secator.ai.actions import ( - ActionContext, check_guardrails, dispatch_action, _run_batch, _decrypt_dict, _build_action_display + ActionContext, check_guardrails, safe_dispatch_action, _run_batch, _decrypt_dict, _build_action_display ) from secator.ai.guardrails import PermissionEngine from secator.ai.interactivity import create_backend, RemoteBackend @@ -25,7 +26,7 @@ load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue ) from secator.ai.tools import build_tool_schemas, tool_call_to_action, TOOL_SCHEMAS -from secator.ai.session import save_history, show_session_picker, replay_session +from secator.ai.session import save_history, show_session_picker, replay_session, restore_history_from_db from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status @@ -147,6 +148,13 @@ def yielder(self) -> Generator: if not self.model: return + # Remote (web) resume: a respawned chat task restores its history from the + # workspace Mongo `_type:"ai"` docs (headless — no local files, no TUI). + if self.interactive == "remote": + restored = yield from self._maybe_resume_remote() + if restored: + return + # Resume session if self.resume and not self.is_subagent: session = show_session_picker() @@ -161,7 +169,7 @@ def yielder(self) -> Generator: self._reports_folder = session['folder'] result = self._prompt_and_redetect([]) if result is None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return self.context["session_name"] = self.session_name yield from result @@ -209,6 +217,91 @@ def yielder(self) -> Generator: # Run loop yield from self._run_loop() + # ------------------------------------------------------------------------- + # Remote (web) session restore + # ------------------------------------------------------------------------- + + def _get_query_engine(self): + """Build a workspace-scoped QueryEngine from the runner context. + + The backend (mongodb/api/local) is resolved from ``context['drivers']`` + via ``QueryEngine._select_backend``. For the remote channel the API + appends the ``mongodb`` driver on dispatch, so this resolves to the + workspace Mongo backend. + """ + from secator.query import QueryEngine + return QueryEngine(self.context.get("workspace_id", ""), context=dict(self.context)) + + def _maybe_resume_remote(self): + """Restore chat history from Mongo when a remote session has prior docs. + + Returns True (via generator return) if this turn was fully handled as a + respawn (history restored, loop run), False to fall through to a fresh + conversation. Yields any items produced along the way. + """ + query_engine = self._get_query_engine() + + # Guard: remote interactivity requires a Mongo-backed query engine, else + # the RemoteBackend poll can never see the web answer (and restore can't + # read the channel docs). Warn loudly but don't hard-fail a fresh run. + backend_name = getattr(query_engine.backend, "name", "") + if backend_name not in ("mongodb", "api"): + yield Warning( + message=f'interactive="remote" but query engine resolved to "{backend_name}" backend ' + '(expected mongodb/api). The web answer channel will not work — check that the ' + '`mongodb` driver is in the runner context.' + ) + + # Look for prior `_type:"ai"` docs for this session + try: + prior = query_engine.search({"_type": "ai", "_context.session_id": self.session_id}, limit=1) + except Exception as e: # noqa: BLE001 - backend errors must not crash the worker + self.debug(f'remote resume: failed to query prior docs: {e}', sub='llm') + prior = None + + if not prior: + # Fresh conversation: nothing to restore, fall through to normal start. + return False + + # Resolve the user's new prompt (the message that triggered this respawn) + self.prompt = self.run_opts.get("prompt", "") + if self.prompt and Path(self.prompt).is_file(): + self.prompt = Path(self.prompt).read_text().strip() + + # Session metadata + if not self.session_name: + self.session_name = (self.prompt[:80] + '...') if self.prompt and len(self.prompt) > 80 else self.prompt + self.context["session_name"] = self.session_name + + # Detect mode (defaults to chat) and build the system prompt + tools + self._detect_mode() + self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + + # Rebuild history from the channel docs (text-only; see restore_history_from_db) + self.history = restore_history_from_db( + self.session_id, query_engine, model=self.model, + encryptor=self.encryptor, system_prompt=self.system_prompt) + self.history.model = self.model + + # Append the new user message that respawned the conversation + if self.prompt: + self.history.add_user(maybe_encrypt(self.prompt, self.encryptor)) + yield Ai(content=self.prompt, ai_type="prompt", session_id=self.session_id) + + yield Info(message=f"Resumed session from DB ({len(self.history.messages)} messages), model: {self.model}, mode: {self.mode}") # noqa: E501 + yield from self._run_loop() + return True + + def _save_history(self): + """Persist chat history to the local reports folder, unless on the remote path. + + For the remote (web) channel the workspace Mongo `_type:"ai"` docs are the + source of truth, so the local `history.json` write is skipped. + """ + if self.interactive == "remote": + return + save_history(self.history, self.reports_folder, debug_fn=self.debug) + # ------------------------------------------------------------------------- # _run_loop: main LLM interaction loop # ------------------------------------------------------------------------- @@ -249,6 +342,11 @@ def _run_loop(self) -> Generator: iteration += 1 try: + # Mid-flight steering: drain any user messages sent WHILE the agent + # was running and inject them into history so the next turn redirects. + # Cheap query per iteration; robust (never crashes the loop). + yield from self._drain_steers() + # Auto-summarize when context > 85% threshold yield from self._summarize_auto() @@ -286,7 +384,7 @@ def _run_loop(self) -> Generator: yield Warning(message="LLM returned empty response") if empty_streak >= 3: yield Error(message="3 consecutive empty responses - the model may not support tool calling. Stopping.") - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return continue @@ -315,7 +413,7 @@ def _run_loop(self) -> Generator: # Process tool calls → validated actions follow_up_choices = None stop_reason = None - follow_up_ai = None + follow_up_prompt_uuid = None if tool_calls: actions = yield from self._process_tool_calls(tool_calls, ctx) @@ -331,7 +429,7 @@ def _run_loop(self) -> Generator: dispatch_result = yield from self._dispatch_and_collect(actions, ctx) follow_up_choices = dispatch_result.get("follow_up_choices") stop_reason = dispatch_result.get("stop_reason") - follow_up_ai = dispatch_result.get("follow_up_ai") + follow_up_prompt_uuid = dispatch_result.get("follow_up_prompt_uuid") if len(actions) > 1: yield Info(message=f"Executed {len(actions)} actions.") @@ -343,20 +441,20 @@ def _run_loop(self) -> Generator: # Stop tool → save and exit if stop_reason is not None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return # Follow-up / content-only / max_iter → prompt user if follow_up_choices is not None or not tool_calls or iteration == self.max_iterations: - # For remote follow-up, yield the pending Ai so frontend can show it - if follow_up_ai and isinstance(self.backend, RemoteBackend): - follow_up_ai.status = "pending" - follow_up_ai.session_id = self.session_id - yield follow_up_ai + # Remote follow-up: the pending Ai (status="pending" + top-level choices + + # session_id) was already stamped and persisted as a single doc in + # _dispatch_and_collect (add_result dedupes by _uuid, so persistence can + # only happen once). Nothing to re-yield here — the frontend reads the + # persisted doc. - result = self._prompt_and_redetect(follow_up_choices or []) + result = self._prompt_and_redetect(follow_up_choices or [], prompt_uuid=follow_up_prompt_uuid) if result is None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return yield from result continue @@ -370,7 +468,7 @@ def _run_loop(self) -> Generator: yield Warning(message="Interrupted by user.") result = self._prompt_and_redetect([]) if result is None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return yield from result continue @@ -384,7 +482,7 @@ def _run_loop(self) -> Generator: elif isinstance(e, litellm.AuthenticationError): yield Error(message=str(e)) yield Error(message='Please set a valid API key with `secator config set addons.ai.api_key `') - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return elif isinstance(e, litellm.APIConnectionError) or ( isinstance(e, litellm.InternalServerError) and 'connection error' in str(e).lower() @@ -395,13 +493,13 @@ def _run_loop(self) -> Generator: # to avoid swallowing unrelated upstream 500 errors. yield Error(message=f"Cannot connect to model '{self.model}': {e}") yield Error(message='Check api_base and connectivity: `secator config set addons.ai.api_base `') - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return yield Error.from_exception(e) - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() yield Info(message=f"Reached max iterations ({iteration}/{self.max_iterations})") # ------------------------------------------------------------------------- @@ -455,8 +553,19 @@ def _init_options(self): workspace=self.reports_folder or "" ) - # Create interactivity backend - self.session_id = self.session_name or str(self.id) + # Create interactivity backend. + # For the remote (web) channel, the UI generates a stable session_id and + # reuses it verbatim on respawn so a respawned task finds its prior + # `_type:"ai"` docs. It arrives on the runner context (self.context) — + # the dispatcher sends self.context to the worker (task.py build_celery) + # and pops run_opts['context'], so self.context is authoritative here; + # run_opts['context'] only carries it for local/sync runs. + self.session_id = ( + self.passed_context.get("session_id") + or (self.context or {}).get("session_id") + or self.session_name + or str(self.id) + ) self.backend = create_backend(self.interactive, timeout=CONFIG.addons.ai.user_response_timeout) # Auto-approve workspace targets @@ -554,6 +663,48 @@ def _auto_approve_workspace_targets(self): except Exception as e: self.debug(f'[workspace] failed to query targets: {e}', sub='guardrail') + # ------------------------------------------------------------------------- + # Mid-flight steering + # ------------------------------------------------------------------------- + + def _drain_steers(self): + """Drain pending mid-flight steers and inject them into the LLM history. + + A "steer" is a user message sent WHILE the agent is running (over the + remote/web channel: a pending ``_type:"ai", ai_type:"steer"`` doc written by + ``POST /ai/conversations/{id}/steer``). At the top of each loop iteration we + drain any pending steers for this session and append each to the history as + a ``[User interjected]: …`` user message so the model sees them on the next + turn. Cooperative — not a hard cancel (Stop already does that). + + The steer doc the API wrote is itself the persisted transcript entry (it + carries ``_context.session_id``, so the UI's transcript poll surfaces it as + an "interjected" user bubble). We deliberately do NOT yield a second + ``Ai(ai_type="steer")`` echo here — that would persist a duplicate doc with + the same content and double-render in the UI. ``poll_steers`` flips the + drained doc to ``status:"consumed"`` so it injects exactly once. + + Only the RemoteBackend has a channel to poll; for every other backend this + is a no-op. Robust: a steer must never crash the run, so all backend access + is best-effort and swallowed. + + Generator (``yield from``-compatible with the loop) — currently yields no + items, but kept a generator so future transcript echoes can be added without + changing the call site. + """ + if not isinstance(self.backend, RemoteBackend): + return + try: + steers = self.backend.poll_steers(self.session_id) + except Exception as e: # noqa: BLE001 - a steer must never crash the run + self.debug(f'steer: failed to poll steers: {e}', sub='llm') + return + for content in steers: + self.debug(f'steer: injecting user interjection: {content[:120]}', sub='llm') + self.history.add_user(maybe_encrypt(f"[User interjected]: {content}", self.encryptor)) + return + yield # noqa: unreachable - keeps this a generator for `yield from` + # ------------------------------------------------------------------------- # Summarization / compaction # ------------------------------------------------------------------------- @@ -695,9 +846,15 @@ def _dispatch_and_collect(self, actions, ctx): follow_up_choices = None stop_reason = None follow_up_ai = None + follow_up_prompt_uuid = None is_batch = len(actions) > 1 - action_iter = _run_batch(actions, ctx) if is_batch else dispatch_action(actions[0], ctx) + # safe_dispatch_action wraps each action's dispatch so a Python error during + # a handler (e.g. a malformed LLM action/opts raising TypeError) becomes an + # Error item fed back to the LLM as that tool call's result, instead of + # propagating out and killing the main loop. _run_batch already wraps each + # of its actions the same way internally. + action_iter = _run_batch(actions, ctx) if is_batch else safe_dispatch_action(actions[0], ctx) collected = [] for result in action_iter: @@ -708,12 +865,32 @@ def _dispatch_and_collect(self, actions, ctx): is_from_subagent = isinstance(result, OutputType) and bool(result._context.get('subagent')) if isinstance(result, Ai): - self.add_result(result, print=not is_from_subagent) if result.ai_type == "follow_up": follow_up_ai = result follow_up_choices = result.choices or (result.extra_data or {}).get("choices", []) + # Persist the follow-up doc in its FINAL renderable state. add_result() + # dedupes by _uuid, so once persisted here it can never be re-persisted + # (the later `yield follow_up_ai` in the main loop is dropped). For a + # remote run, stamp status="pending" + top-level choices + session_id + # BEFORE the single add_result, so the one persisted doc is what the web + # UI needs: status=="pending" (clears "thinking") and non-empty choices. + if isinstance(self.backend, RemoteBackend): + follow_up_ai.status = "pending" + follow_up_ai.session_id = self.session_id + if not follow_up_ai.choices and follow_up_choices: + follow_up_ai.choices = list(follow_up_choices) + # Stamp a unique correlation id so the poll resolves ONLY this + # prompt's own answer (not a stale answered follow_up from a + # prior turn, which would loop). Generated here (not reusing + # _uuid, which mongo may reassign to its _id on insert) and + # persisted in extra_data so it round-trips on read. + follow_up_prompt_uuid = str(uuid.uuid4()) + follow_up_ai.extra_data = { + **(follow_up_ai.extra_data or {}), "prompt_uuid": follow_up_prompt_uuid} + self.add_result(result, print=not is_from_subagent) continue - elif result.ai_type == "stopped": + self.add_result(result, print=not is_from_subagent) + if result.ai_type == "stopped": stop_reason = result.content continue if result.ai_type not in ("shell_output", "response"): @@ -754,7 +931,12 @@ def _dispatch_and_collect(self, actions, ctx): tool_result_str = maybe_encrypt(tool_result_str, self.encryptor) self.history.add_tool_result(tc_name, tc_id, tool_result_str) - return {"follow_up_choices": follow_up_choices, "stop_reason": stop_reason, "follow_up_ai": follow_up_ai} + return { + "follow_up_choices": follow_up_choices, + "stop_reason": stop_reason, + "follow_up_ai": follow_up_ai, + "follow_up_prompt_uuid": follow_up_prompt_uuid, + } # ------------------------------------------------------------------------- # History helpers @@ -782,12 +964,16 @@ def _add_assistant_to_history(self, content, tool_calls): # Follow-up / prompt # ------------------------------------------------------------------------- - def _prompt_and_redetect(self, choices): + def _prompt_and_redetect(self, choices, prompt_uuid=None): """Prompt user via backend and re-detect intent. Works for all backends: CLIBackend shows rich menus, RemoteBackend polls DB, AutoBackend returns None (exits). + ``prompt_uuid`` correlates the (remote) poll to the SPECIFIC pending + follow_up doc this call raised, so a stale answered follow_up from a prior + turn can't resolve it (which would re-inject the old prompt and loop). + Returns list of items to yield, or None to exit. """ response = self.backend.ask_user( @@ -800,6 +986,7 @@ def _prompt_and_redetect(self, choices): max_iterations=self.max_iterations, mode=self.mode, model=self.model, + prompt_uuid=prompt_uuid, ) if response is None: return None diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 05ff6c99b..735867076 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -9,7 +9,8 @@ if ADDONS_ENABLED['ai']: from secator.ai.actions import ( ActionContext, dispatch_action, _handle_follow_up, _handle_shell, - _handle_query, _handle_add_finding, _run_runner, _decrypt_dict + _handle_query, _handle_add_finding, _run_runner, _decrypt_dict, + _build_hooks_from_context, _coerce_finding_fields ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -119,6 +120,9 @@ def test_follow_up_with_choices(self): self.assertEqual(results[0].ai_type, 'follow_up') self.assertEqual(results[0].content, 'What next?') self.assertEqual(results[0].extra_data['choices'], ['Scan deeper', 'Try SQL injection']) + # Choices must also land on the top-level `choices` field (what the web UI reads), + # not only in extra_data — otherwise the persisted follow-up doc renders no buttons. + self.assertEqual(results[0].choices, ['Scan deeper', 'Try SQL injection']) @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') @@ -316,6 +320,132 @@ def test_run_runner_uses_ctx_targets_as_default(self): self.assertIn('default.com', results[0].message) + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_propagates_hooks_and_emits_runner_id(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """Sub-runner must receive driver hooks (so its results persist) and the + emitted action Ai must carry the created runner's id + type for the UI.""" + sentinel_hooks = {'fake': ['hook']} + mock_build_hooks.return_value = sentinel_hooks + + # Fake runner: an iterable whose id is populated (mimics on_init stamping it) + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + ctx = ActionContext( + targets=['t.com'], model='m', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}, + ) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + + results = list(_run_runner(action, ctx, 'task')) + + # Runner constructed with hooks= from the context drivers + _, kwargs = mock_task_cls.call_args + self.assertEqual(kwargs.get('hooks'), sentinel_hooks) + self.assertEqual(kwargs.get('context', {}).get('workspace_id'), 'ws1') + + # Action Ai item carries runner_id + runner_type + ai_items = [r for r in results if isinstance(r, Ai) and r.ai_type == 'task'] + self.assertEqual(len(ai_items), 1) + self.assertEqual(ai_items[0].extra_data.get('runner_id'), 'runner123') + self.assertEqual(ai_items[0].extra_data.get('runner_type'), 'task') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_propagates_session_id(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """The dispatched sub-runner's context must carry the ai task's session_id + (the conversation id) so its persisted runner doc is queryable by the + conversation. session_id may be derived (not already in ctx.context), so + it must be stamped from ctx.session_id.""" + mock_build_hooks.return_value = {} + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + # session_id lives on the ActionContext but NOT in context (it is derived) + ctx = ActionContext( + targets=['t.com'], model='m', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}, + session_id='conv-abc-123', + ) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + sub_context = kwargs.get('context', {}) + self.assertEqual(sub_context.get('session_id'), 'conv-abc-123') + self.assertEqual(sub_context.get('workspace_id'), 'ws1') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_preserves_existing_session_id(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """A session_id already present in ctx.context must not be overwritten.""" + mock_build_hooks.return_value = {} + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + ctx = ActionContext( + targets=['t.com'], model='m', + context={'workspace_id': 'ws1', 'session_id': 'from-context'}, + session_id='from-ctx-field', + ) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + self.assertEqual(kwargs.get('context', {}).get('session_id'), 'from-context') + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestBuildHooksFromContext(unittest.TestCase): + """Tests for _build_hooks_from_context (driver name -> hooks dict).""" + + def test_no_drivers_returns_empty(self): + self.assertEqual(_build_hooks_from_context({}), {}) + self.assertEqual(_build_hooks_from_context({'drivers': []}), {}) + + @patch('secator.loader.get_available_drivers') + @patch('secator.loader.order_drivers') + @patch('secator.loader.discover_external_drivers') + @patch('secator.utils.import_dynamic') + def test_builds_hooks_from_driver_names(self, mock_import, _disc, mock_order, mock_avail): + from secator.runners import Task + mock_order.side_effect = lambda d: d + mock_avail.return_value = ['mongodb', 'api'] + mongo_hooks = {Task: {'on_init': ['update_runner']}} + mock_import.return_value = mongo_hooks + + hooks = _build_hooks_from_context({'drivers': ['mongodb']}) + + mock_import.assert_called_once_with('secator.hooks.mongodb', 'HOOKS') + self.assertIn(Task, hooks) + self.assertIn('on_init', hooks[Task]) + + @patch('secator.loader.get_available_drivers') + @patch('secator.loader.order_drivers') + @patch('secator.loader.discover_external_drivers') + @patch('secator.utils.import_dynamic') + def test_skips_unsupported_driver(self, mock_import, _disc, mock_order, mock_avail): + mock_order.side_effect = lambda d: d + mock_avail.return_value = ['mongodb'] + hooks = _build_hooks_from_context({'drivers': ['bogus']}) + self.assertEqual(hooks, {}) + mock_import.assert_not_called() + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestGetQueryEngine(unittest.TestCase): @@ -647,6 +777,77 @@ def test_add_finding_decrypts_values(self): self.assertIsInstance(results[1], Vulnerability) self.assertEqual(results[1].matched_at, 'http://t.com/search') + def test_coerce_finding_fields_scalar_types(self): + # LLMs send wrong-typed scalars (bool as "true", float/int as strings). + # The coercion helper fixes them to the declared field types. + data = _coerce_finding_fields( + Vulnerability, + { + 'name': 'SQL Injection', + 'verified': 'true', + 'cvss_score': '7.5', + 'severity_nb': '3', + }, + ) + self.assertIs(data['verified'], True) + self.assertIsInstance(data['verified'], bool) + self.assertEqual(data['cvss_score'], 7.5) + self.assertIsInstance(data['cvss_score'], float) + self.assertEqual(data['severity_nb'], 3) + self.assertIsInstance(data['severity_nb'], int) + # str fields are left untouched. + self.assertEqual(data['name'], 'SQL Injection') + # Coerced data validates clean. + self.assertEqual(Vulnerability.validate_fields(data), []) + + def test_add_finding_coerces_scalar_types(self): + # End-to-end: wrong-typed scalars flow through the handler and validate + # clean, producing a Vulnerability with the coerced bool/float values. + ctx = ActionContext(targets=['t.com'], model='m') + results = list( + _handle_add_finding( + { + 'action': 'add_finding', + '_type': 'vulnerability', + 'name': 'SQL Injection', + 'matched_at': 'http://t.com/login', + 'verified': 'true', + 'cvss_score': '7.5', + 'severity_nb': '3', + }, + ctx, + ) + ) + + # No validation Error: the sloppy types were coerced before validation. + self.assertEqual(len(results), 2) + vuln = results[1] + self.assertIsInstance(vuln, Vulnerability) + self.assertIs(vuln.verified, True) + self.assertIsInstance(vuln.verified, bool) + self.assertEqual(vuln.cvss_score, 7.5) + self.assertIsInstance(vuln.cvss_score, float) + + def test_add_finding_unparseable_bool_surfaces_error(self): + # An unparseable value must NOT be silently dropped; validation reports it. + ctx = ActionContext(targets=['t.com'], model='m') + results = list( + _handle_add_finding( + { + 'action': 'add_finding', + '_type': 'vulnerability', + 'name': 'SQL Injection', + 'matched_at': 'http://t.com/login', + 'verified': 'maybe', + }, + ctx, + ) + ) + + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Error) + self.assertIn('verified', results[0].message) + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestRunBatch(unittest.TestCase): diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 6060c5f16..bec54f3be 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -100,14 +100,62 @@ def test_ask_user_polls_until_timeout(self, mock_sleep): # Should have called update to set timed_out mock_engine.update.assert_called_once() + def test_poll_scopes_query_to_prompt_uuid(self): + """The poll must correlate on the specific prompt's uuid. + + Regression test for the infinite-respawn loop: without scoping on + prompt_uuid, a stale answered follow_up from a prior turn resolves the + current wait immediately, the worker re-injects that old answer as a new + prompt and re-runs the turn forever. The query MUST include + extra_data.prompt_uuid so only THIS prompt's own answer resolves it. + """ + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [{"answer": "the right answer"}] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="abc-123") + + self.assertEqual(result["answer"], "the right answer") + # The search query must be scoped to this prompt's uuid (else a stale + # answered follow_up from a prior turn would match -> loop). + search_query = mock_engine.search.call_args[0][0] + self.assertEqual(search_query.get("extra_data.prompt_uuid"), "abc-123") + self.assertEqual(search_query.get("status"), "answered") + + @patch('secator.ai.interactivity.sleep') + def test_timeout_update_scoped_to_prompt_uuid(self, mock_sleep): + """On timeout, only THIS prompt's pending doc is flipped to timed_out.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [] # never answered + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="abc-123") + + self.assertIsNone(result) + mock_engine.update.assert_called_once() + update_query = mock_engine.update.call_args[0][0] + self.assertEqual(update_query.get("extra_data.prompt_uuid"), "abc-123") + self.assertEqual(update_query.get("status"), "pending") + @patch('secator.ai.interactivity.sleep') def test_ask_user_returns_on_second_poll(self, mock_sleep): from secator.ai.interactivity import RemoteBackend mock_engine = MagicMock() - mock_engine.search.side_effect = [ - [], # first poll: not answered - [{"answer": "option B"}], # second poll: answered - ] + # Query-aware: the follow-up answer poll (ai_type=="follow_up") returns the + # answer on the second call; the interleaved steer poll (ai_type=="steer") + # always returns nothing — so the steer-break never fires here. + answer_calls = {"n": 0} + + def search(query, limit=1): + if query.get("ai_type") == "steer": + return [] + answer_calls["n"] += 1 + return [] if answer_calls["n"] == 1 else [{"answer": "option B"}] + + mock_engine.search.side_effect = search backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=5) result = backend.ask_user("What next?", [], "session1") @@ -117,6 +165,76 @@ def test_ask_user_returns_on_second_poll(self, mock_sleep): self.assertEqual(mock_sleep.call_count, 1) +class TestRemoteBackendSteer(unittest.TestCase): + """Verify mid-flight steer draining + the blocked-wait break.""" + + def test_poll_steers_returns_and_consumes(self): + """poll_steers returns pending steer content and marks them consumed.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"content": "actually focus on the API", "_timestamp": 2}, + {"content": "and skip port 80", "_timestamp": 1}, + ] + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + steers = backend.poll_steers("session1") + + # Oldest-first by _timestamp + self.assertEqual(steers, ["and skip port 80", "actually focus on the API"]) + # Query scoped to pending steer docs for this session + search_query = mock_engine.search.call_args[0][0] + self.assertEqual(search_query.get("ai_type"), "steer") + self.assertEqual(search_query.get("status"), "pending") + self.assertEqual(search_query.get("_context.session_id"), "session1") + # Pending steers flipped to consumed (inject exactly once) + mock_engine.update.assert_called_once() + update_set = mock_engine.update.call_args[0][1] + self.assertEqual(update_set["$set"]["status"], "consumed") + + def test_poll_steers_no_pending_returns_empty(self): + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + self.assertEqual(backend.poll_steers("session1"), []) + # Nothing to consume when nothing is pending + mock_engine.update.assert_not_called() + + def test_poll_steers_robust_on_backend_error(self): + """A steer must never crash the run: backend errors return [].""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.side_effect = RuntimeError("mongo down") + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + self.assertEqual(backend.poll_steers("session1"), []) + + def test_poll_steers_no_query_engine(self): + from secator.ai.interactivity import RemoteBackend + backend = RemoteBackend(timeout=60, query_engine=None, poll_interval=0.01) + self.assertEqual(backend.poll_steers("session1"), []) + + def test_steer_breaks_blocked_follow_up_wait(self): + """A steer arriving during a follow-up wait returns as the answer.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + # No follow-up answer ever; a steer arrives on the first poll. + mock_engine.search.side_effect = [ + [], # answered? no + [{"content": "change course now", "_timestamp": 1}], # poll_steers -> steer + ] + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="uuid-1") + + # The steer content resolves the blocked wait (returned as the answer). + self.assertEqual(result["answer"], "change course now") + + class TestCreateBackend(unittest.TestCase): """Verify create_backend factory.""" diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index 39e179cfd..2f89ad40c 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -307,6 +307,82 @@ def test_stop_yields_ai_stopped(self): self.assertIn("completed", ai_results[0].content) +# ============================================================================= +# UNIT TESTS: Remote follow-up persistence (status + top-level choices) +# ============================================================================= + +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestRemoteFollowUpPersistence(unittest.TestCase): + """In remote mode, the single persisted follow-up doc must be renderable: + status=="pending" + non-empty top-level `choices` (what the web UI reads).""" + + def _run_dispatch(self, backend): + """Drive the real ai._dispatch_and_collect with a minimal fake self. + + Returns (yielded_items, persisted_items) where persisted_items are what + add_result() received (i.e. what the mongodb on_item hook would persist). + """ + from secator.tasks.ai import ai as AiTask + + choices = ["Fuzz parameters", "Run nuclei", "Deep crawl"] + follow_up = Ai( + content="Presenting actionable next steps", + ai_type="follow_up", + extra_data={"choices": choices}, + _context={"tool_call_id": "tc_fu", "tool_call_name": "follow_up"}, + ) + + persisted = [] + + class _FakeHistory: + def get_action_budget(self, model): + return 10000 + + def add_tool_result(self, *a, **k): + pass + + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-123" + fake_self.model = "test-model" + fake_self.reports_folder = None + fake_self.history = _FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + ctx = MagicMock() + ctx.results = [] + + def _fake_dispatch_action(action, c): + yield follow_up + + with patch("secator.tasks.ai.safe_dispatch_action", _fake_dispatch_action): + gen = AiTask._dispatch_and_collect(fake_self, [{"tool_call_id": "tc_fu"}], ctx) + yielded = list(gen) + return yielded, persisted, follow_up + + def test_remote_follow_up_persisted_pending_with_choices(self): + backend = RemoteBackend(timeout=60, query_engine=MagicMock()) + yielded, persisted, follow_up = self._run_dispatch(backend) + + # Exactly one follow_up Ai is persisted (no duplicate display + pending docs). + fu_docs = [p for p in persisted if isinstance(p, Ai) and p.ai_type == "follow_up"] + self.assertEqual(len(fu_docs), 1) + doc = fu_docs[0] + self.assertEqual(doc.status, "pending") + self.assertEqual(doc.choices, ["Fuzz parameters", "Run nuclei", "Deep crawl"]) + self.assertEqual(doc.session_id, "sess-123") + # Same object → single doc by _uuid. + self.assertIs(doc, follow_up) + + def test_local_follow_up_not_stamped_pending(self): + """CLI/local mode must NOT stamp status=pending (drives the TUI menu directly).""" + backend = CLIBackend() + yielded, persisted, follow_up = self._run_dispatch(backend) + fu_docs = [p for p in persisted if isinstance(p, Ai) and p.ai_type == "follow_up"] + self.assertEqual(len(fu_docs), 1) + self.assertNotEqual(fu_docs[0].status, "pending") + + # ============================================================================= # UNIT TESTS: Backend and tool schema behavior # ============================================================================= @@ -939,5 +1015,201 @@ def test_multi_turn_auto_loop(self): self.assertIsNotNone(stop_reason) +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestLoopResilientToActionErrors(unittest.TestCase): + """A Python error during an action dispatch must NOT kill the main loop. + + It must be caught, turned into an Error item fed back to the LLM as that + tool call's result, and the loop must continue. + """ + + def test_safe_dispatch_catches_exception_and_feeds_back(self): + """safe_dispatch_action converts a raised Exception into an Error item + carrying the action's tool_call_id, instead of propagating.""" + from secator.ai.actions import safe_dispatch_action + from secator.output_types import Error + + ctx = _make_ctx(interactive="auto") + action = { + "action": "shell", + "command": "curl http://10.0.0.1", + "tool_call_id": "tc_err", + "tool_call_name": "run_shell", + } + + # Make the shell handler raise the exact failure from the spec. + def _boom(*a, **k): + raise TypeError("'str' object is not a mapping") + + with patch("secator.ai.actions._handle_shell", _boom): + # Must NOT raise. + results = list(safe_dispatch_action(action, ctx)) + + errors = [r for r in results if isinstance(r, Error)] + self.assertEqual(len(errors), 1, "expected exactly one Error item") + err = errors[0] + # LLM-facing feedback phrasing + the exception type/message. + self.assertIn("Action failed with error", err.message) + self.assertIn("TypeError", err.message) + self.assertIn("'str' object is not a mapping", err.message) + self.assertIn("try again", err.message.lower()) + # Attributed to the failing tool call so it groups into that tool result. + self.assertEqual(err._context.get("tool_call_id"), "tc_err") + self.assertEqual(err._context.get("tool_call_name"), "run_shell") + + def test_does_not_catch_keyboardinterrupt(self): + """Control-flow exceptions (BaseException) must propagate, not be swallowed.""" + from secator.ai.actions import safe_dispatch_action + + ctx = _make_ctx(interactive="auto") + action = {"action": "shell", "command": "x", "tool_call_id": "tc", "tool_call_name": "run_shell"} + + def _interrupt(*a, **k): + raise KeyboardInterrupt() + yield # pragma: no cover - make it a generator + + with patch("secator.ai.actions._handle_shell", _interrupt): + with self.assertRaises(KeyboardInterrupt): + list(safe_dispatch_action(action, ctx)) + + def test_dispatch_and_collect_continues_and_feeds_history(self): + """Drive the real _dispatch_and_collect: a raising action yields an Error, + appends an error result to history (LLM-visible), and does NOT raise.""" + from secator.tasks.ai import ai as AiTask + from secator.output_types import Error + + tool_results = [] # (name, tc_id, content) tuples appended to history + + class _FakeHistory: + def get_action_budget(self, model): + return 10000 + + def add_tool_result(self, name, tc_id, content): + tool_results.append((name, tc_id, content)) + + persisted = [] + fake_self = MagicMock() + fake_self.backend = CLIBackend() + fake_self.session_id = "sess-err" + fake_self.model = "test-model" + fake_self.reports_folder = None + fake_self.encryptor = None + fake_self.history = _FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + ctx = MagicMock() + ctx.results = [] + + action = { + "action": "shell", + "command": "curl http://10.0.0.1", + "tool_call_id": "tc_err", + "tool_call_name": "run_shell", + } + + def _boom(*a, **k): + raise TypeError("'str' object is not a mapping") + yield # pragma: no cover + + with patch("secator.ai.actions._handle_shell", _boom): + # Single action → safe_dispatch_action path. Must not raise. + gen = AiTask._dispatch_and_collect(fake_self, [action], ctx) + yielded = list(gen) + + # An Error item was yielded to the caller (visible in console / persisted). + errors = [r for r in yielded if isinstance(r, Error)] + self.assertEqual(len(errors), 1) + + # The error reached the LLM-visible history as this tool call's result. + self.assertEqual(len(tool_results), 1) + name, tc_id, content = tool_results[0] + self.assertEqual(tc_id, "tc_err") + self.assertIn("error", content.lower()) + self.assertIn("'str' object is not a mapping", content) + + +# ============================================================================= +# Mid-flight steering: _drain_steers injects pending steers into history +# ============================================================================= + +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestDrainSteers(unittest.TestCase): + """The loop's `_drain_steers` drains pending steers and injects them. + + A pending `ai_type:"steer"` doc (user message sent WHILE the agent runs) must + be drained at the loop checkpoint, appended to the LLM history as a + `[User interjected]: …` user message, echoed as a steer Ai item, and marked + consumed — without breaking the loop or the existing follow-up flow. + """ + + def _make_task(self, backend, history=None): + """Build a minimal `ai` task with only what _drain_steers reads.""" + from secator.tasks.ai import ai + task = object.__new__(ai) + task.backend = backend + task.session_id = "steer-sess" + task.encryptor = None + task.history = history or ChatHistory() + task.debug = lambda *a, **k: None + return task + + def test_steer_drained_injected_and_consumed(self): + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"content": "actually focus on the API", "_timestamp": 1}, + ] + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + task = self._make_task(backend) + + yielded = list(task._drain_steers()) + + # Injected into history as a user "interjected" message. + user_msgs = [m for m in task.history.to_messages() if m["role"] == "user"] + self.assertEqual(len(user_msgs), 1) + self.assertEqual(user_msgs[-1]["content"], "[User interjected]: actually focus on the API") + + # No echo doc is yielded: the API's pending steer doc is itself the + # persisted transcript entry, so a second steer Ai would double-render. + self.assertEqual(yielded, []) + + # Marked consumed so it injects exactly once. + update_set = mock_engine.update.call_args[0][1] + self.assertEqual(update_set["$set"]["status"], "consumed") + + def test_no_steer_is_noop_and_preserves_loop(self): + """No pending steer -> nothing injected, history untouched (loop intact).""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + history = ChatHistory() + history.add_user("original prompt") + task = self._make_task(backend, history=history) + + yielded = list(task._drain_steers()) + + self.assertEqual(yielded, []) + user_msgs = [m for m in task.history.to_messages() if m["role"] == "user"] + self.assertEqual([m["content"] for m in user_msgs], ["original prompt"]) + + def test_non_remote_backend_is_noop(self): + """Local/auto backends have no channel -> drain is a no-op (no crash).""" + task = self._make_task(create_backend("auto")) + self.assertEqual(list(task._drain_steers()), []) + + def test_steer_poll_error_never_crashes_loop(self): + """A backend error during drain is swallowed (run must not crash).""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.side_effect = RuntimeError("mongo down") + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + task = self._make_task(backend) + # Should not raise, yields nothing, history untouched. + self.assertEqual(list(task._drain_steers()), []) + self.assertEqual(task.history.to_messages(), []) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py new file mode 100644 index 000000000..b371213d8 --- /dev/null +++ b/tests/unit/test_ai_session.py @@ -0,0 +1,188 @@ +"""Tests for secator.ai.session restore_history_from_db + remote resume branch.""" +import tempfile +import unittest +from unittest.mock import MagicMock, patch + + +class TestRestoreHistoryFromDB(unittest.TestCase): + """Verify restore_history_from_db rebuilds an equivalent ChatHistory from Mongo docs.""" + + def _docs(self): + # Intentionally out of timestamp order to verify sorting. + return [ + {"_type": "ai", "ai_type": "response", "content": "Hi, how can I help?", "_timestamp": 2}, + {"_type": "ai", "ai_type": "prompt", "content": "Hello", "_timestamp": 1}, + {"_type": "ai", "ai_type": "shell", "content": "nmap -p- host", "_timestamp": 3}, + {"_type": "ai", "ai_type": "prompt", "content": "Scan the target", "_timestamp": 4}, + {"_type": "ai", "ai_type": "follow_up", "content": "What next?", "_timestamp": 5}, + {"_type": "ai", "ai_type": "response", "content": "Found 2 open ports.", "_timestamp": 6}, + ] + + def test_rebuilds_order_roles_and_system(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = self._docs() + + history = restore_history_from_db( + "session1", engine, model="gpt-4o", system_prompt="SYSTEM PROMPT") + + # Query was scoped to the session + engine.search.assert_called_once_with({"_type": "ai", "session_id": "session1"}) + + # System prompt set, conversation turns in timestamp order, non-turn docs skipped + self.assertEqual(history.messages, [ + {"role": "system", "content": "SYSTEM PROMPT"}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi, how can I help?"}, + {"role": "user", "content": "Scan the target"}, + {"role": "assistant", "content": "Found 2 open ports."}, + ]) + self.assertEqual(history.model, "gpt-4o") + + def test_no_prior_docs_returns_system_only(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [] + + history = restore_history_from_db("s2", engine, system_prompt="SYS") + self.assertEqual(history.messages, [{"role": "system", "content": "SYS"}]) + + def test_no_system_prompt_yields_empty_when_no_docs(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [] + + history = restore_history_from_db("s3", engine) + self.assertEqual(history.messages, []) + + def test_empty_content_docs_skipped(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [ + {"ai_type": "prompt", "content": "", "_timestamp": 1}, + {"ai_type": "response", "content": "Real answer", "_timestamp": 2}, + ] + history = restore_history_from_db("s4", engine) + self.assertEqual(history.messages, [{"role": "assistant", "content": "Real answer"}]) + + def test_search_failure_returns_system_only(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.side_effect = RuntimeError("backend down") + + history = restore_history_from_db("s5", engine, system_prompt="SYS") + # Failure must not crash; returns just the system prompt + self.assertEqual(history.messages, [{"role": "system", "content": "SYS"}]) + + def test_encryptor_reencrypts_restored_turns(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [ + {"ai_type": "prompt", "content": "scan 10.0.0.1", "_timestamp": 1}, + ] + encryptor = MagicMock() + encryptor.encrypt.side_effect = lambda t: f"ENC({t})" + + history = restore_history_from_db("s6", engine, encryptor=encryptor) + self.assertEqual(history.messages, [{"role": "user", "content": "ENC(scan 10.0.0.1)"}]) + + +class TestRemoteResumeBranch(unittest.TestCase): + """Verify the yielder remote-resume branch picks Mongo restore vs fresh start.""" + + def _make_task(self, prior_docs, backend_name="mongodb"): + from secator.tasks.ai import ai + + task = ai.__new__(ai) + # Minimal attributes the branch touches + task.interactive = "remote" + task.session_id = "sess-123" + task.session_name = "" + task.mode = "chat" + task.model = "gpt-4o" + task.encryptor = None + task.context = {"workspace_id": "ws1", "drivers": ["mongodb"]} + task.run_opts = {"prompt": "Tell me about this workspace"} + # An existing dir short-circuits the reports_folder property (no dir creation) + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task.backend = MagicMock() + task.debug = MagicMock() + task.history = MagicMock() + + # Stub query engine + engine = MagicMock() + engine.backend = MagicMock() + engine.backend.name = backend_name + + def _search(query, limit=0): + if query.get("_type") == "ai" and "session_id" in query: + return prior_docs + return [] + engine.search.side_effect = _search + task._get_query_engine = MagicMock(return_value=engine) + return task, engine + + def test_fresh_when_no_prior_docs(self): + task, engine = self._make_task(prior_docs=[]) + # Generator return value is the StopIteration value. + gen = task._maybe_resume_remote() + restored = None + try: + while True: + next(gen) + except StopIteration as e: + restored = e.value + self.assertFalse(restored) + + @patch("secator.tasks.ai.restore_history_from_db") + @patch("secator.tasks.ai.get_system_prompt", return_value="SYS") + def test_restores_when_prior_docs(self, mock_sys, mock_restore): + mock_history = MagicMock() + mock_history.messages = [{"role": "system", "content": "SYS"}] + mock_restore.return_value = mock_history + + task, engine = self._make_task(prior_docs=[{"ai_type": "prompt", "content": "hi"}]) + # Stub the heavy methods the branch calls + task._detect_mode = MagicMock() + task._run_loop = MagicMock(return_value=iter([])) + + gen = task._maybe_resume_remote() + restored = None + try: + while True: + next(gen) + except StopIteration as e: + restored = e.value + + self.assertTrue(restored) + mock_restore.assert_called_once() + # Restored from Mongo via the resolved query engine + _, kwargs = mock_restore.call_args + self.assertEqual(mock_restore.call_args[0][0], "sess-123") + task._run_loop.assert_called_once() + + @patch("secator.tasks.ai.restore_history_from_db") + @patch("secator.tasks.ai.get_system_prompt", return_value="SYS") + def test_warns_on_non_mongo_backend(self, mock_sys, mock_restore): + from secator.output_types import Warning as WarningType + mock_restore.return_value = MagicMock(messages=[]) + + task, engine = self._make_task( + prior_docs=[{"ai_type": "prompt", "content": "hi"}], backend_name="local") + task._detect_mode = MagicMock() + task._run_loop = MagicMock(return_value=iter([])) + + items = [] + gen = task._maybe_resume_remote() + try: + while True: + items.append(next(gen)) + except StopIteration: + pass + + warnings = [i for i in items if isinstance(i, WarningType)] + self.assertTrue(any("remote" in w.message for w in warnings)) + + +if __name__ == "__main__": + unittest.main()