SSE 调整

internal/chat/service_message_test.go

@@ -704,6 +704,35 @@ func (s *closedSubscription) Next(_ context.Context, lastID int64) ([]SSEEvent,
 	return nil, true, nil
 }
 
+type recordingSubscription struct {
+	mu      sync.Mutex
+	calls   []int64
+	batches []subscriptionBatch
+}
+
+type subscriptionBatch struct {
+	events []SSEEvent
+	closed bool
+}
+
+func (s *recordingSubscription) Next(_ context.Context, nextID int64) ([]SSEEvent, bool, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.calls = append(s.calls, nextID)
+	i := len(s.calls) - 1
+	if i >= len(s.batches) {
+		return nil, true, nil
+	}
+	batch := s.batches[i]
+	return append([]SSEEvent(nil), batch.events...), batch.closed, nil
+}
+
+func (s *recordingSubscription) callsSnapshot() []int64 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return append([]int64(nil), s.calls...)
+}
+
 func TestMessageService_Stream_HappyPath_OK(t *testing.T) {
 	repo := newMsgFakeRepo()
 	hub := &fakeHub{}
@@ -785,6 +814,60 @@ func TestMessageService_Stream_PendingFallbackToError(t *testing.T) {
 	repo.mu.Unlock()
 }
 
+func TestMessageService_Stream_PendingSubscriptionStartsAtFirstEvent(t *testing.T) {
+	repo := newMsgFakeRepo()
+	userID := uuid.New()
+	conv := newConv(userID, uuid.New())
+	repo.addConversation(conv)
+
+	pending := &Message{
+		ID:             1,
+		ConversationID: conv.ID,
+		Status:         StatusPending,
+		Role:           RoleAssistant,
+	}
+	repo.addMessage(pending)
+
+	sub := &recordingSubscription{
+		batches: []subscriptionBatch{
+			{
+				events: []SSEEvent{
+					{ID: 0, Event: "text_delta", Data: map[string]string{"text": "hi"}},
+				},
+				closed: false,
+			},
+			{
+				events: []SSEEvent{
+					{ID: 1, Event: "done", Data: map[string]string{"finish_reason": "stop"}},
+				},
+				closed: true,
+			},
+		},
+	}
+	hub := &fakeHub{
+		subscribeFn: func(id int64) (Subscription, bool) {
+			require.Equal(t, pending.ID, id)
+			return sub, true
+		},
+	}
+
+	svc := buildMessageService(repo, hub, defaultMsgConfig())
+	ch, err := svc.Stream(context.Background(), userID, conv.ID, 0)
+
+	require.NoError(t, err)
+
+	var events []SSEEvent
+	for ev := range ch {
+		events = append(events, ev)
+	}
+
+	require.Equal(t, []int64{0, 1}, sub.callsSnapshot())
+	require.Len(t, events, 3)
+	require.Equal(t, "message_meta", events[0].Event)
+	require.Equal(t, "text_delta", events[1].Event)
+	require.Equal(t, "done", events[2].Event)
+}
+
 func TestMessageService_Stream_NotOwner_Forbidden(t *testing.T) {
 	repo := newMsgFakeRepo()
 	hub := &fakeHub{}