Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 36 additions & 40 deletions workflow/agentworkflow/message_merger.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ func (m *messageMerger) AddUpdate(update *agent.ResponseUpdate) {
return
}
if update.ResponseID == "" {
m.danglingState.addDangling(update)
m.danglingState.addUpdate(update)
return
}
state, ok := m.states[update.ResponseID]
Expand All @@ -51,7 +51,6 @@ func (m *messageMerger) ComputeMerged(primaryResponseID string, primaryAgentID s
for _, responseID := range m.stateOrder {
state := m.states[responseID]
responses := state.computeResponses()
slices.SortFunc(responses, compareResponsesByCreatedAt)
merged := mergeResponseList(responses)
if merged == nil {
continue
Expand Down Expand Up @@ -89,37 +88,49 @@ func (m *messageMerger) ComputeMerged(primaryResponseID string, primaryAgentID s
}

type responseMergeState struct {
responseID string
updatesByMessageID map[string][]*agent.ResponseUpdate
messageOrder []string
danglingUpdates []*agent.ResponseUpdate
responseID string
messageStatesByID map[string]*messageMergeState
messageStateOrder []*messageMergeState
lastObservedState *messageMergeState
}

type messageMergeState struct {
messageID string
identifierless bool
updates []*agent.ResponseUpdate
}

func (s *responseMergeState) addUpdate(update *agent.ResponseUpdate) {
if update.MessageID == "" {
s.addDangling(update)
return
state := s.getOrCreateMessageState(update.MessageID)
state.updates = append(state.updates, update)
s.lastObservedState = state
}

func (s *responseMergeState) getOrCreateMessageState(messageID string) *messageMergeState {
if messageID == "" {
if s.lastObservedState != nil && s.lastObservedState.identifierless {
return s.lastObservedState
}
state := &messageMergeState{identifierless: true}
s.messageStateOrder = append(s.messageStateOrder, state)
return state
}
if s.updatesByMessageID == nil {
s.updatesByMessageID = make(map[string][]*agent.ResponseUpdate)
if s.messageStatesByID == nil {
s.messageStatesByID = make(map[string]*messageMergeState)
}
if _, ok := s.updatesByMessageID[update.MessageID]; !ok {
s.messageOrder = append(s.messageOrder, update.MessageID)
if existing, ok := s.messageStatesByID[messageID]; ok {
return existing
}
s.updatesByMessageID[update.MessageID] = append(s.updatesByMessageID[update.MessageID], update)
}

func (s *responseMergeState) addDangling(update *agent.ResponseUpdate) {
s.danglingUpdates = append(s.danglingUpdates, update)
state := &messageMergeState{messageID: messageID}
s.messageStatesByID[messageID] = state
s.messageStateOrder = append(s.messageStateOrder, state)
return state
}

func (s *responseMergeState) computeResponses() []*agent.Response {
responses := make([]*agent.Response, 0, len(s.messageOrder)+1)
for _, messageID := range s.messageOrder {
responses = append(responses, responseFromUpdates(s.updatesByMessageID[messageID]))
}
if len(s.danglingUpdates) > 0 {
responses = append(responses, responseFromUpdates(s.danglingUpdates))
responses := make([]*agent.Response, 0, len(s.messageStateOrder))
for _, state := range s.messageStateOrder {
responses = append(responses, responseFromUpdates(state.updates))
}
return responses
}
Expand All @@ -141,21 +152,6 @@ func responseFromUpdates(updates []*agent.ResponseUpdate) *agent.Response {
return response
}

func compareResponsesByCreatedAt(left, right *agent.Response) int {
leftZero := left == nil || left.CreatedAt.IsZero()
rightZero := right == nil || right.CreatedAt.IsZero()
switch {
case leftZero && rightZero:
return 0
case leftZero:
return 1
case rightZero:
return -1
default:
return left.CreatedAt.Compare(right.CreatedAt)
}
}

func mergeResponseList(responses []*agent.Response) *agent.Response {
var current *agent.Response
for _, incoming := range responses {
Expand Down Expand Up @@ -188,7 +184,7 @@ func messagesWithCreatedAt(response *agent.Response) []*message.Message {
messages := make([]*message.Message, 0, len(response.Messages))
Comment on lines 180 to 184
for _, msg := range response.Messages {
clone := msg.Clone()
if clone != nil && !response.CreatedAt.IsZero() {
if clone != nil && clone.CreatedAt.IsZero() && !response.CreatedAt.IsZero() {
clone.CreatedAt = response.CreatedAt
}
messages = append(messages, clone)
Expand Down
168 changes: 168 additions & 0 deletions workflow/agentworkflow/message_merger_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,168 @@
// Copyright (c) Microsoft. All rights reserved.

package agentworkflow

import (
"testing"
"time"

"github.com/microsoft/agent-framework-go/agent"
"github.com/microsoft/agent-framework-go/message"
)

func TestMessageMerger_PreservesFirstSeenMessageOrder(t *testing.T) {
responseID := "response"
now := time.Date(2026, 7, 15, 12, 0, 0, 0, time.UTC)

merger := newMessageMerger()
addTextUpdate(merger, responseID, "first", "message-1", now.Add(time.Minute))
addTextUpdate(merger, responseID, "second", "message-2", time.Time{})
addTextUpdate(merger, responseID, "third", "message-3", now.Add(-time.Minute))
addTextUpdate(merger, responseID, "fourth", "message-4", now.Add(-time.Minute))

response := merger.ComputeMerged(responseID, "", "")

assertMessageTexts(t, response.Messages, "first", "second", "third", "fourth")
if got := response.Messages[0].CreatedAt; !got.Equal(now.Add(time.Minute)) {
t.Fatalf("first message CreatedAt = %v, want %v", got, now.Add(time.Minute))
}
if got := response.Messages[2].CreatedAt; !got.Equal(now.Add(-time.Minute)) {
t.Fatalf("third message CreatedAt = %v, want %v", got, now.Add(-time.Minute))
}
}

func TestMessageMerger_KeepsResponsesContiguousInFirstSeenOrder(t *testing.T) {
merger := newMessageMerger()

addTextUpdate(merger, "response-1", "A1", "message-a1", time.Time{})
addTextUpdate(merger, "response-2", "B1", "message-b1", time.Time{})
addTextUpdate(merger, "response-1", "A2", "message-a2", time.Time{})
addTextUpdate(merger, "response-2", "B2", "message-b2", time.Time{})

response := merger.ComputeMerged("response-1", "", "")

assertMessageTexts(t, response.Messages, "A1", "A2", "B1", "B2")
}

func TestMessageMerger_PreservesFunctionCallResultOrder(t *testing.T) {
const (
responseID = "response"
callID = "call"
)

merger := newMessageMerger()
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: responseID,
MessageID: "call-message",
Role: message.RoleAssistant,
Contents: []message.Content{&message.FunctionCallContent{CallID: callID, Name: "handoff"}},
})
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: responseID,
MessageID: "result-message",
Role: message.RoleTool,
CreatedAt: time.Date(2026, 7, 15, 12, 0, 0, 0, time.UTC),
Contents: []message.Content{&message.FunctionResultContent{CallID: callID, Result: "Transferred."}},
})

response := merger.ComputeMerged(responseID, "", "")

if len(response.Messages) != 2 {
t.Fatalf("message count = %d, want 2", len(response.Messages))
}
if _, ok := response.Messages[0].Contents[0].(*message.FunctionCallContent); !ok {
t.Fatalf("first content = %T, want *message.FunctionCallContent", response.Messages[0].Contents[0])
}
if _, ok := response.Messages[1].Contents[0].(*message.FunctionResultContent); !ok {
t.Fatalf("second content = %T, want *message.FunctionResultContent", response.Messages[1].Contents[0])
}
}

func TestMessageMerger_PreservesIdentifierlessMessageOrder(t *testing.T) {
const (
responseID = "response"
callID = "call"
)

merger := newMessageMerger()
addTextUpdate(merger, responseID, "before", "before-message", time.Time{})
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: responseID,
Role: message.RoleAssistant,
Contents: []message.Content{&message.FunctionCallContent{CallID: callID, Name: "handoff"}},
})
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: responseID,
MessageID: "result-message",
Role: message.RoleTool,
CreatedAt: time.Date(2026, 7, 15, 12, 0, 0, 0, time.UTC),
Contents: []message.Content{&message.FunctionResultContent{CallID: callID, Result: "Transferred."}},
})

response := merger.ComputeMerged(responseID, "", "")

if len(response.Messages) != 3 {
t.Fatalf("message count = %d, want 3", len(response.Messages))
}
if got := response.Messages[0].String(); got != "before" {
t.Fatalf("first message = %q, want %q", got, "before")
}
if _, ok := response.Messages[1].Contents[0].(*message.FunctionCallContent); !ok {
t.Fatalf("second content = %T, want *message.FunctionCallContent", response.Messages[1].Contents[0])
}
if _, ok := response.Messages[2].Contents[0].(*message.FunctionResultContent); !ok {
t.Fatalf("third content = %T, want *message.FunctionResultContent", response.Messages[2].Contents[0])
}
}

func TestMessageMerger_SeparatesIdentifierlessSegments(t *testing.T) {
merger := newMessageMerger()
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: "response",
MessageID: "message",
Role: message.RoleAssistant,
Contents: []message.Content{&message.TextContent{Text: "A"}},
})
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: "response",
Role: message.RoleTool,
Contents: []message.Content{&message.TextContent{Text: "X"}},
})
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: "response",
MessageID: "message",
Role: message.RoleAssistant,
Contents: []message.Content{&message.TextContent{Text: "B"}},
})
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: "response",
Role: message.RoleTool,
Contents: []message.Content{&message.TextContent{Text: "Y"}},
})

response := merger.ComputeMerged("response", "", "")

assertMessageTexts(t, response.Messages, "AB", "X", "Y")
}

func addTextUpdate(merger *messageMerger, responseID string, text string, messageID string, createdAt time.Time) {
merger.AddUpdate(&agent.ResponseUpdate{
ResponseID: responseID,
MessageID: messageID,
Role: message.RoleAssistant,
CreatedAt: createdAt,
Contents: []message.Content{&message.TextContent{Text: text}},
})
}

func assertMessageTexts(t *testing.T, messages []*message.Message, want ...string) {
t.Helper()
if len(messages) != len(want) {
t.Fatalf("message count = %d, want %d", len(messages), len(want))
}
for i, msg := range messages {
if got := msg.String(); got != want[i] {
t.Fatalf("message[%d] = %q, want %q", i, got, want[i])
}
}
}