roboco/tests/unit/runtime/test_orchestrator_write_hooks.py

"""
Unit tests for orchestrator write-hooks:
    _finalize_spawn_session  — closes the agent_spawn_sessions DB row on stop
    _sweep_token_snapshots   — polls active agents and upserts token snapshots

These tests mock the httpx transport and the SQLAlchemy session factory so no
real network or database is required.

Coverage:
  1. _finalize_spawn_session success — SDK returns token data → DB update
     carries those exact values to calculate_cost and the UPDATE statement.
  2. _finalize_spawn_session HTTP error — SDK unreachable → DB update proceeds
     with all-zero token counts (finalization must not raise).
  3. _sweep_token_snapshots active agent — non-zero tokens → snapshot row
     inserted and session row updated.
  4. _sweep_token_snapshots per-agent HTTP error — ConnectError on one agent
     is caught; the sweep continues and the next agent is still processed.
"""

from __future__ import annotations

from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import UUID, uuid4

import httpx
from roboco.models.runtime import (
    AgentInstance,
    OrchestratorAgentConfig,
    OrchestratorAgentState,
)
from roboco.runtime.orchestrator import AgentOrchestrator

# ---------------------------------------------------------------------------
# Module-level constants (ruff PLR2004: no magic values in comparisons)
# ---------------------------------------------------------------------------

_AGENT_ID = "be-dev-1"
_AGENT_ID_2 = "be-dev-2"

# Token counts used in success-path assertions
_TI = 111  # tokens_input
_TO = 222  # tokens_output
_TCR = 33  # tokens_cache_read
_TCW = 44  # tokens_cache_write

# Token counts for the snapshot test
_SNAP_TI = 50
_SNAP_TO = 100
_SNAP_TCR = 10
_SNAP_TCW = 5

# Token counts for the loop-continues test (agent-2)
_LOOP_TI = 25
_LOOP_TO = 75

# Expected number of DB execute() calls for a normal finalize (SELECT + UPDATE)
_FINALIZE_EXEC_CALLS = 2


# ---------------------------------------------------------------------------
# Test helpers
# ---------------------------------------------------------------------------


def _make_orchestrator() -> AgentOrchestrator:
    """Minimal AgentOrchestrator — no background tasks, no real DB."""
    return AgentOrchestrator(mcp_config_dir=Path("/tmp"), project_root=Path("/tmp"))


def _make_instance(
    agent_id: str = _AGENT_ID,
    usage_session_id: UUID | None = None,
) -> AgentInstance:
    """Return an ACTIVE AgentInstance with a running container."""
    return AgentInstance(
        agent_id=agent_id,
        state=OrchestratorAgentState.ACTIVE,
        container_id="abc123def456",
        config=OrchestratorAgentConfig(
            agent_id=agent_id,
            blueprint_path=Path("/tmp/blueprint.md"),
            model="sonnet",
        ),
        usage_session_id=usage_session_id,
    )


def _mock_response(
    status: int = 200,
    json_data: dict[str, Any] | None = None,
) -> MagicMock:
    """Build a mock httpx.Response."""
    resp = MagicMock(spec=httpx.Response)
    resp.status_code = status
    resp.json = MagicMock(return_value=json_data or {})
    return resp


def _make_db_factory(
    session_row: Any = None,
    add_list: list[Any] | None = None,
    execute_list: list[Any] | None = None,
) -> Any:
    """Return a callable that acts like get_session_factory().

    The returned callable, when called with no arguments, returns an async
    context manager yielding a mock AsyncSession whose execute() returns a
    result whose scalar_one_or_none() returns *session_row*.
    """

    @asynccontextmanager
    async def _db_context() -> Any:
        db = MagicMock()

        result = MagicMock()
        result.scalar_one_or_none = MagicMock(return_value=session_row)

        async def _exec(stmt: Any) -> MagicMock:
            if execute_list is not None:
                execute_list.append(stmt)
            return result

        db.execute = AsyncMock(side_effect=_exec)
        db.commit = AsyncMock()

        def _add(obj: Any) -> None:
            if add_list is not None:
                add_list.append(obj)

        db.add = _add if add_list is not None else MagicMock()
        yield db

    return _db_context


class _FakeHTTPClient:
    """Drop-in replacement for ``httpx.AsyncClient`` in tests.

    Accepts a *handler* callable ``(url: str) -> httpx.Response | raises``
    that is invoked by ``get()``.  Supports the ``async with`` protocol.
    """

    def __init__(self, handler: Any, **_: Any) -> None:
        self._handler = handler

    async def __aenter__(self) -> _FakeHTTPClient:
        return self

    async def __aexit__(self, *_: Any) -> None:
        pass

    async def get(self, url: str, **_: Any) -> Any:
        return self._handler(url)


# ---------------------------------------------------------------------------
# _finalize_spawn_session — success path
# ---------------------------------------------------------------------------


async def test_finalize_spawn_session_success_calls_calculate_cost() -> None:
    """Token values returned by the SDK /usage/status are passed to calculate_cost.

    This verifies the full data-flow: SDK response → token vars → cost calc.
    """
    orch = _make_orchestrator()
    session_uuid = uuid4()
    orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)

    token_data = {
        "tokens_input": _TI,
        "tokens_output": _TO,
        "tokens_cache_read": _TCR,
        "tokens_cache_write": _TCW,
    }

    def _handler(_url: str) -> Any:
        return _mock_response(200, token_data)

    session_row = MagicMock()
    session_row.id = session_uuid
    db_factory = _make_db_factory(session_row=session_row)

    def _client_cls(**_kw: Any) -> _FakeHTTPClient:
        return _FakeHTTPClient(_handler)

    with (
        patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
        patch("roboco.db.base.get_session_factory", return_value=db_factory),
        patch("roboco.billing.pricing.calculate_cost", return_value=0.001) as mock_cost,
    ):
        await orch._finalize_spawn_session(_AGENT_ID, exit_reason="stopped")

    mock_cost.assert_called_once_with(
        model="sonnet",
        tokens_input=_TI,
        tokens_output=_TO,
        tokens_cache_read=_TCR,
        tokens_cache_write=_TCW,
    )


async def test_finalize_spawn_session_success_executes_select_and_update() -> None:
    """When a session row exists the function calls execute() twice: SELECT + UPDATE."""
    orch = _make_orchestrator()
    session_uuid = uuid4()
    orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)

    def _handler(_url: str) -> Any:
        return _mock_response(
            200,
            {
                "tokens_input": 10,
                "tokens_output": 20,
                "tokens_cache_read": 0,
                "tokens_cache_write": 0,
            },
        )

    session_row = MagicMock()
    session_row.id = session_uuid
    execute_calls: list[Any] = []
    db_factory = _make_db_factory(session_row=session_row, execute_list=execute_calls)

    def _client_cls(**_kw: Any) -> _FakeHTTPClient:
        return _FakeHTTPClient(_handler)

    with (
        patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
        patch("roboco.db.base.get_session_factory", return_value=db_factory),
        patch("roboco.billing.pricing.calculate_cost", return_value=0.0),
    ):
        await orch._finalize_spawn_session(_AGENT_ID, exit_reason="completed")

    # SELECT (find the row) + UPDATE (write the values) = 2 execute() calls
    assert len(execute_calls) == _FINALIZE_EXEC_CALLS


# ---------------------------------------------------------------------------
# _finalize_spawn_session — HTTP-error path
# ---------------------------------------------------------------------------


async def test_finalize_spawn_session_http_error_uses_zero_tokens() -> None:
    """When the SDK endpoint is unreachable, finalization uses zero tokens.

    The function must not raise; cost must be calculated with all-zero counts.
    """
    orch = _make_orchestrator()
    session_uuid = uuid4()
    orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)

    def _boom(_url: str) -> Any:
        raise httpx.ConnectError("container not reachable")

    session_row = MagicMock()
    session_row.id = session_uuid
    db_factory = _make_db_factory(session_row=session_row)

    def _client_cls(**_kw: Any) -> _FakeHTTPClient:
        return _FakeHTTPClient(_boom)

    with (
        patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
        patch.object(orch, "_usage_from_transcript", return_value=(0, 0, 0, 0)),
        patch("roboco.db.base.get_session_factory", return_value=db_factory),
        patch("roboco.billing.pricing.calculate_cost", return_value=0.0) as mock_cost,
    ):
        # Must not raise even though the SDK is unreachable
        await orch._finalize_spawn_session(_AGENT_ID, exit_reason="stopped")

    mock_cost.assert_called_once_with(
        model="sonnet",
        tokens_input=0,
        tokens_output=0,
        tokens_cache_read=0,
        tokens_cache_write=0,
    )


async def test_finalize_spawn_session_non_200_uses_zero_tokens() -> None:
    """A non-200 SDK response results in zero-token finalization, no exception."""
    orch = _make_orchestrator()
    session_uuid = uuid4()
    orch._instances[_AGENT_ID] = _make_instance(usage_session_id=session_uuid)

    def _handler(_url: str) -> Any:
        return _mock_response(503)

    session_row = MagicMock()
    session_row.id = session_uuid
    db_factory = _make_db_factory(session_row=session_row)

    def _client_cls(**_kw: Any) -> _FakeHTTPClient:
        return _FakeHTTPClient(_handler)

    with (
        patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
        patch.object(orch, "_usage_from_transcript", return_value=(0, 0, 0, 0)),
        patch("roboco.db.base.get_session_factory", return_value=db_factory),
        patch("roboco.billing.pricing.calculate_cost", return_value=0.0) as mock_cost,
    ):
        await orch._finalize_spawn_session(_AGENT_ID, exit_reason="stopped")

    mock_cost.assert_called_once_with(
        model="sonnet",
        tokens_input=0,
        tokens_output=0,
        tokens_cache_read=0,
        tokens_cache_write=0,
    )


# ---------------------------------------------------------------------------
# _sweep_token_snapshots — active agent
# ---------------------------------------------------------------------------


async def test_sweep_token_snapshots_inserts_snapshot_for_active_agent() -> None:
    """An active agent with non-zero tokens gets a snapshot row added to the DB."""
    orch = _make_orchestrator()
    instance = _make_instance(_AGENT_ID)
    instance.state = OrchestratorAgentState.ACTIVE
    orch._instances[_AGENT_ID] = instance

    token_data = {
        "tokens_input": _SNAP_TI,
        "tokens_output": _SNAP_TO,
        "tokens_cache_read": _SNAP_TCR,
        "tokens_cache_write": _SNAP_TCW,
    }

    def _handler(_url: str) -> Any:
        return _mock_response(200, token_data)

    session_row = MagicMock()
    session_row.id = uuid4()
    added: list[Any] = []
    db_factory = _make_db_factory(session_row=session_row, add_list=added)

    def _client_cls(**_kw: Any) -> _FakeHTTPClient:
        return _FakeHTTPClient(_handler)

    with (
        patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
        patch("roboco.db.base.get_session_factory", return_value=db_factory),
    ):
        await orch._sweep_token_snapshots()

    # Exactly one snapshot row must have been passed to db.add()
    assert len(added) == 1
    snap = added[0]
    assert snap.tokens_input == _SNAP_TI
    assert snap.tokens_output == _SNAP_TO
    assert snap.tokens_cache_read == _SNAP_TCR
    assert snap.tokens_cache_write == _SNAP_TCW


async def test_sweep_token_snapshots_skips_zero_token_agents() -> None:
    """An agent whose SDK reports all-zero tokens is skipped (no DB writes)."""
    orch = _make_orchestrator()
    instance = _make_instance(_AGENT_ID)
    instance.state = OrchestratorAgentState.ACTIVE
    orch._instances[_AGENT_ID] = instance

    def _handler(_url: str) -> Any:
        return _mock_response(
            200,
            {
                "tokens_input": 0,
                "tokens_output": 0,
                "tokens_cache_read": 0,
                "tokens_cache_write": 0,
            },
        )

    added: list[Any] = []
    db_factory = _make_db_factory(add_list=added)

    def _client_cls(**_kw: Any) -> _FakeHTTPClient:
        return _FakeHTTPClient(_handler)

    with (
        patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
        patch.object(orch, "_usage_from_transcript", return_value=(0, 0, 0, 0)),
        patch("roboco.db.base.get_session_factory", return_value=db_factory),
    ):
        await orch._sweep_token_snapshots()

    assert added == []


async def test_sweep_token_snapshots_per_agent_error_does_not_abort_loop() -> None:
    """A ConnectError for one agent is caught; the next agent is still processed."""
    orch = _make_orchestrator()

    # Agent 1: HTTP error
    inst1 = _make_instance(_AGENT_ID)
    inst1.state = OrchestratorAgentState.ACTIVE
    orch._instances[_AGENT_ID] = inst1

    # Agent 2: success with non-zero tokens
    inst2 = _make_instance(_AGENT_ID_2)
    inst2.state = OrchestratorAgentState.ACTIVE
    orch._instances[_AGENT_ID_2] = inst2

    def _handler(url: str) -> Any:
        if _AGENT_ID in url and _AGENT_ID_2 not in url:
            raise httpx.ConnectError("agent-1 unreachable")
        return _mock_response(
            200,
            {
                "tokens_input": _LOOP_TI,
                "tokens_output": _LOOP_TO,
                "tokens_cache_read": 0,
                "tokens_cache_write": 0,
            },
        )

    session_row = MagicMock()
    session_row.id = uuid4()
    added: list[Any] = []
    db_factory = _make_db_factory(session_row=session_row, add_list=added)

    def _client_cls(**_kw: Any) -> _FakeHTTPClient:
        return _FakeHTTPClient(_handler)

    with (
        patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
        patch("roboco.db.base.get_session_factory", return_value=db_factory),
    ):
        await orch._sweep_token_snapshots()

    # Only agent-2's snapshot should be present; agent-1's error was caught.
    assert len(added) == 1
    assert added[0].tokens_input == _LOOP_TI
    assert added[0].tokens_output == _LOOP_TO


# ---------------------------------------------------------------------------
# _sweep_daily_rollup — inserts new row when none exists
# ---------------------------------------------------------------------------

# Token counts for the rollup test
_ROLLUP_TI = 200
_ROLLUP_TO = 300
_ROLLUP_TCR = 20
_ROLLUP_TCW = 10


async def test_sweep_daily_rollup_inserts_new_row() -> None:
    """When no existing DailyUsageRollupTable row exists, db.add() is called
    with the correct aggregated token values."""
    orch = _make_orchestrator()

    # Build a fake aggregate result row
    agg_row = MagicMock()
    agg_row.date = "2026-06-10"
    agg_row.agent_slug = _AGENT_ID
    agg_row.team = "backend"
    agg_row.model = "sonnet"
    agg_row.tokens_input = _ROLLUP_TI
    agg_row.tokens_output = _ROLLUP_TO
    agg_row.tokens_cache_read = _ROLLUP_TCR
    agg_row.tokens_cache_write = _ROLLUP_TCW
    agg_row.total_cost_usd = 0.0
    agg_row.session_count = 1

    added: list[Any] = []
    call_count = 0

    @asynccontextmanager
    async def _db_context() -> Any:
        nonlocal call_count

        db = MagicMock()
        db.commit = AsyncMock()

        def _add(obj: Any) -> None:
            added.append(obj)

        db.add = _add

        async def _exec(_stmt: Any) -> MagicMock:
            nonlocal call_count
            call_count += 1
            result = MagicMock()
            if call_count == 1:
                # First call: aggregate SELECT — return one agg_row via fetchall()
                result.fetchall = MagicMock(return_value=[agg_row])
                result.scalar_one_or_none = MagicMock(return_value=None)
            else:
                # Second call: lookup SELECT for existing row — return None
                result.fetchall = MagicMock(return_value=[])
                result.scalar_one_or_none = MagicMock(return_value=None)
            return result

        db.execute = AsyncMock(side_effect=_exec)
        yield db

    with patch("roboco.db.base.get_session_factory", return_value=_db_context):
        await orch._sweep_daily_rollup()

    # Exactly one new DailyUsageRollupTable row must have been added
    assert len(added) == 1
    row = added[0]
    assert row.tokens_input == _ROLLUP_TI
    assert row.tokens_output == _ROLLUP_TO
    assert row.tokens_cache_read == _ROLLUP_TCR
    assert row.tokens_cache_write == _ROLLUP_TCW


# ---------------------------------------------------------------------------
# stop_agent — _finalize_spawn_session is awaited before acquiring the lock
# ---------------------------------------------------------------------------


async def test_stop_agent_finalizes_before_lock() -> None:
    """stop_agent awaits _finalize_spawn_session when the instance has a
    running container_id (the finalization must happen before the lock)."""
    orch = _make_orchestrator()
    instance = _make_instance(_AGENT_ID)
    instance.container_id = "abc123def456"  # non-None → finalize must be called
    orch._instances[_AGENT_ID] = instance

    finalized: list[str] = []

    async def _fake_finalize(agent_id: str, **_kwargs: object) -> None:
        finalized.append(agent_id)

    # Stub out the Docker subprocess so stop_agent doesn't actually run Docker
    mock_proc = MagicMock()
    mock_proc.wait = AsyncMock()

    with (
        patch.object(orch, "_finalize_spawn_session", side_effect=_fake_finalize),
        patch("asyncio.create_subprocess_exec", AsyncMock(return_value=mock_proc)),
        patch.object(orch, "_remove_container", AsyncMock()),
    ):
        await orch.stop_agent(_AGENT_ID, graceful=True)

    # _finalize_spawn_session must have been called exactly once with our agent id
    assert finalized == [_AGENT_ID]


# ---------------------------------------------------------------------------
# _handle_stopped_container — self-exits finalize (stop_agent was not called)
# ---------------------------------------------------------------------------


async def test_handle_stopped_container_graceful_finalizes() -> None:
    """A graceful self-exit (exit 0) finalizes the spawn session.

    The agent calls i_am_idle and its container exits 0 without stop_agent
    being invoked, so _handle_stopped_container must finalize to capture the
    token usage; otherwise the session row is left open with zero tokens.
    """
    orch = _make_orchestrator()
    instance = _make_instance(_AGENT_ID)
    instance.container_id = "abc123def456"
    orch._instances[_AGENT_ID] = instance

    calls: list[tuple[str, str]] = []

    async def _fake_finalize(agent_id: str, exit_reason: str = "stopped") -> None:
        calls.append((agent_id, exit_reason))

    with patch.object(orch, "_finalize_spawn_session", side_effect=_fake_finalize):
        await orch._handle_stopped_container(_AGENT_ID, instance, 0)

    assert calls == [(_AGENT_ID, "completed")]
    assert instance.state is OrchestratorAgentState.OFFLINE


async def test_handle_stopped_container_crash_finalizes_then_restarts() -> None:
    """A non-zero exit finalizes (exit_reason='crashed') before auto-restart."""
    orch = _make_orchestrator()
    instance = _make_instance(_AGENT_ID)
    instance.container_id = "abc123def456"
    instance.error_count = 0
    orch._instances[_AGENT_ID] = instance

    calls: list[tuple[str, str]] = []

    async def _fake_finalize(agent_id: str, exit_reason: str = "stopped") -> None:
        calls.append((agent_id, exit_reason))

    with (
        patch.object(orch, "_finalize_spawn_session", side_effect=_fake_finalize),
        patch.object(orch, "spawn_agent", AsyncMock()) as mock_spawn,
    ):
        await orch._handle_stopped_container(_AGENT_ID, instance, 1)

    assert calls == [(_AGENT_ID, "crashed")]
    mock_spawn.assert_awaited_once()


# ---------------------------------------------------------------------------
# _resolve_active_tokens — SDK first, transcript fallback for live agents
# ---------------------------------------------------------------------------


async def test_resolve_active_tokens_falls_back_to_transcript() -> None:
    """When the SDK reports all-zero, live resolution uses the transcript."""
    orch = _make_orchestrator()

    def _handler(_url: str) -> Any:
        return _mock_response(
            200,
            {
                "tokens_input": 0,
                "tokens_output": 0,
                "tokens_cache_read": 0,
                "tokens_cache_write": 0,
            },
        )

    client = _FakeHTTPClient(_handler)
    with patch.object(orch, "_usage_from_transcript", return_value=(6, 514, 100, 50)):
        tokens = await orch._resolve_active_tokens(client, _AGENT_ID)

    assert tokens == (6, 514, 100, 50)


async def test_resolve_active_tokens_prefers_sdk() -> None:
    """A non-zero SDK response is used directly — no transcript fallback."""
    orch = _make_orchestrator()

    def _handler(_url: str) -> Any:
        return _mock_response(
            200,
            {
                "tokens_input": 10,
                "tokens_output": 20,
                "tokens_cache_read": 0,
                "tokens_cache_write": 0,
            },
        )

    client = _FakeHTTPClient(_handler)
    with patch.object(
        orch, "_usage_from_transcript", return_value=(999, 999, 999, 999)
    ) as mock_tx:
        tokens = await orch._resolve_active_tokens(client, _AGENT_ID)

    assert tokens == (10, 20, 0, 0)
    mock_tx.assert_not_called()