llmgw/internal/router/images_test.go

package router

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// imageRouter has a ComfyUI model, flux, beside a text model, and returns
// what the ComfyUI server was sent.
func imageRouter(t *testing.T) (*Router, *map[string]any, *sync.Mutex) {
	t.Helper()
	return imageRouterWith(t, nil)
}

func imageRouterWith(t *testing.T, edit func(*config.Server)) (*Router, *map[string]any, *sync.Mutex) {
	t.Helper()
	var mu sync.Mutex
	var queued map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/metrics":
			t.Error("ComfyUI was scraped for metrics")
		case "/prompt":
			var b struct {
				Prompt map[string]any `json:"prompt"`
			}
			_ = json.NewDecoder(r.Body).Decode(&b)
			mu.Lock()
			queued = b.Prompt
			mu.Unlock()
			fmt.Fprint(w, `{"prompt_id":"p1"}`)
		case "/history/p1":
			mu.Lock()
			n := 1
			if latent, ok := queued["4"].(map[string]any); ok {
				n = int(latent["inputs"].(map[string]any)["batch_size"].(float64))
			}
			mu.Unlock()
			var images []string
			for i := range n {
				images = append(images, fmt.Sprintf(`{"filename":"%d.png","type":"temp"}`, i))
			}
			fmt.Fprintf(w, `{"p1":{"status":{"status_str":"success","completed":true},"outputs":{"7":{"images":[%s]}}}}`,
				strings.Join(images, ","))
		case "/view":
			w.Header().Set("Content-Type", "image/png")
			fmt.Fprint(w, "png-bytes")
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(srv.Close)

	cfg := baseConfig()
	cfg.Models = []catalog.Model{catModel("chat", nil), {Name: "flux", Source: "external",
		External: &catalog.External{URL: srv.URL, Model: "flux1-schnell-fp8.safetensors", API: catalog.APIComfyUI}}}
	if edit != nil {
		edit(&cfg)
	}
	rt := newTestRouter(t, cfg)
	rt.comfy = &comfy.Client{Poll: time.Millisecond}
	rt.SetEndpoints(t.Context(), []endpoint.Info{
		{ID: "external:flux", URL: srv.URL, Models: []string{"flux"}, Engine: endpoint.EngineComfyUI,
			Pool: "external", Weight: 1, UpstreamModel: "flux1-schnell-fp8.safetensors"},
	})
	return rt, &queued, &mu
}

func postImages(t *testing.T, h http.Handler, path string, body any) *httptest.ResponseRecorder {
	t.Helper()
	b, _ := json.Marshal(body)
	req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestAPictureIsMadeByComfyUIAndComesBackAsBase64(t *testing.T) {
	rt, queued, mu := imageRouter(t)
	rec := postImages(t, rt.Handler(), imagesPath, map[string]any{
		"model": "flux", "prompt": "a lighthouse at dusk", "size": "768x512", "seed": 7})
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	var out struct {
		Data []struct {
			B64 string `json:"b64_json"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || len(out.Data) != 1 {
		t.Fatalf("body %s: %v", rec.Body.String(), err)
	}
	if b, _ := base64.StdEncoding.DecodeString(out.Data[0].B64); string(b) != "png-bytes" {
		t.Errorf("picture = %q", b)
	}
	if rec.Header().Get("x-llmgw-seed") != "7" {
		t.Errorf("seed header = %q", rec.Header().Get("x-llmgw-seed"))
	}
	mu.Lock()
	defer mu.Unlock()
	wf, _ := json.Marshal(*queued)
	for _, want := range []string{`"ckpt_name":"flux1-schnell-fp8.safetensors"`, `"text":"a lighthouse at dusk"`,
		`"width":768`, `"height":512`, `"seed":7`} {
		if !strings.Contains(string(wf), want) {
			t.Errorf("%s not in the workflow sent: %s", want, wf)
		}
	}
}

func TestPicturesAndTextAreAskedForAtTheirOwnPaths(t *testing.T) {
	rt, _, _ := imageRouter(t)
	h := rt.Handler()
	for name, c := range map[string]struct {
		path string
		body map[string]any
		want string
	}{
		"text of flux":     {"/v1/chat/completions", map[string]any{"model": "flux", "messages": []any{}}, "ask for them at " + imagesPath},
		"picture of chat":  {imagesPath, map[string]any{"model": "chat", "prompt": "x"}, "does not make pictures"},
		"no prompt":        {imagesPath, map[string]any{"model": "flux"}, "prompt is required"},
		"odd size":         {imagesPath, map[string]any{"model": "flux", "prompt": "x", "size": "1000x1000"}, "multiple of 16"},
		"url wanted":       {imagesPath, map[string]any{"model": "flux", "prompt": "x", "response_format": "url"}, "b64_json"},
		"too many at once": {imagesPath, map[string]any{"model": "flux", "prompt": "x", "n": 9}, "n is between"},
	} {
		rec := postImages(t, h, c.path, c.body)
		if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), c.want) {
			t.Errorf("%s: %d %s", name, rec.Code, rec.Body.String())
		}
	}
}

func TestAComfyUIModelSaysItMakesPictures(t *testing.T) {
	rt, _, _ := imageRouter(t)
	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)
	if !strings.Contains(rec.Body.String(), `"images":true`) {
		t.Errorf("/v1/models: %s", rec.Body.String())
	}
}

func TestAComfyUIModelIsCheckedInTheCatalog(t *testing.T) {
	ok := catalog.Model{Name: "f", Source: "external", External: &catalog.External{
		URL: "http://10.0.0.147:8188", Model: "flux1-dev.safetensors", API: catalog.APIComfyUI, Workflow: "flux-dev"}}
	if _, err := catalog.New([]catalog.Model{ok}); err != nil {
		t.Fatalf("refused: %v", err)
	}
	for name, e := range map[string]catalog.External{
		"unknown api":       {URL: "http://a:1", Model: "m", API: "a1111"},
		"unknown workflow":  {URL: "http://a:1", Model: "m", API: catalog.APIComfyUI, Workflow: "sdxl"},
		"workflow, openai":  {URL: "http://a:1", Model: "m", Workflow: "flux-dev"},
		"comfyui from peer": {URL: "http://a:1", Model: "m", API: catalog.APIComfyUI, Peer: "other"},
	} {
		e := e
		if _, err := catalog.New([]catalog.Model{{Name: "f", Source: "external", External: &e}}); err == nil {
			t.Errorf("%s: accepted", name)
		}
	}
}

func TestPicturesAreChargedByThePictureAndRecorded(t *testing.T) {
	rt, _, _ := imageRouterWith(t, func(cfg *config.Server) {
		cfg.Auth.Enabled = true
		cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
		cfg.Auth.Pricing = map[string]config.Price{"flux": {PromptPer1k: 1000, PerImage: 50000}}
		cfg.Auth.FlushInterval = 20 * time.Millisecond
		cfg.Auth.BatchSize = 1
	})
	org, _ := rt.store.CreateOrg(t.Context(), "acme")
	team, _ := rt.store.CreateTeam(t.Context(), org, "design")
	project, err := rt.store.CreateProject(t.Context(), team, "posters")
	if err != nil {
		t.Fatal(err)
	}
	secret := store.KeyPrefix + strings.Repeat("b", 43)
	sum := sha256.Sum256([]byte(secret))
	if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-1", Hash: sum[:], Prefix: "bbbb", Name: "t",
		ProjectID: project, CreatedAt: time.Now()}); err != nil {
		t.Fatal(err)
	}

	b, _ := json.Marshal(map[string]any{"model": "flux", "prompt": "three posters", "n": 3})
	req := httptest.NewRequest(http.MethodPost, imagesPath, bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+secret)
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	id := rec.Header().Get("X-Request-Id")
	deadline := time.Now().Add(5 * time.Second)
	for {
		u, ok, err := rt.store.UsageByRequest(t.Context(), id)
		if err != nil {
			t.Fatal(err)
		}
		if ok {
			if u.Images != 3 {
				t.Errorf("recorded %d pictures, want 3", u.Images)
			}
			if want := (u.PromptTokens*1000+999)/1000 + 3*50000; u.CostMicros != want {
				t.Errorf("charged %d micros, want %d (%d prompt tokens and three pictures)", u.CostMicros, want, u.PromptTokens)
			}
			return
		}
		if time.Now().After(deadline) {
			t.Fatal("no usage was recorded")
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestThePicturesPageOffersPictureModelsAndTheChatTheOthers(t *testing.T) {
	e := newFakeEngine(t, "chat")
	rt, _, _ := imageRouterWith(t, func(cfg *config.Server) {
		cfg.Auth.Enabled = true
		cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
	})
	rt.SetEndpoints(t.Context(), append(rt.reg.Infos(), endpoint.Info{ID: "c", URL: e.srv.URL, Models: []string{"chat"}}))
	org, _ := rt.store.CreateOrg(t.Context(), "acme")
	team, _ := rt.store.CreateTeam(t.Context(), org, "design")
	project, err := rt.store.CreateProject(t.Context(), team, "posters")
	if err != nil {
		t.Fatal(err)
	}
	user := store.User{ID: "usr-1", ProjectID: project}
	names := func(pictures bool) []string {
		var out []string
		for _, m := range rt.modelsOf(httptest.NewRequest("GET", "/", nil), user, pictures) {
			out = append(out, m.Name)
		}
		return out
	}
	if got := names(false); len(got) != 1 || got[0] != "chat" {
		t.Errorf("the chat offers %v", got)
	}
	if got := names(true); len(got) != 1 || got[0] != "flux" {
		t.Errorf("the pictures page offers %v", got)
	}

	// The page's request is rebuilt: a field it has no business sending
	// does not reach the model.
	body := `{"model":"flux","prompt":"a kite","size":"512x512","n":1,"seed":3,"response_format":"url"}`
	rec := httptest.NewRecorder()
	rt.imagesSend(rec, httptest.NewRequest("POST", "/images/generate", strings.NewReader(body)), user)
	if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), "b64_json") {
		t.Fatalf("%d %s", rec.Code, rec.Body.String())
	}
	if rec.Header().Get("x-llmgw-seed") != "3" {
		t.Errorf("seed %q", rec.Header().Get("x-llmgw-seed"))
	}
}