Cyrene_v2/internal/services/gateway/ws_bridge_subagent_result_test.go

package gateway

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

	"github.com/gorilla/websocket"

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

// 本文件锁住 02.4.3 的结果回填链路：编排层在
// orchestration.SubagentResultSubject 上发布的后台批次结果，必须经
// ws_bridge 的中继变成 `subagent_result` 下行帧、并**按父会话归属**
// （帧壳 sessionId），从而出现在父会话的实时流里（design/05 §9）；
// 它是内容帧，因此只发 owner + admin（design/11 §7.1，fail-closed）。
//
// 断言走真路径：真 bus Publish → hub 的真订阅与中继 → 真 /ws/chat 连接
// （真 JWT）。viewer 的负数用例（收不到）与 owner/admin 的正数用例
// （收得到 + 载荷形状）在同一次发布上一起判定。

// subagentResultReadWindow 是排空一条连接时的单帧等待窗口。广播是异步的
// （bus worker → relay → hub → 连接写协程），窗口内应收到全部相关帧；
// 读超时即视为排空完毕。
const subagentResultReadWindow = 1500 * time.Millisecond

// subagentFrame 是测试侧的线上帧视图（等价于客户端解码）：帧壳 type/sessionId 直接
// 解出，data 保留原始 JSON，使断言既能按形状解码、又能锁住键集。
type subagentFrame struct {
	Type      string `json:"type"`
	SessionID string `json:"sessionId"`
	// TurnID 是帧壳上的回合标识（design/15 §2.5.9）：结构化流帧、`turn_begin` 与
	// 带帧壳 turnId 的 `steer_ack` 才有它，其余帧为空。
	TurnID string          `json:"turnId"`
	Data   json.RawMessage `json:"data"`
}

// subagentFrameData 是 api-client §3.10 登记的 data 形状。
type subagentFrameData struct {
	SessionID string `json:"sessionId"`
	Status    string `json:"status"`
	Summary   string `json:"summary"`
	Agents    []struct {
		ID      string `json:"id"`
		Goal    string `json:"goal"`
		Role    string `json:"role"`
		Mode    string `json:"mode"`
		Status  string `json:"status"`
		Summary string `json:"summary"`
		Text    string `json:"text"`
		Steps   int    `json:"steps"`
	} `json:"agents"`
}

// TestSubagentResultReachesOwnerAndAdminOnly 是后台结果回填的真 WS 帧断言：
// 真 bus 上发布编排层的真实载荷类型 → owner/admin 收到 `subagent_result`
// 帧（带父会话与逐代理结果），viewer **收不到**（内容帧默认拒）。
func TestSubagentResultReachesOwnerAndAdminOnly(t *testing.T) {
	t.Parallel()
	// 新帧属内容帧：不在那 7 个状态机信号白名单里，必须落 fail-closed 的
	// owner+admin 档（design/11 §7.1「要放进白名单必须有明确理由并在本表登记」）。
	if frameReachesAll(wsSubagentResult) {
		t.Fatalf("%s reached every connection; a result frame carries content and must stay owner/admin only", wsSubagentResult)
	}

	const parentSessionID = "web:subagent-dispatch"
	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))

	// 编排层发布的真实载荷（与 SpawnTool.await 同类型同键集）。
	payload := &orchestration.SubagentResultEvent{
		SessionID: parentSessionID,
		SpawnResult: &orchestration.SpawnResult{
			Status:  "done",
			Summary: "子代理汇总: 共 1 个, 1 成功, 0 失败, 0 中止",
			Agents: []orchestration.SpawnAgentResult{{
				ID: "sub-a1", Goal: "调研 X", Role: "cyrene", Mode: "engineer",
				Status: "done", Summary: "完成 (步骤=2): 结论",
				Text: "结论", Steps: 2,
			}},
		},
	}
	if err := env.bus.Publish(context.Background(), orchestration.SubagentResultSubject, payload); err != nil {
		t.Fatalf("publish %s: %v", orchestration.SubagentResultSubject, err)
	}
	// 哨兵：一个白名单状态信号收尾。它让每条连接的「本批帧已到齐」从猜测的
	// 空闲窗口变成可判定的边界；viewer 也必定收到它，故「viewer 什么都没收到」
	// 不会与「连接根本没工作」混淆。
	if err := env.bus.Publish(context.Background(), wsPrefix+scopeSentinel, map[string]any{"sessionId": parentSessionID}); err != nil {
		t.Fatalf("publish sentinel: %v", err)
	}

	for _, role := range []auth.Role{auth.RoleOwner, auth.RoleAdmin} {
		frames := drainFrames(t, conns[role])
		assertFrameCount(t, string(role), frames, wsSubagentResult, 1)
		assertFrameCount(t, string(role), frames, scopeSentinel, 1)
		if len(frames) != 2 {
			t.Fatalf("%s received %d frames, want exactly the result + the sentinel: %v", role, len(frames), frameTypesOf(frames))
		}
		assertSubagentResultFrame(t, string(role), findFrame(t, string(role), frames, wsSubagentResult), parentSessionID, payload)
	}

	viewerFrames := drainFrames(t, conns[auth.RoleViewer])
	if len(viewerFrames) != 1 || viewerFrames[0].Type != scopeSentinel {
		t.Fatalf("viewer received %v, want only the whitelisted %q: the result frame is owner/admin only (design/11 §7.1)",
			frameTypesOf(viewerFrames), scopeSentinel)
	}
}

// assertSubagentResultFrame 按 api-client §3.10 登记的帧形状逐项断言一帧
// `subagent_result`：帧壳 sessionId 必须是被回填的父会话，data 的**键集**必须
// 恰好是登记的那四个，其汇总与逐代理字段必须与该批次发布的载荷一致。
func assertSubagentResultFrame(t *testing.T, role string, frame subagentFrame, parentSessionID string, payload *orchestration.SubagentResultEvent) {
	t.Helper()
	if frame.SessionID != parentSessionID {
		t.Fatalf("%s frame sessionId = %q, want the parent session %q (the result must land in that session's stream)",
			role, frame.SessionID, parentSessionID)
	}
	wantKeys := []string{"sessionId", "status", "summary", "agents"}
	var keys map[string]json.RawMessage
	if err := json.Unmarshal(frame.Data, &keys); err != nil {
		t.Fatalf("%s frame data is not a JSON object: %v (%s)", role, err, frame.Data)
	}
	if len(keys) != len(wantKeys) {
		t.Fatalf("%s frame data keys = %v, want exactly %v (api-client §3.10)", role, keys, wantKeys)
	}
	for _, k := range wantKeys {
		if _, ok := keys[k]; !ok {
			t.Fatalf("%s frame data lacks key %q (api-client §3.10): %v", role, k, keys)
		}
	}
	var data subagentFrameData
	if err := json.Unmarshal(frame.Data, &data); err != nil {
		t.Fatalf("%s frame data = %s: %v", role, frame.Data, err)
	}
	if data.SessionID != parentSessionID || data.Status != "done" {
		t.Fatalf("%s frame data = %+v, want the parent session and a settled batch", role, data)
	}
	if data.Summary != payload.Summary {
		t.Fatalf("%s frame summary = %q, want %q", role, data.Summary, payload.Summary)
	}
	if len(data.Agents) != 1 {
		t.Fatalf("%s frame agents = %+v, want one entry", role, data.Agents)
	}
	if got := data.Agents[0]; got.ID != "sub-a1" || got.Goal != "调研 X" || got.Role != "cyrene" ||
		got.Mode != "engineer" || got.Status != "done" || got.Text != "结论" || got.Steps != 2 {
		t.Fatalf("%s frame agent = %+v, want the dispatched sub-agent's outcome", role, got)
	}
}

// drainFrames 读取连接直到读超时，返回按到达顺序收到的帧。它只用于
// 已发布的批次之后：第一批帧到达后剩余窗口用于确认「没有别的帧了」，
// 这正是 viewer 负数用例所需的判定。
func drainFrames(t *testing.T, conn *websocket.Conn) []subagentFrame {
	t.Helper()
	var out []subagentFrame
	for {
		_ = conn.SetReadDeadline(time.Now().Add(subagentResultReadWindow))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			return out
		}
		var frame subagentFrame
		if uerr := json.Unmarshal(msg, &frame); uerr != nil {
			t.Fatalf("decode frame %s: %v", msg, uerr)
		}
		out = append(out, frame)
	}
}

// assertFrameCount 断言角色收到的某类型帧恰好 want 条。
func assertFrameCount(t *testing.T, role string, frames []subagentFrame, typ string, want int) {
	t.Helper()
	got := 0
	for _, f := range frames {
		if f.Type == typ {
			got++
		}
	}
	if got != want {
		t.Fatalf("%s received %d %q frames, want %d (all frames: %v)", role, got, typ, want, frameTypesOf(frames))
	}
}

// findFrame 返回角色收到的第一条 typ 帧。
func findFrame(t *testing.T, role string, frames []subagentFrame, typ string) subagentFrame {
	t.Helper()
	for _, f := range frames {
		if f.Type == typ {
			return f
		}
	}
	t.Fatalf("%s received no %q frame (all frames: %v)", role, typ, frameTypesOf(frames))
	return subagentFrame{}
}

// frameTypesOf 列出收到的帧类型，用于失败信息。
func frameTypesOf(frames []subagentFrame) []string {
	out := make([]string, 0, len(frames))
	for _, f := range frames {
		out = append(out, f.Type)
	}
	return out
}