Agent Zero
1426 lines
54 KiB
Python
1426 lines
54 KiB
Python
"""Phase F — Agent system tests (F-TEST).
|
|
|
|
Comprehensive tests for the Phase F agent system. All tests use mocks
|
|
(AsyncMock / MagicMock) — no real DB, no real LLM, no Redis required.
|
|
|
|
Covers:
|
|
1. ReAct loop (multi-step, max_steps, timeout, error recovery, cost, dry-run, audit)
|
|
2. Agent permissions (intersection, system admin, visibility, execute, optimistic lock)
|
|
3. Approval requests (create, approve/reject/expire, list filters)
|
|
4. Skill registry (registration, get_by_names, no permission grant)
|
|
5. Context builder (system prompt, ReAct format, sensitive fields, tool descriptions)
|
|
6. Data policy (sensitive field removal, provider compliance)
|
|
7. Transparency (AI-generated marking, AI participant detection)
|
|
8. Workstream integration (message, step, result)
|
|
9. Budget limits (run stops on budget exceeded, cost accumulation)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from app.ai.agent_loop import ReActResult, ReActStep, run_react_loop
|
|
from app.ai.agent_permissions import (
|
|
AgentPermissionContext,
|
|
check_agent_execute_permission,
|
|
check_entity_lock,
|
|
filter_visible_agents,
|
|
resolve_agent_permissions,
|
|
)
|
|
from app.ai.context_builder import ReActSystemPromptBuilder, build_agent_context
|
|
from app.ai.data_policy import enforce_data_policy
|
|
from app.ai.skill_registry import SkillDefinition, SkillRegistry, get_skill_registry
|
|
from app.ai.transparency import is_ai_participant, mark_as_ai_generated
|
|
from app.core.approval import (
|
|
ApprovalRequest,
|
|
create_approval_request,
|
|
expire_approval_request,
|
|
resolve_approval_request,
|
|
)
|
|
from app.core.error_codes import ApiError
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
# No-op overrides of conftest DB fixtures — these tests use mocks only
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="session")
|
|
def db_setup():
|
|
"""No-op override of conftest db_setup (no real DB needed)."""
|
|
yield
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clean_tables(db_setup):
|
|
"""No-op override of conftest clean_tables."""
|
|
yield
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
# Shared helpers
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@dataclass
|
|
class MockAgentDefinition:
|
|
"""Minimal stand-in for AgentDefinition used across tests."""
|
|
|
|
id: uuid.UUID = field(default_factory=uuid.uuid4)
|
|
name: str = "Test Agent"
|
|
description: str = "A test agent"
|
|
llm_model: str = "gpt-4o"
|
|
system_prompt: str = "You are a helpful assistant."
|
|
api_key: str | None = None
|
|
api_base: str | None = None
|
|
provider: str | None = None
|
|
max_tokens: int = 1000
|
|
max_steps: int = 20
|
|
tool_ids: list[str] = field(default_factory=list)
|
|
skill_ids: list[str] = field(default_factory=list)
|
|
capabilities: list[str] = field(default_factory=list)
|
|
budget_limit_usd: float = 1.0
|
|
trace_mode: str = "standard"
|
|
ai_use_case_metadata: dict[str, Any] | None = None
|
|
|
|
|
|
@dataclass
|
|
class MockTool:
|
|
"""Minimal stand-in for AITool."""
|
|
|
|
name: str
|
|
description: str = ""
|
|
required_permission: str | None = None
|
|
|
|
def to_openai_schema(self) -> dict[str, Any]:
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"parameters": {},
|
|
},
|
|
}
|
|
|
|
|
|
class MockToolRegistry:
|
|
"""In-memory tool registry with get_by_names / get / get_all / list_tools."""
|
|
|
|
def __init__(self, tools: list[MockTool] | None = None) -> None:
|
|
self._tools = {t.name: t for t in (tools or [])}
|
|
|
|
def get_by_names(self, names: list[str]) -> list[MockTool]:
|
|
return [self._tools[n] for n in names if n in self._tools]
|
|
|
|
def get(self, name: str) -> MockTool | None:
|
|
return self._tools.get(name)
|
|
|
|
def get_all(self) -> list[MockTool]:
|
|
return list(self._tools.values())
|
|
|
|
def list_tools(self) -> list[dict[str, Any]]:
|
|
return [
|
|
{"id": t.name, "name": t.name, "description": t.description, "plugin": "test"}
|
|
for t in self._tools.values()
|
|
]
|
|
|
|
|
|
def _make_tool_call(
|
|
call_id: str = "call_1",
|
|
name: str = "search",
|
|
arguments: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Build a mock LiteLLM tool-call object."""
|
|
args_str = json.dumps(arguments or {})
|
|
function_mock = MagicMock()
|
|
function_mock.name = name
|
|
function_mock.arguments = args_str
|
|
tc = MagicMock()
|
|
tc.id = call_id
|
|
tc.function = function_mock
|
|
return tc
|
|
|
|
|
|
def _make_llm_response(
|
|
content: str = "",
|
|
tool_calls: list[Any] | None = None,
|
|
cost_usd: float = 0.001,
|
|
) -> dict[str, Any]:
|
|
"""Build a mock llm_complete() return value."""
|
|
message = MagicMock()
|
|
message.content = content
|
|
message.tool_calls = tool_calls or None
|
|
|
|
choice = MagicMock()
|
|
choice.message = message
|
|
|
|
raw_response = MagicMock()
|
|
raw_response.choices = [choice]
|
|
|
|
return {
|
|
"content": content,
|
|
"usage": {"total_tokens": 100},
|
|
"cost_usd": cost_usd,
|
|
"model": "gpt-4o",
|
|
"raw_response": raw_response,
|
|
}
|
|
|
|
|
|
def _make_tool_registry(tools: dict[str, AsyncMock] | None = None) -> MagicMock:
|
|
"""Build a mock ToolRegistry for the ReAct loop."""
|
|
registry = MagicMock()
|
|
_tools = tools or {}
|
|
|
|
def _get(name: str) -> Any:
|
|
if name in _tools:
|
|
tool = MagicMock()
|
|
tool.handler = _tools[name]
|
|
return tool
|
|
return None
|
|
|
|
registry.get = _get
|
|
return registry
|
|
|
|
|
|
def _permissions(
|
|
permissions: list[str] | None = None,
|
|
denied: list[str] | None = None,
|
|
is_system_admin: bool = False,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"permissions": set(permissions or []),
|
|
"denied": set(denied or []),
|
|
"field_permissions": {},
|
|
"is_system_admin": is_system_admin,
|
|
"version": 0,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def tenant_id() -> uuid.UUID:
|
|
return uuid.uuid4()
|
|
|
|
|
|
@pytest.fixture
|
|
def user_id() -> uuid.UUID:
|
|
return uuid.uuid4()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_db() -> AsyncMock:
|
|
return AsyncMock()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _mock_hooks():
|
|
"""Patch do_action so the loop doesn't try to fire real hooks."""
|
|
with patch("app.core.hooks.do_action", new_callable=AsyncMock):
|
|
yield
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clean_skill_registry():
|
|
"""Reset the singleton skill registry between tests."""
|
|
registry = get_skill_registry()
|
|
for skill in list(registry.list_all()):
|
|
registry.unregister(skill.name)
|
|
yield
|
|
for skill in list(registry.list_all()):
|
|
registry.unregister(skill.name)
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 1. ReAct Loop tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestReActLoop:
|
|
@pytest.mark.asyncio
|
|
async def test_multi_step_loop_with_3_tool_calls(self, tenant_id, user_id, mock_db):
|
|
"""Multi-step loop with 3+ tool calls then a final answer."""
|
|
agent = MockAgentDefinition()
|
|
handler = AsyncMock(return_value="result")
|
|
registry = _make_tool_registry({"search": handler})
|
|
|
|
responses = [
|
|
_make_llm_response(
|
|
content="Step 1",
|
|
tool_calls=[_make_tool_call(name="search", arguments={"q": "1"})],
|
|
cost_usd=0.001,
|
|
),
|
|
_make_llm_response(
|
|
content="Step 2",
|
|
tool_calls=[_make_tool_call(name="search", arguments={"q": "2"})],
|
|
cost_usd=0.001,
|
|
),
|
|
_make_llm_response(
|
|
content="Step 3",
|
|
tool_calls=[_make_tool_call(name="search", arguments={"q": "3"})],
|
|
cost_usd=0.001,
|
|
),
|
|
_make_llm_response(content="Final answer", cost_usd=0.001),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Do it"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.steps_taken == 4
|
|
assert len(result.steps) == 4
|
|
assert [s.action for s in result.steps] == ["search", "search", "search", None]
|
|
assert result.final_content == "Final answer"
|
|
assert handler.await_count == 3
|
|
assert mock_llm.await_count == 4
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graceful_stop_on_max_steps(self, tenant_id, user_id, mock_db):
|
|
"""Loop stops gracefully when max_steps is reached."""
|
|
agent = MockAgentDefinition()
|
|
handler = AsyncMock(return_value="result")
|
|
registry = _make_tool_registry({"search": handler})
|
|
response = _make_llm_response(
|
|
content="Searching",
|
|
tool_calls=[_make_tool_call(name="search")],
|
|
cost_usd=0.001,
|
|
)
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.return_value = response
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Keep going"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
max_steps=3,
|
|
)
|
|
|
|
assert result.status == "stopped_max_steps"
|
|
assert result.steps_taken == 3
|
|
assert "max_steps" in (result.error or "")
|
|
assert mock_llm.await_count == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graceful_stop_on_timeout(self, tenant_id, user_id, mock_db):
|
|
"""Loop stops gracefully when timeout is exceeded."""
|
|
agent = MockAgentDefinition()
|
|
registry = _make_tool_registry()
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.return_value = _make_llm_response(content="thinking", cost_usd=0.001)
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
timeout_seconds=0,
|
|
)
|
|
|
|
assert result.status == "stopped_timeout"
|
|
assert "timeout" in (result.error or "").lower()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transient_error_retry_then_success(self, tenant_id, user_id, mock_db):
|
|
"""TRANSIENT errors are retried, then the loop succeeds."""
|
|
agent = MockAgentDefinition()
|
|
registry = _make_tool_registry()
|
|
|
|
responses: list[Any] = [
|
|
Exception("rate limit exceeded"),
|
|
_make_llm_response(content="Success after retry", cost_usd=0.002),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
with patch("app.ai.agent_loop.asyncio.sleep", new_callable=AsyncMock):
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.final_content == "Success after retry"
|
|
assert mock_llm.await_count == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_permanent_error_stops_immediately(self, tenant_id, user_id, mock_db):
|
|
"""PERMANENT errors stop the loop immediately (no retry)."""
|
|
agent = MockAgentDefinition()
|
|
registry = _make_tool_registry()
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = Exception("authentication error: invalid api key")
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "stopped_error"
|
|
assert "Permanent error" in (result.error or "")
|
|
assert mock_llm.await_count == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_accumulation_across_steps(self, tenant_id, user_id, mock_db):
|
|
"""Cost accumulates across multiple LLM calls."""
|
|
agent = MockAgentDefinition()
|
|
handler = AsyncMock(return_value="result")
|
|
registry = _make_tool_registry({"search": handler})
|
|
|
|
responses = [
|
|
_make_llm_response(
|
|
content="S1", tool_calls=[_make_tool_call(name="search")], cost_usd=0.01
|
|
),
|
|
_make_llm_response(
|
|
content="S2", tool_calls=[_make_tool_call(name="search")], cost_usd=0.02
|
|
),
|
|
_make_llm_response(content="Final", cost_usd=0.03),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Search twice"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.total_cost_usd == pytest.approx(0.06)
|
|
assert result.steps_taken == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dry_run_does_not_execute_tools(self, tenant_id, user_id, mock_db):
|
|
"""Dry-run mode returns mock results and never calls tool handlers."""
|
|
agent = MockAgentDefinition()
|
|
handler = AsyncMock(return_value="REAL RESULT")
|
|
registry = _make_tool_registry({"search": handler})
|
|
|
|
responses = [
|
|
_make_llm_response(
|
|
content="Calling search",
|
|
tool_calls=[_make_tool_call(name="search", arguments={"q": "x"})],
|
|
cost_usd=0.001,
|
|
),
|
|
_make_llm_response(content="Done", cost_usd=0.001),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Search"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
dry_run=True,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
handler.assert_not_awaited()
|
|
obs = result.steps[0].observation
|
|
assert obs is not None
|
|
assert "dry_run" in obs
|
|
assert "REAL RESULT" not in obs
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_audit_log_created_for_each_tool_call(self, tenant_id, user_id, mock_db):
|
|
"""An audit log entry is created for every tool call."""
|
|
agent = MockAgentDefinition()
|
|
handler = AsyncMock(return_value="result")
|
|
registry = _make_tool_registry({"search": handler})
|
|
|
|
responses = [
|
|
_make_llm_response(
|
|
content="Calling",
|
|
tool_calls=[_make_tool_call(name="search", arguments={"q": "1"})],
|
|
cost_usd=0.001,
|
|
),
|
|
_make_llm_response(content="Done", cost_usd=0.001),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
with patch("app.core.audit.log_audit", new_callable=AsyncMock) as mock_audit:
|
|
await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Search"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
agent_run_id=uuid.uuid4(),
|
|
)
|
|
|
|
mock_audit.assert_awaited_once()
|
|
call_kwargs = mock_audit.await_args.kwargs
|
|
assert call_kwargs["action"] == "agent.tool_call"
|
|
assert call_kwargs["entity_type"] == "agent_run"
|
|
assert call_kwargs["details"]["tool_name"] == "search"
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 2. Agent Permission tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestAgentPermissions:
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_agent_permissions_intersection(self, tenant_id, user_id, mock_db):
|
|
"""resolve_agent_permissions returns the correct intersection."""
|
|
agent = MockAgentDefinition(
|
|
tool_ids=["mail_read", "contact_list"],
|
|
skill_ids=["mail_skill"],
|
|
)
|
|
skill_registry = get_skill_registry()
|
|
skill_registry.register(
|
|
SkillDefinition(
|
|
name="mail_skill",
|
|
description="mail",
|
|
instructions="use mail tools",
|
|
allowed_tool_ids=["mail_read"],
|
|
)
|
|
)
|
|
tool_registry = MockToolRegistry([
|
|
MockTool("mail_read", required_permission="mail:read"),
|
|
MockTool("contact_list"),
|
|
])
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(permissions=["mail:read"])
|
|
with patch(
|
|
"app.ai.tool_registry.get_tool_registry",
|
|
return_value=tool_registry,
|
|
):
|
|
ctx = await resolve_agent_permissions(
|
|
mock_db, tenant_id, user_id, agent
|
|
)
|
|
|
|
assert isinstance(ctx, AgentPermissionContext)
|
|
assert ctx.is_system_admin is False
|
|
assert set(ctx.effective_tool_ids) == {"mail_read", "contact_list"}
|
|
assert ctx.can_use_tool("mail_read") is True
|
|
assert ctx.can_use_tool("contact_list") is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_system_admin_gets_all_agent_tools(self, tenant_id, user_id, mock_db):
|
|
"""System admins get all tools configured on the agent."""
|
|
agent = MockAgentDefinition(tool_ids=["mail_read", "secret_tool"])
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(is_system_admin=True)
|
|
ctx = await resolve_agent_permissions(
|
|
mock_db, tenant_id, user_id, agent
|
|
)
|
|
|
|
assert ctx.is_system_admin is True
|
|
assert set(ctx.effective_tool_ids) == {"mail_read", "secret_tool"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_admin_only_gets_matching_permissions(self, tenant_id, user_id, mock_db):
|
|
"""Non-admin only gets tools whose required permission they hold."""
|
|
agent = MockAgentDefinition(tool_ids=["mail_read", "secret_tool"])
|
|
tool_registry = MockToolRegistry([
|
|
MockTool("mail_read", required_permission="mail:read"),
|
|
MockTool("secret_tool", required_permission="admin:all"),
|
|
])
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(permissions=["mail:read"])
|
|
with patch(
|
|
"app.ai.tool_registry.get_tool_registry",
|
|
return_value=tool_registry,
|
|
):
|
|
ctx = await resolve_agent_permissions(
|
|
mock_db, tenant_id, user_id, agent
|
|
)
|
|
|
|
assert set(ctx.effective_tool_ids) == {"mail_read"}
|
|
assert ctx.can_use_tool("secret_tool") is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filter_visible_agents_respects_agents_read(self, tenant_id, user_id, mock_db):
|
|
"""filter_visible_agents returns [] without agents:read."""
|
|
agents = [MockAgentDefinition(), MockAgentDefinition()]
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(permissions=[])
|
|
visible = await filter_visible_agents(mock_db, tenant_id, user_id, agents)
|
|
|
|
assert visible == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filter_visible_agents_system_admin_sees_all(self, tenant_id, user_id, mock_db):
|
|
"""System admins see all agents."""
|
|
agents = [MockAgentDefinition(), MockAgentDefinition()]
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(is_system_admin=True)
|
|
visible = await filter_visible_agents(mock_db, tenant_id, user_id, agents)
|
|
|
|
assert len(visible) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filter_visible_agents_filters_by_visible_ids(self, tenant_id, user_id, mock_db):
|
|
"""filter_visible_agents filters by get_visible_ids when agents:read is held."""
|
|
a1 = MockAgentDefinition()
|
|
a2 = MockAgentDefinition()
|
|
agents = [a1, a2]
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(permissions=["agents:read"])
|
|
with patch(
|
|
"app.services.permission_resolver.get_visible_ids",
|
|
new_callable=AsyncMock,
|
|
) as mock_visible:
|
|
mock_visible.return_value = ({a1.id}, {a1.id: "read"})
|
|
visible = await filter_visible_agents(mock_db, tenant_id, user_id, agents)
|
|
|
|
assert visible == [a1]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_agent_execute_permission_respects_agents_execute(self, tenant_id, user_id, mock_db):
|
|
"""check_agent_execute_permission returns False without agents:execute."""
|
|
agent_id = uuid.uuid4()
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(permissions=[])
|
|
allowed = await check_agent_execute_permission(
|
|
mock_db, tenant_id, user_id, agent_id
|
|
)
|
|
|
|
assert allowed is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_agent_execute_permission_system_admin(self, tenant_id, user_id, mock_db):
|
|
"""System admins always pass agents:execute."""
|
|
agent_id = uuid.uuid4()
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(is_system_admin=True)
|
|
allowed = await check_agent_execute_permission(
|
|
mock_db, tenant_id, user_id, agent_id
|
|
)
|
|
|
|
assert allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_agent_execute_permission_entity_access(self, tenant_id, user_id, mock_db):
|
|
"""With agents:execute, entity access decides."""
|
|
agent_id = uuid.uuid4()
|
|
|
|
with patch(
|
|
"app.ai.agent_permissions.resolve_permissions",
|
|
new_callable=AsyncMock,
|
|
) as mock_resolve:
|
|
mock_resolve.return_value = _permissions(permissions=["agents:execute"])
|
|
with patch(
|
|
"app.services.permission_resolver.check_entity_access",
|
|
new_callable=AsyncMock,
|
|
) as mock_access:
|
|
mock_access.return_value = True
|
|
allowed = await check_agent_execute_permission(
|
|
mock_db, tenant_id, user_id, agent_id
|
|
)
|
|
assert allowed is True
|
|
mock_access.return_value = False
|
|
denied = await check_agent_execute_permission(
|
|
mock_db, tenant_id, user_id, agent_id
|
|
)
|
|
assert denied is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_optimistic_lock_detects_version_mismatch(self, tenant_id, mock_db):
|
|
"""check_entity_lock raises ApiError conflict on version mismatch."""
|
|
entity_id = uuid.uuid4()
|
|
|
|
from sqlalchemy import column
|
|
|
|
class MockModel:
|
|
version = column("version")
|
|
id = column("id")
|
|
|
|
with patch(
|
|
"app.services.entity_permission_service.ENTITY_MODELS",
|
|
{"agent_definition": MockModel},
|
|
):
|
|
# Version match → passes
|
|
result = MagicMock()
|
|
result.scalar_one_or_none.return_value = 3
|
|
mock_db.execute.return_value = result
|
|
ok = await check_entity_lock(mock_db, "agent_definition", entity_id, 3)
|
|
assert ok is True
|
|
|
|
# Version mismatch → conflict
|
|
result.scalar_one_or_none.return_value = 5
|
|
with pytest.raises(ApiError) as exc_info:
|
|
await check_entity_lock(mock_db, "agent_definition", entity_id, 3)
|
|
assert exc_info.value.code == "conflict"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_optimistic_lock_no_version_column_is_noop(self, tenant_id, mock_db):
|
|
"""Models without a version column are treated as unlocked."""
|
|
entity_id = uuid.uuid4()
|
|
|
|
class NoVersionModel:
|
|
id = "id"
|
|
|
|
with patch(
|
|
"app.services.entity_permission_service.ENTITY_MODELS",
|
|
{"agent_definition": NoVersionModel},
|
|
):
|
|
ok = await check_entity_lock(mock_db, "agent_definition", entity_id, 1)
|
|
assert ok is True
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 3. Approval tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestApproval:
|
|
@pytest.mark.asyncio
|
|
async def test_create_approval_request(self, tenant_id, mock_db):
|
|
"""create_approval_request creates a pending request."""
|
|
entity_id = uuid.uuid4()
|
|
requested_by = uuid.uuid4()
|
|
|
|
req = await create_approval_request(
|
|
mock_db,
|
|
tenant_id,
|
|
entity_type="agent_definition",
|
|
entity_id=entity_id,
|
|
action="execute",
|
|
requested_by=requested_by,
|
|
requested_by_type="agent",
|
|
metadata={"agent_run_id": "run-1"},
|
|
)
|
|
|
|
assert isinstance(req, ApprovalRequest)
|
|
assert req.status == "pending"
|
|
assert req.entity_type == "agent_definition"
|
|
assert req.entity_id == entity_id
|
|
assert req.action == "execute"
|
|
assert req.requested_by == requested_by
|
|
assert req.request_metadata == {"agent_run_id": "run-1"}
|
|
mock_db.add.assert_called_once()
|
|
mock_db.flush.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_approve_approval_request(self, tenant_id, mock_db):
|
|
"""resolve_approval_request approves a pending request."""
|
|
request_id = uuid.uuid4()
|
|
approver_id = uuid.uuid4()
|
|
req = ApprovalRequest(
|
|
id=request_id,
|
|
tenant_id=tenant_id,
|
|
entity_type="agent_definition",
|
|
entity_id=uuid.uuid4(),
|
|
action="execute",
|
|
requested_by=uuid.uuid4(),
|
|
status="pending",
|
|
)
|
|
|
|
result = MagicMock()
|
|
result.scalar_one_or_none.return_value = req
|
|
mock_db.execute.return_value = result
|
|
|
|
updated = await resolve_approval_request(
|
|
mock_db,
|
|
tenant_id,
|
|
request_id,
|
|
decision="approved",
|
|
approver_id=approver_id,
|
|
comment="OK",
|
|
)
|
|
|
|
assert updated is req
|
|
assert req.status == "approved"
|
|
assert req.approver_id == approver_id
|
|
assert req.comment == "OK"
|
|
assert req.resolved_at is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reject_approval_request(self, tenant_id, mock_db):
|
|
"""resolve_approval_request rejects a pending request."""
|
|
request_id = uuid.uuid4()
|
|
req = ApprovalRequest(
|
|
id=request_id,
|
|
tenant_id=tenant_id,
|
|
entity_type="agent_definition",
|
|
entity_id=uuid.uuid4(),
|
|
action="execute",
|
|
requested_by=uuid.uuid4(),
|
|
status="pending",
|
|
)
|
|
result = MagicMock()
|
|
result.scalar_one_or_none.return_value = req
|
|
mock_db.execute.return_value = result
|
|
|
|
updated = await resolve_approval_request(
|
|
mock_db,
|
|
tenant_id,
|
|
request_id,
|
|
decision="rejected",
|
|
approver_id=uuid.uuid4(),
|
|
comment="No",
|
|
)
|
|
|
|
assert updated is req
|
|
assert req.status == "rejected"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expire_approval_request(self, tenant_id, mock_db):
|
|
"""expire_approval_request marks a pending request as expired."""
|
|
request_id = uuid.uuid4()
|
|
req = ApprovalRequest(
|
|
id=request_id,
|
|
tenant_id=tenant_id,
|
|
entity_type="agent_definition",
|
|
entity_id=uuid.uuid4(),
|
|
action="execute",
|
|
requested_by=uuid.uuid4(),
|
|
status="pending",
|
|
)
|
|
result = MagicMock()
|
|
result.scalar_one_or_none.return_value = req
|
|
mock_db.execute.return_value = result
|
|
|
|
updated = await expire_approval_request(mock_db, tenant_id, request_id)
|
|
|
|
assert updated is req
|
|
assert req.status == "expired"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_non_pending_returns_none(self, tenant_id, mock_db):
|
|
"""Resolving a non-pending request returns None."""
|
|
request_id = uuid.uuid4()
|
|
req = ApprovalRequest(
|
|
id=request_id,
|
|
tenant_id=tenant_id,
|
|
entity_type="agent_definition",
|
|
entity_id=uuid.uuid4(),
|
|
action="execute",
|
|
requested_by=uuid.uuid4(),
|
|
status="approved",
|
|
)
|
|
result = MagicMock()
|
|
result.scalar_one_or_none.return_value = req
|
|
mock_db.execute.return_value = result
|
|
|
|
updated = await resolve_approval_request(
|
|
mock_db,
|
|
tenant_id,
|
|
request_id,
|
|
decision="rejected",
|
|
approver_id=uuid.uuid4(),
|
|
)
|
|
|
|
assert updated is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_approvals_with_filters(self, tenant_id, mock_db):
|
|
"""List approvals applies status/entity filters via SQLAlchemy query."""
|
|
from sqlalchemy import select
|
|
|
|
# Build the same query the route uses and verify filters are applied.
|
|
query = select(ApprovalRequest).where(ApprovalRequest.tenant_id == tenant_id)
|
|
query = query.where(ApprovalRequest.status == "pending")
|
|
query = query.where(ApprovalRequest.entity_type == "agent_definition")
|
|
|
|
compiled = str(query)
|
|
assert "approval_requests" in compiled
|
|
assert "status" in compiled
|
|
assert "entity_type" in compiled
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 4. Skill tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestSkills:
|
|
def test_skill_definition_registration(self):
|
|
"""SkillDefinition registers and is retrievable."""
|
|
registry = get_skill_registry()
|
|
skill = SkillDefinition(
|
|
name="mail_reader",
|
|
description="Read emails",
|
|
instructions="Use mail tools",
|
|
allowed_tool_ids=["mail_read"],
|
|
category="mail",
|
|
)
|
|
registry.register(skill)
|
|
assert registry.get("mail_reader") is skill
|
|
assert skill.to_dict()["category"] == "mail"
|
|
|
|
def test_get_by_names_skips_unknown(self):
|
|
"""get_by_names skips unknown skill names."""
|
|
registry = get_skill_registry()
|
|
registry.register(SkillDefinition(name="a", description="d", instructions="i"))
|
|
registry.register(SkillDefinition(name="b", description="d", instructions="i"))
|
|
resolved = registry.get_by_names(["a", "b", "missing"])
|
|
assert {s.name for s in resolved} == {"a", "b"}
|
|
|
|
def test_skills_never_grant_permissions_not_in_user_role(self):
|
|
"""A skill cannot grant a tool the user lacks permission for."""
|
|
from app.ai.agent_tools import get_agent_tools
|
|
|
|
registry = MockToolRegistry([
|
|
MockTool("mail_read", required_permission="mail:read"),
|
|
])
|
|
skill_registry = get_skill_registry()
|
|
skill_registry.register(
|
|
SkillDefinition(
|
|
name="mail_skill",
|
|
description="mail",
|
|
instructions="use mail",
|
|
allowed_tool_ids=["mail_read"],
|
|
)
|
|
)
|
|
agent = MockAgentDefinition(tool_ids=["mail_read"], skill_ids=["mail_skill"])
|
|
|
|
# User does NOT have mail:read
|
|
schemas, skills = get_agent_tools(
|
|
agent, registry, skill_registry, _permissions(permissions=[])
|
|
)
|
|
|
|
assert schemas == []
|
|
assert len(skills) == 1 # skill resolved but grants no tools
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 5. Context builder tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestContextBuilder:
|
|
@pytest.mark.asyncio
|
|
async def test_build_agent_context_includes_system_prompt(self, tenant_id, user_id):
|
|
"""build_agent_context includes the agent's system prompt."""
|
|
agent = MockAgentDefinition(name="SalesBot", system_prompt="You are SalesBot.")
|
|
messages = await build_agent_context(
|
|
agent_definition=agent,
|
|
user_message="Hello",
|
|
db=None,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
system_content = messages[0]["content"]
|
|
assert "You are SalesBot." in system_content
|
|
assert messages[-1] == {"role": "user", "content": "Hello"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_react_format_instructions_present(self, tenant_id, user_id):
|
|
"""The system prompt contains ReAct format instructions."""
|
|
agent = MockAgentDefinition()
|
|
messages = await build_agent_context(
|
|
agent_definition=agent,
|
|
user_message="Hi",
|
|
db=None,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
system_content = messages[0]["content"]
|
|
assert "Thought:" in system_content
|
|
assert "Action:" in system_content
|
|
assert "Final Answer:" in system_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sensitive_fields_excluded(self, tenant_id, user_id):
|
|
"""Sensitive fields (api_key) are not included in the built context."""
|
|
agent = MockAgentDefinition(
|
|
name="Agent",
|
|
system_prompt="You are an agent.",
|
|
api_key="sk-super-secret-key-12345",
|
|
)
|
|
messages = await build_agent_context(
|
|
agent_definition=agent,
|
|
user_message="Hi",
|
|
db=None,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
all_content = "\n".join(m["content"] for m in messages)
|
|
assert "sk-super-secret-key-12345" not in all_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_descriptions_included(self, tenant_id, user_id):
|
|
"""Tool descriptions are included in the system prompt."""
|
|
agent = MockAgentDefinition(tool_ids=["mail_read"])
|
|
tool_registry = MockToolRegistry([
|
|
MockTool("mail_read", description="Read emails from the mailbox"),
|
|
])
|
|
|
|
with patch(
|
|
"app.ai.tool_registry.get_tool_registry",
|
|
return_value=tool_registry,
|
|
):
|
|
messages = await build_agent_context(
|
|
agent_definition=agent,
|
|
user_message="Hi",
|
|
db=None,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
system_content = messages[0]["content"]
|
|
assert "mail_read" in system_content
|
|
assert "Read emails from the mailbox" in system_content
|
|
|
|
def test_react_system_prompt_builder_constraints(self):
|
|
"""ReActSystemPromptBuilder includes max_steps and budget constraints."""
|
|
agent = MockAgentDefinition(max_steps=5, budget_limit_usd=2.5)
|
|
builder = ReActSystemPromptBuilder(agent_definition=agent)
|
|
prompt = builder.build()
|
|
assert "Maximum 5 reasoning steps per run." in prompt
|
|
assert "Budget limit: $2.50 per run." in prompt
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 6. Data policy tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestDataPolicy:
|
|
@pytest.mark.asyncio
|
|
async def test_enforce_data_policy_removes_sensitive_fields(self, tenant_id):
|
|
"""enforce_data_policy strips sensitive fields from dict content."""
|
|
agent = MockAgentDefinition()
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"email": "a@b.com",
|
|
"smtp_password": "secret",
|
|
"imap_password": "secret2",
|
|
"api_key": "sk-123",
|
|
"name": "Alice",
|
|
},
|
|
}
|
|
]
|
|
|
|
with patch(
|
|
"app.ai.llm_client.get_provider_compliance",
|
|
new_callable=AsyncMock,
|
|
) as mock_compliance:
|
|
mock_compliance.return_value = None
|
|
filtered = await enforce_data_policy(
|
|
mock_db := AsyncMock(), tenant_id, messages, agent
|
|
)
|
|
|
|
content = filtered[0]["content"]
|
|
assert "smtp_password" not in content
|
|
assert "imap_password" not in content
|
|
assert "api_key" not in content
|
|
assert content["email"] == "a@b.com"
|
|
assert content["name"] == "Alice"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enforce_data_policy_provider_compliance(self, tenant_id):
|
|
"""Provider compliance blocks fields the provider may not process."""
|
|
agent = MockAgentDefinition()
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"email": "a@b.com",
|
|
"smtp_password": "secret",
|
|
"notes": "internal note",
|
|
},
|
|
}
|
|
]
|
|
|
|
with patch(
|
|
"app.ai.llm_client.get_provider_compliance",
|
|
new_callable=AsyncMock,
|
|
) as mock_compliance:
|
|
# Provider only allowed to process 'internal' data — critical fields blocked
|
|
mock_compliance.return_value = {"allowed_data_classes": ["internal"]}
|
|
filtered = await enforce_data_policy(
|
|
AsyncMock(), tenant_id, messages, agent
|
|
)
|
|
|
|
content = filtered[0]["content"]
|
|
# smtp_password is critical → blocked
|
|
assert "smtp_password" not in content
|
|
# email is guessed as 'internal' (no sensitive marker) → kept
|
|
assert content.get("email") == "a@b.com"
|
|
# notes is internal → kept
|
|
assert content.get("notes") == "internal note"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enforce_data_policy_allowed_categories(self, tenant_id):
|
|
"""Allowed data categories filter disallowed entity payloads."""
|
|
agent = MockAgentDefinition(
|
|
ai_use_case_metadata={"data_categories": ["contact_data"]}
|
|
)
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"first_name": "Alice",
|
|
"last_name": "Smith",
|
|
"company_id": "c1",
|
|
},
|
|
}
|
|
]
|
|
|
|
with patch(
|
|
"app.ai.llm_client.get_provider_compliance",
|
|
new_callable=AsyncMock,
|
|
) as mock_compliance:
|
|
mock_compliance.return_value = None
|
|
filtered = await enforce_data_policy(
|
|
AsyncMock(), tenant_id, messages, agent
|
|
)
|
|
|
|
# contact_data is allowed → contact payload kept
|
|
content = filtered[0]["content"]
|
|
assert content.get("first_name") == "Alice"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enforce_data_policy_blocks_disallowed_category(self, tenant_id):
|
|
"""Disallowed data categories are removed."""
|
|
agent = MockAgentDefinition(
|
|
ai_use_case_metadata={"data_categories": ["contact_data"]}
|
|
)
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"email": "a@b.com",
|
|
"smtp_password": "secret",
|
|
},
|
|
}
|
|
]
|
|
|
|
with patch(
|
|
"app.ai.llm_client.get_provider_compliance",
|
|
new_callable=AsyncMock,
|
|
) as mock_compliance:
|
|
mock_compliance.return_value = None
|
|
filtered = await enforce_data_policy(
|
|
AsyncMock(), tenant_id, messages, agent
|
|
)
|
|
|
|
# mail_account payload is not in allowed categories → removed
|
|
content = filtered[0]["content"]
|
|
assert content == {}
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 7. Transparency tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestTransparency:
|
|
def test_mark_as_ai_generated_adds_metadata(self):
|
|
"""mark_as_ai_generated adds ai_generated flag and metadata."""
|
|
result = mark_as_ai_generated(
|
|
"Hello from AI",
|
|
{"model": "gpt-4o", "provider": "openai"},
|
|
)
|
|
assert result["content"] == "Hello from AI"
|
|
assert result["ai_generated"] is True
|
|
assert result["ai_metadata"]["model"] == "gpt-4o"
|
|
assert result["ai_metadata"]["provider"] == "openai"
|
|
assert "timestamp" in result["ai_metadata"]
|
|
|
|
def test_mark_as_ai_generated_defaults(self):
|
|
"""mark_as_ai_generated uses unknown defaults when metadata is empty."""
|
|
result = mark_as_ai_generated("content")
|
|
assert result["ai_metadata"]["model"] == "unknown"
|
|
assert result["ai_metadata"]["provider"] == "unknown"
|
|
|
|
def test_is_ai_participant_detects_ai(self):
|
|
"""is_ai_participant detects AI participant types."""
|
|
assert is_ai_participant("a1", "agent") is True
|
|
assert is_ai_participant("a2", "ai") is True
|
|
assert is_ai_participant("a3", "system_ai") is True
|
|
assert is_ai_participant("u1", "user") is False
|
|
assert is_ai_participant("u2", "human") is False
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 8. Workstream tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestWorkstream:
|
|
@pytest.mark.asyncio
|
|
async def test_post_agent_message_creates_message(self, tenant_id, mock_db):
|
|
"""post_agent_message calls send_message and returns the message ID."""
|
|
agent_id = uuid.uuid4()
|
|
run_id = uuid.uuid4()
|
|
conv_id = uuid.uuid4()
|
|
message_id = uuid.uuid4()
|
|
|
|
with patch(
|
|
"app.plugins.builtins.kommunikation.services.send_message",
|
|
new_callable=AsyncMock,
|
|
) as mock_send:
|
|
mock_send.return_value = message_id
|
|
with patch(
|
|
"app.ai.agent_workstream._get_or_create_agent_channel",
|
|
new_callable=AsyncMock,
|
|
) as mock_channel:
|
|
mock_channel.return_value = conv_id
|
|
from app.ai.agent_workstream import post_agent_message
|
|
|
|
result = await post_agent_message(
|
|
mock_db,
|
|
tenant_id,
|
|
agent_id,
|
|
run_id,
|
|
"Hello from agent",
|
|
)
|
|
|
|
assert result == message_id
|
|
mock_send.assert_awaited_once()
|
|
call_kwargs = mock_send.await_args.kwargs
|
|
assert call_kwargs["sender_type"] == "agent"
|
|
assert call_kwargs["sender_id"] == agent_id
|
|
assert call_kwargs["conversation_id"] == conv_id
|
|
assert call_kwargs["content"] == "Hello from agent"
|
|
# AI-generated metadata attached
|
|
assert call_kwargs["metadata"]["ai_generated"] is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_agent_step_creates_action_card(self, tenant_id, mock_db):
|
|
"""post_agent_step creates an action_card message."""
|
|
agent_id = uuid.uuid4()
|
|
run_id = uuid.uuid4()
|
|
message_id = uuid.uuid4()
|
|
|
|
with patch(
|
|
"app.plugins.builtins.kommunikation.services.send_message",
|
|
new_callable=AsyncMock,
|
|
) as mock_send:
|
|
mock_send.return_value = message_id
|
|
with patch(
|
|
"app.ai.agent_workstream._get_or_create_agent_channel",
|
|
new_callable=AsyncMock,
|
|
) as mock_channel:
|
|
mock_channel.return_value = uuid.uuid4()
|
|
from app.ai.agent_workstream import post_agent_step
|
|
|
|
result = await post_agent_step(
|
|
mock_db,
|
|
tenant_id,
|
|
agent_id,
|
|
run_id,
|
|
step_number=1,
|
|
thought="Thinking...",
|
|
action="search",
|
|
observation="Found 3",
|
|
)
|
|
|
|
assert result == message_id
|
|
call_kwargs = mock_send.await_args.kwargs
|
|
assert call_kwargs["blocks"][0]["type"] == "action_card"
|
|
assert call_kwargs["blocks"][0]["data"]["step_number"] == 1
|
|
assert call_kwargs["blocks"][0]["data"]["action"] == "search"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_agent_result_includes_cost_and_status(self, tenant_id, mock_db):
|
|
"""post_agent_result includes cost and status in block data."""
|
|
agent_id = uuid.uuid4()
|
|
run_id = uuid.uuid4()
|
|
message_id = uuid.uuid4()
|
|
|
|
with patch(
|
|
"app.plugins.builtins.kommunikation.services.send_message",
|
|
new_callable=AsyncMock,
|
|
) as mock_send:
|
|
mock_send.return_value = message_id
|
|
with patch(
|
|
"app.ai.agent_workstream._get_or_create_agent_channel",
|
|
new_callable=AsyncMock,
|
|
) as mock_channel:
|
|
mock_channel.return_value = uuid.uuid4()
|
|
from app.ai.agent_workstream import post_agent_result
|
|
|
|
result = await post_agent_result(
|
|
mock_db,
|
|
tenant_id,
|
|
agent_id,
|
|
run_id,
|
|
final_content="Done",
|
|
total_cost_usd=0.123456,
|
|
steps_taken=4,
|
|
status="completed",
|
|
)
|
|
|
|
assert result == message_id
|
|
call_kwargs = mock_send.await_args.kwargs
|
|
assert call_kwargs["content"] == "Done"
|
|
block_data = call_kwargs["blocks"][0]["data"]
|
|
assert block_data["status"] == "completed"
|
|
assert block_data["steps_taken"] == 4
|
|
assert block_data["total_cost_usd"] == pytest.approx(0.123456)
|
|
assert block_data["run_id"] == str(run_id)
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
# 9. Budget limit tests
|
|
# ══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestBudgetLimits:
|
|
@pytest.mark.asyncio
|
|
async def test_agent_run_stops_when_budget_exceeded(self):
|
|
"""run_agent returns budget_exceeded when cumulative cost >= budget."""
|
|
from app.plugins.builtins.automation.agent_runner import run_agent
|
|
|
|
agent_id = uuid.uuid4()
|
|
agent = MagicMock()
|
|
agent.id = agent_id
|
|
agent.is_active = True
|
|
agent.max_executions_per_hour = 0
|
|
agent.budget_limit_usd = 1.0
|
|
agent.tenant_id = uuid.uuid4()
|
|
agent.created_by = uuid.uuid4()
|
|
agent.mode = "reactive"
|
|
agent.tool_ids = []
|
|
agent.skill_ids = []
|
|
agent.name = "Budget Agent"
|
|
agent.max_duration_seconds = 300
|
|
|
|
# Mock session factory: first query returns agent, second returns cost sum
|
|
mock_db = AsyncMock()
|
|
|
|
def _execute_side_effect(query):
|
|
result = MagicMock()
|
|
if "agent_definitions" in str(query).lower():
|
|
result.scalar_one_or_none.return_value = agent
|
|
else:
|
|
# cost sum query
|
|
result.scalar.return_value = 1.5 # >= budget 1.0
|
|
return result
|
|
|
|
mock_db.execute.side_effect = _execute_side_effect
|
|
|
|
class MockFactory:
|
|
def __init__(self, db):
|
|
self._db = db
|
|
|
|
async def __aenter__(self):
|
|
return self._db
|
|
|
|
async def __aexit__(self, *args):
|
|
return False
|
|
|
|
factory = MagicMock()
|
|
factory.return_value = MockFactory(mock_db)
|
|
|
|
with patch(
|
|
"app.plugins.builtins.automation.agent_runner.get_session_factory",
|
|
return_value=factory,
|
|
):
|
|
result = await run_agent(
|
|
ctx={"tenant_id": str(agent.tenant_id), "user_id": str(agent.created_by)},
|
|
agent_id=str(agent_id),
|
|
)
|
|
|
|
assert result["status"] == "budget_exceeded"
|
|
assert "Budget limit exceeded" in result["error"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_tracking_accumulates_in_react_loop(self, tenant_id, user_id, mock_db):
|
|
"""ReAct loop cost tracking accumulates across steps (budget input)."""
|
|
agent = MockAgentDefinition()
|
|
handler = AsyncMock(return_value="result")
|
|
registry = _make_tool_registry({"search": handler})
|
|
|
|
responses = [
|
|
_make_llm_response(
|
|
content="S1", tool_calls=[_make_tool_call(name="search")], cost_usd=0.4
|
|
),
|
|
_make_llm_response(
|
|
content="S2", tool_calls=[_make_tool_call(name="search")], cost_usd=0.4
|
|
),
|
|
_make_llm_response(
|
|
content="S3", tool_calls=[_make_tool_call(name="search")], cost_usd=0.4
|
|
),
|
|
_make_llm_response(content="Final", cost_usd=0.1),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
result = await run_react_loop(
|
|
agent_definition=agent,
|
|
messages=[{"role": "user", "content": "Run"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
# 0.4 + 0.4 + 0.4 + 0.1 = 1.3 — exceeds a 1.0 budget
|
|
assert result.total_cost_usd == pytest.approx(1.3)
|
|
assert result.total_cost_usd > 1.0
|