test(agent-chat): audit guardrail and error handling

This commit is contained in:
2026-07-06 01:20:12 +03:00
parent 49e4ac0fe2
commit 082d6af3ab
23 changed files with 1541 additions and 354 deletions

View File

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