Skip to content

feat: send NewCheckpoint msgs on tool use generations

What does this merge request do and why?

This reverts this MR: chore: send only non-empty new-checkpoints gprc (!2938 - merged)

We don't perform a round-trip on NewCheckpoint anymore, so it should be safe to send it on token generation to keep the channel busy.

When a tool-use is streamed, we can send NewCheckpoint with the snapshot of the state made before the tool-use started the generation.

Combined with fine-grained-tool-streaming-2025-05-14 beta feature:

  • It'll help to keep the channel busy while the tool-use is generated
  • It the first step towards tool-use message streaming

Potential tool streaming

Click to expand
diff --git a/duo_workflow_service/checkpointer/notifier.py b/duo_workflow_service/checkpointer/notifier.py
index 078fa4b48..2a00d0b6c 100644
--- a/duo_workflow_service/checkpointer/notifier.py
+++ b/duo_workflow_service/checkpointer/notifier.py
@@ -14,6 +14,7 @@ from duo_workflow_service.checkpointer.gitlab_workflow import (
 )
 from duo_workflow_service.entities.state import (
     MessageTypeEnum,
+    ToolInfo,
     UiChatLog,
     WorkflowStatusEnum,
 )
@@ -33,6 +34,7 @@ class UserInterface:
         self.ui_chat_log: list[UiChatLog] = []
         self.status = WorkflowStatusEnum.NOT_STARTED
         self.steps: list[dict] = []
+        self.last_message_chunk = None
 
     async def send_event(
         self,
@@ -95,15 +97,13 @@ class UserInterface:
         Args:
             message (BaseMessage): The message chunk to be processed and added to the log.
         """
-        content = StrOutputParser().invoke(message) or ""
-        if not content:
-            return
 
         if (
             not self.ui_chat_log
-            or self.ui_chat_log[-1]["message_type"] != MessageTypeEnum.AGENT
+            or not self.ui_chat_log[-1]["message_type"] in [MessageTypeEnum.AGENT, MessageTypeEnum.TOOL]
             or self.ui_chat_log[-1]["status"]
         ):
+            self.last_message_chunk = message
             last_message = UiChatLog(
                 status=None,
                 correlation_id=None,
@@ -116,6 +116,12 @@ class UserInterface:
             )
             self.ui_chat_log.append(last_message)
         else:
+            self.last_message_chunk += message
             last_message = self.ui_chat_log[-1]
 
-        last_message["content"] = last_message["content"] + content
+        last_message["content"] = StrOutputParser().invoke(self.last_message_chunk) or ""
+
+        if self.last_message_chunk.tool_calls:
+            tool_call = self.last_message_chunk.tool_calls[0]
+            last_message["tool_info"] = ToolInfo(name=tool_call["name"], args=tool_call["args"])
+            last_message["message_type"] = MessageTypeEnum.TOOL

output

Related to gitlab-org/gitlab#570575 (closed)

Edited by Dylan Griffith

Merge request reports

Loading