fix(agent): route streaming finalization through channel modules

This commit is contained in:
jxxghp
2026-05-13 00:50:53 +08:00
parent d86d24fc4f
commit 77aa65bfdc
5 changed files with 82 additions and 28 deletions
+8 -17
View File
@@ -12,7 +12,7 @@ from app.schemas.message import (
ChannelCapabilityManager, ChannelCapabilityManager,
ChannelCapability, ChannelCapability,
) )
from app.schemas.types import MessageChannel from app.schemas.types import MessageChannel, NotificationType
class _StreamChain(ChainBase): class _StreamChain(ChainBase):
@@ -212,22 +212,11 @@ class StreamingHandler:
await self._flush() await self._flush()
message_response = self._message_response message_response = self._message_response
if ( if message_response:
message_response await run_in_threadpool(
and message_response.channel == MessageChannel.Feishu _StreamChain().finalize_message,
and isinstance(message_response.metadata, dict) message_response,
): )
stream_meta = message_response.metadata.get("feishu_streaming") or {}
card_id = str(stream_meta.get("card_id") or "").strip()
sequence = int(stream_meta.get("sequence") or 1) + 1
if card_id:
await run_in_threadpool(
_StreamChain().run_module,
"close_feishu_streaming_card",
card_id=card_id,
sequence=sequence,
source=message_response.source,
)
# 检查是否所有缓冲内容都已发送 # 检查是否所有缓冲内容都已发送
with self._lock: with self._lock:
@@ -480,6 +469,7 @@ class StreamingHandler:
Notification( Notification(
channel=self._channel, channel=self._channel,
source=self._source, source=self._source,
mtype=NotificationType.Agent,
userid=self._user_id, userid=self._user_id,
username=self._username, username=self._username,
title=self._title, title=self._title,
@@ -522,6 +512,7 @@ class StreamingHandler:
Notification( Notification(
channel=self._channel, channel=self._channel,
source=self._source, source=self._source,
mtype=NotificationType.Agent,
userid=self._user_id, userid=self._user_id,
username=self._username, username=self._username,
title=self._title, title=self._title,
+10
View File
@@ -1585,6 +1585,16 @@ class ChainBase(metaclass=ABCMeta):
message=self._normalize_notification_for_dispatch(message), message=self._normalize_notification_for_dispatch(message),
) )
def finalize_message(
self,
response: MessageResponse,
) -> bool:
"""
对已发送消息执行渠道收尾动作。
例如关闭流式卡片状态;无特殊收尾的渠道直接返回 False。
"""
return self.run_module("finalize_message", response=response)
def metadata_img( def metadata_img(
self, self,
mediainfo: MediaInfo, mediainfo: MediaInfo,
+9 -7
View File
@@ -305,16 +305,18 @@ class FeishuModule(_ModuleBase, _MessageBase[Feishu]):
return False return False
return client.delete_message_reaction(message_id=message_id, reaction_id=reaction_id) return client.delete_message_reaction(message_id=message_id, reaction_id=reaction_id)
def close_feishu_streaming_card( def finalize_message(self, response: MessageResponse) -> bool:
self, if response.channel != self._channel or not isinstance(response.metadata, dict):
card_id: str, return False
sequence: int, stream_meta = response.metadata.get("feishu_streaming") or {}
source: str, card_id = str(stream_meta.get("card_id") or "").strip()
) -> bool: if not card_id:
client_config = self.get_config(source) return False
client_config = self.get_config(response.source)
if not client_config: if not client_config:
return False return False
client = self.get_instance(client_config.name) client = self.get_instance(client_config.name)
if not client: if not client:
return False return False
sequence = int(stream_meta.get("sequence") or 1) + 1
return client.close_streaming_card(card_id=card_id, sequence=sequence) return client.close_streaming_card(card_id=card_id, sequence=sequence)
+37 -1
View File
@@ -12,7 +12,7 @@ from app.agent.tools.base import MoviePilotTool
from app.api.endpoints.openai import _OpenAIStreamingHandler from app.api.endpoints.openai import _OpenAIStreamingHandler
from app.core.config import settings from app.core.config import settings
from app.schemas.message import MessageResponse from app.schemas.message import MessageResponse
from app.schemas.types import MessageChannel from app.schemas.types import MessageChannel, NotificationType
class DummyTool(MoviePilotTool): class DummyTool(MoviePilotTool):
@@ -159,6 +159,10 @@ class TestAgentToolStreaming(unittest.TestCase):
self.assertEqual( self.assertEqual(
run_in_threadpool_mock.await_args.args[0].__name__, "send_direct_message" run_in_threadpool_mock.await_args.args[0].__name__, "send_direct_message"
) )
self.assertEqual(
run_in_threadpool_mock.await_args.args[1].mtype,
NotificationType.Agent,
)
self.assertTrue(handler.has_sent_message) self.assertTrue(handler.has_sent_message)
def test_flush_edits_message_via_threadpool(self): def test_flush_edits_message_via_threadpool(self):
@@ -188,6 +192,38 @@ class TestAgentToolStreaming(unittest.TestCase):
) )
self.assertEqual(handler._sent_text, "hello world") self.assertEqual(handler._sent_text, "hello world")
def test_stop_streaming_uses_generic_finalize_message(self):
handler = StreamingHandler()
handler._message_response = MessageResponse(
message_id="om_stream",
chat_id="oc_stream",
channel=MessageChannel.Feishu,
source="feishu-main",
metadata={"feishu_streaming": {"card_id": "card_stream", "sequence": 2}},
success=True,
)
handler._sent_text = "hello"
handler._buffer = "hello"
handler._streaming_enabled = True
with patch(
"app.agent.callback.run_in_threadpool", new_callable=AsyncMock
) as run_in_threadpool_mock, patch.object(
handler, "_cancel_flush_task", new_callable=AsyncMock
), patch.object(
handler, "_flush", new_callable=AsyncMock
):
asyncio.run(handler.stop_streaming())
self.assertEqual(run_in_threadpool_mock.await_count, 1)
self.assertEqual(
run_in_threadpool_mock.await_args.args[0].__name__, "finalize_message"
)
self.assertEqual(
run_in_threadpool_mock.await_args.args[1].message_id,
"om_stream",
)
def test_flush_without_channel_context_does_not_send_direct_message(self): def test_flush_without_channel_context_does_not_send_direct_message(self):
handler = StreamingHandler() handler = StreamingHandler()
handler._streaming_enabled = True handler._streaming_enabled = True
+18 -3
View File
@@ -21,7 +21,7 @@ if "Pinyin2Hanzi" not in sys.modules:
from app.modules.feishu import FeishuModule from app.modules.feishu import FeishuModule
from app.modules.feishu.feishu import Feishu from app.modules.feishu.feishu import Feishu
from app.schemas import Notification from app.schemas import Notification
from app.schemas.message import ChannelCapability, ChannelCapabilityManager from app.schemas.message import ChannelCapability, ChannelCapabilityManager, MessageResponse
from app.schemas.types import MessageChannel, NotificationType from app.schemas.types import MessageChannel, NotificationType
@@ -586,15 +586,30 @@ class TestFeishu(unittest.TestCase):
self.assertEqual(reaction_id, "reaction_2") self.assertEqual(reaction_id, "reaction_2")
self.assertTrue(deleted) self.assertTrue(deleted)
def test_module_close_streaming_card_delegates_to_client(self): def test_module_finalize_message_closes_streaming_card(self):
module = FeishuModule() module = FeishuModule()
module._channel = MessageChannel.Feishu
client = MagicMock() client = MagicMock()
client.close_streaming_card.return_value = True client.close_streaming_card.return_value = True
with patch.object(module, "get_config", return_value=SimpleNamespace(name="feishu-main")), patch.object( with patch.object(module, "get_config", return_value=SimpleNamespace(name="feishu-main")), patch.object(
module, "get_instance", return_value=client module, "get_instance", return_value=client
): ):
success = module.close_feishu_streaming_card("card_stream", 4, "feishu-main") success = module.finalize_message(
MessageResponse(
message_id="om_stream",
chat_id="oc_stream",
channel=MessageChannel.Feishu,
source="feishu-main",
metadata={
"feishu_streaming": {
"card_id": "card_stream",
"sequence": 3,
}
},
success=True,
)
)
self.assertTrue(success) self.assertTrue(success)
client.close_streaming_card.assert_called_once_with(card_id="card_stream", sequence=4) client.close_streaming_card.assert_called_once_with(card_id="card_stream", sequence=4)