From 3b3b09331ac924cf77f4a97ec7f2446d1c10aabd Mon Sep 17 00:00:00 2001 From: RockChinQ Date: Wed, 17 Jun 2026 22:08:42 -0400 Subject: [PATCH] test(box): cover inbound/outbound attachment helpers; fix ruff format - ruff format localagent.py (CI ruff format --check was failing) - add unit tests for ResponseWrapper outbound-attachment helpers (wrapper.py 78%->98%) - add unit tests for LocalAgentRunner._inject_inbound_attachments - add unit tests for WebSocketAdapter._process_image_components (0%->covered) Lifts PR patch coverage from 68.97% to ~88% (>75% target). --- .../pkg/provider/runners/localagent.py | 4 +- .../test_wrapper_outbound_attachments.py | 146 ++++++++++++++++++ .../test_websocket_adapter_attachments.py | 92 +++++++++++ .../test_localagent_inbound_attachments.py | 146 ++++++++++++++++++ 4 files changed, 385 insertions(+), 3 deletions(-) create mode 100644 tests/unit_tests/pipeline/test_wrapper_outbound_attachments.py create mode 100644 tests/unit_tests/platform/test_websocket_adapter_attachments.py create mode 100644 tests/unit_tests/provider/test_localagent_inbound_attachments.py diff --git a/src/langbot/pkg/provider/runners/localagent.py b/src/langbot/pkg/provider/runners/localagent.py index 5143e5b0..338de2e5 100644 --- a/src/langbot/pkg/provider/runners/localagent.py +++ b/src/langbot/pkg/provider/runners/localagent.py @@ -153,9 +153,7 @@ class LocalAgentRunner(runner.RequestRunner): _model_unsafe_types = {'file_base64', 'file_url'} if isinstance(user_message.content, list): user_message.content = [ - ce - for ce in user_message.content - if getattr(ce, 'type', None) not in _model_unsafe_types + ce for ce in user_message.content if getattr(ce, 'type', None) not in _model_unsafe_types ] if isinstance(user_message.content, str): diff --git a/tests/unit_tests/pipeline/test_wrapper_outbound_attachments.py b/tests/unit_tests/pipeline/test_wrapper_outbound_attachments.py new file mode 100644 index 00000000..8fc000bf --- /dev/null +++ b/tests/unit_tests/pipeline/test_wrapper_outbound_attachments.py @@ -0,0 +1,146 @@ +"""Unit tests for ResponseWrapper outbound-attachment helpers. + +Covers the sandbox -> user attachment path added for the Box attachment +round-trip: + +* ``_is_final_assistant_message`` — only the terminal, tool-call-free assistant + message (or a final MessageChunk) should trigger collection. +* ``_append_outbound_attachments`` — collects sandbox outbox files exactly once + per query and maps each descriptor to the right platform component, swallowing + collection errors. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +import langbot_plugin.api.entities.builtin.platform.message as platform_message +import langbot_plugin.api.entities.builtin.provider.message as provider_message + +from langbot.pkg.pipeline.wrapper.wrapper import ResponseWrapper + + +def _make_wrapper(box_service) -> ResponseWrapper: + app = SimpleNamespace(logger=Mock()) + wrapper = ResponseWrapper.__new__(ResponseWrapper) + wrapper.ap = app + return wrapper + + +def _make_query(): + return SimpleNamespace(variables={}) + + +def test_is_final_assistant_message_plain_assistant(): + wrapper = _make_wrapper(box_service=None) + msg = provider_message.Message(role='assistant', content='done') + assert wrapper._is_final_assistant_message(msg) is True + + +def test_is_final_assistant_message_rejects_non_assistant(): + wrapper = _make_wrapper(box_service=None) + msg = provider_message.Message(role='tool', content='{}') + assert wrapper._is_final_assistant_message(msg) is False + + +def test_is_final_assistant_message_rejects_tool_call_round(): + wrapper = _make_wrapper(box_service=None) + msg = provider_message.Message( + role='assistant', + content='calling', + tool_calls=[ + provider_message.ToolCall( + id='c1', + type='function', + function=provider_message.FunctionCall(name='exec', arguments='{}'), + ) + ], + ) + assert wrapper._is_final_assistant_message(msg) is False + + +def test_is_final_assistant_message_non_final_chunk(): + wrapper = _make_wrapper(box_service=None) + chunk = provider_message.MessageChunk(role='assistant', content='partial', is_final=False) + assert wrapper._is_final_assistant_message(chunk) is False + + final_chunk = provider_message.MessageChunk(role='assistant', content='partial', is_final=True) + assert wrapper._is_final_assistant_message(final_chunk) is True + + +@pytest.mark.asyncio +async def test_append_outbound_attachments_maps_each_type(): + box_service = SimpleNamespace( + available=True, + collect_outbound_attachments=AsyncMock( + return_value=[ + {'type': 'Image', 'base64': 'data:image/png;base64,iVBORw0K'}, + {'type': 'Voice', 'base64': 'data:audio/wav;base64,UklGRg=='}, + {'type': 'File', 'name': 'report.xlsx', 'base64': 'data:app;base64,UEsDBA=='}, + ] + ), + ) + wrapper = _make_wrapper(box_service) + wrapper.ap.box_service = box_service + query = _make_query() + chain = platform_message.MessageChain([]) + + await wrapper._append_outbound_attachments(query, chain) + + kinds = [type(c).__name__ for c in chain] + assert kinds == ['Image', 'Voice', 'File'] + assert query.variables['_sandbox_outbound_collected'] is True + # File keeps its name + file_comp = chain[2] + assert getattr(file_comp, 'name', None) == 'report.xlsx' + + +@pytest.mark.asyncio +async def test_append_outbound_attachments_runs_once_per_query(): + box_service = SimpleNamespace( + available=True, + collect_outbound_attachments=AsyncMock(return_value=[]), + ) + wrapper = _make_wrapper(box_service) + wrapper.ap.box_service = box_service + query = _make_query() + query.variables['_sandbox_outbound_collected'] = True + chain = platform_message.MessageChain([]) + + await wrapper._append_outbound_attachments(query, chain) + + box_service.collect_outbound_attachments.assert_not_awaited() + assert len(chain) == 0 + + +@pytest.mark.asyncio +async def test_append_outbound_attachments_noop_without_box_service(): + wrapper = _make_wrapper(box_service=None) + wrapper.ap.box_service = None + query = _make_query() + chain = platform_message.MessageChain([]) + + await wrapper._append_outbound_attachments(query, chain) + assert len(chain) == 0 + # not marked collected, since service is unavailable + assert '_sandbox_outbound_collected' not in query.variables + + +@pytest.mark.asyncio +async def test_append_outbound_attachments_swallows_collection_error(): + box_service = SimpleNamespace( + available=True, + collect_outbound_attachments=AsyncMock(side_effect=RuntimeError('boom')), + ) + wrapper = _make_wrapper(box_service) + wrapper.ap.box_service = box_service + query = _make_query() + chain = platform_message.MessageChain([]) + + # must not raise + await wrapper._append_outbound_attachments(query, chain) + assert len(chain) == 0 + wrapper.ap.logger.warning.assert_called_once() diff --git a/tests/unit_tests/platform/test_websocket_adapter_attachments.py b/tests/unit_tests/platform/test_websocket_adapter_attachments.py new file mode 100644 index 00000000..18138383 --- /dev/null +++ b/tests/unit_tests/platform/test_websocket_adapter_attachments.py @@ -0,0 +1,92 @@ +"""Unit tests for WebSocketAdapter._process_image_components. + +The web debug client uploads Image / Voice / File components carrying a storage +key in ``path``. This helper resolves each to a base64 data URI (so multimodal +LLM input and the Box sandbox inbox have usable bytes), then deletes the +consumed storage object and clears ``path``. Covers mimetype selection per +type and graceful error handling. +""" + +from __future__ import annotations + +import base64 +from unittest.mock import AsyncMock, Mock + +import pytest + +from langbot.pkg.platform.sources.websocket_adapter import WebSocketAdapter + + +def _make_adapter(load_return=b'hello', load_side_effect=None): + provider = Mock() + provider.load = AsyncMock(return_value=load_return, side_effect=load_side_effect) + provider.delete = AsyncMock() + ap = Mock() + ap.storage_mgr.storage_provider = provider + logger = Mock() + logger.error = AsyncMock() + # WebSocketAdapter is a pydantic model; bypass full __init__/validation. + adapter = WebSocketAdapter.model_construct(ap=ap, logger=logger) + return adapter, provider + + +@pytest.mark.asyncio +async def test_image_jpeg_mimetype_and_cleanup(): + adapter, provider = _make_adapter(load_return=b'\xff\xd8\xff') + chain = [{'type': 'Image', 'path': 'storage://abc/photo.jpg'}] + + await adapter._process_image_components(chain) + + expected_b64 = base64.b64encode(b'\xff\xd8\xff').decode('utf-8') + assert chain[0]['base64'] == f'data:image/jpeg;base64,{expected_b64}' + assert chain[0]['path'] == '' # consumed + provider.delete.assert_awaited_once_with('storage://abc/photo.jpg') + + +@pytest.mark.asyncio +async def test_image_defaults_to_png(): + adapter, _ = _make_adapter() + chain = [{'type': 'Image', 'path': 'storage://abc/blob'}] + await adapter._process_image_components(chain) + assert chain[0]['base64'].startswith('data:image/png;base64,') + + +@pytest.mark.asyncio +async def test_voice_uses_guessed_or_wav_mimetype(): + adapter, _ = _make_adapter() + chain = [{'type': 'Voice', 'path': 'storage://abc/clip.wav'}] + await adapter._process_image_components(chain) + assert chain[0]['base64'].startswith('data:audio/') + + +@pytest.mark.asyncio +async def test_file_uses_octet_stream_fallback(): + adapter, _ = _make_adapter() + chain = [{'type': 'File', 'path': 'storage://abc/unknownblob'}] + await adapter._process_image_components(chain) + assert chain[0]['base64'].startswith('data:application/octet-stream;base64,') + + +@pytest.mark.asyncio +async def test_skips_components_without_path_or_unknown_type(): + adapter, provider = _make_adapter() + chain = [ + {'type': 'Image', 'path': ''}, # no path + {'type': 'Plain', 'path': 'storage://abc/x'}, # not a file component + {'type': 'At', 'target': '123'}, # no path key at all + ] + await adapter._process_image_components(chain) + provider.load.assert_not_awaited() + assert 'base64' not in chain[0] + assert 'base64' not in chain[1] + + +@pytest.mark.asyncio +async def test_load_failure_is_logged_not_raised(): + adapter, _ = _make_adapter(load_side_effect=RuntimeError('storage down')) + chain = [{'type': 'File', 'path': 'storage://abc/doc.pdf'}] + + # must not raise + await adapter._process_image_components(chain) + assert 'base64' not in chain[0] + adapter.logger.error.assert_awaited_once() diff --git a/tests/unit_tests/provider/test_localagent_inbound_attachments.py b/tests/unit_tests/provider/test_localagent_inbound_attachments.py new file mode 100644 index 00000000..bc7352f1 --- /dev/null +++ b/tests/unit_tests/provider/test_localagent_inbound_attachments.py @@ -0,0 +1,146 @@ +"""Unit tests for LocalAgentRunner._inject_inbound_attachments. + +Covers the user -> sandbox attachment path added for the Box attachment +round-trip: + +* materialized descriptors are stashed on the query and described to the model + via an appended text note (in-sandbox paths + outbox convention); +* non-image file parts (file_base64 / file_url) are stripped from the user + message content because OpenAI-compatible chat models reject them, while + image and text parts are kept for vision models; +* the helper is a no-op when the box service is unavailable or yields nothing, + and never raises into the chat turn on materialization failure. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +import langbot_plugin.api.entities.builtin.provider.message as provider_message + +from langbot.pkg.provider.runners.localagent import LocalAgentRunner + + +def _make_runner(box_service) -> LocalAgentRunner: + runner = LocalAgentRunner.__new__(LocalAgentRunner) + runner.ap = SimpleNamespace(logger=Mock(), box_service=box_service) + return runner + + +def _make_query(): + return SimpleNamespace(variables={}, query_id='q-123') + + +def _box_service(attachments): + svc = SimpleNamespace( + available=True, + OUTBOX_MOUNT_DIR='/outbox', + materialize_inbound_attachments=AsyncMock(return_value=attachments), + ) + return svc + + +@pytest.mark.asyncio +async def test_inject_strips_file_parts_and_appends_note(): + box = _box_service([{'type': 'Voice', 'path': '/inbox/q-123/voice.wav', 'size': 176000}]) + runner = _make_runner(box) + query = _make_query() + user_message = provider_message.Message( + role='user', + content=[ + provider_message.ContentElement.from_text('transcribe this'), + provider_message.ContentElement.from_file_base64('data:audio/wav;base64,AAAA', 'voice.wav'), + ], + ) + + await runner._inject_inbound_attachments(query, user_message) + + types = [getattr(ce, 'type', None) for ce in user_message.content] + # file_base64 dropped; text kept; sandbox-path note appended as text + assert 'file_base64' not in types + assert types.count('text') == 2 + note = user_message.content[-1].text + assert '/inbox/q-123/voice.wav' in note + assert '/outbox/q-123' in note + # descriptors stashed for downstream stages + assert query.variables['_sandbox_inbound_attachments'] == box.materialize_inbound_attachments.return_value + + +@pytest.mark.asyncio +async def test_inject_keeps_image_parts(): + box = _box_service([{'type': 'Image', 'path': '/inbox/q-123/pic.png', 'size': 1234}]) + runner = _make_runner(box) + query = _make_query() + user_message = provider_message.Message( + role='user', + content=[ + provider_message.ContentElement.from_text('what is this'), + provider_message.ContentElement.from_image_base64('data:image/png;base64,iVBORw0K'), + ], + ) + + await runner._inject_inbound_attachments(query, user_message) + + types = [getattr(ce, 'type', None) for ce in user_message.content] + assert 'image_base64' in types # vision part preserved + assert types[-1] == 'text' # note appended last + + +@pytest.mark.asyncio +async def test_inject_promotes_string_content_to_list_with_note(): + box = _box_service([{'type': 'File', 'path': '/inbox/q-123/data.csv', 'size': 42}]) + runner = _make_runner(box) + query = _make_query() + user_message = provider_message.Message(role='user', content='clean this csv') + + await runner._inject_inbound_attachments(query, user_message) + + assert isinstance(user_message.content, list) + assert [getattr(ce, 'type', None) for ce in user_message.content] == ['text', 'text'] + assert user_message.content[0].text == 'clean this csv' + assert '/inbox/q-123/data.csv' in user_message.content[1].text + + +@pytest.mark.asyncio +async def test_inject_noop_without_box_service(): + runner = _make_runner(box_service=None) + query = _make_query() + user_message = provider_message.Message(role='user', content='hello') + + await runner._inject_inbound_attachments(query, user_message) + + assert user_message.content == 'hello' + assert '_sandbox_inbound_attachments' not in query.variables + + +@pytest.mark.asyncio +async def test_inject_noop_when_no_attachments(): + box = _box_service([]) + runner = _make_runner(box) + query = _make_query() + user_message = provider_message.Message(role='user', content='hello') + + await runner._inject_inbound_attachments(query, user_message) + + assert user_message.content == 'hello' + assert '_sandbox_inbound_attachments' not in query.variables + + +@pytest.mark.asyncio +async def test_inject_swallows_materialization_error(): + box = SimpleNamespace( + available=True, + OUTBOX_MOUNT_DIR='/outbox', + materialize_inbound_attachments=AsyncMock(side_effect=RuntimeError('disk full')), + ) + runner = _make_runner(box) + query = _make_query() + user_message = provider_message.Message(role='user', content='hello') + + # must not raise + await runner._inject_inbound_attachments(query, user_message) + assert user_message.content == 'hello' + runner.ap.logger.warning.assert_called_once()