tests/unit/test_session_manager_persistence.py

"""Regression tests for server-side session persistence restore/access."""

from __future__ import annotations

import asyncio
import sys
import threading
import uuid
from datetime import UTC, datetime, timedelta
from pathlib import Path
from types import SimpleNamespace
from typing import Any

import pytest

_BACKEND_DIR = Path(__file__).resolve().parent.parent.parent / "backend"
if str(_BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(_BACKEND_DIR))

from agent.core.model_ids import KIMI_K26_MODEL_ID  # noqa: E402
from agent.core.session_persistence import NoopSessionStore  # noqa: E402
from agent.core.usage_thresholds import USAGE_THRESHOLD_TOOL_NAME  # noqa: E402
from agent.core.yolo_budget import YOLO_BUDGET_TOOL_NAME  # noqa: E402
from session_manager import (  # noqa: E402
    AgentSession,
    SessionManager,
    new_inference_billing_session_id,
)


class FakeRuntimeSession:
    def __init__(
        self,
        *,
        hf_token: str | None = None,
        user_plan: str | None = None,
        model: str = "test-model",
    ):
        self.hf_token = hf_token
        self.user_plan = user_plan
        self.context_manager = SimpleNamespace(items=[])
        self.pending_approval = None
        self.turn_count = 0
        self.config = SimpleNamespace(model_name=model)
        self.notification_destinations = []
        self.auto_approval_enabled = False
        self.auto_approval_cost_cap_usd = None
        self.auto_approval_estimated_spend_usd = 0.0
        self.usage_warning_next_threshold_usd = 5.0
        self.usage_threshold_checker = None
        self.yolo_budget_checker = None
        self.logged_events = []
        self.sandbox = None
        self.sandbox_hardware = None
        self.sandbox_preload_task = None
        self.sandbox_preload_cancel_event = None
        self.events = []
        self.session_id = "s1"

    async def send_event(self, event):
        self.events.append(event)

    def auto_approval_policy_summary(self):
        cap = self.auto_approval_cost_cap_usd
        remaining = (
            None
            if cap is None
            else max(0, cap - self.auto_approval_estimated_spend_usd)
        )
        return {
            "enabled": self.auto_approval_enabled,
            "cost_cap_usd": cap,
            "estimated_spend_usd": self.auto_approval_estimated_spend_usd,
            "remaining_usd": remaining,
        }

    def set_auto_approval_policy(self, *, enabled, cost_cap_usd):
        self.auto_approval_enabled = enabled
        self.auto_approval_cost_cap_usd = cost_cap_usd


class RestoreStore(NoopSessionStore):
    enabled = True

    def __init__(
        self,
        *,
        metadata: dict[str, Any] | None = None,
        messages: list[dict[str, Any]] | None = None,
        delay: float = 0,
    ) -> None:
        self.metadata = metadata or {
            "session_id": "persisted-session",
            "user_id": "owner",
            "model": "test-model",
            "created_at": datetime.now(UTC),
        }
        self.messages = messages or []
        self.delay = delay
        self.load_calls = 0
        self.updated_fields: list[tuple[str, dict[str, Any]]] = []

    async def load_session(self, session_id: str, **_: Any) -> dict[str, Any] | None:
        self.load_calls += 1
        if self.delay:
            await asyncio.sleep(self.delay)
        metadata = dict(self.metadata)
        metadata.setdefault("session_id", session_id)
        metadata.setdefault("_id", session_id)
        return {"metadata": metadata, "messages": self.messages}

    async def update_session_fields(self, session_id: str, **fields: Any) -> None:
        self.updated_fields.append((session_id, fields))
        self.metadata.update(fields)


class CloseableResource:
    def __init__(self) -> None:
        self.closed = False

    async def close(self) -> None:
        self.closed = True


def _manager_with_store(store: NoopSessionStore) -> SessionManager:
    manager = object.__new__(SessionManager)
    manager.config = SimpleNamespace(model_name="test-model")
    manager.sessions = {}
    manager._lock = asyncio.Lock()
    manager.persistence_store = store
    manager.messaging_gateway = CloseableResource()
    manager._pending_creates = 0
    manager._reaper_task = None
    return manager


def _runtime_agent_session(
    session_id: str,
    *,
    user_id: str = "owner",
    hf_token: str | None = "owner-token",
    user_plan: str | None = None,
) -> AgentSession:
    runtime_session = FakeRuntimeSession(hf_token=hf_token, user_plan=user_plan)
    return AgentSession(
        session_id=session_id,
        session=runtime_session,  # type: ignore[arg-type]
        tool_router=object(),  # type: ignore[arg-type]
        submission_queue=asyncio.Queue(),
        user_id=user_id,
        hf_token=hf_token,
        user_plan=user_plan,
    )


def test_inference_billing_session_id_is_uuid_for_long_agent_session_ids():
    billing_session_id = new_inference_billing_session_id(
        "s" * 300,
        datetime(2026, 6, 5, 12, 30, tzinfo=UTC),
    )

    assert str(uuid.UUID(billing_session_id)) == billing_session_id


def test_agent_session_replaces_non_uuid_inference_billing_session_id():
    runtime_session = FakeRuntimeSession()
    agent_session = AgentSession(
        session_id="s1",
        session=runtime_session,  # type: ignore[arg-type]
        tool_router=object(),  # type: ignore[arg-type]
        submission_queue=asyncio.Queue(),
        inference_billing_session_id="not-a-uuid",
    )

    assert str(uuid.UUID(agent_session.inference_billing_session_id)) == (
        agent_session.inference_billing_session_id
    )
    assert runtime_session.inference_billing_session_id == (
        agent_session.inference_billing_session_id
    )


@pytest.mark.asyncio
async def test_reset_session_usage_window_updates_runtime_and_store():
    store = RestoreStore()
    manager = _manager_with_store(store)
    agent_session = _runtime_agent_session("s1")
    manager.sessions["s1"] = agent_session
    started_at = datetime(2026, 6, 5, 12, 30, tzinfo=UTC)
    original_billing_session_id = agent_session.inference_billing_session_id
    agent_session.usage_warning_spend_cache = {"spend_usd": 12.0}

    info = await manager.reset_session_usage_window(
        "s1",
        started_at=started_at,
    )

    assert agent_session.usage_window_started_at == started_at
    assert agent_session.inference_billing_session_id is not None
    assert agent_session.inference_billing_session_id != original_billing_session_id
    assert str(uuid.UUID(agent_session.inference_billing_session_id)) == (
        agent_session.inference_billing_session_id
    )
    assert (
        agent_session.session.inference_billing_session_id
        == agent_session.inference_billing_session_id
    )
    assert agent_session.usage_warning_spend_cache == {}
    assert info is not None
    assert info["usage_window_started_at"] == started_at.isoformat()
    session_id, fields = store.updated_fields[-1]
    assert session_id == "s1"
    assert fields == {
        "usage_window_started_at": started_at,
        "inference_billing_session_id": agent_session.inference_billing_session_id,
        "last_active_at": agent_session.last_active_at,
    }


def test_usage_threshold_pending_approval_serializes_and_restores():
    manager = _manager_with_store(NoopSessionStore())
    pending = {
        "kind": USAGE_THRESHOLD_TOOL_NAME,
        "tool_call_id": "usage-threshold-1",
        "threshold_usd": 5.0,
        "current_spend_usd": 5.25,
        "next_threshold_usd": 10.0,
        "billing_source": "app_telemetry_session",
        "continuation": "continue_agent",
        "history_size": 3,
    }
    runtime = FakeRuntimeSession()
    runtime.pending_approval = pending

    assert manager._serialize_pending_approval(runtime) == [pending]
    assert manager._pending_tools_for_api(runtime) == [
        {
            "tool": USAGE_THRESHOLD_TOOL_NAME,
            "tool_call_id": "usage-threshold-1",
            "arguments": {
                "threshold_usd": 5.0,
                "current_spend_usd": 5.25,
                "next_threshold_usd": 10.0,
                "billing_source": "app_telemetry_session",
            },
        }
    ]

    restored = FakeRuntimeSession()
    manager._restore_pending_approval(restored, [pending])

    assert restored.pending_approval == pending


def test_usage_spend_prefers_hf_current_session_over_telemetry():
    spend, source = SessionManager._usage_spend_from_response(
        {
            "hf_account": {
                "current_session": {
                    "total_usd": 7.5,
                },
            },
            "session": {
                "total_usd": 2.0,
            },
        }
    )

    assert spend == 7.5
    assert source == "hf_billing_current_session"


def test_usage_spend_combines_hf_inference_with_telemetry_estimates():
    spend, source = SessionManager._usage_spend_from_response(
        {
            "hf_account": {
                "current_session": {
                    "inference_providers_usd": 7.5,
                },
            },
            "session": {
                "total_usd": 99.0,
                "hf_jobs_estimated_usd": 1.25,
                "sandbox_estimated_usd": 0.5,
            },
        }
    )

    assert spend == 9.25
    assert source == "hf_billing_current_session"


def test_usage_spend_falls_back_to_app_telemetry():
    spend, source = SessionManager._usage_spend_from_response(
        {
            "hf_account": {
                "current_session": None,
            },
            "session": {
                "total_usd": 2.0,
            },
        }
    )

    assert spend == 2.0
    assert source == "app_telemetry_session"


def test_usage_spend_falls_back_when_hf_total_is_unavailable():
    spend, source = SessionManager._usage_spend_from_response(
        {
            "hf_account": {
                "current_session": {
                    "total_usd": None,
                },
            },
            "session": {
                "total_usd": 3.5,
            },
        }
    )

    assert spend == 3.5
    assert source == "app_telemetry_session"


@pytest.mark.asyncio
async def test_refresh_usage_metrics_uses_hf_billing_plus_sandbox(monkeypatch):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1", hf_token="owner-token")
    agent_session.session.logged_events = [
        {
            "timestamp": "2026-06-01T12:00:00+00:00",
            "event_type": "llm_call",
            "data": {"cost_usd": 0.5, "total_tokens": 42},
        },
        {
            "timestamp": "2026-06-01T12:05:00+00:00",
            "event_type": "hf_job_complete",
            "data": {"estimated_cost_usd": 1.0},
        },
        {
            "timestamp": "2026-06-01T12:10:00+00:00",
            "event_type": "sandbox_create",
            "data": {"sandbox_id": "owner/sandbox-1", "hardware": "t4-small"},
        },
        {
            "timestamp": "2026-06-01T12:40:00+00:00",
            "event_type": "sandbox_destroy",
            "data": {"sandbox_id": "owner/sandbox-1", "lifetime_s": 1800},
        },
    ]

    async def fake_billing_snapshot(_manager, *, hf_token, session_id, timezone_name):
        assert hf_token == "owner-token"
        assert session_id == "s1"
        assert timezone_name == "UTC"
        return {
            "billing_scope": "account_window_delta",
            "hf_billing": {
                "source": "hf_billing_usage_v2",
                "available": True,
                "current_session": {
                    "window_start": "2026-06-01T12:00:00Z",
                    "window_end": "2026-06-01T12:40:00Z",
                    "timezone": "UTC",
                    "total_usd": 4.0,
                    "inference_providers_usd": 3.0,
                    "hf_jobs_usd": 1.0,
                    "inference_provider_requests": 6,
                    "hf_jobs_minutes": 2.0,
                    "access_token": "must-not-persist",
                },
            },
            "month": {"total_usd": 999},
            "inference_providers_credits": {"limit_usd": 999},
        }

    monkeypatch.setattr("usage.build_hf_billing_snapshot", fake_billing_snapshot)

    metrics = await manager.refresh_session_usage_metrics(agent_session)

    assert metrics["total_usd"] == 4.3
    assert metrics["total_usd_source"] == "hf_billing_plus_sandbox_estimate"
    assert metrics["app_total_usd"] == 1.8
    assert metrics["hf_billing_total_usd"] == 4.0
    assert agent_session.session.usage_metrics == metrics
    assert agent_session.session.usage_hf_billing_snapshot == {
        "billing_scope": "account_window_delta",
        "hf_billing": {
            "source": "hf_billing_usage_v2",
            "available": True,
            "error": None,
            "current_session": {
                "window_start": "2026-06-01T12:00:00Z",
                "window_end": "2026-06-01T12:40:00Z",
                "timezone": "UTC",
                "total_usd": 4.0,
                "inference_providers_usd": 3.0,
                "hf_jobs_usd": 1.0,
                "inference_provider_requests": 6,
                "hf_jobs_minutes": 2.0,
            },
        },
    }


@pytest.mark.asyncio
async def test_refresh_usage_metrics_missing_token_falls_back_to_app_telemetry():
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1", hf_token=None)
    agent_session.session.hf_token = None
    agent_session.session.logged_events = [
        {
            "timestamp": "2026-06-01T12:00:00+00:00",
            "event_type": "llm_call",
            "data": {"cost_usd": 2.0, "total_tokens": 10},
        }
    ]

    metrics = await manager.refresh_session_usage_metrics(agent_session)

    assert metrics["total_usd"] == 2.0
    assert metrics["total_usd_source"] == "app_telemetry_fallback"
    assert metrics["hf_billing"] == {
        "source": "hf_billing_usage_v2",
        "available": False,
        "error": "missing_hf_token",
        "current_session": None,
    }


@pytest.mark.asyncio
async def test_refresh_usage_metrics_failure_records_error_code(monkeypatch):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1", hf_token="owner-token")
    agent_session.session.logged_events = [
        {
            "timestamp": "2026-06-01T12:00:00+00:00",
            "event_type": "llm_call",
            "data": {"cost_usd": 2.0, "total_tokens": 10},
        }
    ]

    async def fail_billing_snapshot(*_args, **_kwargs):
        raise RuntimeError("boom")

    monkeypatch.setattr("usage.build_hf_billing_snapshot", fail_billing_snapshot)

    metrics = await manager.refresh_session_usage_metrics(
        agent_session,
        error_code="unit_billing_error",
    )

    assert metrics["total_usd"] == 2.0
    assert metrics["total_usd_source"] == "app_telemetry_fallback"
    assert metrics["hf_billing"] == {
        "source": "hf_billing_usage_v2",
        "available": False,
        "error": "unit_billing_error",
        "current_session": None,
    }


@pytest.mark.asyncio
async def test_refresh_usage_metrics_timeout_records_error_code(monkeypatch):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1", hf_token="owner-token")
    agent_session.session.logged_events = [
        {
            "timestamp": "2026-06-01T12:00:00+00:00",
            "event_type": "llm_call",
            "data": {"cost_usd": 2.0, "total_tokens": 10},
        }
    ]

    async def slow_billing_snapshot(*_args, **_kwargs):
        await asyncio.sleep(0.05)
        return {
            "hf_billing": {
                "available": True,
                "current_session": {"total_usd": 999},
            }
        }

    monkeypatch.setattr("usage.build_hf_billing_snapshot", slow_billing_snapshot)

    metrics = await manager.refresh_session_usage_metrics(
        agent_session,
        error_code="unit_billing_timeout",
        billing_timeout_s=0.001,
    )

    assert metrics["total_usd"] == 2.0
    assert metrics["total_usd_source"] == "app_telemetry_fallback"
    assert metrics["hf_billing"] == {
        "source": "hf_billing_usage_v2",
        "available": False,
        "error": "unit_billing_timeout",
        "current_session": None,
    }


@pytest.mark.asyncio
async def test_usage_threshold_checker_creates_synthetic_pending_approval(monkeypatch):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    agent_session.session.logged_events.append(
        {
            "timestamp": datetime.now(UTC).isoformat(),
            "event_type": "llm_call",
            "data": {"cost_usd": 6.0},
        }
    )
    manager.sessions["s1"] = agent_session

    async def fake_current_session_usage_spend(_agent_session, **_kwargs):
        return 12.25, "app_telemetry_session"

    monkeypatch.setattr(
        manager,
        "_current_session_usage_spend",
        fake_current_session_usage_spend,
    )
    manager._install_usage_threshold_checker(agent_session)

    paused = await agent_session.session.usage_threshold_checker(
        {"continuation": "continue_agent", "history_size": 4}
    )

    assert paused is True
    pending = agent_session.session.pending_approval
    assert pending["kind"] == USAGE_THRESHOLD_TOOL_NAME
    assert pending["threshold_usd"] == 5.0
    assert pending["current_spend_usd"] == 12.25
    assert pending["next_threshold_usd"] == 20.0
    assert pending["billing_source"] == "app_telemetry_session"
    assert agent_session.session.events[-1].event_type == "approval_required"


@pytest.mark.asyncio
async def test_usage_threshold_checker_skips_billing_when_local_spend_below_threshold(
    monkeypatch,
):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    manager.sessions["s1"] = agent_session
    calls = 0

    async def fake_current_session_usage_spend(_agent_session, **_kwargs):
        nonlocal calls
        calls += 1
        return 12.25, "hf_billing_current_session"

    monkeypatch.setattr(
        manager,
        "_current_session_usage_spend",
        fake_current_session_usage_spend,
    )
    manager._install_usage_threshold_checker(agent_session)

    paused = await agent_session.session.usage_threshold_checker(
        {"continuation": "continue_agent", "history_size": 4}
    )

    assert paused is False
    assert calls == 0
    assert agent_session.session.pending_approval is None


@pytest.mark.asyncio
async def test_usage_threshold_checker_forces_billing_at_complete_turn(monkeypatch):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    manager.sessions["s1"] = agent_session
    calls = 0

    async def fake_current_session_usage_spend(_agent_session, **_kwargs):
        nonlocal calls
        calls += 1
        return 12.25, "hf_billing_current_session"

    monkeypatch.setattr(
        manager,
        "_current_session_usage_spend",
        fake_current_session_usage_spend,
    )
    manager._install_usage_threshold_checker(agent_session)

    paused = await agent_session.session.usage_threshold_checker(
        {
            "continuation": "complete_turn",
            "force_check": True,
            "history_size": 4,
        }
    )

    assert paused is True
    assert calls == 1
    assert agent_session.session.pending_approval["continuation"] == "complete_turn"


@pytest.mark.asyncio
async def test_usage_threshold_checker_uses_cached_spend_after_local_crossing(
    monkeypatch,
):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    agent_session.session.logged_events.append(
        {
            "timestamp": datetime.now(UTC).isoformat(),
            "event_type": "llm_call",
            "data": {"cost_usd": 6.0},
        }
    )
    manager.sessions["s1"] = agent_session

    async def fail_build_usage_response(*_args, **_kwargs):
        raise AssertionError("fresh billing usage should not be fetched")

    monkeypatch.setattr("usage.build_usage_response", fail_build_usage_response)
    agent_session.usage_warning_spend_cache = {
        "spend_usd": 4.5,
        "billing_source": "hf_billing_current_session",
        "expires_at": datetime.now(UTC) + timedelta(seconds=30),
    }
    manager._install_usage_threshold_checker(agent_session)

    first = await agent_session.session.usage_threshold_checker(
        {"continuation": "continue_agent", "history_size": 4}
    )
    second = await agent_session.session.usage_threshold_checker(
        {"continuation": "continue_agent", "history_size": 4}
    )

    assert first is False
    assert second is False
    assert agent_session.session.pending_approval is None


@pytest.mark.asyncio
async def test_yolo_budget_checker_pauses_after_fresh_session_spend_crosses_cap(
    monkeypatch,
):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    agent_session.session.auto_approval_enabled = True
    agent_session.session.auto_approval_cost_cap_usd = 1.0
    agent_session.session.auto_approval_estimated_spend_usd = 0.5
    manager.sessions["s1"] = agent_session
    calls = 0

    async def fake_current_session_usage_spend(_agent_session, **kwargs):
        nonlocal calls
        calls += 1
        assert kwargs["use_cache"] is False
        return 1.25, "hf_billing_current_session"

    monkeypatch.setattr(
        manager,
        "_current_session_usage_spend",
        fake_current_session_usage_spend,
    )
    manager._install_yolo_budget_checker(agent_session)

    paused = await agent_session.session.yolo_budget_checker(
        {"spend_kind": "llm_call", "history_size": 4}
    )

    assert paused is True
    assert calls == 1
    assert agent_session.session.auto_approval_estimated_spend_usd == 1.25
    pending = agent_session.session.pending_approval
    assert pending["kind"] == YOLO_BUDGET_TOOL_NAME
    assert pending["current_spend_usd"] == 1.25
    assert pending["cap_usd"] == 1.0
    assert pending["estimated_next_usd"] is None
    assert pending["billing_source"] == "hf_billing_current_session"
    assert agent_session.session.events[-1].event_type == "approval_required"


@pytest.mark.asyncio
async def test_yolo_budget_checker_emits_session_update_under_cap(monkeypatch):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    agent_session.session.auto_approval_enabled = True
    agent_session.session.auto_approval_cost_cap_usd = 1.0
    agent_session.session.auto_approval_estimated_spend_usd = 0.0
    manager.sessions["s1"] = agent_session

    async def fake_current_session_usage_spend(_agent_session, **kwargs):
        assert kwargs["use_cache"] is False
        return 0.25, "hf_billing_current_session"

    monkeypatch.setattr(
        manager,
        "_current_session_usage_spend",
        fake_current_session_usage_spend,
    )
    manager._install_yolo_budget_checker(agent_session)

    paused = await agent_session.session.yolo_budget_checker(
        {"spend_kind": "llm_call", "history_size": 4}
    )

    assert paused is False
    assert agent_session.session.auto_approval_estimated_spend_usd == 0.25
    assert agent_session.session.pending_approval is None
    event = agent_session.session.events[-1]
    assert event.event_type == "session_update"
    assert event.data == {
        "session_id": "s1",
        "auto_approval": {
            "enabled": True,
            "cost_cap_usd": 1.0,
            "estimated_spend_usd": 0.25,
            "remaining_usd": 0.75,
        },
    }


@pytest.mark.asyncio
async def test_yolo_budget_checker_emits_session_update_for_observed_spend(
    monkeypatch,
):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    agent_session.session.auto_approval_enabled = True
    agent_session.session.auto_approval_cost_cap_usd = 1.0
    agent_session.session.auto_approval_estimated_spend_usd = 0.03
    manager.sessions["s1"] = agent_session

    async def fake_current_session_usage_spend(_agent_session, **kwargs):
        assert kwargs["use_cache"] is False
        return 0.0, "hf_billing_current_session"

    monkeypatch.setattr(
        manager,
        "_current_session_usage_spend",
        fake_current_session_usage_spend,
    )
    manager._install_yolo_budget_checker(agent_session)

    paused = await agent_session.session.yolo_budget_checker(
        {
            "spend_kind": "llm_call",
            "observed_cost_usd": 0.03,
            "history_size": 4,
        }
    )

    assert paused is False
    event = agent_session.session.events[-1]
    assert event.event_type == "session_update"
    assert event.data["auto_approval"] == {
        "enabled": True,
        "cost_cap_usd": 1.0,
        "estimated_spend_usd": 0.03,
        "remaining_usd": 0.97,
    }


@pytest.mark.asyncio
async def test_usage_fetch_reconciles_yolo_ledger_from_hf_billing():
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    agent_session.session.auto_approval_enabled = True
    agent_session.session.auto_approval_cost_cap_usd = 1.0
    agent_session.session.auto_approval_estimated_spend_usd = 0.03
    manager.sessions["s1"] = agent_session

    summary = await manager.reconcile_session_auto_approval_from_usage(
        "s1",
        {
            "session": {
                "hf_jobs_estimated_usd": 0.0,
                "sandbox_estimated_usd": 0.0,
            },
            "hf_account": {
                "current_session": {
                    "inference_providers_usd": 0.16,
                }
            },
        },
    )

    assert summary == {
        "enabled": True,
        "cost_cap_usd": 1.0,
        "estimated_spend_usd": 0.16,
        "remaining_usd": 0.84,
    }
    assert agent_session.session.auto_approval_estimated_spend_usd == 0.16


@pytest.mark.asyncio
async def test_yolo_budget_checker_uses_local_ledger_when_billing_lags(monkeypatch):
    manager = _manager_with_store(NoopSessionStore())
    agent_session = _runtime_agent_session("s1")
    agent_session.session.auto_approval_enabled = True
    agent_session.session.auto_approval_cost_cap_usd = 1.0
    agent_session.session.auto_approval_estimated_spend_usd = 1.1
    manager.sessions["s1"] = agent_session

    async def fake_current_session_usage_spend(_agent_session, **kwargs):
        assert kwargs["use_cache"] is False
        return 0.75, "hf_billing_current_session"

    monkeypatch.setattr(
        manager,
        "_current_session_usage_spend",
        fake_current_session_usage_spend,
    )
    manager._install_yolo_budget_checker(agent_session)

    paused = await agent_session.session.yolo_budget_checker(
        {"spend_kind": "llm_call", "history_size": 4}
    )

    assert paused is True
    assert agent_session.session.auto_approval_estimated_spend_usd == 1.1
    pending = agent_session.session.pending_approval
    assert pending["kind"] == YOLO_BUDGET_TOOL_NAME
    assert pending["current_spend_usd"] == 1.1
    assert pending["cap_usd"] == 1.0
    assert pending["billing_source"] == "hf_billing_current_session"


def test_unknown_saved_model_defaults_to_kimi():
    model = SessionManager._model_from_saved_metadata("unsupported/model")

    assert model == KIMI_K26_MODEL_ID


@pytest.mark.asyncio
async def test_update_session_auto_approval_defaults_to_five_dollars():
    manager = _manager_with_store(NoopSessionStore())
    existing = _runtime_agent_session("s1", user_id="owner")
    manager.sessions["s1"] = existing

    summary = await manager.update_session_auto_approval(
        "s1",
        enabled=True,
        cost_cap_usd=None,
        cap_provided=False,
    )

    assert summary["enabled"] is True
    assert summary["cost_cap_usd"] == 5.0
    assert summary["remaining_usd"] == 5.0


def _install_fake_runtime(manager: SessionManager) -> asyncio.Event:
    stop = asyncio.Event()
    manager.run_calls = 0  # type: ignore[attr-defined]

    def fake_create_session_sync(**kwargs: Any):
        return object(), FakeRuntimeSession(
            hf_token=kwargs.get("hf_token"),
            user_plan=kwargs.get("user_plan"),
            model=kwargs.get("model") or "test-model",
        )

    async def fake_run_session(*_: Any) -> None:
        manager.run_calls += 1  # type: ignore[attr-defined]
        await stop.wait()

    manager._create_session_sync = fake_create_session_sync  # type: ignore[method-assign]
    manager._run_session = fake_run_session  # type: ignore[method-assign]
    return stop


async def _cancel_runtime_tasks(manager: SessionManager) -> None:
    tasks = [
        agent_session.task
        for agent_session in manager.sessions.values()
        if agent_session.task and not agent_session.task.done()
    ]
    for task in tasks:
        task.cancel()
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True)


@pytest.mark.asyncio
async def test_close_cancels_preload_and_deletes_owned_sandbox(monkeypatch):
    deleted: list[str] = []

    async def fake_record_sandbox_destroy(*args, **kwargs):
        pass

    monkeypatch.setattr(
        "agent.core.telemetry.record_sandbox_destroy",
        fake_record_sandbox_destroy,
    )

    store = NoopSessionStore()
    manager = _manager_with_store(store)
    gateway = CloseableResource()
    persistence = CloseableResource()
    manager.messaging_gateway = gateway  # type: ignore[assignment]
    manager.persistence_store = persistence  # type: ignore[assignment]

    cancel_event = asyncio.Event()
    preload_cancel_event = threading.Event()

    async def preload():
        while not preload_cancel_event.is_set():
            await asyncio.sleep(0)
        cancel_event.set()

    session = FakeRuntimeSession(hf_token="token")
    session.session_id = "s1"
    session.persistence_store = NoopSessionStore()
    session.sandbox = SimpleNamespace(
        space_id="owner/sandbox-12345678",
        _owns_space=True,
        delete=lambda log=None: deleted.append("owner/sandbox-12345678"),
    )
    session.sandbox_hardware = "cpu-basic"
    session.sandbox_preload_cancel_event = preload_cancel_event
    session.sandbox_preload_task = asyncio.create_task(preload())
    manager.sessions["s1"] = AgentSession(
        session_id="s1",
        session=session,  # type: ignore[arg-type]
        tool_router=object(),  # type: ignore[arg-type]
        submission_queue=asyncio.Queue(),
        user_id="owner",
        hf_token="token",
    )

    await manager.close()

    assert preload_cancel_event.is_set()
    assert cancel_event.is_set()
    assert deleted == ["owner/sandbox-12345678"]
    assert gateway.closed is True
    assert persistence.closed is True


@pytest.mark.asyncio
async def test_close_closes_resources_when_sandbox_cleanup_fails():
    manager = _manager_with_store(NoopSessionStore())
    gateway = CloseableResource()
    persistence = CloseableResource()
    manager.messaging_gateway = gateway  # type: ignore[assignment]
    manager.persistence_store = persistence  # type: ignore[assignment]
    manager.sessions["s1"] = _runtime_agent_session("s1")
    manager.sessions["s2"] = _runtime_agent_session("s2")
    cleaned: list[str] = []

    async def fake_cleanup(session):
        cleaned.append(session.hf_token)
        if session.hf_token == "owner-token":
            raise RuntimeError("boom")

    manager._cleanup_sandbox = fake_cleanup  # type: ignore[method-assign]

    await manager.close()

    assert cleaned == ["owner-token", "owner-token"]
    assert gateway.closed is True
    assert persistence.closed is True


@pytest.mark.asyncio
async def test_existing_session_rejects_cross_user_token_overwrite():
    manager = _manager_with_store(NoopSessionStore())
    existing = _runtime_agent_session("s1", user_id="victim", hf_token="victim-token")
    manager.sessions["s1"] = existing

    result = await manager.ensure_session_loaded(
        "s1", user_id="attacker", hf_token="attacker-token"
    )

    assert result is None
    assert existing.hf_token == "victim-token"
    assert existing.session.hf_token == "victim-token"


@pytest.mark.asyncio
async def test_existing_session_updates_token_after_access_check():
    manager = _manager_with_store(NoopSessionStore())
    existing = _runtime_agent_session("s1", user_id="owner", hf_token="old-token")
    manager.sessions["s1"] = existing

    result = await manager.ensure_session_loaded(
        "s1", user_id="owner", hf_token="new-token"
    )

    assert result is existing
    assert existing.hf_token == "new-token"
    assert existing.session.hf_token == "new-token"


@pytest.mark.asyncio
async def test_existing_session_updates_plan_after_access_check():
    manager = _manager_with_store(NoopSessionStore())
    existing = _runtime_agent_session("s1", user_id="owner", user_plan="free")
    manager.sessions["s1"] = existing

    result = await manager.ensure_session_loaded("s1", user_id="owner", user_plan="pro")

    assert result is existing
    assert existing.user_plan == "pro"
    assert existing.session.user_plan == "pro"


@pytest.mark.asyncio
async def test_existing_session_retries_preload_after_token_recovered():
    manager = _manager_with_store(NoopSessionStore())
    existing = _runtime_agent_session("s1", user_id="owner", hf_token=None)
    done_task = asyncio.get_running_loop().create_future()
    done_task.set_result(None)
    existing.session.sandbox_preload_task = done_task
    existing.session.sandbox_preload_error = (
        "No HF token available. Cannot create sandbox."
    )
    manager.sessions["s1"] = existing
    started: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session):
        started.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    result = await manager.ensure_session_loaded(
        "s1",
        user_id="owner",
        hf_token="new-token",
    )

    assert result is existing
    assert existing.hf_token == "new-token"
    assert existing.session.hf_token == "new-token"
    assert existing.session.sandbox_preload_error is None
    assert existing.session.sandbox_preload_task is None
    assert started == ["s1"]


@pytest.mark.asyncio
async def test_existing_session_does_not_retry_preload_when_disabled():
    manager = _manager_with_store(NoopSessionStore())
    existing = _runtime_agent_session("s1", user_id="owner", hf_token=None)
    done_task = asyncio.get_running_loop().create_future()
    done_task.set_result(None)
    existing.session.sandbox_preload_task = done_task
    existing.session.sandbox_preload_error = (
        "No HF token available. Cannot create sandbox."
    )
    manager.sessions["s1"] = existing
    started: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session):
        started.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    result = await manager.ensure_session_loaded(
        "s1",
        user_id="owner",
        hf_token="new-token",
        preload_sandbox=False,
    )

    assert result is existing
    assert existing.hf_token == "new-token"
    assert existing.session.hf_token == "new-token"
    assert existing.session.sandbox_preload_error == (
        "No HF token available. Cannot create sandbox."
    )
    assert started == []


@pytest.mark.asyncio
async def test_existing_session_does_not_restart_preload_after_teardown():
    manager = _manager_with_store(NoopSessionStore())
    existing = _runtime_agent_session("s1", user_id="owner", hf_token="token")
    done_task = asyncio.get_running_loop().create_future()
    done_task.set_result(None)
    existing.session.sandbox = None
    existing.session.sandbox_preload_task = done_task
    existing.session.sandbox_preload_error = None
    manager.sessions["s1"] = existing
    started: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session):
        started.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    result = await manager.ensure_session_loaded(
        "s1",
        user_id="owner",
        hf_token="token",
    )

    assert result is existing
    assert existing.session.sandbox_preload_task is done_task
    assert existing.session.sandbox_preload_error is None
    assert started == []


@pytest.mark.asyncio
async def test_concurrent_lazy_restore_starts_only_one_agent_task():
    store = RestoreStore(delay=0.01)
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)
    scheduled: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session: AgentSession) -> None:
        scheduled.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    try:
        first, second = await asyncio.gather(
            manager.ensure_session_loaded("persisted-session", user_id="owner"),
            manager.ensure_session_loaded("persisted-session", user_id="owner"),
        )
        await asyncio.sleep(0)

        assert first is second
        assert list(manager.sessions) == ["persisted-session"]
        assert manager.run_calls == 1  # type: ignore[attr-defined]
        assert scheduled == ["persisted-session"]
        assert not stop.is_set()
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_create_session_schedules_cpu_sandbox_preload():
    manager = _manager_with_store(NoopSessionStore())
    stop = _install_fake_runtime(manager)
    scheduled: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session: AgentSession) -> None:
        scheduled.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    try:
        session_id = await manager.create_session(
            user_id="owner",
            hf_token="token",
            user_plan="pro",
        )

        assert scheduled == [session_id]
        assert session_id in manager.sessions
        assert manager.sessions[session_id].user_plan == "pro"
        runtime_session = manager.sessions[session_id].session
        assert runtime_session.user_plan == "pro"
        assert not hasattr(runtime_session, "_ml_intern_artifact_collection_task")
        assert not hasattr(runtime_session, "_ml_intern_artifact_collection_slug")
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_lazy_restore_schedules_cpu_sandbox_preload():
    manager = _manager_with_store(RestoreStore())
    stop = _install_fake_runtime(manager)
    scheduled: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session: AgentSession) -> None:
        scheduled.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    try:
        restored = await manager.ensure_session_loaded(
            "persisted-session",
            user_id="owner",
            user_plan="free",
        )

        assert restored is not None
        assert scheduled == ["persisted-session"]
        assert "persisted-session" in manager.sessions
        assert restored.user_plan == "free"
        assert restored.session.user_plan == "free"
        assert not hasattr(restored.session, "_ml_intern_artifact_collection_task")
        assert not hasattr(restored.session, "_ml_intern_artifact_collection_slug")
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_lazy_restore_deletes_persisted_sandbox_before_preload(monkeypatch):
    deleted: list[tuple[str, str, str]] = []

    class FakeApi:
        def __init__(self, token=None):
            self.token = token

        def delete_repo(self, repo_id, repo_type):
            deleted.append((self.token, repo_id, repo_type))

    monkeypatch.setattr("huggingface_hub.HfApi", FakeApi)

    store = RestoreStore(
        metadata={
            "session_id": "persisted-session",
            "user_id": "owner",
            "model": "test-model",
            "created_at": datetime.now(UTC),
            "sandbox_space_id": "owner/sandbox-12345678",
            "sandbox_hardware": "cpu-basic",
            "sandbox_owner": "owner",
            "sandbox_created_at": datetime.now(UTC),
            "sandbox_status": "active",
        }
    )
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)
    scheduled: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session: AgentSession) -> None:
        scheduled.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    try:
        restored = await manager.ensure_session_loaded(
            "persisted-session",
            user_id="owner",
            hf_token="user-token",
        )

        assert restored is not None
        assert deleted == [("user-token", "owner/sandbox-12345678", "space")]
        assert scheduled == ["persisted-session"]
        assert store.metadata["sandbox_space_id"] is None
        assert store.metadata["sandbox_status"] == "destroyed"
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_lazy_restore_can_skip_cpu_sandbox_preload_after_cleanup(monkeypatch):
    deleted: list[str] = []

    class FakeApi:
        def __init__(self, token=None):
            self.token = token

        def delete_repo(self, repo_id, repo_type):
            deleted.append(repo_id)

    monkeypatch.setattr("huggingface_hub.HfApi", FakeApi)

    store = RestoreStore(
        metadata={
            "session_id": "persisted-session",
            "user_id": "owner",
            "model": "test-model",
            "created_at": datetime.now(UTC),
            "sandbox_space_id": "owner/sandbox-87654321",
            "sandbox_status": "active",
        }
    )
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)
    scheduled: list[str] = []

    def fake_start_cpu_sandbox_preload(agent_session: AgentSession) -> None:
        scheduled.append(agent_session.session_id)

    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]

    try:
        restored = await manager.ensure_session_loaded(
            "persisted-session",
            user_id="owner",
            hf_token="user-token",
            preload_sandbox=False,
        )

        assert restored is not None
        assert deleted == ["owner/sandbox-87654321"]
        assert scheduled == []
        assert store.metadata["sandbox_space_id"] is None
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_lazy_restore_preserves_pending_approval_tool_calls():
    store = RestoreStore(
        metadata={
            "session_id": "approval-session",
            "user_id": "owner",
            "model": "test-model",
            "pending_approval": [
                {
                    "id": "call_123",
                    "type": "function",
                    "function": {
                        "name": "create_file",
                        "arguments": '{"path":"app.py"}',
                    },
                }
            ],
        }
    )
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)

    try:
        restored = await manager.ensure_session_loaded(
            "approval-session", user_id="owner"
        )

        assert restored is not None
        tool_calls = restored.session.pending_approval["tool_calls"]
        assert len(tool_calls) == 1
        assert tool_calls[0].id == "call_123"
        assert tool_calls[0].function.name == "create_file"
        assert tool_calls[0].function.arguments == '{"path":"app.py"}'
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_lazy_restore_preserves_auto_approval_policy():
    store = RestoreStore(
        metadata={
            "session_id": "yolo-session",
            "user_id": "owner",
            "model": "test-model",
            "auto_approval_enabled": True,
            "auto_approval_cost_cap_usd": 5.0,
            "auto_approval_estimated_spend_usd": 1.25,
        }
    )
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)

    try:
        restored = await manager.ensure_session_loaded("yolo-session", user_id="owner")

        assert restored is not None
        assert restored.session.auto_approval_enabled is True
        assert restored.session.auto_approval_cost_cap_usd == 5.0
        assert restored.session.auto_approval_estimated_spend_usd == 1.25
        assert restored.session.auto_approval_policy_summary()["remaining_usd"] == 3.75
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_update_session_auto_approval_seeds_existing_session_usage(monkeypatch):
    store = RestoreStore()
    manager = _manager_with_store(store)
    agent_session = _runtime_agent_session("seed-yolo")
    manager.sessions["seed-yolo"] = agent_session

    async def fake_current_spend(*args, **kwargs):
        assert kwargs["use_cache"] is False
        return 2.75, "app_telemetry_session"

    async def fake_persist(*args, **kwargs):
        return None

    monkeypatch.setattr(manager, "_current_session_usage_spend", fake_current_spend)
    monkeypatch.setattr(manager, "persist_session_snapshot", fake_persist)

    summary = await manager.update_session_auto_approval(
        "seed-yolo",
        enabled=True,
        cost_cap_usd=None,
        cap_provided=False,
    )

    assert summary == {
        "enabled": True,
        "cost_cap_usd": 5.0,
        "estimated_spend_usd": 2.75,
        "remaining_usd": 2.25,
    }


@pytest.mark.asyncio
async def test_lazy_restore_injects_sandbox_reset_note_when_session_had_sandbox():
    store = RestoreStore(
        metadata={
            "session_id": "had-sandbox",
            "user_id": "owner",
            "model": "test-model",
            "sandbox_status": "destroyed",
        }
    )
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)

    try:
        restored = await manager.ensure_session_loaded("had-sandbox", user_id="owner")

        assert restored is not None
        items = restored.session.context_manager.items
        assert len(items) == 1
        assert "sandbox was reset" in items[0].content
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_lazy_restore_skips_sandbox_reset_note_when_no_sandbox():
    store = RestoreStore(
        metadata={
            "session_id": "no-sandbox",
            "user_id": "owner",
            "model": "test-model",
        }
    )
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)

    try:
        restored = await manager.ensure_session_loaded("no-sandbox", user_id="owner")

        assert restored is not None
        assert restored.session.context_manager.items == []
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_lazy_restore_skips_sandbox_note_when_pending_approval():
    """The sandbox-reset note must be skipped when an approval is pending: it
    would land between the restored assistant tool-calls and their results,
    orphaning the tool results on approval (the provider rejects the ordering)."""
    store = RestoreStore(
        metadata={
            "session_id": "had-sandbox-pending",
            "user_id": "owner",
            "model": "test-model",
            "sandbox_status": "destroyed",
            "pending_approval": [
                {"tool": "bash", "tool_call_id": "call_1", "arguments": {}}
            ],
        }
    )
    manager = _manager_with_store(store)
    stop = _install_fake_runtime(manager)

    try:
        restored = await manager.ensure_session_loaded(
            "had-sandbox-pending", user_id="owner"
        )

        assert restored is not None
        items = restored.session.context_manager.items
        assert not any("sandbox was reset" in getattr(m, "content", "") for m in items)
        # The pending approval itself is still restored.
        assert restored.session.pending_approval is not None
    finally:
        stop.set()
        await _cancel_runtime_tasks(manager)


@pytest.mark.asyncio
async def test_list_sessions_dev_uses_store_dev_visibility():
    class ListStore(NoopSessionStore):
        enabled = True

        def __init__(self) -> None:
            self.seen_user_id: str | None = None

        async def list_sessions(self, user_id: str, **_: Any) -> list[dict[str, Any]]:
            self.seen_user_id = user_id
            if user_id == "dev":
                return [
                    {
                        "session_id": "s1",
                        "user_id": "alice",
                        "model": "m",
                        "created_at": datetime.now(UTC),
                        "auto_approval_enabled": True,
                        "auto_approval_cost_cap_usd": 5.0,
                        "auto_approval_estimated_spend_usd": 2.0,
                    },
                    {
                        "session_id": "s2",
                        "user_id": "bob",
                        "model": "m",
                        "created_at": datetime.now(UTC),
                    },
                ]
            return []

    store = ListStore()
    manager = _manager_with_store(store)

    sessions = await manager.list_sessions(user_id="dev")

    assert store.seen_user_id == "dev"
    assert {session["session_id"] for session in sessions} == {"s1", "s2"}
    yolo = next(session for session in sessions if session["session_id"] == "s1")
    assert yolo["auto_approval"] == {
        "enabled": True,
        "cost_cap_usd": 5.0,
        "estimated_spend_usd": 2.0,
        "remaining_usd": 3.0,
    }