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

View File

@@ -12,7 +12,7 @@ from app.agent.tools.base import MoviePilotTool
from app.api.endpoints.openai import _OpenAIStreamingHandler
from app.core.config import settings
from app.schemas.message import MessageResponse
from app.schemas.types import MessageChannel
from app.schemas.types import MessageChannel, NotificationType
class DummyTool(MoviePilotTool):
@@ -159,6 +159,10 @@ class TestAgentToolStreaming(unittest.TestCase):
self.assertEqual(
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)
def test_flush_edits_message_via_threadpool(self):
@@ -188,6 +192,38 @@ class TestAgentToolStreaming(unittest.TestCase):
)
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):
handler = StreamingHandler()
handler._streaming_enabled = True

View File

@@ -21,7 +21,7 @@ if "Pinyin2Hanzi" not in sys.modules:
from app.modules.feishu import FeishuModule
from app.modules.feishu.feishu import Feishu
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
@@ -586,15 +586,30 @@ class TestFeishu(unittest.TestCase):
self.assertEqual(reaction_id, "reaction_2")
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._channel = MessageChannel.Feishu
client = MagicMock()
client.close_streaming_card.return_value = True
with patch.object(module, "get_config", return_value=SimpleNamespace(name="feishu-main")), patch.object(
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)
client.close_streaming_card.assert_called_once_with(card_id="card_stream", sequence=4)