Reconstruct original order of assistant and tool messages

[carloscco]

The current reducer only creates a single assistant message where all text content is included:

ag-ui/sdks/typescript/packages/client/src/apply/default.ts

Line 135 in 872665d

// If message already exists, we don't need to create a new one

This way, despite having references to toolCalls, is not possible to reconstruct the order of tool-call vs assistant-message

Eg.: With the following events:

• tool-1 call
• assistant message
• tool-2 call
• assistant message

data: {"type":"TOOL_CALL_START","parentMessageId":"m-1","toolCallId":"call_1","toolCallName":"tool-1"}

data: {"type":"TOOL_CALL_ARGS","toolCallId":"call_1","delta":"{...}"}

data: {"type":"TOOL_CALL_END","toolCallId":"call_1"}

data: {"type":"TOOL_CALL_RESULT","toolCallId":"call_1","content":"{}" }

data: {"type":"TEXT_MESSAGE_START","messageId":"m-1", "role":"assistant"}

data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"m-1","delta":"<delta-1>"}

data: {"type":"TEXT_MESSAGE_END","messageId":"m-1"}

data: {"type":"TOOL_CALL_START","parentMessageId":"m-1","toolCallId":"call_2","toolCallName":"tool-2"}

data: {"type":"TOOL_CALL_ARGS","toolCallId":"call_2","delta":"{...}"}

data: {"type":"TOOL_CALL_END","toolCallId":"call_2"}

data: {"type":"TOOL_CALL_RESULT","toolCallId":"call_2","content":"{}" }

data: {"type":"TEXT_MESSAGE_START","messageId":"m-1", "role":"assistant"}

data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"m-1","delta":"<delta-2>"}

data: {"type":"TEXT_MESSAGE_END","messageId":"m-1"}

Because a single message is generated to include all assistant content, the resulting messages include:

[{ id: "m-1", role: "assistant", toolCalls: [...], content: "<delta-1><delta-2>"}, ...]
 

Is this intended behavior, or is it a limitation?

[carloscco]

The solution I found is writing an AgentSubscriber:

• override onTextMessageStartEvent to always create a new message (instead of reusing the same for all content)
• override onTextMessageContentEvent to check the last message matching the id, instead of the first

onTextMessageStartEvent: ({ event, messages}) => {
    const { messageId: id, role = "assistant", name } = event;
    // Always create a new message
    messages.push({
      id,
      role,
      content: "",
      ...(name !== undefined && { name}),
    } as Message);
    return {
      messages,
      stopPropagation: true
    }
},
onTextMessageContentEvent({ event, messages}){
  const targetMessage = [...messages].reverse().find((m) => m.id === event.messageId)!;

    const existingContent = typeof targetMessage.content === "string" ? targetMessage.content : "";
    targetMessage.content = `${existingContent}${event.delta}`;

    return {
      messages,
      stopPropagation: true
    }
}