test(agent-chat): audit guardrail and error handling
This commit is contained in:
@@ -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