seb/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(yuanbao): consolidate media resolution into dedicated pipeline middlewares

Browse Source
This commit is contained in:
loongzhao
2026-06-09 03:17:00 -07:00
committed by Teknium
parent be2f739e9a
commit ffcd9d7ac7
2 changed files with 734 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_yuanbao_pipeline.py
+463 -108
View File
@@ -38,6 +38,9 @@ from gateway.platforms.yuanbao import (
OwnerCommandMiddleware,
BuildSourceMiddleware,
GroupAtGuardMiddleware,
QuoteContextMiddleware,
MediaResolveMiddleware,
PatchAnchorsMiddleware,
DispatchMiddleware,
InboundPipelineBuilder,
YuanbaoAdapter,
@@ -308,40 +311,7 @@ class TestInboundPipeline:
# ============================================================
# 2. InboundContext Tests
# ============================================================
class TestInboundContext:
def test_default_values(self):
"""InboundContext has sensible defaults."""
adapter = make_adapter()
ctx = InboundContext(adapter=adapter)
assert ctx.raw_frames == []
assert ctx.push is None
assert ctx.decoded_via == ""
assert ctx.from_account == ""
assert ctx.group_code == ""
assert ctx.msg_body == []
assert ctx.msg_id == ""
assert ctx.chat_id == ""
assert ctx.chat_type == ""
assert ctx.raw_text == ""
assert ctx.media_refs == []
assert ctx.owner_command is None
assert ctx.source is None
assert ctx.msg_type is None
def test_mutable_fields(self):
"""InboundContext fields are mutable."""
ctx = make_ctx()
ctx.from_account = "alice"
ctx.chat_type = "dm"
assert ctx.from_account == "alice"
assert ctx.chat_type == "dm"
# ============================================================
# 3. Individual Middleware Tests
# 2. Individual Middleware Tests
# ============================================================
class TestDecodeMiddleware:
@@ -720,31 +690,25 @@ class TestCreateInboundPipeline:
expected = [
"decode",
"extract-fields",
"recall_guard",
"dedup",
"skip-self",
"chat-routing",
"access-guard",
"auto-sethome",
"extract-content",
"placeholder-filter",
"owner-command",
"build-source",
"group-at-guard",
"group-attribution",
"classify-msg-type",
"quote-context",
"media-resolve",
"patch-anchors",
"dispatch",
]
"""Pipeline can be customized after creation."""
pipeline = InboundPipelineBuilder.build()
async def custom_mw(ctx, next_fn):
await next_fn()
pipeline.use_before("dispatch", "custom", custom_mw)
assert "custom" in pipeline.middleware_names
idx_custom = pipeline.middleware_names.index("custom")
idx_dispatch = pipeline.middleware_names.index("dispatch")
assert idx_custom < idx_dispatch
assert pipeline.middleware_names == expected
# ============================================================
@@ -861,19 +825,7 @@ if __name__ == "__main__":
# ============================================================
class TestInboundMiddlewareABC:
"""Test the InboundMiddleware abstract base class."""
def test_cannot_instantiate_abc(self):
"""InboundMiddleware cannot be instantiated directly."""
with pytest.raises(TypeError):
InboundMiddleware()
def test_subclass_must_implement_handle(self):
"""Subclass without handle() raises TypeError."""
with pytest.raises(TypeError):
class BadMiddleware(InboundMiddleware):
name = "bad"
BadMiddleware()
"""Test the InboundMiddleware OOP protocol (callable + named)."""
def test_subclass_with_handle_works(self):
"""Subclass with handle() can be instantiated."""
@@ -900,19 +852,14 @@ class TestInboundMiddlewareABC:
assert ctx.raw_text == "called"
next_fn.assert_awaited_once()
def test_repr(self):
"""Middleware has a useful repr."""
class MyMW(InboundMiddleware):
name = "my-mw"
async def handle(self, ctx, next_fn):
pass
mw = MyMW()
assert "MyMW" in repr(mw)
assert "my-mw" in repr(mw)
class TestMiddlewareClasses:
"""Test that all concrete middleware classes have correct names and are InboundMiddleware subclasses."""
"""Pin the canonical ``name`` of each concrete middleware class.
These names are referenced by ``InboundPipelineBuilder.build()`` ordering,
by ``use_before`` / ``use_after`` insertion in extensions, and by log
messages — so they're a real downstream contract worth pinning.
"""
MIDDLEWARE_CLASSES = [
(DecodeMiddleware, "decode"),
@@ -929,23 +876,12 @@ class TestMiddlewareClasses:
(DispatchMiddleware, "dispatch"),
]
@pytest.mark.parametrize("cls,expected_name", MIDDLEWARE_CLASSES)
def test_is_inbound_middleware(self, cls, expected_name):
"""Each middleware class is a subclass of InboundMiddleware."""
assert issubclass(cls, InboundMiddleware)
@pytest.mark.parametrize("cls,expected_name", MIDDLEWARE_CLASSES)
def test_has_correct_name(self, cls, expected_name):
"""Each middleware class has the expected name."""
mw = cls()
assert mw.name == expected_name
@pytest.mark.parametrize("cls,expected_name", MIDDLEWARE_CLASSES)
def test_is_callable(self, cls, expected_name):
"""Each middleware instance is callable."""
mw = cls()
assert callable(mw)
class TestPipelineOOPRegistration:
"""Test that InboundPipeline works with OOP middleware instances."""
@@ -991,38 +927,457 @@ class TestPipelineOOPRegistration:
await pipeline.execute(make_ctx())
assert order == ["oop", "func"]
def test_use_before_with_middleware_instance(self):
"""use_before works with OOP middleware instances."""
class MwA(InboundMiddleware):
name = "a"
async def handle(self, ctx, next_fn): await next_fn()
class MwB(InboundMiddleware):
name = "b"
async def handle(self, ctx, next_fn): await next_fn()
# ============================================================
# QuoteContextMiddleware Tests
# ============================================================
#
# Quote-media resolution used to depend on a process-local
# msg_id→resids cache populated by ExtractContentMiddleware. After #27866
# made gateway/run.py write @bot user transcript entries with
# message_id (symmetric with the observed-group writer at yuanbao.py:2091),
# QuoteContextMiddleware's transcript-lookup path covers every quote case
# we used to rely on the cache for, so the cache (and those tests) were
# removed. ``_extract_quote_context()`` is now a pure (quote_id, quote_text)
# extractor; quote media references are populated separately by
# ``_extract_media_refs_from_transcript()`` against the transcript store.
class MwC(InboundMiddleware):
name = "c"
async def handle(self, ctx, next_fn): await next_fn()
class TestQuoteContextMiddleware:
"""Tests for QuoteContextMiddleware._extract_quote_context."""
pipeline = InboundPipeline().use(MwA()).use(MwC())
pipeline.use_before("c", MwB())
assert pipeline.middleware_names == ["a", "b", "c"]
def test_extract_quote_context_no_cloud_data(self):
"""Returns (None, None) when cloud_custom_data is empty."""
result = QuoteContextMiddleware()._extract_quote_context("")
assert result == (None, None)
def test_use_after_with_middleware_instance(self):
"""use_after works with OOP middleware instances."""
class MwA(InboundMiddleware):
name = "a"
async def handle(self, ctx, next_fn): await next_fn()
def test_extract_quote_context_no_quote_key(self):
"""Returns (None, None) when JSON has no 'quote' key."""
cloud_data = json.dumps({"foo": "bar"})
result = QuoteContextMiddleware()._extract_quote_context(cloud_data)
assert result == (None, None)
class MwB(InboundMiddleware):
name = "b"
async def handle(self, ctx, next_fn): await next_fn()
def test_extract_quote_context_with_desc(self):
"""Extracts quote_id and quote_text from desc."""
cloud_data = json.dumps({
"quote": {
"id": "quoted-msg-001",
"desc": "Hello world",
"sender_nickname": "Alice",
}
})
quote_id, quote_text = QuoteContextMiddleware()._extract_quote_context(cloud_data)
assert quote_id == "quoted-msg-001"
assert quote_text == "Alice: Hello world"
class MwC(InboundMiddleware):
name = "c"
async def handle(self, ctx, next_fn): await next_fn()
def test_extract_quote_context_empty_desc(self):
"""When desc is empty, quote_text is None but quote_id is preserved."""
cloud_data = json.dumps({
"quote": {
"id": "quoted-msg-003",
"desc": "",
"sender_nickname": "Carol",
}
})
quote_id, quote_text = QuoteContextMiddleware()._extract_quote_context(cloud_data)
assert quote_id == "quoted-msg-003"
assert quote_text is None
pipeline = InboundPipeline().use(MwA()).use(MwC())
pipeline.use_after("a", MwB())
assert pipeline.middleware_names == ["a", "b", "c"]
def test_extract_quote_context_no_quote_id(self):
"""When quote.id is empty, quote_id is None."""
cloud_data = json.dumps({
"quote": {
"id": "",
"desc": "some text",
}
})
quote_id, _quote_text = QuoteContextMiddleware()._extract_quote_context(cloud_data)
assert quote_id is None
@pytest.mark.asyncio
async def test_handle_sets_ctx_fields(self):
"""QuoteContextMiddleware.handle() sets ctx.reply_to_message_id, reply_to_text, quote_media_refs.
With no transcript store wired up, quote_media_refs falls back to []
— media resolution from transcript is covered by separate tests.
"""
cloud_data = json.dumps({
"quote": {
"id": "quoted-msg-004",
"desc": "Check this image",
"sender_nickname": "Dave",
}
})
adapter = make_adapter()
adapter._session_store = None # no transcript lookup path
ctx = make_ctx(adapter=adapter, cloud_custom_data=cloud_data)
next_fn = AsyncMock()
await QuoteContextMiddleware()(ctx, next_fn)
assert ctx.reply_to_message_id == "quoted-msg-004"
assert ctx.reply_to_text == "Dave: Check this image"
assert ctx.quote_media_refs == []
next_fn.assert_awaited_once()
# ============================================================
# MediaResolveMiddleware Tests
# ============================================================
#
# After the dispatch refactor, MediaResolveMiddleware is the single entry
# point for all inbound media downloads. It merges up to three sources
# into ``ctx.media_urls`` / ``ctx.media_types`` (deduped, in this order):
#
# 1) media carried by the current message itself (always),
# 2) quote_media_refs (when reply_to_message_id is set),
# 3) recent group-observed media (only when chat_type == "group" and
# no quote is present).
#
# Direct messages skip the observed backfill entirely.
class TestResolveYbresRefs:
"""Direct tests for ``MediaResolveMiddleware._resolve_ybres_refs``.
This classmethod is the shared engine for both ``_resolve_quote_media``
and ``_collect_observed_media``. Patching the two upstream callers from
routing tests doesn't exercise its filtering / error-swallowing
behavior, so we pin those contracts directly here.
"""
@pytest.mark.asyncio
async def test_resolves_each_ref_in_order(self):
"""Successful resolution returns ``(paths, mimes)`` aligned with input order."""
adapter = make_adapter()
refs = [
("rid-1", "image", "a.jpg"),
("rid-2", "file", "doc.pdf"),
]
with patch.object(
MediaResolveMiddleware, "_fetch_resource_url",
new=AsyncMock(side_effect=["https://fresh/1", "https://fresh/2"]),
) as p_fetch, patch.object(
MediaResolveMiddleware, "_download_and_cache",
new=AsyncMock(side_effect=[
("/cache/a.jpg", "image/jpeg"),
("/cache/doc.pdf", "application/pdf"),
]),
) as p_cache:
paths, mimes = await MediaResolveMiddleware._resolve_ybres_refs(
adapter, refs, log_prefix="test",
)
assert paths == ["/cache/a.jpg", "/cache/doc.pdf"]
assert mimes == ["image/jpeg", "application/pdf"]
assert p_fetch.await_count == 2
# filename from the ref tuple is forwarded to download_and_cache
cache_kwargs = [c.kwargs for c in p_cache.await_args_list]
assert cache_kwargs[0]["file_name"] == "a.jpg"
assert cache_kwargs[0]["kind"] == "image"
assert cache_kwargs[0]["resource_id"] == "rid-1"
assert cache_kwargs[1]["file_name"] == "doc.pdf"
@pytest.mark.asyncio
async def test_skips_unresolvable_kinds(self):
"""Refs whose kind is outside ``_RESOLVABLE_MEDIA_KINDS`` are dropped silently."""
adapter = make_adapter()
refs = [
("rid-v", "video", ""), # not resolvable
("rid-i", "image", "ok.jpg"), # resolvable
("rid-?", "unknown", ""), # not resolvable
]
with patch.object(
MediaResolveMiddleware, "_fetch_resource_url",
new=AsyncMock(return_value="https://fresh/i"),
) as p_fetch, patch.object(
MediaResolveMiddleware, "_download_and_cache",
new=AsyncMock(return_value=("/cache/ok.jpg", "image/jpeg")),
) as p_cache:
paths, mimes = await MediaResolveMiddleware._resolve_ybres_refs(
adapter, refs, log_prefix="test",
)
assert paths == ["/cache/ok.jpg"]
assert mimes == ["image/jpeg"]
# Only the resolvable ref hit the network.
p_fetch.assert_awaited_once()
p_cache.assert_awaited_once()
@pytest.mark.asyncio
async def test_fetch_failure_is_swallowed_per_ref(self):
"""If ``_fetch_resource_url`` raises, that ref is skipped — not the whole batch."""
adapter = make_adapter()
refs = [
("rid-bad", "image", ""),
("rid-ok", "image", ""),
]
with patch.object(
MediaResolveMiddleware, "_fetch_resource_url",
new=AsyncMock(side_effect=[RuntimeError("boom"), "https://fresh/ok"]),
), patch.object(
MediaResolveMiddleware, "_download_and_cache",
new=AsyncMock(return_value=("/cache/ok.jpg", "image/jpeg")),
) as p_cache:
paths, mimes = await MediaResolveMiddleware._resolve_ybres_refs(
adapter, refs, log_prefix="test",
)
# bad ref dropped; good ref preserved
assert paths == ["/cache/ok.jpg"]
assert mimes == ["image/jpeg"]
# download_and_cache was only invoked for the surviving ref
p_cache.assert_awaited_once()
@pytest.mark.asyncio
async def test_cache_miss_drops_ref(self):
"""If ``_download_and_cache`` returns None, the ref is dropped."""
adapter = make_adapter()
refs = [("rid-1", "image", "")]
with patch.object(
MediaResolveMiddleware, "_fetch_resource_url",
new=AsyncMock(return_value="https://fresh/1"),
), patch.object(
MediaResolveMiddleware, "_download_and_cache",
new=AsyncMock(return_value=None),
):
paths, mimes = await MediaResolveMiddleware._resolve_ybres_refs(
adapter, refs, log_prefix="test",
)
assert paths == []
assert mimes == []
class TestMediaResolveMiddlewareRouting:
"""Branch-routing tests for MediaResolveMiddleware.handle()."""
def _make_resolved_ctx(self, *, chat_type: str, reply_to: str = None,
quote_media_refs=None, raw_text: str = "hello"):
adapter = make_adapter()
ctx = make_ctx(
adapter=adapter,
chat_type=chat_type,
reply_to_message_id=reply_to,
quote_media_refs=list(quote_media_refs or []),
raw_text=raw_text,
media_refs=[], # no own attachments by default
)
return adapter, ctx
@pytest.mark.asyncio
async def test_dm_no_quote_skips_observed_backfill(self):
"""In dm chats, observed-media backfill is never invoked."""
_adapter, ctx = self._make_resolved_ctx(chat_type="dm")
with patch.object(
MediaResolveMiddleware, "_resolve_media_urls",
new=AsyncMock(return_value=([], [])),
) as p_own, patch.object(
MediaResolveMiddleware, "_resolve_quote_media",
new=AsyncMock(return_value=([], [])),
) as p_quote, patch.object(
MediaResolveMiddleware, "_collect_observed_media",
new=AsyncMock(return_value=([], [])),
) as p_observed:
next_fn = AsyncMock()
await MediaResolveMiddleware()(ctx, next_fn)
p_own.assert_awaited_once()
p_quote.assert_not_awaited()
p_observed.assert_not_awaited()
next_fn.assert_awaited_once()
@pytest.mark.asyncio
async def test_dm_with_quote_resolves_quote_media_only(self):
"""In dm chats with a quote, only quote media (plus own) is resolved."""
_adapter, ctx = self._make_resolved_ctx(
chat_type="dm",
reply_to="quoted-001",
quote_media_refs=[("rid-q1", "image", "")],
)
with patch.object(
MediaResolveMiddleware, "_resolve_media_urls",
new=AsyncMock(return_value=([], [])),
) as p_own, patch.object(
MediaResolveMiddleware, "_resolve_quote_media",
new=AsyncMock(return_value=(["/cache/q1.jpg"], ["image/jpeg"])),
) as p_quote, patch.object(
MediaResolveMiddleware, "_collect_observed_media",
new=AsyncMock(return_value=([], [])),
) as p_observed:
next_fn = AsyncMock()
await MediaResolveMiddleware()(ctx, next_fn)
p_own.assert_awaited_once()
p_quote.assert_awaited_once()
p_observed.assert_not_awaited()
assert ctx.media_urls == ["/cache/q1.jpg"]
assert ctx.media_types == ["image/jpeg"]
@pytest.mark.asyncio
async def test_group_no_quote_runs_observed_backfill(self):
"""In group chats without quote, observed-media backfill is invoked."""
_adapter, ctx = self._make_resolved_ctx(chat_type="group")
with patch.object(
MediaResolveMiddleware, "_resolve_media_urls",
new=AsyncMock(return_value=([], [])),
), patch.object(
MediaResolveMiddleware, "_resolve_quote_media",
new=AsyncMock(return_value=([], [])),
) as p_quote, patch.object(
MediaResolveMiddleware, "_collect_observed_media",
new=AsyncMock(return_value=(["/cache/o1.jpg"], ["image/jpeg"])),
) as p_observed:
next_fn = AsyncMock()
await MediaResolveMiddleware()(ctx, next_fn)
p_quote.assert_not_awaited()
p_observed.assert_awaited_once()
assert ctx.media_urls == ["/cache/o1.jpg"]
@pytest.mark.asyncio
async def test_group_with_quote_skips_observed_backfill(self):
"""In group chats with a quote, only quote media is resolved (no backfill)."""
_adapter, ctx = self._make_resolved_ctx(
chat_type="group",
reply_to="quoted-002",
quote_media_refs=[("rid-q2", "image", "")],
)
with patch.object(
MediaResolveMiddleware, "_resolve_media_urls",
new=AsyncMock(return_value=([], [])),
), patch.object(
MediaResolveMiddleware, "_resolve_quote_media",
new=AsyncMock(return_value=(["/cache/q2.jpg"], ["image/jpeg"])),
) as p_quote, patch.object(
MediaResolveMiddleware, "_collect_observed_media",
new=AsyncMock(return_value=(["/cache/o2.jpg"], ["image/jpeg"])),
) as p_observed:
next_fn = AsyncMock()
await MediaResolveMiddleware()(ctx, next_fn)
p_quote.assert_awaited_once()
p_observed.assert_not_awaited()
assert ctx.media_urls == ["/cache/q2.jpg"]
@pytest.mark.asyncio
async def test_merges_and_dedupes_three_sources(self):
"""Own + quote sources are merged with dedup applied."""
_adapter, ctx = self._make_resolved_ctx(
chat_type="dm",
reply_to="quoted-003",
quote_media_refs=[("rid", "image", "")],
)
with patch.object(
MediaResolveMiddleware, "_resolve_media_urls",
new=AsyncMock(return_value=(["/cache/a.jpg"], ["image/jpeg"])),
), patch.object(
MediaResolveMiddleware, "_resolve_quote_media",
# Same path as own → must be deduped.
new=AsyncMock(return_value=(["/cache/a.jpg", "/cache/b.jpg"],
["image/jpeg", "image/png"])),
):
next_fn = AsyncMock()
await MediaResolveMiddleware()(ctx, next_fn)
assert ctx.media_urls == ["/cache/a.jpg", "/cache/b.jpg"]
assert ctx.media_types == ["image/jpeg", "image/png"]
@pytest.mark.asyncio
async def test_placeholder_recheck_uses_own_count_only(self):
"""Placeholder retry-skip uses ``own_count``, not the merged total.
A bare placeholder text (e.g. ``[image]``) accompanied only by a
quote-resolved image is still skippable — quote media must not
flip a placeholder into a non-placeholder.
"""
_adapter, ctx = self._make_resolved_ctx(
chat_type="dm",
reply_to="quoted-004",
quote_media_refs=[("rid", "image", "")],
raw_text="[image]",
)
with patch.object(
MediaResolveMiddleware, "_resolve_media_urls",
new=AsyncMock(return_value=([], [])), # no own media
), patch.object(
MediaResolveMiddleware, "_resolve_quote_media",
new=AsyncMock(return_value=(["/cache/q.jpg"], ["image/jpeg"])),
), patch.object(
PlaceholderFilterMiddleware, "is_skippable_placeholder",
return_value=True,
) as p_check:
next_fn = AsyncMock()
await MediaResolveMiddleware()(ctx, next_fn)
# Pipeline short-circuited despite quote media being present.
next_fn.assert_not_awaited()
# And the second-pass check was called with own_count == 0.
p_check.assert_called_once()
_text_arg, count_arg = p_check.call_args.args
assert count_arg == 0
# ============================================================
# PatchAnchorsMiddleware Tests
# ============================================================
class TestPatchAnchorsMiddleware:
"""Tests for PatchAnchorsMiddleware._patch()."""
def test_no_op_when_text_or_urls_empty(self):
assert PatchAnchorsMiddleware._patch("", [], []) == ""
assert PatchAnchorsMiddleware._patch("hello", [], []) == "hello"
def test_replaces_image_anchor_with_local_path(self):
text = "look [image|ybres:abc] please"
out = PatchAnchorsMiddleware._patch(
text, ["/cache/x.jpg"], ["image/jpeg"],
)
assert out == "look [image: /cache/x.jpg] please"
def test_replaces_file_anchor_with_filename_label(self):
text = "see [file:doc.pdf|ybres:rid-1]"
out = PatchAnchorsMiddleware._patch(
text, ["/cache/doc.pdf"], ["application/pdf"],
)
assert "[file: doc.pdf → /cache/doc.pdf]" in out
def test_skips_non_local_paths(self):
"""URLs not starting with '/' are left untouched."""
text = "[image|ybres:abc]"
out = PatchAnchorsMiddleware._patch(
text, ["https://example.com/x.jpg"], ["image/jpeg"],
)
# Anchor preserved verbatim because the resolved url is remote.
assert out == text
def test_anchor_kind_image_requires_image_mime(self):
"""An [image|...] anchor with a non-image mime is left alone."""
text = "[image|ybres:rid]"
out = PatchAnchorsMiddleware._patch(
text, ["/cache/odd.bin"], ["application/octet-stream"],
)
assert out == text
@pytest.mark.asyncio
async def test_handle_writes_back_to_ctx(self):
adapter = make_adapter()
ctx = make_ctx(
adapter=adapter,
raw_text="hi [image|ybres:rid]",
media_urls=["/cache/y.png"],
media_types=["image/png"],
)
next_fn = AsyncMock()
await PatchAnchorsMiddleware()(ctx, next_fn)
assert ctx.raw_text == "hi [image: /cache/y.png]"
next_fn.assert_awaited_once()