Merge pull request 'fix: update message handling for resend functionality' (#410) from fix-abort-then-resend into development

Reviewed-on: https://git.biggo.com/Funmula/dive-mcp-host/pulls/410

Author: LouisFunmula

dive_mcp_host/host/chat.py

@@ -115,13 +115,16 @@ async def _run_in_context(self) -> AsyncGenerator[Self, None]:
         finally:
             self._agent = None
 
-    async def _get_updates_for_resend(
+    async def _remove_messages_for_resend(
         self,
         resend: list[BaseMessage],
         update: list[BaseMessage],
-    ) -> list[BaseMessage]:
-        if not self._checkpointer:
-            return update
+    ) -> None:
+        if not self._checkpointer or not (
+            self.active_agent.checkpointer
+            and isinstance(self.active_agent.checkpointer, BaseCheckpointSaver)
+        ):
+            return
         resend_map = {msg.id: msg for msg in resend}
         to_update = [i for i in update if i.id not in resend_map]
         if state := await self.active_agent.aget_state(
@@ -134,15 +137,25 @@ async def _get_updates_for_resend(
         ):
             drop_after = False
             if not state.values:
-                return to_update
+                return
 
             for msg in cast(MessagesState, state.values)["messages"]:
                 assert msg.id is not None  # all messages from the agent have an ID
                 if msg.id in resend_map:
                     drop_after = True
                 elif drop_after:
                     to_update.append(RemoveMessage(msg.id))
-            return to_update
+            if to_update:
+                await self.active_agent.aupdate_state(
+                    RunnableConfig(
+                        configurable={
+                            "thread_id": self._chat_id,
+                            "user_id": self._user_id,
+                        },
+                    ),
+                    {"messages": to_update},
+                )
+            return
         raise ThreadNotFoundError(self._chat_id)
 
     def query(
@@ -180,9 +193,7 @@ async def _stream_response() -> AsyncGenerator[dict[str, Any] | Any, None]:
                     isinstance(msg, BaseMessage) and msg.id for msg in query_msgs
                 ):
                     raise MessageTypeError("Resending messages must has an ID")
-                query_msgs += await self._get_updates_for_resend(
-                    query_msgs, modify or []
-                )
+                await self._remove_messages_for_resend(query_msgs, modify or [])
             elif modify:
                 query_msgs = [*query_msgs, *modify]
             signal = asyncio.Event()

tests/test_host.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import os
+import time
 from contextlib import AbstractAsyncContextManager
 from typing import Any, cast
 from unittest import mock
@@ -15,8 +16,10 @@
     ToolCall,
     ToolMessage,
 )
+from langgraph.checkpoint.memory import InMemorySaver
 from pydantic import AnyUrl, SecretStr
 
+from dive_mcp_host.host.chat import Chat
 from dive_mcp_host.host.conf import CheckpointerConfig, HostConfig
 from dive_mcp_host.host.conf.llm import LLMConfig
 from dive_mcp_host.host.custom_events import ToolCallProgress
@@ -561,3 +564,99 @@ async def test_custom_event(
                     assert isinstance(i[1][1], ToolCallProgress)
                     done = True
         assert done
+
+
+@pytest.mark.asyncio
+async def test_resend_after_abort(  # noqa: C901
+    echo_tool_stdio_config: dict[str, ServerConfig],
+) -> None:
+    """Test that resend after abort works."""
+    config = HostConfig(
+        llm=LLMConfig(
+            model="fake",
+            model_provider="dive",
+        ),
+        mcp_servers=echo_tool_stdio_config,
+    )
+
+    async with DiveMcpHost(config) as mcp_host:
+        mcp_host._checkpointer = InMemorySaver()
+        fake_responses = [
+            AIMessage(
+                content="Call echo tool",
+                tool_calls=[
+                    ToolCall(
+                        name="echo",
+                        args={"message": "Hello, world!", "delay_ms": 2000},
+                        id="phase1-tool-1",
+                        type="tool_call",
+                    ),
+                ],
+                id="phase1-1",
+            ),
+            AIMessage(
+                content="Bye",
+                id="phase1-2",
+            ),
+        ]
+
+        ts = time.time()
+        got_msg = False
+
+        async def _abort_task(chat: Chat) -> None:
+            while True:
+                if got_msg and time.time() - ts > 0.5:
+                    chat.abort()
+                    break
+                await asyncio.sleep(0.1)
+
+        fake_model = cast("FakeMessageToolModel", mcp_host.model)
+        fake_model.responses = fake_responses
+        await mcp_host.tools_initialized_event.wait()
+        chat = mcp_host.chat(chat_id="chat_id")
+        async with chat:
+            task = asyncio.create_task(_abort_task(chat))
+            async for r in chat.query(
+                HumanMessage(content="Hello, world!", id="H1"),
+                stream_mode=["messages", "values", "updates"],
+            ):
+                if r[0] == "messages":  # type: ignore
+                    got_msg = True
+                    ts = time.time()
+        await task
+
+        resend = [HumanMessage(content="Resend message!", id="H1")]
+        fake_responses = [
+            AIMessage(
+                content="Call echo tool",
+                tool_calls=[
+                    ToolCall(
+                        name="echo",
+                        args={"message": "Hello, world!"},
+                        id="phase2-tool-1",
+                        type="tool_call",
+                    ),
+                ],
+                id="phase2-1",
+            ),
+            AIMessage(
+                content="Bye",
+                id="phase2-2",
+            ),
+        ]
+        fake_model.i = 0
+        fake_model.responses = fake_responses
+        chat = mcp_host.chat(chat_id="chat_id")
+        async with chat:
+            async for r in chat.query(
+                resend,  # type: ignore
+                is_resend=True,
+            ):
+                r = cast(tuple[str, list[BaseMessage]], r)
+                if r[0] == "messages":
+                    for m in r[1]:
+                        assert m.id
+                        if isinstance(m, HumanMessage):
+                            assert m.id == "H1"
+                        elif isinstance(m, AIMessage):
+                            assert m.id.startswith("phase2-")