Cyrene_v2/internal/services/gateway/ws_stream_events_test.go

package gateway

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

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"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"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
)

// wsEventsAdapter 是一个两步脚本适配器：第一步发出一个文本增量后
// 返回原生工具调用，第二步再发一个增量并收尾 —— 恰好覆盖
// v0.50 事件批的全部帧形态（turn 边界、两次模型流的 streamId、
// tool_execution 三事件中的 start/end）。
type wsEventsAdapter struct{ calls int }

func (a *wsEventsAdapter) Name() string { return "ws-events-test" }

func (a *wsEventsAdapter) Chat(context.Context, model.Request, model.Options) (model.Response, error) {
	return model.Response{Content: "done"}, nil
}

func (a *wsEventsAdapter) ChatStream(_ context.Context, _ model.Request, on func(model.StreamEvent)) (model.Response, error) {
	a.calls++
	if a.calls == 1 {
		if on != nil {
			on(model.StreamEvent{Kind: model.StreamKindTextDelta, Delta: "Hel", Index: 0})
		}
		return model.Response{Model: "m", Content: "", ToolCalls: []model.ToolCall{{
			ID: "call_ws_1", Name: "datetime", Args: map[string]any{},
		}}}, nil
	}
	if on != nil {
		on(model.StreamEvent{Kind: model.StreamKindTextDelta, Delta: "lo", Index: 0})
	}
	return model.Response{Model: "m", Content: "Hello"}, nil
}

// wsEventsTool 是注册进真实 dispatcher 的最小工具，使 runTurn 走过
// 完整的工具执行路径（含 tool_execution.start/end 与 phase）。
type wsEventsTool struct{}

func (wsEventsTool) Name() string                        { return "datetime" }
func (wsEventsTool) Description() string                 { return "test datetime tool" }
func (wsEventsTool) Parameters() (map[string]any, error) { return map[string]any{}, nil }
func (wsEventsTool) Execute(context.Context, any) (any, error) {
	return map[string]any{"now": "2026-10-04 12:00:00"}, nil
}

// newEventsStreamHub 把 hub 接到带工具的真实 loop，OnStream 闭包与
// cmd/cyrene 的接线一致（整个 StreamEvent 发布到 ws.stream），因此
// 测试驱动的是生产路径 OnStream → bus → relay → broadcastStream。
func newEventsStreamHub(t *testing.T) (*hub, *wsConn) {
	t.Helper()
	b := bus.New()
	reg := registry.New()
	td, err := tool.New(reg)
	if err != nil {
		t.Fatalf("tool.New: %v", err)
	}
	if err := td.Register(wsEventsTool{}); err != nil {
		t.Fatalf("register tool: %v", err)
	}
	ag, err := agent.New("You are a helpful assistant.", td, nil, agent.WithModel(&wsEventsAdapter{}, 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
}

// collectFramesUntilIdle 收集帧直到凑满 want 帧（含 stream_end）后再
// 静默 200ms 确认没有迟到的 bus 帧。loop 帧经 bus 异步抵达而
// stream_end 由 runTurn 直接广播，因此不能假定两者相对顺序。
func collectFramesUntilIdle(t *testing.T, c *wsConn, want int) []wsEvent {
	t.Helper()
	var frames []wsEvent
	decode := func(b []byte) wsEvent {
		var ev wsEvent
		if err := json.Unmarshal(b, &ev); err != nil {
			t.Fatalf("decode frame %s: %v", b, err)
		}
		return ev
	}
	deadline := time.After(5 * time.Second)
	for {
		if len(frames) >= want {
			select {
			case b := <-c.send:
				frames = append(frames, decode(b))
			case <-time.After(200 * time.Millisecond):
				return frames
			case <-deadline:
				return frames
			}
			continue
		}
		select {
		case b := <-c.send:
			frames = append(frames, decode(b))
		case <-deadline:
			t.Fatalf("timed out with %d frames, want at least %d: %+v", len(frames), want, frames)
		}
	}
}

// runEventsTurn 驱动一个完整的工具回合并返回按 seq 排序的帧，
// 以及其中的 stream_end（loop 帧经 bus 异步抵达而 stream_end 由
// runTurn 直接广播，因此不能假定两者相对顺序，排序后统一断言）。
func runEventsTurn(t *testing.T) ([]wsEvent, *wsEvent) {
	t.Helper()
	hb, c := newEventsStreamHub(t)
	if _, err := hb.runTurn(context.Background(), "web", []model.Message{model.TextMessage(model.RoleUser, "time?")}, loop.OwnerEnv()); err != nil {
		t.Fatalf("runTurn: %v", err)
	}
	frames := collectFramesUntilIdle(t, c, 7)
	sort.SliceStable(frames, func(i, j int) bool { return frames[i].Seq < frames[j].Seq })
	var streamEnd *wsEvent
	for i := range frames {
		if frames[i].Type == "stream_end" && streamEnd == nil {
			streamEnd = &frames[i]
		}
	}
	return frames, streamEnd
}

// eventsByType 把 loop 帧（排除 stream_end）按类型归组，同时校验
// 每帧都带会话与 turn-* 作用域。
func eventsByType(t *testing.T, frames []wsEvent) map[string][]wsEvent {
	t.Helper()
	byType := map[string][]wsEvent{}
	for _, f := range frames {
		if f.Type == "stream_end" {
			continue
		}
		byType[f.Type] = append(byType[f.Type], f)
		if f.SessionID != "web" || !strings.HasPrefix(f.TurnID, "turn-") {
			t.Fatalf("frame %s scope = session %q turn %q, want web + turn-*", f.Type, f.SessionID, f.TurnID)
		}
	}
	return byType
}

// TestStreamFramesTurnSequenceAndBoundaries 是 v0.50 回合边界与帧序的
// 端到端锁定（裁决 01.4.3）：经真实 loop + bus + hub 驱动的一个工具
// 回合，loop 帧必须按序发出 turn.start → message.delta →
// tool_execution.start → tool_execution.end → message.delta →
// turn.end；turn 边界帧无载荷、整轮 turnId 一致、stream_end 照常收尾。
func TestStreamFramesTurnSequenceAndBoundaries(t *testing.T) {
	t.Parallel()
	frames, streamEnd := runEventsTurn(t)

	var seq []string
	for _, f := range frames {
		if f.Type != "stream_end" {
			seq = append(seq, f.Type)
		}
	}
	want := []string{
		"turn.start", "message.delta",
		"tool_execution.start", "tool_execution.end",
		"message.delta", "turn.end",
	}
	if len(seq) != len(want) {
		t.Fatalf("loop frame types = %v, want %v (frames=%+v)", seq, want, frames)
	}
	for i := range want {
		if seq[i] != want[i] {
			t.Fatalf("loop frame[%d] = %q, want %q (all=%v)", i, seq[i], want[i], seq)
		}
	}
	if streamEnd == nil {
		t.Fatalf("no stream_end frame arrived (frames=%+v)", frames)
	}

	byType := eventsByType(t, frames)
	turnID := byType["turn.start"][0].TurnID
	if byType["turn.end"][0].TurnID != turnID || byType["message.delta"][0].TurnID != turnID {
		t.Fatalf("turnId differs across the turn: %q / %q / %q", turnID, byType["turn.end"][0].TurnID, byType["message.delta"][0].TurnID)
	}
	for _, boundary := range []wsEvent{byType["turn.start"][0], byType["turn.end"][0]} {
		if boundary.Data != nil {
			t.Fatalf("%s carries data %+v; the boundary frames must be payload-free", boundary.Type, boundary.Data)
		}
	}
}

// TestStreamFramesStreamIdAndToolExecutionPayloads 是 v0.50 streamId
// 与工具执行帧载荷的端到端锁定（裁决 01.4.1 / 01.4.2）：模型流帧带
// 每流独立的 streamId（同流共享、跨流相异）；tool_execution 帧带
// 结构化 tool/phase/id，且不带 delta 拼串、不带 streamId。
func TestStreamFramesStreamIdAndToolExecutionPayloads(t *testing.T) {
	t.Parallel()
	frames, _ := runEventsTurn(t)
	byType := eventsByType(t, frames)

	d1 := dataMap(t, byType["message.delta"][0])
	d2 := dataMap(t, byType["message.delta"][1])
	if d1["delta"] != "Hel" || d2["delta"] != "lo" {
		t.Fatalf("deltas = %v / %v, want Hel / lo", d1["delta"], d2["delta"])
	}
	s1, _ := d1["streamId"].(string)
	s2, _ := d2["streamId"].(string)
	if !strings.HasPrefix(s1, "stream-") || !strings.HasPrefix(s2, "stream-") {
		t.Fatalf("streamIds = %q / %q, want two stream-* ids (one per model stream)", s1, s2)
	}
	if s1 == s2 {
		t.Fatalf("the tool turn's two model streams share streamId %q; each ChatStream call must mint a fresh one", s1)
	}

	startData := dataMap(t, byType["tool_execution.start"][0])
	endData := dataMap(t, byType["tool_execution.end"][0])
	if startData["tool"] != "datetime" || endData["tool"] != "datetime" {
		t.Fatalf("tool_execution tool = %v / %v, want datetime", startData["tool"], endData["tool"])
	}
	if startData["phase"] != model.ToolPhaseRunning || endData["phase"] != model.ToolPhaseDone {
		t.Fatalf("tool_execution phases = %v / %v, want running / done", startData["phase"], endData["phase"])
	}
	if startData["id"] != "call_ws_1" || endData["id"] != "call_ws_1" {
		t.Fatalf("tool_execution ids = %v / %v, want the provider call id", startData["id"], endData["id"])
	}
	for typ, data := range map[string]map[string]any{"start": startData, "end": endData} {
		if _, ok := data["delta"]; ok {
			t.Fatalf("tool_execution.%s carries a delta key; phase must not ride in delta: %v", typ, data)
		}
		if _, ok := data["streamId"]; ok {
			t.Fatalf("tool_execution.%s carries a streamId; tool execution is not a model stream: %v", typ, data)
		}
	}
}

// dataMap 取帧 data 的映射形态，形状不符时直接失败。
func dataMap(t *testing.T, ev wsEvent) map[string]any {
	t.Helper()
	data, ok := ev.Data.(map[string]any)
	if !ok {
		t.Fatalf("frame %s data = %+v, want a map", ev.Type, ev.Data)
	}
	return data
}

// TestStreamFramesToolExecutionErrorPhase 是失败路径的帧形状锁定：
// 直接经生产映射路径（broadcastStream，即 relay 消费 ws.stream 的
// 同一入口）广播一个 error phase 的 end 事件，验证帧携带
// phase=error（裁决 03.4.4：前端/流程图据此把 tool 节点翻成
// error 态）。真实的失败产生路径由 loop 包的
// TestRun_ToolExecutionErrorPhase 端到端覆盖。
func TestStreamFramesToolExecutionErrorPhase(t *testing.T) {
	t.Parallel()
	hb, c := newEventsStreamHub(t)

	hb.broadcastStream(model.StreamEvent{
		Kind:      model.StreamKindToolExecutionEnd,
		Tool:      "datetime",
		Phase:     model.ToolPhaseError,
		ID:        "call_ws_1",
		SessionID: "web",
		TurnID:    "turn-err",
	})
	ev := readFakeFrame(t, c)
	if ev.Type != "tool_execution.end" {
		t.Fatalf("frame type = %q, want tool_execution.end (frame=%+v)", ev.Type, ev)
	}
	data := dataMap(t, ev)
	if data["phase"] != model.ToolPhaseError || data["tool"] != "datetime" {
		t.Fatalf("frame data = %+v, want phase=error tool=datetime", data)
	}
}