byteclaw/internal/server/server.go

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

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

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

	"github.com/crypto/byteclaw/internal/agent"
	"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
}

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"}) })
	api := http.NewServeMux()
	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.originCheck(s.auth(api)))
	return securityHeaders(mux)
}

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)
	})
}

// auth requires "Authorization: Bearer <token>" (constant-time compare).
func (s *Server) auth(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
		if !ok || subtle.ConstantTimeCompare([]byte(got), []byte(s.Token)) != 1 {
			writeErr(w, http.StatusUnauthorized, "unauthorized")
			return
		}
		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, 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.ListSessions(r.Context(), 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.CreateSession(r.Context(), 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) {
	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 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 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 _, err := s.Store.GetSession(r.Context(), id); err != nil {
		s.storeErr(w, err)
		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"`
}

// 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 != "" {
		_, err := s.Store.GetSession(ctx, id)
		return id, err
	}
	title := strings.TrimSpace(strings.ReplaceAll(text, "\n", " "))
	if utf8.RuneCountInString(title) > 48 {
		title = string([]rune(title)[:48]) + "…"
	}
	ss, err := s.Store.CreateSession(ctx, 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.ensureSession(r.Context(), req.SessionID, req.Text)
	if err != nil {
		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 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":"..."}
// server streams agent.Event JSON objects ending with type "done" or "error".
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 := r.Context()
	for {
		var req chatReq
		if err := wsjson.Read(ctx, c, &req); err != nil {
			return
		}
		if req.Type == "ping" {
			wsjson.Write(ctx, c, agent.Event{Type: "pong"})
			continue
		}
		if strings.TrimSpace(req.Text) == "" {
			wsjson.Write(ctx, c, agent.Event{Type: "error", Text: "text required"})
			continue
		}
		id, err := s.ensureSession(ctx, req.SessionID, req.Text)
		if err != nil {
			wsjson.Write(ctx, c, agent.Event{Type: "error", Text: "unknown session"})
			continue
		}
		wsjson.Write(ctx, c, agent.Event{Type: "session", Session: id})
		runCtx, cancel := context.WithTimeout(ctx, s.RunTimeout)
		reply, err := s.Agent.Run(runCtx, id, req.Text, func(e agent.Event) {
			e.Session = id
			wsjson.Write(runCtx, c, e)
		})
		cancel()
		if err != nil {
			s.Log.Warn("agent run", "err", err)
			wsjson.Write(ctx, c, agent.Event{Type: "error", Text: err.Error(), Session: id})
			continue
		}
		wsjson.Write(ctx, c, agent.Event{Type: "done", Text: reply, Session: id})
	}
}