Luminary/internal/handler/replicate/handler.go

package replicate

import (
	"bufio"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/prediction"
)

type Handler struct {
	svc     *prediction.Service
	baseURL func() string
}

func New(svc *prediction.Service, baseURL func() string) *Handler {
	return &Handler{svc: svc, baseURL: baseURL}
}

func (h *Handler) Register(r *gin.RouterGroup) {
	r.POST("/predictions", h.Create)
	r.GET("/predictions", h.List)
	r.GET("/predictions/:id", h.Get)
	r.POST("/predictions/:id/cancel", h.Cancel)
	r.GET("/predictions/:id/stream", h.Stream)
}

func (h *Handler) ingress(c *gin.Context) (*model.IngressKey, bool) {
	v, ok := c.Get("ingress_key")
	if !ok {
		return nil, false
	}
	k, ok := v.(*model.IngressKey)
	return k, ok
}

func (h *Handler) Create(c *gin.Context) {
	ingress, ok := h.ingress(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"detail": "unauthorized"})
		return
	}
	var req prediction.CreateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
		return
	}
	wait := strings.Contains(c.GetHeader("Prefer"), "wait")
	p, err := h.svc.Create(c.Request.Context(), ingress, req, c.ClientIP(), wait)
	if err != nil {
		code := http.StatusBadRequest
		if strings.Contains(err.Error(), "not found") {
			code = http.StatusNotFound
		}
		c.JSON(code, gin.H{"detail": err.Error()})
		return
	}
	c.JSON(http.StatusCreated, h.svc.APIView(p, h.baseURL()))
}

func (h *Handler) Get(c *gin.Context) {
	ingress, ok := h.ingress(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"detail": "unauthorized"})
		return
	}
	p, err := h.svc.Get(c.Request.Context(), c.Param("id"), ingress)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()})
		return
	}
	c.JSON(http.StatusOK, h.svc.APIView(p, h.baseURL()))
}

func (h *Handler) List(c *gin.Context) {
	ingress, ok := h.ingress(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"detail": "unauthorized"})
		return
	}
	items, err := h.svc.List(c.Request.Context(), ingress, c.Query("cursor"), 100)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"detail": err.Error()})
		return
	}
	results := make([]map[string]any, 0, len(items))
	for _, p := range items {
		results = append(results, h.svc.APIView(p, h.baseURL()))
	}
	c.JSON(http.StatusOK, gin.H{"results": results})
}

func (h *Handler) Cancel(c *gin.Context) {
	ingress, ok := h.ingress(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"detail": "unauthorized"})
		return
	}
	p, err := h.svc.Cancel(c.Request.Context(), c.Param("id"), ingress)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()})
		return
	}
	c.JSON(http.StatusOK, h.svc.APIView(p, h.baseURL()))
}

func (h *Handler) Stream(c *gin.Context) {
	ingress, ok := h.ingress(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"detail": "unauthorized"})
		return
	}
	id := c.Param("id")
	if _, err := h.svc.Get(c.Request.Context(), id, ingress); err != nil {
		c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()})
		return
	}

	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Status(http.StatusOK)

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		return
	}
	w := bufio.NewWriter(c.Writer)
	ch := h.svc.Hub().Subscribe(id)
	defer h.svc.Hub().Unsubscribe(id, ch)

	for {
		select {
		case <-c.Request.Context().Done():
			return
		case ev, open := <-ch:
			if !open {
				return
			}
			payload, _ := json.Marshal(map[string]any{
				"id": id, "status": ev.Status, "logs": ev.Logs,
				"metrics": ev.Metrics, "output": ev.Output, "error": ev.Error,
			})
			_, _ = fmt.Fprintf(w, "data: %s\n\n", payload)
			w.Flush()
			flusher.Flush()
			if ev.Done {
				_, _ = fmt.Fprintf(w, "event: done\ndata: {}\n\n")
				w.Flush()
				flusher.Flush()
				return
			}
		}
	}
}