"""Tests for per-model concurrency limit settings and enforcement.""" from __future__ import annotations import json from pathlib import Path from types import SimpleNamespace import pytest from app.gateway.routers import system_settings as ss_router from app.gateway.services import ( _get_model_concurrency_limit, _resolve_effective_run_model_name, ) from deerflow.config import system_settings as ss from deerflow.runtime.runs.manager import ModelConcurrencyLimitError, RunManager @pytest.fixture def tmp_runtime_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path: def fake_home() -> Path: return tmp_path monkeypatch.setattr(ss, "runtime_home", fake_home) return tmp_path def _patch_admin(monkeypatch: pytest.MonkeyPatch, user: object | None = None) -> None: admin = user or SimpleNamespace( id="admin1", system_role="admin", email="admin@example.com", ) async def _optional(_request): return admin monkeypatch.setattr(ss_router, "get_optional_user_from_request", _optional) @pytest.mark.asyncio async def test_model_concurrency_settings_round_trip( tmp_runtime_home: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: _patch_admin(monkeypatch) req = SimpleNamespace() saved = await ss_router.update_model_concurrency( ss_router.ModelConcurrencyUpdateRequest( enabled=True, default_limit=8, model_limits={"glm5": 10, "deepseek": 4}, ), req, ) assert saved.enabled is True assert saved.default_limit == 8 assert saved.model_limits == {"glm5": 10, "deepseek": 4} file_path = tmp_runtime_home / "system_settings.json" payload = json.loads(file_path.read_text(encoding="utf-8")) assert payload["model_concurrency"]["enabled"] is True assert payload["model_concurrency"]["default_limit"] == 8 assert payload["model_concurrency"]["model_limits"]["glm5"] == 10 loaded = await ss_router.get_model_concurrency(req) assert loaded.enabled is True assert loaded.default_limit == 8 assert loaded.model_limits["deepseek"] == 4 def test_resolve_effective_run_model_name_prefers_requested_model( monkeypatch: pytest.MonkeyPatch, ) -> None: app_config = SimpleNamespace( models=[SimpleNamespace(name="default-model")], get_model_config=lambda name: {"name": name} if name in {"default-model", "glm5", "agent-model"} else None, ) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr( "app.gateway.services.load_agent_config", lambda _agent_id: SimpleNamespace(model="agent-model"), ) body = SimpleNamespace( assistant_id="custom-agent", context={"model_name": "glm5"}, config=None, ) assert _resolve_effective_run_model_name(body) == "glm5" def test_resolve_effective_run_model_name_uses_agent_default_when_request_missing( monkeypatch: pytest.MonkeyPatch, ) -> None: app_config = SimpleNamespace( models=[SimpleNamespace(name="default-model")], get_model_config=lambda name: {"name": name} if name in {"default-model", "agent-model"} else None, ) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr( "app.gateway.services.load_agent_config", lambda _agent_id: SimpleNamespace(model="agent-model"), ) body = SimpleNamespace( assistant_id="custom-agent", context={}, config=None, ) assert _resolve_effective_run_model_name(body) == "agent-model" def test_resolve_effective_run_model_name_uses_context_agent_id( monkeypatch: pytest.MonkeyPatch, ) -> None: app_config = SimpleNamespace( models=[SimpleNamespace(name="default-model")], get_model_config=lambda name: {"name": name} if name in {"default-model", "agent-model"} else None, ) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr( "app.gateway.services.load_agent_config", lambda agent_id: SimpleNamespace(model="agent-model") if agent_id == "custom-agent" else SimpleNamespace(model="default-model"), ) body = SimpleNamespace( assistant_id="lead_agent", context={"agent_id": "custom-agent"}, config=None, ) assert _resolve_effective_run_model_name(body) == "agent-model" def test_resolve_effective_run_model_name_context_agent_overrides_assistant_id( monkeypatch: pytest.MonkeyPatch, ) -> None: app_config = SimpleNamespace( models=[SimpleNamespace(name="default-model")], get_model_config=lambda name: {"name": name} if name in {"default-model", "context-model", "assistant-model"} else None, ) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr( "app.gateway.services.load_agent_config", lambda agent_id: SimpleNamespace(model="context-model" if agent_id == "context-agent" else "assistant-model"), ) body = SimpleNamespace( assistant_id="assistant-agent", context={"agent_id": "context-agent"}, config=None, ) assert _resolve_effective_run_model_name(body) == "context-model" @pytest.mark.asyncio async def test_run_manager_rejects_model_limit_atomically( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( "app.gateway.services.load_system_settings", lambda: SimpleNamespace( model_concurrency=SimpleNamespace( enabled=True, default_limit=5, model_limits={"glm5": 2}, ) ), ) assert _get_model_concurrency_limit("glm5") == 2 run_mgr = RunManager() await run_mgr.create_or_reject("t1", metadata={"model_name": "glm5"}) await run_mgr.create_or_reject("t2", metadata={"model_name": "glm5"}) with pytest.raises(ModelConcurrencyLimitError): await run_mgr.create_or_reject( "t3", metadata={"model_name": "glm5"}, model_concurrency_limit=("glm5", 2), ) def test_get_model_concurrency_limit_skips_when_disabled( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( "app.gateway.services.load_system_settings", lambda: SimpleNamespace( model_concurrency=SimpleNamespace( enabled=False, default_limit=1, model_limits={"glm5": 1}, ) ), ) assert _get_model_concurrency_limit("glm5") is None