chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,390 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _load_context_module():
|
||||
return importlib.import_module("yuxi.agents.context")
|
||||
|
||||
|
||||
context_module = _load_context_module()
|
||||
BaseContext = context_module.BaseContext
|
||||
filter_config_by_role = context_module.filter_config_by_role
|
||||
normalize_agent_context_config = context_module.normalize_agent_context_config
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ChatBotContext(BaseContext):
|
||||
subagents: list[str] | None = field(default=None, metadata={"kind": "subagents"})
|
||||
|
||||
|
||||
@dataclass
|
||||
class SuperAdminOnlyContext(BaseContext):
|
||||
secret_setting: str = field(default="hidden", metadata={"name": "Secret", "auth": "superadmin"})
|
||||
|
||||
|
||||
def test_get_configurable_items_filters_admin_fields_for_user():
|
||||
items = BaseContext.get_configurable_items(user_role="user")
|
||||
|
||||
assert "system_prompt" in items
|
||||
assert "summary_threshold" not in items
|
||||
assert "summary_keep_messages" not in items
|
||||
assert "summary_prompt" not in items
|
||||
assert "summary_tool_result_token_limit" not in items
|
||||
assert "max_execution_steps" not in items
|
||||
|
||||
|
||||
def test_get_configurable_items_allows_admin_and_superadmin_fields():
|
||||
admin_items = BaseContext.get_configurable_items(user_role="admin")
|
||||
superadmin_items = SuperAdminOnlyContext.get_configurable_items(user_role="superadmin")
|
||||
|
||||
assert "summary_threshold" in admin_items
|
||||
assert "summary_keep_messages" in admin_items
|
||||
assert "summary_prompt" in admin_items
|
||||
assert "summary_tool_result_token_limit" in admin_items
|
||||
assert "max_execution_steps" in admin_items
|
||||
assert "secret_setting" in superadmin_items
|
||||
|
||||
|
||||
def test_filter_config_by_role_removes_unauthorized_context_values():
|
||||
config_json = {
|
||||
"context": {
|
||||
"system_prompt": "visible",
|
||||
"summary_threshold": 10,
|
||||
"summary_keep_messages": 8,
|
||||
"summary_prompt": "custom summary",
|
||||
"summary_tool_result_token_limit": 500,
|
||||
"max_execution_steps": 50,
|
||||
"secret_setting": "nope",
|
||||
},
|
||||
"other": {"keep": True},
|
||||
}
|
||||
|
||||
filtered = filter_config_by_role(config_json, "user", context_schema=SuperAdminOnlyContext)
|
||||
|
||||
assert filtered == {"context": {"system_prompt": "visible"}, "other": {"keep": True}}
|
||||
assert config_json["context"]["summary_threshold"] == 10
|
||||
|
||||
|
||||
def test_filter_config_by_role_keeps_admin_context_values_for_admin():
|
||||
filtered = filter_config_by_role(
|
||||
{
|
||||
"context": {
|
||||
"summary_threshold": 10,
|
||||
"summary_keep_messages": 8,
|
||||
"summary_prompt": "custom summary",
|
||||
"summary_tool_result_token_limit": 500,
|
||||
"max_execution_steps": 50,
|
||||
"secret_setting": "nope",
|
||||
}
|
||||
},
|
||||
"admin",
|
||||
context_schema=SuperAdminOnlyContext,
|
||||
)
|
||||
|
||||
assert filtered == {
|
||||
"context": {
|
||||
"summary_threshold": 10,
|
||||
"summary_keep_messages": 8,
|
||||
"summary_prompt": "custom summary",
|
||||
"summary_tool_result_token_limit": 500,
|
||||
"max_execution_steps": 50,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_agent_resource_options_empty_fields_loads_nothing(monkeypatch):
|
||||
async def fail_if_loaded(*_args, **_kwargs):
|
||||
raise AssertionError("empty resource_fields should not load resources")
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.knowledge",
|
||||
types.SimpleNamespace(knowledge_base=types.SimpleNamespace(get_databases_by_user=fail_if_loaded)),
|
||||
)
|
||||
|
||||
assert await context_module.resolve_agent_resource_options(set(), db=object(), user=object()) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normalize_agent_context_config_expands_null_and_filters_explicit_lists(monkeypatch):
|
||||
async def fake_get_databases_by_user(_user):
|
||||
return {"databases": [{"kb_id": "kb-a"}, {"kb_id": "kb-b"}]}
|
||||
|
||||
async def fake_get_all_mcp_servers(_db):
|
||||
return [
|
||||
types.SimpleNamespace(slug="mcp-a", name="MCP A", description="", enabled=True),
|
||||
types.SimpleNamespace(slug="mcp-b", name="MCP B", description="", enabled=True),
|
||||
]
|
||||
|
||||
async def fake_list_skills(_db, _user):
|
||||
return [
|
||||
types.SimpleNamespace(slug="skill-a", name="Skill A", description=""),
|
||||
types.SimpleNamespace(slug="skill-b", name="Skill B", description=""),
|
||||
]
|
||||
|
||||
class FakeAgentRepository:
|
||||
def __init__(self, _db):
|
||||
pass
|
||||
|
||||
async def list_visible_subagents(self, *, user):
|
||||
assert user.uid == "u1"
|
||||
return [
|
||||
types.SimpleNamespace(slug="research-agent", name="Research", description=""),
|
||||
types.SimpleNamespace(slug="critique-agent", name="Critique", description=""),
|
||||
]
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.toolkits.service",
|
||||
types.SimpleNamespace(
|
||||
get_tool_metadata=lambda category=None: [
|
||||
{"slug": "ask_user_question", "name": "Ask User", "description": ""},
|
||||
{"slug": "tavily_search", "name": "Tavily", "description": ""},
|
||||
]
|
||||
),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.knowledge",
|
||||
types.SimpleNamespace(knowledge_base=types.SimpleNamespace(get_databases_by_user=fake_get_databases_by_user)),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.mcp.service",
|
||||
types.SimpleNamespace(get_all_mcp_servers=fake_get_all_mcp_servers),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.skills.service",
|
||||
types.SimpleNamespace(list_accessible_skills=fake_list_skills),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.repositories.agent_repository",
|
||||
types.SimpleNamespace(AgentRepository=FakeAgentRepository),
|
||||
)
|
||||
|
||||
normalized = await normalize_agent_context_config(
|
||||
{
|
||||
"tools": None,
|
||||
"knowledges": ["kb-b", "missing", "kb-b"],
|
||||
"mcps": None,
|
||||
"skills": [],
|
||||
"subagents": ["research-agent", "missing"],
|
||||
"summary_threshold": 10,
|
||||
"summary_keep_messages": 8,
|
||||
"summary_prompt": "custom summary",
|
||||
"summary_tool_result_token_limit": 500,
|
||||
"max_execution_steps": 50,
|
||||
},
|
||||
db=object(),
|
||||
user=types.SimpleNamespace(role="user", uid="u1", department_id=None),
|
||||
context_schema=ChatBotContext,
|
||||
)
|
||||
|
||||
assert normalized["tools"] == ["ask_user_question", "tavily_search"]
|
||||
assert normalized["knowledges"] == ["kb-b"]
|
||||
assert normalized["mcps"] == ["mcp-a", "mcp-b"]
|
||||
assert normalized["skills"] == []
|
||||
assert normalized["subagents"] == ["research-agent"]
|
||||
assert "summary_threshold" not in normalized
|
||||
assert "summary_keep_messages" not in normalized
|
||||
assert "summary_prompt" not in normalized
|
||||
assert "summary_tool_result_token_limit" not in normalized
|
||||
assert "max_execution_steps" not in normalized
|
||||
|
||||
empty_subagents_normalized = await normalize_agent_context_config(
|
||||
{"tools": [], "knowledges": [], "mcps": [], "skills": [], "subagents": []},
|
||||
db=object(),
|
||||
user=types.SimpleNamespace(role="user", uid="u1", department_id=None),
|
||||
context_schema=ChatBotContext,
|
||||
)
|
||||
|
||||
assert empty_subagents_normalized["subagents"] == ["research-agent", "critique-agent"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prepare_agent_runtime_context_filters_resources_and_derives_runtime_scope(monkeypatch):
|
||||
async def fake_get_databases_by_user(_user):
|
||||
return {"databases": [{"kb_id": "kb-a"}, {"kb_id": "kb-b"}]}
|
||||
|
||||
async def fake_get_all_mcp_servers(_db):
|
||||
return [types.SimpleNamespace(slug="mcp-a", name="MCP A", description="", enabled=True)]
|
||||
|
||||
async def fake_list_skills(_db, _user):
|
||||
return [
|
||||
types.SimpleNamespace(slug="skill-a", name="Skill A", description=""),
|
||||
types.SimpleNamespace(slug="skill-b", name="Skill B", description=""),
|
||||
]
|
||||
|
||||
async def fake_resolve_visible_knowledge_bases(context):
|
||||
assert context.knowledges == ["kb-a"]
|
||||
context._visible_knowledge_bases = [{"slug": "kb-a", "name": "Docs A"}]
|
||||
return context._visible_knowledge_bases
|
||||
|
||||
async def fake_resolve_runtime_skills_for_context(context, *, db=None, user=None):
|
||||
del db
|
||||
assert user.uid == "u1"
|
||||
assert context.skills == ["skill-a"]
|
||||
return {
|
||||
"context_skills": ["skill-a"],
|
||||
"prompt_skills": ["skill-a", "skill-b"],
|
||||
"readable_skills": ["skill-a", "skill-b"],
|
||||
"runtime_skill_metadata": {"skill-a": {"name": "Skill A"}},
|
||||
"runtime_skill_dependency_map": {"skill-a": {"skills": ["skill-b"]}},
|
||||
}
|
||||
|
||||
class FakeSessionContext:
|
||||
async def __aenter__(self):
|
||||
return object()
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return None
|
||||
|
||||
class FakeUserRepository:
|
||||
async def get_by_uid_with_db(self, _db, uid):
|
||||
assert uid == "u1"
|
||||
return types.SimpleNamespace(role="user", uid="u1", department_id=None)
|
||||
|
||||
class FakeAgentRepository:
|
||||
def __init__(self, _db):
|
||||
pass
|
||||
|
||||
async def list_visible_subagents(self, *, user):
|
||||
assert user.uid == "u1"
|
||||
return [types.SimpleNamespace(slug="research-agent", name="Research", description="")]
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.backends.knowledge_base_backend",
|
||||
types.SimpleNamespace(resolve_visible_knowledge_bases_for_context=fake_resolve_visible_knowledge_bases),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.middlewares.skills",
|
||||
types.SimpleNamespace(resolve_runtime_skills_for_context=fake_resolve_runtime_skills_for_context),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.repositories.user_repository",
|
||||
types.SimpleNamespace(UserRepository=FakeUserRepository),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.storage.postgres.manager",
|
||||
types.SimpleNamespace(pg_manager=types.SimpleNamespace(get_async_session_context=lambda: FakeSessionContext())),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.toolkits.service",
|
||||
types.SimpleNamespace(
|
||||
get_tool_metadata=lambda category=None: [
|
||||
{"slug": "ask_user_question", "name": "Ask User", "description": ""}
|
||||
]
|
||||
),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.knowledge",
|
||||
types.SimpleNamespace(knowledge_base=types.SimpleNamespace(get_databases_by_user=fake_get_databases_by_user)),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.mcp.service",
|
||||
types.SimpleNamespace(get_all_mcp_servers=fake_get_all_mcp_servers),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.skills.service",
|
||||
types.SimpleNamespace(list_accessible_skills=fake_list_skills),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.repositories.agent_repository",
|
||||
types.SimpleNamespace(AgentRepository=FakeAgentRepository),
|
||||
)
|
||||
context = ChatBotContext(
|
||||
uid="u1",
|
||||
tools=["ask_user_question", "missing"],
|
||||
knowledges=["kb-a", "missing"],
|
||||
mcps=None,
|
||||
skills=["skill-a", "missing"],
|
||||
subagents=[],
|
||||
)
|
||||
|
||||
prepared = await context_module.prepare_agent_runtime_context(context)
|
||||
|
||||
assert prepared.tools == ["ask_user_question"]
|
||||
assert prepared.knowledges == ["kb-a"]
|
||||
assert prepared.mcps == ["mcp-a"]
|
||||
assert prepared.skills == ["skill-a"]
|
||||
assert prepared.subagents == ["research-agent"]
|
||||
assert prepared._visible_knowledge_bases == [{"slug": "kb-a", "name": "Docs A"}]
|
||||
assert prepared._prompt_skills == ["skill-a", "skill-b"]
|
||||
assert prepared._readable_skills == ["skill-a", "skill-b"]
|
||||
assert prepared._runtime_skill_metadata == {"skill-a": {"name": "Skill A"}}
|
||||
assert prepared._runtime_skill_dependency_map == {"skill-a": {"skills": ["skill-b"]}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prepare_agent_runtime_context_clears_resources_for_missing_user(monkeypatch):
|
||||
class FakeSessionContext:
|
||||
async def __aenter__(self):
|
||||
return object()
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return None
|
||||
|
||||
class FakeUserRepository:
|
||||
async def get_by_uid_with_db(self, _db, _uid):
|
||||
return None
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.backends.knowledge_base_backend",
|
||||
types.SimpleNamespace(resolve_visible_knowledge_bases_for_context=lambda _context: None),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.agents.middlewares.skills",
|
||||
types.SimpleNamespace(resolve_runtime_skills_for_context=lambda _context, db=None, user=None: None),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.repositories.user_repository",
|
||||
types.SimpleNamespace(UserRepository=FakeUserRepository),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"yuxi.storage.postgres.manager",
|
||||
types.SimpleNamespace(pg_manager=types.SimpleNamespace(get_async_session_context=lambda: FakeSessionContext())),
|
||||
)
|
||||
|
||||
context = ChatBotContext(
|
||||
uid="missing",
|
||||
tools=["tool"],
|
||||
knowledges=["kb"],
|
||||
mcps=["mcp"],
|
||||
skills=["skill"],
|
||||
subagents=["agent"],
|
||||
)
|
||||
|
||||
prepared = await context_module.prepare_agent_runtime_context(context)
|
||||
|
||||
assert prepared.tools == []
|
||||
assert prepared.knowledges == []
|
||||
assert prepared.mcps == []
|
||||
assert prepared.skills == []
|
||||
assert prepared.subagents == []
|
||||
assert prepared._visible_knowledge_bases == []
|
||||
assert prepared._prompt_skills == []
|
||||
assert prepared._readable_skills == []
|
||||
assert prepared._runtime_skill_metadata == {}
|
||||
assert prepared._runtime_skill_dependency_map == {}
|
||||
Reference in New Issue
Block a user