byre/internal/server/server.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"log/slog"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"g.fc.run/misc/byre/internal/agents"
	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/git"
	"g.fc.run/misc/byre/internal/push"
	"g.fc.run/misc/byre/internal/store"
)

type Server struct {
	store     *store.Store
	agents    *agents.Manager
	notifier  *push.Notifier
	hub       *hub
	startedAt time.Time
	spa       fs.FS

	usersMu sync.Mutex
	users   map[string]string
}

type ctxKey struct{}

func New(dataDir string, spa fs.FS) (*Server, error) {
	err := os.MkdirAll(dataDir, 0o755)
	if err != nil {
		return nil, err
	}

	s := &Server{
		hub:       newHub(),
		startedAt: time.Now().UTC(),
		spa:       spa,
		users:     map[string]string{},
	}

	s.store, err = store.Open(filepath.Join(dataDir, "byre.db"), s.publish)
	if err != nil {
		return nil, err
	}

	s.notifier, err = push.New(s.store, pushWindow)
	if err != nil {
		s.store.Close()
		return nil, fmt.Errorf("start push notifier: %w", err)
	}

	s.agents = agents.New(s.store, agents.DefaultLimits, filepath.Join(dataDir, "worktrees"))
	err = s.agents.Recover(context.Background())
	if err != nil {
		s.notifier.Close()
		s.store.Close()
		return nil, fmt.Errorf("recover sessions: %w", err)
	}
	go s.agents.Watch()

	return s, nil
}

const pushWindow = 3 * time.Second

func (s *Server) publish(ev api.Event) {
	s.hub.publish(ev)
	s.notifier.Observe(ev)
}

func (s *Server) Close() error {
	s.agents.Close()
	s.notifier.Close()
	return s.store.Close()
}

func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /sw.js", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFileFS(w, r, s.spa, "dist/sw.js")
	})
	api.HandlerWithOptions(s, api.StdHTTPServerOptions{
		BaseRouter:  mux,
		Middlewares: []api.MiddlewareFunc{s.identify},
		ErrorHandlerFunc: func(w http.ResponseWriter, r *http.Request, err error) {
			writeError(w, http.StatusBadRequest, err)
		},
	})
	mux.Handle("/", http.FileServerFS(s.spa))
	return mux
}

func (s *Server) identify(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			writeError(w, http.StatusForbidden, err)
			return
		}
		ip := net.ParseIP(host)
		if ip == nil || !ip.IsLoopback() {
			writeError(w, http.StatusForbidden, fmt.Errorf("request from %s did not arrive via tailscale serve", host))
			return
		}

		login := r.Header.Get("Tailscale-User-Login")
		name := r.Header.Get("Tailscale-User-Name")
		if login == "" || name == "" {
			writeError(w, http.StatusUnauthorized, errors.New("missing Tailscale-User-Login or Tailscale-User-Name"))
			return
		}

		err = s.ensureUser(r.Context(), login, name)
		if err != nil {
			writeError(w, http.StatusInternalServerError, err)
			return
		}

		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), ctxKey{}, login)))
	})
}

func (s *Server) ensureUser(ctx context.Context, login, name string) error {
	s.usersMu.Lock()
	defer s.usersMu.Unlock()

	if s.users[login] == name {
		return nil
	}

	_, err := s.store.EnsureUser(ctx, login, name)
	if err != nil {
		return err
	}
	s.users[login] = name
	return nil
}

func actor(r *http.Request) string {
	return r.Context().Value(ctxKey{}).(string)
}

func (s *Server) GetMe(w http.ResponseWriter, r *http.Request) {
	u, err := s.store.GetUser(r.Context(), actor(r))
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, u)
}

func (s *Server) UpdateMe(w http.ResponseWriter, r *http.Request) {
	req := api.UserUpdate{}
	if !readJSON(w, r, &req) {
		return
	}

	if strings.TrimSpace(req.GitName) == "" || !strings.Contains(req.GitEmail, "@") {
		writeError(w, http.StatusBadRequest, errors.New("git_name and a valid git_email are required"))
		return
	}

	u, err := s.store.UpdateUserGit(r.Context(), actor(r), req.GitName, req.GitEmail)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, u)
}

func (s *Server) ListProjects(w http.ResponseWriter, r *http.Request) {
	projects, err := s.store.ListProjects(r.Context())
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, projects)
}

func (s *Server) CreateProject(w http.ResponseWriter, r *http.Request) {
	req := api.ProjectCreate{}
	if !readJSON(w, r, &req) {
		return
	}

	if strings.TrimSpace(req.Name) == "" {
		writeError(w, http.StatusBadRequest, errors.New("name is required"))
		return
	}

	info, err := git.Inspect(req.RepoPath)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}

	p := api.Project{
		Id:            store.NewID(),
		Name:          req.Name,
		RepoPath:      info.Root,
		DefaultBranch: info.DefaultBranch,
		Remote:        info.Remote,
		CreatedBy:     actor(r),
		CreatedAt:     time.Now().UTC(),
	}

	err = s.store.CreateProject(r.Context(), p)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	writeJSON(w, http.StatusCreated, p)
}

func (s *Server) GetProject(w http.ResponseWriter, r *http.Request, projectID string) {
	p, err := s.store.GetProject(r.Context(), projectID)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, p)
}

func (s *Server) ListTasks(w http.ResponseWriter, r *http.Request, projectID string) {
	_, err := s.store.GetProject(r.Context(), projectID)
	if err != nil {
		writeStoreError(w, err)
		return
	}

	tasks, err := s.store.ListTasks(r.Context(), projectID)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, tasks)
}

func (s *Server) CreateTask(w http.ResponseWriter, r *http.Request, projectID string) {
	req := api.TaskCreate{}
	if !readJSON(w, r, &req) {
		return
	}

	if strings.TrimSpace(req.Title) == "" {
		writeError(w, http.StatusBadRequest, errors.New("title is required"))
		return
	}

	p, err := s.store.GetProject(r.Context(), projectID)
	if err != nil {
		writeStoreError(w, err)
		return
	}

	t, err := s.agents.CreateTask(r.Context(), p, req.Title, actor(r), p.DefaultBranch)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err)
		return
	}
	writeJSON(w, http.StatusCreated, t)
}

func (s *Server) GetTask(w http.ResponseWriter, r *http.Request, taskID string) {
	t, err := s.store.GetTask(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, t)
}

func (s *Server) GetUnmerged(w http.ResponseWriter, r *http.Request, taskID string) {
	t, p, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}

	u, err := git.UnmergedChanges(t.WorktreePath, p.DefaultBranch)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err)
		return
	}
	writeJSON(w, http.StatusOK, api.Unmerged{
		Commits:     u.Commits,
		Files:       u.Files,
		Insertions:  u.Insertions,
		Deletions:   u.Deletions,
		Uncommitted: u.Uncommitted,
	})
}

func (s *Server) MergeTask(w http.ResponseWriter, r *http.Request, taskID string) {
	t, p, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}

	u, err := s.store.GetUser(r.Context(), actor(r))
	if err != nil {
		writeStoreError(w, err)
		return
	}

	sess, err := s.agents.Merge(r.Context(), t, p, u)
	if err != nil {
		writeAgentError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, sess)
}

func (s *Server) MergeIn(w http.ResponseWriter, r *http.Request, taskID string) {
	t, p, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}

	u, err := s.store.GetUser(r.Context(), actor(r))
	if err != nil {
		writeStoreError(w, err)
		return
	}

	sess, err := s.agents.MergeIn(r.Context(), t, u, p.DefaultBranch)
	if err != nil {
		writeAgentError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, sess)
}

func (s *Server) DiscardTask(w http.ResponseWriter, r *http.Request, taskID string) {
	t, p, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}

	err := s.agents.RemoveTask(r.Context(), t, p, actor(r))
	if err != nil {
		writeError(w, http.StatusInternalServerError, err)
		return
	}

	t, err = s.store.DiscardTask(r.Context(), taskID, actor(r))
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, t)
}

func (s *Server) activeTask(w http.ResponseWriter, r *http.Request, taskID string) (api.Task, api.Project, bool) {
	t, err := s.store.GetTask(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return t, api.Project{}, false
	}

	if t.Status != api.TaskStatusActive {
		writeError(w, http.StatusConflict, fmt.Errorf("task %s is %s", t.Id, t.Status))
		return t, api.Project{}, false
	}

	p, err := s.store.GetProject(r.Context(), t.ProjectId)
	if err != nil {
		writeStoreError(w, err)
		return t, p, false
	}
	return t, p, true
}

func (s *Server) GetStatus(w http.ResponseWriter, r *http.Request) {
	c, err := s.store.Counts(r.Context())
	if err != nil {
		writeStoreError(w, err)
		return
	}

	usage, err := s.agents.Usage(r.Context())
	if err != nil {
		writeStoreError(w, err)
		return
	}

	health, err := s.agents.Health(r.Context())
	if err != nil {
		writeError(w, http.StatusInternalServerError, err)
		return
	}

	dbBytes := int64(0)
	for _, suffix := range []string{"", "-wal", "-shm"} {
		fi, err := os.Stat(s.store.Path() + suffix)
		if err == nil {
			dbBytes += fi.Size()
		}
	}

	writeJSON(w, http.StatusOK, api.Status{
		StartedAt:    s.startedAt,
		DbBytes:      dbBytes,
		LastEventSeq: c.LastEventSeq,
		Projects:     c.Projects,
		ActiveTasks:  c.ActiveTasks,
		SseClients:   s.hub.count(),
		LiveAgents:   s.agents.LiveCount(),
		RunningTurns: usage.Running,
		QueuedTurns:  usage.Queued,
		Caps:         usage.Caps,
		RateLimits:   usage.RateLimits,

		AgentProcesses:    health.Processes,
		Zombies:           health.Zombies,
		StuckAfterSeconds: int(agents.StuckAfter / time.Second),
		ProtocolErrors:    s.agents.ProtocolErrors(),
	})
}

func (s *Server) ListTaskEvents(w http.ResponseWriter, r *http.Request, taskID string, params api.ListTaskEventsParams) {
	_, err := s.store.GetTask(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return
	}

	after := int64(0)
	if params.After != nil {
		after = *params.After
	}
	limit := 500
	if params.Limit != nil {
		limit = *params.Limit
	}

	events, err := s.store.TaskEvents(r.Context(), taskID, after, limit)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, events)
}

func (s *Server) ListSnapshots(w http.ResponseWriter, r *http.Request, taskID string) {
	_, err := s.store.GetTask(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return
	}

	snaps, err := s.store.ListSnapshots(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, snaps)
}

func (s *Server) ListSessions(w http.ResponseWriter, r *http.Request, taskID string) {
	_, err := s.store.GetTask(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return
	}

	sessions, err := s.store.ListSessions(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, sessions)
}

func (s *Server) CreateSession(w http.ResponseWriter, r *http.Request, taskID string) {
	req := api.SessionCreate{}
	if !readJSON(w, r, &req) {
		return
	}

	if strings.TrimSpace(req.Prompt) == "" {
		writeError(w, http.StatusBadRequest, errors.New("prompt is required"))
		return
	}

	t, err := s.store.GetTask(r.Context(), taskID)
	if err != nil {
		writeStoreError(w, err)
		return
	}

	u, err := s.store.GetUser(r.Context(), actor(r))
	if err != nil {
		writeStoreError(w, err)
		return
	}

	model := ""
	if req.Model != nil {
		model = *req.Model
	}

	images, ok := s.images(w, r, req.Images)
	if !ok {
		return
	}

	sess, err := s.agents.Create(r.Context(), t, u, req.Prompt, model, images...)
	if err != nil {
		writeAgentError(w, err)
		return
	}
	writeJSON(w, http.StatusCreated, sess)
}

func (s *Server) GetSession(w http.ResponseWriter, r *http.Request, sessionID string) {
	sess, err := s.store.GetSession(r.Context(), sessionID)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, sess)
}

func (s *Server) SendMessage(w http.ResponseWriter, r *http.Request, sessionID string) {
	req := api.MessageCreate{}
	if !readJSON(w, r, &req) {
		return
	}

	images, ok := s.images(w, r, req.Images)
	if !ok {
		return
	}

	if strings.TrimSpace(req.Text) == "" && len(images) == 0 {
		writeError(w, http.StatusBadRequest, errors.New("text or images are required"))
		return
	}

	model := ""
	if req.Model != nil {
		model = *req.Model
	}

	sess, err := s.agents.Send(r.Context(), sessionID, actor(r), req.Text, model, images...)
	if err != nil {
		writeAgentError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, sess)
}

func (s *Server) CompactSession(w http.ResponseWriter, r *http.Request, sessionID string) {
	sess, err := s.agents.Compact(r.Context(), sessionID)
	if err != nil {
		writeAgentError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, sess)
}

func (s *Server) InterruptSession(w http.ResponseWriter, r *http.Request, sessionID string) {
	sess, err := s.agents.Interrupt(r.Context(), sessionID)
	if err != nil {
		writeAgentError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, sess)
}

func writeAgentError(w http.ResponseWriter, err error) {
	switch {
	case errors.Is(err, store.ErrNotFound):
		writeError(w, http.StatusNotFound, err)
	case errors.Is(err, agents.ErrConflict):
		writeError(w, http.StatusConflict, err)
	case errors.Is(err, agents.ErrNoGitIdentity):
		writeError(w, http.StatusBadRequest, err)
	default:
		writeError(w, http.StatusInternalServerError, err)
	}
}

func readJSON(w http.ResponseWriter, r *http.Request, v any) bool {
	dec := json.NewDecoder(r.Body)
	dec.DisallowUnknownFields()
	err := dec.Decode(v)
	if err != nil {
		writeError(w, http.StatusBadRequest, fmt.Errorf("decode request: %w", err))
		return false
	}
	return true
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	err := json.NewEncoder(w).Encode(v)
	if err != nil {
		slog.Warn("write response", "err", err)
	}
}

func writeError(w http.ResponseWriter, status int, err error) {
	writeJSON(w, status, api.Error{Message: err.Error()})
}

func writeStoreError(w http.ResponseWriter, err error) {
	if errors.Is(err, store.ErrNotFound) {
		writeError(w, http.StatusNotFound, err)
		return
	}
	writeError(w, http.StatusInternalServerError, err)
}