fix(send_message): avoid shared schema mutation and support sync enricher handlers

This commit is contained in:
Kevin Anderson 2026-07-15 22:35:32 -06:00 committed by Teknium
parent 482682db78
commit 274214d3c9
3 changed files with 127 additions and 42 deletions

View File

@ -1288,10 +1288,11 @@ class PluginContext:
Args:
platform_name: Platform key used in target strings
(e.g. ``"myplatform"`` ``myplatform:chat_id``).
handler: Async or sync callable receiving
``(args, chat_id, platform_name, pconfig)`` and returning
a dict like ``{"success": True, "message_id": "..."}``
or ``{"error": "..."}``.
handler: Callable receiving ``(args, chat_id, platform_name, pconfig)``
and returning a dict like ``{"success": True, "message_id": "..."}``
or ``{"error": "..."}``. May be ``async def`` or a regular
function the dispatcher detects coroutine functions via
``inspect.iscoroutinefunction`` and awaits as needed.
schema_fragment: Optional dict of JSON-schema properties to merge
into ``send_message``'s parameter schema so the LLM sees the
custom fields.

View File

@ -2,15 +2,17 @@
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, MagicMock
import pytest
from tools.send_message_tool import (
SEND_MESSAGE_SCHEMA,
_SEND_MESSAGE_ENRICHERS,
_SEND_MESSAGE_SCHEMA_FRAGMENTS,
_parse_target_ref,
_send_to_platform,
get_send_message_schema,
register_send_message_enricher,
)
@ -19,27 +21,10 @@ from tools.send_message_tool import (
def _reset_enrichers():
"""Clear the enricher registry before and after every test."""
_SEND_MESSAGE_ENRICHERS.clear()
# Also clear any schema fragments injected by previous tests
for key in list(SEND_MESSAGE_SCHEMA["parameters"]["properties"]):
if key not in {
"action",
"target",
"message",
"emoji",
"message_id",
}:
SEND_MESSAGE_SCHEMA["parameters"]["properties"].pop(key, None)
_SEND_MESSAGE_SCHEMA_FRAGMENTS.clear()
yield
_SEND_MESSAGE_ENRICHERS.clear()
for key in list(SEND_MESSAGE_SCHEMA["parameters"]["properties"]):
if key not in {
"action",
"target",
"message",
"emoji",
"message_id",
}:
SEND_MESSAGE_SCHEMA["parameters"]["properties"].pop(key, None)
_SEND_MESSAGE_SCHEMA_FRAGMENTS.clear()
class TestParseTargetRef:
@ -69,27 +54,45 @@ class TestParseTargetRef:
class TestSchemaMerge:
def test_schema_fragment_merged(self):
"""Schema fragments are injected into SEND_MESSAGE_SCHEMA."""
def test_schema_fragment_stored(self):
"""Schema fragments are stored in _SEND_MESSAGE_SCHEMA_FRAGMENTS."""
register_send_message_enricher(
"myplatform",
AsyncMock(),
schema_fragment={"voice": {"type": "string", "description": "Voice setting"}},
)
props = SEND_MESSAGE_SCHEMA["parameters"]["properties"]
assert "voice" in props
assert props["voice"]["type"] == "string"
assert "myplatform" in _SEND_MESSAGE_SCHEMA_FRAGMENTS
assert _SEND_MESSAGE_SCHEMA_FRAGMENTS["myplatform"]["voice"]["type"] == "string"
def test_get_send_message_schema_assembles_fragments(self):
"""get_send_message_schema returns a copy with fragments merged."""
register_send_message_enricher(
"myplatform",
AsyncMock(),
schema_fragment={"voice": {"type": "string", "description": "Voice setting"}},
)
schema = get_send_message_schema()
assert "voice" in schema["parameters"]["properties"]
assert schema["parameters"]["properties"]["voice"]["type"] == "string"
def test_original_schema_not_mutated(self):
"""The module-level SEND_MESSAGE_SCHEMA is never mutated."""
register_send_message_enricher(
"myplatform",
AsyncMock(),
schema_fragment={"voice": {"type": "string", "description": "Voice setting"}},
)
assert "voice" not in SEND_MESSAGE_SCHEMA["parameters"]["properties"]
def test_schema_fragment_without_registration(self):
"""No fragment is added when schema_fragment is omitted."""
register_send_message_enricher("myplatform", AsyncMock())
props = SEND_MESSAGE_SCHEMA["parameters"]["properties"]
assert "voice" not in props
assert "myplatform" not in _SEND_MESSAGE_SCHEMA_FRAGMENTS
class TestHandlerRouting:
def test_enricher_handler_invoked(self):
"""The enricher handler receives correct arguments."""
def test_async_handler_invoked(self):
"""The async enricher handler receives correct arguments."""
handler = AsyncMock(return_value={"success": True, "message_id": "msg_1"})
register_send_message_enricher("myplatform", handler)
@ -112,8 +115,32 @@ class TestHandlerRouting:
assert call_args[2] == "myplatform"
assert call_args[3] is pconfig
def test_enricher_handler_result(self):
"""Enricher result is returned verbatim."""
def test_sync_handler_invoked(self):
"""The sync enricher handler is called directly (not awaited)."""
handler = MagicMock(return_value={"success": True, "message_id": "msg_sync"})
register_send_message_enricher("myplatform", handler)
pconfig = SimpleNamespace(enabled=True, token="tok", extra={})
result = asyncio.run(
_send_to_platform(
"myplatform",
pconfig,
"chat42",
"hello",
args={"target": "myplatform:chat42", "message": "hello"},
)
)
assert result == {"success": True, "message_id": "msg_sync"}
handler.assert_called_once()
call_args = handler.call_args.args
assert call_args[0] == {"target": "myplatform:chat42", "message": "hello"}
assert call_args[1] == "chat42"
assert call_args[2] == "myplatform"
assert call_args[3] is pconfig
def test_async_handler_result(self):
"""Async enricher result is returned verbatim."""
handler = AsyncMock(return_value={"success": True, "message_id": "msg_1"})
register_send_message_enricher("myplatform", handler)
@ -127,8 +154,23 @@ class TestHandlerRouting:
)
assert result == {"success": True, "message_id": "msg_1"}
def test_enricher_handler_error(self):
"""Enricher exceptions are caught and surfaced as error dicts."""
def test_sync_handler_result(self):
"""Sync enricher result is returned verbatim."""
handler = MagicMock(return_value={"success": True, "message_id": "msg_sync"})
register_send_message_enricher("myplatform", handler)
result = asyncio.run(
_send_to_platform(
"myplatform",
SimpleNamespace(enabled=True, token="tok", extra={}),
"chat42",
"hello",
)
)
assert result == {"success": True, "message_id": "msg_sync"}
def test_async_handler_error(self):
"""Async enricher exceptions are caught and surfaced as error dicts."""
handler = AsyncMock(side_effect=RuntimeError("boom"))
register_send_message_enricher("myplatform", handler)
@ -143,6 +185,22 @@ class TestHandlerRouting:
assert "error" in result
assert "boom" in result["error"]
def test_sync_handler_error(self):
"""Sync enricher exceptions are caught and surfaced as error dicts."""
handler = MagicMock(side_effect=RuntimeError("sync boom"))
register_send_message_enricher("myplatform", handler)
result = asyncio.run(
_send_to_platform(
"myplatform",
SimpleNamespace(enabled=True, token="tok", extra={}),
"chat42",
"hello",
)
)
assert "error" in result
assert "sync boom" in result["error"]
class TestFallback:
def test_no_enricher_falls_through(self):

View File

@ -90,12 +90,19 @@ _DEFAULT_CAPTION_LIMIT = 4096
# Plugin enricher registry for send_message
# ---------------------------------------------------------------------------
SendMessageEnricher = Callable[[dict, str, str, Any], Awaitable[dict]]
"""Callable receiving (args, chat_id, platform_name, pconfig) -> dict result."""
SendMessageEnricher = Callable[[dict, str, str, Any], Awaitable[dict] | dict]
"""Callable receiving (args, chat_id, platform_name, pconfig) -> dict result.
May be sync or async the dispatcher detects coroutine functions via
``inspect.iscoroutinefunction`` and awaits as needed.
"""
_SEND_MESSAGE_ENRICHERS: dict[str, SendMessageEnricher] = {}
"""platform_name -> enricher handler."""
_SEND_MESSAGE_SCHEMA_FRAGMENTS: dict[str, dict] = {}
"""platform_name -> schema fragment dict (JSON Schema properties)."""
def register_send_message_enricher(
platform_name: str,
@ -109,8 +116,22 @@ def register_send_message_enricher(
"""
_SEND_MESSAGE_ENRICHERS[platform_name] = handler
if schema_fragment:
for key, spec in schema_fragment.items():
SEND_MESSAGE_SCHEMA["parameters"]["properties"][key] = spec
_SEND_MESSAGE_SCHEMA_FRAGMENTS[platform_name] = schema_fragment
def get_send_message_schema() -> dict:
"""Return a fresh copy of the send_message schema with plugin fragments merged.
Fragments are assembled on demand so deferred plugin registration never
mutates the shared ``SEND_MESSAGE_SCHEMA`` dict. Callers that need the
current wire schema (e.g. tool-search builders, MCP catalogues) should use
this instead of reading ``SEND_MESSAGE_SCHEMA`` directly.
"""
import copy
schema = copy.deepcopy(SEND_MESSAGE_SCHEMA)
for fragment in _SEND_MESSAGE_SCHEMA_FRAGMENTS.values():
schema["parameters"]["properties"].update(fragment)
return schema
def _media_caption_split(text, media_files, *, max_caption_len):
@ -1184,7 +1205,12 @@ async def _send_to_platform(platform, pconfig, chat_id, message, thread_id=None,
if args is None:
args = {}
try:
return await enricher(args, chat_id, platform_name, pconfig)
import inspect
if inspect.iscoroutinefunction(enricher):
result = await enricher(args, chat_id, platform_name, pconfig)
else:
result = enricher(args, chat_id, platform_name, pconfig)
return result
except Exception as e:
return {"error": f"Enricher send failed: {e}"}
# Plugin platform: route through the gateway's live adapter if