You've already forked agentic-coding-workflow
SSE 调整
This commit is contained in:
@@ -327,5 +327,5 @@ type StreamParams struct {
|
||||
|
||||
// Subscription 是订阅者读取 ring buffer 的句柄。
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -258,15 +258,15 @@ func (s *messageService) Stream(
|
||||
out <- SSEEvent{Event: "error", Data: map[string]string{"code": "server_restart", "message": e}}
|
||||
return
|
||||
}
|
||||
lastID := int64(-1)
|
||||
nextEventID := int64(0)
|
||||
for {
|
||||
events, closed, err := sub.Next(ctx, lastID)
|
||||
events, closed, err := sub.Next(ctx, nextEventID)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, ev := range events {
|
||||
out <- ev
|
||||
lastID = ev.ID
|
||||
nextEventID = ev.ID + 1
|
||||
}
|
||||
if closed {
|
||||
return
|
||||
|
||||
@@ -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{}
|
||||
|
||||
@@ -66,14 +66,14 @@ type subscription struct {
|
||||
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.
|
||||
// It returns all events from lastEventID to the current tail.
|
||||
func (sub *subscription) Next(ctx context.Context, lastEventID int64) ([]SSEEvent, bool, error) {
|
||||
// It returns all events from nextEventID to the current tail.
|
||||
func (sub *subscription) Next(ctx context.Context, nextEventID int64) ([]SSEEvent, bool, error) {
|
||||
sub.st.mu.Lock()
|
||||
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{})
|
||||
go func() {
|
||||
select {
|
||||
@@ -88,7 +88,7 @@ func (sub *subscription) Next(ctx context.Context, lastEventID int64) ([]SSEEven
|
||||
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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user