diff --git a/workflow/agentworkflow/message_merger.go b/workflow/agentworkflow/message_merger.go index df7bb5ce..6abe16e7 100644 --- a/workflow/agentworkflow/message_merger.go +++ b/workflow/agentworkflow/message_merger.go @@ -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] @@ -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 @@ -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 } @@ -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 { @@ -171,7 +167,7 @@ func mergeResponseList(responses []*agent.Response) *agent.Response { } current.AgentID = cmp.Or(incoming.AgentID, current.AgentID) current.AdditionalProperties = mergeProperties(current.AdditionalProperties, incoming.AdditionalProperties) - if !incoming.CreatedAt.IsZero() { + if !incoming.CreatedAt.IsZero() && (current.CreatedAt.IsZero() || incoming.CreatedAt.After(current.CreatedAt)) { current.CreatedAt = incoming.CreatedAt } current.FinishReason = cmp.Or(incoming.FinishReason, current.FinishReason) @@ -188,7 +184,7 @@ func messagesWithCreatedAt(response *agent.Response) []*message.Message { messages := make([]*message.Message, 0, len(response.Messages)) 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) diff --git a/workflow/agentworkflow/message_merger_test.go b/workflow/agentworkflow/message_merger_test.go new file mode 100644 index 00000000..08db3781 --- /dev/null +++ b/workflow/agentworkflow/message_merger_test.go @@ -0,0 +1,188 @@ +// 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_StampsZeroCreatedAtWithMaxTimestamp(t *testing.T) { + responseID := "response" + older := time.Date(2026, 7, 15, 12, 0, 0, 0, time.UTC) + newer := older.Add(time.Hour) + + merger := newMessageMerger() + // message-1 has the newest timestamp; message-2 has no timestamp; message-3 is older. + // Without max-tracking, "last-wins" would stamp message-2 with `older` instead of `newer`. + addTextUpdate(merger, responseID, "first", "message-1", newer) + addTextUpdate(merger, responseID, "second", "message-2", time.Time{}) + addTextUpdate(merger, responseID, "third", "message-3", older) + + response := merger.ComputeMerged(responseID, "", "") + + assertMessageTexts(t, response.Messages, "first", "second", "third") + if got := response.Messages[1].CreatedAt; !got.Equal(newer) { + t.Fatalf("zero-timestamp message stamped with %v, want max timestamp %v", got, newer) + } +} + +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]) + } + } +}