Cyrene_v2/internal/services/gateway/ws_frame_scope_test.go

package gateway

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

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// These tests lock the design/11 §7.1 frame gate: a viewer connection may follow
// the session state machine but must never receive content, while owner and admin
// keep both classes. frameScopeCases is the design/11 §7.1 table in code — every
// frame type named there, published on the subject that really produces it — so a
// classification mistake (a content frame made public, or a state signal made
// private) fails here instead of leaking silently.

// frameScopeCase is one row of the design/11 §7.1 table.
type frameScopeCase struct {
	typ     string // downstream frame type
	subject string // bus subject that produces it
	payload any    // payload published on that subject
	public  bool   // 「状态机信号」: no content → every authenticated connection
}

// scopeSentinel is published last: the bus runs one worker per subscription
// sequentially, so a connection that has seen the sentinel has already been
// offered every frame published before it. That turns "the viewer received nothing
// else" into an exact count instead of a guessed idle window.
const scopeSentinel = EventSessionIdle

// The batch pacing used by publishFrameScopeCases, kept well under the hub's
// wsSendBuffer depth so a burst never trips the slow-client drop.
const (
	scopePaceFrames = 4
	scopePacePause  = 50 * time.Millisecond
)

// scopeRead is the outcome of one connection's reader goroutine.
type scopeRead struct {
	role  auth.Role
	types []string
	err   error
}

// frameScopeCases returns the design/11 §7.1 table. Relayed frames use their real
// ws.<type> subject; stream frames publish the real model.StreamEvent the loop
// emits on ws.stream, so the test drives the production path
// bus → relay → broadcast instead of a hand-built frame.
func frameScopeCases() []frameScopeCase {
	relayed := func(typ string, public bool) frameScopeCase {
		return frameScopeCase{typ: typ, subject: wsPrefix + typ, payload: map[string]any{"sessionId": "web"}, public: public}
	}
	streamed := func(kind model.StreamKind, typ string) frameScopeCase {
		return frameScopeCase{
			typ:     typ,
			subject: wsStreamSubject,
			payload: model.StreamEvent{Kind: kind, Delta: "x", SessionID: "web", TurnID: "turn-scope"},
		}
	}
	return []frameScopeCase{
		// 状态机信号（白名单，不含任何内容）→ 所有已认证连接。
		relayed(EventSessionBusy, true),
		relayed(EventPendingAck, true),
		relayed(EventPendingDrain, true),
		relayed(EventSessionIdle, true),
		relayed(EventSessionInterrupt, true),
		relayed(EventPendingFull, true),
		relayed(wsSessionRead, true),

		// 流式内容（含角色回复正文与推理内容）→ 仅 owner/admin。
		streamed(model.StreamKindTextStart, "message.start"),
		streamed(model.StreamKindTextDelta, "message.delta"),
		streamed(model.StreamKindTextEnd, "message.end"),
		streamed(model.StreamKindThinkingStart, "thinking.start"),
		streamed(model.StreamKindThinkingDelta, "thinking.delta"),
		streamed(model.StreamKindThinkingEnd, "thinking.end"),
		streamed(model.StreamKindToolCallStart, "tool_call.start"),
		streamed(model.StreamKindToolCallDelta, "tool_call.delta"),
		streamed(model.StreamKindToolCallEnd, "tool_call.end"),
		streamed(model.StreamKindToolExecution, "tool_execution.update"),
		relayed("stream_chunk", false),
		relayed("stream_end", false),
		relayed("thinking", false),

		// 持久化内容 → 仅 owner/admin。
		relayed("message_persisted", false),
		relayed("message_edited", false),
		relayed("message_deleted", false),
		relayed("channel_persisted", false),

		// 系统内部 → 仅 owner/admin。
		relayed("graph_update", false),
		relayed("subagent_event", false),
		relayed("subagent.result", false),
		relayed("task_output", false),
		relayed("task_done", false),
		relayed("automation_result", false),
		relayed("tool_used", false),
		relayed("capacity_awaken", false),
		relayed("mode_switch", false),

		// 私有状态 → 仅 owner/admin。
		relayed("mood", false),
		relayed("device_update", false),
		relayed("reminder", false),
	}
}

// TestWSFrameScopePerRole drives the whole design/11 §7.1 table through real
// WebSocket connections with real role tokens: owner and admin receive every frame
// of the table, while the viewer receives exactly the whitelisted state signals, in
// publish order. The viewer assertions the design calls out by name (message.delta /
// thinking.delta / message_persisted / graph_update / mood denied, session_idle
// delivered) are re-stated explicitly so the contract is readable in the test.
func TestWSFrameScopePerRole(t *testing.T) {
	t.Parallel()
	cases := frameScopeCases()
	env := newEnv(t, nil, false)

	conns := map[auth.Role]*websocket.Conn{}
	for _, role := range []auth.Role{auth.RoleOwner, auth.RoleAdmin, auth.RoleViewer} {
		conn := dialWS(t, env.ts.URL, env.token(t, role))
		defer conn.Close()
		conns[role] = conn
	}
	waitConnCount(t, env.srv, len(conns))

	all := append(frameTypes(cases, false), scopeSentinel)
	public := append(frameTypes(cases, true), scopeSentinel)
	wants := map[auth.Role]int{
		auth.RoleOwner:  len(all),
		auth.RoleAdmin:  len(all),
		auth.RoleViewer: len(public),
	}

	// The clients drain on their own goroutines while the batch is published: the
	// hub drops frames for a client whose send buffer is full (its slow-client
	// policy), so a client that only started reading afterwards would lose the tail.
	results := make(chan scopeRead, len(conns))
	for role, want := range wants {
		role, want := role, want
		go func() {
			types, err := readFrameTypes(conns[role], want)
			results <- scopeRead{role: role, types: types, err: err}
		}()
	}
	publishFrameScopeCases(t, env, cases)

	got := make(map[auth.Role][]string, len(conns))
	for range conns {
		res := <-results
		if res.err != nil {
			t.Fatalf("%s reader: %v", res.role, res.err)
		}
		got[res.role] = res.types
	}

	for _, role := range []auth.Role{auth.RoleOwner, auth.RoleAdmin} {
		assertFrameTypes(t, string(role), got[role], all)
	}
	gotViewer := got[auth.RoleViewer]
	assertFrameTypes(t, string(auth.RoleViewer), gotViewer, public)
	for _, denied := range []string{"message.delta", "thinking.delta", "message_persisted", "graph_update", "mood"} {
		if containsFrameType(gotViewer, denied) {
			t.Fatalf("viewer received %q; it must not (design/11 §7.1)", denied)
		}
	}
	if !containsFrameType(gotViewer, EventSessionIdle) {
		t.Fatalf("viewer missing %q, a whitelisted state signal (got %v)", EventSessionIdle, gotViewer)
	}
}

// TestHubBroadcastScopeMatrix is the per-frame-type × per-role matrix, driven
// directly through the delivery entry so every row of the table is asserted in
// isolation (a leaked frame would otherwise be hard to attribute). claims are set
// explicitly per role, so a default-deny regression shows up as the viewer losing
// the frame it is entitled to.
func TestHubBroadcastScopeMatrix(t *testing.T) {
	t.Parallel()
	hb := newTestHub(nil, nil)
	owner := newFakeConn(hb, 4)
	owner.claims = &auth.Claims{Role: auth.RoleOwner}
	admin := newFakeConn(hb, 4)
	admin.claims = &auth.Claims{Role: auth.RoleAdmin}
	viewer := newFakeConn(hb, 4)
	viewer.claims = &auth.Claims{Role: auth.RoleViewer}
	for _, c := range []*wsConn{owner, admin, viewer} {
		hb.register(c)
	}

	for _, tc := range frameScopeCases() {
		tc := tc
		t.Run(tc.typ, func(t *testing.T) {
			hb.broadcast(wsEvent{Type: tc.typ, SessionID: "web"})
			for _, c := range []*wsConn{owner, admin} {
				if ev := readFakeFrame(t, c); ev.Type != tc.typ {
					t.Fatalf("owner/admin frame = %q, want %q", ev.Type, tc.typ)
				}
			}
			if tc.public {
				if ev := readFakeFrame(t, viewer); ev.Type != tc.typ {
					t.Fatalf("viewer frame = %q, want %q (whitelisted state signal)", ev.Type, tc.typ)
				}
				return
			}
			if err := connHasNoFrame(viewer); err != nil {
				t.Fatalf("viewer received the restricted frame %q: %v", tc.typ, err)
			}
		})
	}
}

// TestFrameScopeAllowsExactlyTheStateSignals pins the two properties that keep the
// gate from drifting: an unclassified frame type is never public (so a frame added
// by a later milestone cannot leak by omission), and the public set is exactly the
// state signals design/11 §7.1 registers — widening it means editing the design, the
// gate and this table together.
func TestFrameScopeAllowsExactlyTheStateSignals(t *testing.T) {
	t.Parallel()
	if frameReachesAll("frame_type_nobody_classified") {
		t.Fatal("an unclassified frame type reached every connection; the default must be owner/admin only")
	}
	public := make([]string, 0, len(frameScopeCases()))
	for _, tc := range frameScopeCases() {
		if got := frameReachesAll(tc.typ); got != tc.public {
			t.Fatalf("frame %q: gate says public=%v, design/11 §7.1 says public=%v", tc.typ, got, tc.public)
		}
		if tc.public {
			public = append(public, tc.typ)
		}
	}
	want := []string{
		EventSessionBusy,
		EventPendingAck,
		EventPendingDrain,
		EventSessionIdle,
		EventSessionInterrupt,
		EventPendingFull,
		wsSessionRead,
	}
	sort.Strings(public)
	sort.Strings(want)
	assertFrameTypes(t, "design/11 §7.1 whitelist", public, want)
}

// frameTypes returns the table's frame types in table order (all of them, or only
// the public ones).
func frameTypes(cases []frameScopeCase, publicOnly bool) []string {
	out := make([]string, 0, len(cases))
	for _, tc := range cases {
		if publicOnly && !tc.public {
			continue
		}
		out = append(out, tc.typ)
	}
	return out
}

// publishFrameScopeCases publishes every row on its production subject, followed by
// the sentinel that closes the batch. The batch is paced: a client whose send buffer
// is full has its frames dropped (the hub's documented slow-client policy), so a
// 37-frame burst would exercise the drop policy rather than the gate.
func publishFrameScopeCases(t *testing.T, env *testEnv, cases []frameScopeCase) {
	t.Helper()
	for i, tc := range cases {
		if err := env.bus.Publish(context.Background(), tc.subject, tc.payload); err != nil {
			t.Fatalf("publish %s: %v", tc.subject, err)
		}
		if (i+1)%scopePaceFrames == 0 {
			time.Sleep(scopePacePause)
		}
	}
	if err := env.bus.Publish(context.Background(), wsPrefix+scopeSentinel, map[string]any{"sessionId": "web"}); err != nil {
		t.Fatalf("publish sentinel: %v", err)
	}
}

// readFrameTypes reads exactly want frames and returns their types in arrival order.
// It reports an error rather than failing the test so it can run on the reader
// goroutines that drain each client while the batch is published.
func readFrameTypes(conn *websocket.Conn, want int) ([]string, error) {
	got := make([]string, 0, want)
	for len(got) < want {
		_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			return got, fmt.Errorf("read frame %d/%d: %w (got %v)", len(got)+1, want, err, got)
		}
		var ev wsEvent
		if err := json.Unmarshal(msg, &ev); err != nil {
			return got, fmt.Errorf("decode frame %s: %w", msg, err)
		}
		got = append(got, ev.Type)
	}
	return got, nil
}

// assertFrameTypes compares the received frame types to the expected sequence.
func assertFrameTypes(t *testing.T, what string, got, want []string) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("%s frame count = %d, want %d (got %v)", what, len(got), len(want), got)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("%s frame[%d] = %q, want %q (got %v)", what, i, got[i], want[i], got)
		}
	}
}

// containsFrameType reports whether typ is among the received frame types.
func containsFrameType(types []string, typ string) bool {
	for _, got := range types {
		if got == typ {
			return true
		}
	}
	return false
}