fix: preserve checkpoints after failed LLM requests

astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py

@@ -461,7 +461,33 @@ async def _save_to_history(
         if not req or not req.conversation:
             return
 
-        if not llm_response and not user_aborted:
+        messages_to_save: list[Message] = []
+        skipped_initial_system = False
+        for message in all_messages:
+            if message.role == "system" and not skipped_initial_system:
+                skipped_initial_system = True
+                continue
+            if message.role in ["assistant", "user"] and message._no_save:
+                continue
+            messages_to_save.append(message)
+
+        checkpoint_id = event.get_extra("llm_checkpoint_id")
+        message_to_save = dump_messages_with_checkpoints(messages_to_save)
+        if not user_aborted and (
+            llm_response is None or llm_response.role != "assistant"
+        ):
+            if isinstance(checkpoint_id, str) and checkpoint_id:
+                message_to_save.append(
+                    CheckpointMessageSegment(
+                        content=CheckpointData(id=checkpoint_id),
+                    ).model_dump()
+                )
+                await self.conv_manager.update_conversation(
+                    event.unified_msg_origin,
+                    req.conversation.cid,
+                    history=message_to_save,
+                    token_usage=None,
+                )
             return
 
         if llm_response and llm_response.role != "assistant":
@@ -482,18 +508,6 @@ async def _save_to_history(
             logger.debug("The LLM response is empty; not saving a record.")
             return
 
-        messages_to_save: list[Message] = []
-        skipped_initial_system = False
-        for message in all_messages:
-            if message.role == "system" and not skipped_initial_system:
-                skipped_initial_system = True
-                continue
-            if message.role in ["assistant", "user"] and message._no_save:
-                continue
-            messages_to_save.append(message)
-
-        checkpoint_id = event.get_extra("llm_checkpoint_id")
-        message_to_save = dump_messages_with_checkpoints(messages_to_save)
         if isinstance(checkpoint_id, str) and checkpoint_id:
             message_to_save.append(
                 CheckpointMessageSegment(

tests/test_conversation_checkpoint.py

@@ -1,3 +1,6 @@
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
 import pytest
 
 from astrbot.core.agent.message import (
@@ -10,7 +13,11 @@
     get_checkpoint_id,
     strip_checkpoint_messages,
 )
-from astrbot.core.provider.entities import ProviderRequest
+from astrbot.core.db.po import Conversation
+from astrbot.core.pipeline.process_stage.method.agent_sub_stages.internal import (
+    InternalAgentSubStage,
+)
+from astrbot.core.provider.entities import LLMResponse, ProviderRequest
 from astrbot.core.provider.provider import Provider
 from astrbot.dashboard.services.chat_service import find_turn_range
 
@@ -163,3 +170,39 @@ def test_chat_service_find_turn_range():
 
     assert find_turn_range(history, "cp-2") == (3, 5)
     assert find_turn_range(history, "missing") is None
+
+
+@pytest.mark.asyncio
+async def test_failed_llm_response_persists_checkpoint_for_retry():
+    conversation_manager = AsyncMock()
+    stage = InternalAgentSubStage()
+    stage.conv_manager = conversation_manager
+    event = SimpleNamespace(
+        unified_msg_origin="webchat:FriendMessage:test",
+        get_extra=lambda key: {"llm_checkpoint_id": "cp-1"}.get(key),
+    )
+    request = ProviderRequest(
+        conversation=Conversation(
+            platform_id="webchat",
+            user_id="webchat:FriendMessage:test",
+            cid="conversation-1",
+        )
+    )
+
+    await stage._save_to_history(
+        event,
+        request,
+        LLMResponse(role="err", completion_text="upstream failed"),
+        [Message(role="user", content="hello")],
+        runner_stats=None,
+    )
+
+    conversation_manager.update_conversation.assert_awaited_once_with(
+        "webchat:FriendMessage:test",
+        "conversation-1",
+        history=[
+            {"role": "user", "content": "hello"},
+            {"role": "_checkpoint", "content": {"id": "cp-1"}},
+        ],
+        token_usage=None,
+    )