localchat/internal/web/server.go

// Package web serves the chat UI: Templ pages, HTMX fragments and an SSE
// endpoint that streams model replies token by token.
package web

import (
	"bytes"
	"context"
	"embed"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"github.com/a-h/templ"

	"git.b0b.be/bdeb/localchat/internal/chat"
	"git.b0b.be/bdeb/localchat/internal/config"
	"git.b0b.be/bdeb/localchat/internal/llm"
	"git.b0b.be/bdeb/localchat/internal/web/views"
)

//go:embed static
var staticFS embed.FS

const (
	sessionCookie = "localchat_sid"
	maxMessageLen = 32 << 10
)

// Server is the chat web app.
type Server struct {
	cfg   config.Config
	llm   *llm.Client
	store *chat.Store
	log   *slog.Logger
}

// New wires a server from config.
func New(cfg config.Config, log *slog.Logger) *Server {
	return &Server{
		cfg: cfg,
		llm: &llm.Client{
			BaseURL: cfg.BaseURL,
			APIKey:  cfg.APIKey,
			Model:   cfg.Model,
			HTTP:    &http.Client{}, // per-request deadlines come from contexts
		},
		store: chat.NewStore(),
		log:   log,
	}
}

// Handler returns the app's routes.
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("GET /static/", http.FileServerFS(staticFS))
	mux.HandleFunc("GET /{$}", s.index)
	mux.HandleFunc("GET /health", s.health)
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { w.Write([]byte("ok")) })
	mux.HandleFunc("POST /chat", s.send)
	mux.HandleFunc("POST /chat/reset", s.reset)
	mux.HandleFunc("GET /chat/stream/{id}", s.stream)
	return s.withSession(mux)
}

// Run serves until ctx is cancelled, then shuts down gracefully.
func (s *Server) Run(ctx context.Context) error {
	srv := &http.Server{
		Addr:              s.cfg.Addr,
		Handler:           s.Handler(),
		ReadHeaderTimeout: 10 * time.Second,
		// No WriteTimeout: replies stream for as long as the model talks.
		// LLM_TIMEOUT bounds each reply instead.
	}

	go func() {
		t := time.NewTicker(10 * time.Minute)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				s.store.Prune(24 * time.Hour)
			}
		}
	}()

	errc := make(chan error, 1)
	go func() { errc <- srv.ListenAndServe() }()
	s.log.Info("localchat listening", "url", "http://"+s.cfg.Addr, "llm", s.cfg.BaseURL, "model", s.cfg.Model)

	select {
	case err := <-errc:
		return err
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return srv.Shutdown(shutdownCtx)
	}
}

type ctxKey struct{}

// withSession gives every browser a random session cookie.
func (s *Server) withSession(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sid := ""
		if c, err := r.Cookie(sessionCookie); err == nil && len(c.Value) == 32 {
			sid = c.Value
		} else {
			sid = chat.NewID()
			http.SetCookie(w, &http.Cookie{
				Name:     sessionCookie,
				Value:    sid,
				Path:     "/",
				HttpOnly: true,
				SameSite: http.SameSiteLaxMode,
			})
		}
		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), ctxKey{}, sid)))
	})
}

func session(r *http.Request) string {
	sid, _ := r.Context().Value(ctxKey{}).(string)
	return sid
}

func (s *Server) render(w http.ResponseWriter, r *http.Request, c templ.Component) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	if err := c.Render(r.Context(), w); err != nil {
		s.log.Error("render", "err", err)
	}
}

func (s *Server) index(w http.ResponseWriter, r *http.Request) {
	msgs := s.store.Get(session(r)).Messages()
	s.render(w, r, views.Page(s.cfg.Title, msgs))
}

func (s *Server) health(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
	defer cancel()
	model, err := s.llm.ResolveModel(ctx)
	if err == nil && s.cfg.Model != "" {
		_, err = s.llm.Models(ctx) // configured model: still check the server is up
	}
	h := views.HealthInfo{OK: err == nil, Model: model}
	if err != nil {
		h.Err = err.Error()
	}
	s.render(w, r, views.Health(h))
}

func (s *Server) send(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxMessageLen)
	text := strings.TrimSpace(r.FormValue("message"))
	if text == "" {
		http.Error(w, "message is empty or too long", http.StatusBadRequest)
		return
	}
	user, reply := s.store.Get(session(r)).Ask(text)
	s.render(w, r, views.Exchange(user, reply))
}

func (s *Server) reset(w http.ResponseWriter, r *http.Request) {
	s.store.Reset(session(r))
	s.render(w, r, views.Messages(nil))
}

// stream generates a reply and pushes it to the browser as server-sent
// events. EventSource reconnects automatically, so a reply that is already
// claimed or finished is answered with its current state instead of being
// generated twice.
func (s *Server) stream(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	conv := s.store.Get(session(r))

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("X-Accel-Buffering", "no") // don't let reverse proxies buffer
	sse := &sseWriter{w: w, rc: http.NewResponseController(w)}

	prompt, msg, ok := conv.Claim(id, s.cfg.SystemPrompt, s.cfg.MaxHistory)
	if !ok {
		if msg.ID == "" {
			msg = chat.Message{ID: id, Role: "assistant", State: chat.Done,
				Err: "This reply was lost, probably because the server restarted. Please ask again."}
		}
		sse.component(r.Context(), "done", views.MessageView(msg))
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), s.cfg.Timeout)
	defer cancel()
	start := time.Now()
	err := s.llm.Stream(ctx, prompt, func(delta string) error {
		conv.Append(id, delta)
		return sse.component(ctx, "token", views.Token(delta))
	})

	if r.Context().Err() != nil { // browser went away; keep what we have
		conv.Finish(id, errors.New("interrupted"))
		return
	}
	if errors.Is(err, context.DeadlineExceeded) {
		err = fmt.Errorf("reply took longer than %s and was cut off", s.cfg.Timeout)
	}
	if err != nil {
		s.log.Warn("generation failed", "err", err)
	} else {
		s.log.Debug("reply done", "id", id, "took", time.Since(start))
	}
	sse.component(r.Context(), "done", views.MessageView(conv.Finish(id, err)))
}

type sseWriter struct {
	w  http.ResponseWriter
	rc *http.ResponseController
}

// component sends one rendered Templ component as a named SSE event.
func (s *sseWriter) component(ctx context.Context, event string, c templ.Component) error {
	var buf bytes.Buffer
	if err := c.Render(ctx, &buf); err != nil {
		return err
	}
	var b strings.Builder
	b.WriteString("event: " + event + "\n")
	for _, line := range strings.Split(buf.String(), "\n") {
		b.WriteString("data: " + line + "\n")
	}
	b.WriteString("\n")
	if _, err := s.w.Write([]byte(b.String())); err != nil {
		return err
	}
	return s.rc.Flush()
}