server/agui: correct id for assistant message and tool result event (#703)

Author: Flash-LHR

server/agui/runner/messagessnapshot.go

@@ -139,7 +139,7 @@ func (r *runner) convertToMessagesSnapshotEvent(ctx context.Context, userID stri
 					messages = append(messages, *r.convertToUserMessage(event.ID, userID, choice))
 				}
 			case model.RoleAssistant:
-				messages = append(messages, *r.convertToAssistantMessage(event.ID, choice))
+				messages = append(messages, *r.convertToAssistantMessage(event.Response.ID, choice))
 			case model.RoleTool:
 				messages = append(messages, *r.convertToToolMessage(event.ID, choice))
 			default:

server/agui/runner/messagessnapshot_test.go

@@ -326,6 +326,26 @@ func TestConvertToMessagesSnapshotEventDeduplicatesUserMessages(t *testing.T) {
 	assert.Equal(t, "next", *snapshot.Messages[1].Content)
 }
 
+func TestConvertToMessagesSnapshotEventUsesResponseIDForAssistant(t *testing.T) {
+	r := &runner{}
+	assistantEvent := event.Event{
+		ID: "event-id",
+		Response: &model.Response{
+			ID:     "response-id",
+			Object: model.ObjectTypeChatCompletion,
+			Choices: []model.Choice{
+				{Message: model.Message{Role: model.RoleAssistant, Content: "reply"}},
+			},
+		},
+	}
+	snapshot, err := r.convertToMessagesSnapshotEvent(context.Background(), "user-id", []event.Event{assistantEvent})
+	assert.NoError(t, err)
+	assert.NotNil(t, snapshot)
+	if assert.Len(t, snapshot.Messages, 1) {
+		assert.Equal(t, "response-id", snapshot.Messages[0].ID)
+	}
+}
+
 func TestIgnoreEvent(t *testing.T) {
 	tests := []struct {
 		name string

server/agui/translator/translator.go

@@ -67,7 +67,7 @@ func (t *translator) Translate(event *agentevent.Event) ([]aguievents.Event, err
 		events = append(events, toolCallEvents...)
 	}
 	if rsp.IsToolResultResponse() {
-		toolResultEvents, err := t.toolResultEvent(rsp)
+		toolResultEvents, err := t.toolResultEvent(rsp, event.ID)
 		if err != nil {
 			return nil, err
 		}
@@ -165,7 +165,7 @@ func (t *translator) toolCallEvent(rsp *model.Response) ([]aguievents.Event, err
 }
 
 // toolResultEvent translates a tool result trpc-agent-go event to AG-UI events.
-func (t *translator) toolResultEvent(rsp *model.Response) ([]aguievents.Event, error) {
+func (t *translator) toolResultEvent(rsp *model.Response, messageID string) ([]aguievents.Event, error) {
 	if rsp == nil || len(rsp.Choices) == 0 {
 		return nil, nil
 	}
@@ -174,7 +174,7 @@ func (t *translator) toolResultEvent(rsp *model.Response) ([]aguievents.Event, e
 		// Tool call end event.
 		eventas = append(eventas, aguievents.NewToolCallEndEvent(choice.Message.ToolID))
 		// Tool call result event.
-		eventas = append(eventas, aguievents.NewToolCallResultEvent(t.lastMessageID,
+		eventas = append(eventas, aguievents.NewToolCallResultEvent(messageID,
 			choice.Message.ToolID, choice.Message.Content))
 	}
 	return eventas, nil

server/agui/translator/translator_test.go

@@ -280,15 +280,15 @@ func TestToolCallAndResultEvents(t *testing.T) {
 			Message: model.Message{ToolID: "call-1", Content: "done"},
 		}},
 	}
-	resultEvents, err := translator.toolResultEvent(resultRsp)
+	resultEvents, err := translator.toolResultEvent(resultRsp, "event-tool-result")
 	assert.NoError(t, err)
 	assert.Len(t, resultEvents, 2)
 	end, ok := resultEvents[0].(*aguievents.ToolCallEndEvent)
 	assert.True(t, ok)
 	assert.Equal(t, "call-1", end.ToolCallID)
 	res, ok := resultEvents[1].(*aguievents.ToolCallResultEvent)
 	assert.True(t, ok)
-	assert.Equal(t, "msg-tool", res.MessageID)
+	assert.Equal(t, "event-tool-result", res.MessageID)
 	assert.Equal(t, "call-1", res.ToolCallID)
 	assert.Equal(t, "done", res.Content)
 }
@@ -416,17 +416,20 @@ func TestTranslateToolResultResponse(t *testing.T) {
 	}})
 	assert.NoError(t, err)
 
-	events, err := translator.Translate(&agentevent.Event{Response: &model.Response{
-		Choices: []model.Choice{{
-			Message: model.Message{ToolID: "tool-1", Content: "done"},
-		}},
-	}})
+	events, err := translator.Translate(&agentevent.Event{
+		ID: "evt-tool-1",
+		Response: &model.Response{
+			Choices: []model.Choice{{
+				Message: model.Message{ToolID: "tool-1", Content: "done"},
+			}},
+		},
+	})
 	assert.NoError(t, err)
 	assert.Len(t, events, 2)
 	assert.IsType(t, (*aguievents.ToolCallEndEvent)(nil), events[0])
 	result, ok := events[1].(*aguievents.ToolCallResultEvent)
 	assert.True(t, ok)
-	assert.Equal(t, "msg-1", result.MessageID)
+	assert.Equal(t, "evt-tool-1", result.MessageID)
 	assert.Equal(t, "tool-1", result.ToolCallID)
 	assert.Equal(t, "done", result.Content)
 }
@@ -471,13 +474,13 @@ func TestTranslateSequentialEvents(t *testing.T) {
 			Message: model.Message{ToolID: "call-1", Content: "success"},
 		}},
 	}
-	events, err = translator.Translate(&agentevent.Event{Response: toolResultRsp})
+	events, err = translator.Translate(&agentevent.Event{ID: "evt-call-1-result", Response: toolResultRsp})
 	assert.NoError(t, err)
 	assert.Len(t, events, 2)
 	assert.IsType(t, (*aguievents.ToolCallEndEvent)(nil), events[0])
 	res, ok := events[1].(*aguievents.ToolCallResultEvent)
 	assert.True(t, ok)
-	assert.Equal(t, "msg-1", res.MessageID)
+	assert.Equal(t, "evt-call-1-result", res.MessageID)
 	assert.Equal(t, "call-1", res.ToolCallID)
 
 	finalRsp := &model.Response{
@@ -545,7 +548,7 @@ func TestToolNilResponse(t *testing.T) {
 	events, err := translator.toolCallEvent(nil)
 	assert.Empty(t, events)
 	assert.NoError(t, err)
-	events, err = translator.toolResultEvent(nil)
+	events, err = translator.toolResultEvent(nil, "")
 	assert.Empty(t, events)
 	assert.NoError(t, err)
 }