/tests/runtime/test_provider_manager.py

import asyncio
from typing import cast
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from free_claude_code.application.errors import ApplicationUnavailableError
from free_claude_code.application.model_metadata import ProviderModelInfo
from free_claude_code.config.settings import Settings
from free_claude_code.providers.base import BaseProvider
from free_claude_code.providers.nvidia_nim import NvidiaNimProvider
from free_claude_code.providers.runtime import ProviderRuntime
from free_claude_code.runtime.provider_manager import ProviderRuntimeManager


class FakeRuntime(ProviderRuntime):
    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.cleanup_calls = 0
        self.cleanup_error: Exception | None = None
        self.cleanup_started: asyncio.Event | None = None
        self.cleanup_release: asyncio.Event | None = None
        self.provider = MagicMock()
        self.provider.list_model_infos = AsyncMock(return_value=frozenset())

    def is_cached(self, provider_id: str) -> bool:
        return provider_id == "cached"

    def resolve_provider(self, provider_id: str) -> BaseProvider:
        return cast(BaseProvider, self.provider)

    async def cleanup(self) -> None:
        self.cleanup_calls += 1
        if self.cleanup_started is not None:
            self.cleanup_started.set()
        if self.cleanup_release is not None:
            await self.cleanup_release.wait()
        if self.cleanup_error is not None:
            raise self.cleanup_error


class RuntimeFactory:
    def __init__(self) -> None:
        self.runtimes: list[FakeRuntime] = []
        self.error: Exception | None = None

    def __call__(self, settings: Settings) -> ProviderRuntime:
        if self.error is not None:
            raise self.error
        runtime = FakeRuntime(settings)
        self.runtimes.append(runtime)
        return runtime


def _settings(model: str) -> Settings:
    return Settings().model_copy(update={"model": model})


@pytest.mark.asyncio
async def test_startup_generation_lease_and_shutdown_close_exactly_once() -> None:
    factory = RuntimeFactory()
    settings = _settings("nvidia_nim/one")
    manager = ProviderRuntimeManager(settings, runtime_factory=factory)

    lease = await manager.acquire()

    assert lease.generation_id == 1
    assert lease.settings is settings
    assert lease.is_provider_cached("cached") is True
    assert lease.resolve_provider("nvidia_nim") is factory.runtimes[0].provider
    await lease.release()
    await lease.release()
    await manager.close()
    await manager.close()

    assert factory.runtimes[0].cleanup_calls == 1
    with pytest.raises(ApplicationUnavailableError, match="shutting down"):
        await manager.acquire()


@pytest.mark.asyncio
async def test_replacement_keeps_leased_generation_open_until_final_release() -> None:
    factory = RuntimeFactory()
    first_settings = _settings("nvidia_nim/one")
    second_settings = _settings("nvidia_nim/two")
    manager = ProviderRuntimeManager(first_settings, runtime_factory=factory)
    old_lease = await manager.acquire()
    committed: list[str] = []

    generation_id = await manager.replace(
        second_settings,
        commit=lambda: committed.append("persisted"),
    )
    new_lease = await manager.acquire()

    assert generation_id == 2
    assert committed == ["persisted"]
    assert new_lease.generation_id == 2
    assert new_lease.settings is second_settings
    assert factory.runtimes[0].cleanup_calls == 0
    await new_lease.release()
    await old_lease.release()
    assert factory.runtimes[0].cleanup_calls == 1
    await manager.close()
    assert factory.runtimes[1].cleanup_calls == 1


@pytest.mark.asyncio
async def test_hot_replacement_owns_admission_per_provider_generation() -> None:
    first_settings = _settings("nvidia_nim/one")
    second_settings = _settings("nvidia_nim/two")
    clients: list[MagicMock] = []

    def create_client(*_args: object, **_kwargs: object) -> MagicMock:
        client = MagicMock()
        client.close = AsyncMock()
        clients.append(client)
        return client

    with patch(
        "free_claude_code.providers.openai_chat.provider.AsyncOpenAI",
        side_effect=create_client,
    ):
        manager = ProviderRuntimeManager(first_settings)
        old_lease = await manager.acquire()
        old_provider = old_lease.resolve_provider("nvidia_nim")
        refresh = AsyncMock()

        with patch.object(manager, "_refresh_generation", refresh):
            await manager.replace(second_settings, commit=lambda: None)
            new_lease = await manager.acquire()
            new_provider = new_lease.resolve_provider("nvidia_nim")
            await asyncio.sleep(0)

            assert isinstance(old_provider, NvidiaNimProvider)
            assert isinstance(new_provider, NvidiaNimProvider)
            assert new_provider is not old_provider
            assert new_provider._admission is not old_provider._admission
            assert old_lease.resolve_provider("nvidia_nim") is old_provider
            clients[0].close.assert_not_awaited()

            await new_lease.release()
            await old_lease.release()

            clients[0].close.assert_awaited_once()
            clients[1].close.assert_not_awaited()
            refresh.assert_awaited_once()
            await manager.close()

        clients[1].close.assert_awaited_once()


@pytest.mark.asyncio
async def test_replacement_closes_unleased_generation_immediately() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )

    await manager.replace(
        _settings("nvidia_nim/two"),
        commit=lambda: None,
    )

    assert factory.runtimes[0].cleanup_calls == 1
    await manager.close()


@pytest.mark.asyncio
async def test_cancelled_replacement_does_not_cancel_owned_generation_cleanup() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    cleanup_started = asyncio.Event()
    cleanup_release = asyncio.Event()
    refresh_started = asyncio.Event()
    factory.runtimes[0].cleanup_started = cleanup_started
    factory.runtimes[0].cleanup_release = cleanup_release

    async def refresh(*_args: object, **_kwargs: object) -> None:
        refresh_started.set()
        await asyncio.Event().wait()

    with patch.object(manager, "_refresh_generation", side_effect=refresh):
        replace_task = asyncio.create_task(
            manager.replace(
                _settings("nvidia_nim/two"),
                commit=lambda: None,
            )
        )
        await cleanup_started.wait()
        await refresh_started.wait()

        replace_task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await replace_task

        retired = manager._retired[1]
        assert manager.current_generation_id == 2
        assert retired.cleanup_task is not None
        assert not retired.cleanup_task.cancelled()
        assert factory.runtimes[0].cleanup_calls == 1

        close_task = asyncio.create_task(manager.close())
        await asyncio.sleep(0)
        assert not close_task.done()
        cleanup_release.set()
        await close_task

    assert factory.runtimes[0].cleanup_calls == 1
    assert factory.runtimes[1].cleanup_calls == 1
    assert manager._retired == {}


@pytest.mark.asyncio
async def test_cancelled_final_lease_release_keeps_owned_cleanup_running() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    lease = await manager.acquire()
    await manager.replace(
        _settings("nvidia_nim/two"),
        commit=lambda: None,
    )
    cleanup_started = asyncio.Event()
    cleanup_release = asyncio.Event()
    factory.runtimes[0].cleanup_started = cleanup_started
    factory.runtimes[0].cleanup_release = cleanup_release

    release_task = asyncio.create_task(lease.release())
    await cleanup_started.wait()
    release_task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await release_task

    retired = manager._retired[1]
    assert retired.active_leases == 0
    assert retired.cleanup_task is not None
    assert not retired.cleanup_task.cancelled()
    assert factory.runtimes[0].cleanup_calls == 1

    close_task = asyncio.create_task(manager.close())
    await asyncio.sleep(0)
    assert not close_task.done()
    cleanup_release.set()
    await close_task

    assert factory.runtimes[0].cleanup_calls == 1
    assert manager._retired == {}


@pytest.mark.asyncio
async def test_hot_cleanup_failure_keeps_published_replacement() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    factory.runtimes[0].cleanup_error = RuntimeError("cleanup failed")

    generation_id = await manager.replace(
        _settings("nvidia_nim/two"),
        commit=lambda: None,
    )

    assert generation_id == 2
    assert manager.current_generation_id == 2
    assert factory.runtimes[0].cleanup_calls == 1
    assert 1 in manager._retired

    factory.runtimes[0].cleanup_error = None
    await manager.close()

    assert factory.runtimes[0].cleanup_calls == 2
    assert manager._retired == {}


@pytest.mark.asyncio
async def test_candidate_construction_failure_preserves_current_generation() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    factory.error = RuntimeError("cannot construct")

    with pytest.raises(RuntimeError, match="cannot construct"):
        await manager.replace(
            _settings("nvidia_nim/two"),
            commit=lambda: None,
        )

    assert manager.current_generation_id == 1
    assert manager.current_settings().model == "nvidia_nim/one"
    assert factory.runtimes[0].cleanup_calls == 0
    await manager.close()


@pytest.mark.asyncio
async def test_failed_candidate_cleanup_is_retried_at_shutdown() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )

    def fail_commit() -> None:
        factory.runtimes[1].cleanup_error = RuntimeError("private cleanup detail")
        raise OSError("disk full")

    with pytest.raises(OSError, match="disk full"):
        await manager.replace(
            _settings("nvidia_nim/two"),
            commit=fail_commit,
        )

    assert manager.current_generation_id == 1
    assert factory.runtimes[0].cleanup_calls == 0
    assert factory.runtimes[1].cleanup_calls == 1
    assert manager._unpublished == {factory.runtimes[1]}

    manager.cache_model_infos("lmstudio", {ProviderModelInfo("cached")})
    with pytest.raises(
        RuntimeError,
        match="One or more provider runtimes failed to close",
    ) as exc_info:
        await manager.close()

    assert "private cleanup detail" not in str(exc_info.value)
    assert manager._closed is False
    assert manager._unpublished == {factory.runtimes[1]}
    assert manager.cached_model_ids() == {"lmstudio": frozenset({"cached"})}

    factory.runtimes[1].cleanup_error = None
    await manager.close()

    assert factory.runtimes[0].cleanup_calls == 1
    assert factory.runtimes[1].cleanup_calls == 3
    assert manager._unpublished == set()
    assert manager.cached_model_ids() == {}
    assert manager._closed is True


@pytest.mark.asyncio
async def test_later_replacement_retries_failed_unpublished_candidate() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )

    def fail_commit() -> None:
        factory.runtimes[1].cleanup_error = RuntimeError("cleanup failed")
        raise OSError("disk full")

    with pytest.raises(OSError, match="disk full"):
        await manager.replace(
            _settings("nvidia_nim/two"),
            commit=fail_commit,
        )

    factory.runtimes[1].cleanup_error = None
    generation_id = await manager.replace(
        _settings("nvidia_nim/three"),
        commit=lambda: None,
    )

    assert generation_id == 2
    assert factory.runtimes[1].cleanup_calls == 2
    assert manager._unpublished == set()
    await manager.close()


@pytest.mark.asyncio
async def test_cancelled_candidate_cleanup_remains_owned_until_shutdown() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    cleanup_started = asyncio.Event()
    cleanup_release = asyncio.Event()

    def fail_commit() -> None:
        candidate = factory.runtimes[1]
        candidate.cleanup_started = cleanup_started
        candidate.cleanup_release = cleanup_release
        raise OSError("disk full")

    replace_task = asyncio.create_task(
        manager.replace(
            _settings("nvidia_nim/two"),
            commit=fail_commit,
        )
    )
    await cleanup_started.wait()
    replace_task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await replace_task

    assert manager._unpublished == {factory.runtimes[1]}
    assert factory.runtimes[1].cleanup_calls == 1

    cleanup_release.set()
    await manager.close()

    assert factory.runtimes[1].cleanup_calls == 2
    assert manager._unpublished == set()


@pytest.mark.asyncio
async def test_concurrent_replacements_are_serialized_in_call_order() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    first_entered = asyncio.Event()
    release_first = asyncio.Event()
    cancel_calls = 0
    original_cancel = manager._cancel_refresh

    async def controlled_cancel() -> None:
        nonlocal cancel_calls
        cancel_calls += 1
        if cancel_calls == 1:
            first_entered.set()
            await release_first.wait()
        await original_cancel()

    with patch.object(manager, "_cancel_refresh", side_effect=controlled_cancel):
        first = asyncio.create_task(
            manager.replace(
                _settings("nvidia_nim/two"),
                commit=lambda: None,
            )
        )
        await first_entered.wait()
        second = asyncio.create_task(
            manager.replace(
                _settings("nvidia_nim/three"),
                commit=lambda: None,
            )
        )
        await asyncio.sleep(0)
        assert len(factory.runtimes) == 1
        assert not second.done()
        release_first.set()
        assert await asyncio.gather(first, second) == [2, 3]

    assert manager.current_settings().model == "nvidia_nim/three"
    assert [runtime.cleanup_calls for runtime in factory.runtimes[:2]] == [1, 1]
    await manager.close()


@pytest.mark.asyncio
async def test_shutdown_waits_for_active_lease_then_rejects_new_work() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    lease = await manager.acquire()

    close_task = asyncio.create_task(manager.close())
    await asyncio.sleep(0)

    assert not close_task.done()
    with pytest.raises(ApplicationUnavailableError, match="shutting down"):
        await manager.acquire()
    await lease.release()
    await close_task
    assert factory.runtimes[0].cleanup_calls == 1


@pytest.mark.asyncio
async def test_cancelled_shutdown_retains_generation_for_retry() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    lease = await manager.acquire()
    close_task = asyncio.create_task(manager.close())
    await asyncio.sleep(0)

    close_task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await close_task

    with pytest.raises(ApplicationUnavailableError, match="shutting down"):
        await manager.acquire()
    with pytest.raises(ApplicationUnavailableError, match="shutting down"):
        await manager.replace(_settings("nvidia_nim/two"), commit=lambda: None)
    assert factory.runtimes[0].cleanup_calls == 0

    await lease.release()
    await manager.close()

    assert factory.runtimes[0].cleanup_calls == 1
    assert manager._closed is True


@pytest.mark.asyncio
async def test_cancelled_shutdown_reuses_the_same_owned_cleanup_task() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    cleanup_started = asyncio.Event()
    cleanup_release = asyncio.Event()
    cleanup_calls = 0

    async def cleanup() -> None:
        nonlocal cleanup_calls
        cleanup_calls += 1
        cleanup_started.set()
        await cleanup_release.wait()

    with patch.object(factory.runtimes[0], "cleanup", side_effect=cleanup):
        close_task = asyncio.create_task(manager.close())
        await cleanup_started.wait()

        close_task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await close_task

        assert manager._closed is False
        assert manager._retired

        retry_task = asyncio.create_task(manager.close())
        await asyncio.sleep(0)
        assert not retry_task.done()
        cleanup_release.set()
        await retry_task

    assert cleanup_calls == 1
    assert manager._retired == {}
    assert manager._closed is True


@pytest.mark.asyncio
async def test_failed_shutdown_cleanup_is_retryable() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    manager.cache_model_infos("lmstudio", {ProviderModelInfo("cached")})
    factory.runtimes[0].cleanup_error = RuntimeError("private provider detail")

    with pytest.raises(
        RuntimeError,
        match="One or more provider runtimes failed to close",
    ) as exc_info:
        await manager.close()

    assert "private provider detail" not in str(exc_info.value)
    assert manager._closed is False
    assert 1 in manager._retired
    assert manager.cached_model_ids() == {"lmstudio": frozenset({"cached"})}
    assert factory.runtimes[0].cleanup_calls == 1

    factory.runtimes[0].cleanup_error = None
    await manager.close()

    assert factory.runtimes[0].cleanup_calls == 2
    assert manager._retired == {}
    assert manager.cached_model_ids() == {}
    assert manager._closed is True


@pytest.mark.asyncio
async def test_application_catalog_survives_generation_replacement() -> None:
    factory = RuntimeFactory()
    manager = ProviderRuntimeManager(
        _settings("nvidia_nim/one"),
        runtime_factory=factory,
    )
    manager.cache_model_infos(
        "lmstudio",
        {ProviderModelInfo("persisted", supports_thinking=True)},
    )

    await manager.replace(
        _settings("nvidia_nim/two"),
        commit=lambda: None,
    )

    assert manager.cached_model_ids() == {"lmstudio": frozenset({"persisted"})}
    assert manager.cached_model_supports_thinking("lmstudio", "persisted") is True
    await manager.close()


@pytest.mark.asyncio
async def test_replacement_prunes_and_rejects_removed_remote_provider_cache() -> None:
    factory = RuntimeFactory()
    first_settings = _settings("open_router/one").model_copy(
        update={"open_router_api_key": "open-router-key"}
    )
    manager = ProviderRuntimeManager(first_settings, runtime_factory=factory)
    manager.cache_model_infos("open_router", {ProviderModelInfo("old-model")})

    await manager.replace(
        _settings("nvidia_nim/two"),
        commit=lambda: None,
    )
    manager.cache_model_infos("open_router", {ProviderModelInfo("late-old-model")})

    assert "open_router" not in manager.cached_model_ids()
    await manager.close()


@pytest.mark.asyncio
async def test_generation_lifecycle_traces_contain_minimal_correlation_fields() -> None:
    factory = RuntimeFactory()

    with patch("free_claude_code.runtime.provider_manager.trace_event") as trace:
        manager = ProviderRuntimeManager(
            _settings("nvidia_nim/one"),
            runtime_factory=factory,
        )
        lease = await manager.acquire()
        await manager.replace(
            _settings("nvidia_nim/two"),
            commit=lambda: None,
            reason="test_replace",
        )
        await lease.release()
        await manager.close()

    events = [call.kwargs for call in trace.call_args_list]
    names = [event["event"] for event in events]
    assert names == [
        "provider_generation.published",
        "provider_generation.published",
        "provider_generation.retired",
        "provider_generation.closed",
        "provider_generation.retired",
        "provider_generation.closed",
    ]
    assert events[1]["generation_id"] == 2
    assert events[1]["previous_generation_id"] == 1
    assert events[1]["reason"] == "test_replace"
    assert events[2]["active_leases"] == 1
    assert events[3]["forced"] is False