"""Helpers for launching session-local background investigations.""" from __future__ import annotations import threading from collections.abc import Callable from contextlib import nullcontext from typing import Any from uuid import uuid4 from prompt_toolkit.patch_stdout import patch_stdout from rich.console import Console from rich.markup import escape from platform.analytics.cli import track_investigation from platform.analytics.source import EntrypointSource, TriggerMode from platform.common.errors import OpenSREError from surfaces.interactive_shell.runtime import ( BackgroundInvestigationRecord, Session, TaskKind, ) from surfaces.interactive_shell.runtime.background.notifications import ( deliver_background_notifications, ) from surfaces.interactive_shell.ui import DIM, ERROR, HIGHLIGHT, WARNING from surfaces.interactive_shell.utils.error_handling.exception_reporting import report_exception BackgroundRunFn = Callable[..., dict[str, Any]] def _safe_console_print(console: Console, message: str) -> None: isatty = getattr(console.file, "isatty", None) stdout_context = patch_stdout(raw=True) if callable(isatty) and isatty() else nullcontext() with stdout_context: console.print(message) def drain_background_notices(session: Session, console: Console) -> None: """Print queued background investigation status lines on the main REPL thread.""" for message in session.terminal.drain_background_notices(): _safe_console_print(console, message) def _build_record( *, task_id: str, command: str, investigation_id: str, ) -> BackgroundInvestigationRecord: return BackgroundInvestigationRecord( task_id=task_id, status="running", command=command, investigation_id=investigation_id, ) def _top_analysis(final_state: dict[str, Any]) -> tuple[str, ...]: claims = final_state.get("validated_claims", []) if not isinstance(claims, list): return () lines: list[str] = [] for entry in claims: if not isinstance(entry, dict): continue claim = str(entry.get("claim") or "").strip() if claim: lines.append(claim) if len(lines) >= 3: break return tuple(lines) def _next_steps(final_state: dict[str, Any]) -> tuple[str, ...]: steps = final_state.get("remediation_steps", []) if not isinstance(steps, list): return () values: list[str] = [] for step in steps[:3]: text = str(step).strip() if text: values.append(text) return tuple(values) def _stats(final_state: dict[str, Any]) -> dict[str, Any]: tool_calls = final_state.get("evidence_entries", []) loops = final_state.get("investigation_loop_count", 0) validity = final_state.get("validity_score", 0.0) return { "tool_call_count": len(tool_calls) if isinstance(tool_calls, list) else 0, "investigation_loop_count": int(loops) if isinstance(loops, int | float) else 0, "validity_score": float(validity) if isinstance(validity, int | float) else 0.0, } def _start_background_investigation( *, session: Session, console: Console, display_command: str, run_fn: BackgroundRunFn, kwargs: dict[str, Any], investigation_target: str = "", input_path: str | None = None, ) -> str: investigation_id = str(uuid4()) session.last_investigation_id = investigation_id task = session.task_registry.create(TaskKind.INVESTIGATION, command=display_command) task.mark_running() record = _build_record( task_id=task.task_id, command=display_command, investigation_id=investigation_id, ) session.terminal.background_investigations[task.task_id] = record def _worker() -> None: try: with track_investigation( entrypoint=EntrypointSource.CLI_REPL_FILE, trigger_mode=TriggerMode.FILE, input_path=input_path, interactive=True, investigation_id=investigation_id, investigation_target=investigation_target or None, session=session, ) as tracker: final_state = run_fn(cancel_requested=task.cancel_requested, **kwargs) tracker.record_loop_metrics_from_state(final_state) root = str(final_state.get("root_cause") or "") record.status = "completed" record.root_cause = root record.top_analysis = _top_analysis(final_state) record.next_steps = _next_steps(final_state) record.stats = _stats(final_state) record.final_state = dict(final_state) record.notification_results = deliver_background_notifications( record=record, channels=session.terminal.background_notification_preferences.channels, ) task.mark_completed(result=root) session.terminal.enqueue_background_notice( f"[{HIGHLIGHT}]background investigation complete[/] " f"[{DIM}]— task {escape(task.task_id)} ready; " f"use[/] [{HIGHLIGHT}]/background show {escape(task.task_id)}[/]", ) except KeyboardInterrupt: record.status = "cancelled" task.mark_cancelled() session.terminal.enqueue_background_notice( f"[{WARNING}]background investigation cancelled[/] " f"[{DIM}]for task {escape(task.task_id)}.[/]", ) except OpenSREError as exc: record.status = "failed" task.mark_failed(str(exc)) session.terminal.enqueue_background_notice( f"[{ERROR}]background investigation failed[/] " f"[{DIM}]for task {escape(task.task_id)}:[/] {escape(str(exc))}", ) except Exception as exc: # noqa: BLE001 record.status = "failed" task.mark_failed(str(exc)) report_exception(exc, context="surfaces.interactive_shell.background_investigation") session.terminal.enqueue_background_notice( f"[{ERROR}]background investigation failed[/] " f"[{DIM}]for task {escape(task.task_id)}:[/] {escape(str(exc))}", ) thread = threading.Thread( target=_worker, daemon=True, name=f"background-investigation-{task.task_id}", ) thread.start() _safe_console_print( console, f"[{DIM}]background investigation started — task[/] [bold]{escape(task.task_id)}[/bold]. " f"[{HIGHLIGHT}]/background list[/] [{DIM}]to monitor, " f"[/][{HIGHLIGHT}]/cancel {escape(task.task_id)}[/] [{DIM}]to stop.[/]", ) return task.task_id def start_background_text_investigation( *, alert_text: str, session: Session, console: Console, display_command: str = "background free-text investigation", investigation_target: str = "", ) -> str: from surfaces.interactive_shell.runtime.investigation_adapter import ( run_investigation_for_session_background, ) return _start_background_investigation( session=session, console=console, display_command=display_command, run_fn=run_investigation_for_session_background, kwargs={ "alert_text": alert_text, "context_overrides": session.accumulated_context or None, }, investigation_target=investigation_target, input_path=display_command, ) def start_background_template_investigation( *, template_name: str, session: Session, console: Console, display_command: str, investigation_target: str = "", ) -> str: from surfaces.interactive_shell.runtime.investigation_adapter import ( run_sample_alert_for_session_background, ) return _start_background_investigation( session=session, console=console, display_command=display_command, run_fn=run_sample_alert_for_session_background, kwargs={ "template_name": template_name, "context_overrides": session.accumulated_context or None, }, investigation_target=investigation_target, input_path=f"template:{template_name}", )