llmgw/internal/comfy/gate.go

package comfy

import (
	"context"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// The gate is how a node runs ComfyUI as one of llmgw's engines.
//
// ComfyUI asks nobody for a key: whoever reaches its port can queue work,
// upload files and install nodes. Every other engine here demands the key
// the control plane made for its replica, so ComfyUI is put behind a gate
// that does: ComfyUI listens on the loopback only, and the gate on the
// replica's port, letting through only the requests the router makes. The
// gate also answers what the agent asks every engine — /health once
// ComfyUI is up, and /v1/models naming the model — and lays the weights out
// where ComfyUI looks for them.
//
// ComfyUI runs as the gate's child, in its process group, so stopping the
// replica stops both, the card memory ComfyUI holds is counted as the
// replica's, and either one exiting ends the other.

// GateConfig is one replica's gate.
type GateConfig struct {
	// Listen is the replica's address, the one the router is given.
	Listen string
	// Model is the name the agent asks /v1/models for.
	Model string
	// Key is what every request must carry as a bearer token; empty lets
	// every request through, as an engine started without one does.
	Key string
	// Weights is the model as the agent fetched it: a file, or a directory
	// of files.
	Weights string
	// State is the replica's own directory: ComfyUI's base directory, with
	// its output, temp and user directories, and the models it is shown.
	State string
	// Python runs ComfyUI's main.py, found in Dir.
	Python, Dir string
	// Args are passed on to main.py, after the gate's own.
	Args []string
	// Output is where ComfyUI's output goes.
	Output io.Writer
	Log    *slog.Logger
}

// The requests the router makes of ComfyUI: queueing, waiting, fetching and
// cancelling. Nothing that uploads, installs or changes ComfyUI.
var gatePaths = []string{"/prompt", "/history/", "/view", "/queue", "/interrupt", "/system_stats", "/object_info",
	// The input picture of a 3D model: into ComfyUI's input directory,
	// which is the replica's own.
	"/upload/image"}

// modelFolders are the kinds of file ComfyUI's loaders look for. Each is
// shown the same directory: a workflow names its files, so a file being
// listed where it is not wanted costs nothing.
var modelFolders = []string{"checkpoints", "diffusion_models", "text_encoders", "vae", "loras",
	"clip_vision", "controlnet", "upscale_models", "embeddings", "style_models"}

// RunGate lays out the replica, starts ComfyUI and serves the gate until
// ctx ends or ComfyUI exits.
func RunGate(ctx context.Context, cfg GateConfig) error {
	log := cfg.Log
	if log == nil {
		log = slog.Default()
	}
	state, err := filepath.Abs(cfg.State)
	if err != nil {
		return err
	}
	// What ComfyUI expects to find in its base directory, and dies
	// without: custom_nodes is listed before anything else is done.
	for _, d := range []string{"custom_nodes", "input", "output", "temp", "user"} {
		if err := os.MkdirAll(filepath.Join(state, d), 0o755); err != nil {
			return err
		}
	}
	out := cfg.Output
	if out == nil {
		out = os.Stdout
	}
	if err := installOwnNodes(state); err != nil {
		return fmt.Errorf("llmgw's own ComfyUI nodes: %w", err)
	}
	nodes, err := parseNodes(os.Getenv(NodesEnv))
	if err != nil {
		return err
	}
	site, err := installNodes(ctx, state, cfg.Python, nodes, findUV(), func(s string) { fmt.Fprintln(out, "[gate] "+s) })
	if err != nil {
		return err
	}
	models := filepath.Join(state, "models")
	if err := LinkWeights(cfg.Weights, models); err != nil {
		return err
	}
	if repo := os.Getenv(HFRepoEnv); repo != "" {
		if err := LinkHFCache(cfg.Weights, repo, hfHubDir()); err != nil {
			return fmt.Errorf("showing the weights as %s: %w", repo, err)
		}
	}
	extra := filepath.Join(state, "extra_model_paths.yaml")
	if err := os.WriteFile(extra, []byte(extraModelPaths(models)), 0o644); err != nil { //nolint:gosec // not secret
		return err
	}

	inner, err := freePort()
	if err != nil {
		return err
	}
	args := append([]string{filepath.Join(cfg.Dir, "main.py"),
		"--listen", "127.0.0.1", "--port", strconv.Itoa(inner),
		"--base-directory", state, "--extra-model-paths-config", extra,
		"--disable-auto-launch"}, cfg.Args...)
	child := exec.Command(cfg.Python, args...) //nolint:gosec // the agent's own configuration
	child.Dir = cfg.Dir
	child.Env = os.Environ()
	if site != "" {
		pp := site
		if old := os.Getenv("PYTHONPATH"); old != "" {
			pp += string(os.PathListSeparator) + old
		}
		child.Env = append(child.Env, "PYTHONPATH="+pp)
	}
	// ComfyUI writes to a pipe of the gate's, never to what the gate was
	// given. That is the agent's, and an agent that restarts and adopts the
	// replica leaves it with no reader: ComfyUI's next progress bar would
	// fail with a broken pipe, and take the workflow it belongs to with it.
	// The gate reads the pipe for as long as ComfyUI writes, and passes on
	// what it can.
	signal.Ignore(syscall.SIGPIPE)
	pr, err := child.StdoutPipe()
	if err != nil {
		return err
	}
	child.Stderr = child.Stdout
	if err := child.Start(); err != nil {
		return fmt.Errorf("starting ComfyUI: %w", err)
	}
	relayed := make(chan struct{})
	go func() { relay(pr, out); close(relayed) }()
	exited := make(chan error, 1)
	go func() {
		// Wait closes the pipe; what is in it is read first.
		<-relayed
		exited <- child.Wait()
	}()

	upstream := &url.URL{Scheme: "http", Host: net.JoinHostPort("127.0.0.1", strconv.Itoa(inner))}
	// The workflows' progress, heard on ComfyUI's websocket under an id of
	// the gate's own.
	book := newProgressBook()
	cid := make([]byte, 8)
	_, _ = rand.Read(cid)
	clientID := "llmgw-gate-" + hex.EncodeToString(cid)
	go book.watch(ctx, upstream, clientID, log)
	srv := &http.Server{
		Addr:              cfg.Listen,
		Handler:           newGate(upstream, cfg.Key, cfg.Model, book, clientID),
		ReadHeaderTimeout: 10 * time.Second,
	}
	served := make(chan error, 1)
	go func() { served <- srv.ListenAndServe() }()
	log.Info("ComfyUI gate up", "listen", cfg.Listen, "comfyui", upstream.Host, "weights", cfg.Weights)

	stop := func() {
		_ = child.Process.Signal(syscall.SIGTERM)
		select {
		case <-exited:
		case <-time.After(20 * time.Second):
			_ = child.Process.Kill()
			<-exited
		}
	}
	select {
	case <-ctx.Done():
		shutdown(srv)
		stop()
		return nil
	case err := <-exited:
		shutdown(srv)
		if err == nil {
			err = errors.New("exited")
		}
		return fmt.Errorf("ComfyUI %w", err)
	case err := <-served:
		stop()
		return fmt.Errorf("the gate stopped serving: %w", err)
	}
}

func shutdown(srv *http.Server) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = srv.Shutdown(ctx)
}

// NewGate is the gate's handler, in front of the ComfyUI at upstream.
func NewGate(upstream *url.URL, key, model string) http.Handler {
	return newGate(upstream, key, model, newProgressBook(), "")
}

// newGate is the gate with the book of its workflows' progress, kept under
// clientID; an empty one leaves the workflows' own client ids alone.
func newGate(upstream *url.URL, key, model string, book *progressBook, clientID string) http.Handler {
	proxy := httputil.NewSingleHostReverseProxy(upstream)
	base := proxy.Director
	proxy.Director = func(r *http.Request) {
		base(r)
		r.Header.Del("Authorization")
		// ComfyUI refuses a request whose Origin is not its own host; the
		// router sends none, and nothing else gets this far.
		r.Header.Del("Origin")
		r.Host = upstream.Host
	}
	client := &http.Client{Timeout: 3 * time.Second}
	up := func(ctx context.Context) error {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, upstream.String()+"/system_stats", nil)
		resp, err := client.Do(req)
		if err != nil {
			return err
		}
		resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			return fmt.Errorf("ComfyUI answered %s", resp.Status)
		}
		return nil
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if key != "" {
			got, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
			if !ok || subtle.ConstantTimeCompare([]byte(got), []byte(key)) != 1 {
				http.Error(w, "this engine wants its key", http.StatusUnauthorized)
				return
			}
		}
		switch r.URL.Path {
		case "/health":
			if err := up(r.Context()); err != nil {
				http.Error(w, "ComfyUI is not up yet: "+err.Error(), http.StatusServiceUnavailable)
				return
			}
			w.WriteHeader(http.StatusOK)
			return
		case "/v1/models":
			if err := up(r.Context()); err != nil {
				http.Error(w, "ComfyUI is not up yet: "+err.Error(), http.StatusServiceUnavailable)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			fmt.Fprintf(w, `{"object":"list","data":[{"id":%q,"object":"model","owned_by":"comfyui"}]}`, model)
			return
		}
		if id, ok := strings.CutPrefix(r.URL.Path, "/progress/"); ok && r.Method == http.MethodGet {
			p, found := book.get(id)
			if !found {
				http.Error(w, "no word of that workflow", http.StatusNotFound)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_ = json.NewEncoder(w).Encode(p)
			return
		}
		if r.URL.Path == "/prompt" && r.Method == http.MethodPost && clientID != "" {
			book.promptThroughGate(proxy, clientID)(w, r)
			return
		}
		for _, p := range gatePaths {
			if r.URL.Path == p || strings.HasSuffix(p, "/") && strings.HasPrefix(r.URL.Path, p) {
				proxy.ServeHTTP(w, r)
				return
			}
		}
		http.NotFound(w, r)
	})
}

// cachedName is a file the agent fetched by itself: its cache names it by
// the source's hash, then the name it had.
var cachedName = regexp.MustCompile(`^[0-9a-f]{16}-`)

// FileName is what a weights file is called where ComfyUI sees it: its own
// name, without the prefix the agent's cache gives it.
func FileName(path string) string {
	return cachedName.ReplaceAllString(filepath.Base(path), "")
}

// LinkWeights shows ComfyUI the weights: a link in dir to the file, or to
// every file of the directory. Links made before are taken away first, so
// a replica started again on other weights sees only those.
func LinkWeights(weights, dir string) error {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	old, err := os.ReadDir(dir)
	if err != nil {
		return err
	}
	for _, e := range old {
		if e.Type()&os.ModeSymlink != 0 {
			_ = os.Remove(filepath.Join(dir, e.Name()))
		}
	}
	weights, err = filepath.Abs(weights)
	if err != nil {
		return err
	}
	info, err := os.Stat(weights)
	if err != nil {
		return fmt.Errorf("the weights: %w", err)
	}
	files := []string{weights}
	if info.IsDir() {
		entries, err := os.ReadDir(weights)
		if err != nil {
			return err
		}
		files = files[:0]
		for _, e := range entries {
			if e.Type().IsRegular() && !strings.HasPrefix(e.Name(), ".") {
				files = append(files, filepath.Join(weights, e.Name()))
			}
		}
		if len(files) == 0 {
			return fmt.Errorf("the weights directory %s has no files", weights)
		}
	}
	for _, f := range files {
		if err := os.Symlink(f, filepath.Join(dir, FileName(f))); err != nil {
			return err
		}
	}
	return nil
}

func extraModelPaths(models string) string {
	var b strings.Builder
	fmt.Fprintf(&b, "llmgw:\n  base_path: %q\n  is_default: true\n", models)
	for _, f := range modelFolders {
		fmt.Fprintf(&b, "  %s: .\n", f)
	}
	return b.String()
}

func freePort() (int, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}

// relay copies r to w for as long as r has anything, and goes on reading
// after w fails: a writer that stops reading must not stop the one writing.
func relay(r io.Reader, w io.Writer) {
	buf := make([]byte, 32<<10)
	broken := false
	for {
		n, err := r.Read(buf)
		if n > 0 && !broken {
			if _, werr := w.Write(buf[:n]); werr != nil {
				broken = true
			}
		}
		if err != nil {
			return
		}
	}
}

// HFRepoEnv is where the agent hands the gate the repository a custom node
// loads the weights by name from.
const HFRepoEnv = "LLMGW_COMFYUI_HF_REPO"

// hfHubDir is the Hugging Face cache of the user the gate runs as.
func hfHubDir() string {
	if d := os.Getenv("HF_HUB_CACHE"); d != "" {
		return d
	}
	if d := os.Getenv("HF_HOME"); d != "" {
		return filepath.Join(d, "hub")
	}
	home, _ := os.UserHomeDir()
	return filepath.Join(home, ".cache", "huggingface", "hub")
}

// LinkHFCache shows the weights to a node that loads them by repository
// name: a snapshot of repo in the Hugging Face cache at hub, each file a link
// to the weights, refs/main naming it. huggingface_hub then answers from the
// cache, and a gated repository needs no token on the node.
func LinkHFCache(weights, repo, hub string) error {
	owner, name, ok := strings.Cut(repo, "/")
	if !ok || owner == "" || name == "" || strings.Contains(name, "/") || strings.Contains(repo, "..") || strings.Contains(repo, "\\") {
		return fmt.Errorf("%q is not owner/name", repo)
	}
	weights, err := filepath.Abs(weights)
	if err != nil {
		return err
	}
	// The snapshot is named as a commit is, which is what refs/main must
	// hold, and by the weights, so other weights are another snapshot.
	sum := sha1.Sum([]byte(weights)) //nolint:gosec // a name, not a secret
	revision := hex.EncodeToString(sum[:])
	info, err := os.Stat(weights)
	if err != nil {
		return err
	}
	files := []string{weights}
	if info.IsDir() {
		entries, err := os.ReadDir(weights)
		if err != nil {
			return err
		}
		files = files[:0]
		for _, e := range entries {
			if !e.IsDir() && !strings.HasPrefix(e.Name(), ".") {
				files = append(files, filepath.Join(weights, e.Name()))
			}
		}
	}
	root := filepath.Join(hub, "models--"+owner+"--"+name)
	snap := filepath.Join(root, "snapshots", revision)
	// Made again each time: links to weights that moved must not stay.
	_ = os.RemoveAll(snap)
	if err := os.MkdirAll(snap, 0o755); err != nil {
		return err
	}
	for _, f := range files {
		if err := os.Symlink(f, filepath.Join(snap, FileName(f))); err != nil {
			return err
		}
	}
	if err := os.MkdirAll(filepath.Join(root, "refs"), 0o755); err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(root, "refs", "main"), []byte(revision), 0o644) //nolint:gosec // not secret
}