mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-05 05:16:03 +00:00
feat(agent-runner): persist created artifacts
This commit is contained in:
huanghuoguoguo
860
tests/unit_tests/agent/test_orchestrator_artifact.py
Normal file
860
tests/unit_tests/agent/test_orchestrator_artifact.py
Normal file
@@ -0,0 +1,860 @@
|
||||
"""Tests for artifact.created handling in orchestrator."""
|
||||
import pytest
|
||||
import base64
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
import uuid
|
||||
|
||||
from langbot.pkg.agent.runner.orchestrator import (
|
||||
AgentRunOrchestrator,
|
||||
MAX_ARTIFACT_INLINE_BYTES,
|
||||
)
|
||||
from langbot.pkg.agent.runner.descriptor import AgentRunnerDescriptor
|
||||
from langbot.pkg.agent.runner.host_models import AgentEventEnvelope, AgentBinding
|
||||
from langbot.pkg.agent.runner.errors import RunnerProtocolError
|
||||
from langbot_plugin.api.entities.builtin.agent_runner.input import AgentInput
|
||||
from langbot_plugin.api.entities.builtin.agent_runner.event import ActorContext
|
||||
from langbot.pkg.core import app
|
||||
|
||||
|
||||
class TestArtifactCreatedValidation:
|
||||
"""Test artifact.created validation and protocol errors."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
"""Create mock application."""
|
||||
ap = MagicMock(spec=app.Application)
|
||||
ap.logger = MagicMock()
|
||||
ap.plugin_connector = MagicMock()
|
||||
ap.plugin_connector.is_enable_plugin = True
|
||||
ap.persistence_mgr = MagicMock()
|
||||
ap.persistence_mgr.get_db_engine = MagicMock()
|
||||
return ap
|
||||
|
||||
@pytest.fixture
|
||||
def mock_registry(self):
|
||||
"""Create mock registry."""
|
||||
registry = MagicMock()
|
||||
registry.get = AsyncMock()
|
||||
return registry
|
||||
|
||||
@pytest.fixture
|
||||
def mock_event(self):
|
||||
"""Create mock event envelope."""
|
||||
event = MagicMock(spec=AgentEventEnvelope)
|
||||
event.event_id = str(uuid.uuid4())
|
||||
event.event_type = 'message.received'
|
||||
event.source = 'test'
|
||||
event.bot_id = str(uuid.uuid4())
|
||||
event.workspace_id = str(uuid.uuid4())
|
||||
event.conversation_id = str(uuid.uuid4())
|
||||
event.thread_id = None
|
||||
event.event_time = 1700000000
|
||||
event.actor = MagicMock(spec=ActorContext)
|
||||
event.actor.actor_type = 'user'
|
||||
event.actor.actor_id = 'user-123'
|
||||
event.actor.actor_name = 'Test User'
|
||||
event.subject = None
|
||||
event.input = MagicMock(spec=AgentInput)
|
||||
event.input.text = 'Hello'
|
||||
event.input.contents = []
|
||||
event.input.attachments = []
|
||||
return event
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_id_mismatch_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that run_id mismatch raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
wrong_run_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': wrong_run_id,
|
||||
'data': {
|
||||
'artifact_type': 'image',
|
||||
},
|
||||
}
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
)
|
||||
|
||||
assert 'run_id mismatch' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_artifact_type_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that missing artifact_type raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_id': str(uuid.uuid4()),
|
||||
# missing artifact_type
|
||||
},
|
||||
}
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
)
|
||||
|
||||
assert 'missing required field' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_base64_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that invalid base64 raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_type': 'image',
|
||||
'content_base64': '!!!invalid-base64!!!',
|
||||
},
|
||||
}
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
)
|
||||
|
||||
assert 'invalid base64' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_content_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that content exceeding limit raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
# Create content larger than limit
|
||||
oversized_content = b'x' * (MAX_ARTIFACT_INLINE_BYTES + 1)
|
||||
content_base64 = base64.b64encode(oversized_content).decode('utf-8')
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_type': 'image',
|
||||
'content_base64': content_base64,
|
||||
},
|
||||
}
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
)
|
||||
|
||||
assert 'exceeds limit' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_artifact_store_failure_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that ArtifactStore failure raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_type': 'image',
|
||||
},
|
||||
}
|
||||
|
||||
with patch('langbot.pkg.agent.runner.artifact_store.ArtifactStore') as MockArtifactStore:
|
||||
mock_artifact_store = MagicMock()
|
||||
mock_artifact_store.register_artifact = AsyncMock(
|
||||
side_effect=Exception('DB connection failed')
|
||||
)
|
||||
MockArtifactStore.return_value = mock_artifact_store
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
)
|
||||
|
||||
assert 'failed to register artifact' in str(exc_info.value)
|
||||
|
||||
|
||||
class TestArtifactCreatedSuccess:
|
||||
"""Test successful artifact.created handling."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
"""Create mock application."""
|
||||
ap = MagicMock(spec=app.Application)
|
||||
ap.logger = MagicMock()
|
||||
ap.plugin_connector = MagicMock()
|
||||
ap.plugin_connector.is_enable_plugin = True
|
||||
ap.persistence_mgr = MagicMock()
|
||||
ap.persistence_mgr.get_db_engine = MagicMock()
|
||||
return ap
|
||||
|
||||
@pytest.fixture
|
||||
def mock_registry(self):
|
||||
"""Create mock registry."""
|
||||
registry = MagicMock()
|
||||
registry.get = AsyncMock()
|
||||
return registry
|
||||
|
||||
@pytest.fixture
|
||||
def mock_event(self):
|
||||
"""Create mock event envelope."""
|
||||
event = MagicMock(spec=AgentEventEnvelope)
|
||||
event.event_id = str(uuid.uuid4())
|
||||
event.event_type = 'message.received'
|
||||
event.source = 'test'
|
||||
event.bot_id = str(uuid.uuid4())
|
||||
event.workspace_id = str(uuid.uuid4())
|
||||
event.conversation_id = str(uuid.uuid4())
|
||||
event.thread_id = None
|
||||
event.event_time = 1700000000
|
||||
event.actor = MagicMock(spec=ActorContext)
|
||||
event.actor.actor_type = 'user'
|
||||
event.actor.actor_id = 'user-123'
|
||||
event.actor.actor_name = 'Test User'
|
||||
event.subject = None
|
||||
return event
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_artifact_created_registers_artifact(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that artifact.created registers artifact via ArtifactStore."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
runner_id = 'test-runner'
|
||||
|
||||
# Create artifact.created result
|
||||
content = b'test artifact content'
|
||||
content_base64 = base64.b64encode(content).decode('utf-8')
|
||||
artifact_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_id': artifact_id,
|
||||
'artifact_type': 'image',
|
||||
'mime_type': 'image/png',
|
||||
'name': 'test.png',
|
||||
'size_bytes': len(content),
|
||||
'content_base64': content_base64,
|
||||
},
|
||||
}
|
||||
|
||||
with patch('langbot.pkg.agent.runner.artifact_store.ArtifactStore') as MockArtifactStore:
|
||||
with patch('langbot.pkg.agent.runner.event_log_store.EventLogStore') as MockEventLogStore:
|
||||
mock_artifact_store = MagicMock()
|
||||
mock_artifact_store.register_artifact = AsyncMock(return_value=artifact_id)
|
||||
MockArtifactStore.return_value = mock_artifact_store
|
||||
|
||||
mock_event_log_store = MagicMock()
|
||||
mock_event_log_store.append_event = AsyncMock()
|
||||
MockEventLogStore.return_value = mock_event_log_store
|
||||
|
||||
# Call _handle_artifact_created
|
||||
result = await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id=runner_id,
|
||||
)
|
||||
|
||||
# Verify artifact was registered
|
||||
mock_artifact_store.register_artifact.assert_called_once()
|
||||
call_kwargs = mock_artifact_store.register_artifact.call_args.kwargs
|
||||
assert call_kwargs['artifact_id'] == artifact_id
|
||||
assert call_kwargs['artifact_type'] == 'image'
|
||||
assert call_kwargs['mime_type'] == 'image/png'
|
||||
assert call_kwargs['name'] == 'test.png'
|
||||
assert call_kwargs['content'] == content
|
||||
assert call_kwargs['conversation_id'] == mock_event.conversation_id
|
||||
assert call_kwargs['run_id'] == run_id
|
||||
assert call_kwargs['runner_id'] == runner_id
|
||||
|
||||
# Verify EventLog was written
|
||||
mock_event_log_store.append_event.assert_called_once()
|
||||
event_kwargs = mock_event_log_store.append_event.call_args.kwargs
|
||||
assert event_kwargs['event_type'] == 'artifact.created'
|
||||
assert event_kwargs['run_id'] == run_id
|
||||
|
||||
# Verify artifact ref returned
|
||||
assert result is not None
|
||||
assert result['artifact_id'] == artifact_id
|
||||
assert result['artifact_type'] == 'image'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_artifact_created_metadata_only(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test artifact.created without content (metadata-only)."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
artifact_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_id': artifact_id,
|
||||
'artifact_type': 'file',
|
||||
'mime_type': 'application/pdf',
|
||||
'name': 'document.pdf',
|
||||
'size_bytes': 1024,
|
||||
'sha256': 'abc123',
|
||||
'metadata': {'source': 'external'},
|
||||
},
|
||||
}
|
||||
|
||||
with patch('langbot.pkg.agent.runner.artifact_store.ArtifactStore') as MockArtifactStore:
|
||||
with patch('langbot.pkg.agent.runner.event_log_store.EventLogStore') as MockEventLogStore:
|
||||
mock_artifact_store = MagicMock()
|
||||
mock_artifact_store.register_artifact = AsyncMock(return_value=artifact_id)
|
||||
MockArtifactStore.return_value = mock_artifact_store
|
||||
|
||||
mock_event_log_store = MagicMock()
|
||||
mock_event_log_store.append_event = AsyncMock()
|
||||
MockEventLogStore.return_value = mock_event_log_store
|
||||
|
||||
result = await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
)
|
||||
|
||||
# Verify artifact was registered without content
|
||||
call_kwargs = mock_artifact_store.register_artifact.call_args.kwargs
|
||||
assert call_kwargs['content'] is None
|
||||
assert call_kwargs['sha256'] == 'abc123'
|
||||
assert call_kwargs['metadata'] == {'source': 'external'}
|
||||
|
||||
assert result is not None
|
||||
assert result['artifact_id'] == artifact_id
|
||||
|
||||
|
||||
class TestArtifactRefsLifecycle:
|
||||
"""Test artifact refs lifecycle in event-first flow."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
"""Create mock application."""
|
||||
ap = MagicMock(spec=app.Application)
|
||||
ap.logger = MagicMock()
|
||||
ap.plugin_connector = MagicMock()
|
||||
ap.plugin_connector.is_enable_plugin = True
|
||||
ap.persistence_mgr = MagicMock()
|
||||
ap.persistence_mgr.get_db_engine = MagicMock()
|
||||
return ap
|
||||
|
||||
@pytest.fixture
|
||||
def mock_registry(self):
|
||||
"""Create mock registry."""
|
||||
registry = MagicMock()
|
||||
registry.get = AsyncMock()
|
||||
return registry
|
||||
|
||||
def test_merge_artifact_refs_deduplicates(
|
||||
self, mock_app, mock_registry
|
||||
):
|
||||
"""Test that _merge_artifact_refs deduplicates by artifact_id."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
|
||||
pending_refs = [
|
||||
{'artifact_id': 'artifact-1', 'artifact_type': 'image'},
|
||||
{'artifact_id': 'artifact-2', 'artifact_type': 'file'},
|
||||
]
|
||||
|
||||
result_dict = {
|
||||
'type': 'message.completed',
|
||||
'data': {
|
||||
'message': {
|
||||
'content': 'Hello',
|
||||
'artifact_refs': [
|
||||
{'artifact_id': 'artifact-2', 'artifact_type': 'file'}, # duplicate
|
||||
{'artifact_id': 'artifact-3', 'artifact_type': 'voice'},
|
||||
],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
merged = orchestrator._merge_artifact_refs(pending_refs, result_dict)
|
||||
|
||||
# Should have 3 unique artifacts
|
||||
assert len(merged) == 3
|
||||
artifact_ids = {ref['artifact_id'] for ref in merged}
|
||||
assert artifact_ids == {'artifact-1', 'artifact-2', 'artifact-3'}
|
||||
|
||||
def test_merge_artifact_refs_empty_pending(
|
||||
self, mock_app, mock_registry
|
||||
):
|
||||
"""Test merge with empty pending refs."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
|
||||
pending_refs = []
|
||||
|
||||
result_dict = {
|
||||
'type': 'message.completed',
|
||||
'data': {
|
||||
'message': {
|
||||
'content': 'Hello',
|
||||
'artifact_refs': [
|
||||
{'artifact_id': 'artifact-1', 'artifact_type': 'image'},
|
||||
],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
merged = orchestrator._merge_artifact_refs(pending_refs, result_dict)
|
||||
|
||||
assert len(merged) == 1
|
||||
assert merged[0]['artifact_id'] == 'artifact-1'
|
||||
|
||||
def test_merge_artifact_refs_empty_message_refs(
|
||||
self, mock_app, mock_registry
|
||||
):
|
||||
"""Test merge with no message artifact_refs."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
|
||||
pending_refs = [
|
||||
{'artifact_id': 'artifact-1', 'artifact_type': 'image'},
|
||||
]
|
||||
|
||||
result_dict = {
|
||||
'type': 'message.completed',
|
||||
'data': {
|
||||
'message': {
|
||||
'content': 'Hello',
|
||||
# no artifact_refs
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
merged = orchestrator._merge_artifact_refs(pending_refs, result_dict)
|
||||
|
||||
assert len(merged) == 1
|
||||
assert merged[0]['artifact_id'] == 'artifact-1'
|
||||
|
||||
|
||||
class TestArtifactCreatedQueryFlow:
|
||||
"""Test artifact.created handling in legacy Query-based flow."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
"""Create mock application."""
|
||||
ap = MagicMock(spec=app.Application)
|
||||
ap.logger = MagicMock()
|
||||
ap.plugin_connector = MagicMock()
|
||||
ap.plugin_connector.is_enable_plugin = True
|
||||
ap.persistence_mgr = MagicMock()
|
||||
ap.persistence_mgr.get_db_engine = MagicMock()
|
||||
return ap
|
||||
|
||||
@pytest.fixture
|
||||
def mock_registry(self):
|
||||
"""Create mock registry."""
|
||||
registry = MagicMock()
|
||||
registry.get = AsyncMock()
|
||||
return registry
|
||||
|
||||
@pytest.fixture
|
||||
def mock_query(self):
|
||||
"""Create mock Query."""
|
||||
from langbot_plugin.api.entities.builtin.pipeline import query as pipeline_query
|
||||
query = MagicMock(spec=pipeline_query.Query)
|
||||
query.query_id = str(uuid.uuid4())
|
||||
query.pipeline_config = {'runner': {'id': 'test-runner'}}
|
||||
query.variables = {}
|
||||
return query
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_flow_run_id_mismatch_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_query
|
||||
):
|
||||
"""Test run_id mismatch in Query flow raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
wrong_run_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': wrong_run_id,
|
||||
'data': {'artifact_type': 'image'},
|
||||
}
|
||||
|
||||
mock_descriptor = MagicMock()
|
||||
mock_descriptor.id = 'test-runner'
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created_query(
|
||||
result_dict=result_dict,
|
||||
query=mock_query,
|
||||
descriptor=mock_descriptor,
|
||||
run_id=run_id,
|
||||
conversation_id=str(uuid.uuid4()),
|
||||
)
|
||||
|
||||
assert 'run_id mismatch' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_flow_invalid_base64_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_query
|
||||
):
|
||||
"""Test invalid base64 in Query flow raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_type': 'image',
|
||||
'content_base64': '!!!invalid!!!',
|
||||
},
|
||||
}
|
||||
|
||||
mock_descriptor = MagicMock()
|
||||
mock_descriptor.id = 'test-runner'
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created_query(
|
||||
result_dict=result_dict,
|
||||
query=mock_query,
|
||||
descriptor=mock_descriptor,
|
||||
run_id=run_id,
|
||||
conversation_id=str(uuid.uuid4()),
|
||||
)
|
||||
|
||||
assert 'invalid base64' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_flow_missing_artifact_type_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_query
|
||||
):
|
||||
"""Test missing artifact_type in Query flow raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {}, # missing artifact_type
|
||||
}
|
||||
|
||||
mock_descriptor = MagicMock()
|
||||
mock_descriptor.id = 'test-runner'
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created_query(
|
||||
result_dict=result_dict,
|
||||
query=mock_query,
|
||||
descriptor=mock_descriptor,
|
||||
run_id=run_id,
|
||||
conversation_id=str(uuid.uuid4()),
|
||||
)
|
||||
|
||||
assert 'missing required field' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_flow_oversized_content_raises_protocol_error(
|
||||
self, mock_app, mock_registry, mock_query
|
||||
):
|
||||
"""Test oversized content in Query flow raises RunnerProtocolError."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
oversized_content = b'x' * (MAX_ARTIFACT_INLINE_BYTES + 1)
|
||||
content_base64 = base64.b64encode(oversized_content).decode('utf-8')
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_type': 'image',
|
||||
'content_base64': content_base64,
|
||||
},
|
||||
}
|
||||
|
||||
mock_descriptor = MagicMock()
|
||||
mock_descriptor.id = 'test-runner'
|
||||
|
||||
with pytest.raises(RunnerProtocolError) as exc_info:
|
||||
await orchestrator._handle_artifact_created_query(
|
||||
result_dict=result_dict,
|
||||
query=mock_query,
|
||||
descriptor=mock_descriptor,
|
||||
run_id=run_id,
|
||||
conversation_id=str(uuid.uuid4()),
|
||||
)
|
||||
|
||||
assert 'exceeds limit' in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_flow_register_success(
|
||||
self, mock_app, mock_registry, mock_query
|
||||
):
|
||||
"""Test successful artifact registration in Query flow."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
conversation_id = str(uuid.uuid4())
|
||||
artifact_id = str(uuid.uuid4())
|
||||
|
||||
content = b'test content'
|
||||
content_base64 = base64.b64encode(content).decode('utf-8')
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_id': artifact_id,
|
||||
'artifact_type': 'voice',
|
||||
'mime_type': 'audio/mp3',
|
||||
'content_base64': content_base64,
|
||||
},
|
||||
}
|
||||
|
||||
mock_descriptor = MagicMock()
|
||||
mock_descriptor.id = 'test-runner'
|
||||
|
||||
with patch('langbot.pkg.agent.runner.artifact_store.ArtifactStore') as MockArtifactStore:
|
||||
mock_artifact_store = MagicMock()
|
||||
mock_artifact_store.register_artifact = AsyncMock(return_value=artifact_id)
|
||||
MockArtifactStore.return_value = mock_artifact_store
|
||||
|
||||
await orchestrator._handle_artifact_created_query(
|
||||
result_dict=result_dict,
|
||||
query=mock_query,
|
||||
descriptor=mock_descriptor,
|
||||
run_id=run_id,
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
|
||||
# Verify artifact was registered
|
||||
mock_artifact_store.register_artifact.assert_called_once()
|
||||
call_kwargs = mock_artifact_store.register_artifact.call_args.kwargs
|
||||
assert call_kwargs['artifact_id'] == artifact_id
|
||||
assert call_kwargs['artifact_type'] == 'voice'
|
||||
assert call_kwargs['content'] == content
|
||||
assert call_kwargs['conversation_id'] == conversation_id
|
||||
|
||||
|
||||
class TestResultNormalizerArtifactCreated:
|
||||
"""Test ResultNormalizer handling of artifact.created."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
"""Create mock application."""
|
||||
ap = MagicMock(spec=app.Application)
|
||||
ap.logger = MagicMock()
|
||||
return ap
|
||||
|
||||
@pytest.fixture
|
||||
def mock_descriptor(self):
|
||||
"""Create mock descriptor."""
|
||||
descriptor = MagicMock()
|
||||
descriptor.id = 'test-runner'
|
||||
return descriptor
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normalize_artifact_created_returns_none(
|
||||
self, mock_app, mock_descriptor
|
||||
):
|
||||
"""Test that artifact.created is consumed (returns None)."""
|
||||
from langbot.pkg.agent.runner.result_normalizer import AgentResultNormalizer
|
||||
|
||||
normalizer = AgentResultNormalizer(mock_app)
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': 'test-run-id',
|
||||
'data': {
|
||||
'artifact_id': 'artifact-123',
|
||||
'artifact_type': 'image',
|
||||
},
|
||||
}
|
||||
|
||||
result = await normalizer.normalize(result_dict, mock_descriptor)
|
||||
|
||||
# Should return None (consumed)
|
||||
assert result is None
|
||||
|
||||
# Debug log should be written
|
||||
mock_app.logger.debug.assert_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normalize_unknown_type_warning(
|
||||
self, mock_app, mock_descriptor
|
||||
):
|
||||
"""Test that unknown result types still produce warnings."""
|
||||
from langbot.pkg.agent.runner.result_normalizer import AgentResultNormalizer
|
||||
|
||||
normalizer = AgentResultNormalizer(mock_app)
|
||||
|
||||
result_dict = {
|
||||
'type': 'unknown.type',
|
||||
'data': {},
|
||||
}
|
||||
|
||||
result = await normalizer.normalize(result_dict, mock_descriptor)
|
||||
|
||||
# Should return None
|
||||
assert result is None
|
||||
|
||||
# Warning should be logged
|
||||
mock_app.logger.warning.assert_called()
|
||||
|
||||
|
||||
class TestEventLogTranscriptIntegration:
|
||||
"""Test EventLog and Transcript integration with artifact.created."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
"""Create mock application."""
|
||||
ap = MagicMock(spec=app.Application)
|
||||
ap.logger = MagicMock()
|
||||
ap.plugin_connector = MagicMock()
|
||||
ap.plugin_connector.is_enable_plugin = True
|
||||
ap.persistence_mgr = MagicMock()
|
||||
ap.persistence_mgr.get_db_engine = MagicMock()
|
||||
return ap
|
||||
|
||||
@pytest.fixture
|
||||
def mock_registry(self):
|
||||
"""Create mock registry."""
|
||||
registry = MagicMock()
|
||||
registry.get = AsyncMock()
|
||||
return registry
|
||||
|
||||
@pytest.fixture
|
||||
def mock_event(self):
|
||||
"""Create mock event envelope."""
|
||||
event = MagicMock(spec=AgentEventEnvelope)
|
||||
event.event_id = str(uuid.uuid4())
|
||||
event.event_type = 'message.received'
|
||||
event.source = 'test'
|
||||
event.bot_id = str(uuid.uuid4())
|
||||
event.workspace_id = str(uuid.uuid4())
|
||||
event.conversation_id = str(uuid.uuid4())
|
||||
event.thread_id = None
|
||||
event.event_time = 1700000000
|
||||
event.actor = MagicMock(spec=ActorContext)
|
||||
event.actor.actor_type = 'user'
|
||||
event.actor.actor_id = 'user-123'
|
||||
event.actor.actor_name = 'Test User'
|
||||
event.subject = None
|
||||
return event
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_log_written_with_correct_event_type(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that EventLog is written with event_type='artifact.created'."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
artifact_id = str(uuid.uuid4())
|
||||
|
||||
result_dict = {
|
||||
'type': 'artifact.created',
|
||||
'run_id': run_id,
|
||||
'data': {
|
||||
'artifact_id': artifact_id,
|
||||
'artifact_type': 'image',
|
||||
},
|
||||
}
|
||||
|
||||
with patch('langbot.pkg.agent.runner.artifact_store.ArtifactStore') as MockArtifactStore:
|
||||
with patch('langbot.pkg.agent.runner.event_log_store.EventLogStore') as MockEventLogStore:
|
||||
mock_artifact_store = MagicMock()
|
||||
mock_artifact_store.register_artifact = AsyncMock(return_value=artifact_id)
|
||||
MockArtifactStore.return_value = mock_artifact_store
|
||||
|
||||
mock_event_log_store = MagicMock()
|
||||
mock_event_log_store.append_event = AsyncMock()
|
||||
MockEventLogStore.return_value = mock_event_log_store
|
||||
|
||||
await orchestrator._handle_artifact_created(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
)
|
||||
|
||||
# Verify EventLog.append_event was called with correct event_type
|
||||
mock_event_log_store.append_event.assert_called_once()
|
||||
call_kwargs = mock_event_log_store.append_event.call_args.kwargs
|
||||
assert call_kwargs['event_type'] == 'artifact.created'
|
||||
assert call_kwargs['source'] == 'runner'
|
||||
assert call_kwargs['conversation_id'] == mock_event.conversation_id
|
||||
assert call_kwargs['run_id'] == run_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assistant_transcript_receives_artifact_refs(
|
||||
self, mock_app, mock_registry, mock_event
|
||||
):
|
||||
"""Test that assistant transcript receives artifact refs from artifact.created."""
|
||||
orchestrator = AgentRunOrchestrator(mock_app, mock_registry)
|
||||
run_id = str(uuid.uuid4())
|
||||
artifact_id = str(uuid.uuid4())
|
||||
|
||||
# Create pending artifact refs
|
||||
pending_refs = [
|
||||
{'artifact_id': artifact_id, 'artifact_type': 'image', 'mime_type': 'image/png'},
|
||||
]
|
||||
|
||||
result_dict = {
|
||||
'type': 'message.completed',
|
||||
'data': {
|
||||
'message': {
|
||||
'content': 'Here is your image',
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
with patch('langbot.pkg.agent.runner.transcript_store.TranscriptStore') as MockTranscriptStore:
|
||||
mock_transcript_store = MagicMock()
|
||||
mock_transcript_store.append_transcript = AsyncMock()
|
||||
MockTranscriptStore.return_value = mock_transcript_store
|
||||
|
||||
await orchestrator._write_assistant_transcript(
|
||||
result_dict=result_dict,
|
||||
event=mock_event,
|
||||
run_id=run_id,
|
||||
runner_id='test-runner',
|
||||
artifact_refs=pending_refs,
|
||||
)
|
||||
|
||||
# Verify transcript was written with artifact_refs
|
||||
mock_transcript_store.append_transcript.assert_called_once()
|
||||
call_kwargs = mock_transcript_store.append_transcript.call_args.kwargs
|
||||
assert call_kwargs['artifact_refs'] == pending_refs
|
||||
Reference in New Issue
Block a user