deerflow-code/offline-backend-20260512/backend/tests/test_model_concurrency.py

"""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