tracer-cloud--opensre/tools/investigation/capability.py

"""Public runner API — wraps the pipeline for CLI and external callers."""

from __future__ import annotations

import asyncio
import contextlib
import contextvars
import logging
import queue
import threading
from collections.abc import AsyncIterator, Callable
from typing import TYPE_CHECKING, Any, cast

from core.domain.stream import StreamEvent
from core.state import AgentState
from platform.observability.errors.boundary import report_and_reraise
from platform.observability.errors.sentry import init_sentry
from platform.observability.trace.spans import stage_span
from tools.investigation.state_factory import make_initial_state
from tools.investigation.streaming import (
    InvestigationPipelineStreamError,
    resolved_integrations_stream_payload,
)

if TYPE_CHECKING:
    # Type-only — avoids paying the agent module's heavy import cost at
    # runner load while still letting static type-checkers validate
    # ``agent_class`` injections.
    from tools.investigation.stages.gather_evidence import ConnectedInvestigationAgent

logger = logging.getLogger(__name__)

_SENTRY_CAPTURED_ATTR = "_opensre_sentry_captured"


def _exception_was_captured(exc: BaseException) -> bool:
    return bool(getattr(exc, _SENTRY_CAPTURED_ATTR, False))


def _mark_exception_captured(exc: BaseException) -> None:
    with contextlib.suppress(Exception):
        setattr(exc, _SENTRY_CAPTURED_ATTR, True)


def _capture_exception_once(
    exc: BaseException,
    *,
    context: str,
    tags: dict[str, str] | None = None,
) -> None:
    if _exception_was_captured(exc):
        return
    from platform.observability.errors.sentry import capture_exception

    capture_exception(exc, context=context, tags=tags)
    _mark_exception_captured(exc)


def _traced_node(node_name: str, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
    with stage_span(node_name):
        try:
            return fn(*args, **kwargs)
        except Exception as exc:
            _capture_exception_once(
                exc,
                context=f"node.{node_name}",
                tags={"surface": "node", "node": node_name},
            )
            raise


def run_investigation(
    raw_alert: str | dict[str, Any],
    *,
    resolved_integrations: dict[str, Any] | None = None,
    openclaw_context: dict[str, Any] | None = None,
    opensre_evaluate: bool = False,
    investigation_metadata: tuple[str, str, str] | None = None,
    agent_class: type[ConnectedInvestigationAgent] | None = None,
) -> AgentState:
    """Run the investigation from a raw alert payload. Pure function: inputs in, state out.

    Args:
        raw_alert: The original alert payload or free-text alert description.
        resolved_integrations: Optional pre-resolved integrations dict. When provided,
            integration resolution is skipped — useful for synthetic testing where a
            FixtureGrafanaBackend should be injected without real credential resolution.
        investigation_metadata: Optional ``(alert_name, pipeline_name, severity)`` for
            initial state; avoids copying those fields onto ``raw_alert``.
        agent_class: Optional override for the investigation agent class. Defaults
            to ``ConnectedInvestigationAgent``. Callers that need a custom
            termination policy, structured-stage progression, or other
            agent-level extensions can pass a subclass instead.
    """
    init_sentry(entrypoint="pipeline")
    from tools.investigation.lifecycle import run_connected_investigation as _run

    initial = make_initial_state(
        raw_alert=raw_alert,
        opensre_evaluate=opensre_evaluate,
        investigation_metadata=investigation_metadata,
    )
    if resolved_integrations is not None:
        cast(dict[str, Any], initial)["resolved_integrations"] = resolved_integrations
    if openclaw_context:
        cast(dict[str, Any], initial)["openclaw_context"] = dict(openclaw_context)

    with report_and_reraise(
        logger=logger,
        message="run_investigation failed",
        tags={"surface": "pipeline", "component": "tools.investigation.capability"},
    ):
        from platform.analytics.investigation_loop import bind_investigation_loop_metrics_from_state

        state = _run(initial, agent_class=agent_class)
        bind_investigation_loop_metrics_from_state(state)
        return state


def resolve_investigation_context(
    *,
    raw_alert: dict[str, Any],
    alert_name: str | None,
    pipeline_name: str | None,
    severity: str | None,
) -> tuple[str, str, str]:
    """Resolve ``(alert_name, pipeline_name, severity)`` from overrides and payload defaults.

    Pure helper shared by every delivery surface (CLI, HTTP server, MCP); overrides win,
    then the raw alert's own fields, then common labels, then sensible fallbacks.
    """
    labels = raw_alert.get("commonLabels") or raw_alert.get("labels") or {}
    labels = labels if isinstance(labels, dict) else {}
    canonical = raw_alert.get("canonical_alert")
    canonical = canonical if isinstance(canonical, dict) else {}
    return (
        alert_name
        or raw_alert.get("alert_name")
        or raw_alert.get("title")
        or canonical.get("alert_name")
        or labels.get("alertname")
        or "Incident",
        pipeline_name
        or raw_alert.get("pipeline_name")
        or canonical.get("pipeline_name")
        or labels.get("pipeline_name")
        or labels.get("pipeline")
        or labels.get("service")
        or "unknown",
        severity
        or raw_alert.get("severity")
        or canonical.get("severity")
        or labels.get("severity")
        or "warning",
    )


def build_investigation_payload(
    state: AgentState,
    *,
    opensre_evaluate: bool = False,
) -> dict[str, Any]:
    """Project a finished investigation ``AgentState`` into the public result payload.

    Shared by every delivery surface so the serializable result shape stays identical
    regardless of how the run was triggered (CLI, HTTP server, MCP, integration webhook).
    """
    out: dict[str, Any] = {
        "report": state["slack_message"],
        "problem_md": state["problem_md"],
        "root_cause": state["root_cause"],
        "is_noise": state.get("is_noise", False),
        "validity_score": state.get("validity_score", 0.0),
    }
    if state.get("evidence_entries"):
        out["tool_calls"] = state["evidence_entries"]
    if opensre_evaluate:
        ev = state.get("opensre_llm_eval")
        if isinstance(ev, dict) and ev:
            out["opensre_llm_eval"] = ev
        elif not (state.get("opensre_eval_rubric") or "").strip():
            out["opensre_llm_eval"] = {
                "skipped": True,
                "reason": (
                    "No scoring_points on this alert — nothing to judge against "
                    "(not a scoring_points rubric payload, or field missing)."
                ),
            }
        else:
            out["opensre_llm_eval"] = {
                "skipped": True,
                "reason": "Evaluate was enabled but no judge output was recorded.",
            }
    return out


def run_investigation_payload(
    *,
    raw_alert: str | dict[str, Any],
    opensre_evaluate: bool = False,
    investigation_metadata: tuple[str, str, str] | None = None,
) -> dict[str, Any]:
    """Run an investigation and return the serializable result payload.

    The headless counterpart used by surfaces that do not render a live terminal
    stream (HTTP server, MCP, integration webhooks). It returns the same ``dict`` the
    CLI produces without depending on the ``cli`` package, so callers no longer have
    to reach up into ``cli.investigation`` to run an investigation.

    ``investigation_metadata`` is an optional ``(alert_name, pipeline_name, severity)``
    tuple for initial state (e.g. HTTP request overrides) without mutating ``raw_alert``.
    """
    state = run_investigation(
        raw_alert,
        opensre_evaluate=opensre_evaluate,
        investigation_metadata=investigation_metadata,
    )
    return build_investigation_payload(state, opensre_evaluate=opensre_evaluate)


async def astream_investigation(
    raw_alert: str | dict[str, Any],
    *,
    opensre_evaluate: bool = False,
    investigation_metadata: tuple[str, str, str] | None = None,
) -> AsyncIterator[Any]:
    """Stream investigation events in real time.

    Runs the pipeline in a background thread and yields StreamEvents as each
    stage and tool call happens. The renderer sees individual tool_start /
    tool_end events and shows them as spinner subtext, just like Claude Code.
    """
    init_sentry(entrypoint="pipeline")

    initial = make_initial_state(
        raw_alert=raw_alert,
        opensre_evaluate=opensre_evaluate,
        investigation_metadata=investigation_metadata,
    )

    # Silence the global ProgressTracker before starting the background thread
    # so pipeline internals (extract_alert, resolve_integrations, etc.) don't
    # open their own Rich Live display — the StreamRenderer drives it instead.
    from platform.observability import silence_progress_tracker

    silence_progress_tracker()

    event_queue: queue.Queue[StreamEvent | BaseException | None] = queue.Queue()
    loop = asyncio.get_running_loop()

    def _put(evt: StreamEvent) -> None:
        with contextlib.suppress(RuntimeError):  # loop already closed; consumer is gone
            loop.call_soon_threadsafe(event_queue.put_nowait, evt)

    def _make_node_event(kind: str, node: str, data: dict[str, Any]) -> StreamEvent:
        return StreamEvent(
            event_type="events",
            data={"event": kind, "name": node, "data": data},
            node_name=node,
            kind=kind,
            run_id="",
            tags=["graph:step:0"],
        )

    def _make_tool_event(kind: str, name: str, data: dict[str, Any]) -> StreamEvent:
        # Tool events carry the name in data so the renderer can extract it.
        payload = dict(data)
        payload["name"] = name
        payload["event"] = kind
        return StreamEvent(
            event_type="events",
            data=payload,
            node_name="investigation_agent",
            kind=kind,
            run_id="",
            tags=[],
        )

    def _on_agent_event(event_kind: str, data: dict[str, Any]) -> None:
        if event_kind == "agent_start":
            _put(_make_node_event("on_chain_start", "investigation_agent", data))
        elif event_kind == "tool_start":
            _put(_make_tool_event("on_tool_start", data.get("name", "tool"), data))
        elif event_kind == "tool_end":
            _put(_make_tool_event("on_tool_end", data.get("name", "tool"), data))
        elif event_kind == "llm_start":
            # Forward LLM iterations so the renderer can print "analyzing…" hints
            # during the silent gap between tool batches and during synthesis.
            _put(_make_tool_event("on_llm_start", "investigation_agent", data))
        elif event_kind == "agent_end":
            _put(
                _make_node_event(
                    "on_chain_end",
                    "investigation_agent",
                    {"output": data},
                )
            )

    def _run_pipeline() -> None:
        try:
            from core.state.updates import apply_state_updates
            from platform.analytics.investigation_loop import loop_metrics_from_state
            from tools.investigation.reporting.node import generate_report
            from tools.investigation.stages.diagnose import diagnose
            from tools.investigation.stages.gather_evidence import ConnectedInvestigationAgent
            from tools.investigation.stages.intake import extract_alert
            from tools.investigation.stages.plan_evidence import plan_actions
            from tools.investigation.stages.resolve_integrations import resolve_integrations

            state = initial

            # --- resolve_integrations ---
            _put(_make_node_event("on_chain_start", "resolve_integrations", {}))
            resolved_updates = _traced_node("resolve_integrations", resolve_integrations, state)
            apply_state_updates(state, resolved_updates)
            resolved = resolved_updates.get("resolved_integrations") or {}
            _put(
                _make_node_event(
                    "on_chain_end",
                    "resolve_integrations",
                    {
                        "output": {
                            "resolved_integrations": resolved_integrations_stream_payload(resolved)
                        }
                    },
                )
            )

            # --- extract_alert ---
            _put(_make_node_event("on_chain_start", "extract_alert", {}))
            apply_state_updates(state, _traced_node("extract_alert", extract_alert, state))
            _put(
                _make_node_event(
                    "on_chain_end",
                    "extract_alert",
                    {
                        "output": {
                            k: state.get(k) for k in ("alert_name", "pipeline_name", "severity")
                        }
                    },
                )
            )

            if state.get("is_noise"):
                with contextlib.suppress(RuntimeError):  # loop closed (consumer cancelled)
                    loop.call_soon_threadsafe(event_queue.put_nowait, None)
                return

            # --- plan_actions ---
            _put(_make_node_event("on_chain_start", "plan_actions", {}))
            apply_state_updates(
                state,
                _traced_node("plan_actions", plan_actions, state),
            )
            _put(
                _make_node_event(
                    "on_chain_end",
                    "plan_actions",
                    {
                        "output": {
                            "planned_actions": state.get("planned_actions", []),
                            "plan_rationale": state.get("plan_rationale", ""),
                            "plan_audit": state.get("plan_audit", {}),
                        }
                    },
                )
            )

            # --- investigation agent (with real tool events) ---
            apply_state_updates(
                state,
                _traced_node(
                    "investigation_agent",
                    ConnectedInvestigationAgent().run,
                    state,
                    on_event=_on_agent_event,
                ),
            )

            # --- diagnose ---
            _put(_make_node_event("on_chain_start", "diagnose", {}))
            apply_state_updates(state, _traced_node("diagnose", diagnose, state))
            _put(
                _make_node_event(
                    "on_chain_end",
                    "diagnose",
                    {
                        "output": {
                            "root_cause": state.get("root_cause", ""),
                            "root_cause_category": state.get("root_cause_category", ""),
                            "validity_score": state.get("validity_score"),
                            "validated_claims": state.get("validated_claims", []),
                            "remediation_steps": state.get("remediation_steps", []),
                        }
                    },
                )
            )

            # --- upstream correlation ---
            from tools.investigation.reporting.upstream_correlation import (
                enrich_upstream_correlation,
            )

            _put(
                _make_node_event(
                    "on_chain_start",
                    "correlate_upstream",
                    {},
                )
            )

            apply_state_updates(
                state,
                _traced_node(
                    "correlate_upstream",
                    enrich_upstream_correlation,
                    state,
                ),
            )

            _put(
                _make_node_event(
                    "on_chain_end",
                    "correlate_upstream",
                    {
                        "output": {
                            "correlation": state.get("correlation", {}),
                        }
                    },
                )
            )

            # --- deliver / publish (skip terminal/editor render; StreamRenderer owns output) ---
            _put(_make_node_event("on_chain_start", "publish_findings", {}))
            apply_state_updates(
                state,
                _traced_node(
                    "publish_findings",
                    generate_report,
                    state,
                    render_terminal=False,
                    open_editor=False,
                ),
            )

            _put(
                _make_node_event(
                    "on_chain_end",
                    "publish_findings",
                    {
                        "output": {
                            "root_cause": state.get("root_cause", ""),
                            "root_cause_category": state.get("root_cause_category", ""),
                            "validity_score": state.get("validity_score"),
                            "report": state.get("report", ""),
                            "slack_message": state.get("slack_message", ""),
                            "problem_md": state.get("problem_md", ""),
                            "validated_claims": state.get("validated_claims", []),
                            "remediation_steps": state.get("remediation_steps", []),
                        }
                    },
                )
            )

        except Exception as exc:
            loop_count, iteration_cap = loop_metrics_from_state(state)
            _capture_exception_once(exc, context="pipeline.astream_investigation")
            with contextlib.suppress(RuntimeError):
                loop.call_soon_threadsafe(
                    event_queue.put_nowait,
                    InvestigationPipelineStreamError(
                        cause=exc,
                        loop_count=loop_count,
                        iteration_cap=iteration_cap,
                    ),
                )
        finally:
            with contextlib.suppress(RuntimeError):
                loop.call_soon_threadsafe(event_queue.put_nowait, None)

    # Copy the caller's context so ContextVar bindings (session trace) reach the thread.
    thread = threading.Thread(
        target=contextvars.copy_context().run, args=(_run_pipeline,), daemon=True
    )
    thread.start()

    while True:
        # Drain the queue without blocking the event loop
        try:
            item = event_queue.get_nowait()
        except queue.Empty:
            await asyncio.sleep(0.01)
            continue

        if item is None:
            break
        if isinstance(item, InvestigationPipelineStreamError):
            raise item
        if isinstance(item, BaseException):
            raise item
        yield item

    thread.join()