byteclaw/internal/server/server.go

// Package server exposes the HTTP API and the WebSocket streaming endpoint.
package server

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"

	"github.com/crypto/byteclaw/internal/agent"
	"github.com/crypto/byteclaw/internal/providers"
	"github.com/crypto/byteclaw/internal/store"
)

const maxBody = 256 << 10

type Server struct {
	Agent          *agent.Agent
	Store          *store.Store
	Token          string
	AllowedOrigins []string
	Log            *slog.Logger
	RunTimeout     time.Duration
	// Info is returned by GET /api/info (never contains secrets).
	Info Info
	// Providers enables runtime provider switching (optional).
	Providers *providers.Manager
	// AllowRegistration enables POST /api/auth/register.
	AllowRegistration bool
	// TokenTTL is the lifetime of user tokens (default 30 days).
	TokenTTL time.Duration
	// TrustedProxies whose X-Forwarded-For is honoured (empty = loopback only).
	TrustedProxies []*net.IPNet

	authOnce sync.Once
	limiter  *limiter
}

type Info struct {
	Version  string `json:"version"`
	Provider string `json:"provider"`
	Model    string `json:"model"`
	BaseURL  string `json:"base_url"`
}

func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { writeJSON(w, 200, map[string]string{"status": "ok"}) })
	s.initAuth()
	mux.HandleFunc("POST /api/auth/register", s.register)
	mux.HandleFunc("POST /api/auth/login", s.login)
	api := http.NewServeMux()
	api.HandleFunc("POST /api/auth/logout", s.logout)
	api.HandleFunc("GET /api/auth/me", s.me)
	api.HandleFunc("GET /api/info", s.info)
	api.HandleFunc("GET /api/providers", s.listProviders)
	api.HandleFunc("PUT /api/providers/active", s.setActiveProvider)
	api.HandleFunc("GET /api/providers/{name}/models", s.listModels)
	api.HandleFunc("PUT /api/sessions/{id}/provider", s.setSessionProvider)
	api.HandleFunc("POST /api/sessions/{id}/cancel", s.cancel)
	api.HandleFunc("POST /api/sessions/{id}/summarize", s.summarize)
	api.HandleFunc("GET /api/sessions/{id}/summary", s.getSummary)
	api.HandleFunc("GET /api/sessions", s.listSessions)
	api.HandleFunc("POST /api/sessions", s.createSession)
	api.HandleFunc("GET /api/sessions/{id}", s.getSession)
	api.HandleFunc("PATCH /api/sessions/{id}", s.renameSession)
	api.HandleFunc("DELETE /api/sessions/{id}", s.deleteSession)
	api.HandleFunc("GET /api/sessions/{id}/messages", s.listMessages)
	api.HandleFunc("POST /api/chat", s.chat)
	api.HandleFunc("GET /api/ws", s.ws)
	mux.Handle("/api/", s.auth(api))
	h := http.Handler(mux)
	h = s.originCheck(h)
	return securityHeaders(h)
}

func securityHeaders(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("Cache-Control", "no-store")
		h.ServeHTTP(w, r)
	})
}

// OriginAllowed: requests without Origin (native apps, curl) pass; browser
// requests must come from the same host or an explicitly configured origin.
func (s *Server) OriginAllowed(r *http.Request) bool {
	o := r.Header.Get("Origin")
	if o == "" {
		return true
	}
	u, err := url.Parse(o)
	if err != nil || u.Host == "" {
		return false
	}
	if strings.EqualFold(u.Host, r.Host) {
		return true
	}
	for _, a := range s.AllowedOrigins {
		if a == "*" || strings.EqualFold(strings.TrimRight(a, "/"), o) {
			return true
		}
	}
	return false
}

func (s *Server) originCheck(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !s.OriginAllowed(r) {
			writeErr(w, http.StatusForbidden, "origin not allowed")
			return
		}
		if o := r.Header.Get("Origin"); o != "" {
			w.Header().Set("Access-Control-Allow-Origin", o)
			w.Header().Set("Vary", "Origin")
			if r.Method == http.MethodOptions {
				w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
				w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE")
				w.WriteHeader(http.StatusNoContent)
				return
			}
		}
		h.ServeHTTP(w, r)
	})
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(v)
}

func writeErr(w http.ResponseWriter, code int, msg string) {
	writeJSON(w, code, map[string]string{"error": msg})
}

func decode(w http.ResponseWriter, r *http.Request, v any) bool {
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxBody)).Decode(v); err != nil {
		writeErr(w, http.StatusBadRequest, "invalid JSON body")
		return false
	}
	return true
}

func (s *Server) storeErr(w http.ResponseWriter, err error) {
	if errors.Is(err, store.ErrNotFound) {
		writeErr(w, http.StatusNotFound, "not found")
		return
	}
	s.Log.Error("store", "err", err)
	writeErr(w, http.StatusInternalServerError, "internal error")
}

func (s *Server) listSessions(w http.ResponseWriter, r *http.Request) {
	list, err := s.Store.ListSessionsFor(r.Context(), principal(r.Context()).Scope(), 200)
	if err != nil {
		s.storeErr(w, err)
		return
	}
	writeJSON(w, 200, list)
}

func (s *Server) createSession(w http.ResponseWriter, r *http.Request) {
	var b struct {
		Title string `json:"title"`
	}
	if r.ContentLength != 0 && !decode(w, r, &b) {
		return
	}
	ss, err := s.Store.CreateSessionFor(r.Context(), principal(r.Context()).OwnerID(), b.Title)
	if err != nil {
		s.storeErr(w, err)
		return
	}
	writeJSON(w, http.StatusCreated, ss)
}

func (s *Server) getSession(w http.ResponseWriter, r *http.Request) {
	if !s.owns(w, r, r.PathValue("id")) {
		return
	}
	ss, err := s.Store.GetSession(r.Context(), r.PathValue("id"))
	if err != nil {
		s.storeErr(w, err)
		return
	}
	writeJSON(w, 200, ss)
}

func (s *Server) renameSession(w http.ResponseWriter, r *http.Request) {
	var b struct {
		Title string `json:"title"`
	}
	if !decode(w, r, &b) {
		return
	}
	if !s.owns(w, r, r.PathValue("id")) {
		return
	}
	if err := s.Store.RenameSession(r.Context(), r.PathValue("id"), b.Title); err != nil {
		s.storeErr(w, err)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) deleteSession(w http.ResponseWriter, r *http.Request) {
	if !s.owns(w, r, r.PathValue("id")) {
		return
	}
	if err := s.Store.DeleteSession(r.Context(), r.PathValue("id")); err != nil {
		s.storeErr(w, err)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) listMessages(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	if !s.owns(w, r, id) {
		return
	}
	msgs, err := s.Store.Messages(r.Context(), id, 0)
	if err != nil {
		s.storeErr(w, err)
		return
	}
	writeJSON(w, 200, msgs)
}

type chatReq struct {
	Type      string `json:"type,omitempty"`
	SessionID string `json:"session_id"`
	Text      string `json:"text"`
	// Optional provider/model for a NEW session, so the very first answer
	// already uses the per-session choice. Ignored for existing sessions.
	Provider string `json:"provider,omitempty"`
	Model    string `json:"model,omitempty"`
}

// startSession is ensureSession plus applying an optional per-session provider
// to a freshly created session before the first agent run.
func (s *Server) startSession(ctx context.Context, req chatReq) (string, error) {
	var a providers.Active
	if req.SessionID == "" && req.Provider != "" && s.Providers != nil {
		var err error
		if a, err = s.Providers.Validate(req.Provider, req.Model); err != nil {
			return "", err
		}
	}
	id, err := s.ensureSession(ctx, req.SessionID, req.Text)
	if err != nil || a.Provider == "" {
		return id, err
	}
	return id, s.Store.SetSessionProvider(ctx, id, a.Provider, a.Model)
}

// ensureSession returns an existing session id or creates one titled after the message.
func (s *Server) ensureSession(ctx context.Context, id, text string) (string, error) {
	if id != "" {
		return id, s.Store.SessionOwnedBy(ctx, principal(ctx).Scope(), id)
	}
	title := strings.TrimSpace(strings.ReplaceAll(text, "\n", " "))
	if utf8.RuneCountInString(title) > 48 {
		title = string([]rune(title)[:48]) + "…"
	}
	ss, err := s.Store.CreateSessionFor(ctx, principal(ctx).OwnerID(), title)
	return ss.ID, err
}

func (s *Server) chat(w http.ResponseWriter, r *http.Request) {
	var req chatReq
	if !decode(w, r, &req) {
		return
	}
	if strings.TrimSpace(req.Text) == "" {
		writeErr(w, http.StatusBadRequest, "text required")
		return
	}
	id, err := s.startSession(r.Context(), req)
	if err != nil {
		if req.SessionID == "" && req.Provider != "" {
			providerErr(w, err)
			return
		}
		s.storeErr(w, err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), s.RunTimeout)
	defer cancel()
	reply, err := s.Agent.Run(ctx, id, req.Text, nil)
	if errors.Is(err, agent.ErrCancelled) {
		writeJSON(w, 200, map[string]any{"session_id": id, "reply": reply, "cancelled": true})
		return
	}
	if err != nil {
		s.Log.Warn("agent run", "err", err)
		writeJSON(w, http.StatusBadGateway, map[string]string{"session_id": id, "error": err.Error()})
		return
	}
	writeJSON(w, 200, map[string]string{"session_id": id, "reply": reply})
}

// ws protocol: client sends {"type":"message","session_id":"...","text":"..."};
// the server streams agent.Event JSON objects ending with "done", "cancelled"
// or "error". While an answer is streaming the client may send
// {"type":"cancel","session_id":"..."} (or call POST /api/sessions/{id}/cancel).
func (s *Server) ws(w http.ResponseWriter, r *http.Request) {
	// Origin was already validated by originCheck; the library's own check is
	// therefore disabled here to avoid duplicating (and diverging from) policy.
	c, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true})
	if err != nil {
		return
	}
	defer c.CloseNow()
	c.SetReadLimit(maxBody)
	ctx, stop := context.WithCancel(r.Context())
	defer stop()
	var wmu sync.Mutex
	send := func(e agent.Event) {
		wmu.Lock()
		defer wmu.Unlock()
		wctx, cancel := context.WithTimeout(ctx, 10*time.Second)
		defer cancel()
		wsjson.Write(wctx, c, e)
	}
	var running sync.WaitGroup
	defer running.Wait()
	busy := make(chan struct{}, 1) // one answer at a time per connection
	for {
		var req chatReq
		if err := wsjson.Read(ctx, c, &req); err != nil {
			stop() // connection gone: abort any running answer
			return
		}
		switch req.Type {
		case "ping":
			send(agent.Event{Type: "pong"})
			continue
		case "cancel":
			if req.SessionID == "" || s.Store.SessionOwnedBy(ctx, principal(ctx).Scope(), req.SessionID) != nil || !s.Agent.Cancel(req.SessionID) {
				send(agent.Event{Type: "error", Text: "nothing to cancel", Session: req.SessionID})
			}
			continue
		}
		if strings.TrimSpace(req.Text) == "" {
			send(agent.Event{Type: "error", Text: "text required"})
			continue
		}
		select {
		case busy <- struct{}{}:
		default:
			send(agent.Event{Type: "error", Text: "an answer is still streaming; cancel it first"})
			continue
		}
		id, err := s.startSession(ctx, req)
		if err != nil {
			<-busy
			msg := "unknown session"
			if req.SessionID == "" && req.Provider != "" {
				msg = "provider: " + err.Error()
			}
			send(agent.Event{Type: "error", Text: msg})
			continue
		}
		send(agent.Event{Type: "session", Session: id})
		running.Add(1)
		go func(id, text string) {
			defer running.Done()
			defer func() { <-busy }()
			runCtx, cancel := context.WithTimeout(ctx, s.RunTimeout)
			defer cancel()
			reply, err := s.Agent.Run(runCtx, id, text, func(e agent.Event) {
				e.Session = id
				send(e)
			})
			switch {
			case errors.Is(err, agent.ErrCancelled):
				send(agent.Event{Type: "cancelled", Text: reply, Session: id})
			case err != nil:
				s.Log.Warn("agent run", "err", err)
				send(agent.Event{Type: "error", Text: err.Error(), Session: id})
			default:
				send(agent.Event{Type: "done", Text: reply, Session: id})
			}
		}(id, req.Text)
	}
}