From 2de2a87d36e61e3fdecc45ca41c0004fa3554a2e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 23 Jun 2026 18:46:42 +0200 Subject: [PATCH 01/18] feat: headless Mongo session restore for remote AI chat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the two secator-core gaps for the Workspace AI Assistant (Mongo-channel chat), repo 1/3: - `restore_history_from_db(session_id, query_engine, model, encryptor, system_prompt)` in `secator/ai/session.py`: rebuilds a `ChatHistory` from the workspace `_type:"ai"` docs (queried by session_id, ordered by `_timestamp`) — `prompt`→user, `response`→assistant, system prompt set, re-encrypted when an encryptor is active. Headless: no local files, no TUI. - Wire a remote-resume branch in `ai.py:yielder`: when `interactive="remote"` and the session has prior `_type:"ai"` docs, restore from Mongo and continue; fresh conversations (no docs) start as before. The local CLI `replay_session`/`show_session_picker` path is untouched. - `session_id` now prefers `run_opts.context.session_id` so a respawned task finds its prior docs. - `save_history` (local `history.json`) is skipped on the remote path via a `_save_history()` helper — the Mongo docs are the source of truth. - Query-engine guard: warn when `interactive="remote"` but the resolved query backend is not mongodb/api (the web answer channel can't work otherwise). History-fidelity finding: persisted `_type:"ai"` docs capture only text turns (prompt/response) plus action *display* records — not the litellm assistant `tool_calls` messages or their `tool` results. Restore is therefore text-only. This is valid and sufficient for `mode="chat"` continuation; fabricating partial tool-call messages would produce a malformed transcript providers reject, so tool activity is deliberately collapsed. Richer assistant persistence for `mode="attack"` replay is a documented follow-up. Tests: `tests/unit/test_ai_session.py` — restore rebuilds equivalent History (order/roles/system/encryption/empty-docs/search-failure), and the remote-resume branch picks Mongo restore for prior docs / fresh otherwise / warns on non-Mongo backend. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/session.py | 65 ++++++++++++ secator/tasks/ai.py | 119 ++++++++++++++++++--- tests/unit/test_ai_session.py | 188 ++++++++++++++++++++++++++++++++++ 3 files changed, 360 insertions(+), 12 deletions(-) create mode 100644 tests/unit/test_ai_session.py diff --git a/secator/ai/session.py b/secator/ai/session.py index 3023b533a..ea1ffae48 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -177,3 +177,68 @@ 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', '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)) + # 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/tasks/ai.py b/secator/tasks/ai.py index 0f7ca0e2d..93dd4be4b 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -25,7 +25,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 +147,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 +168,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 +216,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", "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 # ------------------------------------------------------------------------- @@ -286,7 +378,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 @@ -343,7 +435,7 @@ 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 @@ -356,7 +448,7 @@ def _run_loop(self) -> Generator: result = self._prompt_and_redetect(follow_up_choices or []) 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 +462,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 +476,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 +487,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 +547,11 @@ 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 (passed in run_opts.context.session_id); + # prefer it so a respawned task can find its prior `_type:"ai"` docs. + self.session_id = self.passed_context.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 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() From 94c43f3f1d0e2c926f9b51edc6b43469bf043360 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 09:07:15 +0200 Subject: [PATCH 02/18] fix(ai): stamp session_id on every Ai item for the remote-channel transcript MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The web UI correlates an AI chat conversation by session_id (across respawns), but only the resume-prompt and follow_up items set it — the prompt/response/ token_usage/chat_compacted message items did not, so they persisted to Mongo without session_id and the UI's {_type:"ai", session_id} query returned nothing (empty transcript despite the task running fine). Wrap yielder to stamp session_id on every Ai item centrally (self.session_id is set in _init_options before the first yield). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/tasks/ai.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 93dd4be4b..bededda96 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -119,6 +119,17 @@ def requires_local_execution(cls, inputs, run_opts): # ------------------------------------------------------------------------- def yielder(self) -> Generator: + """Stamp every Ai item with the session_id so the remote-channel transcript + is queryable by session_id. The web UI correlates the whole conversation + (across respawns) by session_id, so a message item without it is invisible. + _init_options() sets self.session_id before the first yield, so the stamp + is always valid here.""" + for _item in self._yielder(): + if isinstance(_item, Ai) and not getattr(_item, "session_id", ""): + _item.session_id = self.session_id + yield _item + + def _yielder(self) -> Generator: """Execute AI task.""" # Addon / setup check if self.inputs == ['setup']: From 787f480bf911e559419af292cbf4c9ad63bf26af Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 10:32:06 +0200 Subject: [PATCH 03/18] fix(ai): read session_id from self.context (dispatch drops run_opts.context) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The web UI's session_id arrives on the runner context, but the Task dispatcher sends self.context (not run_opts['context']) to the worker and pops run_opts['context'] — so in the worker run_opts.context is empty and session_id fell back to the prompt label, never matching the UI's UUID (empty transcript). Prefer self.context for session_id. Pairs with secator-api adding session_id to the RunnerContext model so it survives validation into self.context. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/tasks/ai.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index bededda96..6592af41e 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -560,9 +560,17 @@ def _init_options(self): # Create interactivity backend. # For the remote (web) channel, the UI generates a stable session_id and - # reuses it verbatim on respawn (passed in run_opts.context.session_id); - # prefer it so a respawned task can find its prior `_type:"ai"` docs. - self.session_id = self.passed_context.get("session_id") or self.session_name or str(self.id) + # 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 From 870ab2a4a9a777fb50ea646b6e8fd72d50bfb54e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 11:38:49 +0200 Subject: [PATCH 04/18] fix(ai): correlate chat channel by _context.session_id (top-level was empty) Persisted _type:"ai" docs had session_id="" but _context.session_id=: the runner auto-stamps item._context = self.context, so _context.session_id is reliably present, while the top-level session_id field never landed. Query _context.session_id in _poll_for_answer, the timeout update, restore_history_from_db and the resume check; drop the now-pointless yielder session_id stamp. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/interactivity.py | 7 +++++-- secator/ai/session.py | 2 +- secator/tasks/ai.py | 13 +------------ 3 files changed, 7 insertions(+), 15 deletions(-) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 7744ae2e1..98c2d7d5d 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -142,7 +142,10 @@ def _poll_for_answer(self, session_id, prompt_type): results = self.query_engine.search({ "_type": "ai", "ai_type": prompt_type, - "session_id": session_id, + # 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, "status": "answered" }, limit=1) if results: @@ -151,7 +154,7 @@ def _poll_for_answer(self, session_id, prompt_type): elapsed += self.poll_interval # Timeout: update finding status self.query_engine.update( - {"_type": "ai", "ai_type": prompt_type, "session_id": session_id, "status": "pending"}, + {"_type": "ai", "ai_type": prompt_type, "_context.session_id": session_id, "status": "pending"}, {"$set": {"status": "timed_out"}} ) return None diff --git a/secator/ai/session.py b/secator/ai/session.py index ea1ffae48..3af15fe63 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -222,7 +222,7 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None history.set_system(maybe_encrypt(system_prompt, encryptor)) try: - docs = query_engine.search({'_type': 'ai', 'session_id': session_id}) + 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 diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 6592af41e..05c161595 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -119,17 +119,6 @@ def requires_local_execution(cls, inputs, run_opts): # ------------------------------------------------------------------------- def yielder(self) -> Generator: - """Stamp every Ai item with the session_id so the remote-channel transcript - is queryable by session_id. The web UI correlates the whole conversation - (across respawns) by session_id, so a message item without it is invisible. - _init_options() sets self.session_id before the first yield, so the stamp - is always valid here.""" - for _item in self._yielder(): - if isinstance(_item, Ai) and not getattr(_item, "session_id", ""): - _item.session_id = self.session_id - yield _item - - def _yielder(self) -> Generator: """Execute AI task.""" # Addon / setup check if self.inputs == ['setup']: @@ -264,7 +253,7 @@ def _maybe_resume_remote(self): # Look for prior `_type:"ai"` docs for this session try: - prior = query_engine.search({"_type": "ai", "session_id": self.session_id}, limit=1) + 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 From 01cff02a14c4c6f48919b2d32d9c86c03de70c5f Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 15:16:52 +0200 Subject: [PATCH 05/18] fix(ai): make remote follow-up doc renderable (status=pending + top-level choices) In the web AI chat, when the worker hit a follow_up the persisted `_type:"ai"` doc had `status:""` and empty top-level `choices`, so the UI (which gates on `status=="pending"` and reads `m.choices`) stayed stuck on "thinking" with no question/buttons. Two root causes: 1. `_handle_follow_up` (ai/actions.py) stored choices ONLY in `extra_data["choices"]`, never on the top-level `Ai.choices` field the UI reads -> persisted `choices: []`. Now populate both. 2. `_dispatch_and_collect` (tasks/ai.py) persisted the follow_up Ai via `add_result()` (status="") BEFORE the main loop mutated it to `status="pending"`. Since `add_result` dedupes by `_uuid`, the later re-yield could never re-persist the pending state. Now, for a RemoteBackend run, stamp `status="pending"` + top-level `choices` + `session_id` on the single Ai BEFORE the one `add_result`, so the one persisted doc is renderable. The redundant re-stamp/yield in the main loop is removed. Local/CLI follow-up is untouched (remote-only branch). No secator-ui change needed: the doc now carries top-level `choices` and `status=="pending"`. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/actions.py | 5 ++- secator/tasks/ai.py | 26 ++++++++---- tests/unit/test_ai_actions.py | 3 ++ tests/unit/test_ai_loop.py | 76 +++++++++++++++++++++++++++++++++++ 4 files changed, 102 insertions(+), 8 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 63c010d2c..7c880eaac 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -444,7 +444,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: diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 05c161595..242a900e5 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -440,11 +440,11 @@ def _run_loop(self) -> Generator: # 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 []) if result is None: @@ -811,12 +811,24 @@ 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) + 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"): diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 05ff6c99b..1d1122a6b 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -119,6 +119,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') diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index 39e179cfd..bbc3a0cd1 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.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 # ============================================================================= From 6665fd8cd617204a26706048ddbfd8c305fea115 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 17:11:43 +0200 Subject: [PATCH 06/18] fix(ai): persist sub-runner results to workspace + emit runner id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ai task dispatches task/workflow sub-runners in-process and runs them synchronously. The runner framework only re-registers driver hooks (mongodb/api) from context['drivers'] on the pickle path (__setstate__, used by Celery workers) — a sync sub-runner never hits that path. So the sub-runner inherited the ai task's workspace_id/drivers in its context but registered no driver hooks: its update_runner/update_finding hooks never fired, its runner doc + findings were never persisted, and the sub-runs were absent from the workspace History. Build the hooks dict from context['drivers'] (mirroring the CLI entrypoint in cli_helper) and pass hooks= to each dispatched sub-runner, so its results are workspace-scoped and appear in History exactly like a normal runner. Also emit the created runner's id on the action Ai item (extra_data.runner_id + extra_data.runner_type) so the UI can link the action to a RunnerCard. The Ai item is now emitted after the runner is constructed (its on_init hook stamps the id into context), and is emitted even in batch/silent mode so the action doc is always persisted. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/actions.py | 79 +++++++++++++++++++++++++++++++++-- tests/unit/test_ai_actions.py | 75 ++++++++++++++++++++++++++++++++- 2 files changed, 149 insertions(+), 5 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 7c880eaac..aa46f6009 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. @@ -292,9 +339,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 +359,38 @@ 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. + runner_id = runner.id or context.get(f"{runner_type}_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 diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 1d1122a6b..4449f8c47 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 ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -319,6 +320,78 @@ 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') + + +@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): From 7a447aa7611eb52ba7c5830a63ce6a05d3da2db6 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 17:35:35 +0200 Subject: [PATCH 07/18] feat(ai): stamp created finding on add_finding action item (extra_data.finding) So the web UI can render the finding's FindingCard (VulnerabilityCard/etc.) for an add_finding action. The finding is serialized (toDict, includes _type for routing). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/actions.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index aa46f6009..7de178137 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -593,6 +593,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 From eeea43605a86fe1ae045a2a7222d0b0f54ae710e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 18:36:45 +0200 Subject: [PATCH 08/18] fix(ai): coerce add_finding scalars to declared field types before validation LLMs frequently emit wrong-typed scalars in add_finding (a bool field as the string "true", an int as "3"), which validate_fields then rejected, dropping the finding. Add _coerce_finding_fields(cls, data), called before validate_fields, that fixes obvious type mismatches (bool/int/float/list) while leaving valid values, unknown keys, and unparseable values untouched so real errors still surface. Field type resolution is robust to both actual-type and string annotations (from __future__ import annotations), mirroring validate_fields. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/actions.py | 104 ++++++++++++++++++++++++++++++++++ tests/unit/test_ai_actions.py | 73 +++++++++++++++++++++++- 2 files changed, 176 insertions(+), 1 deletion(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 7de178137..a9c2c2b10 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -528,6 +528,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. @@ -581,6 +681,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: diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 4449f8c47..1b351fdfb 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -10,7 +10,7 @@ from secator.ai.actions import ( ActionContext, dispatch_action, _handle_follow_up, _handle_shell, _handle_query, _handle_add_finding, _run_runner, _decrypt_dict, - _build_hooks_from_context + _build_hooks_from_context, _coerce_finding_fields ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -723,6 +723,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): From 6541e6c388de18f9beef350762deccdf30f293e3 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 18:38:35 +0200 Subject: [PATCH 09/18] fix(ai): stamp persisted runner id ({type}_id) on action item, not runner.id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The UI's getRunner queries the persisted runner doc by its _id, which equals context.{type}_id (stamped by the on_init mongodb hook) — not runner.id (secator's internal id). So the RunnerCard showed "Runner not found" for ai-dispatched sub-runners even though they appear in History. Prefer the context id. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/actions.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index a9c2c2b10..e503957ca 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -378,7 +378,11 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator # 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. - runner_id = runner.id or context.get(f"{runner_type}_id", "") + # 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, From 2ba00c6b7c480cff0b5b06292c6ee0b2aeedbdd1 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 18:52:22 +0200 Subject: [PATCH 10/18] fix(ai): scope remote follow_up poll to its own prompt to stop respawn loop RemoteBackend._poll_for_answer matched ANY answered follow_up doc in the session ({_type:"ai", ai_type:"follow_up", _context.session_id, status: "answered"}, limit:1, no sort). Across a multi-turn chat, previously answered follow_up docs accumulate, so the poll for a NEW follow_up immediately matched a STALE answered doc from a prior turn and returned its old answer. The loop then set that old answer as self.prompt, re-yielded Ai(ai_type="prompt") (the original prompt reappears), re-ran the whole turn, asked the follow_up again, re-matched the same stale doc -> an infinite respawn that re-runs scans and burns tokens. (On the very first turn with no prior answered docs it instead timed out cleanly, masking the deeper stale-match bug.) Fix: correlate the poll AND the timeout update to the SPECIFIC pending doc the worker is blocked on. A unique prompt_uuid is stamped into the pending follow_up's extra_data before persist and threaded _dispatch_and_collect -> _run_loop -> _prompt_and_redetect -> ask_user -> _poll_for_answer, which now filters on extra_data.prompt_uuid. A timeout flips only that doc to timed_out. The turn ends cleanly and nothing re-dispatches until the user explicitly sends a new message. The secator-ui AiChatPanel side was investigated and is clean: spawn() is only called from the explicit user send(); there is no watch/effect that re-spawns on done/timed_out. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/interactivity.py | 44 ++++++++++++++++++++--------- secator/tasks/ai.py | 28 ++++++++++++++++-- tests/unit/test_ai_interactivity.py | 40 ++++++++++++++++++++++++++ 3 files changed, 95 insertions(+), 17 deletions(-) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 98c2d7d5d..0610251b2 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,26 +135,42 @@ 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_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). + """ + 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, - # 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, - "status": "answered" - }, limit=1) + results = self.query_engine.search({**base, "status": "answered"}, limit=1) if results: return results[0].get("answer") 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, "_context.session_id": session_id, "status": "pending"}, + {**base, "status": "pending"}, {"$set": {"status": "timed_out"}} ) return None diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 242a900e5..3d7423426 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 @@ -408,6 +409,7 @@ def _run_loop(self) -> Generator: 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) @@ -424,6 +426,7 @@ def _run_loop(self) -> Generator: 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.") @@ -446,7 +449,7 @@ def _run_loop(self) -> Generator: # 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: self._save_history() return @@ -798,6 +801,7 @@ 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) @@ -825,6 +829,14 @@ def _dispatch_and_collect(self, actions, ctx): 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 self.add_result(result, print=not is_from_subagent) @@ -869,7 +881,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 @@ -897,12 +914,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( @@ -915,6 +936,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_interactivity.py b/tests/unit/test_ai_interactivity.py index 6060c5f16..8a1117ce5 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -100,6 +100,46 @@ 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 From f02bef9736e3e99b21d709f149cbabf2211e165d Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 19:15:59 +0200 Subject: [PATCH 11/18] feat(ai): stamp conversation session_id onto AI-spawned sub-runners AI-spawned sub-runners (task/workflow/scan) need context.session_id set so their persisted runner docs are queryable by conversation. The ai task's session_id is often derived (from session_name / the runner id) and is not guaranteed to live in self.context, so sub-runners did NOT carry it. Stamp it in _get_result_context from ActionContext.session_id (without overwriting an existing one). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/actions.py | 17 +++++++++-- tests/unit/test_ai_actions.py | 54 +++++++++++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index e503957ca..03a117edd 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -404,15 +404,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: diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 1b351fdfb..735867076 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -355,6 +355,60 @@ def test_run_runner_propagates_hooks_and_emits_runner_id(self, mock_build_hooks, 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): From 8c731a3a2c5f8d5e5ff29ba21e861a945f2c1c32 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 18:44:08 +0200 Subject: [PATCH 12/18] fix(ai): keep the AI loop alive when an action dispatch raises MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A Python error during an iteration (e.g. TypeError: 'str' object is not a mapping from a malformed LLM action/opts) previously propagated out of _dispatch_and_collect, was caught by the loop's broad except Exception, and killed the task. Now each action's dispatch is wrapped so the failure becomes that tool call's result fed back to the LLM, and the loop continues. - Add safe_dispatch_action(): wraps dispatch_action and, on Exception, yields an Error carrying the action's tool_call_id/tool_call_name in _context. Only Exception is caught — KeyboardInterrupt/SystemExit/GeneratorExit propagate. - The Error groups into a tool result via the existing format_tool_result / add_tool_result path, so the model sees "Action failed with error: : \n. Fix the issue and try again." next turn. - Use safe_dispatch_action for the single-action path in _dispatch_and_collect and inside _run_batch's run_single, so one action's failure no longer aborts the turn or the other batch actions. - max_iterations still bounds a persistently-erroring model: each failed turn increments the iteration counter as before. - Drop a pre-existing unused follow_up_ai assignment to keep flake8 green. - Tests: a raising handler yields an Error, appends the error to history (LLM-visible), and continues without raising; KeyboardInterrupt propagates. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/actions.py | 60 ++++++++++++++++++- secator/tasks/ai.py | 11 ++-- tests/unit/test_ai_loop.py | 115 ++++++++++++++++++++++++++++++++++++- 3 files changed, 180 insertions(+), 6 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 03a117edd..110e935f0 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -306,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. @@ -789,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/tasks/ai.py b/secator/tasks/ai.py index 3d7423426..cc1e9f74b 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -16,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 @@ -408,7 +408,6 @@ 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: @@ -425,7 +424,6 @@ 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: @@ -804,7 +802,12 @@ def _dispatch_and_collect(self, actions, ctx): 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: diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index bbc3a0cd1..b628bc1f6 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -355,7 +355,7 @@ def add_tool_result(self, *a, **k): def _fake_dispatch_action(action, c): yield follow_up - with patch("secator.tasks.ai.dispatch_action", _fake_dispatch_action): + 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 @@ -1015,5 +1015,118 @@ 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) + + if __name__ == "__main__": unittest.main() From 5968c95ffb4225cbf3b8ea75b35d1acc4cca23fd Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:03:58 +0200 Subject: [PATCH 13/18] feat(ai): mid-flight steering (interrupt + redirect) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add cooperative mid-flight steering to the Workspace AI Assistant: a user can send a message WHILE the agent is running, and the worker picks it up at the next loop checkpoint to redirect the next turn. Distinct from the hard Stop button (which revokes the Celery task). - RemoteBackend.poll_steers(session_id): drains pending `ai_type:"steer"` channel docs, returns their content oldest-first, marks them consumed so each injects exactly once. Robust — backend errors return [] (never crash). - _poll_for_answer: a steer breaks a blocked follow-up wait (returns the steer content as the answer) so the loop redirects instead of stalling; follow-up semantics intact for the no-steer case. - _run_loop: _drain_steers() at the top of each iteration appends each steer to history as `[User interjected]: …` and echoes a steer Ai item (with session_id so it persists in the transcript). - output_types/ai.py: render `steer` ai_type in the CLI transcript. - Tests: poll_steers drain/consume/robustness, steer-breaks-wait, _drain_steers inject-into-history, no-steer no-op, non-remote no-op. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/interactivity.py | 58 +++++++++++++++++++ secator/output_types/ai.py | 1 + secator/tasks/ai.py | 38 +++++++++++++ tests/unit/test_ai_interactivity.py | 86 +++++++++++++++++++++++++++-- tests/unit/test_ai_loop.py | 85 ++++++++++++++++++++++++++++ 5 files changed, 264 insertions(+), 4 deletions(-) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 0610251b2..78e314145 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -135,6 +135,53 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con # follow_up: return the answer text return {"answer": answer} + 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. @@ -148,6 +195,12 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): 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", @@ -165,6 +218,11 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): 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: flip ONLY this prompt's still-pending doc to timed_out, so a 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 cc1e9f74b..46b410297 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -342,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() @@ -658,6 +663,39 @@ 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). At the + top of each loop iteration we drain any pending steers for this session, + append each to the history as a ``[User interjected]: …`` user message so + the model sees them on the next turn, and echo a steer Ai item (with + ``_context`` so it persists in the transcript). Cooperative — not a hard + cancel (Stop already does that). + + 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. + """ + 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)) + # Echo into the transcript (persisted via _context.session_id) so the + # UI shows the steer as an interjected user bubble. + yield Ai(content=content, ai_type="steer", session_id=self.session_id) + # ------------------------------------------------------------------------- # Summarization / compaction # ------------------------------------------------------------------------- diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 8a1117ce5..bec54f3be 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -144,10 +144,18 @@ def test_timeout_update_scoped_to_prompt_uuid(self, mock_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") @@ -157,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 b628bc1f6..783d8e422 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -1128,5 +1128,90 @@ def _boom(*a, **k): 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") + + # Echoed as a steer Ai item carrying the session_id (so it persists). + steer_items = [r for r in yielded if isinstance(r, Ai) and r.ai_type == "steer"] + self.assertEqual(len(steer_items), 1) + self.assertEqual(steer_items[0].content, "actually focus on the API") + self.assertEqual(steer_items[0].session_id, "steer-sess") + + # 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() From 56ec8bcec1bed48606fdf8315230f209ae4b4185 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:09:25 +0200 Subject: [PATCH 14/18] refactor(ai): steer doc is the transcript entry; restore steers on respawn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the worker's redundant `Ai(ai_type="steer")` echo: the API's pending steer doc already carries `_context.session_id` and is itself the persisted transcript entry, so a second echo would double-render in the UI. Keep `_drain_steers` a generator (no items yielded) so the loop call site is unchanged and future echoes can be added without churn. Also restore steers as user turns in `restore_history_from_db` (framed `[User interjected]: …`) so a mid-flight redirect survives a respawn/history restore. Update the drain test to assert no echo doc is yielded. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/session.py | 5 +++++ secator/tasks/ai.py | 27 ++++++++++++++++++--------- tests/unit/test_ai_loop.py | 8 +++----- 3 files changed, 26 insertions(+), 14 deletions(-) diff --git a/secator/ai/session.py b/secator/ai/session.py index 3af15fe63..a6e803404 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -237,6 +237,11 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None 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. diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 46b410297..57d04f654 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -671,16 +671,26 @@ 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). At the - top of each loop iteration we drain any pending steers for this session, - append each to the history as a ``[User interjected]: …`` user message so - the model sees them on the next turn, and echo a steer Ai item (with - ``_context`` so it persists in the transcript). Cooperative — not a hard - cancel (Stop already does that). + 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 @@ -692,9 +702,8 @@ def _drain_steers(self): 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)) - # Echo into the transcript (persisted via _context.session_id) so the - # UI shows the steer as an interjected user bubble. - yield Ai(content=content, ai_type="steer", session_id=self.session_id) + return + yield # noqa: unreachable - keeps this a generator for `yield from` # ------------------------------------------------------------------------- # Summarization / compaction diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index 783d8e422..2f89ad40c 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -1170,11 +1170,9 @@ def test_steer_drained_injected_and_consumed(self): self.assertEqual(len(user_msgs), 1) self.assertEqual(user_msgs[-1]["content"], "[User interjected]: actually focus on the API") - # Echoed as a steer Ai item carrying the session_id (so it persists). - steer_items = [r for r in yielded if isinstance(r, Ai) and r.ai_type == "steer"] - self.assertEqual(len(steer_items), 1) - self.assertEqual(steer_items[0].content, "actually focus on the API") - self.assertEqual(steer_items[0].session_id, "steer-sess") + # 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] From 00d3d2ae69d73ad44442db2c8c833b2fb44dcc49 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 01:11:02 +0200 Subject: [PATCH 15/18] refactor(ai): auto-register driver hooks from context (drop _build_hooks_from_context) Runner.__init__ now auto-registers driver hooks from context['drivers'] (via _apply_context_drivers, added on main), so the ai task's manual hook-building for sub-runners is redundant. Delete the helper and the hooks= kwarg it fed into Task/Workflow construction. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 56 ++------------------------------ tests/unit/test_ai_actions.py | 60 +++++------------------------------ 2 files changed, 10 insertions(+), 106 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 110e935f0..16e2eed5c 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -70,53 +70,6 @@ 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. @@ -414,14 +367,9 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator 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) + # Driver hooks (mongodb/api) auto-register from context['drivers'] in Runner.__init__. try: - runner = runner_cls(tpl, targets, run_opts=run_opts, hooks=hooks, context=context) + runner = runner_cls(tpl, targets, run_opts=run_opts, context=context) except TaskNotFoundError as e: yield Error(message=str(e), _context=context) return diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 735867076..365687df8 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -10,7 +10,7 @@ from secator.ai.actions import ( ActionContext, dispatch_action, _handle_follow_up, _handle_shell, _handle_query, _handle_add_finding, _run_runner, _decrypt_dict, - _build_hooks_from_context, _coerce_finding_fields + _coerce_finding_fields ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -322,13 +322,10 @@ def test_run_runner_uses_ctx_targets_as_default(self): @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 + def test_run_runner_no_manual_hooks_and_emits_runner_id(self, mock_task_cls, _mock_tpl): + """Sub-runner must NOT receive a manual hooks= kwarg (driver hooks now + auto-register from context['drivers'] in Runner.__init__), 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' @@ -344,9 +341,9 @@ def test_run_runner_propagates_hooks_and_emits_runner_id(self, mock_build_hooks, results = list(_run_runner(action, ctx, 'task')) - # Runner constructed with hooks= from the context drivers + # Runner constructed without a hooks= kwarg (auto-registered from context) _, kwargs = mock_task_cls.call_args - self.assertEqual(kwargs.get('hooks'), sentinel_hooks) + self.assertNotIn('hooks', kwargs) self.assertEqual(kwargs.get('context', {}).get('workspace_id'), 'ws1') # Action Ai item carries runner_id + runner_type @@ -357,13 +354,11 @@ def test_run_runner_propagates_hooks_and_emits_runner_id(self, mock_build_hooks, @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): + def test_run_runner_propagates_session_id(self, 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 @@ -387,10 +382,8 @@ def test_run_runner_propagates_session_id(self, mock_build_hooks, mock_task_cls, @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): + def test_run_runner_preserves_existing_session_id(self, 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 @@ -410,43 +403,6 @@ def test_run_runner_preserves_existing_session_id(self, mock_build_hooks, mock_t 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): """Tests for ActionContext.get_query_engine caching and backend selection.""" From bb7fb71d56bc14618d239d9956ba3228b148f154 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 01:11:38 +0200 Subject: [PATCH 16/18] refactor(ai): use Error.from_exception for action-dispatch errors Replace the hand-rolled _format_action_error formatter with the framework's Error.from_exception, already used elsewhere in this file, so action-dispatch failures build their LLM-facing message the same way as every other error path. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 34 +--------------------------------- tests/unit/test_ai_loop.py | 4 +--- 2 files changed, 2 insertions(+), 36 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 16e2eed5c..dcb03e5a2 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -259,33 +259,6 @@ 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. @@ -301,16 +274,11 @@ def safe_dispatch_action(action: Dict, ctx: ActionContext) -> Generator: ``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, - ) + yield Error.from_exception(e, _context=context) def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index 2f89ad40c..15d847183 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -1048,11 +1048,9 @@ def _boom(*a, **k): 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) + # LLM-facing feedback carries the exception type/message (Error.from_exception). 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") From 0539f24e4a25614623bb97bf653d23190689f130 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 01:12:37 +0200 Subject: [PATCH 17/18] refactor(output_types): share field-type resolution via OutputType.field_types Extract the field-name -> concrete-type resolution duplicated between OutputType.validate_fields and actions._resolve_field_type onto a single OutputType.field_types() classmethod, reused by both validate_fields and _coerce_finding_fields. Behavior unchanged. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 28 +------------------------- secator/output_types/_base.py | 38 ++++++++++++++++++++++++++++++----- 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index dcb03e5a2..f3b141860 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -513,32 +513,6 @@ 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. @@ -550,7 +524,7 @@ def _coerce_finding_fields(cls, data: Dict) -> Dict: 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)} + field_types = cls.field_types() for key, value in list(data.items()): if key.startswith('_'): continue diff --git a/secator/output_types/_base.py b/secator/output_types/_base.py index 0a7c08620..a14c77e48 100644 --- a/secator/output_types/_base.py +++ b/secator/output_types/_base.py @@ -147,6 +147,37 @@ def toDict(self, exclude=[]): return {k: v for k, v in data.items() if k not in exclude} return data + @classmethod + def field_types(cls) -> dict: + """Resolve each non-underscore field's declared type to a concrete builtin type. + + ``f.type`` may be an actual type (``bool``), a typing generic (``List[str]``, + resolved via ``__origin__``), or — under ``from __future__ import annotations`` + — a string annotation (``'bool'``, ``'List[str]'``). Returns + ``{field_name: concrete_type}``, omitting fields that can't be resolved. + """ + type_map = { + 'bool': bool, 'int': int, 'float': float, + 'str': str, 'list': list, 'dict': dict, + } + resolved = {} + for f in fields(cls): + if f.name.startswith('_'): + continue + t = f.type + if isinstance(t, type): + resolved[f.name] = t + continue + origin = getattr(t, '__origin__', None) + if origin is not None: + resolved[f.name] = origin + continue + if isinstance(t, str): + name = t.split('[', 1)[0].strip().lower() + if name in type_map: + resolved[f.name] = type_map[name] + return resolved + @classmethod def validate_fields(cls, data: dict) -> list: """Validate data types against dataclass field definitions. @@ -155,17 +186,14 @@ def validate_fields(cls, data: dict) -> list: """ errors = [] type_names = {str: 'str', int: 'int', float: 'float', dict: 'dict', list: 'list', bool: 'bool'} + expected_types = cls.field_types() for f in fields(cls): if f.name.startswith('_') or f.name not in data: continue value = data[f.name] if value is None: continue - expected_type = f.type if isinstance(f.type, type) else None - if expected_type is None: - origin = getattr(f.type, '__origin__', None) - if origin is not None: - expected_type = origin + expected_type = expected_types.get(f.name) if expected_type and not isinstance(value, expected_type): expected_name = type_names.get(expected_type, getattr(expected_type, '__name__', str(expected_type))) actual_name = type_names.get(type(value), type(value).__name__) From d7202f65821cd3b40dd9e61228621e8eafd6e406 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 01:14:51 +0200 Subject: [PATCH 18/18] style(ai): trim verbose core comments per review Delete comments that restate the obvious or describe UI behavior; reduce the rest to one terse line capturing the non-obvious why, per reviewer request. No logic changes. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 33 ++------------------------------ secator/ai/interactivity.py | 38 +++++-------------------------------- secator/tasks/ai.py | 24 +++-------------------- 3 files changed, 10 insertions(+), 85 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index f3b141860..de920a1fb 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -342,16 +342,7 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator 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". + # Prefer the persisted doc id ({type}_id from on_init) over runner.id. runner_id = context.get(f"{runner_type}_id", "") or runner.id yield Ai( content=name, @@ -374,16 +365,7 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator def _get_result_context(action, ctx): - """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. - """ + """Derive a sub-runner result context, stamping the conversation 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 @@ -500,9 +482,6 @@ def _handle_follow_up(action: Dict, ctx: ActionContext) -> Generator: context = _get_result_context(action, ctx) reason = action.get("reason", "completed") choices = action.get("choices", []) - # 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) @@ -640,8 +619,6 @@ 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 @@ -656,8 +633,6 @@ 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 ) @@ -733,10 +708,6 @@ 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 safe_dispatch_action(act, batch_ctx): if isinstance(item, Ai) and item.ai_type == "token_usage": diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 78e314145..77c2b09f4 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -136,26 +136,14 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con return {"answer": answer} 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. + """Drain pending steer docs for ``session_id`` and mark them consumed + (oldest-first). Any backend error returns ``[]`` — a steer must 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", } @@ -185,29 +173,13 @@ def poll_steers(self, session_id): 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. + Scoped by ``prompt_uuid`` (not just session_id + status:"answered"), else a + multi-turn conversation matches a stale answered doc and respawn-loops. A + pending steer also breaks the wait early and is returned as the answer. """ 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: diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 57d04f654..a4ffb865e 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -342,9 +342,7 @@ 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). + # Drain mid-flight steer messages and inject them so the next turn redirects. yield from self._drain_steers() # Auto-summarize when context > 85% threshold @@ -553,13 +551,7 @@ def _init_options(self): workspace=self.reports_folder or "" ) - # 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. + # Remote channel: UI sends a stable session_id on self.context; local uses self.id. self.session_id = ( self.passed_context.get("session_id") or (self.context or {}).get("session_id") @@ -849,11 +841,6 @@ def _dispatch_and_collect(self, actions, ctx): follow_up_prompt_uuid = None is_batch = len(actions) > 1 - # 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 = [] @@ -868,12 +855,7 @@ def _dispatch_and_collect(self, actions, ctx): 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. + # Persist the follow-up doc once, in its final renderable state (add_result dedupes by _uuid). if isinstance(self.backend, RemoteBackend): follow_up_ai.status = "pending" follow_up_ai.session_id = self.session_id