feat(gateway): 网关(REST /api/v1 + WS Hub /ws/chat + admin /api/v1/admin/* + JWT/RBAC + 审计,单端口 :8080)

This commit is contained in:
Aska committed 2026-09-06 07:32:40 +08:00
1 parent ff34b62079
commit 05e23ea590
7 files changed
+1535

No files matched your search

+170
View File
@@ -0,0 +1,170 @@
package gateway
import (
"context"
"encoding/json"
"errors"
"net/http"
capreg "git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors"
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)
// PluginSummary is the wire shape of a capability/plugin in the admin list.
type PluginSummary struct {
Name string `json:"name"`
Kind string `json:"kind"`
Description string `json:"description"`
AllowInGroup bool `json:"allowInGroup"`
AllowInStranger bool `json:"allowInStranger"`
}
// ChannelSummary is the wire shape of a registered channel adapter.
type ChannelSummary struct {
Name string `json:"name"`
Connected bool `json:"connected"`
}
// policyRequest is the POST /api/v1/admin/plugins/policy body.
type policyRequest struct {
Name string `json:"name"`
AllowInGroup bool `json:"allowInGroup"`
AllowInStranger bool `json:"allowInStranger"`
}
// adminListPersonas lists persona summaries (admin read).
func (s *Server) adminListPersonas(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, s.collectPersonas())
}
// adminListSessions lists channel sessions (admin read).
func (s *Server) adminListSessions(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, s.collectSessions(r.Context()))
}
// adminListMemories lists memory entries (admin read).
func (s *Server) adminListMemories(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, s.collectMemories(r.Context()))
}
// adminListPlugins lists every registered capability (admin read).
func (s *Server) adminListPlugins(w http.ResponseWriter, r *http.Request) {
out := make([]PluginSummary, 0)
if s.caps != nil {
for _, cap := range s.caps.List("") {
out = append(out, PluginSummary{
Name: cap.Name,
Kind: string(cap.Kind),
Description: cap.Description,
AllowInGroup: cap.Policy.AllowInGroup,
AllowInStranger: cap.Policy.AllowInStranger,
})
}
}
writeJSON(w, http.StatusOK, out)
}
// adminListChannels lists the registered channel adapters (admin read).
func (s *Server) adminListChannels(w http.ResponseWriter, r *http.Request) {
out := make([]ChannelSummary, 0)
if s.caps != nil {
for _, cap := range s.caps.List(capreg.KindChannel) {
ch, ok := cap.Value.(sdk.Channel)
if !ok {
continue
}
out = append(out, ChannelSummary{Name: ch.Name(), Connected: ch.IsConnected()})
}
}
writeJSON(w, http.StatusOK, out)
}
// adminListIoT lists IoT devices (admin read).
func (s *Server) adminListIoT(w http.ResponseWriter, r *http.Request) {
out := make([]iotDeviceSummary, 0)
if s.iot != nil {
devices, err := s.iot.List(r.Context())
if err != nil {
s.log.Error("list iot", "err", err)
writeError(w, apperr.Wrap(apperr.ErrInternal, "iot_failed", "failed to list devices"))
return
}
for _, d := range devices {
out = append(out, iotDeviceSummary{ID: d.ID, Name: d.Name, Type: d.Type, State: d.State, Online: d.Online})
}
}
writeJSON(w, http.StatusOK, out)
}
// adminSetPolicy upserts the reachability markers for a capability. It is the
// audited admin write example (owner-only, enforced by the route registration).
func (s *Server) adminSetPolicy(w http.ResponseWriter, r *http.Request) {
var req policyRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "invalid request body"))
return
}
if req.Name == "" {
writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "name is required"))
return
}
if s.caps == nil {
s.auditClaims(r, "policy.set", "error")
writeError(w, apperr.Wrap(apperr.ErrInternal, "policy_unavailable", "capability service is not configured"))
return
}
err := s.caps.SetPolicy(req.Name, capreg.ToolPolicy{AllowInGroup: req.AllowInGroup, AllowInStranger: req.AllowInStranger})
if err != nil {
s.auditClaims(r, "policy.set", "error")
if errors.Is(err, capreg.ErrNotFound) {
writeError(w, apperr.Wrap(apperr.ErrNotFound, "not_found", "capability %q not found", req.Name))
return
}
writeError(w, apperr.Wrap(apperr.ErrInternal, "policy_failed", "failed to set policy"))
return
}
s.auditClaims(r, "policy.set", "ok")
writeJSON(w, http.StatusOK, map[string]any{"name": req.Name, "allowInGroup": req.AllowInGroup, "allowInStranger": req.AllowInStranger})
}
// collectSessions returns a non-nil session slice from the configured store.
func (s *Server) collectSessions(ctx context.Context) []SessionSummary {
out := make([]SessionSummary, 0)
if s.sessions == nil {
return out
}
list, err := s.sessions.List(ctx)
if err != nil {
s.log.Error("list sessions", "err", err)
return out
}
for _, sess := range list {
out = append(out, summaryOfSession(sess))
}
return out
}
// collectMemories returns a non-nil memory slice from the configured store.
func (s *Server) collectMemories(ctx context.Context) []MemorySummary {
out := make([]MemorySummary, 0)
if s.memories == nil {
return out
}
list, err := s.memories.List(ctx)
if err != nil {
s.log.Error("list memories", "err", err)
return out
}
out = append(out, list...)
return out
}
// iotDeviceSummary is the wire shape of an IoT device.
type iotDeviceSummary struct {
ID string `json:"id"`
Name string `json:"name"`
Type string `json:"type"`
State map[string]any `json:"state"`
Online bool `json:"online"`
}
+30
View File
@@ -0,0 +1,30 @@
package gateway
import (
"net/http"
"time"
"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
)
// auditClaims records an admin write with the caller identity read from the
// request context (populated by auth.Authorize). It is the who/what/when/result
// audit trail required for admin operations; the "主代理代审" user-facing note is
// the frontend's job, not the gateway's.
func (s *Server) auditClaims(r *http.Request, what, result string) {
who := "anonymous"
if claims, ok := auth.ClaimsFromContext(r.Context()); ok {
who = claims.Sub + ":" + string(claims.Role)
}
s.audit(who, what, result)
}
// audit emits one structured audit record on the gateway logger.
func (s *Server) audit(who, what, result string) {
s.log.Info("admin audit",
"who", who,
"what", what,
"when", time.Now().UTC().Format(time.RFC3339),
"result", result,
)
}
+399
View File
@@ -0,0 +1,399 @@
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},
{name: "admin viewer", url: "/api/v1/admin/plugins", 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)
}
}
// 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
}
+113
View File
@@ -0,0 +1,113 @@
package gateway
import (
"bufio"
"fmt"
"net"
"net/http"
"time"
apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors"
)
// cors sets the Access-Control-* headers for origins in the configured allow
// list and short-circuits OPTIONS preflight requests. An empty origin (same-origin
// or non-browser client) is passed through untouched.
func (s *Server) cors(next http.Handler) http.Handler {
allowed, allowAll := s.allowedOrigins()
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if origin := r.Header.Get("Origin"); origin != "" && (allowAll || allowed[origin]) {
if allowAll {
w.Header().Set("Access-Control-Allow-Origin", "*")
} else {
w.Header().Set("Access-Control-Allow-Origin", origin)
}
w.Header().Set("Vary", "Origin")
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
w.Header().Set("Access-Control-Max-Age", "600")
}
if r.Method == http.MethodOptions {
w.WriteHeader(http.StatusNoContent)
return
}
next.ServeHTTP(w, r)
})
}
// allowedOrigins resolves the configured origin list into a lookup set plus an
// allow-all flag for the "*" wildcard.
func (s *Server) allowedOrigins() (map[string]bool, bool) {
list := s.cfg.AllowedOrigins
if len(list) == 0 {
return map[string]bool{}, false
}
allowAll := false
out := make(map[string]bool, len(list))
for _, o := range list {
if o == "*" {
allowAll = true
continue
}
out[o] = true
}
return out, allowAll
}
// recover converts a panic in a downstream handler into a 500 {error:...}
// response instead of crashing the process, and logs the panic.
func (s *Server) recover(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer func() {
if rec := recover(); rec != nil {
s.log.Error("panic recovered", "panic", rec, "method", r.Method, "path", r.URL.Path)
writeError(w, apperr.ErrInternal)
}
}()
next.ServeHTTP(w, r)
})
}
// logging records one structured line per request (method, path, status, dur).
func (s *Server) logging(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
next.ServeHTTP(rec, r)
s.log.Info("request",
"method", r.Method,
"path", r.URL.Path,
"status", rec.status,
"duration_ms", time.Since(start).Milliseconds(),
)
})
}
// statusRecorder captures the written status so the logging middleware can report
// it. It forwards the optional Hijacker/Flusher interfaces so the WebSocket
// upgrader (which requires http.Hijacker) and streaming handlers still work.
type statusRecorder struct {
http.ResponseWriter
status int
}
// WriteHeader records the status before delegating.
func (r *statusRecorder) WriteHeader(code int) {
r.status = code
r.ResponseWriter.WriteHeader(code)
}
// Hijack forwards to the underlying writer when it supports connection hijacking.
func (r *statusRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if hj, ok := r.ResponseWriter.(http.Hijacker); ok {
return hj.Hijack()
}
return nil, nil, fmt.Errorf("gateway: hijack unsupported by %T", r.ResponseWriter)
}
// Flush forwards to the underlying writer when it supports flushing.
func (r *statusRecorder) Flush() {
if f, ok := r.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
+193
View File
@@ -0,0 +1,193 @@
package gateway
import (
"encoding/json"
"net/http"
"strings"
"time"
"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
"git.yeij.top/AskaEth/Cyrene_v2/internal/channel/session"
apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors"
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)
// tokenTTL is how long an issued JWT remains valid.
const tokenTTL = 24 * time.Hour
// PersonaSummary is the wire shape of a persona in REST/admin list responses.
type PersonaSummary struct {
ID string `json:"id"`
Name string `json:"name"`
Character string `json:"character"`
IsDefault bool `json:"isDefault"`
Mode string `json:"mode"`
MemoryNS string `json:"memoryNamespace"`
}
// MemorySummary is the wire shape of a memory entry in the admin list response.
type MemorySummary struct {
ID string `json:"id"`
Scope string `json:"scope"`
Text string `json:"text"`
}
// SessionSummary is the wire shape of a channel session in list responses.
type SessionSummary struct {
Platform string `json:"platform"`
ChannelID string `json:"channelId"`
SessionID string `json:"sessionId"`
ActorRole string `json:"actorRole"`
Owner bool `json:"owner"`
IsGroup bool `json:"isGroup"`
MemWrite bool `json:"memWritable"`
}
// loginRequest is the POST /api/v1/auth/login body.
type loginRequest struct {
Username string `json:"username"`
Password string `json:"password"`
}
// loginResponse is the login success payload.
type loginResponse struct {
Token string `json:"token"`
ExpiresAt int64 `json:"expiresAt"`
Role string `json:"role"`
}
// chatRequest is the POST /api/v1/chat body. When Stream is true the gateway
// triggers streaming over the WebSocket channel and returns immediately.
type chatRequest struct {
SessionID string `json:"sessionId"`
Text string `json:"text"`
Stream bool `json:"stream"`
}
// chatResponse is the synchronous chat reply payload.
type chatResponse struct {
Reply string `json:"reply"`
Steps int `json:"steps"`
StopReason string `json:"stopReason"`
}
// handleHealth reports liveness. It is unauthenticated because it is a probe.
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"status": "ok"})
}
// handleLogin exchanges the admin username/password for a signed JWT.
func (s *Server) handleLogin(w http.ResponseWriter, r *http.Request) {
var req loginRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "invalid request body"))
return
}
cfg := s.cfg.Auth
if req.Username != cfg.AdminUser || req.Password != cfg.AdminPass {
writeError(w, apperr.Wrap(apperr.ErrUnauthorized, "unauthorized", "invalid username or password"))
return
}
token, err := s.auth.Sign(req.Username, auth.RoleOwner, tokenTTL)
if err != nil {
s.log.Error("login sign", "err", err)
writeError(w, apperr.Wrap(apperr.ErrInternal, "login_failed", "could not sign token"))
return
}
writeJSON(w, http.StatusOK, loginResponse{
Token: token,
ExpiresAt: time.Now().Add(tokenTTL).Unix(),
Role: string(auth.RoleOwner),
})
}
// handleChat runs one synchronous agent turn (or triggers streaming).
func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
var req chatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "invalid request body"))
return
}
if strings.TrimSpace(req.Text) == "" {
writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "text is required"))
return
}
if req.Stream {
s.publishChatStream(r.Context(), req.SessionID, req.Text)
writeJSON(w, http.StatusOK, map[string]any{"stream": true})
return
}
if s.runner == nil {
writeError(w, apperr.Wrap(apperr.ErrInternal, "chat_unavailable", "chat service is not configured"))
return
}
if req.SessionID == "" {
req.SessionID = "web"
}
res, err := s.runner.Run(r.Context(), req.SessionID, []model.Message{model.TextMessage(model.RoleUser, req.Text)})
if err != nil {
s.log.Error("chat run", "err", err, "session", req.SessionID)
writeError(w, apperr.Wrap(apperr.ErrInternal, "chat_failed", "chat failed"))
return
}
writeJSON(w, http.StatusOK, chatResponse{Reply: res.Text, Steps: res.Steps, StopReason: res.StopReason})
}
// handleListSessions returns the known sessions (or [] when no store is wired).
func (s *Server) handleListSessions(w http.ResponseWriter, r *http.Request) {
out := make([]SessionSummary, 0)
if s.sessions != nil {
list, err := s.sessions.List(r.Context())
if err != nil {
s.log.Error("list sessions", "err", err)
writeError(w, apperr.Wrap(apperr.ErrInternal, "sessions_failed", "failed to list sessions"))
return
}
for _, sess := range list {
out = append(out, summaryOfSession(sess))
}
}
writeJSON(w, http.StatusOK, out)
}
// handleListPersonas returns the personas (or [] when no store is wired).
func (s *Server) handleListPersonas(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, s.collectPersonas())
}
// collectPersonas returns a non-nil persona slice from the configured store.
func (s *Server) collectPersonas() []PersonaSummary {
out := make([]PersonaSummary, 0)
if s.personas != nil {
out = append(out, s.personas.List()...)
}
return out
}
// summaryOfSession maps a session.Session to its wire shape.
func summaryOfSession(sess session.Session) SessionSummary {
return SessionSummary{
Platform: sess.Platform,
ChannelID: sess.ChannelID,
SessionID: sess.SessionID,
ActorRole: string(sess.ActorRole),
Owner: sess.Owner,
IsGroup: sess.IsGroup,
MemWrite: sess.MemWritable,
}
}
// writeJSON encodes v as the {data: v} envelope with the given status.
func writeJSON(w http.ResponseWriter, status int, v any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(map[string]any{"data": v})
}
// writeError encodes err as the {error:{code,message}} envelope with the status
// resolved by apperr.HTTPStatus.
func writeError(w http.ResponseWriter, err error) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(apperr.HTTPStatus(err))
_ = json.NewEncoder(w).Encode(apperr.ToJSON(err))
}
+247
View File
@@ -0,0 +1,247 @@
// Package gateway implements the Cyrene v2 HTTP gateway: the only socket the
// single binary opens. It serves the versioned REST API (/api/v1), the realtime
// WebSocket hub (/ws/chat) and the RBAC-protected admin API (/api/v1/admin/*)
// all on one port (:8080).
//
// It is deliberately lean. All backend capability flows through injected
// dependencies (construction-time dependency injection) and the in-process bus,
// never through cross-process IPC, so the package never opens a listener of its
// own beyond the gateway port.
//
// Client contract (DevDocs/02-development-constraints.md §九, §十一):
// - REST responses are {data} on success and {error:{code,message}} on error.
// - Authentication is JWT bearer (internal/auth); authorization is the
// owner>admin>viewer RBAC hierarchy (internal/auth).
// - Admin writes are audited (who/what/when/result).
// - Realtime events follow DevDocs/design/07 §3 (WS protocol).
package gateway
import (
"context"
"fmt"
"io"
"log/slog"
"net"
"net/http"
"strconv"
"time"
"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/session"
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
"git.yeij.top/AskaEth/Cyrene_v2/internal/iot"
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)
// source is the log provenance tag for every record this package emits.
const source = "gateway"
// Authenticator abstracts the credential + RBAC service. *auth.JWTService is the
// production implementation; it is an interface so tests substitute a mock.
type Authenticator interface {
// Sign issues a token for sub with the given role and lifetime.
Sign(sub string, role auth.Role, ttl time.Duration) (string, error)
// Verify validates a token and returns its claims.
Verify(token string) (*auth.Claims, error)
// Authorize returns an HTTP middleware enforcing at least one required role.
Authorize(required ...auth.Role) func(http.Handler) http.Handler
}
// CapabilityStore abstracts the global capability registry used to enumerate
// plugins and channels in the admin API. *capability.Registry is the production
// implementation.
type CapabilityStore interface {
// List returns every capability whose Kind matches, or all when kind is
// empty.
List(kind capreg.Kind) []capreg.Capability
// SetPolicy replaces the reachability markers for the named capability.
SetPolicy(name string, p capreg.ToolPolicy) error
}
// Chatter runs a single agent turn. *loop.Loop is the production implementation.
type Chatter interface {
// Run drives one turn for sessionID and returns the final answer.
Run(ctx context.Context, sessionID string, messages []model.Message) (*loop.Result, error)
}
// SessionStore lists known channel sessions. The concrete session manager does
// not yet expose a list method, so production wiring may leave Sessions nil and
// the handlers return [].
type SessionStore interface {
List(ctx context.Context) ([]session.Session, error)
}
// PersonaStore lists persona summaries. The concrete persona registry does not
// yet expose a list method, so production wiring may leave Personas nil.
type PersonaStore interface {
List() []PersonaSummary
}
// MemoryStore lists memory entries. The concrete memory store does not yet
// expose a list method, so production wiring may leave Memories nil.
type MemoryStore interface {
List(ctx context.Context) ([]MemorySummary, error)
}
// IoTStore lists IoT devices. *iot.Manager is the production implementation.
type IoTStore interface {
List(ctx context.Context) ([]iot.Device, error)
}
// Deps carries the gateway's injected dependencies. Every field is nil-tolerant
// except Auth, which is required (authentication backs login and every RBAC
// decision). Registry/Loop are satisfied by the concrete *capability.Registry and
// *loop.Loop; Sessions/Personas/Memories remain nil until their managers expose a
// list method, in which case the handlers return [].
type Deps struct {
// Log is the base logger. Nil uses a discard logger.
Log *slog.Logger
// Auth is the JWT/RBAC service. Required.
Auth Authenticator
// Bus is the in-process event bus the WS hub subscribes to. Optional.
Bus bus.Bus
// Registry enumerates capabilities/plugins/channels. Optional.
Registry CapabilityStore
// Loop runs chat turns. Optional (chat returns an error when unset).
Loop Chatter
// Session lists channel sessions. Optional.
Session SessionStore
// Personas lists personas. Optional.
Personas PersonaStore
// Memories lists memory entries. Optional.
Memories MemoryStore
// IoT lists IoT devices. Optional.
IoT IoTStore
}
// Server is the gateway. It owns a ServeMux aggregating REST + WS + admin and a
// WebSocket hub. It is safe for concurrent use once constructed.
type Server struct {
cfg *config.Config
auth Authenticator
bus bus.Bus
caps CapabilityStore
runner Chatter
sessions SessionStore
personas PersonaStore
memories MemoryStore
iot IoTStore
log *slog.Logger
hub *hub
}
// New validates cfg and deps and returns a ready Server. Auth is required; the
// WebSocket hub subscribes to deps.Bus immediately so events published before
// ListenAndServe are still relayed to clients.
func New(cfg *config.Config, deps Deps) (*Server, error) {
if cfg == nil {
return nil, fmt.Errorf("gateway: nil config")
}
if deps.Auth == nil {
return nil, fmt.Errorf("gateway: auth dependency required")
}
s := &Server{
cfg: cfg,
auth: deps.Auth,
bus: deps.Bus,
caps: deps.Registry,
runner: deps.Loop,
sessions: deps.Session,
personas: deps.Personas,
memories: deps.Memories,
iot: deps.IoT,
log: sourcedLog(deps.Log),
}
s.hub = newHub(s)
s.hub.start()
return s, nil
}
// Handler returns the fully assembled http.Handler: the ServeMux routing
// REST + WS + admin, wrapped in recovery, CORS and request logging.
func (s *Server) Handler() http.Handler {
mux := http.NewServeMux()
// Public, no auth.
mux.HandleFunc("GET /api/v1/health", s.handleHealth)
mux.HandleFunc("POST /api/v1/auth/login", s.handleLogin)
// User side, authenticated. Chat is a write (owner/admin); reads accept any
// authenticated role so a viewer can still list personas/sessions.
mux.Handle("POST /api/v1/chat", s.auth.Authorize(auth.RoleOwner, auth.RoleAdmin)(http.HandlerFunc(s.handleChat)))
readAuth := s.auth.Authorize()
mux.Handle("GET /api/v1/sessions", readAuth(http.HandlerFunc(s.handleListSessions)))
mux.Handle("GET /api/v1/personas", readAuth(http.HandlerFunc(s.handleListPersonas)))
// Realtime.
mux.HandleFunc("/ws/chat", s.handleWS)
// Admin: reads require owner or admin; the audited policy write requires owner.
adminAuth := s.auth.Authorize(auth.RoleOwner, auth.RoleAdmin)
mux.Handle("GET /api/v1/admin/personas", adminAuth(http.HandlerFunc(s.adminListPersonas)))
mux.Handle("GET /api/v1/admin/sessions", adminAuth(http.HandlerFunc(s.adminListSessions)))
mux.Handle("GET /api/v1/admin/memories", adminAuth(http.HandlerFunc(s.adminListMemories)))
mux.Handle("GET /api/v1/admin/plugins", adminAuth(http.HandlerFunc(s.adminListPlugins)))
mux.Handle("GET /api/v1/admin/channels", adminAuth(http.HandlerFunc(s.adminListChannels)))
mux.Handle("GET /api/v1/admin/iot", adminAuth(http.HandlerFunc(s.adminListIoT)))
mux.Handle("POST /api/v1/admin/plugins/policy", s.auth.Authorize(auth.RoleOwner)(http.HandlerFunc(s.adminSetPolicy)))
var h http.Handler = mux
h = s.recover(h)
h = s.cors(h)
h = s.logging(h)
return h
}
// Start launches the gateway in a goroutine, returning a channel that carries the
// ListenAndServe result. The server runs until ctx is canceled.
func (s *Server) Start(ctx context.Context) (<-chan error, error) {
errCh := make(chan error, 1)
go func() { errCh <- s.ListenAndServe(ctx) }()
return errCh, nil
}
// ListenAndServe binds the configured host:port (default :8080) and blocks until
// ctx is canceled, then gracefully drains in-flight requests.
func (s *Server) ListenAndServe(ctx context.Context) error {
if ctx == nil {
ctx = context.Background()
}
addr := net.JoinHostPort(s.cfg.Server.Host, strconv.Itoa(s.cfg.Server.Port))
httpSrv := &http.Server{Addr: addr, Handler: s.Handler()}
done := make(chan struct{})
go func() {
<-ctx.Done()
s.log.Info("gateway shutting down", "addr", addr)
s.hub.close()
shCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
_ = httpSrv.Shutdown(shCtx)
close(done)
}()
s.log.Info("gateway listening", "addr", addr)
err := httpSrv.ListenAndServe()
if err != nil && err != http.ErrServerClosed {
return fmt.Errorf("gateway: listen: %w", err)
}
<-done
return nil
}
// Close frees the hub's bus subscription and closes live connections.
func (s *Server) Close() { s.hub.close() }
// sourcedLog returns a logger tagged with source "gateway", substituting a
// discard logger for nil so Server can always log safely.
func sourcedLog(l *slog.Logger) *slog.Logger {
if l == nil {
l = slog.New(slog.NewTextHandler(io.Discard, nil))
}
return logger.WithSource(l, source)
}
+383
View File
@@ -0,0 +1,383 @@
package gateway
import (
"context"
"encoding/json"
"log/slog"
"net/http"
"strings"
"sync"
"time"
"github.com/gorilla/websocket"
"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors"
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)
// WebSocket tuning constants. The ping/pong cadence keeps intermediary proxies
// from dropping idle connections.
const (
wsEventsPattern = "ws.>"
wsPrefix = "ws."
wsChatSubject = "ws.chat"
wsReadLimit = 4096
wsWriteWait = 10 * time.Second
wsPongWait = 60 * time.Second
wsPingPeriod = 54 * time.Second
wsSendBuffer = 32
)
// wsEvent is the outbound (server→client) frame.
type wsEvent struct {
Type string `json:"type"`
Data any `json:"data,omitempty"`
}
// wsInbound is the client→server frame. Data carries the per-type payload.
type wsInbound struct {
Type string `json:"type"`
ReqID string `json:"reqId,omitempty"`
Data json.RawMessage `json:"data,omitempty"`
}
// wsConn wraps one WebSocket connection. Writes are serialized through the
// buffered send channel drained by a single writePump, so the gorilla conn is
// only ever touched by one goroutine at a time.
type wsConn struct {
conn *websocket.Conn
send chan []byte
done chan struct{}
h *hub
claims *auth.Claims
once sync.Once
}
// hub owns the connection set and relays bus events to every client.
type hub struct {
mu sync.RWMutex
conns map[*wsConn]struct{}
bus bus.Bus
runner Chatter
log *slog.Logger
upgrader websocket.Upgrader
origins map[string]bool
allowAll bool
cancel context.CancelFunc
}
// newHub builds a hub bound to the server's dependencies.
func newHub(s *Server) *hub {
origins, allowAll := s.allowedOrigins()
return &hub{
conns: make(map[*wsConn]struct{}),
bus: s.bus,
runner: s.runner,
log: s.log,
origins: origins,
allowAll: allowAll,
upgrader: websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: originChecker(origins, allowAll),
},
}
}
// start subscribes the hub to the bus event stream. Publishing to ws.> arrives
// here and is fanned out to clients; nothing happens when no bus is wired.
func (h *hub) start() {
if h.bus == nil {
return
}
ctx, cancel := context.WithCancel(context.Background())
h.cancel = cancel
if _, err := h.bus.Subscribe(ctx, wsEventsPattern, h.relay); err != nil {
h.log.Warn("bus subscribe", "pattern", wsEventsPattern, "err", err)
}
}
// close cancels the bus subscription and shuts every live connection down.
func (h *hub) close() {
if h.cancel != nil {
h.cancel()
}
h.mu.Lock()
conns := make([]*wsConn, 0, len(h.conns))
for c := range h.conns {
conns = append(conns, c)
}
h.conns = make(map[*wsConn]struct{})
h.mu.Unlock()
for _, c := range conns {
c.shutdown()
}
}
// register adds a connection to the set.
func (h *hub) register(c *wsConn) {
h.mu.Lock()
h.conns[c] = struct{}{}
h.mu.Unlock()
}
// unregister removes a connection from the set.
func (h *hub) unregister(c *wsConn) {
h.mu.Lock()
if _, ok := h.conns[c]; ok {
delete(h.conns, c)
}
h.mu.Unlock()
}
// connCount reports the number of live connections (observability/tests).
func (h *hub) connCount() int {
h.mu.RLock()
defer h.mu.RUnlock()
return len(h.conns)
}
// broadcast fans an event out to every registered connection without blocking on
// any single (possibly slow) client: a full send buffer drops the frame.
func (h *hub) broadcast(ev wsEvent) {
b, err := json.Marshal(ev)
if err != nil {
h.log.Warn("marshal broadcast", "err", err)
return
}
h.mu.RLock()
defer h.mu.RUnlock()
for c := range h.conns {
c.enqueue(b)
}
}
// relay is the bus handler for the ws.> pattern. The ws.chat subject is a
// streaming trigger (run the loop and emit chunks); every other subject becomes
// an outbound event of its design type.
func (h *hub) relay(ctx context.Context, e bus.Event) error {
if e.Subject == wsChatSubject {
go h.streamChat(ctx, e.Payload)
return nil
}
h.broadcast(wsEvent{Type: eventType(e.Subject), Data: e.Payload})
return nil
}
// streamChat runs one agent turn for a chat trigger and broadcasts the result as
// stream_chunk/stream_end. It runs on its own goroutine so a slow model call does
// not stall the bus subscription worker.
func (h *hub) streamChat(ctx context.Context, payload any) {
var req chatStreamRequest
if b, err := json.Marshal(payload); err == nil {
_ = json.Unmarshal(b, &req)
}
if h.runner == nil || strings.TrimSpace(req.Text) == "" {
h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": "idle"}})
return
}
if req.SessionID == "" {
req.SessionID = "web"
}
res, err := h.runner.Run(ctx, req.SessionID, []model.Message{model.TextMessage(model.RoleUser, req.Text)})
if err != nil {
h.log.Error("ws chat run", "err", err, "session", req.SessionID)
h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": "error", "error": err.Error()}})
return
}
h.broadcast(wsEvent{Type: "stream_chunk", Data: map[string]any{"delta": res.Text}})
h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": res.StopReason, "steps": res.Steps}})
}
// chatStreamRequest is the decoded ws.chat payload.
type chatStreamRequest struct {
SessionID string `json:"sessionId"`
Text string `json:"text"`
}
// handleWS upgrades /ws/chat after authenticating the JWT from the query token or
// the Authorization header. A failed upgrade writes a 401 {error:...} before the
// handshake, so a browser/IM client sees a clean reject.
func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "websocket requires GET"))
return
}
claims, err := s.auth.Verify(wsToken(r))
if err != nil {
s.log.Debug("ws auth", "err", err)
writeError(w, apperr.ErrUnauthorized)
return
}
conn, err := s.hub.upgrader.Upgrade(w, r, nil)
if err != nil {
s.log.Debug("ws upgrade", "err", err)
return
}
c := &wsConn{
conn: conn,
send: make(chan []byte, wsSendBuffer),
done: make(chan struct{}),
h: s.hub,
claims: claims,
}
s.hub.register(c)
go c.writePump()
go c.readPump()
}
// readPump reads client frames until the connection closes, routing each by type.
func (c *wsConn) readPump() {
defer func() {
c.h.unregister(c)
c.shutdown()
}()
c.conn.SetReadLimit(wsReadLimit)
_ = c.conn.SetReadDeadline(time.Now().Add(wsPongWait))
c.conn.SetPongHandler(func(string) error {
return c.conn.SetReadDeadline(time.Now().Add(wsPongWait))
})
for {
_, msg, err := c.conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
c.h.log.Debug("ws read", "err", err)
}
return
}
var in wsInbound
if err := json.Unmarshal(msg, &in); err != nil {
c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "invalid json"}})
continue
}
c.route(in)
}
}
// writePump drains the send channel (and sends keepalive pings) until the
// connection closes. It is the only writer on the gorilla conn.
func (c *wsConn) writePump() {
ticker := time.NewTicker(wsPingPeriod)
defer func() {
ticker.Stop()
_ = c.conn.Close()
}()
for {
select {
case <-c.done:
return
case msg := <-c.send:
_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
return
}
case <-ticker.C:
_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
// route dispatches one inbound frame by its type.
func (c *wsConn) route(in wsInbound) {
switch in.Type {
case "mode_switch":
c.publishPublic("ws.mode_switch", rawData(in.Data))
case "chat":
c.publishPublic(wsChatSubject, rawData(in.Data))
case "subagent_ask":
c.enqueueJSON(wsEvent{Type: "subagent_event", Data: map[string]any{"state": "pending"}})
case "task_output_query":
c.enqueueJSON(wsEvent{Type: "task_output", Data: map[string]any{}})
default:
c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "unknown event type"}})
}
}
// publishPublic publishes a bus event without any per-connection reply.
func (c *wsConn) publishPublic(subject string, payload any) {
if c.h.bus == nil {
return
}
_ = c.h.bus.Publish(context.Background(), subject, payload)
}
// shutdown closes the connection exactly once.
func (c *wsConn) shutdown() {
c.once.Do(func() {
close(c.done)
_ = c.conn.Close()
})
}
// enqueue queues a serialized frame, dropping it when a slow client's buffer is
// full so the producer never blocks.
func (c *wsConn) enqueue(b []byte) {
select {
case c.send <- b:
default:
}
}
// enqueueJSON serializes ev and queues it.
func (c *wsConn) enqueueJSON(ev wsEvent) {
b, err := json.Marshal(ev)
if err != nil {
c.h.log.Warn("marshal ws event", "err", err)
return
}
c.enqueue(b)
}
// publishChatStream publishes a ws.chat streaming trigger (REST stream mode).
func (s *Server) publishChatStream(ctx context.Context, sessionID, text string) {
if s.bus == nil {
return
}
_ = s.bus.Publish(ctx, wsChatSubject, map[string]any{"sessionId": sessionID, "text": text})
}
// wsToken extracts the JWT from the ?token= query or the Authorization header.
func wsToken(r *http.Request) string {
if t := r.URL.Query().Get("token"); t != "" {
return t
}
const prefix = "Bearer "
h := r.Header.Get("Authorization")
if strings.HasPrefix(h, prefix) {
return strings.TrimSpace(h[len(prefix):])
}
return ""
}
// eventType maps a bus subject to its outbound event type, dropping the ws. prefix.
func eventType(subject string) string {
return strings.TrimPrefix(subject, wsPrefix)
}
// rawData converts a possibly-empty raw payload into a value safe to publish.
func rawData(d json.RawMessage) any {
if len(d) == 0 {
return nil
}
return d
}
// originChecker allows same-origin (empty) and whitelisted WebSocket origins,
// mirroring the CORS allow list. A "*" default allows all.
func originChecker(allowed map[string]bool, allowAll bool) func(*http.Request) bool {
return func(r *http.Request) bool {
origin := r.Header.Get("Origin")
if origin == "" {
return true
}
if allowAll {
return true
}
return allowed[origin]
}
}