You've already forked agentic-coding-workflow
SSE 调整
This commit is contained in:
@@ -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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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{}
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user