test(agent-chat): audit guardrail and error handling
This commit is contained in:
@@ -336,8 +336,18 @@ async def _format_tool_output_via_llm(
|
||||
# @SIDE_EFFECT Invokes LangChain tools; modifies _pending_confirmations dict.
|
||||
# @RELATION DEPENDS_ON -> [AgentChat.LangGraph.Setup]
|
||||
# @DATA_CONTRACT Input: (conv_id, action, user_jwt, env_id) -> Output: AsyncGenerator[str]
|
||||
# @RATIONALE Fast-path resume (direct tool execution without re-entering LangGraph) chosen because the HITL checkpoint already contains all necessary context — re-running the agent would be redundant and slow.
|
||||
# @REJECTED Pure streaming without checkpoint was rejected — without a persisted checkpoint, a crash after confirmation but before tool execution would lose the operation entirely with no rollback capability.
|
||||
# @RATIONALE Fast-path resume (direct tool execution via _pending_confirmations dict)
|
||||
# chosen because the HITL confirmation payload already contains serialised tool
|
||||
# name + args — re-entering LangGraph to invoke the same tool is redundant.
|
||||
# Bypasses ~1-3s of LangGraph overhead (agent init, state reconstruction, tool
|
||||
# re-selection) per resume. Falls back to full LangGraph checkpoint resume when
|
||||
# _pending_confirmations is empty (e.g. after container restart).
|
||||
# @REJECTED ALWAYS checkpoint resume via create_agent(interrupt_before=[]) was
|
||||
# rejected — adds 1-3s latency to every resume for no reliability gain when
|
||||
# _pending_confirmations is populated. The full checkpoint path is preserved as
|
||||
# the fallback, providing defense-in-depth for container restart scenarios.
|
||||
# @REJECTED Pure streaming without checkpoint — would lose unconfirmed operations
|
||||
# on crash with no rollback capability.
|
||||
async def handle_resume( # noqa: C901
|
||||
conversation_id: str, action: str,
|
||||
user_jwt: str = "", env_id: str | None = None,
|
||||
|
||||
@@ -13,7 +13,6 @@
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger("superset_tools_app")
|
||||
|
||||
@@ -26,7 +25,7 @@ def _get_descriptions() -> tuple[list[str], list[str]]:
|
||||
"""Return (descriptions, tool_names) from get_all_tools() docstrings.
|
||||
Uses _TOOL_DESCRIPTIONS_OVERRIDES from tools.py for optional RU/EN synonyms.
|
||||
"""
|
||||
from src.agent.tools import get_all_tools, _TOOL_DESCRIPTIONS_OVERRIDES
|
||||
from src.agent.tools import _TOOL_DESCRIPTIONS_OVERRIDES, get_all_tools
|
||||
|
||||
all_tools = get_all_tools()
|
||||
names = []
|
||||
@@ -44,8 +43,8 @@ def _get_descriptions() -> tuple[list[str], list[str]]:
|
||||
# Model state — lazy-loaded on first call
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
_embedding_model: Optional[object] = None
|
||||
_tool_embeddings: Optional[object] = None # torch.Tensor or numpy array
|
||||
_embedding_model: object | None = None
|
||||
_tool_embeddings: object | None = None # torch.Tensor or numpy array
|
||||
_tool_names: list[str] = []
|
||||
|
||||
_THRESHOLD = float(os.getenv("EMBEDDING_SIMILARITY_THRESHOLD", "0.65"))
|
||||
|
||||
@@ -10,7 +10,6 @@
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
_TEMPERATURE_UNSUPPORTED_PREFIXES = (
|
||||
"codex/",
|
||||
"omni/codex/",
|
||||
|
||||
@@ -31,16 +31,16 @@ from jose import JWTError
|
||||
from langchain_core.exceptions import OutputParserException
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langchain_openai import ChatOpenAI
|
||||
from openai import APIConnectionError, APITimeoutError, AuthenticationError
|
||||
from openai import APIConnectionError, APITimeoutError, AuthenticationError, RateLimitError
|
||||
|
||||
from src.agent._config import GRADIO_SERVER_NAME, GRADIO_SERVER_PORT, STORAGE_ROOT as _STORAGE_ROOT
|
||||
from src.agent._llm_params import chat_openai_kwargs
|
||||
from src.agent._confirmation import (
|
||||
_pending_confirmations,
|
||||
confirmation_payload,
|
||||
handle_resume,
|
||||
permission_denied_payload,
|
||||
)
|
||||
from src.agent._llm_params import chat_openai_kwargs
|
||||
from src.agent._persistence import (
|
||||
extract_user_id,
|
||||
generate_llm_title,
|
||||
@@ -261,15 +261,6 @@ async def _check_llm_provider_health() -> str:
|
||||
# #endregion AgentChat.GradioApp.LlmHealthCheck
|
||||
|
||||
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Core streaming handler — runs LangGraph agent, yields ChatMessage tokens with structured metadata.
|
||||
# @PRE JWT valid, user authenticated.
|
||||
# @POST Tokens streamed via yield; HITL interrupts yield confirm_required metadata.
|
||||
# @SIDE_EFFECT Calls LLM, invokes tools, writes checkpoints.
|
||||
# @RATIONALE Async generator pattern chosen for Gradio ChatInterface compatibility — Gradio iterates
|
||||
# @REJECTED Returning a single response (non-streaming) was rejected — violates FR-003 (streaming mandate).
|
||||
|
||||
|
||||
# #region AgentChat.GradioApp.InjectUIContext [C:2] [TYPE Function] [SEMANTICS agent-chat,context,uicontext]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Add UI context block to runtime context string. Informational only, not instructions.
|
||||
@@ -307,7 +298,6 @@ def _inject_env_id_into_tools(tools: list, env_id: str | None) -> list:
|
||||
if not env_id:
|
||||
return tools
|
||||
|
||||
import types
|
||||
|
||||
for tool in tools:
|
||||
# Only wrap tools that accept 'environment_id' parameter
|
||||
@@ -388,7 +378,7 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
# ── Per-user lock ──
|
||||
user_id = user_id_str or (extract_user_id(user_jwt_str) if user_jwt_str else "admin")
|
||||
if _user_locks.get(user_id, False):
|
||||
yield json.dumps({"metadata": {"type": "error", "code": "CONCURRENT_SEND"}})
|
||||
yield json.dumps({"metadata": {"type": "error", "code": "CONCURRENT_SEND", "detail": "Другой запрос уже обрабатывается. Дождитесь завершения перед отправкой нового."}})
|
||||
return
|
||||
_user_locks[user_id] = True
|
||||
conv_id: str | None = None
|
||||
@@ -495,7 +485,7 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
except TimeoutError:
|
||||
yield json.dumps({
|
||||
"content": "❌ Ошибка: предыдущий стрим ещё не завершён",
|
||||
"metadata": {"type": "error", "code": "STREAM_CLEANUP_TIMEOUT"},
|
||||
"metadata": {"type": "error", "code": "STREAM_CLEANUP_TIMEOUT", "detail": "Предыдущий поток не завершился. Повторите запрос."},
|
||||
})
|
||||
return
|
||||
# Capture tool_name BEFORE handle_resume pops the pending confirmation
|
||||
@@ -640,6 +630,7 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
yield json.dumps({
|
||||
"content": "❌ Агент завершился без ответа.",
|
||||
"metadata": {"type": "error", "code": "EMPTY_AGENT_RESPONSE",
|
||||
"detail": "Агент завершил обработку без ответа. Попробуйте переформулировать запрос.",
|
||||
"state_next": repr(getattr(state, "next", None)),
|
||||
"state_tasks": repr(getattr(state, "tasks", None))[:500]},
|
||||
})
|
||||
@@ -699,6 +690,24 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
await save_conversation(conv_id, visible_user_text, user_id, assistant_text="")
|
||||
return
|
||||
|
||||
except RateLimitError as exc:
|
||||
_llm_status["status"] = "unavailable"
|
||||
_llm_status["last_error"] = str(exc)
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
logger.explore("LLM provider rate limited",
|
||||
error=str(exc),
|
||||
extra={"src": "AgentChat.GradioApp.Handler"})
|
||||
yield json.dumps({
|
||||
"content": "❌ Превышен лимит запросов к LLM. Попробуйте позже.",
|
||||
"metadata": {
|
||||
"type": "error", "code": "LLM_RATE_LIMITED",
|
||||
"detail": "Превышена квота LLM провайдера. Повторите запрос через несколько минут.",
|
||||
"retryable": True,
|
||||
},
|
||||
})
|
||||
await save_conversation(conv_id, visible_user_text, user_id, assistant_text="")
|
||||
return
|
||||
|
||||
except OutputParserException as e:
|
||||
if attempt < max_attempts - 1:
|
||||
text = "Respond with valid JSON only. Previous response was malformed.\n\n" + text
|
||||
@@ -721,13 +730,19 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
error=str(exc),
|
||||
extra={"src": "AgentChat.GradioApp.Handler"},
|
||||
)
|
||||
try:
|
||||
state = await agent.aget_state(config)
|
||||
if getattr(state, "next", None):
|
||||
yield confirmation_payload(conv_id, state, visible_user_text, user_role, env_id)
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
# Only attempt HITL recovery if this is NOT a known LLM/API error.
|
||||
# LLM errors (rate limits, connection failures) crash the graph
|
||||
# before tool execution — there is no HITL checkpoint to resume.
|
||||
_llm_error_patterns = ["rate limit", "quota", "429", "api key", "auth", "timeout"]
|
||||
_is_llm_error = any(p in str(exc).lower() for p in _llm_error_patterns)
|
||||
if not _is_llm_error:
|
||||
try:
|
||||
state = await agent.aget_state(config)
|
||||
if getattr(state, "next", None):
|
||||
yield confirmation_payload(conv_id, state, visible_user_text, user_role, env_id)
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
yield json.dumps({
|
||||
"content": f"❌ Ошибка: {exc}",
|
||||
"metadata": {"type": "error", "code": "PROCESSING_ERROR", "detail": str(exc)},
|
||||
|
||||
@@ -9,21 +9,6 @@
|
||||
# @REJECTED Using only environment variables for LLM config was rejected — FastAPI API-based config allows runtime switching without restart.
|
||||
# RunnableWithMessageHistory wrapper is NOT used — PostgresSaver handles history natively.
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import psycopg
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from src.agent._config import FASTAPI_URL, AGENT_CONFIRM_TOOLS, AGENT_INTERRUPT_BEFORE as _INTERRUPT_BEFORE
|
||||
from src.agent._llm_params import chat_openai_kwargs
|
||||
from src.core.logger import logger
|
||||
|
||||
# ── Monkey-patch: OpenAI SDK for Pydantic BaseModel classes ──
|
||||
# LangChain BaseTool objects carry an ``args_schema`` field that is a Pydantic
|
||||
# BaseModel *class* reference (not an instance). When the OpenAI SDK recursively
|
||||
@@ -34,12 +19,24 @@ from src.core.logger import logger
|
||||
# PydanticSerializationError: Unable to serialize unknown type: ModelMetaclass
|
||||
#
|
||||
# The fix: skip model_dump for classes, only dump instances.
|
||||
|
||||
import inspect as _inspect
|
||||
import os
|
||||
|
||||
import httpx
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
import openai._utils._transform as _openai_transform
|
||||
import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
import pydantic as _pydantic
|
||||
import pydantic_core as _pydantic_core
|
||||
|
||||
from src.agent._config import AGENT_CONFIRM_TOOLS, AGENT_INTERRUPT_BEFORE as _INTERRUPT_BEFORE, FASTAPI_URL
|
||||
from src.agent._llm_params import chat_openai_kwargs
|
||||
from src.core.logger import logger
|
||||
|
||||
_original_transform = _openai_transform._async_transform_recursive
|
||||
|
||||
async def _patched_transform(data, *, annotation, inner_type=None):
|
||||
|
||||
@@ -12,7 +12,6 @@ from src.agent.context import get_user_jwt
|
||||
from src.agent.tools import _redact_sensitive_fields
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
# #region AgentChat.Middleware.LoggingMiddleware [C:3] [TYPE Function] [SEMANTICS audit,tool,logging]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Log every tool-call event to assistant_audit table with user context.
|
||||
|
||||
@@ -9,9 +9,10 @@
|
||||
# when GRADIO_ALLOW_PORT_FALLBACK=true and an external proxy is updated separately.
|
||||
# @REJECTED Hardcoding the port was rejected — it must be configurable for different deployment environments.
|
||||
import socket
|
||||
|
||||
import httpx
|
||||
|
||||
from src.agent._config import FASTAPI_URL, SERVICE_JWT, GRADIO_SERVER_NAME, GRADIO_SERVER_PORT, GRADIO_ALLOW_PORT_FALLBACK
|
||||
from src.agent._config import FASTAPI_URL, GRADIO_ALLOW_PORT_FALLBACK, GRADIO_SERVER_NAME, GRADIO_SERVER_PORT, SERVICE_JWT
|
||||
from src.core.cot_logger import seed_trace_id
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -80,6 +81,7 @@ def _fetch_llm_config() -> dict | None:
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
|
||||
from src.agent.app import create_chat_interface
|
||||
from src.agent.context import set_service_jwt
|
||||
from src.agent.langgraph_setup import configure_from_api, init_checkpointer
|
||||
|
||||
241
backend/tests/test_agent/test_agent_confirmation_v2.py
Normal file
241
backend/tests/test_agent/test_agent_confirmation_v2.py
Normal file
@@ -0,0 +1,241 @@
|
||||
# backend/tests/test_agent/test_agent_confirmation_v2.py
|
||||
# #region Test.Agent.ConfirmationV2 [C:3] [TYPE Module] [SEMANTICS test,agent-chat,confirmation,guardrails,hitl]
|
||||
# @BRIEF Contract tests for confirmation_metadata_for_tool and build_confirmation_contract_v2 — three-axis risk, env resolution, permission denied.
|
||||
# @RELATION BINDS_TO -> [AgentChat.Confirmation]
|
||||
# @TEST_FIXTURE: deploy_prod -> guarded + prod env_context
|
||||
# @TEST_FIXTURE: deploy_staging -> guarded + staging env_context
|
||||
# @TEST_FIXTURE: deploy_dev -> guarded + dev env_context
|
||||
# @TEST_FIXTURE: delete_any -> dangerous risk_level
|
||||
# @TEST_FIXTURE: read_only -> safe risk_level
|
||||
# @TEST_FIXTURE: analyst_deploy -> permission_granted=false
|
||||
# @TEST_FIXTURE: env_resolution -> tool_args.env_id > target_env > None
|
||||
|
||||
|
||||
from src.agent._confirmation import (
|
||||
_resolve_env_tier,
|
||||
build_confirmation_contract_v2,
|
||||
confirmation_metadata_for_tool,
|
||||
permission_denied_payload,
|
||||
)
|
||||
|
||||
# ── _resolve_env_tier ───────────────────────────────────────────────
|
||||
|
||||
# #region test_env_resolution_from_tool_args [C:2] [TYPE Function]
|
||||
def test_resolve_env_tier_from_tool_args_env_id():
|
||||
"""env_id in tool_args takes highest priority."""
|
||||
assert _resolve_env_tier({"env_id": "prod-01"}, None) == "prod"
|
||||
assert _resolve_env_tier({"env_id": "prod-01"}, "staging") == "prod" # tool_args wins
|
||||
# #endregion test_env_resolution_from_tool_args
|
||||
|
||||
# #region test_env_resolution_from_environment_id [C:2] [TYPE Function]
|
||||
def test_resolve_env_tier_from_environment_id():
|
||||
"""environment_id in tool_args (alternative key) works."""
|
||||
assert _resolve_env_tier({"environment_id": "staging-v2"}, None) == "staging"
|
||||
# #endregion test_env_resolution_from_environment_id
|
||||
|
||||
# #region test_env_resolution_from_target_env [C:2] [TYPE Function]
|
||||
def test_resolve_env_tier_from_target_env():
|
||||
"""target_env fallback when tool_args has no env."""
|
||||
assert _resolve_env_tier({}, "ss-dev") == "dev"
|
||||
assert _resolve_env_tier({"query": "test"}, "prod-v2") == "prod"
|
||||
# #endregion test_env_resolution_from_target_env
|
||||
|
||||
# #region test_env_resolution_null_when_no_env [C:2] [TYPE Function]
|
||||
def test_resolve_env_tier_null_when_no_env():
|
||||
"""When neither tool_args nor target_env provides env, return None."""
|
||||
assert _resolve_env_tier({}, None) is None
|
||||
assert _resolve_env_tier({"dashboard_id": 42}, None) is None
|
||||
# #endregion test_env_resolution_null_when_no_env
|
||||
|
||||
# #region test_env_resolution_ambiguous_names [C:2] [TYPE Function]
|
||||
def test_resolve_env_tier_ambiguous_names():
|
||||
"""Environment names containing stag/test/local/dev are correctly tiered."""
|
||||
assert _resolve_env_tier({"env_id": "autotest"}, None) == "staging" # "test" in autotest
|
||||
assert _resolve_env_tier({"env_id": "localhost"}, None) == "dev" # "local" in localhost
|
||||
# #endregion test_env_resolution_ambiguous_names
|
||||
|
||||
|
||||
# ── build_confirmation_contract_v2 ───────────────────────────────────
|
||||
|
||||
# #region test_deploy_to_prod_is_guarded_with_prod_context [C:2] [TYPE Function]
|
||||
def test_deploy_to_prod_is_guarded_with_prod_context():
|
||||
"""Deploy to production → guarded risk + prod env_context."""
|
||||
contract = build_confirmation_contract_v2(
|
||||
"deploy_dashboard", {"env_id": "prod-01"}, "admin", "prod-01",
|
||||
)
|
||||
assert contract["risk"] == "write"
|
||||
assert contract["risk_level"] == "guarded"
|
||||
assert contract["dangerous"] is False
|
||||
assert contract["env_context"] == "prod"
|
||||
assert contract["permission_granted"] is True
|
||||
# #endregion test_deploy_to_prod_is_guarded_with_prod_context
|
||||
|
||||
|
||||
# #region test_deploy_to_staging_is_guarded_with_staging_context [C:2] [TYPE Function]
|
||||
def test_deploy_to_staging_is_guarded_with_staging_context():
|
||||
"""Deploy to staging → guarded risk + staging env_context."""
|
||||
contract = build_confirmation_contract_v2(
|
||||
"deploy_dashboard", {"env_id": "staging-v2"}, "admin", "staging-v2",
|
||||
)
|
||||
assert contract["risk"] == "write"
|
||||
assert contract["risk_level"] == "guarded"
|
||||
assert contract["env_context"] == "staging"
|
||||
# #endregion test_deploy_to_staging_is_guarded_with_staging_context
|
||||
|
||||
|
||||
# #region test_deploy_to_dev_is_guarded_with_dev_context [C:2] [TYPE Function]
|
||||
def test_deploy_to_dev_is_guarded_with_dev_context():
|
||||
"""Deploy to dev → guarded risk + dev env_context."""
|
||||
contract = build_confirmation_contract_v2(
|
||||
"deploy_dashboard", {"env_id": "my-dev-env"}, "admin", "my-dev-env",
|
||||
)
|
||||
assert contract["risk"] == "write"
|
||||
assert contract["risk_level"] == "guarded"
|
||||
assert contract["env_context"] == "dev"
|
||||
# #endregion test_deploy_to_dev_is_guarded_with_dev_context
|
||||
|
||||
|
||||
# #region test_delete_operation_is_dangerous [C:2] [TYPE Function]
|
||||
def test_delete_operation_is_dangerous():
|
||||
"""Delete-prefixed tools should be classified as dangerous."""
|
||||
contract = build_confirmation_contract_v2(
|
||||
"delete_dashboard", {}, "admin",
|
||||
)
|
||||
assert contract["risk"] == "write"
|
||||
assert contract["risk_level"] == "dangerous"
|
||||
assert contract["dangerous"] is True
|
||||
# #endregion test_delete_operation_is_dangerous
|
||||
|
||||
|
||||
# #region test_read_only_tool_is_safe [C:2] [TYPE Function]
|
||||
def test_read_only_tool_is_safe():
|
||||
"""Search/list/get tools should be classified as safe/read."""
|
||||
contract = build_confirmation_contract_v2(
|
||||
"search_dashboards", {"env_id": "prod"}, "admin", "prod",
|
||||
)
|
||||
assert contract["risk"] == "read"
|
||||
assert contract["risk_level"] == "safe"
|
||||
assert contract["dangerous"] is False
|
||||
# #endregion test_read_only_tool_is_safe
|
||||
|
||||
|
||||
# #region test_viewer_gets_permission_denied [C:2] [TYPE Function]
|
||||
def test_viewer_permission_denied_for_write_tools():
|
||||
"""Viewer role should be denied permission for write tools."""
|
||||
contract = build_confirmation_contract_v2(
|
||||
"deploy_dashboard", {"env_id": "prod"}, "viewer", "prod",
|
||||
)
|
||||
assert contract["permission_granted"] is False
|
||||
assert contract["required_role"] == "admin"
|
||||
assert contract["alternatives"] is not None
|
||||
assert len(contract["alternatives"]) >= 1
|
||||
# #endregion test_viewer_gets_permission_denied
|
||||
|
||||
|
||||
# #region test_analyst_permission_denied_specific_tools [C:2] [TYPE Function]
|
||||
def test_analyst_permission_denied_for_admin_tools():
|
||||
"""Analyst/editor role should be denied for admin-only tools."""
|
||||
for tool_name in ("deploy_dashboard", "execute_migration", "run_backup"):
|
||||
contract = build_confirmation_contract_v2(tool_name, {}, "analyst")
|
||||
assert contract["permission_granted"] is False, f"{tool_name} should be denied for analyst"
|
||||
# #endregion test_analyst_permission_denied_specific_tools
|
||||
|
||||
|
||||
# #region test_admin_write_tool_permission_granted [C:2] [TYPE Function]
|
||||
def test_admin_write_tool_permission_granted():
|
||||
"""Admin role should always have permission_granted=True for guarded tools."""
|
||||
write_tools = ["deploy_dashboard", "commit_changes", "execute_migration"]
|
||||
for tool_name in write_tools:
|
||||
contract = build_confirmation_contract_v2(tool_name, {}, "admin")
|
||||
assert contract["permission_granted"] is True, f"{tool_name} should be allowed for admin"
|
||||
# #endregion test_admin_write_tool_permission_granted
|
||||
|
||||
|
||||
# ── confirmation_metadata_for_tool ────────────────────────────────────
|
||||
|
||||
# #region test_metadata_for_tool_contains_required_fields [C:2] [TYPE Function]
|
||||
def test_metadata_for_tool_contains_all_required_fields():
|
||||
"""confirmation_metadata_for_tool should produce a complete metadata dict."""
|
||||
meta = confirmation_metadata_for_tool(
|
||||
"conv-123", "deploy_dashboard",
|
||||
{"env_id": "prod-01"}, "admin", "prod-01",
|
||||
)
|
||||
required_fields = [
|
||||
"type", "thread_id", "prompt", "tool_name", "tool_args",
|
||||
"risk", "risk_level", "requires_confirmation",
|
||||
"dangerous", "env_context",
|
||||
]
|
||||
for field in required_fields:
|
||||
assert field in meta, f"Missing field '{field}' in metadata"
|
||||
assert meta["type"] == "confirm_required"
|
||||
assert meta["thread_id"] == "conv-123"
|
||||
assert meta["tool_name"] == "deploy_dashboard"
|
||||
assert meta["risk"] == "write"
|
||||
assert meta["risk_level"] == "guarded"
|
||||
assert meta["requires_confirmation"] is True
|
||||
# #endregion test_metadata_for_tool_contains_required_fields
|
||||
|
||||
|
||||
# ── permission_denied_payload ─────────────────────────────────────────
|
||||
|
||||
# #region test_permission_denied_payload_structure [C:2] [TYPE Function]
|
||||
def test_permission_denied_payload_structure():
|
||||
"""permission_denied_payload should produce valid JSON with correct type."""
|
||||
import json
|
||||
|
||||
payload_str = permission_denied_payload("deploy_dashboard", "admin", "viewer")
|
||||
payload = json.loads(payload_str)
|
||||
|
||||
assert payload["content"] == "⛔ Недостаточно прав для deploy_dashboard"
|
||||
assert payload["metadata"]["type"] == "permission_denied"
|
||||
assert payload["metadata"]["tool_name"] == "deploy_dashboard"
|
||||
assert payload["metadata"]["required_role"] == "admin"
|
||||
assert payload["metadata"]["user_role"] == "viewer"
|
||||
assert payload["metadata"]["alternatives"] == []
|
||||
# #endregion test_permission_denied_payload_structure
|
||||
|
||||
|
||||
# #region test_permission_denied_payload_with_alternatives [C:2] [TYPE Function]
|
||||
def test_permission_denied_payload_with_alternatives():
|
||||
"""Alternatives list should be preserved in the payload."""
|
||||
import json
|
||||
|
||||
alternatives = [{"action": "get_health_summary", "prompt": "Check system health"}]
|
||||
payload_str = permission_denied_payload(
|
||||
"deploy_dashboard", "admin", "viewer", alternatives,
|
||||
)
|
||||
payload = json.loads(payload_str)
|
||||
|
||||
assert payload["metadata"]["alternatives"] == alternatives
|
||||
# #endregion test_permission_denied_payload_with_alternatives
|
||||
|
||||
|
||||
# ── Unknown/null tool ─────────────────────────────────────────────────
|
||||
|
||||
# #region test_null_tool_name_handled [C:2] [TYPE Function]
|
||||
def test_null_tool_name_handled_gracefully():
|
||||
"""Null tool_name should produce 'unknown_action' fallback."""
|
||||
contract = build_confirmation_contract_v2(None)
|
||||
assert contract["operation"] == "unknown_action"
|
||||
assert contract["risk_level"] == "safe"
|
||||
# #endregion test_null_tool_name_handled
|
||||
|
||||
|
||||
# #region test_execute_migration_is_guarded [C:2] [TYPE Function]
|
||||
def test_execute_migration_is_guarded():
|
||||
"""execute_migration should be classified as guarded (write)."""
|
||||
contract = build_confirmation_contract_v2("execute_migration", {})
|
||||
assert contract["risk"] == "write"
|
||||
assert contract["risk_level"] == "guarded"
|
||||
# #endregion test_execute_migration_is_guarded
|
||||
|
||||
|
||||
# #region test_commit_changes_is_guarded [C:2] [TYPE Function]
|
||||
def test_commit_changes_is_guarded():
|
||||
"""commit_changes should be classified as guarded (write)."""
|
||||
contract = build_confirmation_contract_v2("commit_changes", {})
|
||||
assert contract["risk"] == "write"
|
||||
assert contract["risk_level"] == "guarded"
|
||||
# #endregion test_commit_changes_is_guarded
|
||||
|
||||
# #endregion Test.Agent.ConfirmationV2
|
||||
218
backend/tests/test_agent/test_agent_context.py
Normal file
218
backend/tests/test_agent/test_agent_context.py
Normal file
@@ -0,0 +1,218 @@
|
||||
# backend/tests/test_agent/test_agent_context.py
|
||||
# #region Test.Agent.Context [C:3] [TYPE Module] [SEMANTICS test,agent-chat,context,validation,uicontext]
|
||||
# @BRIEF Contract tests for UIContext validation — enum checks, size limits, field lengths, boundary values.
|
||||
# @RELATION BINDS_TO -> [AgentChat.Context.Validate]
|
||||
# @TEST_FIXTURE: null -> returns {}
|
||||
# @TEST_FIXTURE: dashboard+id+name -> passes
|
||||
# @TEST_FIXTURE: malformed_objectType -> raises UIContextValidationError
|
||||
# @TEST_FIXTURE: objectName>256 -> raises
|
||||
# @TEST_FIXTURE: oversized>4KB -> raises
|
||||
# @TEST_FIXTURE: contextVersion!=1 -> raises
|
||||
# @TEST_FIXTURE: dataset_context -> passes
|
||||
# @TEST_FIXTURE: migration_context -> passes
|
||||
# @TEST_FIXTURE: non_numeric_objectId -> raises
|
||||
# @TEST_FIXTURE: empty_route -> passes
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from src.agent._context import UIContextValidationError, validate_uicontext
|
||||
|
||||
|
||||
# #region test_null_payload_returns_empty [C:2] [TYPE Function]
|
||||
def test_null_payload_returns_empty_dict():
|
||||
"""Null payloads should safely return an empty dict."""
|
||||
assert validate_uicontext(None) == {}
|
||||
# #endregion test_null_payload_returns_empty
|
||||
|
||||
|
||||
# #region test_valid_dashboard_context_passes [C:2] [TYPE Function]
|
||||
def test_valid_dashboard_context_passes():
|
||||
"""Full dashboard UIContext with all fields should validate cleanly."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "42",
|
||||
"objectName": "Energy Report",
|
||||
"envId": "ss-dev",
|
||||
"route": "/dashboards/42",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
result = validate_uicontext(payload)
|
||||
assert result == payload
|
||||
# #endregion test_valid_dashboard_context_passes
|
||||
|
||||
|
||||
# #region test_valid_dataset_context_passes [C:2] [TYPE Function]
|
||||
def test_valid_dataset_context_passes():
|
||||
"""Dataset UIContext should pass validation."""
|
||||
payload = {
|
||||
"objectType": "dataset",
|
||||
"objectId": "15",
|
||||
"objectName": None,
|
||||
"envId": "staging",
|
||||
"route": "/datasets/15",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
assert validate_uicontext(payload) == payload
|
||||
# #endregion test_valid_dataset_context_passes
|
||||
|
||||
|
||||
# #region test_valid_migration_context_passes [C:2] [TYPE Function]
|
||||
def test_valid_migration_context_passes():
|
||||
"""Migration UIContext should pass validation."""
|
||||
payload = {
|
||||
"objectType": "migration",
|
||||
"objectId": None,
|
||||
"objectName": None,
|
||||
"envId": None,
|
||||
"route": "/migrations",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
assert validate_uicontext(payload) == payload
|
||||
# #endregion test_valid_migration_context_passes
|
||||
|
||||
|
||||
# #region test_invalid_object_type_raises [C:2] [TYPE Function]
|
||||
def test_invalid_object_type_raises_validation_error():
|
||||
"""Unknown objectType values should be rejected."""
|
||||
payload = {
|
||||
"objectType": "chart",
|
||||
"objectId": "42",
|
||||
"envId": "ss-dev",
|
||||
"route": "/dashboards/42",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
with pytest.raises(UIContextValidationError, match="objectType"):
|
||||
validate_uicontext(payload)
|
||||
# #endregion test_invalid_object_type_raises
|
||||
|
||||
|
||||
# #region test_invalid_object_id_raises [C:2] [TYPE Function]
|
||||
def test_non_numeric_object_id_raises():
|
||||
"""ObjectId must be a numeric string or None."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "abc-not-a-number",
|
||||
"route": "/dashboards/42",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
with pytest.raises(UIContextValidationError, match="objectId"):
|
||||
validate_uicontext(payload)
|
||||
# #endregion test_invalid_object_id_raises
|
||||
|
||||
|
||||
# #region test_object_name_exceeds_limit_raises [C:2] [TYPE Function]
|
||||
def test_object_name_exceeds_256_chars_raises():
|
||||
"""ObjectName longer than 256 characters should be rejected."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "42",
|
||||
"objectName": "A" * 257,
|
||||
"route": "/dashboards/42",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
with pytest.raises(UIContextValidationError, match="objectName"):
|
||||
validate_uicontext(payload)
|
||||
# #endregion test_object_name_exceeds_limit_raises
|
||||
|
||||
|
||||
# #region test_object_name_at_boundary_passes [C:2] [TYPE Function]
|
||||
def test_object_name_at_256_chars_passes():
|
||||
"""ObjectName at exactly 256 characters should pass."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "42",
|
||||
"objectName": "A" * 256,
|
||||
"route": "/dashboards/42",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
assert validate_uicontext(payload) == payload
|
||||
# #endregion test_object_name_at_boundary_passes
|
||||
|
||||
|
||||
# #region test_invalid_context_version_raises [C:2] [TYPE Function]
|
||||
def test_invalid_context_version_raises():
|
||||
"""Only contextVersion=1 is currently supported."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "42",
|
||||
"route": "/dashboards/42",
|
||||
"contextVersion": 2,
|
||||
}
|
||||
with pytest.raises(UIContextValidationError, match="contextVersion"):
|
||||
validate_uicontext(payload)
|
||||
# #endregion test_invalid_context_version_raises
|
||||
|
||||
|
||||
# #region test_missing_context_version_raises [C:2] [TYPE Function]
|
||||
def test_missing_context_version_raises():
|
||||
"""ContextVersion should always be present and equal 1."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "42",
|
||||
"route": "/dashboards/42",
|
||||
}
|
||||
with pytest.raises(UIContextValidationError):
|
||||
validate_uicontext(payload)
|
||||
# #endregion test_missing_context_version_raises
|
||||
|
||||
|
||||
# #region test_payload_exceeds_4kb_raises [C:2] [TYPE Function]
|
||||
def test_payload_exceeds_4kb_raises():
|
||||
"""Payloads larger than 4 KB should be rejected to prevent prompt injection."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "42",
|
||||
"route": "/dashboards/42",
|
||||
"contextVersion": 1,
|
||||
# Create a field that pushes the serialized JSON over 4096 bytes
|
||||
"padding": "X" * 4096,
|
||||
}
|
||||
with pytest.raises(UIContextValidationError, match="exceeds 4 KB"):
|
||||
validate_uicontext(payload)
|
||||
# #endregion test_payload_exceeds_4kb_raises
|
||||
|
||||
|
||||
# #region test_route_exceeds_512_chars_raises [C:2] [TYPE Function]
|
||||
def test_route_exceeds_512_chars_raises():
|
||||
"""Route length should be capped at 512 characters."""
|
||||
payload = {
|
||||
"objectType": "dashboard",
|
||||
"objectId": "42",
|
||||
"route": "/" + "a" * 512,
|
||||
"contextVersion": 1,
|
||||
}
|
||||
with pytest.raises(UIContextValidationError, match="route"):
|
||||
validate_uicontext(payload)
|
||||
# #endregion test_route_exceeds_512_chars_raises
|
||||
|
||||
|
||||
# #region test_env_id_none_and_string_accepted [C:2] [TYPE Function]
|
||||
def test_env_id_none_and_string_accepted():
|
||||
"""envId should accept None and string values."""
|
||||
assert validate_uicontext({
|
||||
"objectType": "dashboard", "objectId": "42",
|
||||
"envId": None, "route": "/dashboards/42", "contextVersion": 1,
|
||||
})["envId"] is None
|
||||
|
||||
assert validate_uicontext({
|
||||
"objectType": "dashboard", "objectId": "42",
|
||||
"envId": "ss-dev", "route": "/dashboards/42", "contextVersion": 1,
|
||||
})["envId"] == "ss-dev"
|
||||
# #endregion test_env_id_none_and_string_accepted
|
||||
|
||||
|
||||
# #region test_without_object_type_passes [C:2] [TYPE Function]
|
||||
def test_no_object_type_with_env_passes():
|
||||
"""Context without objectType (general mode) should pass validation."""
|
||||
payload = {
|
||||
"objectType": None,
|
||||
"objectId": None,
|
||||
"envId": "ss-dev",
|
||||
"route": "/settings",
|
||||
"contextVersion": 1,
|
||||
}
|
||||
assert validate_uicontext(payload) == payload
|
||||
# #endregion test_without_object_type_passes
|
||||
|
||||
# #endregion Test.Agent.Context
|
||||
255
backend/tests/test_agent/test_agent_tool_filter.py
Normal file
255
backend/tests/test_agent/test_agent_tool_filter.py
Normal file
@@ -0,0 +1,255 @@
|
||||
# backend/tests/test_agent/test_agent_tool_filter.py
|
||||
# #region Test.Agent.ToolFilter [C:3] [TYPE Module] [SEMANTICS test,agent-chat,tools,filter,pipeline,rbac]
|
||||
# @BRIEF Contract tests for build_tool_pipeline and enforce_tool_permission — two-layer RBAC + context affinity.
|
||||
# @RELATION BINDS_TO -> [AgentChat.ToolFilter]
|
||||
# @TEST_FIXTURE: null_object_type -> all RBAC-allowed tools pass
|
||||
# @TEST_FIXTURE: dashboard+admin -> all dashboard tools + capabilities kept
|
||||
# @TEST_FIXTURE: dashboard+analyst -> admin-only tools removed from dashboard
|
||||
# @TEST_FIXTURE: dataset+admin -> all dataset tools + capabilities kept
|
||||
# @TEST_FIXTURE: dataset+viewer -> admin-only tools removed from dataset
|
||||
# @TEST_FIXTURE: migration+admin -> all migration tools + capabilities kept
|
||||
# @TEST_FIXTURE: unknown_object_type -> graceful fallback to full list
|
||||
# @TEST_FIXTURE: invocation_guard_admin -> deploy_dashboard allowed for admin
|
||||
# @TEST_FIXTURE: invocation_guard_viewer -> deploy_dashboard rejected for viewer
|
||||
# @TEST_FIXTURE: unknown_tool_invocation -> always allowed
|
||||
# @TEST_FIXTURE: show_capabilities_always_included -> mandatory tool
|
||||
# @TEST_FIXTURE: idempotent -> build_tool_pipeline never mutates input list
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
from src.agent._tool_filter import (
|
||||
_CONTEXT_TOOL_AFFINITY,
|
||||
_TOOL_PERMISSIONS,
|
||||
build_tool_pipeline,
|
||||
enforce_tool_permission,
|
||||
)
|
||||
|
||||
|
||||
# #region _tools_helper [C:1] [TYPE Function] [SEMANTICS test,helper]
|
||||
# @BRIEF Create SimpleNamespace-based tool mimics with .name attribute for pipeline tests.
|
||||
def _tools(names: list[str]) -> list[SimpleNamespace]:
|
||||
return [SimpleNamespace(name=name) for name in names]
|
||||
# #endregion _tools_helper
|
||||
|
||||
|
||||
# ── build_tool_pipeline ──────────────────────────────────────────────
|
||||
|
||||
# #region test_null_object_type_returns_all_rbac_allowed [C:2] [TYPE Function]
|
||||
def test_null_object_type_returns_all_rbac_allowed_tools():
|
||||
"""With no object_type, all tools pass except those blocked by RBAC (viewer)."""
|
||||
tools = _tools(["search_dashboards", "deploy_dashboard", "show_capabilities"])
|
||||
|
||||
# Admin: all tools pass
|
||||
result_admin = [t.name for t in build_tool_pipeline(tools, "admin", None)]
|
||||
assert "deploy_dashboard" in result_admin
|
||||
assert "search_dashboards" in result_admin
|
||||
assert "show_capabilities" in result_admin
|
||||
|
||||
# Viewer: deploy_dashboard blocked by RBAC
|
||||
result_viewer = [t.name for t in build_tool_pipeline(tools, "viewer", None)]
|
||||
assert "deploy_dashboard" not in result_viewer
|
||||
assert "search_dashboards" in result_viewer
|
||||
assert "show_capabilities" in result_viewer
|
||||
# #endregion test_null_object_type_returns_all_rbac_allowed
|
||||
|
||||
|
||||
# #region test_dashboard_context_admin [C:2] [TYPE Function]
|
||||
def test_dashboard_context_admin_keeps_affinity_tools():
|
||||
"""Dashboard context + admin role: keep all dashboard tools + capabilities."""
|
||||
tools = _tools([
|
||||
"search_dashboards", "get_health_summary", "deploy_dashboard",
|
||||
"run_llm_validation", "run_llm_documentation", "execute_migration",
|
||||
"create_branch", "commit_changes", "show_capabilities",
|
||||
# Non-dashboard tools (should be excluded)
|
||||
"run_backup", "superset_execute_sql", "list_environments",
|
||||
])
|
||||
result = [t.name for t in build_tool_pipeline(tools, "admin", "dashboard")]
|
||||
assert "search_dashboards" in result
|
||||
assert "get_health_summary" in result
|
||||
assert "deploy_dashboard" in result
|
||||
assert "show_capabilities" in result
|
||||
assert "run_backup" not in result
|
||||
assert "superset_execute_sql" not in result
|
||||
# #endregion test_dashboard_context_admin
|
||||
|
||||
|
||||
# #region test_dashboard_context_viewer [C:2] [TYPE Function]
|
||||
def test_dashboard_context_viewer_removes_admin_only():
|
||||
"""Dashboard context + viewer: admin-only tools removed even from dashboard affinity."""
|
||||
tools = _tools([
|
||||
"search_dashboards", "deploy_dashboard", "execute_migration",
|
||||
"commit_changes", "run_llm_validation", "show_capabilities",
|
||||
])
|
||||
result = [t.name for t in build_tool_pipeline(tools, "viewer", "dashboard")]
|
||||
assert "search_dashboards" in result
|
||||
assert "run_llm_validation" in result
|
||||
assert "show_capabilities" in result
|
||||
assert "deploy_dashboard" not in result
|
||||
assert "execute_migration" not in result
|
||||
assert "commit_changes" not in result
|
||||
# #endregion test_dashboard_context_viewer
|
||||
|
||||
|
||||
# #region test_dataset_context_admin [C:2] [TYPE Function]
|
||||
def test_dataset_context_admin_keeps_dataset_tools():
|
||||
"""Dataset context + admin: keep dataset affinity tools, exclude non-dataset."""
|
||||
tools = _tools([
|
||||
"superset_explore_database", "superset_format_sql", "superset_execute_sql",
|
||||
"superset_audit_permissions", "superset_create_dataset",
|
||||
"search_dashboards", "get_task_status", "list_environments",
|
||||
"show_capabilities",
|
||||
# Non-dataset tools
|
||||
"deploy_dashboard", "run_backup", "start_maintenance",
|
||||
])
|
||||
result = [t.name for t in build_tool_pipeline(tools, "admin", "dataset")]
|
||||
assert "superset_explore_database" in result
|
||||
assert "superset_execute_sql" in result
|
||||
assert "superset_format_sql" in result
|
||||
assert "show_capabilities" in result
|
||||
assert "deploy_dashboard" not in result
|
||||
assert "run_backup" not in result
|
||||
# #endregion test_dataset_context_admin
|
||||
|
||||
|
||||
# #region test_dataset_context_viewer [C:2] [TYPE Function]
|
||||
def test_dataset_context_viewer_removes_admin_tools():
|
||||
"""Dataset context + viewer: admin-only tools in dataset affinity removed."""
|
||||
tools = _tools([
|
||||
"superset_explore_database", "superset_execute_sql",
|
||||
"search_dashboards", "show_capabilities",
|
||||
"start_maintenance",
|
||||
])
|
||||
result = [t.name for t in build_tool_pipeline(tools, "viewer", "dataset")]
|
||||
assert "superset_explore_database" in result
|
||||
assert "superset_execute_sql" in result
|
||||
assert "show_capabilities" in result
|
||||
assert "start_maintenance" not in result
|
||||
# #endregion test_dataset_context_viewer
|
||||
|
||||
|
||||
# #region test_migration_context_admin [C:2] [TYPE Function]
|
||||
def test_migration_context_admin_keeps_migration_tools():
|
||||
"""Migration context + admin: keep migration affinity tools, exclude others."""
|
||||
tools = _tools([
|
||||
"execute_migration", "search_dashboards", "get_health_summary",
|
||||
"deploy_dashboard", "list_environments", "show_capabilities",
|
||||
# Non-migration tools
|
||||
"run_backup", "superset_execute_sql",
|
||||
])
|
||||
result = [t.name for t in build_tool_pipeline(tools, "admin", "migration")]
|
||||
assert "execute_migration" in result
|
||||
assert "search_dashboards" in result
|
||||
assert "show_capabilities" in result
|
||||
assert "run_backup" not in result
|
||||
assert "superset_execute_sql" not in result
|
||||
# #endregion test_migration_context_admin
|
||||
|
||||
|
||||
# #region test_unknown_object_type_falls_back [C:2] [TYPE Function]
|
||||
def test_unknown_object_type_falls_back_to_full_list():
|
||||
"""Unknown object_type should not filter — behaves like null context."""
|
||||
tools = _tools([
|
||||
"search_dashboards", "deploy_dashboard", "run_backup",
|
||||
"superset_execute_sql", "show_capabilities",
|
||||
])
|
||||
result = [t.name for t in build_tool_pipeline(tools, "admin", "unknown_type")]
|
||||
# All RBAC-allowed tools pass (admin sees all)
|
||||
assert "search_dashboards" in result
|
||||
assert "deploy_dashboard" in result
|
||||
assert "run_backup" in result
|
||||
assert "superset_execute_sql" in result
|
||||
assert "show_capabilities" in result
|
||||
# #endregion test_unknown_object_type_falls_back
|
||||
|
||||
|
||||
# #region test_show_capabilities_always_included [C:2] [TYPE Function]
|
||||
def test_show_capabilities_always_included():
|
||||
"""show_capabilities must survive all filtering stages regardless of context."""
|
||||
tools = _tools(["show_capabilities"])
|
||||
# Must appear in any context+role combination
|
||||
for obj_type in (None, "dashboard", "dataset", "migration"):
|
||||
for role in ("admin", "editor", "viewer"):
|
||||
result = [t.name for t in build_tool_pipeline(tools, role, obj_type)]
|
||||
assert "show_capabilities" in result, (
|
||||
f"show_capabilities missing for role={role}, object_type={obj_type}"
|
||||
)
|
||||
# #endregion test_show_capabilities_always_included
|
||||
|
||||
|
||||
# #region test_pipeline_is_idempotent [C:2] [TYPE Function]
|
||||
def test_pipeline_does_not_mutate_input_list():
|
||||
"""build_tool_pipeline must return a new list and not mutate the input."""
|
||||
tools = _tools(["search_dashboards", "show_capabilities"])
|
||||
original_ids = [id(t) for t in tools]
|
||||
_ = build_tool_pipeline(tools, "admin", "dashboard")
|
||||
# Input list unchanged
|
||||
assert [id(t) for t in tools] == original_ids
|
||||
assert len(tools) == 2
|
||||
# #endregion test_pipeline_is_idempotent
|
||||
|
||||
|
||||
# ── enforce_tool_permission ──────────────────────────────────────────
|
||||
|
||||
# #region test_invocation_guard_admin_allowed [C:2] [TYPE Function]
|
||||
def test_enforce_tool_permission_admin_allowed():
|
||||
"""Admin role should be allowed to invoke all restricted tools."""
|
||||
restricted = ["deploy_dashboard", "commit_changes", "create_branch",
|
||||
"run_backup", "execute_migration", "start_maintenance", "end_maintenance"]
|
||||
for tool_name in restricted:
|
||||
assert enforce_tool_permission(tool_name, "admin") is True, (
|
||||
f"Admin should be allowed to invoke '{tool_name}'"
|
||||
)
|
||||
# #endregion test_invocation_guard_admin_allowed
|
||||
|
||||
|
||||
# #region test_invocation_guard_viewer_denied [C:2] [TYPE Function]
|
||||
def test_enforce_tool_permission_viewer_denied():
|
||||
"""Viewer role should be denied for all restricted tools."""
|
||||
restricted = ["deploy_dashboard", "commit_changes", "create_branch",
|
||||
"run_backup", "execute_migration", "start_maintenance", "end_maintenance"]
|
||||
for tool_name in restricted:
|
||||
assert enforce_tool_permission(tool_name, "viewer") is False, (
|
||||
f"Viewer should NOT be allowed to invoke '{tool_name}'"
|
||||
)
|
||||
# #endregion test_invocation_guard_viewer_denied
|
||||
|
||||
|
||||
# #region test_invocation_guard_unknown_tool_allowed [C:2] [TYPE Function]
|
||||
def test_enforce_tool_permission_unknown_tool_always_allowed():
|
||||
"""Unknown tools (not in _TOOL_PERMISSIONS) should always be allowed."""
|
||||
assert enforce_tool_permission("search_dashboards", "viewer") is True
|
||||
assert enforce_tool_permission("show_capabilities", "viewer") is True
|
||||
assert enforce_tool_permission("nonexistent_tool", "viewer") is True
|
||||
# #endregion test_invocation_guard_unknown_tool_allowed
|
||||
|
||||
|
||||
# #region test_context_affinity_coverage [C:2] [TYPE Function]
|
||||
def test_context_affinity_maps_exist_for_all_expected_types():
|
||||
"""Verify that all three known object types have affinity mappings."""
|
||||
assert "dashboard" in _CONTEXT_TOOL_AFFINITY
|
||||
assert "dataset" in _CONTEXT_TOOL_AFFINITY
|
||||
assert "migration" in _CONTEXT_TOOL_AFFINITY
|
||||
# Each mapping should have at least 3 tools
|
||||
for obj_type in ("dashboard", "dataset", "migration"):
|
||||
assert len(_CONTEXT_TOOL_AFFINITY[obj_type]) >= 3, (
|
||||
f"Context affinity for '{obj_type}' should have ≥3 tools"
|
||||
)
|
||||
# #endregion test_context_affinity_coverage
|
||||
|
||||
|
||||
# #region test_rbac_maps_exist_for_write_tools [C:2] [TYPE Function]
|
||||
def test_rbac_permissions_for_write_tools():
|
||||
"""Verify all admin-only tools are explicitly listed in _TOOL_PERMISSIONS."""
|
||||
expected_admin_tools = [
|
||||
"deploy_dashboard", "commit_changes", "create_branch",
|
||||
"run_backup", "execute_migration", "start_maintenance", "end_maintenance",
|
||||
]
|
||||
for tool_name in expected_admin_tools:
|
||||
assert tool_name in _TOOL_PERMISSIONS, (
|
||||
f"'{tool_name}' should be in _TOOL_PERMISSIONS"
|
||||
)
|
||||
assert "admin" in _TOOL_PERMISSIONS[tool_name]
|
||||
# #endregion test_rbac_permissions_for_write_tools
|
||||
|
||||
|
||||
# #endregion Test.Agent.ToolFilter
|
||||
@@ -7,10 +7,10 @@ import sys
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
||||
def _make_async_iter(items):
|
||||
@@ -24,15 +24,20 @@ def _make_async_iter(items):
|
||||
try:
|
||||
return next(self._iter)
|
||||
except StopIteration:
|
||||
raise StopAsyncIteration
|
||||
raise StopAsyncIteration from None
|
||||
return AsyncIter(items)
|
||||
|
||||
|
||||
def _skip_pipeline(results: list[str]) -> list[str]:
|
||||
"""Filter out pipeline_result events emitted before the agent stream loop."""
|
||||
return [r for r in results if json.loads(r).get("metadata", {}).get("type") != "pipeline_result"]
|
||||
|
||||
|
||||
def _make_agent_mock(stream_events=None, raise_on_call=False):
|
||||
"""Create a properly mocked agent for async iteration."""
|
||||
agent = MagicMock()
|
||||
if raise_on_call:
|
||||
agent.astream_events = MagicMock(side_effect=stream_events if isinstance(stream_events, list) else stream_events)
|
||||
agent.astream_events = MagicMock(side_effect=stream_events)
|
||||
elif stream_events is not None:
|
||||
agent.astream_events = MagicMock(return_value=_make_async_iter(stream_events))
|
||||
else:
|
||||
@@ -160,7 +165,7 @@ class TestConfirmationMetadata:
|
||||
assert meta["risk_level"] == "safe"
|
||||
|
||||
def test_fast_resume_deny_closes_without_langgraph(self):
|
||||
from src.agent._confirmation import handle_resume, _pending_confirmations
|
||||
from src.agent._confirmation import _pending_confirmations, handle_resume
|
||||
|
||||
_pending_confirmations["conv-fast-deny"] = {
|
||||
"tool_name": "list_environments",
|
||||
@@ -175,7 +180,7 @@ class TestConfirmationMetadata:
|
||||
assert data["metadata"] == {"type": "confirm_resolved", "result": "denied"}
|
||||
|
||||
def test_fast_resume_confirm_executes_tool_directly(self):
|
||||
from src.agent._confirmation import handle_resume, _pending_confirmations
|
||||
from src.agent._confirmation import _pending_confirmations, handle_resume
|
||||
|
||||
tool = MagicMock()
|
||||
tool.ainvoke = AsyncMock(return_value='[{"id":"ss-dev"}]')
|
||||
@@ -188,7 +193,7 @@ class TestConfirmationMetadata:
|
||||
with patch("src.agent._confirmation.find_tool", return_value=tool), \
|
||||
patch("src.agent._confirmation._format_tool_output_via_llm") as mock_format:
|
||||
# Make the mock format helper yield nothing — don't need LLM in unit test
|
||||
async def _empty_format(*args, **kwargs):
|
||||
async def _empty_format(*_args, **_kwargs):
|
||||
return
|
||||
yield # pragma: no cover — async generator requires at least one yield
|
||||
mock_format.side_effect = _empty_format
|
||||
@@ -205,7 +210,7 @@ class TestConfirmationMetadata:
|
||||
class TestAgentHandler:
|
||||
@pytest.mark.asyncio
|
||||
async def test_handles_concurrent_send(self, mock_request):
|
||||
from src.agent.app import agent_handler, _user_locks
|
||||
from src.agent.app import _user_locks, agent_handler
|
||||
# Handler resolves user_id as "admin" when no user_id_str or JWT is provided
|
||||
_user_locks["admin"] = True
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
@@ -215,7 +220,7 @@ class TestAgentHandler:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handles_file_too_large(self, mock_request):
|
||||
from src.agent.app import agent_handler, MAX_FILE_SIZE_BYTES
|
||||
from src.agent.app import MAX_FILE_SIZE_BYTES, agent_handler
|
||||
message = {"text": "analyze", "files": ["fake_path"]}
|
||||
with patch("os.path.exists", return_value=True), \
|
||||
patch("os.path.getsize", return_value=MAX_FILE_SIZE_BYTES + 1):
|
||||
@@ -267,8 +272,9 @@ class TestAgentHandler:
|
||||
patch("src.agent.app.get_all_tools", return_value=[]), \
|
||||
patch("src.agent.app.save_conversation", AsyncMock()):
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
assert len(results) > 0
|
||||
token_data = json.loads(results[0])
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) > 0
|
||||
token_data = json.loads(filtered[0])
|
||||
assert token_data["metadata"]["type"] == "stream_token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -297,7 +303,8 @@ class TestAgentHandler:
|
||||
"ss-dev",
|
||||
)]
|
||||
|
||||
assert len(results) == 1
|
||||
# pipeline_result + stream_token = 2 events
|
||||
assert len(results) == 2
|
||||
messages_arg = agent.astream_events.call_args.args[0]["messages"]
|
||||
llm_text = messages_arg[0].content
|
||||
assert "[RUNTIME CONTEXT]" in llm_text
|
||||
@@ -320,9 +327,10 @@ class TestAgentHandler:
|
||||
patch("src.agent.app.save_conversation", AsyncMock()), \
|
||||
patch("src.agent.app.log_tool_event", AsyncMock()):
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
assert len(results) == 2
|
||||
assert json.loads(results[0])["metadata"]["type"] == "tool_start"
|
||||
assert json.loads(results[1])["metadata"]["type"] == "tool_end"
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 2
|
||||
assert json.loads(filtered[0])["metadata"]["type"] == "tool_start"
|
||||
assert json.loads(filtered[1])["metadata"]["type"] == "tool_end"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_error_event(self, mock_request):
|
||||
@@ -337,20 +345,22 @@ class TestAgentHandler:
|
||||
patch("src.agent.app.save_conversation", AsyncMock()), \
|
||||
patch("src.agent.app.log_tool_event", AsyncMock()):
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
assert len(results) == 1
|
||||
assert json.loads(results[0])["metadata"]["type"] == "tool_error"
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 1
|
||||
assert json.loads(filtered[0])["metadata"]["type"] == "tool_error"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_parser_exception_retry(self, mock_request):
|
||||
from src.agent.app import agent_handler
|
||||
from langchain_core.exceptions import OutputParserException
|
||||
|
||||
from src.agent.app import agent_handler
|
||||
|
||||
mock_stream_after = _make_async_iter([
|
||||
{"event": "on_chat_model_stream", "data": {"chunk": MagicMock(content="retried")}}
|
||||
])
|
||||
call_count = [0]
|
||||
|
||||
def mock_astream(*args, **kwargs):
|
||||
def mock_astream(*_args, **_kwargs):
|
||||
call_count[0] += 1
|
||||
if call_count[0] == 1:
|
||||
raise OutputParserException("bad output")
|
||||
@@ -363,16 +373,18 @@ class TestAgentHandler:
|
||||
patch("src.agent.app.get_all_tools", return_value=[]), \
|
||||
patch("src.agent.app.save_conversation", AsyncMock()):
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
assert len(results) > 0
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_parser_exception_final_failure(self, mock_request):
|
||||
from src.agent.app import agent_handler
|
||||
from langchain_core.exceptions import OutputParserException
|
||||
|
||||
from src.agent.app import agent_handler
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def mock_astream(*args, **kwargs):
|
||||
def mock_astream(*_args, **_kwargs):
|
||||
call_count[0] += 1
|
||||
raise OutputParserException(f"bad {call_count[0]}")
|
||||
|
||||
@@ -383,15 +395,16 @@ class TestAgentHandler:
|
||||
patch("src.agent.app.get_all_tools", return_value=[]), \
|
||||
patch("src.agent.app.save_conversation", AsyncMock()):
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
assert len(results) == 1
|
||||
data = json.loads(results[0])
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 1
|
||||
data = json.loads(filtered[0])
|
||||
assert data["metadata"]["code"] == "LLM_MALFORMED_OUTPUT"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_llm_error_saves_conversation(self, mock_request):
|
||||
from src.agent.app import agent_handler
|
||||
|
||||
def mock_astream(*args, **kwargs):
|
||||
def mock_astream(*_args, **_kwargs):
|
||||
raise RuntimeError("API connection error")
|
||||
|
||||
agent = MagicMock()
|
||||
@@ -401,17 +414,75 @@ class TestAgentHandler:
|
||||
with patch("src.agent.app.create_agent", return_value=agent), \
|
||||
patch("src.agent.app.get_all_tools", return_value=[]), \
|
||||
patch("src.agent.app.save_conversation", save_mock):
|
||||
# Handler catches RuntimeError, yields an error chunk, and saves conversation
|
||||
# Handler catches RuntimeError, yields a pipeline result first, then error chunk
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
assert len(results) == 1
|
||||
data = json.loads(results[0])
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 1
|
||||
data = json.loads(filtered[0])
|
||||
assert data["metadata"]["code"] == "PROCESSING_ERROR"
|
||||
save_mock.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_error_yields_code(self, mock_request):
|
||||
"""RateLimitError (429) yields LLM_RATE_LIMITED, not guardrails card."""
|
||||
from openai import RateLimitError
|
||||
|
||||
from src.agent.app import agent_handler
|
||||
|
||||
def mock_astream(*_args, **_kwargs):
|
||||
raise RateLimitError(
|
||||
"429 quota exceeded",
|
||||
response=MagicMock(status_code=429),
|
||||
body={"error": {"message": "quota exceeded"}},
|
||||
)
|
||||
|
||||
agent = MagicMock()
|
||||
agent.astream_events = mock_astream
|
||||
|
||||
save_mock = AsyncMock()
|
||||
with patch("src.agent.app.create_agent", return_value=agent), \
|
||||
patch("src.agent.app.get_all_tools", return_value=[]), \
|
||||
patch("src.agent.app.save_conversation", save_mock):
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 1
|
||||
data = json.loads(filtered[0])
|
||||
assert data["metadata"]["code"] == "LLM_RATE_LIMITED"
|
||||
assert "Превышен лимит" in data["content"]
|
||||
# save_conversation should be called even on error
|
||||
save_mock.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_exception_with_llm_keyword_yields_error_not_guardrails(self, mock_request):
|
||||
"""Generic exception containing 'quota' keyword yields error, not confirmation."""
|
||||
from src.agent.app import agent_handler
|
||||
|
||||
def mock_astream(*_args, **_kwargs):
|
||||
raise RuntimeError("All codex accounts reached configured quota threshold")
|
||||
|
||||
agent = MagicMock()
|
||||
agent.astream_events = mock_astream
|
||||
# Even though aget_state may return a truthy 'next', the error
|
||||
# contains LLM keywords and should skip the HITL recovery path.
|
||||
agent.aget_state = AsyncMock(return_value=MagicMock(next=("tools",)))
|
||||
|
||||
save_mock = AsyncMock()
|
||||
with patch("src.agent.app.create_agent", return_value=agent), \
|
||||
patch("src.agent.app.get_all_tools", return_value=[]), \
|
||||
patch("src.agent.app.save_conversation", save_mock):
|
||||
results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 1
|
||||
data = json.loads(filtered[0])
|
||||
# Must be PROCESSING_ERROR (not confirm_required)
|
||||
assert data["metadata"]["code"] == "PROCESSING_ERROR"
|
||||
save_mock.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_jwt_passes_gracefully(self):
|
||||
from src.agent.app import agent_handler
|
||||
from jose import JWTError
|
||||
|
||||
from src.agent.app import agent_handler
|
||||
req = MagicMock()
|
||||
req.headers = {"authorization": "Bearer invalid.jwt"}
|
||||
mock_event_stream = [
|
||||
@@ -423,7 +494,8 @@ class TestAgentHandler:
|
||||
patch("src.agent.app.get_all_tools", return_value=[]), \
|
||||
patch("src.agent.app.save_conversation", AsyncMock()):
|
||||
results = [r async for r in agent_handler("hi", [], req, None, None)]
|
||||
assert len(results) == 1
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_valid_user_jwt_with_audience_is_kept_for_tool_auth(self, mock_request):
|
||||
@@ -442,7 +514,8 @@ class TestAgentHandler:
|
||||
patch("src.agent.app.save_conversation", AsyncMock()):
|
||||
results = [r async for r in agent_handler("hi", [], mock_request, None, None, None, token)]
|
||||
|
||||
assert len(results) == 1
|
||||
filtered = _skip_pipeline(results)
|
||||
assert len(filtered) == 1
|
||||
assert get_user_jwt() == token
|
||||
# #endregion test_agent_handler
|
||||
|
||||
@@ -559,11 +632,12 @@ class TestFileUploadParsing:
|
||||
])
|
||||
mock_create.return_value = agent
|
||||
results = [r async for r in agent_handler(message, [], mock_request, None, None)]
|
||||
assert len(results) >= 2
|
||||
# First result is file upload metadata, second is stream token
|
||||
# Now: file_uploaded, pipeline_result, stream_token (3 events)
|
||||
assert len(results) >= 3
|
||||
# First result is file upload metadata, third (after pipeline) is stream token
|
||||
file_data = json.loads(results[0])
|
||||
assert file_data["metadata"]["type"] == "file_uploaded"
|
||||
token_data = json.loads(results[1])
|
||||
token_data = json.loads(results[2])
|
||||
assert token_data["metadata"]["type"] == "stream_token"
|
||||
# #endregion test_file_upload_parsing
|
||||
|
||||
|
||||
Reference in New Issue
Block a user