SSE 调整

This commit is contained in:
2026-06-10 15:02:58 +08:00
parent 5f015a5c75
commit 41f2a84979
4 changed files with 92 additions and 9 deletions
+1 -1
View File
@@ -327,5 +327,5 @@ type StreamParams struct {
// Subscription 是订阅者读取 ring buffer 的句柄。 // Subscription 是订阅者读取 ring buffer 的句柄。
type Subscription interface { type Subscription interface {
Next(ctx context.Context, lastEventID int64) (events []SSEEvent, closed bool, err error) Next(ctx context.Context, nextEventID int64) (events []SSEEvent, closed bool, err error)
} }
+3 -3
View File
@@ -258,15 +258,15 @@ func (s *messageService) Stream(
out <- SSEEvent{Event: "error", Data: map[string]string{"code": "server_restart", "message": e}} out <- SSEEvent{Event: "error", Data: map[string]string{"code": "server_restart", "message": e}}
return return
} }
lastID := int64(-1) nextEventID := int64(0)
for { for {
events, closed, err := sub.Next(ctx, lastID) events, closed, err := sub.Next(ctx, nextEventID)
if err != nil { if err != nil {
return return
} }
for _, ev := range events { for _, ev := range events {
out <- ev out <- ev
lastID = ev.ID nextEventID = ev.ID + 1
} }
if closed { if closed {
return return
+83
View File
@@ -704,6 +704,35 @@ func (s *closedSubscription) Next(_ context.Context, lastID int64) ([]SSEEvent,
return nil, true, nil 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) { func TestMessageService_Stream_HappyPath_OK(t *testing.T) {
repo := newMsgFakeRepo() repo := newMsgFakeRepo()
hub := &fakeHub{} hub := &fakeHub{}
@@ -785,6 +814,60 @@ func TestMessageService_Stream_PendingFallbackToError(t *testing.T) {
repo.mu.Unlock() 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) { func TestMessageService_Stream_NotOwner_Forbidden(t *testing.T) {
repo := newMsgFakeRepo() repo := newMsgFakeRepo()
hub := &fakeHub{} hub := &fakeHub{}
+5 -5
View File
@@ -66,14 +66,14 @@ type subscription struct {
st *streamer st *streamer
} }
// Next blocks until at least one event with ID >= lastEventID is available, // Next blocks until at least one event with ID >= nextEventID is available,
// the stream closes, or ctx is cancelled. // the stream closes, or ctx is cancelled.
// It returns all events from lastEventID to the current tail. // It returns all events from nextEventID to the current tail.
func (sub *subscription) Next(ctx context.Context, lastEventID int64) ([]SSEEvent, bool, error) { func (sub *subscription) Next(ctx context.Context, nextEventID int64) ([]SSEEvent, bool, error) {
sub.st.mu.Lock() sub.st.mu.Lock()
defer sub.st.mu.Unlock() defer sub.st.mu.Unlock()
for int64(len(sub.st.events)) <= lastEventID && !sub.st.closed { for int64(len(sub.st.events)) <= nextEventID && !sub.st.closed {
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
select { select {
@@ -88,7 +88,7 @@ func (sub *subscription) Next(ctx context.Context, lastEventID int64) ([]SSEEven
return nil, false, ctx.Err() return nil, false, ctx.Err()
} }
} }
out := append([]SSEEvent(nil), sub.st.events[lastEventID:]...) out := append([]SSEEvent(nil), sub.st.events[nextEventID:]...)
return out, sub.st.closed, nil return out, sub.st.closed, nil
} }