Cyrene_v2/internal/services/gateway/gateway_test.go

package gateway

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	capreg "git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

const testSecret = "0123456789abcdef0123456789abcdef"

// stubRunner is a Chatter whose Run returns a fixed result.
type stubRunner struct {
	res *loop.Result
	err error
}

func (s *stubRunner) Run(ctx context.Context, sessionID string, messages []model.Message) (*loop.Result, error) {
	return s.res, s.err
}

// stubTool is a minimal sdk.Tool for capability-registry tests.
type stubTool struct{}

func (stubTool) Name() string                                       { return "echo" }
func (stubTool) Description() string                                { return "echoes input back" }
func (stubTool) Parameters() (map[string]any, error)                { return nil, nil }
func (stubTool) Execute(ctx context.Context, args any) (any, error) { return args, nil }

// testEnv bundles everything a test needs.
type testEnv struct {
	srv  *Server
	ts   *httptest.Server
	auth *auth.JWTService
	bus  bus.Bus
	caps *capreg.Registry
}

// newEnv builds a Server backed by a real JWT service and in-process bus, plus a
// httptest server. runner may be nil (chat returns 500).
func newEnv(t *testing.T, runner Chatter, withCaps bool) *testEnv {
	t.Helper()
	jwt := auth.New(config.AuthConfig{JWTSecret: testSecret})
	b := bus.New()
	cfg := &config.Config{Auth: config.AuthConfig{AdminUser: "admin", AdminPass: "secret", JWTSecret: testSecret}}
	deps := Deps{Auth: jwt, Bus: b, Loop: runner}
	if withCaps {
		caps := capreg.New()
		if err := caps.RegisterTool(stubTool{}); err != nil {
			t.Fatalf("register tool: %v", err)
		}
		deps.Registry = caps
	}
	srv, err := New(cfg, deps)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	ts := httptest.NewServer(srv.Handler())
	t.Cleanup(ts.Close)
	env := &testEnv{srv: srv, ts: ts, auth: jwt, bus: b}
	if withCaps {
		env.caps = deps.Registry.(*capreg.Registry)
	}
	return env
}

func (e *testEnv) token(t *testing.T, role auth.Role) string {
	t.Helper()
	tok, err := e.auth.Sign("acct", role, time.Hour)
	if err != nil {
		t.Fatalf("sign token: %v", err)
	}
	return tok
}

// doReq performs an HTTP request and returns the response plus the raw body.
func doReq(t *testing.T, method, url string, body any, token string) (*http.Response, []byte) {
	t.Helper()
	var rd io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			t.Fatalf("marshal body: %v", err)
		}
		rd = bytes.NewReader(b)
	}
	req, err := http.NewRequest(method, url, rd)
	if err != nil {
		t.Fatalf("new request: %v", err)
	}
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("do request: %v", err)
	}
	defer resp.Body.Close()
	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read response: %v", err)
	}
	return resp, bodyBytes
}

// decodeBody decodes a response body into a generic map.
func decodeBody(t *testing.T, b []byte) map[string]json.RawMessage {
	t.Helper()
	var m map[string]json.RawMessage
	if err := json.Unmarshal(b, &m); err != nil {
		t.Fatalf("decode body %s: %v", b, err)
	}
	return m
}

// numOrZero returns the numeric value of a JSON number in the data envelope.
func wantErrorShape(t *testing.T, b []byte) {
	t.Helper()
	env := decodeBody(t, b)
	var eobj struct {
		Code    string `json:"code"`
		Message string `json:"message"`
	}
	if err := json.Unmarshal(env["error"], &eobj); err != nil {
		t.Fatalf("error envelope: %v (body=%s)", err, b)
	}
	if eobj.Code == "" || eobj.Message == "" {
		t.Fatalf("error envelope missing code/message: %s", b)
	}
}

func TestHealth(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, false)
	resp, body := doReq(t, http.MethodGet, env.ts.URL+"/api/v1/health", nil, "")
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200", resp.StatusCode)
	}
	data := string(decodeBody(t, body)["data"])
	if !strings.Contains(data, `"ok"`) {
		t.Fatalf("health data = %s, want ok", data)
	}
}

func TestLoginIssuesJWT(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, false)
	resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/auth/login", map[string]any{"username": "admin", "password": "secret"}, "")
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body)
	}
	var res struct {
		Token string `json:"token"`
		Role  string `json:"role"`
	}
	if err := json.Unmarshal(decodeBody(t, body)["data"], &res); err != nil {
		t.Fatalf("decode login data: %v", err)
	}
	if res.Token == "" {
		t.Fatal("token is empty")
	}
	if res.Role != "owner" {
		t.Fatalf("role = %q, want owner", res.Role)
	}
}

func TestLoginRejectsBadCredentials(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, false)
	resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/auth/login", map[string]any{"username": "admin", "password": "nope"}, "")
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("status = %d, want 401 (body=%s)", resp.StatusCode, body)
	}
	wantErrorShape(t, body)
}

func TestChatReturnsReply(t *testing.T) {
	t.Parallel()
	res := &loop.Result{Text: "hello world", Steps: 1, StopReason: "stop"}
	env := newEnv(t, &stubRunner{res: res}, false)
	tok := env.token(t, auth.RoleOwner)
	resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{"sessionId": "s1", "text": "hi"}, tok)
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body)
	}
	var resData struct {
		Reply      string `json:"reply"`
		Steps      int    `json:"steps"`
		StopReason string `json:"stopReason"`
	}
	if err := json.Unmarshal(decodeBody(t, body)["data"], &resData); err != nil {
		t.Fatalf("decode chat data: %v", err)
	}
	if resData.Reply != "hello world" {
		t.Fatalf("reply = %q, want hello world", resData.Reply)
	}
	if resData.StopReason != "stop" {
		t.Fatalf("stopReason = %q, want stop", resData.StopReason)
	}
}

func TestChatMissingText(t *testing.T) {
	t.Parallel()
	env := newEnv(t, &stubRunner{}, false)
	tok := env.token(t, auth.RoleOwner)
	resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{"text": "   "}, tok)
	if resp.StatusCode != http.StatusBadRequest {
		t.Fatalf("status = %d, want 400 (body=%s)", resp.StatusCode, body)
	}
	wantErrorShape(t, body)
}

func TestChatUnavailableWhenNoRunner(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, false)
	tok := env.token(t, auth.RoleOwner)
	resp, _ := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{"text": "hi"}, tok)
	if resp.StatusCode != http.StatusInternalServerError {
		t.Fatalf("status = %d, want 500", resp.StatusCode)
	}
}

func TestListSessionsAndPersonasEmpty(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, true)
	tok := env.token(t, auth.RoleOwner)
	for _, path := range []string{"/api/v1/sessions", "/api/v1/personas", "/api/v1/admin/sessions", "/api/v1/admin/personas", "/api/v1/admin/memories", "/api/v1/admin/iot"} {
		resp, body := doReq(t, http.MethodGet, env.ts.URL+path, nil, tok)
		if resp.StatusCode != http.StatusOK {
			t.Fatalf("%s status = %d, want 200 (body=%s)", path, resp.StatusCode, body)
		}
		if got := string(decodeBody(t, body)["data"]); got != "[]" {
			t.Fatalf("%s data = %s, want []", path, got)
		}
	}
}

func TestListPluginsFromRegistry(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, true)
	tok := env.token(t, auth.RoleOwner)
	resp, body := doReq(t, http.MethodGet, env.ts.URL+"/api/v1/admin/plugins", nil, tok)
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body)
	}
	var items []struct {
		Name string `json:"name"`
		Kind string `json:"kind"`
	}
	if err := json.Unmarshal(decodeBody(t, body)["data"], &items); err != nil {
		t.Fatalf("decode plugins: %v", err)
	}
	found := false
	for _, it := range items {
		if it.Name == "echo" && it.Kind == "tool" {
			found = true
		}
	}
	if !found {
		t.Fatalf("echo/tool not in plugins: %+v", items)
	}
}

func TestRBACForbiddenAndUnauthorized(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, true)
	owner := env.token(t, auth.RoleOwner)
	viewer := env.token(t, auth.RoleViewer)

	cases := []struct {
		name  string
		url   string
		token string
		want  int
	}{
		{name: "chat no token", url: "/api/v1/chat", token: "", want: http.StatusUnauthorized},
		{name: "chat viewer", url: "/api/v1/chat", token: viewer, want: http.StatusForbidden},
		{name: "admin no token", url: "/api/v1/admin/plugins", token: "", want: http.StatusUnauthorized},
		// A viewer reads the monitoring surface (/admin/plugins) but not the
		// restricted reads (/admin/memories) — gateway/rbac_viewer_test.go
		// enumerates both sets route by route.
		{name: "admin plugin viewer ok", url: "/api/v1/admin/plugins", token: viewer, want: http.StatusOK},
		{name: "admin memories viewer forbidden", url: "/api/v1/admin/memories", token: viewer, want: http.StatusForbidden},
		{name: "admin owner ok", url: "/api/v1/admin/plugins", token: owner, want: http.StatusOK},
		{name: "sessions viewer ok", url: "/api/v1/sessions", token: viewer, want: http.StatusOK},
	}
	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			method := http.MethodGet
			if tc.url == "/api/v1/chat" {
				method = http.MethodPost
			}
			resp, body := doReq(t, method, env.ts.URL+tc.url, map[string]any{"text": "hi"}, tc.token)
			if resp.StatusCode != tc.want {
				t.Fatalf("status = %d, want %d (body=%s)", resp.StatusCode, tc.want, body)
			}
			if tc.want == http.StatusUnauthorized || tc.want == http.StatusForbidden {
				wantErrorShape(t, body)
			}
		})
	}
}

func TestAdminSetPolicyAndAudit(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, true)
	owner := env.token(t, auth.RoleOwner)
	resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/admin/plugins/policy", map[string]any{"name": "echo", "allowInGroup": true}, owner)
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body)
	}
	data := string(decodeBody(t, body)["data"])
	if !strings.Contains(data, `"allowInGroup":true`) {
		t.Fatalf("policy data = %s, want allowInGroup true", data)
	}

	// Unknown capability → 404.
	resp2, _ := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/admin/plugins/policy", map[string]any{"name": "nope"}, owner)
	if resp2.StatusCode != http.StatusNotFound {
		t.Fatalf("status = %d, want 404", resp2.StatusCode)
	}

	// Viewer is forbidden from the write.
	viewer := env.token(t, auth.RoleViewer)
	resp3, _ := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/admin/plugins/policy", map[string]any{"name": "echo"}, viewer)
	if resp3.StatusCode != http.StatusForbidden {
		t.Fatalf("viewer status = %d, want 403", resp3.StatusCode)
	}
}

func TestWSUpgradeRejectsNoToken(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, false)
	wsURL := "ws" + strings.TrimPrefix(env.ts.URL, "http") + "/ws/chat"
	_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err == nil {
		t.Fatal("dial without token should fail")
	}
	if resp == nil || resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("status = %d, want 401", resp.StatusCode)
	}
}

func TestWSUpgradeAndBroadcast(t *testing.T) {
	t.Parallel()
	env := newEnv(t, nil, false)
	tok := env.token(t, auth.RoleOwner)
	wsURL := "ws" + strings.TrimPrefix(env.ts.URL, "http") + "/ws/chat?token=" + tok

	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer conn.Close()
	_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))

	// Round-trip a mode_switch first: the relay re-broadcasts it, which
	// guarantees the connection is registered before we trigger a broadcast.
	if err := conn.WriteJSON(wsInbound{Type: "mode_switch", Data: json.RawMessage(`{"mode":"chat"}`)}); err != nil {
		t.Fatalf("write mode_switch: %v", err)
	}
	if got := readWS(t, conn); got.Type != "mode_switch" {
		t.Fatalf("first frame type = %q, want mode_switch (frame=%+v)", got.Type, got)
	}

	// Publish an outbound event; the relay must forward it to the client.
	if err := env.bus.Publish(context.Background(), "ws.thinking", map[string]any{"text": "hmm"}); err != nil {
		t.Fatalf("publish: %v", err)
	}
	if got := readWS(t, conn); got.Type != "thinking" {
		t.Fatalf("broadcast type = %q, want thinking (frame=%+v)", got.Type, got)
	}
}

func TestWsEventSessionIDField(t *testing.T) {
	t.Parallel()
	// When set, sessionId is marshalled so the client can filter by session.
	set, err := json.Marshal(wsEvent{Type: "stream_end", SessionID: "web", Data: map[string]any{"stopReason": "stop"}})
	if err != nil {
		t.Fatalf("marshal set: %v", err)
	}
	if !strings.Contains(string(set), `"sessionId":"web"`) {
		t.Fatalf("set frame = %s, want sessionId web", set)
	}
	// When empty, omitempty drops it so existing clients/tests see no change.
	unset, err := json.Marshal(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": "stop"}})
	if err != nil {
		t.Fatalf("marshal unset: %v", err)
	}
	if strings.Contains(string(unset), "sessionId") {
		t.Fatalf("unset frame = %s, want sessionId omitted", unset)
	}
	var ev wsEvent
	if err := json.Unmarshal(unset, &ev); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if ev.SessionID != "" {
		t.Fatalf("decoded SessionID = %q, want empty", ev.SessionID)
	}
}

// readWS reads one text frame and decodes it into a wsEvent.
func readWS(t *testing.T, conn *websocket.Conn) wsEvent {
	t.Helper()
	_, msg, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("read ws: %v", err)
	}
	var ev wsEvent
	if err := json.Unmarshal(msg, &ev); err != nil {
		t.Fatalf("decode ws frame %s: %v", msg, err)
	}
	return ev
}