tracer-cloud--opensre/surfaces/cli/investigation/investigate.py

"""Shared investigation helpers for CLI entrypoints."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import threading
from collections.abc import Generator
from typing import TYPE_CHECKING, Any, NoReturn

from core.domain.stream import StreamEvent
from platform.observability.trace.hook import traceable
from surfaces.cli.error_mapping import reraise_cli_runtime_error
from tools.investigation.session_runner import InvestigationPumpCancelled, check_llm_settings

_logger = logging.getLogger(__name__)

_SESSION_EVENT_POLL_S = 0.25

if TYPE_CHECKING:
    from platform.analytics.cli import InvestigationTracker


def _reraise_cli_investigation_failure(exc: BaseException) -> NoReturn:
    """Map investigation runtime failures to structured CLI errors."""
    from tools.investigation.session_runner import reraise_investigation_failure

    if isinstance(exc, InvestigationPumpCancelled):
        reraise_investigation_failure(exc)
    reraise_cli_runtime_error(exc)


@traceable(name="investigation")
def run_investigation_cli(
    *,
    raw_alert: dict[str, Any],
    opensre_evaluate: bool = False,
    investigation_metadata: tuple[str, str, str] | None = None,
) -> dict[str, Any]:
    """Run the investigation and return the CLI-facing JSON payload.

    Thin CLI wrapper over :func:`tools.investigation.capability.run_investigation_payload`:
    it adds the CLI-only precondition check (LLM settings) and maps runtime failures to
    structured ``OpenSREError`` messages. The run itself and the result shaping live in
    ``core`` so non-CLI surfaces can reuse them without importing ``cli``.

    ``investigation_metadata`` is an optional ``(alert_name, pipeline_name, severity)``
    tuple for initial state (e.g. HTTP request overrides) without mutating ``raw_alert``.
    """
    check_llm_settings()
    from tools.investigation.capability import run_investigation_payload

    try:
        return run_investigation_payload(
            raw_alert=raw_alert,
            opensre_evaluate=opensre_evaluate,
            investigation_metadata=investigation_metadata,
        )
    except Exception as exc:
        _reraise_cli_investigation_failure(exc)


def stream_investigation_cli(
    *,
    raw_alert: dict[str, Any],
) -> Generator[StreamEvent]:
    """Stream investigation events locally via the async pipeline stream.

    Bridges the async streaming API into a synchronous iterator
    using a background thread + queue so events are yielded in real time
    (not batched).  The same ``StreamRenderer`` used for remote
    investigations can render local runs identically.

    On :exc:`KeyboardInterrupt` the background asyncio task is cancelled
    and the thread is joined so Ctrl+C terminates cleanly instead of
    leaving an orphaned investigation task in flight.
    """
    import queue

    from tools.investigation.capability import astream_investigation

    check_llm_settings()

    event_queue: queue.Queue[StreamEvent | BaseException | None] = queue.Queue()
    loop_ref: dict[str, asyncio.AbstractEventLoop] = {}
    pump_task_ref: dict[str, asyncio.Task[None]] = {}

    def _run_async() -> None:
        loop = asyncio.new_event_loop()
        loop_ref["loop"] = loop
        try:

            async def _pump() -> None:
                async for evt in astream_investigation(
                    raw_alert=raw_alert,
                ):
                    event_queue.put(evt)

            task = loop.create_task(_pump())
            pump_task_ref["task"] = task
            try:
                loop.run_until_complete(task)
            except asyncio.CancelledError:
                event_queue.put(InvestigationPumpCancelled())
        except Exception as exc:
            event_queue.put(exc)
        finally:
            event_queue.put(None)
            loop.close()

    thread = threading.Thread(target=_run_async, daemon=True)
    thread.start()

    def _cancel_pump() -> None:
        loop = loop_ref.get("loop")
        task = pump_task_ref.get("task")
        if loop is None or task is None or loop.is_closed():
            return
        with contextlib.suppress(RuntimeError):
            loop.call_soon_threadsafe(task.cancel)

    try:
        while True:
            try:
                item = event_queue.get(timeout=_SESSION_EVENT_POLL_S)
            except queue.Empty:
                continue
            if isinstance(item, BaseException):
                from platform.analytics.investigation_loop import (
                    publish_loop_metrics_from_stream_failure,
                )

                thread.join(timeout=5)
                _reraise_cli_investigation_failure(publish_loop_metrics_from_stream_failure(item))
            if item is None:
                break
            yield item
    finally:
        _cancel_pump()
        thread.join(timeout=5)
        if thread.is_alive():
            _logger.warning(
                "investigation thread did not terminate within 5s after cancellation; "
                "an LLM call may still be in flight"
            )


def run_investigation_cli_streaming(
    *,
    raw_alert: dict[str, Any],
    tracker: InvestigationTracker | None = None,
) -> dict[str, Any]:
    """Run the investigation with real-time streaming UI and return the result.

    Uses async pipeline streaming + ``StreamRenderer`` so the local CLI shows
    the same live tool-call and reasoning updates as a remote investigation.
    """
    from surfaces.cli.ui.renderer import StreamRenderer

    events = stream_investigation_cli(
        raw_alert=raw_alert,
    )
    renderer = StreamRenderer(local=True)
    try:
        final_state = renderer.render_stream(events)
    except KeyboardInterrupt:
        events.close()
        raise

    from surfaces.interactive_shell.ui.components.key_reader import restore_stdin_terminal
    from surfaces.interactive_shell.ui.feedback import prompt_investigation_feedback

    restore_stdin_terminal()
    prompt_investigation_feedback(final_state)
    if tracker is not None:
        tracker.record_loop_metrics_from_state(final_state)
    return {
        "report": final_state.get("slack_message", final_state.get("report", "")),
        "problem_md": final_state.get("problem_md", ""),
        "root_cause": final_state.get("root_cause", ""),
        "is_noise": final_state.get("is_noise", False),
        "tool_calls": final_state.get("evidence_entries", []),
    }