tracer-cloud--opensre/tests/platform/observability/test_session_trace.py

"""Tests for process stats and session trace sink."""

from __future__ import annotations

import contextvars
import json
import threading
import time
from pathlib import Path
from typing import Any

import pytest

from core.agent_harness.session.persistence.jsonl_storage import JsonlSessionStorage
from platform.observability.trace import process_stats
from platform.observability.trace.process_stats import (
    sample_resource_snapshot,
    sample_thread_snapshot,
    sample_turn_boundary_stats,
)
from platform.observability.trace.spans import (
    NoopSessionTraceSink,
    bind_session_trace,
    component_span,
    current_trace_session_id,
    emit_route,
    emit_span,
    emit_thread_boundary,
    get_session_trace_sink,
    is_session_trace_active,
    llm_span,
    mark_span_outcome,
    set_session_trace_sink,
    stage_span,
    timed_span,
    tool_span,
    traced_session,
)
from surfaces.interactive_shell.session.trace_sink import (
    JsonlSessionTraceSink,
    jsonl_trace_sink_for_session,
)


@pytest.fixture(autouse=True)
def _reset_session_trace_sink() -> Any:
    set_session_trace_sink(NoopSessionTraceSink())
    yield
    set_session_trace_sink(NoopSessionTraceSink())


def _seed_session_jsonl(tmp_path: Path, session_id: str) -> Path:
    path = tmp_path / f"{session_id}.jsonl"
    path.write_text(
        json.dumps({"type": "session", "version": 2, "id": session_id}) + "\n",
        encoding="utf-8",
    )
    return path


def _activate_jsonl_sink(tmp_path: Path, session_id: str, monkeypatch: pytest.MonkeyPatch) -> Path:
    monkeypatch.setattr(
        "core.agent_harness.session.persistence.jsonl_storage.session_path",
        lambda sid: tmp_path / f"{sid}.jsonl",
    )
    path = _seed_session_jsonl(tmp_path, session_id)
    set_session_trace_sink(JsonlSessionTraceSink(storage=JsonlSessionStorage()))
    return path


def _trace_spans(path: Path) -> list[dict[str, Any]]:
    return [
        json.loads(line)
        for line in path.read_text(encoding="utf-8").strip().splitlines()
        if json.loads(line).get("type") == "trace_span"
    ]


def test_sample_thread_snapshot_lists_current_thread() -> None:
    snap = sample_thread_snapshot()
    assert snap["thread_count"] >= 1
    names = {row["name"] for row in snap["threads"]}
    assert threading.current_thread().name in names
    assert "main_thread_ident" in snap


def test_sample_resource_snapshot_includes_gc_counts() -> None:
    snap = sample_resource_snapshot()
    assert {"gc_gen0", "gc_gen1", "gc_gen2"} <= snap.keys()
    # Prefer current RSS (Linux /proc); peak watermark is separate.
    if "rss_mb" in snap:
        assert isinstance(snap["rss_mb"], float)
        assert snap["rss_mb"] > 0
    if process_stats._resource is not None:
        assert "rss_peak_mb" in snap
        assert isinstance(snap["rss_peak_mb"], float)


def test_sample_resource_snapshot_current_rss_not_peak(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """``rss_mb`` must be current RSS; ``rss_peak_mb`` is the high-water mark."""
    monkeypatch.setattr(process_stats, "_current_rss_mb", lambda: 50.0)

    class _FakeResource:
        RUSAGE_SELF = 0

        @staticmethod
        def getrusage(_who: int) -> object:
            # Peak higher than current (bytes on darwin path via normalize)
            return type("RU", (), {"ru_maxrss": 200 * 1024 * 1024})()

    monkeypatch.setattr(process_stats, "_resource", _FakeResource)
    monkeypatch.setattr(process_stats.sys, "platform", "darwin")
    snap = sample_resource_snapshot()
    assert snap["rss_mb"] == 50.0
    assert snap["rss_peak_mb"] == 200.0
    assert snap.get("rss_is_peak") is not True


def test_sample_resource_snapshot_skips_rss_without_backends(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """No resource / no Linux /proc → snapshot still succeeds without rss_mb."""
    monkeypatch.setattr(process_stats, "_resource", None)
    monkeypatch.setattr(process_stats.sys, "platform", "win32")
    snap = sample_resource_snapshot()
    assert "rss_mb" not in snap
    assert {"gc_gen0", "gc_gen1", "gc_gen2"} <= snap.keys()
    combined = sample_turn_boundary_stats()
    assert "rss_mb" not in combined
    assert "thread_count" in combined


def test_emit_span_and_thread_boundary_are_free_when_noop() -> None:
    assert not is_session_trace_active()
    assert emit_span(span_kind="route", name="gather_and_answer", session_id="s") == ""
    assert emit_thread_boundary("s", name="turn_boundary", phase="turn_start") == ""
    with timed_span(span_kind="component", name="x", session_id="s") as attrs:
        attrs["ok"] = True
        time.sleep(0)


def test_noop_emit_paths_skip_sampling_and_io(monkeypatch: pytest.MonkeyPatch) -> None:
    """Production default: emit helpers must not touch process stats or sink I/O."""
    calls: list[str] = []

    def _boom(*_a: Any, **_k: Any) -> Any:
        calls.append("sample")
        raise AssertionError("process sampling must not run on noop sink")

    monkeypatch.setattr(
        "platform.observability.trace.spans.sample_turn_boundary_stats",
        _boom,
    )
    assert emit_thread_boundary("s", name="turn_boundary", phase="turn_start") == ""
    assert emit_span(span_kind="route", name="x", session_id="s") == ""
    with timed_span(span_kind="component", name="y", session_id="s"):
        pass
    assert calls == []


def test_emit_span_without_session_id_is_noop_even_when_sink_active(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    path = _activate_jsonl_sink(tmp_path, "sess-no-sid", monkeypatch)
    assert is_session_trace_active()
    assert emit_span(span_kind="route", name="orphan") == ""
    assert emit_route("orphan") == ""
    with timed_span(span_kind="component", name="orphan"):
        pass
    assert _trace_spans(path) == []


def test_bind_session_trace_sets_and_clears_context() -> None:
    assert current_trace_session_id() is None
    with bind_session_trace("sess-bound"):
        assert current_trace_session_id() == "sess-bound"
    assert current_trace_session_id() is None
    with bind_session_trace(None):
        assert current_trace_session_id() is None
    with bind_session_trace(""):
        assert current_trace_session_id() is None


def test_emit_span_writes_route_when_sink_active(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    session_id = "sess-route-test"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)
    with bind_session_trace(session_id):
        emit_span(
            span_kind="route",
            name="gather_and_answer",
            attributes={"handled": False},
        )
        with timed_span(span_kind="stage", name="extract_alert") as attrs:
            attrs["fields_updated"] = ["alert"]
            time.sleep(0.001)
    kinds = {(rec["span_kind"], rec["name"]) for rec in _trace_spans(path)}
    assert ("route", "gather_and_answer") in kinds
    assert ("stage", "extract_alert") in kinds
    stage = next(r for r in _trace_spans(path) if r.get("name") == "extract_alert")
    assert stage["duration_ms"] >= 0
    assert stage["attributes"]["fields_updated"] == ["alert"]


def test_timed_span_honors_status_override(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    from platform.observability.trace.spans import SPAN_STATUS_ATTR, SPAN_STATUS_ERROR

    session_id = "sess-status"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)
    with timed_span(span_kind="component", name="investigation", session_id=session_id) as attrs:
        attrs[SPAN_STATUS_ATTR] = SPAN_STATUS_ERROR
        attrs["failure_category"] = "user_cancelled"
    rec = _trace_spans(path)[-1]
    assert rec["status"] == SPAN_STATUS_ERROR
    assert SPAN_STATUS_ATTR not in rec.get("attributes", {})
    assert rec["attributes"]["failure_category"] == "user_cancelled"


def test_timed_span_marks_error_on_exception(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    from platform.observability.trace.spans import SPAN_STATUS_ERROR

    def _raise_runtime_error() -> None:
        raise RuntimeError("boom")

    session_id = "sess-exc"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)
    # Outer raises so the exception still propagates through timed_span's
    # finally (status=error); inner span must not be wrapped by pytest.raises
    # or the exception would be swallowed before the span exits.
    with (
        pytest.raises(RuntimeError, match="boom"),
        timed_span(span_kind="component", name="failing", session_id=session_id),
    ):
        _raise_runtime_error()
    rec = _trace_spans(path)[-1]
    assert rec["name"] == "failing"
    assert rec["status"] == SPAN_STATUS_ERROR
    assert rec["duration_ms"] >= 0


def test_timed_span_marks_error_on_base_exception(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    # KeyboardInterrupt / CancelledError derive from BaseException, not Exception.
    from platform.observability.trace.spans import SPAN_STATUS_ERROR

    def _raise_keyboard_interrupt() -> None:
        raise KeyboardInterrupt

    session_id = "sess-cancel"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)
    with (
        pytest.raises(KeyboardInterrupt),
        timed_span(span_kind="tool", name="cancelled", session_id=session_id),
    ):
        _raise_keyboard_interrupt()
    rec = _trace_spans(path)[-1]
    assert rec["name"] == "cancelled"
    assert rec["status"] == SPAN_STATUS_ERROR


def test_timed_span_exception_status_beats_preset_override(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    # A pre-set "ok" override must not mask an exception that propagates out.
    from platform.observability.trace.spans import (
        SPAN_STATUS_ATTR,
        SPAN_STATUS_ERROR,
        SPAN_STATUS_OK,
    )

    def _boom() -> None:
        raise RuntimeError("boom")

    session_id = "sess-override-race"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)
    with (
        pytest.raises(RuntimeError, match="boom"),
        timed_span(span_kind="tool", name="racing", session_id=session_id) as attrs,
    ):
        attrs[SPAN_STATUS_ATTR] = SPAN_STATUS_OK
        _boom()
    rec = _trace_spans(path)[-1]
    assert rec["status"] == SPAN_STATUS_ERROR


def test_stage_span_emits_from_thread_with_copied_context(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    # The investigation pipeline runs stages in a worker thread; a copied context
    # carries the bound session id across the thread boundary so the span emits.
    session_id = "sess-thread"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)

    def _emit_stage() -> None:
        with stage_span("gather_evidence"):
            pass

    with bind_session_trace(session_id):
        thread = threading.Thread(target=contextvars.copy_context().run, args=(_emit_stage,))
        thread.start()
        thread.join()

    assert "gather_evidence" in {rec["name"] for rec in _trace_spans(path)}


def test_stage_span_drops_in_plain_thread_without_context(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    # A plain thread starts with a fresh context, so the bound session id is unseen.
    session_id = "sess-thread-plain"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)

    def _emit_stage() -> None:
        with stage_span("gather_evidence"):
            pass

    with bind_session_trace(session_id):
        thread = threading.Thread(target=_emit_stage)
        thread.start()
        thread.join()

    assert _trace_spans(path) == []


def test_mark_span_outcome_sets_status_and_extra_attrs() -> None:
    from platform.observability.trace.spans import SPAN_STATUS_ATTR, SPAN_STATUS_ERROR

    attrs: dict[str, Any] = {}
    mark_span_outcome(attrs, "ok", source="agent")
    assert attrs == {"outcome": "ok", "source": "agent"}

    attrs = {}
    mark_span_outcome(attrs, "blocked", error=True, reason="deny", ignored=None)
    assert attrs["outcome"] == "blocked"
    assert attrs[SPAN_STATUS_ATTR] == SPAN_STATUS_ERROR
    assert attrs["reason"] == "deny"
    assert "ignored" not in attrs


def test_semantic_helpers_match_span_kinds(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    session_id = "sess-helpers"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)
    with traced_session(session_id, component="gateway_turn") as attrs:
        mark_span_outcome(attrs, "ok")
        emit_route("gather_and_answer", attributes={"handled": False})
        with component_span("action_turn"):
            pass
        with stage_span("intake"):
            pass
        with tool_span("echo", tool_call_id="c1") as tool_attrs:
            mark_span_outcome(tool_attrs, "ok", source="agent")
        with llm_span("model-x", iteration=1):
            pass
    spans = _trace_spans(path)
    kinds = {(rec["span_kind"], rec["name"]) for rec in spans}
    assert ("component", "gateway_turn") in kinds
    assert ("route", "gather_and_answer") in kinds
    assert ("component", "action_turn") in kinds
    assert ("stage", "intake") in kinds
    assert ("tool", "echo") in kinds
    assert ("llm", "model-x") in kinds

    tool = next(r for r in spans if r["span_kind"] == "tool")
    assert tool["attributes"]["tool_call_id"] == "c1"
    assert tool["attributes"]["outcome"] == "ok"
    assert tool["attributes"]["source"] == "agent"

    llm = next(r for r in spans if r["span_kind"] == "llm")
    assert llm["attributes"]["iteration"] == 1

    gateway = next(r for r in spans if r["name"] == "gateway_turn")
    assert gateway["attributes"]["outcome"] == "ok"


def test_traced_session_noop_when_session_id_missing() -> None:
    assert not is_session_trace_active()
    with traced_session(None, component="gateway_turn") as attrs:
        attrs["note"] = "still mutable"
        assert current_trace_session_id() is None
    assert attrs["note"] == "still mutable"


def test_jsonl_trace_sink_writes_trace_span(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    session_id = "sess-thread-test"
    path = _activate_jsonl_sink(tmp_path, session_id, monkeypatch)
    emit_thread_boundary(session_id, name="turn_boundary", phase="turn_start")
    lines = path.read_text(encoding="utf-8").strip().splitlines()
    assert len(lines) == 2
    rec = json.loads(lines[1])
    assert rec["type"] == "trace_span"
    assert rec["span_kind"] == "thread"
    attrs: dict[str, Any] = rec["attributes"]
    assert attrs["phase"] == "turn_start"
    assert attrs["thread_count"] >= 1
    assert isinstance(attrs["threads"], list)
    set_session_trace_sink(NoopSessionTraceSink())
    assert isinstance(get_session_trace_sink(), NoopSessionTraceSink)


def test_jsonl_trace_sink_for_session_uses_jsonl_or_noop() -> None:
    from core.agent_harness.session import InMemorySessionStorage
    from surfaces.interactive_shell.session import Session

    jsonl_session = Session(storage=JsonlSessionStorage())
    assert isinstance(jsonl_trace_sink_for_session(jsonl_session), JsonlSessionTraceSink)

    memory_session = Session(storage=InMemorySessionStorage())
    assert isinstance(jsonl_trace_sink_for_session(memory_session), NoopSessionTraceSink)

    assert isinstance(jsonl_trace_sink_for_session(object()), NoopSessionTraceSink)


def test_set_session_trace_sink_none_restores_noop() -> None:
    class _RecordingSink:
        def emit(self, *_a: Any, **_k: Any) -> str:
            return "id"

    set_session_trace_sink(_RecordingSink())  # type: ignore[arg-type]
    assert is_session_trace_active()
    set_session_trace_sink(None)
    assert not is_session_trace_active()
    assert isinstance(get_session_trace_sink(), NoopSessionTraceSink)