feat: headless Mongo session restore for remote AI chat

secator/ai/actions.py

@@ -259,6 +259,28 @@ def dispatch_action(action: Dict, ctx: ActionContext) -> Generator:
 		yield Warning(message=f"Unknown action: {action_type}", _context=context)
 
 
+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.
+	"""
+	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.from_exception(e, _context=context)
+
+
 def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator:
 	"""Execute a secator task or workflow.
 
@@ -292,9 +314,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 +334,28 @@ 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)
+
+	# Driver hooks (mongodb/api) auto-register from context['drivers'] in Runner.__init__.
 	try:
 		runner = runner_cls(tpl, targets, run_opts=run_opts, context=context)
 	except TaskNotFoundError as e:
 		yield Error(message=str(e), _context=context)
 		return
+
+	# 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,
+		ai_type=runner_type,
+		extra_data={
+			"targets": targets,
+			"opts": opts,
+			"runner_id": runner_id,
+			"runner_type": runner_type,
+		},
+		_context=context,
+	)
+
 	yield from runner
 
 	# Auto-allow reading from the spawned runner's reports folder
@@ -329,15 +365,17 @@ 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()
+	"""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
 	action_context = {}
 	tool_call_id = action.get("tool_call_id")
 	tool_call_name = action.get("tool_call_name")
 	if tool_call_id:
 		action_context["tool_call_id"] = tool_call_id
 		action_context["tool_call_name"] = tool_call_name
-	return {**ctx, **action_context}
+	return {**new_ctx, **action_context}
 
 
 def _handle_task(action: Dict, ctx: ActionContext) -> Generator:
@@ -444,7 +482,7 @@ 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)
+	yield Ai(content=reason, ai_type="follow_up", choices=choices, extra_data={"choices": choices}, _context=context)
 
 
 def _handle_stop(action: Dict, ctx: ActionContext) -> Generator:
@@ -454,6 +492,80 @@ def _handle_stop(action: Dict, ctx: ActionContext) -> Generator:
 	yield Ai(content=reason, ai_type="stopped", _context=context)
 
 
+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 = cls.field_types()
+	for key, value in list(data.items()):
+		if key.startswith('_'):
+			continue
+		expected = field_types.get(key)
+		if expected is None or value is None:
+			continue
+		# Already the right type (note: bool is a subclass of int, so guard it).
+		if isinstance(value, expected) and not (expected is int and isinstance(value, bool)):
+			continue
+
+		if expected is bool:
+			if isinstance(value, bool):
+				continue
+			if isinstance(value, int):
+				data[key] = bool(value)
+			elif isinstance(value, str):
+				s = value.strip().lower()
+				if s in ('true', '1', 'yes', 'on'):
+					data[key] = True
+				elif s in ('false', '0', 'no', 'off', ''):
+					data[key] = False
+		elif expected is int:
+			# Avoid coercing real bools into ints.
+			if isinstance(value, bool):
+				continue
+			if isinstance(value, float):
+				if value.is_integer():
+					data[key] = int(value)
+			elif isinstance(value, str):
+				try:
+					data[key] = int(value)
+				except ValueError:
+					try:
+						f_val = float(value)
+						if f_val.is_integer():
+							data[key] = int(f_val)
+					except ValueError:
+						pass
+		elif expected is float:
+			if isinstance(value, bool):
+				continue
+			if isinstance(value, int):
+				data[key] = float(value)
+			elif isinstance(value, str):
+				try:
+					data[key] = float(value)
+				except ValueError:
+					pass
+		elif expected is list:
+			if isinstance(value, str):
+				s = value.strip()
+				if s.startswith('['):
+					try:
+						parsed = json.loads(s)
+						if isinstance(parsed, list):
+							data[key] = parsed
+					except (json.JSONDecodeError, TypeError):
+						pass
+		# str fields: leave as-is (don't stringify); unknown types: leave untouched.
+	return data
+
+
 def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator:
 	"""Create a secator finding from LLM-provided data.
 
@@ -507,6 +619,8 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator:
 		extra.update(unknown)
 		finding_data['extra_data'] = extra
 
+	finding_data = _coerce_finding_fields(cls, finding_data)
+
 	# Validate field types before instantiation
 	errors = cls.validate_fields(finding_data)
 	if errors:
@@ -519,6 +633,7 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator:
 		yield Ai(
 			content=f'{str(finding)}',
 			ai_type="add_finding",
+			extra_data={"finding": finding.toDict()},
 			_context=context
 		)
 		yield finding
@@ -594,7 +709,7 @@ def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator:
 
 	def run_single(act: Dict, idx: int) -> Dict:
 		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 {}

secator/ai/interactivity.py

@@ -119,7 +119,7 @@ def build_pending_prompt(self, question, choices, session_id, prompt_type="follo
 		)
 
 	def ask_user(self, question, choices, session_id, prompt_type="follow_up", **context):
-		answer = self._poll_for_answer(session_id, prompt_type)
+		answer = self._poll_for_answer(session_id, prompt_type, prompt_uuid=context.get("prompt_uuid"))
 		if answer is None:
 			return None
 
@@ -135,23 +135,72 @@ def ask_user(self, question, choices, session_id, prompt_type="follow_up", **con
 		# follow_up: return the answer text
 		return {"answer": answer}
 
-	def _poll_for_answer(self, session_id, prompt_type):
-		"""Poll DB for user answer until timeout."""
+	def poll_steers(self, session_id):
+		"""Drain pending steer docs for ``session_id`` and mark them consumed
+		(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",
+			"_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.
+
+		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,
+			"_context.session_id": session_id,
+		}
+		if prompt_uuid:
+			base["extra_data.prompt_uuid"] = prompt_uuid
+
 		elapsed = 0
 		while elapsed < self.timeout:
-			results = self.query_engine.search({
-				"_type": "ai",
-				"ai_type": prompt_type,
-				"session_id": session_id,
-				"status": "answered"
-			}, limit=1)
+			results = self.query_engine.search({**base, "status": "answered"}, limit=1)
 			if results:
 				return results[0].get("answer")
+			# A steer breaks the wait: treat the steer as the user's answer so the
+			# blocked follow-up resolves and the next turn redirects.
+			steers = self.poll_steers(session_id)
+			if steers:
+				return "\n".join(steers)
 			sleep(self.poll_interval)
 			elapsed += self.poll_interval
-		# Timeout: update finding status
+		# Timeout: flip ONLY this prompt's still-pending doc to timed_out, so a
+		# concurrent/older pending doc for the same session isn't disturbed.
 		self.query_engine.update(
-			{"_type": "ai", "ai_type": prompt_type, "session_id": session_id, "status": "pending"},
+			{**base, "status": "pending"},
 			{"$set": {"status": "timed_out"}}
 		)
 		return None

secator/ai/session.py

@@ -177,3 +177,73 @@ def replay_session(session):
 	except (json.JSONDecodeError, OSError) as e:
 		console.print(Error(message=f'Failed to load history: {e}'))
 		return None
+
+
+def restore_history_from_db(session_id, query_engine, model=None, encryptor=None, system_prompt=None):
+	"""Rebuild an in-memory ChatHistory from the workspace's `_type:"ai"` Mongo docs.
+
+	Headless equivalent of ``replay_session`` for the remote (web) path: a
+	respawned ``ai`` task on a different worker pod has no local report files, so
+	the conversation is rebuilt from the channel docs themselves (queried by
+	``session_id``, ordered by ``_timestamp``).
+
+	This is a **text-only** restore. Only the user turns (``ai_type="prompt"``)
+	and assistant turns (``ai_type="response"``) are reconstructed as litellm
+	``user``/``assistant`` messages. Intermediate tool-call / tool-result
+	messages are NOT persisted as ``_type:"ai"`` docs (only their human-readable
+	action display is), so they cannot be replayed verbatim. Fabricating
+	assistant ``tool_calls`` messages without their matching ``tool`` results
+	would produce a malformed transcript that most providers reject, so we
+	deliberately collapse tool activity into the surrounding text turns. This is
+	sufficient for ``mode="chat"`` continuation (the assistant text already
+	summarises what it did); for ``mode="attack"`` the intermediate tool I/O is
+	not replayed. See the feature spec for the richer-persistence follow-up.
+
+	Args:
+		session_id: The conversation's session id (UUID generated by the UI).
+		query_engine: A ``QueryEngine`` (must resolve to the workspace Mongo
+			backend for the docs to be visible).
+		model: Optional LLM model name to set on the returned history.
+		encryptor: Optional ``SensitiveDataEncryptor``. Persisted docs hold
+			plaintext (response content is decrypted before it is yielded), so
+			when an encryptor is active we re-encrypt restored turns to keep the
+			in-memory convention (encrypted) consistent with a fresh run.
+		system_prompt: Optional system prompt to set as the first message.
+
+	Returns:
+		ChatHistory: The rebuilt history (possibly with only a system prompt if
+		no prior docs exist).
+	"""
+	from secator.ai.history import ChatHistory
+	from secator.ai.encryption import maybe_encrypt
+
+	history = ChatHistory(model=model)
+	if system_prompt is not None:
+		history.set_system(maybe_encrypt(system_prompt, encryptor))
+
+	try:
+		docs = query_engine.search({'_type': 'ai', '_context.session_id': session_id})
+	except Exception as e:  # noqa: BLE001 - backend errors must not crash the worker
+		console.print(Warning(message=f'Failed to restore session from DB: {e}'))
+		return history
+
+	docs = sorted(docs or [], key=lambda d: d.get('_timestamp', 0))
+	for doc in docs:
+		ai_type = doc.get('ai_type')
+		content = doc.get('content', '')
+		if not content:
+			continue
+		if ai_type == 'prompt':
+			history.add_user(maybe_encrypt(content, encryptor))
+		elif ai_type == 'response':
+			history.add_assistant(maybe_encrypt(content, encryptor))
+		elif ai_type == 'steer':
+			# A mid-flight steer is a real user turn (an interjection that
+			# redirected the run): preserve it as a user message on respawn so the
+			# redirect survives a history restore. Mirror the live-loop framing.
+			history.add_user(maybe_encrypt(f'[User interjected]: {content}', encryptor))
+		# All other ai_types (action displays, follow_up/permission prompts,
+		# shell_output, summaries) are channel/UX artifacts, not conversation
+		# turns — intentionally skipped for a valid litellm transcript.
+
+	return history