feat(ai): mid-flight steering (interrupt + redirect)

[ocervell]

Mid-flight steering (worker side)

Add cooperative mid-flight steering to the Workspace AI Assistant: a user can send a message WHILE the agent runs and it picks it up at the next loop checkpoint to redirect the next turn — distinct from the hard Stop (which revokes the Celery task).

Changes

• secator/ai/interactivity.py — RemoteBackend.poll_steers(session_id): drains pending _type:"ai", ai_type:"steer" channel docs for this session, returns their content oldest-first, flips them to status:"consumed" (inject exactly once). Robust — backend errors return [] so a steer can never crash the run. _poll_for_answer now breaks a blocked follow-up wait when a steer arrives (returns the steer content as the answer) so the loop redirects instead of stalling; no-steer follow-up semantics intact.
• secator/tasks/ai.py — _drain_steers() called at the top of each _run_loop iteration: appends each steer to history as [User interjected]: <content>. No echo doc is emitted (the API's pending steer doc is itself the persisted transcript entry — avoids a double-render in the UI).
• secator/ai/session.py — restore_history_from_db reconstructs steers as user turns so a redirect survives a respawn.
• secator/output_types/ai.py — render a steer ai_type.

Channel contract (steer doc)

{ "_type": "ai", "ai_type": "steer", "status": "pending",
  "content": "<user message>", "session_id": "<sid>",
  "_context": { "session_id": "<sid>", "workspace_id": "<wsid>" },
  "_timestamp": <float> }

Worker drains pending steers → injects → flips to status:"consumed".

Tests

poll_steers drain/consume/robustness + no-query-engine no-op, steer-breaks-blocked-wait, _drain_steers inject-into-history + no-echo + no-steer no-op + non-remote no-op + poll-error-never-crashes. AI unit suites + flake8 (added lines) green.

Cross-repo PRs

• secator-api: https://github.com/freelabz/secator-api/pull/203
• secator-ui: https://github.com/freelabz/secator-ui/pull/271

🤖 Generated with Claude Code

https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm

Summary by CodeRabbit

• New Features

  • Added remote session resume and history restoration for AI conversations.
  • Introduced support for mid-run “steer” messages and improved follow-up handling.

• Bug Fixes

  • Made AI action handling more resilient so individual failures no longer stop the whole run.
  • Improved persistence and scoping so responses, follow-ups, and steers map to the correct session and prompt.
  • Better handles finding field type mismatches and preserves more results for display.

[coderabbitai]

Walkthrough

AI action dispatch now preserves sub-runner hooks and runner metadata, remote sessions can restore history from Mongo with prompt-scoped polling, and the task loop persists history while correlating follow-ups and draining steers.

Changes

AI runtime and remote session flow

Layer / File(s)	Summary
Action dispatch and runner metadata secator/ai/actions.py, tests/unit/test_ai_actions.py	Sub-runner hooks are built from context, action failures become Error outputs, and runner results carry runner_id, runner_type, and tool-call context.
Finding field coercion secator/ai/actions.py, tests/unit/test_ai_actions.py	AI finding values are coerced to declared field types before validation, and created findings are attached to the emitted add-finding item.
Transcript storage and prompt scoping secator/output_types/ai.py, secator/ai/session.py, secator/ai/interactivity.py, tests/unit/test_ai_interactivity.py, tests/unit/test_ai_session.py	steer is added to AI output rendering, history restoration rebuilds chat turns from stored AI docs, and remote answer polling is scoped by prompt_uuid and session.
Remote resume and history persistence secator/tasks/ai.py, tests/unit/test_ai_session.py	Remote sessions can resume from stored history, derive session IDs from runner context, and save or skip history on the remote task loop’s terminal paths.
Follow-up UUIDs and steer handling secator/ai/actions.py, secator/tasks/ai.py, tests/unit/test_ai_loop.py	Follow-up prompts carry UUID correlation, remote follow-up docs persist as pending, steer docs are drained into history, and single-action dispatch uses safe_dispatch_action.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• freelabz/secator#949: Shares the AI dispatch and remote interactivity code paths that are extended by this change.
• freelabz/secator#940: Introduced the AI task modules and runner flow that this PR updates.
• freelabz/secator#1163: Added hook registry support used by _build_hooks_from_context().

Suggested labels

ai

Poem

I hopped through prompts at silver night,
and tucked each steer in ferny light.
The runner hummed, the session stayed,
with UUID trails neatly laid.
🐇✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 56.32% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding mid-flight steering to AI runs.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/ai-steer

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 7

🧹 Nitpick comments (2)

tests/unit/test_ai_actions.py (1)

352-356: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Cover the persisted runner-id path, not only the fallback.

Production now prefers context["task_id"] / context["workflow_id"] over runner.id, but this test only sets mock_runner.id. It would still pass if the context-stamped persisted id were ignored.

Example test adjustment

-		# Fake runner: an iterable whose id is populated (mimics on_init stamping it)
+		# Fake runner: constructor mutates the passed context like on_init hooks do.
 		mock_runner = MagicMock()
-		mock_runner.id = 'runner123'
+		mock_runner.id = 'internal-runner-id'
 		mock_runner.reports_folder = None
 		mock_runner.__iter__.return_value = iter([])
-		mock_task_cls.return_value = mock_runner
+
+		def _make_runner(*_args, **kwargs):
+			kwargs['context']['task_id'] = 'persisted-doc-id'
+			return mock_runner
+
+		mock_task_cls.side_effect = _make_runner
@@
-		self.assertEqual(ai_items[0].extra_data.get('runner_id'), 'runner123')
+		self.assertEqual(ai_items[0].extra_data.get('runner_id'), 'persisted-doc-id')

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/test_ai_actions.py` around lines 352 - 356, The test currently
only verifies the fallback runner.id path, so update the Ai action assertion in
test_ai_actions to cover the persisted runner-id branch by setting
context["task_id"] or context["workflow_id"] and asserting the Ai.extra_data
values come from that context-backed id instead of mock_runner.id. Keep the
existing Ai/task lookup, but change the test setup to exercise the
production-preferred path and validate runner_id/runner_type using the relevant
Ai item.

tests/unit/test_ai_loop.py (1)

363-375: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assert that remote follow-ups persist prompt_uuid.

This test verifies the UI fields but not the new correlation field that prevents stale answered follow-ups from resolving future waits. Add an assertion on doc.extra_data["prompt_uuid"].

Proposed test assertion

 		self.assertEqual(doc.status, "pending")
 		self.assertEqual(doc.choices, ["Fuzz parameters", "Run nuclei", "Deep crawl"])
 		self.assertEqual(doc.session_id, "sess-123")
+		self.assertTrue(doc.extra_data.get("prompt_uuid"))

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/test_ai_loop.py` around lines 363 - 375, The remote follow-up
persistence test is missing coverage for the new correlation field. In
test_remote_follow_up_persisted_pending_with_choices, add an assertion on the
persisted Ai follow-up’s extra_data["prompt_uuid"] alongside the existing
status, choices, and session_id checks, using the doc/follow_up objects already
created by _run_dispatch to verify the prompt UUID is stored and preserved.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@secator/ai/actions.py`:
- Around line 97-107: The driver bootstrap in action loading is performing
external discovery twice per sub-runner. Update the driver setup in
secator.ai.actions so the path that calls get_available_drivers() is the only
place that triggers discovery, and remove the extra discover_external_drivers()
call before order_drivers() while keeping the existing driver ordering and
supported-driver filtering logic intact.
- Around line 651-652: The boolean coercion in the AI data normalization path is
too broad and is converting any integer into a boolean, which can mask invalid
values. Update the logic in the data handling flow around the
`value`/`data[key]` assignment so only `0` and `1` are coerced to booleans, and
leave all other integers unchanged for validation to catch; use the existing
`isinstance(value, int)` check in `secator/ai/actions.py` as the location to
narrow the conversion.

In `@secator/ai/interactivity.py`:
- Around line 221-225: The steer handling in interactivity should not return
consumed steers as plain text from the follow-up path, because downstream code
will misclassify them as normal answers. Update the logic around the poll_steers
branch in the interactivity flow so it preserves steer type, either by returning
a typed result or by limiting steer-breaking to follow-up handling only. Then
ensure the caller injects the steer using the same framing used by _drain_steers
so permission handling and prompt persistence do not treat it as an ordinary
reply.
- Around line 175-180: The session steer consumption update is too broad and can
mark unseen docs as consumed. In the interactivity flow around the
drain-and-consume logic, scope the query_engine.update call to only the exact
documents returned by the earlier drain step using stable identifiers such as
_uuid or _id, rather than the shared base session filter. Keep the update tied
to the drained set so only steer docs actually injected are marked consumed.

In `@tests/unit/test_ai_loop.py`:
- Line 365: The tuple unpacking in the test helpers is creating unused locals,
so update the assignments in the affected test cases to use underscore-prefixed
names for any values that are not referenced afterward. Make this change in the
test code around the _run_dispatch usage and the other reported unpacking sites,
keeping only the variables that are actually used and renaming the rest to
_yielded, _follow_up, or _name as appropriate.

In `@tests/unit/test_ai_session.py`:
- Around line 29-30: The tests for the AI session search contract are still
asserting the old session_id field, so update the expectations and mocks in the
ai session tests to use _context.session_id instead. Fix the assertions around
engine.search in the affected session restore/resume tests so they match the
implementation’s query shape {"_type": "ai", "_context.session_id": ...}, and
adjust any related fixtures or stubs in the same test module to exercise the
correct branch.
- Line 126: The unpacked variables from _make_task in the affected test cases
are unused, so rename the unused binding(s) to use a leading underscore in the
test methods that call _make_task, including the occurrences in the cited test
block and the other matching ones in test_ai_session.py. Keep the same tuple
unpacking shape, but ensure only the used value remains named normally so Ruff
no longer reports unused variables.

---

Nitpick comments:
In `@tests/unit/test_ai_actions.py`:
- Around line 352-356: The test currently only verifies the fallback runner.id
path, so update the Ai action assertion in test_ai_actions to cover the
persisted runner-id branch by setting context["task_id"] or
context["workflow_id"] and asserting the Ai.extra_data values come from that
context-backed id instead of mock_runner.id. Keep the existing Ai/task lookup,
but change the test setup to exercise the production-preferred path and validate
runner_id/runner_type using the relevant Ai item.

In `@tests/unit/test_ai_loop.py`:
- Around line 363-375: The remote follow-up persistence test is missing coverage
for the new correlation field. In
test_remote_follow_up_persisted_pending_with_choices, add an assertion on the
persisted Ai follow-up’s extra_data["prompt_uuid"] alongside the existing
status, choices, and session_id checks, using the doc/follow_up objects already
created by _run_dispatch to verify the prompt UUID is stored and preserved.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 7ad0eb55-57c9-4c18-9082-775da6b4279f

📥 Commits

Reviewing files that changed from the base of the PR and between 8d8ec83 and e99b982.

📒 Files selected for processing (9)

• secator/ai/actions.py
• secator/ai/interactivity.py
• secator/ai/session.py
• secator/output_types/ai.py
• secator/tasks/ai.py
• tests/unit/test_ai_actions.py
• tests/unit/test_ai_interactivity.py
• tests/unit/test_ai_loop.py
• tests/unit/test_ai_session.py

[coderabbitai]

🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win

Avoid discovering external drivers twice per sub-runner.

get_available_drivers() already performs external driver discovery, so Line 103 re-runs external hook module loading before doing it again on Line 107. Since this path runs for spawned runners, duplicated driver side effects and extra I/O can accumulate.

Proposed fix

-	from secator.loader import discover_external_drivers, get_available_drivers, order_drivers
+	from secator.loader import get_available_drivers, order_drivers
 	from secator.utils import import_dynamic, deep_merge_dicts
@@
-	discover_external_drivers()
 	# Order by canonical priority so authoritative backends (e.g. mongodb) register

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	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())
	from secator.loader import get_available_drivers, order_drivers
	from secator.utils import import_dynamic, deep_merge_dicts
	drivers = list(context.get('drivers', []))
	if not drivers:
	return {}
	# 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())

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@secator/ai/actions.py` around lines 97 - 107, The driver bootstrap in action
loading is performing external discovery twice per sub-runner. Update the driver
setup in secator.ai.actions so the path that calls get_available_drivers() is
the only place that triggers discovery, and remove the extra
discover_external_drivers() call before order_drivers() while keeping the
existing driver ordering and supported-driver filtering logic intact.

[coderabbitai]

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Only coerce integer booleans for 0 and 1.

bool(2) and bool(-1) both become True, so malformed AI data can incorrectly set finding flags such as verified, is_false_positive, or is_acknowledged. Leave other integers untouched so validation reports them.

Proposed fix

-			if isinstance(value, int):
+			if isinstance(value, int) and value in (0, 1):
 				data[key] = bool(value)

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	if isinstance(value, int):
	data[key] = bool(value)
	if isinstance(value, int) and value in (0, 1):
	data[key] = bool(value)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@secator/ai/actions.py` around lines 651 - 652, The boolean coercion in the AI
data normalization path is too broad and is converting any integer into a
boolean, which can mask invalid values. Update the logic in the data handling
flow around the `value`/`data[key]` assignment so only `0` and `1` are coerced
to booleans, and leave all other integers unchanged for validation to catch; use
the existing `isinstance(value, int)` check in `secator/ai/actions.py` as the
location to narrow the conversion.

[coderabbitai]

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Consume only the steer docs that were actually drained.

This update matches every pending steer for the session, not just the limit=50 docs returned above; any extra or concurrently inserted steer can be marked consumed without being injected. Scope the update by stable returned identifiers such as _uuid/_id.

Suggested direction

-			self.query_engine.update(
-				{**base},
-				{"$set": {"status": "consumed"}},
-			)
+			identifiers = [
+				{"_uuid": doc["_uuid"]} if doc.get("_uuid") else {"_id": doc["_id"]}
+				for doc in results
+				if doc.get("_uuid") or doc.get("_id")
+			]
+			if identifiers:
+				self.query_engine.update(
+					{**base, "$or": identifiers},
+					{"$set": {"status": "consumed"}},
+				)

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	# Mark this session's pending steers consumed so they inject exactly once.
	try:
	self.query_engine.update(
	{**base},
	{"$set": {"status": "consumed"}},
	)
	# Mark this session's pending steers consumed so they inject exactly once.
	try:
	identifiers = [
	{"_uuid": doc["_uuid"]} if doc.get("_uuid") else {"_id": doc["_id"]}
	for doc in results
	if doc.get("_uuid") or doc.get("_id")
	]
	if identifiers:
	self.query_engine.update(
	{**base, "$or": identifiers},
	{"$set": {"status": "consumed"}},
	)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@secator/ai/interactivity.py` around lines 175 - 180, The session steer
consumption update is too broad and can mark unseen docs as consumed. In the
interactivity flow around the drain-and-consume logic, scope the
query_engine.update call to only the exact documents returned by the earlier
drain step using stable identifiers such as _uuid or _id, rather than the shared
base session filter. Keep the update tied to the drained set so only steer docs
actually injected are marked consumed.

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Preserve steer type instead of returning it as a normal answer.

A consumed steer returned here is indistinguishable from a follow-up/permission answer downstream: follow-up handling will persist a second prompt doc, while permission handling will treat arbitrary steer text as denial and drop the user interjection. Return a typed result or restrict steer-breaking to follow-ups, then let the caller inject it with the same steer framing used by _drain_steers.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@secator/ai/interactivity.py` around lines 221 - 225, The steer handling in
interactivity should not return consumed steers as plain text from the follow-up
path, because downstream code will misclassify them as normal answers. Update
the logic around the poll_steers branch in the interactivity flow so it
preserves steer type, either by returning a typed result or by limiting
steer-breaking to follow-up handling only. Then ensure the caller injects the
steer using the same framing used by _drain_steers so permission handling and
prompt persistence do not treat it as an ordinary reply.

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Prefix unused unpacked variables with _.

Ruff reports these locals as unused; rename them to _yielded, _follow_up, or _name as appropriate.

Also applies to: 380-380, 1125-1125

🧰 Tools

🪛 Ruff (0.15.18)

[warning] 365-365: Unpacked variable yielded is never used

Prefix it with an underscore or any other dummy variable pattern

(RUF059)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/test_ai_loop.py` at line 365, The tuple unpacking in the test
helpers is creating unused locals, so update the assignments in the affected
test cases to use underscore-prefixed names for any values that are not
referenced afterward. Make this change in the test code around the _run_dispatch
usage and the other reported unpacking sites, keeping only the variables that
are actually used and renaming the rest to _yielded, _follow_up, or _name as
appropriate.

Source: Linters/SAST tools

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Update the tests to the _context.session_id query contract.

The implementation queries {"_type": "ai", "_context.session_id": ...}, but these tests still assert/mock session_id, so the restore/resume tests won’t exercise the intended branch.

Proposed fix

-		engine.search.assert_called_once_with({"_type": "ai", "session_id": "session1"})
+		engine.search.assert_called_once_with({"_type": "ai", "_context.session_id": "session1"})
...
-		def _search(query, limit=0):
-			if query.get("_type") == "ai" and "session_id" in query:
+		def _search(query, limit=0):
+			if query.get("_type") == "ai" and query.get("_context.session_id") == task.session_id:
 				return prior_docs

Also applies to: 117-121

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/test_ai_session.py` around lines 29 - 30, The tests for the AI
session search contract are still asserting the old session_id field, so update
the expectations and mocks in the ai session tests to use _context.session_id
instead. Fix the assertions around engine.search in the affected session
restore/resume tests so they match the implementation’s query shape {"_type":
"ai", "_context.session_id": ...}, and adjust any related fixtures or stubs in
the same test module to exercise the correct branch.

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Prefix unused unpacked variables with _.

Ruff flags these unpacked variables as unused; renaming them keeps the tests lint-clean without changing behavior.

Also applies to: 144-144, 160-160, 170-170

🧰 Tools

🪛 Ruff (0.15.18)

[warning] 126-126: Unpacked variable engine is never used

Prefix it with an underscore or any other dummy variable pattern

(RUF059)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/test_ai_session.py` at line 126, The unpacked variables from
_make_task in the affected test cases are unused, so rename the unused
binding(s) to use a leading underscore in the test methods that call _make_task,
including the occurrences in the cited test block and the other matching ones in
test_ai_session.py. Keep the same tuple unpacking shape, but ensure only the
used value remains named normally so Ruff no longer reports unused variables.

Source: Linters/SAST tools

[ocervell]

Folded into the AI-chat branch (feat/ai-chat) via clean fast-forward — these commits now ship with the chat PRs (secator #1215 / secator-api #199 / secator-ui #265). Closing as redundant.