llmgw/internal/router/workflows_test.go

package router

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"image"
	"image/png"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"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/store"
)

// testWAV is the smallest WAV header and a few samples.
func testWAV() []byte {
	var b bytes.Buffer
	data := []byte{0, 0, 1, 0, 2, 0, 3, 0}
	b.WriteString("RIFF")
	_ = binary.Write(&b, binary.LittleEndian, uint32(36+len(data)))
	b.WriteString("WAVEfmt ")
	_ = binary.Write(&b, binary.LittleEndian, []any{uint32(16), uint16(1), uint16(1), uint32(8000), uint32(16000), uint16(2), uint16(16)})
	b.WriteString("data")
	_ = binary.Write(&b, binary.LittleEndian, uint32(len(data)))
	b.Write(data)
	return b.Bytes()
}

// runComfy is a ComfyUI whose workflow saves what it is told to: a song, a
// picture, under whatever key its save node uses.
type runComfy struct {
	srv      *httptest.Server
	mu       sync.Mutex
	uploaded []string
	workflow map[string]any
	out      []byte
	outName  string
}

func newRunComfy(t *testing.T, out []byte, outName string) *runComfy {
	c := &runComfy{out: out, outName: outName}
	c.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/upload/image":
			f, h, err := r.FormFile("image")
			if err != nil {
				http.Error(w, err.Error(), 400)
				return
			}
			b, _ := io.ReadAll(f)
			ext, _ := sniffFile(b)
			c.mu.Lock()
			c.uploaded = append(c.uploaded, h.Filename+":"+ext)
			c.mu.Unlock()
			fmt.Fprintf(w, `{"name":%q,"subfolder":"","type":"input"}`, h.Filename)
		case "/prompt":
			var b struct {
				Prompt map[string]any `json:"prompt"`
			}
			_ = json.NewDecoder(r.Body).Decode(&b)
			c.mu.Lock()
			c.workflow = b.Prompt
			c.mu.Unlock()
			fmt.Fprint(w, `{"prompt_id":"r1"}`)
		case "/history/r1":
			fmt.Fprintf(w, `{"r1":{"status":{"status_str":"success","completed":true},"outputs":{
				"9":{"audio":[{"filename":%q,"subfolder":"","type":"output"}]}}}}`, c.outName)
		case "/view":
			w.Header().Set("Content-Type", "application/octet-stream")
			_, _ = w.Write(c.out)
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(c.srv.Close)
	return c
}

const songWorkflow = `{
  "1": {"class_type": "LoadImage", "inputs": {"image": "{{cover}}"}},
  "2": {"class_type": "TextEncodeAceStepAudio", "inputs": {"tags": "{{tags}}", "lyrics": "{{lyrics|}}", "strength": "{{strength}}"}},
  "3": {"class_type": "EmptyAceStepLatentAudio", "inputs": {"seconds": "{{seconds}}", "style": "{{style}}"}},
  "4": {"class_type": "KSampler", "inputs": {"seed": "{{seed}}", "model": "{{checkpoint}}"}},
  "9": {"class_type": "SaveAudio", "inputs": {"filename_prefix": "audio/song"}}
}`

func f64(v float64) *float64 { return &v }

func runRouter(t *testing.T, c *runComfy, workflow string, output string) (*Router, string, string) {
	t.Helper()
	cfg := baseConfig()
	cfg.Models = []catalog.Model{
		{Name: "ace", Source: "external", License: &catalog.License{ID: "Apache-2.0", Name: "Apache 2.0"},
			External: &catalog.External{URL: c.srv.URL, Model: "ace_step.safetensors", API: catalog.APIComfyUI,
				Output: output, Workflow: "w.json"},
			ComfyUI: &catalog.ComfyUI{Inputs: []catalog.Input{
				{Name: "tags", Type: catalog.InputText, Required: true, Max: f64(200), Doc: "genres and moods"},
				{Name: "lyrics", Type: catalog.InputText},
				{Name: "strength", Type: catalog.InputFloat, Min: f64(0), Max: f64(1), Default: 0.7},
				{Name: "seconds", Type: catalog.InputInt, Min: f64(5), Max: f64(240), Default: 30},
				{Name: "style", Type: catalog.InputChoice, Choices: []string{"pop", "metal"}, Default: "pop"},
				{Name: "cover", Type: catalog.InputImage, Required: true},
			}}},
		{Name: "flux", Source: "external", External: &catalog.External{URL: c.srv.URL,
			Model: "flux1-schnell-fp8.safetensors", API: catalog.APIComfyUI}},
	}
	cfg.Auth.Enabled = true
	cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
	cfg.Auth.FlushInterval = 20 * time.Millisecond
	cfg.Auth.BatchSize = 1
	rt := newTestRouter(t, cfg)
	rt.comfy = &comfy.Client{Poll: time.Millisecond}
	rt.SetEndpoints(t.Context(), []endpoint.Info{
		{ID: "external:ace", URL: c.srv.URL, Models: []string{"ace"}, Engine: endpoint.EngineComfyUI,
			Pool: "external", Weight: 1, UpstreamModel: "ace_step.safetensors", Workflow: workflow},
		{ID: "external:flux", URL: c.srv.URL, Models: []string{"flux"}, Engine: endpoint.EngineComfyUI,
			Pool: "external", Weight: 1, UpstreamModel: "flux1-schnell-fp8.safetensors"},
	})
	org, _ := rt.store.CreateOrg(t.Context(), "studio")
	team, _ := rt.store.CreateTeam(t.Context(), org, "sound")
	project, err := rt.store.CreateProject(t.Context(), team, "score")
	if err != nil {
		t.Fatal(err)
	}
	secret := store.KeyPrefix + strings.Repeat("w", 43)
	sum := sha256.Sum256([]byte(secret))
	if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-w", Hash: sum[:], Prefix: "wwww", Name: "sound",
		ProjectID: project, CreatedAt: time.Now()}); err != nil {
		t.Fatal(err)
	}
	return rt, secret, project
}

func TestAWorkflowRunsFromTheInputsItsCatalogEntryDeclares(t *testing.T) {
	wav := testWAV()
	c := newRunComfy(t, wav, "song_00001_.wav")
	rt, secret, _ := runRouter(t, c, songWorkflow, catalog.OutputAudio)
	h := rt.Handler()

	code, j, raw := call(t, h, secret, http.MethodPost, workflowsPath, map[string]any{"model": "ace", "seed": 5,
		"inputs": map[string]any{"tags": "synthwave, driving", "seconds": 60, "cover": "data:image/png;base64," + testPicture(t)}})
	if code != http.StatusAccepted {
		t.Fatalf("%d %s", code, raw)
	}
	got := waitJob(t, h, secret, j["id"].(string), jobDone)
	if got["status"] != "completed" || got["content_type"] != "audio/wav" {
		t.Fatalf("job %v", got)
	}
	code, _, file := call(t, h, secret, http.MethodGet, "/v1/files/"+got["output_file_id"].(string)+"/content", nil)
	if code != http.StatusOK || file != string(wav) {
		t.Fatalf("the result is not the song: %d", code)
	}
	c.mu.Lock()
	wf, _ := json.Marshal(c.workflow)
	uploaded := strings.Join(c.uploaded, ",")
	c.mu.Unlock()
	if !strings.HasPrefix(uploaded, "llmgw-") || !strings.HasSuffix(uploaded, ".png:.png") {
		t.Errorf("uploaded %q", uploaded)
	}
	name := strings.TrimSuffix(uploaded, ":.png")
	for _, want := range []string{`"image":"` + name + `"`, `"tags":"synthwave, driving"`, `"seconds":60`,
		`"strength":0.7`, `"style":"pop"`, `"lyrics":""`, `"seed":5`, `"model":"ace_step.safetensors"`} {
		if !strings.Contains(string(wf), want) {
			t.Errorf("%s not in the workflow sent: %s", want, wf)
		}
	}
}

func TestARunIsRefusedWhatItsWorkflowDoesNotTake(t *testing.T) {
	c := newRunComfy(t, testWAV(), "song.wav")
	rt, secret, _ := runRouter(t, c, songWorkflow, catalog.OutputAudio)
	h := rt.Handler()
	cover := testPicture(t)
	in := func(extra map[string]any) map[string]any {
		m := map[string]any{"tags": "lofi", "cover": cover}
		for k, v := range extra {
			m[k] = v
		}
		return map[string]any{"model": "ace", "inputs": m}
	}
	for name, c := range map[string]struct {
		body map[string]any
		want string
	}{
		"no tags":             {map[string]any{"model": "ace", "inputs": map[string]any{"cover": cover}}, `input "tags" is required: text — genres and moods`},
		"an unknown input":    {in(map[string]any{"tempo": 120}), `takes no input "tempo"`},
		"a number as text":    {in(map[string]any{"tags": 5}), `input "tags" is text`},
		"too long":            {in(map[string]any{"tags": strings.Repeat("x", 201)}), "past 200"},
		"out of bounds":       {in(map[string]any{"strength": 2}), "outside 0 to 1"},
		"not whole":           {in(map[string]any{"seconds": 10.5}), "whole number"},
		"not a choice":        {in(map[string]any{"style": "jazz"}), "one of pop, metal"},
		"a song as a picture": {in(map[string]any{"cover": base64.StdEncoding.EncodeToString(testWAV())}), "PNG, JPEG"},
		"someone else's file": {in(map[string]any{"cover": "file-nothere"}), `no file "file-nothere"`},
		"a run of flux's":     {map[string]any{"model": "flux", "inputs": map[string]any{"tags": "x"}}, "takes no inputs"},
		"a chat with ace":     {map[string]any{"model": "ace", "messages": []any{}}, "makes audio: ask for it at " + workflowsPath},
	} {
		path := workflowsPath
		if name == "a chat with ace" {
			path = "/v1/chat/completions"
		}
		code, _, raw := call(t, h, secret, http.MethodPost, path, c.body)
		if code != http.StatusBadRequest || !strings.Contains(errorMessage([]byte(raw)), c.want) {
			t.Errorf("%s: %d %s", name, code, raw)
		}
	}
}

func TestARunTakesAFileUploadedBeforeAndNamedByItsID(t *testing.T) {
	c := newRunComfy(t, testWAV(), "song.wav")
	rt, secret, project := runRouter(t, c, songWorkflow, catalog.OutputAudio)
	h := rt.Handler()
	pic, _ := base64.StdEncoding.DecodeString(testPicture(t))
	now := time.Now()
	if err := rt.store.PutFile(t.Context(), store.File{ID: "file-cover", Owner: "project:" + project, Purpose: "user_data",
		Filename: "cover.png", Bytes: int64(len(pic)), CreatedAt: now, ExpiresAt: now.Add(time.Hour)}, pic); err != nil {
		t.Fatal(err)
	}
	if err := rt.store.PutFile(t.Context(), store.File{ID: "file-theirs", Owner: "project:other", Purpose: "user_data",
		Filename: "x.png", Bytes: int64(len(pic)), CreatedAt: now, ExpiresAt: now.Add(time.Hour)}, pic); err != nil {
		t.Fatal(err)
	}
	code, _, raw := call(t, h, secret, http.MethodPost, workflowsPath,
		map[string]any{"model": "ace", "inputs": map[string]any{"tags": "x", "cover": "file-theirs"}})
	if code != http.StatusBadRequest || !strings.Contains(errorMessage([]byte(raw)), `no file "file-theirs"`) {
		t.Errorf("another project's file: %d %s", code, raw)
	}
	code, j, raw := call(t, h, secret, http.MethodPost, workflowsPath,
		map[string]any{"model": "ace", "inputs": map[string]any{"tags": "x", "cover": "file-cover"}})
	if code != http.StatusAccepted {
		t.Fatalf("%d %s", code, raw)
	}
	if got := waitJob(t, h, secret, j["id"].(string), jobDone); got["status"] != "completed" {
		t.Fatalf("job %v", got)
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if len(c.uploaded) != 1 || !strings.HasSuffix(c.uploaded[0], ".png:.png") {
		t.Errorf("uploaded %v", c.uploaded)
	}
}

func TestAWorkflowWantingWhatNoInputGivesFailsSayingSo(t *testing.T) {
	c := newRunComfy(t, testWAV(), "song.wav")
	wf := strings.Replace(songWorkflow, `"{{style}}"`, `"{{mood}}"`, 1)
	rt, secret, _ := runRouter(t, c, wf, catalog.OutputAudio)
	h := rt.Handler()
	code, j, raw := call(t, h, secret, http.MethodPost, workflowsPath,
		map[string]any{"model": "ace", "inputs": map[string]any{"tags": "x", "cover": testPicture(t)}})
	if code != http.StatusAccepted {
		t.Fatalf("%d %s", code, raw)
	}
	got := waitJob(t, h, secret, j["id"].(string), jobDone)
	if e, _ := got["error"].(map[string]any); got["status"] != "failed" || !strings.Contains(fmt.Sprint(e["message"]), "placeholders mood") {
		t.Errorf("job %v", got)
	}
}

func TestARunThatSavedTheWrongKindFailsAndAPictureCarriesItsLicence(t *testing.T) {
	c := newRunComfy(t, testWAV(), "song.wav")
	rt, secret, _ := runRouter(t, c, songWorkflow, catalog.OutputVideo)
	h := rt.Handler()
	_, j, _ := call(t, h, secret, http.MethodPost, workflowsPath,
		map[string]any{"model": "ace", "inputs": map[string]any{"tags": "x", "cover": testPicture(t)}})
	got := waitJob(t, h, secret, j["id"].(string), jobDone)
	if e, _ := got["error"].(map[string]any); !strings.Contains(fmt.Sprint(e["message"]), "saved no video: it saved song.wav (audio/wav)") {
		t.Errorf("job %v", got)
	}

	var b bytes.Buffer
	_ = png.Encode(&b, image.NewRGBA(image.Rect(0, 0, 4, 4)))
	with, err := pngWithLicense(b.Bytes(), catalog.License{ID: "Apache-2.0", Name: "Apache 2.0", Attribution: "Made with ACE"})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := png.Decode(bytes.NewReader(with)); err != nil {
		t.Fatalf("not a PNG any more: %v", err)
	}
	if !bytes.Contains(with, []byte("iTXtLicense\x00\x00\x00\x00\x00{\"id\":\"Apache-2.0\"")) {
		t.Error("no licence chunk in the picture")
	}
	if _, err := pngWithLicense([]byte("GIF89a"), catalog.License{Name: "x"}); err == nil {
		t.Error("a GIF taken for a PNG")
	}
}