llmgw/internal/comfy/progress.go

package comfy

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/websocket"
)

// How far along a workflow is, which ComfyUI says only on its websocket and
// only to the client that queued the workflow. The gate keeps one websocket
// open to its ComfyUI, under a client id of its own, puts that id on every
// workflow it lets through, and keeps the last word on each: the node
// running and its step out of how many. The router asks at
// /progress/{prompt_id} while it waits, and a job says it.

// Progress is a workflow's state as the gate last heard it.
type Progress struct {
	PromptID string `json:"prompt_id,omitempty"`
	// State is queued, running, done or failed.
	State string `json:"state"`
	// Ahead is how many workflows wait before it, while it is queued.
	Ahead int `json:"ahead,omitempty"`
	// Node is the workflow's node running, by its id; NodeClass its kind.
	Node      string `json:"node,omitempty"`
	NodeClass string `json:"node_class,omitempty"`
	// Step of Steps is the running node's own count: a sampler's steps.
	Step  int `json:"step,omitempty"`
	Steps int `json:"steps,omitempty"`
	// Percent is Step of Steps, where a node counts.
	Percent int    `json:"percent,omitempty"`
	Error   string `json:"error,omitempty"`
}

// progressBook is the last word on each of the most recent workflows.
type progressBook struct {
	mu    sync.Mutex
	by    map[string]*Progress
	order []string
}

// keptPrompts bounds the book: a gate serves one replica, which runs one
// workflow at a time and queues a few.
const keptPrompts = 256

func newProgressBook() *progressBook { return &progressBook{by: map[string]*Progress{}} }

func (b *progressBook) update(id string, f func(*Progress)) {
	if id == "" {
		return
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	p, ok := b.by[id]
	if !ok {
		p = &Progress{PromptID: id, State: "queued"}
		b.by[id] = p
		b.order = append(b.order, id)
		for len(b.order) > keptPrompts {
			delete(b.by, b.order[0])
			b.order = b.order[1:]
		}
	}
	f(p)
}

func (b *progressBook) get(id string) (Progress, bool) {
	b.mu.Lock()
	defer b.mu.Unlock()
	p, ok := b.by[id]
	if !ok {
		return Progress{}, false
	}
	return *p, true
}

// read takes one message off ComfyUI's websocket into the book.
func (b *progressBook) read(msg []byte) {
	var m struct {
		Type string `json:"type"`
		Data struct {
			PromptID string          `json:"prompt_id"`
			Node     json.RawMessage `json:"node"`
			Value    int             `json:"value"`
			Max      int             `json:"max"`
			Message  string          `json:"exception_message"`
			NodeType string          `json:"node_type"`
		} `json:"data"`
	}
	if json.Unmarshal(msg, &m) != nil {
		return
	}
	d := m.Data
	var node string
	_ = json.Unmarshal(d.Node, &node)
	switch m.Type {
	case "execution_start":
		b.update(d.PromptID, func(p *Progress) { p.State, p.Ahead = "running", 0 })
	case "executing":
		b.update(d.PromptID, func(p *Progress) {
			if node == "" {
				// No node: the workflow is through.
				p.State, p.Node, p.Step, p.Steps, p.Percent = "done", "", 0, 0, 100
				return
			}
			p.State, p.Ahead, p.Node, p.Step, p.Steps, p.Percent = "running", 0, node, 0, 0, 0
		})
	case "progress":
		b.update(d.PromptID, func(p *Progress) {
			p.State, p.Step, p.Steps = "running", d.Value, d.Max
			if node != "" {
				p.Node = node
			}
			if d.Max > 0 {
				p.Percent = min(100, d.Value*100/d.Max)
			}
		})
	case "execution_success":
		b.update(d.PromptID, func(p *Progress) { p.State, p.Percent = "done", 100 })
	case "execution_error", "execution_interrupted":
		b.update(d.PromptID, func(p *Progress) {
			p.State = "failed"
			p.Error = strings.TrimSpace(d.NodeType + ": " + d.Message)
			if m.Type == "execution_interrupted" {
				p.Error = "interrupted"
			}
		})
	}
}

// watch keeps a websocket open to the ComfyUI at upstream under clientID,
// reading into the book, until ctx ends: opened again when it closes, as it
// does whenever ComfyUI restarts.
func (b *progressBook) watch(ctx context.Context, upstream *url.URL, clientID string, log *slog.Logger) {
	ws := "ws://" + upstream.Host + "/ws?clientId=" + url.QueryEscape(clientID)
	wait := time.Second
	for ctx.Err() == nil {
		conn, err := websocket.Dial(ws, "", "http://"+upstream.Host)
		if err == nil {
			wait = time.Second
			go func() { <-ctx.Done(); conn.Close() }()
			for {
				var msg []byte
				if err := websocket.Message.Receive(conn, &msg); err != nil {
					break
				}
				// Binary frames are previews of pictures; only text is news.
				if len(msg) > 0 && msg[0] == '{' {
					b.read(msg)
				}
			}
			conn.Close()
		} else if log != nil {
			log.Debug("ComfyUI's websocket is not open yet", "err", err)
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(wait):
		}
		wait = min(wait*2, 30*time.Second)
	}
}

// withClient is a /prompt body with the gate's client id in it, so the
// progress of the workflow it queues comes to the gate's websocket.
func withClient(body []byte, clientID string) []byte {
	var m map[string]json.RawMessage
	if json.Unmarshal(body, &m) != nil {
		return body
	}
	id, _ := json.Marshal(clientID)
	m["client_id"] = id
	out, err := json.Marshal(m)
	if err != nil {
		return body
	}
	return out
}

// promptThroughGate forwards a /prompt with the gate's client id, and books
// the workflow it queued as queued.
func (b *progressBook) promptThroughGate(proxy http.Handler, clientID string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(io.LimitReader(r.Body, 64<<20))
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		body = withClient(body, clientID)
		r.Body = io.NopCloser(bytes.NewReader(body))
		r.ContentLength = int64(len(body))
		rec := &captured{ResponseWriter: w}
		proxy.ServeHTTP(rec, r)
		var out struct {
			PromptID string `json:"prompt_id"`
		}
		if json.Unmarshal(rec.body.Bytes(), &out) == nil && out.PromptID != "" {
			b.update(out.PromptID, func(*Progress) {})
		}
	}
}

// captured passes a response on and keeps a copy of the first of it.
type captured struct {
	http.ResponseWriter
	body bytes.Buffer
}

func (c *captured) Write(p []byte) (int, error) {
	if c.body.Len() < 64<<10 {
		c.body.Write(p)
	}
	return c.ResponseWriter.Write(p)
}