og/internal/server/batch_api.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/example/ollama-fair-gateway/internal/auth"
	"github.com/example/ollama-fair-gateway/internal/batch"
)

type batchCreateRequest struct {
	Path string          `json:"path"`
	Body json.RawMessage `json:"body"`
}

func (s *Server) batchAPI(w http.ResponseWriter, r *http.Request, id auth.Identity) {
	if s.batchJobs == nil || !s.batchJobs.Enabled() {
		writeProtocolError(w, r, http.StatusNotFound, "batch_disabled", "durable batch jobs are disabled")
		return
	}
	base := "/gateway/v1/batches"
	rest := strings.TrimPrefix(r.URL.Path, base)
	if rest == "" || rest == "/" {
		switch r.Method {
		case http.MethodGet:
			writeJSON(w, http.StatusOK, map[string]any{"jobs": s.batchJobs.List(id.Tenant, id.Actor(), false)})
		case http.MethodPost:
			s.batchCreate(w, r, id)
		default:
			writeProtocolError(w, r, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
		}
		return
	}
	parts := strings.Split(strings.Trim(rest, "/"), "/")
	if len(parts) == 0 || parts[0] == "" {
		writeProtocolError(w, r, http.StatusNotFound, "not_found", "batch job not found")
		return
	}
	jobID := parts[0]
	if len(parts) == 1 && r.Method == http.MethodGet {
		j, ok := s.batchJobs.Get(jobID, id.Tenant, id.Actor(), false)
		if !ok {
			writeProtocolError(w, r, http.StatusNotFound, "batch_not_found", "batch job not found")
			return
		}
		writeJSON(w, http.StatusOK, j)
		return
	}
	if len(parts) == 2 && parts[1] == "output" && r.Method == http.MethodGet {
		s.batchOutput(w, r, id, jobID, false)
		return
	}
	if len(parts) == 2 && r.Method == http.MethodPost {
		var (
			j   batch.Job
			err error
		)
		switch parts[1] {
		case "pause":
			j, err = s.batchJobs.Pause(jobID, id.Tenant, id.Actor(), false)
		case "resume":
			j, err = s.batchJobs.Resume(jobID, id.Tenant, id.Actor(), false)
		case "cancel":
			j, err = s.batchJobs.Cancel(jobID, id.Tenant, id.Actor(), false)
		default:
			writeProtocolError(w, r, http.StatusNotFound, "not_found", "unknown batch operation")
			return
		}
		if err != nil {
			s.writeBatchError(w, r, err)
			return
		}
		writeJSON(w, http.StatusOK, j)
		return
	}
	writeProtocolError(w, r, http.StatusNotFound, "not_found", "unknown batch endpoint")
}

func (s *Server) batchCreate(w http.ResponseWriter, r *http.Request, id auth.Identity) {
	limit := s.cfg.BatchJobs.MaxInputBytes + 64<<10
	if limit <= 0 {
		limit = 16 << 20
	}
	var req batchCreateRequest
	if err := decodeJSON(r, &req, limit); err != nil {
		writeProtocolError(w, r, http.StatusBadRequest, "bad_batch", err.Error())
		return
	}
	if !s.isCompute(http.MethodPost, req.Path) {
		writeProtocolError(w, r, http.StatusBadRequest, "bad_batch_path", "batch path must be a configured compute POST endpoint")
		return
	}
	if len(req.Body) == 0 || string(req.Body) == "null" || !json.Valid(req.Body) {
		writeProtocolError(w, r, http.StatusBadRequest, "bad_batch_body", "batch body must contain a valid JSON request body")
		return
	}
	model := modelFromBody(req.Body)
	if _, _, err := s.resolveModel(r.Context(), id, model); err != nil {
		if errors.Is(err, ErrModelAccessDenied) {
			writeProtocolError(w, r, http.StatusForbidden, "model_access_denied", err.Error())
		} else {
			writeProtocolError(w, r, http.StatusNotFound, "model_alias_unavailable", err.Error())
		}
		return
	}
	j, err := s.batchJobs.Create(batchIdentitySnapshot(id), req.Path, model, req.Body)
	if err != nil {
		s.writeBatchError(w, r, err)
		return
	}
	w.Header().Set("Location", "/gateway/v1/batches/"+j.ID)
	writeJSON(w, http.StatusAccepted, j)
}

func (s *Server) batchOutput(w http.ResponseWriter, r *http.Request, id auth.Identity, jobID string, all bool) {
	f, j, err := s.batchJobs.OpenOutput(jobID, id.Tenant, id.Actor(), all)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			writeProtocolError(w, r, http.StatusConflict, "batch_output_unavailable", "batch output is not available yet")
			return
		}
		s.writeBatchError(w, r, err)
		return
	}
	defer f.Close()
	ct := strings.TrimSpace(j.ResponseContentType)
	if ct == "" {
		ct = "application/octet-stream"
	}
	w.Header().Set("Content-Type", ct)
	w.Header().Set("Cache-Control", "no-store")
	if st, err := f.Stat(); err == nil {
		w.Header().Set("Content-Length", strconv.FormatInt(st.Size(), 10))
	}
	w.WriteHeader(http.StatusOK)
	_, _ = io.Copy(w, f)
}

func (s *Server) writeBatchError(w http.ResponseWriter, r *http.Request, err error) {
	switch {
	case errors.Is(err, batch.ErrDisabled):
		writeProtocolError(w, r, http.StatusNotFound, "batch_disabled", err.Error())
	case errors.Is(err, batch.ErrNotFound):
		writeProtocolError(w, r, http.StatusNotFound, "batch_not_found", "batch job not found")
	case errors.Is(err, batch.ErrInvalidState):
		writeProtocolError(w, r, http.StatusConflict, "batch_state", err.Error())
	case errors.Is(err, batch.ErrFull):
		writeProtocolError(w, r, http.StatusTooManyRequests, "batch_full", err.Error())
	case errors.Is(err, batch.ErrInputTooLarge):
		writeProtocolError(w, r, http.StatusRequestEntityTooLarge, "batch_input_too_large", err.Error())
	default:
		writeProtocolError(w, r, http.StatusInternalServerError, "batch_error", err.Error())
	}
}

// ExecuteBatch is the runner used by the durable batch manager. The request is
// replayed directly into the normal compute path with the original identity
// metadata but the batch service class, so quotas, ACLs, scheduling, routing,
// metering, alerts and OpenTelemetry stay consistent with interactive traffic.
func (s *Server) ExecuteBatch(ctx context.Context, j batch.Job, input io.Reader, output io.Writer) batch.RunResult {
	id := auth.Identity{
		Tenant:       j.Identity.Tenant,
		Subject:      j.Identity.Subject,
		Application:  j.Identity.Application,
		AuthType:     j.Identity.AuthType,
		ClientIP:     j.Identity.ClientIP,
		Scopes:       make(map[string]bool, len(j.Identity.Scopes)),
		ModelACLSet:  j.Identity.ModelACLSet,
		ModelAccess:  j.Identity.ModelAccess,
		ServiceClass: "batch",
	}
	for _, scope := range j.Identity.Scopes {
		id.Scopes[scope] = true
	}
	r, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://gateway.local"+j.Path, input)
	if err != nil {
		return batch.RunResult{Error: err.Error()}
	}
	r.Header.Set("Content-Type", "application/json")
	r.RemoteAddr = "127.0.0.1:0"
	rw := &batchResponseWriter{header: make(http.Header), out: output}
	s.forward(rw, r, id)
	status := rw.status
	if status == 0 {
		status = http.StatusOK
	}
	return batch.RunResult{HTTPStatus: status, ResponseContentType: rw.header.Get("Content-Type"), RequestID: rw.header.Get("X-Request-ID")}
}

type batchResponseWriter struct {
	header http.Header
	out    io.Writer
	status int
}

func (w *batchResponseWriter) Header() http.Header { return w.header }
func (w *batchResponseWriter) WriteHeader(status int) {
	if w.status == 0 {
		w.status = status
	}
}
func (w *batchResponseWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	return w.out.Write(p)
}
func (w *batchResponseWriter) Flush() {}

func batchIdentitySnapshot(id auth.Identity) batch.IdentitySnapshot {
	scopes := make([]string, 0, len(id.Scopes))
	for scope, ok := range id.Scopes {
		if ok {
			scopes = append(scopes, scope)
		}
	}
	sort.Strings(scopes)
	return batch.IdentitySnapshot{Tenant: id.Tenant, Subject: id.Subject, Actor: id.Actor(), Application: id.Application, AuthType: id.AuthType, ClientIP: id.ClientIP, Scopes: scopes, ModelACLSet: id.ModelACLSet, ModelAccess: id.ModelAccess}
}

func (s *Server) uiBatchJobs(w http.ResponseWriter, r *http.Request, id auth.Identity) {
	base := "/gateway/ui-api/batches"
	if s.batchJobs == nil || !s.batchJobs.Enabled() {
		if r.URL.Path == base && r.Method == http.MethodGet {
			writeJSON(w, http.StatusOK, map[string]any{"enabled": false, "jobs": []batch.Job{}})
			return
		}
		writeProtocolError(w, r, http.StatusNotFound, "batch_disabled", "durable batch jobs are disabled")
		return
	}
	rest := strings.TrimPrefix(r.URL.Path, base)
	if rest == "" || rest == "/" {
		if r.Method != http.MethodGet {
			writeProtocolError(w, r, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
			return
		}
		writeJSON(w, http.StatusOK, map[string]any{
			"enabled":         true,
			"jobs":            s.batchJobs.List("", "", true),
			"retention":       s.cfg.BatchJobs.Retention.Value().String(),
			"max_jobs":        s.cfg.BatchJobs.MaxJobs,
			"max_concurrent":  s.cfg.BatchJobs.MaxConcurrent,
			"max_input_bytes": s.cfg.BatchJobs.MaxInputBytes,
		})
		return
	}
	parts := strings.Split(strings.Trim(rest, "/"), "/")
	if len(parts) == 0 || parts[0] == "" {
		writeProtocolError(w, r, http.StatusNotFound, "batch_not_found", "batch job not found")
		return
	}
	jobID := parts[0]
	if len(parts) == 1 && r.Method == http.MethodGet {
		j, ok := s.batchJobs.Get(jobID, "", "", true)
		if !ok {
			writeProtocolError(w, r, http.StatusNotFound, "batch_not_found", "batch job not found")
			return
		}
		writeJSON(w, http.StatusOK, j)
		return
	}
	if len(parts) == 2 && parts[1] == "output" && r.Method == http.MethodGet {
		s.batchOutput(w, r, id, jobID, true)
		return
	}
	if len(parts) == 2 && r.Method == http.MethodPost {
		var (
			j   batch.Job
			err error
		)
		switch parts[1] {
		case "pause":
			j, err = s.batchJobs.Pause(jobID, "", "", true)
		case "resume":
			j, err = s.batchJobs.Resume(jobID, "", "", true)
		case "cancel":
			j, err = s.batchJobs.Cancel(jobID, "", "", true)
		default:
			writeProtocolError(w, r, http.StatusNotFound, "not_found", "unknown batch operation")
			return
		}
		if err != nil {
			s.writeBatchError(w, r, err)
			return
		}
		s.log.Info("durable batch control", "batch_id", jobID, "action", parts[1], "admin_subject", id.Subject, "admin_auth_type", id.AuthType)
		writeJSON(w, http.StatusOK, j)
		return
	}
	writeProtocolError(w, r, http.StatusNotFound, "not_found", "unknown batch UI endpoint")
}