From 41b55952ffed2c93d1cb0703b5b3ea608987c46d Mon Sep 17 00:00:00 2001 From: Qingsong Li <2731468336@qq.com> Date: Tue, 20 Oct 2026 02:14:36 +0800 Subject: [PATCH] fix(runner): avoid duplicate inline image payloads (#2631) * fix(runner): avoid duplicate inline image payloads * fix(runner): preserve attachments and complete HTTP agent delivery Co-authored-by: multica-agent --------- Co-authored-by: Qingsong Li <201055554+qingshungLI@users.noreply.github.com> Co-authored-by: Chin Rock Co-authored-by: multica-agent --- docs/platforms/http-bot.md | 6 + src/langbot/pkg/box/runner.py | 19 +- src/langbot/pkg/platform/botmgr.py | 1 + src/langbot/pkg/platform/sources/http_bot.py | 23 +- .../modelmgr/requesters/litellmchat.py | 6 +- src/langbot/pkg/utils/image.py | 24 ++ tests/unit_tests/box/test_runner_box.py | 49 +++ .../platform/test_http_bot_agent_boundary.py | 384 ++++++++++++++++++ .../test_image_base64_normalization.py | 41 ++ 9 files changed, 548 insertions(+), 5 deletions(-) create mode 100644 tests/unit_tests/platform/test_http_bot_agent_boundary.py create mode 100644 tests/unit_tests/provider/test_image_base64_normalization.py diff --git a/docs/platforms/http-bot.md b/docs/platforms/http-bot.md index 71fbc6799..94117bca2 100644 --- a/docs/platforms/http-bot.md +++ b/docs/platforms/http-bot.md @@ -202,6 +202,12 @@ This is **lossy** (you lose `sequence` / streaming boundaries) and blocks up to `callback_timeout × 4` seconds. Prefer the callback model for anything real-time or multi-part. Only one in-flight `/sync` per `session_id`. +Collection waits for both a final reply and the inbound listener to finish. An +Agent's text reply does not end collection while that Agent is still exporting +and sending output files. The same timeout bounds the entire wait; timing out +does not cancel the running Agent. Replies sent after timeout use normal callback +delivery. + --- ## 9. Error envelope diff --git a/src/langbot/pkg/box/runner.py b/src/langbot/pkg/box/runner.py index c31a72b4c..21ab66731 100644 --- a/src/langbot/pkg/box/runner.py +++ b/src/langbot/pkg/box/runner.py @@ -36,14 +36,24 @@ def prepare_input_files(query, agent_input): event envelope is metadata, never permission to read the Host filesystem. """ from langbot_plugin.api.entities.builtin.runner.input import InputAttachment - from ..agent.runner.query_entry_adapter import QueryEntryAdapter originals = input_files(query) + # References are scoped to this query's previous preparation, not to the + # platform message order or to any caller-supplied filesystem path. + prepared = getattr(query, '_box_inputs', None) + previous_refs = {f'attachment-{i}': item for i, item in enumerate(prepared or [])} if not agent_input.attachments: + from ..agent.runner.query_entry_adapter import QueryEntryAdapter + agent_input.attachments = [ InputAttachment.model_validate(item) for item in QueryEntryAdapter._build_attachments(query, [c.model_dump() for c in agent_input.contents]) ] + inline_images = { + content.image_base64 + for content in (agent_input.contents or []) + if content.type == 'image_base64' and content.image_base64 + } components = [] for i, attachment in enumerate(agent_input.attachments): kind = (attachment.type or 'file').lower() @@ -52,6 +62,11 @@ def prepare_input_files(query, agent_input): if cls is pm.File: fields.update(name=attachment.name or '', id=attachment.id or '') component = cls(**fields) + referenced = previous_refs.get(attachment.ref) + if type(referenced) is cls and not attachment.content and not attachment.url: + components.append(referenced) + attachment.ref, attachment.path = f'attachment-{i}', None + continue # Match by payload, never by list position (contents and message chains # can describe attachments in different orders). for original in originals: @@ -77,6 +92,8 @@ def prepare_input_files(query, agent_input): break components.append(component) attachment.ref, attachment.path = f'attachment-{i}', None + if kind == 'image' and attachment.content in inline_images: + attachment.content = None object.__setattr__(query, '_box_inputs', components) diff --git a/src/langbot/pkg/platform/botmgr.py b/src/langbot/pkg/platform/botmgr.py index ac425a9c4..ad30f5e7b 100644 --- a/src/langbot/pkg/platform/botmgr.py +++ b/src/langbot/pkg/platform/botmgr.py @@ -34,6 +34,7 @@ from ..agent.runner.host_models import ( ) from ..agent.runner.resource_policy import ResourcePolicyProjector from ..agent.runner.platform_tools import resolve_agent_platform_tool_names +from ..agent.runner.query_entry_adapter import QueryEntryAdapter from ..entity.persistence import workspace as persistence_workspace from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType, RequestContext diff --git a/src/langbot/pkg/platform/sources/http_bot.py b/src/langbot/pkg/platform/sources/http_bot.py index 2cb4f19ca..c7c7652a6 100644 --- a/src/langbot/pkg/platform/sources/http_bot.py +++ b/src/langbot/pkg/platform/sources/http_bot.py @@ -89,7 +89,7 @@ class _SessionOutbound: class _SyncCollector: - """Collects reply parts for a /sync request and resolves when the turn ends.""" + """Collects reply parts and signals receipt of a final message.""" def __init__(self) -> None: self.parts: list = [] @@ -162,6 +162,9 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter): async def is_muted(self, group_id: int) -> bool: return False + def get_supported_apis(self) -> list[str]: + return ['send_message', 'reply_message'] + async def is_stream_output_supported(self) -> bool: return True @@ -582,6 +585,12 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter): 'message': [c.model_dump() if hasattr(c, 'model_dump') else c.__dict__ for c in message], 'timestamp': datetime.now().isoformat(), } + collector = self.sync_waiters.get(str(target_id)) + if collector is not None: + collector.parts.extend(payload['message']) + collector.done.set() + return payload + await self._enqueue_callback(str(target_id), payload) return payload @@ -612,6 +621,9 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter): Lossy by design (drops streaming/ordering nuance); documented as such. Concurrency-safe: routing is via the per-session ``_sync_waiters`` registry that ``_emit_reply`` consults, not by patching methods. + A final message can precede Agent output-file delivery, so collection + also waits for the inbound listener to finish. Legacy queued pipelines + may finish their listener before sending any reply. """ if session_id in self.sync_waiters: return self._err('duplicate', 'a sync request is already in flight for this session') @@ -622,12 +634,19 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter): if listener_task is None: self.sync_waiters.pop(session_id, None) return self._err('overloaded', 'too many inbound messages are already being processed') + + async def wait_for_turn(): + await collector.done.wait() + await asyncio.shield(listener_task) + try: timeout = int(self.config.get('callback_timeout', 15)) * 4 try: - await asyncio.wait_for(collector.done.wait(), timeout=timeout) + await asyncio.wait_for(wait_for_turn(), timeout=timeout) except asyncio.TimeoutError: await self.logger.warning(f'http_bot sync wait timed out for session {session_id}') + except Exception: + await self.logger.warning(f'http_bot sync listener failed for session {session_id}') finally: self.sync_waiters.pop(session_id, None) state = self.outbound_states.get(session_id) diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index 6f1de0e86..c8bb662cd 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.secret_managers.main import get_secret, get_secret_bool from .. import errors, reasoning, requester -from ....utils import httpclient +from ....utils import httpclient, image import langbot_plugin.api.entities.builtin.resource.tool as resource_tool import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query import langbot_plugin.api.entities.builtin.provider.message as provider_message @@ -895,7 +895,9 @@ class LiteLLMRequester(requester.ProviderAPIRequester): if isinstance(raw_image_url, dict): fallback_url = raw_image_url.get('url') if image_b64 or fallback_url: - part['image_url'] = {'url': image_b64 or fallback_url} + part['image_url'] = { + 'url': image.normalize_image_data_url(image_b64) if image_b64 else fallback_url + } part['type'] = 'image_url' part.pop('image_base64', None) else: diff --git a/src/langbot/pkg/utils/image.py b/src/langbot/pkg/utils/image.py index 435671ae8..0b9436661 100644 --- a/src/langbot/pkg/utils/image.py +++ b/src/langbot/pkg/utils/image.py @@ -17,6 +17,30 @@ _INSECURE_SSL_CONTEXT.verify_mode = ssl.CERT_NONE DEFAULT_BASE64_MEDIA_LIMIT = 10 * 1024 * 1024 +def normalize_image_data_url(value: str) -> str: + """Wrap raw image base64 with its detected MIME type, never a guessed one. + + Existing data URLs are preserved for provider-specific compatibility. + Bound decoding before inspecting untrusted image headers. + """ + if value.startswith('data:'): + return value + if len(value) > 4 * ((DEFAULT_BASE64_MEDIA_LIMIT + 2) // 3): + raise ValueError('Base64 image exceeds the media size limit') + try: + decoded = base64.b64decode(value, validate=True) + if len(decoded) > DEFAULT_BASE64_MEDIA_LIMIT: + raise ValueError('Base64 image exceeds the media size limit') + with PIL.Image.open(io.BytesIO(decoded)) as image: + mime = PIL.Image.MIME.get(image.format or '') + image.verify() + if not mime or not mime.startswith('image/'): + raise ValueError('Unknown image MIME type') + except (ValueError, OSError, PIL.Image.DecompressionBombError) as exc: + raise ValueError('Invalid base64 image: expected encoded image bytes') from exc + return f'data:{mime};base64,{value}' + + def _detect_image_format(file_bytes: bytes) -> str: with PIL.Image.open(io.BytesIO(file_bytes)) as image: return str(image.format or 'jpeg').lower() diff --git a/tests/unit_tests/box/test_runner_box.py b/tests/unit_tests/box/test_runner_box.py index f5632fe4a..150a51279 100644 --- a/tests/unit_tests/box/test_runner_box.py +++ b/tests/unit_tests/box/test_runner_box.py @@ -156,6 +156,55 @@ def test_event_attachments_preserved_without_eager_io_or_host_path_access(): assert image.attachments[0].url == 'https://example.invalid/image.png' +def test_duplicate_inline_image_payload_is_replaced_by_attachment_reference(): + from langbot_plugin.api.entities.builtin.provider.message import ContentElement + from langbot_plugin.api.entities.builtin.runner.input import AgentInput, InputAttachment + from langbot.pkg.box.runner import input_files, prepare_input_files + + image_base64 = 'data:image/png;base64,YQ==' + query = SimpleNamespace(message_chain=MessageChain([])) + value = AgentInput( + contents=[ContentElement.from_image_base64(image_base64)], + attachments=[ + InputAttachment(type='image', source='base64', content=image_base64), + InputAttachment(type='image', source='base64', content='data:image/png;base64,Yg=='), + ], + ) + + prepare_input_files(query, value) + + assert value.contents[0].image_base64 == image_base64 + assert value.attachments[0].content is None + assert value.attachments[0].ref == 'attachment-0' + assert value.attachments[1].content == 'data:image/png;base64,Yg==' + assert value.attachments[1].ref == 'attachment-1' + assert input_files(query)[0].base64 == image_base64 + + +def test_repeated_prepare_preserves_deduplicated_images_after_reordering(): + from langbot_plugin.api.entities.builtin.provider.message import ContentElement + from langbot_plugin.api.entities.builtin.runner.input import AgentInput, InputAttachment + from langbot.pkg.box.runner import input_files, prepare_input_files + + images = ['data:image/png;base64,YQ==', 'data:image/png;base64,Yg=='] + query = SimpleNamespace(message_chain=MessageChain([])) + value = AgentInput( + contents=[ContentElement.from_image_base64(image) for image in images], + attachments=[InputAttachment(type='image', content=image) for image in images], + ) + prepare_input_files(query, value) + value = AgentInput.model_validate(value.model_dump()) + value.attachments.reverse() + for _ in range(3): + prepare_input_files(query, value) + assert [item.base64 for item in input_files(query)] == images[::-1] + assert all(item.content is None for item in value.attachments) + + other_query = SimpleNamespace(message_chain=MessageChain([])) + prepare_input_files(other_query, value) + assert all(not item.base64 for item in input_files(other_query)) + + def test_attachment_references_follow_metadata_not_platform_list_order(): from langbot_plugin.api.entities.builtin.runner.input import AgentInput, InputAttachment from langbot.pkg.box.runner import prepare_input_files, input_files diff --git a/tests/unit_tests/platform/test_http_bot_agent_boundary.py b/tests/unit_tests/platform/test_http_bot_agent_boundary.py new file mode 100644 index 000000000..a044054d0 --- /dev/null +++ b/tests/unit_tests/platform/test_http_bot_agent_boundary.py @@ -0,0 +1,384 @@ +from __future__ import annotations + +import asyncio +import base64 +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import quart +import yaml +from langbot_plugin.api.definition.abstract.platform.event_logger import AbstractEventLogger +from langbot_plugin.api.entities.builtin.platform import events, message +from langbot_plugin.box.client import BoxRuntimeClient +from langbot_plugin.api.proxies.runner.box import RunnerBoxAPIMixin +from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction + +from langbot.pkg.agent.runner.execution_context import build_execution_query +from langbot.pkg.agent.runner.platform_tools import ( + build_platform_tool_resources, + execute_platform_tool, + freeze_platform_context, + resolve_agent_platform_tool_names, + resolve_platform_api_call, +) +from langbot.pkg.box.runner import RunBoxBinding, RunnerBoxService, exported_message, prepare_input_files +from langbot.pkg.box.service import BoxService +from langbot.pkg.discover.engine import Component +from langbot.pkg.platform.botmgr import RuntimeBot +from langbot.pkg.platform.sources import http_bot_signing as signing +from langbot.pkg.platform.sources.http_bot import HttpBotAdapter + + +def discovered_adapter(): + relative = 'pkg/platform/sources/http_bot.yaml' + manifest = yaml.safe_load((Path(__file__).parents[3] / 'src/langbot' / relative).read_text()) + adapter_class = Component('builtin', manifest, relative).get_python_component_class() + assert adapter_class is HttpBotAdapter + return adapter_class( + config={'inbound_secret': 'test-only-secret', 'signature_required': True, 'callback_timeout': 1}, + logger=AsyncMock(spec=AbstractEventLogger), + ) + + +def envelope(bot, adapter, inbound): + return bot._eba_event_to_agent_envelope(bot._legacy_message_to_eba_event(inbound, adapter), adapter) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('session_type', ['person', 'group']) +@pytest.mark.parametrize('mixed', [False, True]) +async def test_signed_http_file_reaches_box_and_authorized_agent_reply_reaches_sync( + tmp_path, monkeypatch, session_type, mixed +): + async def inline_thread(function, *args, **kwargs): + return function(*args, **kwargs) + + monkeypatch.setattr(asyncio, 'to_thread', inline_thread) + adapter = discovered_adapter() + adapter._enqueue_callback = AsyncMock(side_effect=AssertionError('sync reply must not use callback')) + bot = object.__new__(RuntimeBot) + bot.bot_entity = SimpleNamespace(uuid='http-test-bot') + app = SimpleNamespace( + logger=Mock(), + instance_config=SimpleNamespace(data={'box': {'enabled': True}}), + platform_mgr=SimpleNamespace(get_bot_by_uuid=AsyncMock(return_value=SimpleNamespace(adapter=adapter))), + ) + box = BoxService(app, client=Mock(spec=BoxRuntimeClient)) + box._available = True + + async def write_files(query, subdir, target_dir, pending): + for name, content in pending: + destination = tmp_path / target_dir.lstrip('/') / name + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(content) + return [f'{target_dir}/{name}' for name, content in pending] + + box._write_files_into_sandbox = AsyncMock(side_effect=write_files) + runner_box = RunnerBoxService(box) + received = [] + listener_errors = [] + proof = b'HTTP inbound proof survives event conversion' + image = b'\x89PNG\r\n\x1a\n' + + async def listener(inbound, live_adapter): + event = envelope(bot, live_adapter, inbound) + received.append(event) + query = build_execution_query(event, []) + prepare_input_files(query, event.input) + object.__setattr__(query, '_box_binding', RunBoxBinding('run-test', 'box-test', {}, 'run-test')) + imported = await runner_box.import_attachments(query, None) + assert len(imported['items']) == (2 if mixed else 1) + imported_file = next(item for item in imported['items'] if item['name'] == 'proof.txt') + assert (tmp_path / imported_file['path'].lstrip('/')).read_bytes() == proof + assert imported_file['size'] == len(proof) + if mixed: + imported_image = next(item for item in imported['items'] if item['type'] == 'Image') + assert (tmp_path / imported_image['path'].lstrip('/')).read_bytes() == image + resources, capabilities = build_platform_tool_resources( + event, resolve_agent_platform_tool_names({}, event.event_type), ['call'] + ) + assert any(tool['tool_name'] == 'event_reply' for tool in resources) + assert capabilities['supported_apis'] == ['send_message', 'reply_message'] + session = { + 'authorization': { + 'bot_id': event.bot_id, + 'resources': {'tools': resources}, + 'platform_context': freeze_platform_context(event), + } + } + chain = message.MessageChain([message.Plain(text=proof.decode())]) + tool, params, rich = resolve_platform_api_call( + session, + event.bot_id, + 'send_message', + {'target_type': session_type, 'target_id': 'proof-session', 'message': chain.model_dump()}, + ) + assert tool == 'event_reply' + await execute_platform_tool(app, object(), session, tool, params, message_chain=rich) + + async def checked_listener(inbound, live_adapter): + try: + await listener(inbound, live_adapter) + except Exception as error: + listener_errors.append(error) + live_adapter.sync_waiters['proof-session'].done.set() + + adapter.register_listener(events.FriendMessage, checked_listener) + adapter.register_listener(events.GroupMessage, checked_listener) + web = quart.Quart(__name__) + + @web.post('/bots/http-test-bot/sync') + async def sync(): + return await adapter.handle_unified_webhook('http-test-bot', 'sync', quart.request) + + components = [ + {'type': 'Plain', 'text': 'Read proof.txt and reply with its contents.'}, + {'type': 'File', 'name': 'proof.txt', 'base64': base64.b64encode(proof).decode()}, + ] + if mixed: + components.append({'type': 'Image', 'base64': base64.b64encode(image).decode()}) + body = json.dumps({'session_id': 'proof-session', 'session_type': session_type, 'message': components}).encode() + timestamp, signature = signing.sign('test-only-secret', body) + client = web.test_client() + rejected = await client.post('/bots/http-test-bot/sync', data=body) + assert rejected.status_code == 401 + assert received == [] + response = await asyncio.wait_for( + client.post( + '/bots/http-test-bot/sync', + data=body, + headers={signing.HEADER_TIMESTAMP: timestamp, signing.HEADER_SIGNATURE: signature}, + ), + timeout=5, + ) + assert response.status_code == 200 + if listener_errors: + raise listener_errors[0] + data = (await response.get_json())['data'] + assert data['message'] == [{'type': 'Plain', 'text': proof.decode()}] + assert data['reply_to'].startswith('in_') + assert len(received) == 1 + adapter._enqueue_callback.assert_not_awaited() + assert adapter.sync_waiters == {} + await asyncio.gather(*adapter.inbound_tasks) + + +@pytest.mark.parametrize('denial', ['bot', 'target', 'grant', 'operation', 'capability']) +def test_live_http_reply_grants_fail_closed(denial): + adapter = discovered_adapter() + bot = object.__new__(RuntimeBot) + bot.bot_entity = SimpleNamespace(uuid='http-test-bot') + inbound, *_ = adapter._build_event({'session_id': 'proof-session', 'message': [{'type': 'Plain', 'text': 'hi'}]}) + event = envelope(bot, adapter, inbound) + if denial == 'capability': + event.delivery.platform_capabilities['supported_apis'] = [] + resources, _ = build_platform_tool_resources( + event, ['event_reply'], ['detail' if denial == 'operation' else 'call'] + ) + session = { + 'authorization': { + 'bot_id': event.bot_id, + 'resources': {'tools': [] if denial == 'grant' else resources}, + 'platform_context': freeze_platform_context(event), + } + } + with pytest.raises(ValueError, match='not authorized'): + resolve_platform_api_call( + session, + 'foreign-bot' if denial == 'bot' else event.bot_id, + 'send_message', + { + 'target_type': 'person', + 'target_id': 'foreign-session' if denial == 'target' else 'proof-session', + 'message': [{'type': 'Plain', 'text': 'reply'}], + }, + ) + + +@pytest.mark.asyncio +async def test_http_send_without_sync_waiter_preserves_callback_delivery(): + adapter = discovered_adapter() + adapter._enqueue_callback = AsyncMock() + chain = message.MessageChain([message.Plain(text='async reply')]) + payload = await adapter.send_message('person', 'async-session', chain) + adapter._enqueue_callback.assert_awaited_once_with('async-session', payload) + assert payload['message'] == chain.model_dump() + assert payload['is_final'] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize('text_api', ['send', 'reply', 'chunk']) +async def test_sync_waits_for_authorized_exported_file_after_final_text(text_api): + adapter = discovered_adapter() + adapter._enqueue_callback = AsyncMock() + bot = object.__new__(RuntimeBot) + bot.bot_entity = SimpleNamespace(uuid='http-test-bot') + inbound, session_id, _, message_id = adapter._build_event( + {'session_id': 'generated-session', 'message': [{'type': 'Plain', 'text': 'generate a file'}]} + ) + text_sent = asyncio.Event() + export_allowed = asyncio.Event() + content = b'generated file bytes' + + async def listener(event, live_adapter): + text = message.MessageChain([message.Plain(text=content.decode())]) + if text_api == 'send': + await live_adapter.send_message('person', session_id, text) + elif text_api == 'reply': + await live_adapter.reply_message(event, text) + else: + await live_adapter.reply_message_chunk(event, None, text, is_final=True) + text_sent.set() + await export_allowed.wait() + agent_event = envelope(bot, live_adapter, event) + resources, _ = build_platform_tool_resources(agent_event, ['event_reply'], ['call']) + session = { + 'authorization': { + 'bot_id': agent_event.bot_id, + 'resources': {'tools': resources}, + 'platform_context': freeze_platform_context(agent_event), + } + } + query = SimpleNamespace(_box_binding=RunBoxBinding('generated-run', 'box', {}, 'generated-run')) + box = SimpleNamespace( + _ATTACHMENT_MAX_TOTAL_BYTES=1000000, + collect_outbound_attachments=AsyncMock( + return_value=[ + { + 'type': 'File', + 'name': 'generated-proof.txt', + 'base64': base64.b64encode(content).decode(), + } + ] + ), + ) + exported = await RunnerBoxService(box).export_files(query) + file_ids = [item['id'] for item in exported['items']] + app = SimpleNamespace( + platform_mgr=SimpleNamespace(get_bot_by_uuid=AsyncMock(return_value=SimpleNamespace(adapter=live_adapter))) + ) + + class FileReplyAPI(RunnerBoxAPIMixin): + async def _box_call(self, rpc_action, **data): + assert rpc_action == PluginToRuntimeAction.CALL_PLATFORM_API + files = exported_message(query, data['file_ids']) + tool, params, rich = resolve_platform_api_call( + session, + agent_event.bot_id, + data['action'], + {**data['params'], 'message': files.model_dump()}, + data['context_tool'], + ) + result = await execute_platform_tool(app, object(), session, tool, params, message_chain=rich) + exported_message(query, data['file_ids'], consume=True) + return {'result': result} + + await FileReplyAPI().reply_files(file_ids) + + web = quart.Quart(__name__) + async with web.app_context(): + request_task = asyncio.create_task(adapter._run_sync(inbound, listener, session_id, message_id)) + try: + await asyncio.wait_for(text_sent.wait(), timeout=1) + done, _ = await asyncio.wait({request_task}, timeout=0.03) + assert not done, 'A final text send must not close the active Agent collector' + assert session_id in adapter.sync_waiters + finally: + export_allowed.set() + response, status = await asyncio.wait_for(request_task, timeout=2) + assert status == 200 + payload = (await response.get_json())['data'] + assert payload['reply_to'] == message_id + assert len(payload['message']) == 2 + assert payload['message'][0] == {'type': 'Plain', 'text': content.decode()} + delivered = payload['message'][1] + assert delivered['type'] == 'File' + assert delivered['name'] == 'generated-proof.txt' + assert base64.b64decode(delivered['base64'], validate=True) == content + adapter._enqueue_callback.assert_not_awaited() + assert adapter.sync_waiters == {} + assert adapter.outbound_states == {} + + +@pytest.mark.asyncio +async def test_sync_still_waits_for_queued_legacy_pipeline_reply(): + adapter = discovered_adapter() + listener_returned = asyncio.Event() + inbound, session_id, _, message_id = adapter._build_event( + {'session_id': 'legacy-session', 'message': [{'type': 'Plain', 'text': 'hi'}]} + ) + + async def listener(event, live_adapter): + listener_returned.set() + + web = quart.Quart(__name__) + async with web.app_context(): + request_task = asyncio.create_task(adapter._run_sync(inbound, listener, session_id, message_id)) + await asyncio.wait_for(listener_returned.wait(), timeout=1) + done, _ = await asyncio.wait({request_task}, timeout=0.03) + assert not done + await adapter.reply_message(inbound, message.MessageChain([message.Plain(text='queued reply')])) + response, status = await asyncio.wait_for(request_task, timeout=1) + assert status == 200 + assert (await response.get_json())['data']['message'] == [{'type': 'Plain', 'text': 'queued reply'}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_request', [False, True]) +async def test_sync_timeout_or_disconnect_does_not_cancel_active_listener(cancel_request): + adapter = discovered_adapter() + adapter.config['callback_timeout'] = 0 if not cancel_request else 1 + adapter._enqueue_callback = AsyncMock() + listener_started = asyncio.Event() + finish = asyncio.Event() + inbound, session_id, _, message_id = adapter._build_event( + {'session_id': 'slow-session', 'message': [{'type': 'Plain', 'text': 'hi'}]} + ) + + async def listener(event, live_adapter): + await live_adapter.send_message('person', session_id, message.MessageChain([message.Plain(text='text')])) + listener_started.set() + await finish.wait() + await live_adapter.send_message('person', session_id, message.MessageChain([message.Plain(text='late')])) + + web = quart.Quart(__name__) + async with web.app_context(): + request_task = asyncio.create_task(adapter._run_sync(inbound, listener, session_id, message_id)) + await asyncio.wait_for(listener_started.wait(), timeout=1) + tasks = tuple(adapter.inbound_tasks) + if cancel_request: + await asyncio.sleep(0) + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + else: + _, status = await asyncio.wait_for(request_task, timeout=1) + assert status == 200 + assert adapter.sync_waiters == {} + assert tasks and all(not task.done() for task in tasks) + finish.set() + await asyncio.gather(*tasks) + adapter._enqueue_callback.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_sync_listener_failure_after_reply_returns_partial_without_leaking_error(): + adapter = discovered_adapter() + inbound, session_id, _, message_id = adapter._build_event( + {'session_id': 'failed-session', 'message': [{'type': 'Plain', 'text': 'hi'}]} + ) + + async def listener(event, live_adapter): + await live_adapter.send_message('person', session_id, message.MessageChain([message.Plain(text='text')])) + raise RuntimeError('private failure detail') + + web = quart.Quart(__name__) + async with web.app_context(): + response, status = await asyncio.wait_for(adapter._run_sync(inbound, listener, session_id, message_id), 1) + assert status == 200 + assert 'private failure detail' not in (await response.get_data()).decode() + assert adapter.sync_waiters == {} + adapter.logger.warning.assert_awaited_once() diff --git a/tests/unit_tests/provider/test_image_base64_normalization.py b/tests/unit_tests/provider/test_image_base64_normalization.py new file mode 100644 index 000000000..a8bb7f5b2 --- /dev/null +++ b/tests/unit_tests/provider/test_image_base64_normalization.py @@ -0,0 +1,41 @@ +import base64 +import io + +import pytest +from PIL import Image +import langbot_plugin.api.entities.builtin.provider.message as provider_message + +from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester +from langbot.pkg.utils import image + + +@pytest.mark.parametrize( + 'format,mime', [('PNG', 'image/png'), ('JPEG', 'image/jpeg'), ('WEBP', 'image/webp'), ('GIF', 'image/gif')] +) +def test_raw_image_becomes_data_url_without_mutating_input(format, mime): + buffer = io.BytesIO() + Image.new('RGB', (4, 4), 'red').save(buffer, format=format) + raw = base64.b64encode(buffer.getvalue()).decode() + msg = provider_message.Message(role='user', content=[provider_message.ContentElement.from_image_base64(raw)]) + requester = LiteLLMRequester.__new__(LiteLLMRequester) + result = requester._convert_messages([msg]) + assert result[0]['content'][0]['image_url']['url'] == f'data:{mime};base64,{raw}' + assert msg.content[0].image_base64 == raw + assert requester._convert_messages([msg]) == result + + +@pytest.mark.parametrize('value', ['!invalid!', 'aGVsbG8=', '']) +def test_invalid_raw_image_is_rejected(value): + with pytest.raises(ValueError, match='Invalid base64 image'): + image.normalize_image_data_url(value) + + +def test_existing_data_url_is_preserved(): + value = 'data:image/jpeg;base64,existing-payload' + assert image.normalize_image_data_url(value) == value + + +def test_oversized_raw_input_rejected_before_decode(monkeypatch): + monkeypatch.setattr(image, 'DEFAULT_BASE64_MEDIA_LIMIT', 3) + with pytest.raises(ValueError, match='size limit'): + image.normalize_image_data_url('A' * 8)