diff --git a/internal/services/gateway/extra_test.go b/internal/services/gateway/extra_test.go new file mode 100644 index 0000000..985d9de --- /dev/null +++ b/internal/services/gateway/extra_test.go @@ -0,0 +1,1002 @@ +package gateway + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "testing/fstest" + "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/channel/adapter" + "git.yeij.top/AskaEth/Cyrene_v2/internal/channel/session" + "git.yeij.top/AskaEth/Cyrene_v2/internal/config" + "git.yeij.top/AskaEth/Cyrene_v2/internal/iot" + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop" + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model" + "git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk" +) + +// cfgDefault builds the minimum config used across tests. +func cfgDefault() *config.Config { + return &config.Config{Auth: config.AuthConfig{AdminUser: "admin", AdminPass: "secret", JWTSecret: testSecret}} +} + +// buildEnv constructs a Server + httptest server from an explicit cfg/deps, +// defaulting Auth to a real JWTService and Bus to a real in-process bus. +func buildEnv(t *testing.T, cfg *config.Config, deps Deps) (*Server, *httptest.Server, *auth.JWTService) { + t.Helper() + if cfg == nil { + cfg = cfgDefault() + } + if deps.Auth == nil { + deps.Auth = auth.New(config.AuthConfig{AdminUser: "admin", AdminPass: "secret", JWTSecret: testSecret}) + } + if deps.Bus == nil { + deps.Bus = bus.New() + } + srv, err := New(cfg, deps) + if err != nil { + t.Fatalf("new server: %v", err) + } + ts := httptest.NewServer(srv.Handler()) + t.Cleanup(ts.Close) + jwt, _ := deps.Auth.(*auth.JWTService) + return srv, ts, jwt +} + +// buildServerOnly constructs a Server without an httptest server (used for +// start/shutdown lifecycle tests). +func buildServerOnly(t *testing.T, cfg *config.Config, deps Deps) *Server { + t.Helper() + srv, err := New(cfg, deps) + if err != nil { + t.Fatalf("new server: %v", err) + } + return srv +} + +// mustJWT builds a real JWTService for a test secret. +func mustJWT(t *testing.T) *auth.JWTService { + t.Helper() + return auth.New(config.AuthConfig{JWTSecret: testSecret}) +} + +// signAny issues a token for sub under the given role via a real JWTService. +func signAny(t *testing.T, jwt *auth.JWTService, sub string, role auth.Role) string { + t.Helper() + tok, err := jwt.Sign(sub, role, time.Hour) + if err != nil { + t.Fatalf("sign: %v", err) + } + return tok +} + +// doRaw performs a request with a raw (already-encoded) body. +func doRaw(t *testing.T, method, url, raw, token string) (*http.Response, []byte) { + t.Helper() + req, err := http.NewRequest(method, url, strings.NewReader(raw)) + 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() + b, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + return resp, b +} + +// --- Mocks ------------------------------------------------------------------ + +// stubAuth wraps a real JWTService but lets Sign/Verify fail on demand. +type stubAuth struct { + jwt *auth.JWTService + signErr error + verifyErr error +} + +func (s *stubAuth) Sign(sub string, role auth.Role, ttl time.Duration) (string, error) { + if s.signErr != nil { + return "", s.signErr + } + return s.jwt.Sign(sub, role, ttl) +} + +func (s *stubAuth) Verify(token string) (*auth.Claims, error) { + if s.verifyErr != nil { + return nil, s.verifyErr + } + return s.jwt.Verify(token) +} + +func (s *stubAuth) Authorize(required ...auth.Role) func(http.Handler) http.Handler { + return s.jwt.Authorize(required...) +} + +// stubCaps mimics CapabilityStore with a fixed list and a forced SetPolicy error. +type stubCaps struct { + list []capreg.Capability + setErr error +} + +func (s *stubCaps) List(kind capreg.Kind) []capreg.Capability { + if kind == "" { + return s.list + } + out := make([]capreg.Capability, 0) + for _, cap := range s.list { + if cap.Kind == kind { + out = append(out, cap) + } + } + return out +} + +func (s *stubCaps) SetPolicy(name string, p capreg.ToolPolicy) error { + if s.setErr != nil { + return s.setErr + } + return nil +} + +// stubSessions mimics SessionStore. +type stubSessions struct { + list []session.Session + err error +} + +func (s *stubSessions) List(ctx context.Context) ([]session.Session, error) { return s.list, s.err } + +// stubPersonas mimics PersonaStore. +type stubPersonas struct { + list []PersonaSummary +} + +func (s *stubPersonas) List() []PersonaSummary { return s.list } + +// stubMemories mimics MemoryStore. +type stubMemories struct { + list []MemorySummary + err error +} + +func (s *stubMemories) List(ctx context.Context) ([]MemorySummary, error) { return s.list, s.err } + +// stubIoT mimics IoTStore. +type stubIoT struct { + list []iot.Device + err error +} + +func (s *stubIoT) List(ctx context.Context) ([]iot.Device, error) { return s.list, s.err } + +// stubChannel implements sdk.Channel. +type stubChannel struct{} + +func (stubChannel) Name() string { return "qq" } +func (stubChannel) Connect(ctx context.Context, c any) error { return nil } +func (stubChannel) Disconnect(ctx context.Context) error { return nil } +func (stubChannel) IsConnected() bool { return true } +func (stubChannel) HealthCheck(ctx context.Context) error { return nil } + +var _ sdk.Channel = stubChannel{} + +// captureRunner records the sessionID each Run call sees. +type captureRunner struct { + got []string + res *loop.Result + err error +} + +func (c *captureRunner) Run(ctx context.Context, sessionID string, messages []model.Message) (*loop.Result, error) { + c.got = append(c.got, sessionID) + return c.res, c.err +} + +// plainWriter is a bare ResponseWriter implementing neither Flusher nor Hijacker. +type plainWriter struct{} + +func (plainWriter) Header() http.Header { return http.Header{} } +func (plainWriter) Write(b []byte) (int, error) { return len(b), nil } +func (plainWriter) WriteHeader(code int) {} + +// --- Server construction / validation -------------------------------------- + +func TestNewValidation(t *testing.T) { + t.Parallel() + jwt := mustJWT(t) + if _, err := New(nil, Deps{Auth: &stubAuth{jwt: jwt}}); err == nil { + t.Fatal("nil config must error") + } + if _, err := New(&config.Config{}, Deps{}); err == nil { + t.Fatal("nil auth dependency must error") + } +} + +// --- Login ------------------------------------------------------------------ + +func TestLoginBadBody(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, false) + resp, body := doRaw(t, http.MethodPost, env.ts.URL+"/api/v1/auth/login", "{", "") + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +func TestLoginSignError(t *testing.T) { + t.Parallel() + jwt := mustJWT(t) + srv, ts, _ := buildEnv(t, cfgDefault(), Deps{Auth: &stubAuth{jwt: jwt, signErr: errors.New("boom")}}) + _ = srv + resp, body := doReq(t, http.MethodPost, ts.URL+"/api/v1/auth/login", map[string]any{"username": "admin", "password": "secret"}, "") + if resp.StatusCode != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +// --- Chat ------------------------------------------------------------------- + +func TestChatBadBody(t *testing.T) { + t.Parallel() + env := newEnv(t, &stubRunner{}, false) + tok := env.token(t, auth.RoleOwner) + resp, body := doRaw(t, http.MethodPost, env.ts.URL+"/api/v1/chat", "{", tok) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +func TestChatStreamMode(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, 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", "stream": true}, tok) + 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, `"stream":true`) { + t.Fatalf("data = %s, want stream true", data) + } +} + +func TestChatDefaultsSessionToWeb(t *testing.T) { + t.Parallel() + cr := &captureRunner{res: &loop.Result{Text: "ok", Steps: 1, StopReason: "stop"}} + env := newEnv(t, cr, false) + tok := env.token(t, auth.RoleOwner) + resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{"text": "hi"}, tok) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body) + } + if len(cr.got) != 1 || cr.got[0] != "web" { + t.Fatalf("sessionIDs = %v, want [web]", cr.got) + } +} + +func TestChatRunnerError(t *testing.T) { + t.Parallel() + env := newEnv(t, &stubRunner{err: errors.New("boom")}, false) + tok := env.token(t, auth.RoleOwner) + resp, body := 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 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +// --- List handlers --------------------------------------------------------- + +func TestListSessionsNonEmpty(t *testing.T) { + t.Parallel() + deps := Deps{Session: &stubSessions{list: []session.Session{{ + Platform: "qq", ChannelID: "c1", SessionID: "s1", + ActorRole: adapter.RoleOwner, Owner: true, IsGroup: false, MemWritable: true, + }}}} + srv, ts, jwt := buildEnv(t, cfgDefault(), deps) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + for _, path := range []string{"/api/v1/sessions", "/api/v1/admin/sessions"} { + resp, body := doReq(t, http.MethodGet, ts.URL+path, nil, tok) + if resp.StatusCode != http.StatusOK { + t.Fatalf("%s status = %d, want 200 (body=%s)", path, resp.StatusCode, body) + } + var items []struct { + Platform string `json:"platform"` + SessionID string `json:"sessionId"` + Owner bool `json:"owner"` + } + if err := json.Unmarshal(decodeBody(t, body)["data"], &items); err != nil { + t.Fatalf("decode %s: %v", path, err) + } + if len(items) != 1 || items[0].Platform != "qq" || items[0].SessionID != "s1" || !items[0].Owner { + t.Fatalf("%s items = %+v, want single qq/s1 owner", path, items) + } + } +} + +func TestListSessionsUserError(t *testing.T) { + t.Parallel() + srv, ts, jwt := buildEnv(t, cfgDefault(), Deps{Session: &stubSessions{err: errors.New("boom")}}) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodGet, ts.URL+"/api/v1/sessions", nil, tok) + if resp.StatusCode != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +func TestListSessionsAdminError(t *testing.T) { + t.Parallel() + srv, ts, jwt := buildEnv(t, cfgDefault(), Deps{Session: &stubSessions{err: errors.New("boom")}}) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodGet, ts.URL+"/api/v1/admin/sessions", nil, tok) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body) + } + if got := string(decodeBody(t, body)["data"]); got != "[]" { + t.Fatalf("data = %s, want []", got) + } +} + +func TestListPersonasNonEmpty(t *testing.T) { + t.Parallel() + deps := Deps{Personas: &stubPersonas{list: []PersonaSummary{{ID: "p1", Name: "苏梦", Mode: "chat"}}}} + srv, ts, jwt := buildEnv(t, cfgDefault(), deps) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + for _, path := range []string{"/api/v1/personas", "/api/v1/admin/personas"} { + resp, body := doReq(t, http.MethodGet, ts.URL+path, nil, tok) + if resp.StatusCode != http.StatusOK { + t.Fatalf("%s status = %d, want 200 (body=%s)", path, resp.StatusCode, body) + } + var items []struct { + ID string `json:"id"` + Name string `json:"name"` + } + if err := json.Unmarshal(decodeBody(t, body)["data"], &items); err != nil { + t.Fatalf("decode %s: %v", path, err) + } + if len(items) != 1 || items[0].Name != "苏梦" { + t.Fatalf("%s items = %+v, want single 苏梦", path, items) + } + } +} + +func TestAdminMemoriesNonEmpty(t *testing.T) { + t.Parallel() + deps := Deps{Memories: &stubMemories{list: []MemorySummary{{ID: "m1", Scope: "project", Text: "note"}}}} + srv, ts, jwt := buildEnv(t, cfgDefault(), deps) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodGet, ts.URL+"/api/v1/admin/memories", nil, tok) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body) + } + var items []struct { + ID string `json:"id"` + Scope string `json:"scope"` + } + if err := json.Unmarshal(decodeBody(t, body)["data"], &items); err != nil { + t.Fatalf("decode: %v", err) + } + if len(items) != 1 || items[0].Scope != "project" { + t.Fatalf("items = %+v, want single project", items) + } +} + +func TestAdminMemoriesErrorStillOK(t *testing.T) { + t.Parallel() + srv, ts, jwt := buildEnv(t, cfgDefault(), Deps{Memories: &stubMemories{err: errors.New("boom")}}) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodGet, ts.URL+"/api/v1/admin/memories", nil, tok) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body) + } + if got := string(decodeBody(t, body)["data"]); got != "[]" { + t.Fatalf("data = %s, want []", got) + } +} + +// --- Admin: IoT ------------------------------------------------------------ + +func TestAdminIoTNonEmpty(t *testing.T) { + t.Parallel() + deps := Deps{IoT: &stubIoT{list: []iot.Device{{ + ID: "d1", Name: "客厅灯", Type: "light", + State: map[string]any{"power": "on"}, Online: true, + }}}} + srv, ts, jwt := buildEnv(t, cfgDefault(), deps) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodGet, ts.URL+"/api/v1/admin/iot", nil, tok) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body) + } + var items []struct { + ID string `json:"id"` + Name string `json:"name"` + Type string `json:"type"` + State map[string]any `json:"state"` + Online bool `json:"online"` + } + if err := json.Unmarshal(decodeBody(t, body)["data"], &items); err != nil { + t.Fatalf("decode: %v", err) + } + if len(items) != 1 || items[0].Name != "客厅灯" || items[0].Type != "light" || !items[0].Online { + t.Fatalf("items = %+v, want single 客厅灯/light/online", items) + } + if items[0].State["power"] != "on" { + t.Fatalf("state = %+v, want power on", items[0].State) + } +} + +func TestAdminIoTError(t *testing.T) { + t.Parallel() + srv, ts, jwt := buildEnv(t, cfgDefault(), Deps{IoT: &stubIoT{err: errors.New("boom")}}) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodGet, ts.URL+"/api/v1/admin/iot", nil, tok) + if resp.StatusCode != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +// --- Admin: channels ------------------------------------------------------- + +func TestAdminChannelsFromRegistry(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, true) + if err := env.caps.RegisterChannel(stubChannel{}); err != nil { + t.Fatalf("register channel: %v", err) + } + // A channel-kind capability whose Value is not an sdk.Channel must be skipped. + if err := env.caps.Register(capreg.Capability{Name: "badchan", Kind: capreg.KindChannel, Value: stubTool{}}); err != nil { + t.Fatalf("register bad channel: %v", err) + } + tok := env.token(t, auth.RoleOwner) + resp, body := doReq(t, http.MethodGet, env.ts.URL+"/api/v1/admin/channels", 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"` + Connected bool `json:"connected"` + } + if err := json.Unmarshal(decodeBody(t, body)["data"], &items); err != nil { + t.Fatalf("decode: %v", err) + } + if len(items) != 1 || items[0].Name != "qq" || !items[0].Connected { + t.Fatalf("items = %+v, want single qq/connected", items) + } +} + +// --- Admin: policy errors -------------------------------------------------- + +func TestAdminSetPolicyUnavailable(t *testing.T) { + t.Parallel() + srv, ts, jwt := buildEnv(t, cfgDefault(), Deps{}) // no Registry + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodPost, ts.URL+"/api/v1/admin/plugins/policy", map[string]any{"name": "echo"}, tok) + if resp.StatusCode != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +func TestAdminSetPolicyOtherError(t *testing.T) { + t.Parallel() + deps := Deps{Registry: &stubCaps{setErr: errors.New("boom")}} + srv, ts, jwt := buildEnv(t, cfgDefault(), deps) + _ = srv + tok := signAny(t, jwt, "acct", auth.RoleOwner) + resp, body := doReq(t, http.MethodPost, ts.URL+"/api/v1/admin/plugins/policy", map[string]any{"name": "echo"}, tok) + if resp.StatusCode != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +func TestAdminSetPolicyBadBody(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, true) + tok := env.token(t, auth.RoleOwner) + resp, body := doRaw(t, http.MethodPost, env.ts.URL+"/api/v1/admin/plugins/policy", "{", tok) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +func TestAdminSetPolicyMissingName(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, true) + tok := env.token(t, auth.RoleOwner) + resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/admin/plugins/policy", map[string]any{"allowInGroup": true}, tok) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +// --- CORS ------------------------------------------------------------------ + +func TestCORSAllowAll(t *testing.T) { + t.Parallel() + cfg := cfgDefault() + cfg.AllowedOrigins = []string{"*"} + srv, ts, _ := buildEnv(t, cfg, Deps{}) + _ = srv + req, _ := http.NewRequest(http.MethodGet, ts.URL+"/api/v1/health", nil) + req.Header.Set("Origin", "http://anywhere.example") + resp := doHandler(t, req) + if got := resp.Header.Get("Access-Control-Allow-Origin"); got != "*" { + t.Fatalf("ACAO = %q, want *", got) + } +} + +func TestCORSWhitelistedOrigin(t *testing.T) { + t.Parallel() + cfg := cfgDefault() + cfg.AllowedOrigins = []string{"http://allowed.example"} + srv, ts, _ := buildEnv(t, cfg, Deps{}) + _ = srv + req, _ := http.NewRequest(http.MethodGet, ts.URL+"/api/v1/health", nil) + req.Header.Set("Origin", "http://allowed.example") + resp := doHandler(t, req) + if got := resp.Header.Get("Access-Control-Allow-Origin"); got != "http://allowed.example" { + t.Fatalf("ACAO = %q, want origin", got) + } +} + +func TestCORSRejectsUnknownOrigin(t *testing.T) { + t.Parallel() + cfg := cfgDefault() + cfg.AllowedOrigins = []string{"http://allowed.example"} + srv, ts, _ := buildEnv(t, cfg, Deps{}) + _ = srv + req, _ := http.NewRequest(http.MethodGet, ts.URL+"/api/v1/health", nil) + req.Header.Set("Origin", "http://evil.example") + resp := doHandler(t, req) + if got := resp.Header.Get("Access-Control-Allow-Origin"); got != "" { + t.Fatalf("ACAO = %q, want empty for unknown origin", got) + } +} + +func TestCORSPreflight(t *testing.T) { + t.Parallel() + cfg := cfgDefault() + cfg.AllowedOrigins = []string{"http://allowed.example"} + srv, ts, _ := buildEnv(t, cfg, Deps{}) + _ = srv + req, _ := http.NewRequest(http.MethodOptions, ts.URL+"/api/v1/chat", nil) + req.Header.Set("Origin", "http://allowed.example") + req.Header.Set("Access-Control-Request-Method", "POST") + resp := doHandler(t, req) + if resp.StatusCode != http.StatusNoContent { + t.Fatalf("status = %d, want 204", resp.StatusCode) + } + if got := resp.Header.Get("Access-Control-Allow-Origin"); got != "http://allowed.example" { + t.Fatalf("ACAO = %q, want origin", got) + } +} + +// --- Recovery + statusRecorder -------------------------------------------- + +func TestRecoverPanicTo500(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, false) + h := env.srv.recover(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + panic("boom") + })) + rr := httptest.NewRecorder() + h.ServeHTTP(rr, httptest.NewRequest(http.MethodGet, "/x", nil)) + if rr.Code != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500", rr.Code) + } + wantErrorShape(t, rr.Body.Bytes()) +} + +func TestStatusRecorderHijackUnsupported(t *testing.T) { + t.Parallel() + rec := &statusRecorder{ResponseWriter: plainWriter{}} + conn, rw, err := rec.Hijack() + if err == nil { + t.Fatal("Hijack on a non-Hijacker must error") + } + if conn != nil || rw != nil { + t.Fatalf("conn/rw = %v/%v, want nil", conn, rw) + } +} + +func TestStatusRecorderFlush(t *testing.T) { + t.Parallel() + // Flusher present → forwarded. + rr := httptest.NewRecorder() + (&statusRecorder{ResponseWriter: rr}).Flush() + if !rr.Flushed { + t.Fatal("Flush should have been forwarded to the underlying Flusher") + } + // No Flusher → no-op, no panic. + (&statusRecorder{ResponseWriter: plainWriter{}}).Flush() +} + +// --- SPA ------------------------------------------------------------------- + +func TestSPA(t *testing.T) { + t.Parallel() + front := fstest.MapFS{ + "index.html": &fstest.MapFile{Data: []byte("root")}, + "assets/app.js": &fstest.MapFile{Data: []byte("var x = 1;")}, + } + srv, ts, _ := buildEnv(t, cfgDefault(), Deps{Frontend: front}) + _ = srv + + // Real asset is served. + resp, body := doReq(t, http.MethodGet, ts.URL+"/assets/app.js", nil, "") + if resp.StatusCode != http.StatusOK || string(body) != "var x = 1;" { + t.Fatalf("asset status = %d body = %q, want 200 with js", resp.StatusCode, body) + } + // Unknown non-API path falls back to index.html. + resp, body = doReq(t, http.MethodGet, ts.URL+"/chat/somewhere", nil, "") + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), "root") { + t.Fatalf("fallback status = %d body = %q, want 200 with index", resp.StatusCode, body) + } + // Root serves index.html. + resp, body = doReq(t, http.MethodGet, ts.URL+"/", nil, "") + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), "root") { + t.Fatalf("root status = %d body = %q, want 200 with index", resp.StatusCode, body) + } + // API route is passed through untouched. + resp, body = doReq(t, http.MethodGet, ts.URL+"/api/v1/health", nil, "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("api status = %d, want 200 (body=%s)", resp.StatusCode, body) + } +} + +// --- WS: validation -------------------------------------------------------- + +func TestWSRequiresGET(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, false) + tok := env.token(t, auth.RoleOwner) + resp, body := doReq(t, http.MethodPost, env.ts.URL+"/ws/chat", nil, tok) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 (body=%s)", resp.StatusCode, body) + } + wantErrorShape(t, body) +} + +func TestWSInvalidJSON(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, false) + tok := env.token(t, auth.RoleOwner) + conn := dialWS(t, env.ts.URL, tok) + defer conn.Close() + _ = conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + if err := conn.WriteMessage(websocket.TextMessage, []byte("{not json")); err != nil { + t.Fatalf("write invalid json: %v", err) + } + ev := readWS(t, conn) + if ev.Type != "error" { + t.Fatalf("type = %q, want error (frame=%+v)", ev.Type, ev) + } +} + +func TestWSAuthViaHeader(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" + hdr := http.Header{} + hdr.Set("Authorization", "Bearer "+tok) + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, hdr) + if err != nil { + t.Fatalf("dial with header: %v (resp=%+v)", err, resp) + } + defer conn.Close() + _ = conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + 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("frame type = %q, want mode_switch", got.Type) + } +} + +func TestWSCloseDeregisters(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, false) + tok := env.token(t, auth.RoleOwner) + conn := dialWS(t, env.ts.URL, tok) + waitConnCount(t, env.srv, 1) + conn.Close() + waitConnCount(t, env.srv, 0) +} + +func TestServerCloseShutsHub(t *testing.T) { + t.Parallel() + env := newEnv(t, nil, false) + tok := env.token(t, auth.RoleOwner) + conn := dialWS(t, env.ts.URL, tok) + waitConnCount(t, env.srv, 1) + env.srv.Close() + waitConnCount(t, env.srv, 0) + conn.Close() +} + +// --- WS: hub unit tests ---------------------------------------------------- + +func TestHubStreamChat(t *testing.T) { + t.Parallel() + hb := newTestHub(&stubRunner{res: &loop.Result{Text: "hi", Steps: 2, StopReason: "stop"}}, nil) + c := newFakeConn(hb, 32) + hb.register(c) + + hb.streamChat(context.Background(), map[string]any{"sessionId": "s1", "text": "hello"}) + if ev := readFakeFrame(t, c); ev.Type != "stream_chunk" { + t.Fatalf("first frame type = %q, want stream_chunk (frame=%+v)", ev.Type, ev) + } + if ev := readFakeFrame(t, c); ev.Type != "stream_end" { + t.Fatalf("second frame type = %q, want stream_end (frame=%+v)", ev.Type, ev) + } +} + +func TestHubStreamChatIdle(t *testing.T) { + t.Parallel() + hb := newTestHub(nil, nil) // nil runner + c := newFakeConn(hb, 32) + hb.register(c) + + hb.streamChat(context.Background(), map[string]any{"sessionId": "s1", "text": ""}) + if ev := readFakeFrame(t, c); ev.Type != "stream_end" { + t.Fatalf("frame type = %q, want stream_end (frame=%+v)", ev.Type, ev) + } +} + +func TestHubStreamChatError(t *testing.T) { + t.Parallel() + hb := newTestHub(&stubRunner{err: errors.New("boom")}, nil) + c := newFakeConn(hb, 32) + hb.register(c) + + hb.streamChat(context.Background(), map[string]any{"sessionId": "s1", "text": "hi"}) + if ev := readFakeFrame(t, c); ev.Type != "stream_end" { + t.Fatalf("frame type = %q, want stream_end (frame=%+v)", ev.Type, ev) + } +} + +func TestHubRelayChatBranch(t *testing.T) { + t.Parallel() + hb := newTestHub(&stubRunner{res: &loop.Result{Text: "hi", Steps: 1, StopReason: "stop"}}, nil) + c := newFakeConn(hb, 32) + hb.register(c) + + if err := hb.relay(context.Background(), bus.Event{Subject: wsChatSubject, Payload: map[string]any{"sessionId": "s1", "text": "hi"}}); err != nil { + t.Fatalf("relay: %v", err) + } + // relay runs streamChat in a goroutine; read with a deadline. + if ev := readFakeFrame(t, c); ev.Type != "stream_chunk" { + t.Fatalf("frame type = %q, want stream_chunk (frame=%+v)", ev.Type, ev) + } + if ev := readFakeFrame(t, c); ev.Type != "stream_end" { + t.Fatalf("frame type = %q, want stream_end (frame=%+v)", ev.Type, ev) + } +} + +func TestHubRoute(t *testing.T) { + t.Parallel() + hb := newTestHub(nil, bus.New()) + c := newFakeConn(hb, 32) + hb.register(c) + + // subagent_ask → subagent_event (enqueueJSON). + c.route(wsInbound{Type: "subagent_ask"}) + if ev := readFakeFrame(t, c); ev.Type != "subagent_event" { + t.Fatalf("subagent frame type = %q, want subagent_event", ev.Type) + } + // task_output_query → task_output (enqueueJSON). + c.route(wsInbound{Type: "task_output_query"}) + if ev := readFakeFrame(t, c); ev.Type != "task_output" { + t.Fatalf("task frame type = %q, want task_output", ev.Type) + } + // unknown → error (enqueueJSON), with non-nil rawData via mode_switch path already. + c.route(wsInbound{Type: "nope"}) + if ev := readFakeFrame(t, c); ev.Type != "error" { + t.Fatalf("unknown frame type = %q, want error", ev.Type) + } + // chat → publishes to bus (publishPublic non-nil path); no frame enqueued. + c.route(wsInbound{Type: "chat", Data: json.RawMessage(`{"text":"hi"}`)}) + if err := connHasNoFrame(c); err != nil { + t.Fatalf("chat should not enqueue a frame directly: %v", err) + } +} + +func TestPublishPublicNilBus(t *testing.T) { + t.Parallel() + hb := newTestHub(nil, nil) // nil bus + c := newFakeConn(hb, 32) + // Must not panic and must not enqueue anything. + c.publishPublic(wsChatSubject, map[string]any{"text": "hi"}) +} + +func TestRawData(t *testing.T) { + t.Parallel() + if rawData(nil) != nil { + t.Fatal("rawData(nil) should be nil") + } + if rawData(json.RawMessage(`{"a":1}`)) == nil { + t.Fatal("rawData(non-empty) should be non-nil") + } +} + +func TestOriginChecker(t *testing.T) { + t.Parallel() + allowed := map[string]bool{"http://ok.example": true} + checker := originChecker(allowed, false) + if !checker(reqWithOrigin("")) { + t.Fatal("empty origin should be allowed") + } + if !checker(reqWithOrigin("http://ok.example")) { + t.Fatal("whitelisted origin should be allowed") + } + if checker(reqWithOrigin("http://bad.example")) { + t.Fatal("unknown origin should be rejected") + } + all := originChecker(allowed, true) + if !all(reqWithOrigin("http://bad.example")) { + t.Fatal("allow-all should permit any origin") + } +} + +func TestWSTokenExtraction(t *testing.T) { + t.Parallel() + if got := wsToken(reqWithQuery("token", "abc")); got != "abc" { + t.Fatalf("query token = %q, want abc", got) + } + r := httptest.NewRequest(http.MethodGet, "/ws/chat", nil) + r.Header.Set("Authorization", "Bearer xyz") + if got := wsToken(r); got != "xyz" { + t.Fatalf("header token = %q, want xyz", got) + } + if got := wsToken(httptest.NewRequest(http.MethodGet, "/ws/chat", nil)); got != "" { + t.Fatalf("empty token = %q, want empty", got) + } +} + +// --- Start / ListenAndServe lifecycle ------------------------------------- + +func TestStartAndShutdown(t *testing.T) { + t.Parallel() + cfg := cfgDefault() + cfg.Server = config.ServerConfig{Host: "127.0.0.1", Port: 0} + srv := buildServerOnly(t, cfg, Deps{Auth: mustJWT(t)}) + ctx, cancel := context.WithCancel(context.Background()) + errCh, err := srv.Start(ctx) + if err != nil { + t.Fatalf("start: %v", err) + } + time.Sleep(100 * time.Millisecond) + cancel() + select { + case err := <-errCh: + if err != nil { + t.Fatalf("listen/serve returned error: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("server did not shut down on cancel") + } +} + +// --- Helpers --------------------------------------------------------------- + +func doHandler(t *testing.T, req *http.Request) *http.Response { + t.Helper() + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("do request: %v", err) + } + defer resp.Body.Close() + _, _ = io.ReadAll(resp.Body) + return resp +} + +func dialWS(t *testing.T, baseURL, token string) *websocket.Conn { + t.Helper() + wsURL := "ws" + strings.TrimPrefix(baseURL, "http") + "/ws/chat?token=" + token + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial ws: %v", err) + } + return conn +} + +func waitConnCount(t *testing.T, srv *Server, want int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if srv.hub.connCount() == want { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("conn count = %d, want %d", srv.hub.connCount(), want) +} + +func newTestHub(runner Chatter, b bus.Bus) *hub { + return &hub{ + conns: make(map[*wsConn]struct{}), + bus: b, + runner: runner, + log: sourcedLog(nil), + } +} + +func newFakeConn(hb *hub, buf int) *wsConn { + return &wsConn{ + send: make(chan []byte, buf), + done: make(chan struct{}), + h: hb, + } +} + +func readFakeFrame(t *testing.T, c *wsConn) wsEvent { + t.Helper() + select { + case b := <-c.send: + var ev wsEvent + if err := json.Unmarshal(b, &ev); err != nil { + t.Fatalf("decode frame %s: %v", b, err) + } + return ev + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for a frame") + return wsEvent{} + } +} + +func connHasNoFrame(c *wsConn) error { + select { + case b := <-c.send: + return errors.New("unexpected frame: " + string(b)) + default: + return nil + } +} + +func reqWithOrigin(origin string) *http.Request { + r := httptest.NewRequest(http.MethodGet, "/", nil) + r.Header.Set("Origin", origin) + return r +} + +func reqWithQuery(k, v string) *http.Request { + return httptest.NewRequest(http.MethodGet, "/ws/chat?"+k+"="+v, nil) +}