Cyrene_v2/internal/services/gateway/ws_hub_stream_scope_test.go

package gateway

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// scopeStreamAdapter 是一个 model.Adapter，发出两个文本增量，使测试
// 能驱动真实 loop 并观察它产出的成帧事件。
type scopeStreamAdapter struct{}

func (scopeStreamAdapter) Name() string { return "scope-test" }

func (scopeStreamAdapter) Chat(context.Context, model.Request, model.Options) (model.Response, error) {
	return model.Response{Content: "Hel lo"}, nil
}

func (scopeStreamAdapter) ChatStream(_ context.Context, _ model.Request, on func(model.StreamEvent)) (model.Response, error) {
	if on != nil {
		on(model.StreamEvent{Kind: model.StreamKindTextDelta, Delta: "Hel"})
		on(model.StreamEvent{Kind: model.StreamKindTextDelta, Delta: "lo"})
	}
	return model.Response{Content: "Hel lo"}, nil
}

// newScopedStreamHub 把 hub 接到真实 loop，其 OnStream 闭包与
// cmd/cyrene 的 streamChunk 一致（把整个 StreamEvent 值发布到
// ws.stream），因此测试走的是生产路径 OnStream → bus → relay → frame。
// 它返回 hub 与一个已注册的假连接。
func newScopedStreamHub(t *testing.T) (*hub, *wsConn) {
	t.Helper()
	b := bus.New()
	ag, err := agent.New("You are a helpful assistant.", nil, nil, agent.WithModel(scopeStreamAdapter{}, nil))
	if err != nil {
		t.Fatalf("agent.New: %v", err)
	}
	lp := loop.New(loop.Options{
		Agent:     ag,
		MaxRounds: 8,
		OnStream: func(ev model.StreamEvent) {
			_ = b.Publish(context.Background(), wsStreamSubject, ev)
		},
	})
	hb := newTestHub(lp, b)
	hb.streaming = true
	hb.start()
	t.Cleanup(hb.close)
	c := newOwnerConn(hb, 64)
	hb.register(c)
	return hb, c
}

// readDeltas 扫描连接的帧，直到 want 个 message.delta 帧到达。增量经
// bus 异步到达该连接，而 runTurn 自己的 stream_end 是直接发给连接
// 的，因此测试不能假定顺序。
func readDeltas(t *testing.T, c *wsConn, want int) []wsEvent {
	t.Helper()
	var deltas []wsEvent
	deadline := time.After(3 * time.Second)
	for len(deltas) < want {
		select {
		case raw := <-c.send:
			var ev wsEvent
			if err := json.Unmarshal(raw, &ev); err != nil {
				t.Fatalf("decode frame %s: %v", raw, err)
			}
			if ev.Type == "message.delta" {
				deltas = append(deltas, ev)
			}
		case <-deadline:
			t.Fatalf("timed out with %d message.delta frames, want %d", len(deltas), want)
		}
	}
	return deltas
}

// TestStreamFramesCarrySessionScope 是 A1 的端到端锁定：经 hub 的
// runTurn 驱动的一轮必须发出携带会话与轮次 id 的 message.delta 帧，
// 使客户端能把它们路由到活跃会话。它跑真实 loop 而不是直接调用
// broadcastStream，这同时证明 bus 原样携带了新字段 —— 无需改动
// 发布方。
func TestStreamFramesCarrySessionScope(t *testing.T) {
	t.Parallel()
	hb, c := newScopedStreamHub(t)

	if _, err := hb.runTurn(context.Background(), "web", []model.Message{model.TextMessage(model.RoleUser, "hi")}, loop.OwnerEnv()); err != nil {
		t.Fatalf("runTurn: %v", err)
	}
	deltas := readDeltas(t, c, 2)

	for i, ev := range deltas {
		if ev.SessionID != "web" {
			t.Fatalf("delta[%d] sessionId = %q, want web (frame=%+v)", i, ev.SessionID, ev)
		}
		if !strings.HasPrefix(ev.TurnID, "turn-") {
			t.Fatalf("delta[%d] turnId = %q, want a turn-* id (frame=%+v)", i, ev.TurnID, ev)
		}
	}
	if deltas[0].TurnID != deltas[1].TurnID {
		t.Fatalf("turnId differs within one turn: %q vs %q", deltas[0].TurnID, deltas[1].TurnID)
	}
	if deltas[0].Seq == 0 || deltas[1].Seq <= deltas[0].Seq {
		t.Fatalf("seq not increasing: %d then %d", deltas[0].Seq, deltas[1].Seq)
	}
	if data, ok := deltas[0].Data.(map[string]any); !ok || data["delta"] != "Hel" {
		t.Fatalf("first delta data = %+v, want delta=Hel", deltas[0].Data)
	}
}

// TestStreamFramesTurnScopeChangesPerTurn 覆盖编排下的续跑形态：在
// 首轮完成之后被排空的待处理轮次是第二次 runTurn 调用，因此它必须
// 取新的 turnId。否则续跑轮次的增量与上一轮的将无法区分。
func TestStreamFramesTurnScopeChangesPerTurn(t *testing.T) {
	t.Parallel()
	hb, c := newScopedStreamHub(t)

	msgs := []model.Message{model.TextMessage(model.RoleUser, "hi")}
	if _, err := hb.runTurn(context.Background(), "web", msgs, loop.OwnerEnv()); err != nil {
		t.Fatalf("first runTurn: %v", err)
	}
	first := readDeltas(t, c, 2)
	if _, err := hb.runTurn(context.Background(), "web", msgs, loop.OwnerEnv()); err != nil {
		t.Fatalf("second runTurn: %v", err)
	}
	second := readDeltas(t, c, 2)

	for i, ev := range second {
		if ev.SessionID != "web" {
			t.Fatalf("second turn delta[%d] sessionId = %q, want web", i, ev.SessionID)
		}
		if ev.TurnID != second[0].TurnID {
			t.Fatalf("second turn delta[%d] turnId = %q, want %q for the whole turn", i, ev.TurnID, second[0].TurnID)
		}
	}
	if first[0].TurnID == second[0].TurnID {
		t.Fatalf("turnId reused across turns: %q", first[0].TurnID)
	}
	if second[0].Seq <= first[1].Seq {
		t.Fatalf("seq not monotonic across turns: %d then %d", first[1].Seq, second[0].Seq)
	}
}