feat: add per-model concurrency limits

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

@@ -0,0 +1,161 @@
+"""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 fastapi import HTTPException
+
+from app.gateway.routers import system_settings as ss_router
+from app.gateway.services import (
+    _enforce_model_concurrency_limit,
+    _resolve_effective_run_model_name,
+)
+from deerflow.config import system_settings as ss
+from deerflow.runtime.runs.manager import 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"
+
+
+@pytest.mark.asyncio
+async def test_enforce_model_concurrency_limit_rejects_when_limit_reached(
+    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},
+            )
+        ),
+    )
+    run_mgr = RunManager()
+    await run_mgr.create("t1", metadata={"model_name": "glm5"})
+    await run_mgr.create("t2", metadata={"model_name": "glm5"})
+
+    with pytest.raises(HTTPException) as exc:
+        await _enforce_model_concurrency_limit(run_mgr, "glm5")
+    assert exc.value.status_code == 429
+    assert "并发已达到上限" in str(exc.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_enforce_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},
+            )
+        ),
+    )
+    run_mgr = RunManager()
+    await run_mgr.create("t1", metadata={"model_name": "glm5"})
+    await _enforce_model_concurrency_limit(run_mgr, "glm5")