"""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 GLM_52_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, } @pytest.mark.asyncio async def test_activate_session_preserves_usage_window_and_billing_id(): store = RestoreStore() manager = _manager_with_store(store) agent_session = _runtime_agent_session("s1") manager.sessions["s1"] = agent_session usage_window_started_at = datetime(2026, 6, 5, 12, 30, tzinfo=UTC) billing_session_id = agent_session.inference_billing_session_id old_last_active_at = datetime(2000, 1, 1) agent_session.usage_window_started_at = usage_window_started_at agent_session.last_active_at = old_last_active_at agent_session.usage_warning_spend_cache = {"spend_usd": 12.0} info = await manager.activate_session("s1") assert agent_session.usage_window_started_at == usage_window_started_at assert agent_session.inference_billing_session_id == billing_session_id assert agent_session.session.inference_billing_session_id == billing_session_id assert agent_session.usage_warning_spend_cache == {"spend_usd": 12.0} assert agent_session.last_active_at > old_last_active_at assert info is not None assert info["usage_window_started_at"] == usage_window_started_at.isoformat() session_id, fields = store.updated_fields[-1] assert session_id == "s1" assert fields == {"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_glm(): model = SessionManager._model_from_saved_metadata("unsupported/model") assert model == GLM_52_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, }