207 lines
8.7 KiB
Python
207 lines
8.7 KiB
Python
|
import asyncio
|
||
|
|
from contextlib import nullcontext
|
||
|
|
from types import SimpleNamespace
|
||
|
|
from unittest.mock import AsyncMock
|
||
|
|
from uuid import uuid4
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from app.agents import chat_agent as chat_agent_module
|
||
|
|
from app.api import widget_chat
|
||
|
|
from app.models.chat_history import ChatHistory
|
||
|
|
from app.models.permission import Permission
|
||
|
|
from app.models.session_to_agent import SessionStatus, SessionToAgent
|
||
|
|
from app.services import agent_rooms
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def conversation(db, test_agent, test_customer, test_organization):
|
||
|
|
chat = SessionToAgent(
|
||
|
|
session_id=uuid4(), agent_id=test_agent.id, customer_id=test_customer.id,
|
||
|
|
organization_id=test_organization.id, status=SessionStatus.OPEN,
|
||
|
|
)
|
||
|
|
db.add(chat)
|
||
|
|
db.commit()
|
||
|
|
return chat
|
||
|
|
|
||
|
|
|
||
|
|
def grant(db, role, name):
|
||
|
|
permission = Permission(name=name)
|
||
|
|
db.add(permission)
|
||
|
|
role.permissions = [*role.permissions, permission]
|
||
|
|
db.commit()
|
||
|
|
|
||
|
|
|
||
|
|
def socket_for(user):
|
||
|
|
return {'user_id': str(user.id), 'organization_id': str(user.organization_id)}
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_watch_room_enforces_org_scope_and_current_permissions(db, test_user, test_role, conversation):
|
||
|
|
room = f'chat_watch_{conversation.session_id}'
|
||
|
|
socket = socket_for(test_user)
|
||
|
|
assert not await agent_rooms.may_watch_chat(db, socket, room)
|
||
|
|
grant(db, test_role, 'view_unassigned_chats')
|
||
|
|
assert await agent_rooms.may_watch_chat(db, socket, room)
|
||
|
|
assert not await agent_rooms.may_watch_chat(db, {**socket, 'organization_id': str(uuid4())}, room)
|
||
|
|
assert not await agent_rooms.may_watch_chat(db, socket, 'chat_watch_invalid')
|
||
|
|
|
||
|
|
# A viewer of the AI queue loses access after another person takes it.
|
||
|
|
conversation.user_id = uuid4()
|
||
|
|
db.commit()
|
||
|
|
assert not await agent_rooms.may_watch_chat(db, socket, room)
|
||
|
|
test_user.is_active = False
|
||
|
|
db.commit()
|
||
|
|
assert not await agent_rooms.may_watch_chat(db, socket, room)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_assigned_reader_can_watch_own_but_not_another_users_chat(db, test_user, test_role, conversation):
|
||
|
|
grant(db, test_role, 'view_assigned_chats')
|
||
|
|
room = f'chat_watch_{conversation.session_id}'
|
||
|
|
conversation.user_id = test_user.id
|
||
|
|
db.commit()
|
||
|
|
assert await agent_rooms.may_watch_chat(db, socket_for(test_user), room)
|
||
|
|
conversation.user_id = uuid4()
|
||
|
|
db.commit()
|
||
|
|
assert not await agent_rooms.may_watch_chat(db, socket_for(test_user), room)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_socket_join_uses_the_same_scope_check(db, test_user, test_role, conversation, monkeypatch):
|
||
|
|
grant(db, test_role, 'view_unassigned_chats')
|
||
|
|
sio = SimpleNamespace(
|
||
|
|
get_session=AsyncMock(return_value=socket_for(test_user)),
|
||
|
|
enter_room=AsyncMock(), emit=AsyncMock(),
|
||
|
|
)
|
||
|
|
monkeypatch.setattr(widget_chat, 'sio', sio)
|
||
|
|
monkeypatch.setattr(widget_chat, 'get_db', lambda: iter([db]))
|
||
|
|
room = f'chat_watch_{conversation.session_id}'
|
||
|
|
await widget_chat.handle_join_room('sid', {'session_id': room})
|
||
|
|
sio.enter_room.assert_awaited_once_with('sid', room, namespace='/agent')
|
||
|
|
conversation.user_id = uuid4()
|
||
|
|
db.commit()
|
||
|
|
await widget_chat.handle_join_room('sid', {'session_id': room})
|
||
|
|
assert sio.enter_room.await_count == 1
|
||
|
|
assert sio.emit.await_args.args[0] == 'error'
|
||
|
|
|
||
|
|
|
||
|
|
def test_organization_room_is_limited_to_active_chat_viewers(db, test_user, test_role):
|
||
|
|
room = f'org_chat_updates_{test_user.organization_id}'
|
||
|
|
assert not agent_rooms.may_join_chat_updates(db, socket_for(test_user), room)
|
||
|
|
grant(db, test_role, 'view_assigned_chats')
|
||
|
|
assert agent_rooms.may_join_chat_updates(db, socket_for(test_user), room)
|
||
|
|
assert not agent_rooms.may_join_chat_updates(db, socket_for(test_user), f'org_chat_updates_{uuid4()}')
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_notification_contains_no_message_or_session_data_and_broker_failure_does_not_break_chat(monkeypatch):
|
||
|
|
emit = AsyncMock(side_effect=[RuntimeError('broker unavailable'), None])
|
||
|
|
monkeypatch.setattr(agent_rooms.sio, 'emit', emit)
|
||
|
|
monkeypatch.setattr(agent_rooms, '_INBOX_COALESCE_SECONDS', 0.001)
|
||
|
|
await agent_rooms.notify_conversation_changed('org-1', 'session-1')
|
||
|
|
await agent_rooms._pending_inbox_updates['org-1'].task
|
||
|
|
assert emit.await_count == 2
|
||
|
|
assert emit.await_args_list[0].args == ('chat_changed', {})
|
||
|
|
assert emit.await_args_list[0].kwargs == {'room': 'chat_watch_session-1', 'namespace': '/agent'}
|
||
|
|
assert emit.await_args_list[1].args == ('inbox_changed', {})
|
||
|
|
assert emit.await_args_list[1].kwargs == {'room': 'org_chat_updates_org-1', 'namespace': '/agent'}
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_org_broadcast_coalesces_a_burst_and_keeps_changes_during_publish(monkeypatch):
|
||
|
|
started, release = asyncio.Event(), asyncio.Event()
|
||
|
|
calls = []
|
||
|
|
|
||
|
|
async def emit(event, data, **kwargs):
|
||
|
|
calls.append(event)
|
||
|
|
if event == 'inbox_changed' and calls.count('inbox_changed') == 1:
|
||
|
|
started.set()
|
||
|
|
await release.wait()
|
||
|
|
|
||
|
|
monkeypatch.setattr(agent_rooms.sio, 'emit', emit)
|
||
|
|
monkeypatch.setattr(agent_rooms, '_INBOX_COALESCE_SECONDS', 0.001)
|
||
|
|
await asyncio.gather(*(agent_rooms.notify_conversation_changed('org-1', 'session-1') for _ in range(1000)))
|
||
|
|
await started.wait()
|
||
|
|
assert calls.count('inbox_changed') == 1
|
||
|
|
await agent_rooms.notify_conversation_changed('org-1', 'session-1')
|
||
|
|
task = agent_rooms._pending_inbox_updates['org-1'].task
|
||
|
|
release.set()
|
||
|
|
await task
|
||
|
|
assert calls.count('chat_changed') == 1001
|
||
|
|
assert calls.count('inbox_changed') == 2
|
||
|
|
assert 'org-1' not in agent_rooms._pending_inbox_updates
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_customer_and_bot_are_notified_after_each_commit_before_any_model_call(
|
||
|
|
db, test_agent, test_customer, test_organization, conversation, monkeypatch
|
||
|
|
):
|
||
|
|
agent = object.__new__(chat_agent_module.ChatAgent)
|
||
|
|
agent.agent = SimpleNamespace(session_id=None)
|
||
|
|
agent._guardrail_ctx = None
|
||
|
|
agent._guardrail_surface = None
|
||
|
|
verdict = SimpleNamespace(block=True, reply='synthetic blocked reply', as_attributes=lambda: {})
|
||
|
|
monkeypatch.setattr(chat_agent_module, 'check_inbound', lambda *args, **kwargs: verdict)
|
||
|
|
monkeypatch.setattr(chat_agent_module, 'SessionLocal', lambda: nullcontext(db))
|
||
|
|
seen = []
|
||
|
|
|
||
|
|
async def on_change(org_id, session_id):
|
||
|
|
seen.append([row.message_type for row in db.query(ChatHistory).order_by(ChatHistory.id).all()])
|
||
|
|
|
||
|
|
monkeypatch.setattr(agent_rooms, 'notify_conversation_changed', on_change)
|
||
|
|
response = await agent.get_response(
|
||
|
|
message='synthetic visitor', session_id=str(conversation.session_id),
|
||
|
|
org_id=str(test_organization.id), agent_id=str(test_agent.id), customer_id=str(test_customer.id),
|
||
|
|
)
|
||
|
|
assert response.message == 'synthetic blocked reply'
|
||
|
|
assert seen == [['user'], ['user', 'bot']]
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_customer_is_visible_while_normal_ai_response_is_still_running(
|
||
|
|
db, test_agent, test_customer, test_organization, conversation, monkeypatch
|
||
|
|
):
|
||
|
|
agent = object.__new__(chat_agent_module.ChatAgent)
|
||
|
|
agent.agent = SimpleNamespace(session_id=None)
|
||
|
|
agent._guardrail_ctx = None
|
||
|
|
agent._guardrail_surface = None
|
||
|
|
agent.knowledge_tool = None
|
||
|
|
agent.agent_data = None
|
||
|
|
agent.transfer_to_human = False
|
||
|
|
agent._use_groq_json_tool = False
|
||
|
|
agent._groq_json_capture = {}
|
||
|
|
agent._response_model = chat_agent_module.ChatResponse
|
||
|
|
verdict = SimpleNamespace(block=False, as_attributes=lambda: {})
|
||
|
|
monkeypatch.setattr(chat_agent_module, 'check_inbound', lambda *args, **kwargs: verdict)
|
||
|
|
monkeypatch.setattr(chat_agent_module, 'check_output', lambda message, **kwargs: (message, None))
|
||
|
|
monkeypatch.setattr(chat_agent_module, 'SessionLocal', lambda: nullcontext(db))
|
||
|
|
response = chat_agent_module.ChatResponse(
|
||
|
|
message='synthetic answer', transfer_to_human=False, end_chat=False,
|
||
|
|
request_rating=False, create_ticket=False,
|
||
|
|
)
|
||
|
|
monkeypatch.setattr(chat_agent_module, 'parse_response_content', lambda *args: response)
|
||
|
|
started, release = asyncio.Event(), asyncio.Event()
|
||
|
|
|
||
|
|
async def run_model(*args):
|
||
|
|
started.set()
|
||
|
|
await release.wait()
|
||
|
|
return object()
|
||
|
|
|
||
|
|
agent._arun = run_model
|
||
|
|
seen = []
|
||
|
|
|
||
|
|
async def on_change(org_id, session_id):
|
||
|
|
seen.append([row.message_type for row in db.query(ChatHistory).order_by(ChatHistory.id).all()])
|
||
|
|
|
||
|
|
monkeypatch.setattr(agent_rooms, 'notify_conversation_changed', on_change)
|
||
|
|
task = asyncio.create_task(agent.get_response(
|
||
|
|
message='synthetic question', session_id=str(conversation.session_id),
|
||
|
|
org_id=str(test_organization.id), agent_id=str(test_agent.id), customer_id=str(test_customer.id),
|
||
|
|
))
|
||
|
|
await started.wait()
|
||
|
|
assert seen == [['user']]
|
||
|
|
release.set()
|
||
|
|
result = await task
|
||
|
|
assert result.message == 'synthetic answer'
|
||
|
|
assert seen == [['user'], ['user', 'bot']]
|