gpu-turnstile/cmd/gpu-turnstile/main.go

// gpu-turnstile is a GPU arbitration proxy that sits in front of Ollama and
// ComfyUI and guarantees only one of them uses the GPU at a time.
package main

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"gpu-turnstile/internal/comfy"
	"gpu-turnstile/internal/config"
	"gpu-turnstile/internal/control"
	"gpu-turnstile/internal/game"
	"gpu-turnstile/internal/lock"
	"gpu-turnstile/internal/metrics"
	"gpu-turnstile/internal/ollama"
	"gpu-turnstile/internal/proxy"
	"gpu-turnstile/internal/service"
	"gpu-turnstile/internal/supervise"
	"gpu-turnstile/internal/update"
)

// version is injected at build time via -ldflags "-X main.version=...".
var version = "dev"

// exitCodeUpdate tells the service recovery configuration to restart the
// process: a signed update has been staged or a changed config reload was
// requested, and the GPU lock is idle.
const exitCodeUpdate = 3

// exitCodeStaged is returned by an elevated --force-update child when it
// staged a new binary, so the non-elevated parent can tell "updated" from
// "up to date" (it cannot see the child's console).
const exitCodeStaged = 4

// stdoutIsTerminal reports whether stdout is a console (char device), as
// opposed to a pipe or file — which is what Docker containers and services
// see.
func stdoutIsTerminal() bool {
	fi, err := os.Stdout.Stat()
	return err == nil && fi.Mode()&os.ModeCharDevice != 0
}

// parseFlags extracts -config <path> (or -config=<path>), the
// --install-service / --remove-service switches, --no-copy, -h/--help,
// -v/--version, --force-update and the hidden --elevated-child marker from
// args.
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild bool, rest []string) {
	rest = args[:0]
	for i := 0; i < len(args); i++ {
		switch {
		case args[i] == "-config" && i+1 < len(args):
			configPath = args[i+1]
			i++
		case strings.HasPrefix(args[i], "-config="):
			configPath = strings.TrimPrefix(args[i], "-config=")
		case args[i] == "--install-service" || args[i] == "-install-service" || args[i] == "-i":
			install = true
		case args[i] == "--remove-service" || args[i] == "-remove-service" || args[i] == "-r":
			remove = true
		case args[i] == "--no-copy" || args[i] == "-no-copy":
			noCopy = true
		case args[i] == "-h" || args[i] == "--help" || args[i] == "-help":
			help = true
		case args[i] == "-v" || args[i] == "--version" || args[i] == "-version":
			showVersion = true
		case args[i] == "--force-update" || args[i] == "-force-update":
			forceUpdate = true
		case args[i] == "--update-now" || args[i] == "-update-now":
			updateNow = true
		case args[i] == "--reload-env" || args[i] == "-reload-env":
			reloadEnv = true
		case args[i] == "--monitor" || args[i] == "-monitor" || args[i] == "-m":
			monitor = true
		case args[i] == "--elevated-child":
			elevatedChild = true
		default:
			rest = append(rest, args[i])
		}
	}
	return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, rest
}

// versionLine is printed at the top of every help and error screen.
func versionLine() string { return "gpu-turnstile " + version }

const usageText = `GPU arbitration proxy for Ollama + ComfyUI

Usage:
  gpu-turnstile -config <path>                                run the proxy
  gpu-turnstile -i | --install-service [--no-copy] [-config path]
                                                              install + start as a service
  gpu-turnstile -r | --remove-service                         stop + uninstall the service
  gpu-turnstile -v | --version                                print just the version
  gpu-turnstile --force-update                                check for a signed update now,
                                                              apply it and restart the service
                                                              (no admin needed when the service runs)
  gpu-turnstile --update-now                                  like --force-update, but only
                                                              through the running service
  gpu-turnstile --reload-env                                  make the service re-read and
                                                              validate its config file, then
                                                              restart onto it if it changed
  gpu-turnstile -m | --monitor                                live status view (downstreams,
                                                              GPU lock, queue); Ctrl+C quits
  gpu-turnstile -h | --help                                   this help

Options:
  -config <path>      config file (default: gpu-turnstile.env next to the exe)
  --install-service   copies the binary into the canonical location
                      (%ProgramFiles%\gpu-turnstile or /var/lib/gpu-turnstile)
                      unless --no-copy; on Windows a UAC prompt appears when
                      the shell is not elevated
  --no-copy           with --install-service: register the current location as-is

All runtime settings are environment variables or KEY=VALUE lines in the
config file (OLLAMA_URL, COMFY_URL, LOGLEVEL, ...); see README.md.
`

// printHelp prints the version header plus the full help text.
func printHelp() {
	fmt.Printf("%s — %s", versionLine(), usageText)
}

// fatalUsage prints the version header, an error message and the one-line
// usage summary, then exits with code 2.
func fatalUsage(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "%s\n\n", versionLine())
	fmt.Fprintf(os.Stderr, format+"\n\n", args...)
	fmt.Fprintln(os.Stderr, "usage: gpu-turnstile [-config path] [--install-service [--no-copy] | --remove-service]")
	fmt.Fprintln(os.Stderr, "       gpu-turnstile service install|remove [-config path]")
	os.Exit(2)
}

func main() {
	configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, args := parseFlags(os.Args[1:])
	if showVersion {
		fmt.Println(version)
		return
	}
	bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !reloadEnv && !monitor && !elevatedChild && len(args) == 0
	if help || (bare && stdoutIsTerminal()) {
		// Bare invocation in a terminal (e.g. double-clicked on Windows)
		// shows the help instead of starting a proxy window with no visible
		// explanation. Without a terminal — Docker containers, services,
		// pipes — a bare invocation starts the proxy as before.
		printHelp()
		return
	}
	if len(args) > 0 && args[0] == "service" {
		// Legacy subcommand form: gpu-turnstile service install|remove.
		if len(args) != 2 || (args[1] != "install" && args[1] != "remove") {
			fatalUsage("error: expected 'service install' or 'service remove'")
		}
		install = args[1] == "install"
		remove = !install
		args = nil
	}
	switch {
	case install && remove:
		fatalUsage("error: --install-service and --remove-service are mutually exclusive")
	case forceUpdate && (install || remove):
		fatalUsage("error: --force-update cannot be combined with --install-service/--remove-service")
	case updateNow && (install || remove || forceUpdate):
		fatalUsage("error: --update-now cannot be combined with other commands")
	case reloadEnv && (install || remove || forceUpdate || updateNow):
		fatalUsage("error: --reload-env cannot be combined with other commands")
	case monitor && (install || remove || forceUpdate || updateNow || reloadEnv):
		fatalUsage("error: --monitor cannot be combined with other commands")
	case install:
		os.Exit(serviceCommand(configPath, true, noCopy, elevatedChild))
	case remove:
		os.Exit(serviceCommand(configPath, false, noCopy, elevatedChild))
	case forceUpdate:
		os.Exit(forceUpdateCommand(configPath, elevatedChild))
	case updateNow:
		os.Exit(updateNowCommand())
	case reloadEnv:
		os.Exit(reloadEnvCommand())
	case monitor:
		os.Exit(monitorCommand())
	}
	if len(args) > 0 {
		fatalUsage("error: unknown arguments: %s", strings.Join(args, " "))
	}

	if exePath, err := os.Executable(); err == nil {
		update.CleanupOld(exePath)
	}

	cfg, err := loadMergedConfig(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: %v\n", versionLine(), err)
		os.Exit(1)
	}
	log, logOut, logCloser := newLogger(cfg)
	defer logCloser.Close()
	syncEnvFile(resolveConfigPath(configPath), cfg.LogFile, log)

	if service.IsService() {
		if err := service.Run(func(ctx context.Context) error { return run(ctx, cfg, log, logOut, true, configPath) }); err != nil {
			log.Error("service failed", "err", err)
			os.Exit(1)
		}
		return
	}
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	if err := run(ctx, cfg, log, logOut, false, configPath); err != nil {
		log.Error("listener failed", "err", err)
		os.Exit(1)
	}
}

// syncEnvFile upgrades an installer-written config file after an update:
// settings added since its CFG_VER are appended (commented out) and CFG_VER
// is bumped. Files not written by the installer (no CFG_VER), up-to-date
// files and dev builds are left untouched; a write failure is logged, not
// fatal.
func syncEnvFile(path, logFile string, log *slog.Logger) {
	data, err := os.ReadFile(path)
	if err != nil {
		return // no config file; nothing to upgrade
	}
	synced, changed := config.SyncSample(string(data), version, logFile)
	if !changed {
		return
	}
	if err := os.WriteFile(path, []byte(synced), 0o644); err != nil {
		log.Warn("could not append new settings to the config file", "path", path, "err", err)
		return
	}
	log.Warn("config file updated: new settings appended", "path", path, "version", version)
}

// defaultConfigPath returns gpu-turnstile.env next to the executable.
func defaultConfigPath() string {
	exe, err := os.Executable()
	if err != nil {
		return "gpu-turnstile.env"
	}
	return filepath.Join(filepath.Dir(exe), "gpu-turnstile.env")
}

// resolveConfigPath applies the precedence: -config flag, then
// GPU_TURNSTILE_CONFIG, then the default next to the executable.
func resolveConfigPath(flagValue string) string {
	if flagValue != "" {
		return flagValue
	}
	if v := os.Getenv("GPU_TURNSTILE_CONFIG"); v != "" {
		return v
	}
	return defaultConfigPath()
}

// loadMergedConfig reads the config file (if present) and overlays process
// environment variables on top. A missing file is fine; an unreadable or
// malformed file is fatal.
func loadMergedConfig(flagValue string) (config.Config, error) {
	path := resolveConfigPath(flagValue)
	values := map[string]string{}
	if f, err := os.Open(path); err == nil {
		defer f.Close()
		parsed, err := config.ParseEnvFile(f)
		if err != nil {
			return config.Config{}, fmt.Errorf("%s: %w", path, err)
		}
		values = parsed
	} else if !errors.Is(err, os.ErrNotExist) {
		return config.Config{}, fmt.Errorf("read config file: %w", err)
	}
	getenv := func(key string) string {
		if v := os.Getenv(key); v != "" {
			return v
		}
		return values[key]
	}
	return config.Load(getenv)
}

// newLogger builds the slog logger and returns the output writer (stderr or
// the opened LOG_FILE) plus a closer for it.
func newLogger(cfg config.Config) (*slog.Logger, io.Writer, io.Closer) {
	out := io.Writer(os.Stderr)
	closer := io.NopCloser(nil)
	if cfg.LogFile != "" {
		if f, err := os.OpenFile(cfg.LogFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644); err == nil {
			out = f
			closer = f
		} else {
			fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot open LOG_FILE %s: %v (logging to stderr)\n", cfg.LogFile, err)
		}
	}
	opts := &slog.HandlerOptions{Level: cfg.LogLevel}
	var handler slog.Handler = slog.NewTextHandler(out, opts)
	if cfg.LogJSON {
		handler = slog.NewJSONHandler(out, opts)
	}
	log := slog.New(handler)
	slog.SetDefault(log)
	return log, out, closer
}

// waitForEnter keeps an elevated child's console window open until the
// user has read the output.
func waitForEnter() {
	fmt.Print("\nPress Enter to close this window...")
	bufio.NewReader(os.Stdin).ReadString('\n')
}

// elevateAndMirror relaunches the current command elevated (UAC) and
// mirrors the child's exit code. verb is used in messages.
func elevateAndMirror(verb string) (int, bool) {
	args := append(append([]string{}, os.Args[1:]...), "--elevated-child")
	code, err := service.RelaunchElevated(args)
	if errors.Is(err, service.ErrUserCancelled) {
		fmt.Fprintln(os.Stderr, "gpu-turnstile: UAC prompt declined")
		return 1, true
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: could not elevate: %v\n", err)
		return 1, true
	}
	if code != 0 {
		fmt.Fprintf(os.Stderr, "gpu-turnstile %s failed in the elevated process (exit %d)\n", verb, code)
		return code, true
	}
	return 0, true
}

// isPermission reports whether err is a permission problem (Windows
// ERROR_ACCESS_DENIED, POSIX EACCES/EPERM, possibly wrapped).
func isPermission(err error) bool {
	return errors.Is(err, fs.ErrPermission) || strings.Contains(strings.ToLower(err.Error()), "access is denied")
}

// forceUpdateCommand checks for a signed update immediately, stages it if
// newer, and restarts the service when it is running so the new binary
// takes effect. Staging into a system directory and restarting a service
// need admin rights; instead of prompting unconditionally, permission
// failures trigger the UAC relaunch so a dev copy in a user-writable
// directory updates without a prompt.
func forceUpdateCommand(configPath string, elevatedChild bool) int {
	if elevatedChild {
		defer waitForEnter()
	}
	cfg, err := loadMergedConfig(configPath)
	if errors.Is(err, config.ErrNoConsumer) {
		err = nil // update settings do not depend on a consumer URL
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: %v\n", versionLine(), err)
		return 1
	}
	exePath, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot locate executable: %v\n", err)
		return 1
	}
	// One-shot CLI: the updater logs to stderr, never to the service's
	// LOG_FILE — that file is ACL'd to the service account, and a CLI run
	// has nothing worth persisting there.
	cfg.LogFile = ""
	log, _, logCloser := newLogger(cfg)
	defer logCloser.Close()
	if cfg.AppVersion == "dev" {
		fmt.Printf("%s: APP_VER=dev, updates disabled\n", versionLine())
		return 0
	}
	// A running service can do the privileged work itself (its account owns
	// the install dir and it knows when the GPU is idle): ask it over the
	// local control channel first, no admin rights needed. Fails fast when
	// no service is listening, in which case we do the direct check below.
	if reply, err := control.Ask(control.CmdUpdateNow); err == nil {
		return printControlReply(reply)
	}
	u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}

	// Single-shot: one attempt, fail fast when the server is unreachable
	// instead of hanging in a TCP connect for minutes.
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	staged, to, err := u.Check(ctx, exePath)
	if err != nil && isPermission(err) && !service.Elevated() {
		code, _ := elevateAndMirror("--force-update")
		reportElevatedUpdate(code, to)
		if code == 0 || code == exitCodeStaged {
			return 0
		}
		return code
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: update check failed: %v\n", versionLine(), err)
		return 1
	}
	if !staged {
		fmt.Printf("%s is up to date\n", versionLine())
		return 0
	}
	fmt.Printf("gpu-turnstile: updated from %s to %s\n", version, to)
	restarted, err := service.RestartIfRunning()
	if err != nil && isPermission(err) && !service.Elevated() {
		code, _ := elevateAndMirror("--force-update")
		reportElevatedUpdate(code, to)
		if code == 0 || code == exitCodeStaged {
			return 0
		}
		return code
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: update staged but service restart failed: %v\n", err)
		return 1
	}
	if restarted {
		fmt.Println("service restarted on " + to)
	} else {
		fmt.Println("no running service; " + to + " applies on the next start")
	}
	if elevatedChild {
		// Tell the non-elevated parent (which cannot see this console)
		// whether anything was staged, so its mirror message is honest.
		return exitCodeStaged
	}
	return 0
}

// printControlReply prints the service's answer to a control-channel
// request: "OK ..." on stdout (exit 0), "ERR ..." on stderr (exit 1).
func printControlReply(reply string) int {
	if msg, ok := strings.CutPrefix(reply, "ERR "); ok {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: %s\n", msg)
		return 1
	}
	fmt.Println("gpu-turnstile: " + strings.TrimPrefix(reply, "OK "))
	return 0
}

// updateNowCommand only goes through the running service's control channel
// (no direct check, no elevation): the unprivileged update trigger.
func updateNowCommand() int {
	reply, err := control.Ask(control.CmdUpdateNow)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: no running service to ask — use --force-update for a direct check\n", versionLine())
		return 1
	}
	return printControlReply(reply)
}

// reloadEnvCommand asks the running service, over the control channel, to
// re-read and validate its config file. The client sends no path — the
// service only ever re-reads its own configured file.
func reloadEnvCommand() int {
	reply, err := control.Ask(control.CmdReloadEnv)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: no running service to ask\n", versionLine())
		return 1
	}
	return printControlReply(reply)
}

// reportElevatedUpdate prints the parent's summary of an elevated
// --force-update child: exitCodeStaged means the child staged a new binary,
// 0 means it found nothing to do. to is the tag the parent's own check
// resolved before it hit the permission wall ("" when it never got that
// far).
func reportElevatedUpdate(code int, to string) {
	switch {
	case code == exitCodeStaged && to != "":
		fmt.Printf("gpu-turnstile: updated from %s to %s (elevated)\n", version, to)
	case code == exitCodeStaged:
		fmt.Println("update applied (elevated)")
	case code == 0:
		fmt.Printf("%s is up to date\n", versionLine())
	}
	// Non-zero, non-staged codes: elevateAndMirror already printed the failure.
}

// serviceCommand installs (copyBin = register the canonical-layout copy)
// or removes the service and reports the result. On Windows, when the
// shell is not elevated, the command relaunches itself through a UAC
// prompt and mirrors the elevated child's exit code. An elevated child
// waits for a keypress so its console window does not flash closed before
// the output can be read.
func serviceCommand(configPath string, install, noCopy, elevatedChild bool) int {
	verb, doneVerb := "remove", "removed"
	if install {
		verb, doneVerb = "install", "installed"
	}
	if elevatedChild {
		defer waitForEnter()
	}
	if !service.Elevated() {
		code, done := elevateAndMirror(verb)
		if done && code != 0 {
			return code
		}
		if done {
			fmt.Printf("service %s: %s (elevated)\n", service.Name, doneVerb)
			return 0
		}
	}
	var err error
	if install {
		path := resolveConfigPath(configPath)
		if abs, absErr := filepath.Abs(path); absErr == nil {
			path = abs
		}
		err = service.Install(path, !noCopy, version)
	} else {
		err = service.Remove()
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile service %s: %v\n", verb, err)
		return 1
	}
	fmt.Printf("service %s: %s\n", service.Name, doneVerb)
	return 0
}

// orDisabled renders an empty URL as "disabled" for the startup dump.
func orDisabled(url string) string {
	if url == "" {
		return "disabled"
	}
	return url
}

// managedComfyCommand resolves how ComfyUI is launched when it is managed:
// COMFY_CMD verbatim, or the standard venv layout under COMFY_DIR. Empty
// when neither is set (unmanaged).
func managedComfyCommand(cfg config.Config) string {
	if cfg.ComfyCmd != "" {
		return cfg.ComfyCmd
	}
	if cfg.ComfyDir != "" {
		return supervise.DefaultComfyCommand(runtime.GOOS, cfg.ComfyDir, cfg.ComfyURL)
	}
	return ""
}

func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Writer, isService bool, configPath string) error {
	// ComfyUI can run as a managed child — COMFY_CMD verbatim, or the
	// standard venv layout derived from COMFY_DIR alone: started on demand
	// by the proxy, stopped after COMFY_IDLE_TIMEOUT idle (and on shutdown)
	// so its VRAM is freed.
	comfyCmdLine := managedComfyCommand(cfg)

	// The startup line carries the version and every setting and is emitted
	// at WARN so it is visible even with the default (quiet) log level.
	log.Log(ctx, slog.LevelWarn, "starting gpu-turnstile",
		"version", version,
		"listen_ollama", cfg.ListenOllama,
		"listen_comfy", cfg.ListenComfy,
		"listen_ui", orDisabled(cfg.ListenUI),
		"ollama_url", orDisabled(cfg.OllamaURL),
		"comfy_url", orDisabled(cfg.ComfyURL),
		"unload_timeout", cfg.UnloadTimeout,
		"job_timeout", cfg.JobTimeout,
		"llm_wait_timeout", cfg.LLMWaitTimeout,
		"llm_busy_mode", cfg.LLMBusyMode,
		"llm_busy_status", cfg.LLMBusyStatus,
		"busy_retry_after", cfg.BusyRetryAfter,
		"unload_poll_interval", cfg.UnloadPollInterval,
		"history_poll_interval", cfg.HistoryPollInterval,
		"probe_timeout", cfg.ProbeTimeout,
		"health_interval", cfg.HealthInterval,
		"free_timeout", cfg.FreeTimeout,
		"warm_timeout", cfg.WarmTimeout,
		"shutdown_timeout", cfg.ShutdownTimeout,
		"backoff_initial", cfg.BackoffInitial,
		"backoff_max", cfg.BackoffMax,
		"prompt_capture_limit", cfg.PromptCaptureLimit,
		"warm_model", cfg.WarmModel,
		"comfy_cmd", orDisabled(comfyCmdLine),
		"comfy_dir", cfg.ComfyDir,
		"comfy_idle_timeout", cfg.ComfyIdleTimeout,
		"comfy_start_timeout", cfg.ComfyStartTimeout,
		"game_procs", cfg.GameProcs,
		"gpu_foreign_vram_mb", cfg.GPUForeignVRAMMB,
		"gpu_foreign_util_pct", cfg.GPUForeignUtilPct,
		"gpu_ignore_procs", cfg.GPUIgnoreProcs,
		"game_poll_interval", cfg.GamePollInterval,
		"auto_update", cfg.AutoUpdate,
		"update_interval", cfg.UpdateInterval,
		"update_repo", cfg.UpdateRepo,
		"update_asset", cfg.UpdateAsset,
		"log_level", cfg.LogLevel,
		"log_format", map[bool]string{true: "json", false: "text"}[cfg.LogJSON],
		"log_file", cfg.LogFile,
	)

	lk := lock.New(log)
	health := newHealthTracker()
	started := time.Now()
	// Each consumer is enabled by setting its URL; a disabled consumer gets
	// no client, no listener and no probe.
	var ollamaClient *ollama.Client
	var err error
	if cfg.OllamaURL != "" {
		if ollamaClient, err = ollama.New(cfg.OllamaURL, log); err != nil {
			return err
		}
	}
	var comfyClient *comfy.Client
	if cfg.ComfyURL != "" {
		if comfyClient, err = comfy.New(cfg.ComfyURL, log); err != nil {
			return err
		}
	}

	var comfySup *supervise.Process
	if comfyCmdLine != "" {
		if cfg.ComfyCmd == "" {
			// Derived from COMFY_DIR: flag a wrong-looking layout early,
			// while the operator is still watching the startup log. Inside
			// a profile the service account may not enter, os.Stat fails
			// with EACCES — that reads as "not found" but means "grant
			// access", so say so.
			python, script := supervise.ComfyLayout(runtime.GOOS, cfg.ComfyDir)
			for _, p := range []string{python, script} {
				if _, err := os.Stat(p); err != nil {
					if isPermission(err) {
						log.Warn("COMFY_DIR: not accessible to the service account; grant access or re-run --install-service", "path", p)
					} else {
						log.Warn("COMFY_DIR: file not found; ComfyUI requests will fail until it exists", "path", p)
					}
				}
			}
		}
		var err error
		comfySup, err = supervise.New("comfy", comfyCmdLine, cfg.ComfyDir, comfyClient.Probe, cfg.ComfyStartTimeout, log)
		if err != nil {
			return err
		}
		defer comfySup.Stop()
		gpuIdle := func() bool {
			state, _, pending := lk.Snapshot()
			return state == lock.StateIdle && !pending
		}
		go comfySup.WatchIdle(ctx, cfg.ComfyIdleTimeout, gpuIdle)
	}

	srv, err := proxy.New(proxy.Config{
		OllamaURL:           cfg.OllamaURL,
		ComfyURL:            cfg.ComfyURL,
		Lock:                lk,
		Ollama:              ollamaClient,
		Comfy:               comfyClient,
		ComfySup:            comfySup,
		Metrics:             metrics.New(),
		Log:                 log,
		LogColor:            !cfg.LogJSON && cfg.LogFile == "" && os.Getenv("NO_COLOR") == "",
		LogWriter:           logOut,
		LLMWaitTimeout:      cfg.LLMWaitTimeout,
		UnloadTimeout:       cfg.UnloadTimeout,
		JobTimeout:          cfg.JobTimeout,
		UnloadPollInterval:  cfg.UnloadPollInterval,
		HistoryPollInterval: cfg.HistoryPollInterval,
		FreeTimeout:         cfg.FreeTimeout,
		WarmTimeout:         cfg.WarmTimeout,
		LLMBusyMode:         cfg.LLMBusyMode,
		LLMBusyStatus:       cfg.LLMBusyStatus,
		BusyRetryAfter:      cfg.BusyRetryAfter,
		BackoffInitial:      cfg.BackoffInitial,
		BackoffMax:          cfg.BackoffMax,
		PromptCaptureLimit:  cfg.PromptCaptureLimit,
		WarmModel:           cfg.WarmModel,
	})
	if err != nil {
		return err
	}

	// Probe the enabled upstreams once; failure is logged, not fatal. A
	// managed ComfyUI is intentionally down at startup — the first request
	// starts it — so neither the startup probe nor the health check treats
	// that as an outage.
	probes := map[string]func(context.Context) error{}
	if ollamaClient != nil {
		probes["ollama"] = ollamaClient.Probe
	}
	if comfyClient != nil {
		if comfySup == nil {
			probes["comfy"] = comfyClient.Probe
		} else {
			// Managed upstream: an idle-stopped or still-starting server is
			// not an outage, so it is skipped until it has answered once
			// (Ready resets on every spawn/stop). After that, a failed
			// probe while the process lives is a real "DOWN".
			probes["comfy"] = func(ctx context.Context) error {
				if !comfySup.Ready() {
					if comfySup.Running() {
						if err := comfyClient.Probe(ctx); err == nil {
							comfySup.MarkReady()
						}
					}
					return errManagedDown
				}
				return comfyClient.Probe(ctx)
			}
		}
	}
	probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
	for name, probe := range probes {
		err := probe(probeCtx)
		health.set(name, err == nil)
		if err != nil && !errors.Is(err, errManagedDown) {
			log.Warn(name+" probe failed", "err", err)
		}
	}
	probeCancel()
	if cfg.HealthInterval > 0 {
		go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes, health)
	}

	// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS,
	// GPU_FOREIGN_VRAM_MB and/or GPU_FOREIGN_UTIL_PCT — hold the lock
	// externally while they run. gw collects the VRAM reading and the last
	// check result for the status channel.
	gw := &gpuWatch{enabled: len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 || cfg.GPUForeignUtilPct > 0}
	if gw.enabled {
		det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUForeignUtilPct, cfg.GPUIgnoreProcs, log)
		go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup, gw)
	}

	// Bind the listeners up front so a port conflict fails fast and the
	// readiness notification below really means "accepting connections".
	var servers []*http.Server
	var listeners []net.Listener
	bind := func(addr string, handler http.Handler, consumer string) error {
		ln, err := net.Listen("tcp", addr)
		if err != nil {
			return fmt.Errorf("listen %s on %s: %w", consumer, addr, err)
		}
		servers = append(servers, &http.Server{Addr: addr, Handler: handler})
		listeners = append(listeners, ln)
		log.Warn("listening", "consumer", consumer, "addr", addr)
		return nil
	}
	if ollamaClient != nil {
		if err := bind(cfg.ListenOllama, srv.OllamaHandler(), "ollama"); err != nil {
			return err
		}
	}
	if comfyClient != nil {
		if err := bind(cfg.ListenComfy, srv.ComfyHandler(), "comfy"); err != nil {
			return err
		}
	}

	errCh := make(chan error, len(servers))
	for i := range servers {
		go func(s *http.Server, ln net.Listener) { errCh <- s.Serve(ln) }(servers[i], listeners[i])
	}

	// Tell systemd we are up and start the watchdog pings; both are no-ops
	// when not running under a notify/watchdog unit.
	service.NotifyReady()
	service.StartWatchdog(ctx)

	var u *update.Updater
	var exePath string
	// restartWhenIdle exits with exitCodeUpdate once the GPU lock is idle;
	// the service recovery configuration brings the process back. Shared by
	// staged updates and config reloads; the once guard makes repeated
	// triggers idempotent.
	var restartOnce sync.Once
	restartWhenIdle := func(reason string) {
		log.Warn("restarting once the GPU is idle", "reason", reason)
		restartOnce.Do(func() {
			go func() {
				if waitForIdle(ctx, lk, 24*time.Hour) {
					log.Warn("restarting now", "reason", reason)
					os.Exit(exitCodeUpdate)
				}
			}()
		})
	}
	applyStaged := func(to string) {}
	if cfg.AutoUpdate {
		if p, err := os.Executable(); err != nil {
			log.Warn("auto-update disabled: cannot locate executable", "err", err)
		} else {
			exePath = p
			u = &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
			applyStaged = func(to string) {
				if !isService {
					log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply", "version", to)
					return
				}
				restartWhenIdle("update to " + to)
			}
			go updateLoop(ctx, cfg.UpdateInterval, log, u, exePath, applyStaged)
		}
	}

	// The command handler backs both the local control channel (served
	// whenever running as a service) and the optional web UI (LISTEN_UI),
	// independent of AUTO_UPDATE.
	handler := controlHandler(ctx, u, exePath, applyStaged,
		statusProvider(cfg, lk, comfySup, health, started, gw, ollamaClient),
		reloadHandler(cfg, configPath, restartWhenIdle))
	if isService {
		serveControl(ctx, log, handler)
	}
	if cfg.ListenUI != "" {
		ln, err := net.Listen("tcp", cfg.ListenUI)
		if err != nil {
			return fmt.Errorf("listen ui on %s: %w", cfg.ListenUI, err)
		}
		uiSrv := &http.Server{Addr: cfg.ListenUI, Handler: uiHandler(handler)}
		servers = append(servers, uiSrv)
		log.Warn("listening", "consumer", "ui", "addr", cfg.ListenUI)
		go func() { errCh <- uiSrv.Serve(ln) }()
	}

	select {
	case err := <-errCh:
		if err != nil && !errors.Is(err, http.ErrServerClosed) {
			return err
		}
	case <-ctx.Done():
		log.Info("shutting down")
	}

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
	defer shutdownCancel()
	for _, s := range servers {
		s.Shutdown(shutdownCtx)
	}
	return nil
}

// errManagedDown marks a managed upstream that is intentionally stopped
// (idle); health checks skip it instead of logging an outage.
var errManagedDown = errors.New("managed upstream intentionally stopped")

// gpuWatch records the latest VRAM reading and game-detector finding from
// the game detector's poll loop, for the status channel. Enabled is false
// when game detection is not configured (no polling happens then); Known is
// false until the first successful nvidia-smi reading.
type gpuWatch struct {
	mu      sync.Mutex
	enabled bool
	usedMB  int
	total   int
	tempC   int
	fanPct  int
	known   bool
	foreign string // last Check result: external holders, "" when none
	at      time.Time
}

func (g *gpuWatch) setVRAM(st game.GPUStats) {
	g.mu.Lock()
	g.usedMB, g.total = st.UsedMB, st.TotalMB
	g.tempC, g.fanPct, g.known = st.TempC, st.FanPct, true
	g.mu.Unlock()
}

func (g *gpuWatch) setCheck(foreign string) {
	g.mu.Lock()
	g.foreign, g.at = foreign, time.Now()
	g.mu.Unlock()
}

func (g *gpuWatch) get() (st game.GPUStats, known bool, foreign string, ageS int64) {
	g.mu.Lock()
	defer g.mu.Unlock()
	ageS = -1
	if !g.at.IsZero() {
		ageS = int64(time.Since(g.at).Seconds())
	}
	return game.GPUStats{UsedMB: g.usedMB, TotalMB: g.total, TempC: g.tempC, FanPct: g.fanPct},
		g.known, g.foreign, ageS
}

// gameLoop polls for foreign GPU holders (a game, another ML job). While one
// is detected it holds the lock externally so new LLM and image requests
// wait (or are rejected per LLM_BUSY_MODE), and — once in-flight work has
// drained — frees VRAM for it: the managed ComfyUI is stopped and Ollama's
// resident models are unloaded.
func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *game.Detector, lk *lock.Lock, ollamaClient *ollama.Client, comfySup *supervise.Process, gw *gpuWatch) {
	ticker := time.NewTicker(cfg.GamePollInterval)
	defer ticker.Stop()
	held, freed := false, false
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		holders, err := det.Check(ctx)
		if err != nil && ctx.Err() == nil {
			log.Warn("game detection failed", "err", err)
		}
		gw.setCheck(summarizeHolders(holders))
		if st, verr := game.QueryGPUStats(ctx); verr == nil {
			gw.setVRAM(st)
		}
		switch {
		case len(holders) > 0 && !held:
			held = true
			lk.SetExternal(summarizeHolders(holders))
			log.Warn("GPU held by an external process; new LLM/image requests wait",
				"holders", summarizeHolders(holders))
		case len(holders) == 0 && held:
			held, freed = false, false
			lk.ClearExternal()
			log.Warn("external process released the GPU; resuming")
		}
		if held && !freed {
			if state, _, _ := lk.Snapshot(); state != lock.StateLLM && state != lock.StateImage {
				freed = true
				freeVRAM(ctx, cfg.UnloadTimeout, log, ollamaClient, comfySup)
			}
		}
	}
}

// summarizeHolders joins holder descriptions for logs and busy responses,
// capping the list so a process name matching dozens of PIDs (system
// services) does not flood the log.
func summarizeHolders(holders []string) string {
	const max = 5
	if len(holders) > max {
		return strings.Join(holders[:max], "; ") + fmt.Sprintf("; +%d more", len(holders)-max)
	}
	return strings.Join(holders, "; ")
}

// freeVRAM stops the managed ComfyUI (never an external server on its port)
// and unloads Ollama's resident models so the foreign process gets the GPU
// memory.
func freeVRAM(ctx context.Context, unloadTimeout time.Duration, log *slog.Logger, ollamaClient *ollama.Client, comfySup *supervise.Process) {
	if comfySup != nil && comfySup.Running() {
		log.Warn("stopping the managed ComfyUI to free VRAM")
		comfySup.Stop()
	}
	if ollamaClient == nil {
		return
	}
	uctx, cancel := context.WithTimeout(ctx, unloadTimeout)
	defer cancel()
	if models, err := ollamaClient.LoadedModels(uctx); err != nil || len(models) == 0 {
		return // nothing resident (or ollama unreachable); nothing to free
	}
	elapsed, err := ollamaClient.UnloadAll(uctx)
	if err != nil {
		log.Warn("ollama unload incomplete; continuing", "err", err)
		return
	}
	log.Warn("ollama models unloaded to free VRAM", "seconds", elapsed.Seconds())
}

// healthLoop probes the enabled upstreams every interval and logs status
// transitions — "is DOWN" when a previously healthy upstream stops
// answering, "recovered" when it comes back. The first round only
// establishes the baseline; the startup probe already reported that state.
// Every result goes into the tracker for the status channel.
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error, tracker *healthTracker) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	up := map[string]bool{}
	managed := map[string]bool{}
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		for name, probe := range probes {
			pctx, cancel := context.WithTimeout(ctx, probeTimeout)
			err := probe(pctx)
			cancel()
			if errors.Is(err, errManagedDown) {
				managed[name] = true // intentionally stopped; not an outage
				tracker.set(name, false)
				continue
			}
			now := err == nil
			tracker.set(name, now)
			if managed[name] {
				// First real probe after an idle stop only re-baselines —
				// an on-demand start is not a "recovery".
				managed[name] = false
				up[name] = now
				continue
			}
			was, seen := up[name]
			if seen && now != was {
				if now {
					log.Warn(name + " upstream recovered")
				} else {
					log.Warn(name+" upstream is DOWN", "err", err)
				}
			}
			up[name] = now
		}
	}
}

// updateLoop checks for signed updates on startup and every interval;
// applyStaged decides what a staged update means (restart when idle as a
// service, log only interactively).
func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string)) {
	for {
		staged, to, err := u.Check(ctx, exePath)
		if err != nil && ctx.Err() == nil {
			log.Warn("auto-update check failed", "err", err)
		}
		if staged {
			applyStaged(to)
			return
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(interval):
		}
	}
}

// controlHandler builds the command handler shared by the control channel
// (named pipe / unix socket) and the web UI's action endpoints: status for
// --monitor / the UI, an update check (--force-update/--update-now) and a
// config reload (--reload-env), all safe for unprivileged local users. The
// update payload is signature-verified regardless of who asks; update
// triggers are rate-limited to one per minute and reloads to one per two
// seconds so neither can be used to spam restarts or disk churn. u is nil
// when AUTO_UPDATE=false.
func controlHandler(ctx context.Context, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) control.Handler {
	var mu sync.Mutex
	var lastTrigger, lastReload time.Time
	return func(cmd string) string {
		switch cmd {
		case control.CmdStatus:
			return "OK " + status()
		case control.CmdReloadEnv:
			// reload re-reads the config file from disk; a short limiter
			// keeps a local flood from turning into disk churn.
			mu.Lock()
			if wait := 2*time.Second - time.Since(lastReload); wait > 0 {
				mu.Unlock()
				return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1)
			}
			lastReload = time.Now()
			mu.Unlock()
			return reload()
		case control.CmdUpdateNow:
		default:
			return "ERR unknown command: " + cmd
		}
		if u == nil {
			return "ERR auto-update is disabled on this instance"
		}
		mu.Lock()
		if wait := time.Minute - time.Since(lastTrigger); wait > 0 {
			mu.Unlock()
			return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1)
		}
		lastTrigger = time.Now()
		mu.Unlock()
		cctx, cancel := context.WithTimeout(ctx, 30*time.Second)
		defer cancel()
		staged, to, err := u.Check(cctx, exePath)
		if err != nil {
			return "ERR update check failed: " + err.Error()
		}
		if !staged {
			return "OK " + version + " is up to date"
		}
		applyStaged(to)
		return "OK updated from " + version + " to " + to + "; the service restarts once the GPU is idle"
	}
}

// serveControl opens the local control channel (named pipe on Windows,
// unix socket on Linux) for unprivileged local users (--monitor,
// --force-update/--update-now, --reload-env).
func serveControl(ctx context.Context, log *slog.Logger, h control.Handler) {
	if err := control.Serve(ctx, h, log); err != nil {
		log.Warn("control channel disabled", "err", err)
	}
}

// healthTracker records the latest probe result per upstream for the
// status channel.
type healthTracker struct {
	mu sync.Mutex
	up map[string]bool
}

func newHealthTracker() *healthTracker {
	return &healthTracker{up: map[string]bool{}}
}

func (h *healthTracker) set(name string, up bool) {
	h.mu.Lock()
	h.up[name] = up
	h.mu.Unlock()
}

func (h *healthTracker) get(name string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.up[name]
}

// statusDownstream/statusLock/statusSnapshot are the JSON the control
// channel serves on CmdStatus; the monitor mode renders them.
type statusDownstream struct {
	Name    string `json:"name"`
	URL     string `json:"url"`
	Up      bool   `json:"up"`
	Managed string `json:"managed,omitempty"`
	// Models lists Ollama's loaded models with their VRAM footprint. Null
	// when unknown (query failed / not applicable); [] means none loaded —
	// deliberately no omitempty so the two stay distinguishable.
	Models []statusModel `json:"models"`
}

// statusModel is one loaded Ollama model.
type statusModel struct {
	Name   string `json:"name"`
	VRAMMB int64  `json:"vram_mb"` // 0 = resident in RAM, not VRAM
}

type statusLock struct {
	State       string `json:"state"`
	Detail      string `json:"detail,omitempty"`
	LLMInflight int    `json:"llm_inflight"`
	LLMWaiting  int    `json:"llm_waiting"`
	ImageQueue  int    `json:"image_queue"`
	External    string `json:"external,omitempty"`
	SinceS      int64  `json:"since_s"`
}

type statusSnapshot struct {
	Version     string             `json:"version"`
	UptimeS     int64              `json:"uptime_s"`
	Downstreams []statusDownstream `json:"downstreams"`
	Lock        statusLock         `json:"lock"`
	// GPU carries the latest VRAM reading and detector finding; Enabled is
	// false when game detection (and with it nvidia-smi polling) is not
	// configured.
	GPU statusGPU `json:"gpu"`
	// MonitorNote is set client-side (never over the wire) when the
	// monitor's own binary differs from the service's version.
	MonitorNote string `json:"-"`
}

type statusGPU struct {
	// Enabled reports whether game detection is configured (and with it
	// VRAM polling); when false the other fields carry no information.
	Enabled bool `json:"enabled"`
	UsedMB  int  `json:"used_mb"`
	TotalMB int  `json:"total_mb"`
	// TempC/FanPct are -1 when unknown (never sampled or nvidia-smi
	// reported N/A).
	TempC  int  `json:"temp_c"`
	FanPct int  `json:"fan_pct"`
	Known  bool `json:"known"`
	// Foreign is the last detector finding (external GPU holders), empty
	// when the last check found none.
	Foreign string `json:"foreign,omitempty"`
	// AgeS is how long ago the last check ran; -1 before the first check.
	AgeS int64 `json:"age_s"`
}

// reloadHandler re-reads and validates the service's config file for
// CmdReloadEnv. An invalid config is reported and the service keeps running
// untouched; a valid, changed config triggers a GPU-idle-gated restart onto
// it (same mechanism as staged updates); unchanged is a no-op.
func reloadHandler(current config.Config, configPath string, restartWhenIdle func(reason string)) func() string {
	return func() string {
		ncfg, err := loadMergedConfig(configPath)
		if err != nil {
			return "ERR config invalid: " + err.Error()
		}
		changed := diffConfig(current, ncfg)
		if len(changed) == 0 {
			return "OK config unchanged"
		}
		restartWhenIdle("config reload (" + strings.Join(changed, ", ") + ")")
		return "OK config valid; restarting once the GPU is idle (changed: " + strings.Join(changed, ", ") + ")"
	}
}

// diffConfig lists the env names of settings whose values differ between
// two configs (from each field's env tag, so users recognize them).
func diffConfig(a, b config.Config) []string {
	va, vb := reflect.ValueOf(a), reflect.ValueOf(b)
	t := va.Type()
	var out []string
	for i := 0; i < t.NumField(); i++ {
		if !reflect.DeepEqual(va.Field(i).Interface(), vb.Field(i).Interface()) {
			name := t.Field(i).Tag.Get("env")
			if name == "" {
				name = t.Field(i).Name
			}
			out = append(out, name)
		}
	}
	return out
}

// statusProvider assembles the one-line JSON snapshot for CmdStatus. The
// loaded-model query to Ollama gets a short timeout so a wedged upstream
// cannot stall the status channel for long.
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time, gw *gpuWatch, ollamaClient *ollama.Client) func() string {
	return func() string {
		snap := statusSnapshot{
			Version: version,
			UptimeS: int64(time.Since(started).Seconds()),
		}
		st, known, foreign, ageS := gw.get()
		snap.GPU = statusGPU{
			Enabled: gw.enabled,
			UsedMB:  st.UsedMB, TotalMB: st.TotalMB, Known: known,
			TempC: -1, FanPct: -1,
			Foreign: foreign, AgeS: ageS,
		}
		if known {
			snap.GPU.TempC, snap.GPU.FanPct = st.TempC, st.FanPct
		}
		if cfg.OllamaURL != "" {
			d := statusDownstream{
				Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
			}
			if ollamaClient != nil {
				mctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
				models, err := ollamaClient.LoadedModelDetails(mctx)
				cancel()
				if err == nil {
					d.Models = make([]statusModel, 0, len(models))
					for _, m := range models {
						d.Models = append(d.Models, statusModel{Name: m.Name, VRAMMB: m.SizeVRAM / (1024 * 1024)})
					}
				}
			}
			snap.Downstreams = append(snap.Downstreams, d)
		}
		if cfg.ComfyURL != "" {
			d := statusDownstream{Name: "comfy", URL: cfg.ComfyURL, Up: health.get("comfy")}
			if comfySup != nil {
				d.Managed = comfySup.Status()
			}
			snap.Downstreams = append(snap.Downstreams, d)
		}
		lst := lk.Status()
		snap.Lock = statusLock{
			State:       string(lst.State),
			Detail:      lst.Detail,
			LLMInflight: lst.LLMInflight,
			LLMWaiting:  lst.LLMWaiting,
			ImageQueue:  lst.ImageQueue,
			External:    lst.External,
			SinceS:      int64(time.Since(lst.Since).Seconds()),
		}
		b, err := json.Marshal(snap)
		if err != nil {
			return `{"version":"` + version + `"}`
		}
		return string(b)
	}
}

// waitForIdle polls the lock until no LLM or image work is active or
// pending, max at most. Returns false on timeout or cancellation.
func waitForIdle(ctx context.Context, lk *lock.Lock, max time.Duration) bool {
	deadline := time.Now().Add(max)
	for {
		if state, _, pending := lk.Snapshot(); state == lock.StateIdle && !pending {
			return true
		}
		if time.Now().After(deadline) {
			return false
		}
		select {
		case <-ctx.Done():
			return false
		case <-time.After(5 * time.Second):
		}
	}
}