llmgw/internal/router/workflows.go

package router

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"hash/crc32"
	"math"
	"net/http"
	"path"
	"slices"
	"sort"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
)

// Workflow runs: whatever a ComfyUI model makes, from the inputs its catalog
// entry declares. Pictures and meshes have fixed shapes of their own
// (images.go, meshes.go); this is the one shape for everything, so a model
// is added to the catalog as data — its workflow, its inputs, its output
// kind — rather than as code. A run is a job like a 3D model (jobs.go): it
// is queued, asked after, and its result is a file.

// workflowsPath is where workflow runs are asked for. The synchronous path
// behind it is served to the job worker only; clients get a job.
const workflowsPath = "/v1/workflows/runs"

// runRequest is a request to run a model's workflow.
type runRequest struct {
	Model  string                     `json:"model"`
	Inputs map[string]json.RawMessage `json:"inputs"`
	Seed   *int64                     `json:"seed"`

	// values are the inputs that are not files, as the workflow gets them.
	values map[string]any
	// files are the file inputs, by input name.
	files map[string]runFile
}

// runFile is a file input: its bytes, or the id of a file in /v1/files to be
// read when the run is served.
type runFile struct {
	data   []byte
	fileID string
	ext    string
}

// maxTextInput bounds a text input that names no max of its own.
const maxTextInput = 32 << 10

// parseRunRequest checks a request against what the model's workflow takes.
// It refuses an input the workflow does not take, a required one left out, a
// value of the wrong type or out of its bounds, and a file that is not what
// the input says.
func parseRunRequest(m catalog.Model, body []byte) (runRequest, error) {
	var in runRequest
	if err := json.Unmarshal(body, &in); err != nil {
		return in, err
	}
	in.values, in.files = map[string]any{}, map[string]runFile{}
	declared := m.ComfyInputs()
	for name := range in.Inputs {
		if !slices.ContainsFunc(declared, func(d catalog.Input) bool { return d.Name == name }) {
			names := make([]string, 0, len(declared))
			for _, d := range declared {
				names = append(names, d.Name)
			}
			if len(names) == 0 {
				return in, fmt.Errorf("%s takes no inputs, and was given %q", m.Name, name)
			}
			return in, fmt.Errorf("%s takes no input %q (it takes %s)", m.Name, name, strings.Join(names, ", "))
		}
	}
	for _, d := range declared {
		raw, given := in.Inputs[d.Name]
		if !given || string(raw) == "null" {
			if d.Required {
				return in, fmt.Errorf("input %q is required: %s", d.Name, describeInput(d))
			}
			if d.Default != nil {
				in.values[d.Name] = d.Default
			}
			continue
		}
		if d.IsFile() {
			f, err := parseFileInput(d, raw)
			if err != nil {
				return in, err
			}
			in.files[d.Name] = f
			continue
		}
		v, err := parseValueInput(d, raw)
		if err != nil {
			return in, err
		}
		in.values[d.Name] = v
	}
	return in, nil
}

func describeInput(d catalog.Input) string {
	s := d.Type
	if d.Type == catalog.InputChoice {
		s += " of " + strings.Join(d.Choices, ", ")
	}
	if d.Doc != "" {
		s += " — " + d.Doc
	}
	return s
}

func parseValueInput(d catalog.Input, raw json.RawMessage) (any, error) {
	bounds := func(f float64) error {
		if d.Min != nil && f < *d.Min || d.Max != nil && f > *d.Max {
			return fmt.Errorf("input %q is %v, outside %s", d.Name, f, boundsText(d))
		}
		return nil
	}
	switch d.Type {
	case catalog.InputText:
		var s string
		if json.Unmarshal(raw, &s) != nil {
			return nil, fmt.Errorf("input %q is text", d.Name)
		}
		limit := maxTextInput
		if d.Max != nil {
			limit = int(*d.Max)
		}
		if n := utf8.RuneCountInString(s); n > limit {
			return nil, fmt.Errorf("input %q is %d characters, past %d", d.Name, n, limit)
		}
		return s, nil
	case catalog.InputInt:
		var f float64
		if json.Unmarshal(raw, &f) != nil || f != math.Trunc(f) || math.Abs(f) > 1<<53 {
			return nil, fmt.Errorf("input %q is a whole number", d.Name)
		}
		return int64(f), bounds(f)
	case catalog.InputFloat:
		var f float64
		if json.Unmarshal(raw, &f) != nil || math.IsNaN(f) || math.IsInf(f, 0) {
			return nil, fmt.Errorf("input %q is a number", d.Name)
		}
		return f, bounds(f)
	case catalog.InputBool:
		var b bool
		if json.Unmarshal(raw, &b) != nil {
			return nil, fmt.Errorf("input %q is true or false", d.Name)
		}
		return b, nil
	case catalog.InputChoice:
		var s string
		if json.Unmarshal(raw, &s) != nil || !slices.Contains(d.Choices, s) {
			return nil, fmt.Errorf("input %q is one of %s", d.Name, strings.Join(d.Choices, ", "))
		}
		return s, nil
	}
	return nil, fmt.Errorf("input %q has type %q, which takes no value", d.Name, d.Type)
}

func boundsText(d catalog.Input) string {
	switch {
	case d.Min != nil && d.Max != nil:
		return fmt.Sprintf("%v to %v", *d.Min, *d.Max)
	case d.Min != nil:
		return fmt.Sprintf("at least %v", *d.Min)
	case d.Max != nil:
		return fmt.Sprintf("at most %v", *d.Max)
	}
	return "its bounds"
}

// parseFileInput reads a file input: base64, a data URL, or a file id.
func parseFileInput(d catalog.Input, raw json.RawMessage) (runFile, error) {
	var s string
	if json.Unmarshal(raw, &s) != nil || strings.TrimSpace(s) == "" {
		return runFile{}, fmt.Errorf("input %q is a file: base64, a data URL, or the id of a file uploaded to /v1/files", d.Name)
	}
	s = strings.TrimSpace(s)
	if strings.HasPrefix(s, "file-") && len(s) < 80 && !strings.ContainsAny(s, " /") {
		return runFile{fileID: s}, nil
	}
	if strings.HasPrefix(s, "data:") {
		_, after, ok := strings.Cut(s, ";base64,")
		if !ok {
			return runFile{}, fmt.Errorf("input %q is a data URL without base64 in it", d.Name)
		}
		s = after
	}
	if int64(base64.StdEncoding.DecodedLen(len(s))) > d.MaxBytes()+3 {
		return runFile{}, fmt.Errorf("input %q is past %d MB: upload it to /v1/files and give its id", d.Name, d.MaxBytes()>>20)
	}
	b, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return runFile{}, fmt.Errorf("input %q is not base64: %v", d.Name, err)
	}
	return checkFileInput(d, b)
}

// checkFileInput is a file input's bytes, if they are what the input says.
func checkFileInput(d catalog.Input, b []byte) (runFile, error) {
	if int64(len(b)) > d.MaxBytes() {
		return runFile{}, fmt.Errorf("input %q is past %d MB", d.Name, d.MaxBytes()>>20)
	}
	ext, kind := sniffFile(b)
	want := map[string]string{catalog.InputImage: "image", catalog.InputMask: "image",
		catalog.InputAudio: "audio", catalog.InputVideo: "video"}[d.Type]
	if want != "" && kind != want {
		return runFile{}, fmt.Errorf("input %q is %s, and this is not one it can read", d.Name, map[string]string{
			"image": "a PNG, JPEG, WebP or GIF picture", "audio": "WAV, MP3, FLAC or Ogg audio",
			"video": "an MP4, WebM or QuickTime video"}[want])
	}
	if ext == "" {
		ext = ".bin"
	}
	return runFile{data: b, ext: ext}, nil
}

// sniffFile is a file's extension and kind (image, audio, video) by its
// first bytes, or empty when it is none of the kinds a workflow loads.
func sniffFile(b []byte) (ext, kind string) {
	has := func(off int, sig string) bool { return len(b) >= off+len(sig) && string(b[off:off+len(sig)]) == sig }
	switch {
	case has(0, "\x89PNG\r\n\x1a\n"):
		return ".png", "image"
	case has(0, "\xff\xd8\xff"):
		return ".jpg", "image"
	case has(0, "RIFF") && has(8, "WEBP"):
		return ".webp", "image"
	case has(0, "GIF87a"), has(0, "GIF89a"):
		return ".gif", "image"
	case has(0, "RIFF") && has(8, "WAVE"):
		return ".wav", "audio"
	case has(0, "fLaC"):
		return ".flac", "audio"
	case has(0, "OggS"):
		return ".ogg", "audio"
	case has(0, "ID3"), len(b) > 2 && b[0] == 0xff && b[1]&0xe0 == 0xe0:
		return ".mp3", "audio"
	case has(4, "ftypqt"):
		return ".mov", "video"
	case has(4, "ftyp"):
		return ".mp4", "video"
	case has(0, "\x1a\x45\xdf\xa3"):
		return ".webm", "video"
	}
	return "", ""
}

// runFileIDs are the file ids a run's inputs name.
func (in runRequest) runFileIDs() []string {
	var ids []string
	for _, f := range in.files {
		if f.fileID != "" {
			ids = append(ids, f.fileID)
		}
	}
	sort.Strings(ids)
	return ids
}

// checkRunFiles refuses a run naming a file its owner does not have.
func (r *Router) checkRunFiles(ctx context.Context, owner string, in runRequest) error {
	for _, id := range in.runFileIDs() {
		f, err := r.store.FileByID(ctx, id)
		if err != nil || f.Owner != owner {
			return fmt.Errorf("no file %q", id)
		}
	}
	return nil
}

// generateRun answers a workflow run at a ComfyUI replica: the file inputs
// put in its input directory, the workflow filled in and run, and the file
// it made written back with its type — what the job keeps.
func (r *Router) generateRun(ctx context.Context, w http.ResponseWriter, req *parsedRequest, ep endpoint.Info, body []byte) (proxy.Outcome, error) {
	m, known := r.modelMap()[req.Model]
	if !known {
		return proxy.Outcome{}, fmt.Errorf("%s: no catalog entry for %q says what its workflow takes", ep.ID, req.Model)
	}
	in, err := parseRunRequest(m, body)
	if err != nil {
		return proxy.Outcome{}, err
	}
	tpl, err := comfy.Load(ep.Workflow)
	if err != nil {
		return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err)
	}
	if err := checkPlaceholders(m, tpl); err != nil {
		return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err)
	}
	if t := r.cfg.Timeouts.Total; t > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, t)
		defer cancel()
	}
	start := time.Now()
	fail := func(err error) (proxy.Outcome, error) {
		var ce *comfy.Error
		switch {
		case errors.As(err, &ce):
			return proxy.Outcome{Retryable: ce.Retryable}, err
		case errors.Is(err, context.Canceled) && ctxCauseIsClient(ctx):
			return proxy.Outcome{}, proxy.ErrClientGone
		case errors.Is(err, context.DeadlineExceeded):
			return proxy.Outcome{}, fmt.Errorf("%s finished nothing within %s", ep.ID, r.cfg.Timeouts.Total)
		}
		return proxy.Outcome{}, err
	}
	var upstream struct {
		Model string `json:"model"`
	}
	_ = json.Unmarshal(body, &upstream)
	p := comfy.Params{Checkpoint: upstream.Model, Seed: comfy.RandomSeed(), Extra: in.values}
	if in.Seed != nil {
		p.Seed = *in.Seed
	}
	if s, ok := in.values["seed"].(int64); ok {
		p.Seed = s
	}
	// Each file under a name of its own: two runs' files must not overwrite
	// each other in the replica's input directory.
	names := make([]string, 0, len(in.files))
	for name := range in.files {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		f := in.files[name]
		if f.fileID != "" {
			pr, _ := auth.FromContext(ctx)
			d := m.ComfyInputs()[slices.IndexFunc(m.ComfyInputs(), func(x catalog.Input) bool { return x.Name == name })]
			meta, err := r.store.FileByID(ctx, f.fileID)
			if err != nil || meta.Owner != ownerOf(pr) {
				return proxy.Outcome{}, fmt.Errorf("input %q: no file %q", name, f.fileID)
			}
			b, err := r.store.FileContent(ctx, f.fileID)
			if err != nil {
				return proxy.Outcome{}, fmt.Errorf("input %q: reading file %q: %w", name, f.fileID, err)
			}
			if f, err = checkFileInput(d, b); err != nil {
				return proxy.Outcome{}, err
			}
			in.files[name] = f
		}
		rnd := make([]byte, 8)
		_, _ = rand.Read(rnd)
		f = in.files[name]
		up, err := r.comfy.Upload(ctx, ep.URL, ep.APIKey, "llmgw-"+hex.EncodeToString(rnd)+f.ext, f.data)
		if err != nil {
			return fail(err)
		}
		p.Extra[name] = up
	}
	workflow, err := tpl.Fill(p)
	if err != nil {
		return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err)
	}
	files, err := r.comfy.GenerateWatch(ctx, ep.URL, ep.APIKey, workflow, r.jobProgress(ctx, req))
	if err != nil {
		return fail(err)
	}
	out, err := pickOutput(m.ComfyOutput(), files)
	if err != nil {
		return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err)
	}
	data, ct := out.Data, out.ContentType
	if ct == "model/obj" {
		if data, err = objToGLB(data); err != nil {
			return proxy.Outcome{}, fmt.Errorf("%s: the mesh it saved: %w", ep.ID, err)
		}
		ct = "model/gltf-binary"
	}
	if m.License != nil {
		if data, err = withLicense(data, ct, *m.License); err != nil {
			// Refused rather than sent without it: the file travels without
			// the API, and the notice is what the terms ask to travel too.
			return proxy.Outcome{}, fmt.Errorf("%s: the licence could not be written into the file: %w", ep.ID, err)
		}
		w.Header().Set("x-llmgw-license", m.License.Name)
	}
	if _, ok := tpl.Placeholders()["seed"]; ok {
		w.Header().Set("x-llmgw-seed", strconv.FormatInt(p.Seed, 10))
	}
	w.Header().Set("Content-Type", ct)
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(data)
	took := time.Since(start)
	o := proxy.Outcome{Wrote: true, Status: http.StatusOK, Bytes: int64(len(data)), TTFT: took, Decode: took}
	switch m.ComfyOutput() {
	case catalog.OutputPictures:
		o.Images = 1
	case catalog.OutputMeshes:
		o.Meshes = 1
	}
	return o, nil
}

// checkPlaceholders refuses a workflow that wants what no request can give:
// a placeholder that is neither an input, nor one llmgw fills itself, nor
// has a default of its own. The catalog's to fix, said as such.
func checkPlaceholders(m catalog.Model, tpl *comfy.Template) error {
	var missing []string
	for name, hasDefault := range tpl.Placeholders() {
		if hasDefault || slices.Contains(comfy.Builtin, name) ||
			slices.ContainsFunc(m.ComfyInputs(), func(d catalog.Input) bool { return d.Name == name }) {
			continue
		}
		missing = append(missing, name)
	}
	if len(missing) > 0 {
		sort.Strings(missing)
		return fmt.Errorf("the workflow has placeholders %s that the catalog declares no input for and that have no default",
			strings.Join(missing, ", "))
	}
	return nil
}

// pickOutput is the file a run answers with: the first the workflow saved of
// the model's output kind, or the first of any kind for files.
func pickOutput(kind string, files []comfy.Image) (comfy.Image, error) {
	is := func(ct string) bool {
		switch kind {
		case catalog.OutputPictures:
			return strings.HasPrefix(ct, "image/") && ct != "image/gif"
		case catalog.OutputMeshes:
			return isMeshType(ct)
		case catalog.OutputVideo:
			return strings.HasPrefix(ct, "video/") || ct == "image/gif" || ct == "image/webp"
		case catalog.OutputAudio:
			return strings.HasPrefix(ct, "audio/")
		}
		return true
	}
	for _, f := range files {
		if is(f.ContentType) {
			return f, nil
		}
	}
	saved := make([]string, 0, len(files))
	for _, f := range files {
		saved = append(saved, path.Base(f.Name)+" ("+f.ContentType+")")
	}
	return comfy.Image{}, fmt.Errorf("the workflow saved no %s: it saved %s", kind, strings.Join(saved, ", "))
}

// withLicense is a file with its licence written in, where its format has a
// place for one: a GLB's asset.extras, a PNG's text. Any other is returned as
// it is, and the licence goes beside it in the job's answer only.
func withLicense(data []byte, ct string, lic catalog.License) ([]byte, error) {
	switch ct {
	case "model/gltf-binary":
		return glbWithLicense(data, lic)
	case "image/png":
		return pngWithLicense(data, lic)
	}
	return data, nil
}

// pngWithLicense is a PNG with an iTXt chunk "License" after its header,
// holding the licence as JSON. The workflow's own text chunks were stripped
// when the file was read; this one is the operator's to say.
func pngWithLicense(png []byte, lic catalog.License) ([]byte, error) {
	const sig = "\x89PNG\r\n\x1a\n"
	if len(png) < len(sig)+25 || string(png[:len(sig)]) != sig || string(png[12:16]) != "IHDR" {
		return nil, errors.New("not a PNG")
	}
	text, err := json.Marshal(lic)
	if err != nil {
		return nil, err
	}
	// keyword, NUL, compression flag 0, method 0, empty language, NUL, empty
	// translated keyword, NUL, then the UTF-8 text.
	var chunk bytes.Buffer
	chunk.WriteString("iTXt")
	chunk.WriteString("License\x00\x00\x00\x00\x00")
	chunk.Write(text)
	body := chunk.Bytes()
	out := make([]byte, 0, len(png)+len(body)+8)
	ihdrEnd := len(sig) + 8 + 13 + 4
	out = append(out, png[:ihdrEnd]...)
	out = binary.BigEndian.AppendUint32(out, uint32(len(body)-4)) //nolint:gosec // a licence's JSON
	out = append(out, body...)
	out = binary.BigEndian.AppendUint32(out, crc32.ChecksumIEEE(body))
	return append(out, png[ihdrEnd:]...), nil
}