agent-zero/tests/test_model_config_api_keys.py

import sys
import threading
import types
from pathlib import Path

import pytest
from flask import Flask


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

sys.modules["giturlparse"] = types.SimpleNamespace(parse=lambda *args, **kwargs: None)
sys.modules["whisper"] = types.SimpleNamespace(load_model=lambda *args, **kwargs: None)


class _DummyObserver:
    def __init__(self):
        self._alive = False

    def is_alive(self):
        return self._alive

    def start(self):
        self._alive = True

    def stop(self):
        self._alive = False

    def join(self, *args, **kwargs):
        return None

    def unschedule_all(self):
        return None

    def schedule(self, *args, **kwargs):
        return None


watchdog = types.ModuleType("watchdog")
watchdog.observers = types.SimpleNamespace(Observer=_DummyObserver)
watchdog.events = types.SimpleNamespace(FileSystemEventHandler=object)
sys.modules["watchdog"] = watchdog
sys.modules["watchdog.observers"] = watchdog.observers
sys.modules["watchdog.events"] = watchdog.events

from plugins._model_config.api.api_keys import ApiKeys
from plugins._model_config.extensions.python.banners import _20_missing_api_key as missing_key_banner
import models


def test_model_config_api_keys_can_be_cleared_via_backend(monkeypatch, tmp_path):
    from helpers import dotenv

    env_file = tmp_path / ".env"
    monkeypatch.setattr(dotenv, "get_dotenv_file_path", lambda: str(env_file))

    for key in ("API_KEY_OPENROUTER", "OPENROUTER_API_KEY", "OPENROUTER_API_TOKEN"):
        monkeypatch.delenv(key, raising=False)

    handler = ApiKeys(Flask(__name__), threading.Lock())

    assert handler._set_keys({"keys": {"openrouter": "sk-test-openrouter"}}) == {"ok": True}
    assert models.get_api_key("openrouter") == "sk-test-openrouter"

    assert handler._set_keys({"keys": {"openrouter": ""}}) == {"ok": True}
    assert models.get_api_key("openrouter") == "None"
    assert handler._reveal_key({"provider": "openrouter"}) == {"ok": True, "value": ""}


@pytest.mark.asyncio
async def test_missing_api_key_banner_exposes_missing_providers(monkeypatch):
    from plugins._model_config.helpers import model_config

    fake = [{"model_type": "Chat Model", "provider": "openai"}]
    monkeypatch.setattr(model_config, "get_missing_api_key_providers", lambda: fake)

    banners = []
    await missing_key_banner.MissingApiKeyCheck(agent=None).execute(
        banners=banners, frontend_context={}
    )
    row = next(b for b in banners if b.get("id") == "missing-api-key")
    assert row.get("missing_providers") == fake


def test_model_config_frontend_tracks_inline_api_key_edits():
    store_path = PROJECT_ROOT / "plugins" / "_model_config" / "webui" / "model-config-store.js"
    api_keys_mixin_path = PROJECT_ROOT / "plugins" / "_model_config" / "webui" / "api-keys-mixin.js"
    composer_store_path = PROJECT_ROOT / "webui" / "components" / "chat" / "input" / "composer-banner-store.js"
    config_path = PROJECT_ROOT / "plugins" / "_model_config" / "webui" / "config.html"
    model_field_path = PROJECT_ROOT / "plugins" / "_model_config" / "webui" / "model-field.html"
    modal_path = PROJECT_ROOT / "plugins" / "_model_config" / "webui" / "api-keys.html"

    store_content = (
        store_path.read_text(encoding="utf-8")
        + "\n"
        + api_keys_mixin_path.read_text(encoding="utf-8")
    )
    composer_store_content = composer_store_path.read_text(encoding="utf-8")
    config_content = (
        config_path.read_text(encoding="utf-8")
        + "\n"
        + model_field_path.read_text(encoding="utf-8")
    )
    modal_content = modal_path.read_text(encoding="utf-8")

    assert "apiKeyDirty" in store_content
    assert "resetApiKeyDrafts()" in store_content
    assert "!provider || seen.has(provider) || !this.apiKeyDirty[provider]" in store_content
    assert "normalized[provider] = value.trim() ? value : '';" in store_content
    assert '"missing-api-key"' in composer_store_content
    assert 'callJsonApi("/banners"' in composer_store_content
    assert "/plugins/_model_config/missing_api_key_status" not in composer_store_content
    assert "$store.modelConfig.resetApiKeyDrafts();" in config_content
    assert '@input="$store.modelConfig.setApiKeyValue(_prov, $el.value)"' in config_content
    assert "persistAllDirtyApiKeys()" in modal_content
    assert "$store.modelConfig.resetApiKeyDrafts();" in modal_content


def test_ollama_cloud_provider_config_requires_key_and_base_url():
    import yaml

    provider_path = PROJECT_ROOT / "conf/model_providers.yaml"
    provider_config = yaml.safe_load(provider_path.read_text(encoding="utf-8"))
    ollama_cloud = provider_config["chat"]["ollama_cloud"]

    assert ollama_cloud["name"] == "Ollama Cloud"
    assert ollama_cloud["kwargs"]["api_base"] == "https://ollama.com/v1"
    assert ollama_cloud["models_list"]["endpoint_url"] == "/models"
    assert "api_key_mode" not in ollama_cloud


def test_missing_api_key_banner_includes_auto_modal_metadata(monkeypatch):
    from plugins._model_config.helpers import model_config

    fake = [{"model_type": "Chat Model", "provider": "openai"}]
    monkeypatch.setattr(model_config, "get_missing_api_key_providers", lambda: fake)

    async def run():
        banners = []
        await missing_key_banner.MissingApiKeyCheck(agent=None).execute(
            banners=banners, frontend_context={}
        )
        return next(b for b in banners if b.get("id") == "missing-api-key")

    import asyncio
    row = asyncio.run(run())

    assert row["auto_modal_path"] == "/plugins/_onboarding/webui/onboarding.html"
    assert row["auto_modal_reason"] == "missing-api-key"
    assert row["auto_modal_priority"] == 100
    assert row["type"] == "warning"
    assert row["dismissible"] is False
    assert row["missing_providers"] == fake


def test_provider_key_modes_for_local_and_ollama_cloud():
    from plugins._model_config.helpers import model_config

    assert model_config.provider_requires_api_key("ollama") is False
    assert model_config.provider_requires_api_key("lm_studio") is False
    assert model_config.provider_requires_api_key("other") is False
    assert model_config.provider_requires_api_key("ollama_cloud") is True