gpulock: GPU arbitration proxy for Ollama + ComfyUI

Implements SPEC.md: two listeners, one writer-preferring two-mode lock,
Ollama unload before image jobs, ComfyUI history polling + VRAM free,
optional model warm-up, healthz/metrics endpoints, streaming-safe
reverse proxies, Dockerfile and Gitea Actions CI.

.gitea/workflows/ci.yml

@@ -0,0 +1,62 @@
+name: ci
+
+on:
+  push:
+  workflow_dispatch:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
+        with:
+          go-version: "1.23"
+      - run: go vet ./...
+      - run: go test -race ./...
+      - name: golangci-lint
+        run: |
+          if command -v golangci-lint >/dev/null 2>&1; then
+            golangci-lint run
+          else
+            echo "golangci-lint not available in runner image, skipping"
+          fi
+
+  docker:
+    needs: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Compute image tags and version
+        id: meta
+        run: |
+          SHORT=$(echo "${{ gitea.sha }}" | cut -c1-7)
+          REPO=git.rambossek.at/${{ gitea.repository }}
+          TAGS="$REPO:sha-$SHORT"
+          VERSION="sha-$SHORT"
+          if [ "${{ gitea.ref_type }}" = "branch" ] && [ "${{ gitea.ref_name }}" = "main" ]; then
+            TAGS="$TAGS,$REPO:latest"
+          fi
+          if [ "${{ gitea.ref_type }}" = "tag" ]; then
+            TAGS="$TAGS,$REPO:${{ gitea.ref_name }}"
+            VERSION="${{ gitea.ref_name }}"
+          fi
+          echo "tags=$TAGS" >> "$GITHUB_OUTPUT"
+          echo "version=$VERSION" >> "$GITHUB_OUTPUT"
+
+      - uses: docker/setup-buildx-action@v3
+
+      - uses: docker/login-action@v3
+        with:
+          registry: git.rambossek.at
+          username: ${{ gitea.actor }}
+          password: ${{ secrets.GITEA_TOKEN }}
+
+      - uses: docker/build-push-action@v6
+        with:
+          context: .
+          push: true
+          build-args: |
+            VERSION=${{ steps.meta.outputs.version }}
+          tags: ${{ steps.meta.outputs.tags }}

.gitignore

@@ -0,0 +1,2 @@
+/gpulock
+/gpulock.exe

Dockerfile

@@ -0,0 +1,13 @@
+FROM golang:1.23 AS build
+WORKDIR /src
+COPY go.mod ./
+COPY cmd ./cmd
+COPY internal ./internal
+ARG VERSION=dev
+RUN CGO_ENABLED=0 go build -ldflags="-s -w -X main.version=${VERSION}" -o /gpulock ./cmd/gpulock
+
+FROM gcr.io/distroless/static:nonroot
+COPY --from=build /gpulock /gpulock
+EXPOSE 8189 11435
+USER nonroot
+ENTRYPOINT ["/gpulock"]

README.md

@@ -0,0 +1,87 @@
+# gpulock
+
+GPU arbitration proxy for Ollama + ComfyUI. One consumer GPU is shared by an
+LLM server (Ollama) and an image generator (ComfyUI); gpulock sits in front of
+both and guarantees the GPU is always in exactly one of three states: `idle`,
+`llm` (N ≥ 1 Ollama requests in flight), or `image` (exactly one ComfyUI job,
+Ollama models unloaded). See [SPEC.md](SPEC.md) for the full design.
+
+```
+LiteLLM / Open WebUI ──► :11435 ─┐                     ┌─► Ollama  :11434
+                                 ├── gpulock (1 lock) ──┤
+Open WebUI / n8n ────► :8189  ───┘                     └─► ComfyUI :8188
+```
+
+- LLM endpoints (`/api/generate`, `/api/chat`, `/api/embed`, `/v1/*`) take the
+  LLM lock: concurrent requests allowed, but blocked while an image job is
+  active or waiting (image priority).
+- `POST /prompt` on the ComfyUI listener takes the image lock: new LLM
+  requests block, in-flight LLMs drain, Ollama models are unloaded, the prompt
+  is forwarded, and the lock is held until the job finishes and ComfyUI frees
+  its VRAM.
+- Everything else (including websockets and all streaming) passes through
+  transparently and unbuffered.
+
+## Configuration
+
+All configuration is via environment variables; invalid values fail at
+startup.
+
+| Var | Default | Meaning |
+|---|---|---|
+| `LISTEN_OLLAMA` | `:11435` | Ollama-facing listener |
+| `LISTEN_COMFY` | `:8189` | ComfyUI-facing listener |
+| `OLLAMA_URL` | `http://127.0.0.1:11434` | Ollama upstream |
+| `COMFY_URL` | `http://127.0.0.1:8188` | ComfyUI upstream |
+| `UNLOAD_TIMEOUT` | `60s` | Wait for Ollama to unload before an image job |
+| `JOB_TIMEOUT` | `15m` | Wait for a ComfyUI job to finish |
+| `LLM_WAIT_TIMEOUT` | `10m` | Max lock wait for an LLM request before 503 |
+| `WARM_MODEL` | _(empty)_ | Model to reload after an image job (off by default) |
+| `LOG_LEVEL` | `info` | `debug` logs every lock transition |
+| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
+
+## Observability
+
+- `GET /healthz` (both listeners): `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}`
+- `GET /metrics` (Ollama listener): Prometheus text format — `gpulock_state`,
+  `gpulock_llm_inflight`, `gpulock_image_pending`, `gpulock_image_jobs_total`,
+  `gpulock_lock_wait_seconds` (histogram, `kind="llm|image"`),
+  `gpulock_unload_seconds`.
+
+## Build and run
+
+```sh
+go build ./cmd/gpulock
+./gpulock
+```
+
+```sh
+docker build -t gpulock .
+docker run --rm -p 11435:11435 -p 8189:8189 \
+  -e OLLAMA_URL=http://<workstation-ip>:11434 \
+  -e COMFY_URL=http://<workstation-ip>:8188 \
+  gpulock
+```
+
+Releases are built by Gitea Actions (`.gitea/workflows/ci.yml`): pushes run
+`go vet` and `go test -race` and publish
+`git.rambossek.at/<owner>/gpulock:sha-<short>`; `main` additionally gets
+`:latest`, and a git tag `vX.Y.Z` produces the versioned image.
+
+## Development
+
+```sh
+go vet ./...
+go test -race ./...
+```
+
+Stdlib only, Go 1.23+. Layout:
+
+```
+cmd/gpulock/main.go        wiring, config, listeners
+internal/lock/             two-mode lock (LLM readers / image writer, FIFO)
+internal/ollama/           ps / unload / warm client
+internal/comfy/            history poll / free client
+internal/proxy/            handlers for both listeners
+internal/metrics/          Prometheus exposition, no dependencies
+```

SPEC.md

@@ -0,0 +1,211 @@
+# gpulock — GPU arbitration proxy for Ollama + ComfyUI
+
+## Problem
+
+One consumer GPU (RTX 5080, 16 GB) is shared by an LLM server (Ollama) and an
+image generator (ComfyUI). Both assume they own the card. When both hold models
+at once, the NVIDIA Windows driver falls back to system memory and everything
+becomes very slow; on Linux it would OOM instead.
+
+## Goal
+
+A single Go binary that sits in front of **both** services and guarantees that
+at any moment the GPU is in exactly one of three states:
+
+- `idle` — nothing in flight
+- `llm` — N ≥ 1 Ollama inference requests in flight (concurrency allowed)
+- `image` — exactly one ComfyUI job in flight, Ollama models unloaded
+
+Clients (LiteLLM, Open WebUI, n8n) point at gpulock instead of at the services.
+gpulock is transparent for everything that does not touch the GPU.
+
+## Non-goals
+
+- Not a scheduler across multiple GPUs or hosts. One lock, one card.
+- No auth, TLS, rate limiting. Runs on an internal network behind Traefik or a
+  Docker bridge.
+- No request rewriting, caching, or protocol translation.
+- No persistence. Restart = idle state.
+
+## Architecture
+
+```
+LiteLLM / Open WebUI ──► :11435 ─┐                     ┌─► Ollama  :11434
+                                 ├── gpulock (1 lock) ──┤
+Open WebUI / n8n ────► :8189  ───┘                     └─► ComfyUI :8188
+```
+
+Two listeners, one process, one lock. Each listener is an
+`httputil.ReverseProxy` to its upstream. Websocket upgrades (ComfyUI `/ws`)
+and streaming bodies (Ollama NDJSON / SSE) must pass through unbuffered
+(`FlushInterval = -1`).
+
+### Lock semantics
+
+Two-mode lock with image priority (writer-preferring RW lock, where "readers"
+are LLM requests and the single "writer" is an image job):
+
+- **LLM request** (see endpoint list): `AcquireLLM()` blocks while state is
+  `image` **or while an image job is waiting**. Then state := `llm`, n++.
+  On completion (response fully written, including streamed bodies, or client
+  disconnect) n--; if n == 0 state := `idle`.
+- **Image job**: `AcquireImage()` marks "image pending" (so no new LLM
+  requests start), waits until n == 0, sets state := `image`. Released after
+  the ComfyUI job finished and models were freed.
+- Concurrent image jobs queue FIFO behind each other.
+- All waits are context-aware: a client that disconnects while waiting is
+  removed from the queue.
+
+### Endpoint classification
+
+Ollama listener (`:11435` → `OLLAMA_URL`):
+
+| Path | Handling |
+|---|---|
+| `POST /api/generate`, `/api/chat`, `/api/embed`, `/api/embeddings` | LLM lock |
+| `POST /v1/chat/completions`, `/v1/completions`, `/v1/embeddings` | LLM lock |
+| everything else (`/api/tags`, `/api/ps`, `/api/show`, `/api/version`, `/v1/models`, `/api/pull`, …) | pass-through, no lock |
+
+ComfyUI listener (`:8189` → `COMFY_URL`):
+
+| Path | Handling |
+|---|---|
+| `POST /prompt` | image lock (see flow below) |
+| everything else (`/ws`, `/history/*`, `/view`, `/system_stats`, `/queue`, `/free`, …) | pass-through, no lock |
+
+### Image job flow (`POST /prompt`)
+
+1. `AcquireImage()`.
+2. Unload Ollama: `GET /api/ps`; for each model `POST /api/generate
+   {"model":M,"keep_alive":0}`; if that returns non-2xx (embedding-only
+   models), `POST /api/embed {"model":M,"input":"x","keep_alive":0}`. Poll
+   `/api/ps` every 500 ms until empty or `UNLOAD_TIMEOUT`. On timeout: log and
+   continue (degrade, don't fail the user's request).
+3. Forward the original request body to ComfyUI `/prompt`, return status,
+   headers and body to the caller unchanged, flush.
+4. If the response is 200 and contains `prompt_id`: in a goroutine, poll
+   `GET /history/<prompt_id>` every 1 s until the entry has
+   `status.completed == true`, `status.status_str == "error"`, or
+   `JOB_TIMEOUT`. Then `POST /free {"unload_models":true,"free_memory":true}`.
+   Then release the image lock.
+5. If the response is not 200 or has no `prompt_id`: release the lock
+   immediately.
+
+Optional (config flag `WARM_MODEL`): after releasing the image lock, if the
+state is `idle`, send `POST /api/generate {"model":WARM_MODEL,"keep_alive":-1}`
+with empty prompt to reload the chat model so the next chat doesn't pay the
+load time. Off by default.
+
+## Configuration (env)
+
+| Var | Default | Meaning |
+|---|---|---|
+| `LISTEN_OLLAMA` | `:11435` | Ollama-facing listener |
+| `LISTEN_COMFY` | `:8189` | ComfyUI-facing listener |
+| `OLLAMA_URL` | `http://127.0.0.1:11434` | upstream |
+| `COMFY_URL` | `http://127.0.0.1:8188` | upstream |
+| `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload |
+| `JOB_TIMEOUT` | `15m` | wait for ComfyUI job |
+| `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 |
+| `WARM_MODEL` | `` | optional model to reload after an image job |
+| `LOG_LEVEL` | `info` | `debug` logs every lock transition |
+
+Startup fails fast on unparsable values. Both upstreams are probed once at
+start (`/api/version`, `/system_stats`); failure is logged, not fatal.
+
+## Observability
+
+- `GET /healthz` on both listeners: 200 with JSON
+  `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}`.
+- `GET /metrics` on the Ollama listener: Prometheus text format, no external
+  dependency needed:
+  `gpulock_state{state="…"} 1`, `gpulock_llm_inflight`,
+  `gpulock_image_jobs_total`, `gpulock_lock_wait_seconds` (histogram, label
+  `kind="llm|image"`), `gpulock_unload_seconds`.
+- Structured logs (`log/slog`, JSON when `LOG_FORMAT=json`), one line per
+  state transition and per image job phase with `prompt_id`.
+
+## Edge cases to handle
+
+- Client disconnects while streaming an Ollama response: request context is
+  cancelled, proxy aborts upstream, in-flight counter still decrements.
+- Client disconnects while waiting for the lock: removed from wait, no
+  counter change.
+- ComfyUI job finishes but `/history` never shows it (e.g. ComfyUI restarted):
+  `JOB_TIMEOUT` releases the lock; log at warn.
+- Ollama unreachable during unload: continue with the image job; the whole
+  point is not to block users on a misbehaving neighbour.
+- `POST /prompt` with a body that ComfyUI rejects (400): lock released
+  immediately, body passed back.
+- Websocket `/ws` connections are long-lived and never take the lock.
+- The Ollama OpenAI-compatible endpoints stream SSE; the proxy must not buffer.
+
+## Repository layout
+
+```
+gpulock/
+  cmd/gpulock/main.go        # wiring, config, listeners
+  internal/lock/lock.go      # two-mode lock + tests
+  internal/ollama/client.go  # ps / unload / warm
+  internal/comfy/client.go   # history poll / free
+  internal/proxy/            # handlers for both listeners
+  Dockerfile
+  .gitea/workflows/ci.yml
+  README.md
+  SPEC.md                    # this file
+```
+
+`main.go` from the first prototype (ComfyUI-only) is the starting point for
+`internal/comfy` and the `/prompt` handler; the lock and the Ollama listener
+are new.
+
+## Testing
+
+- `internal/lock`: table tests plus a race test (`go test -race`) with
+  goroutines: image waits for LLMs to drain; new LLMs block while image is
+  pending; FIFO for images; context cancellation removes waiters.
+- `internal/proxy`: `httptest.Server` fakes for Ollama (`/api/ps`,
+  `/api/generate`) and ComfyUI (`/prompt`, `/history/:id`, `/free`); assert
+  the call sequence for one image job and that a concurrent `/api/chat` is
+  held until `/free` was called.
+- Streaming test: fake Ollama emits chunks with delays; assert the client
+  receives the first chunk before the last is sent (no buffering).
+
+## Build and CI
+
+- Go 1.23+, stdlib only. `CGO_ENABLED=0`, `-ldflags="-s -w"`, version from
+  `git describe` injected via `-X main.version=`.
+- Dockerfile: multi-stage, final image `gcr.io/distroless/static` (or
+  `scratch`), non-root user, `EXPOSE 8189 11435`, `ENTRYPOINT ["/gpulock"]`.
+- `.gitea/workflows/ci.yml` (Gitea Actions):
+  1. on push and tag: `go vet`, `go test -race ./...`, `golangci-lint` if
+     available in the runner image
+  2. build image with buildx, tags `:sha-<short>` and `:latest` on main,
+     `:<tag>` on tags
+  3. push to the Gitea registry `git.rambossek.at/<owner>/gpulock` using the
+     workflow token (`${{ secrets.GITEA_TOKEN }}` / `gitea.actor`)
+- Release: a git tag `vX.Y.Z` produces the versioned image; the Open WebUI
+  compose pins that tag.
+
+## Deployment (target)
+
+```yaml
+  gpulock:
+    image: git.rambossek.at/<owner>/gpulock:v0.1.0
+    environment:
+      OLLAMA_URL: http://<workstation-ip>:11434
+      COMFY_URL:  http://<workstation-ip>:8188
+    networks: [internal]
+```
+
+LiteLLM `api_base` → `http://gpulock:11435`; Open WebUI
+`COMFYUI_BASE_URL` → `http://gpulock:8189`. Nothing else talks to the
+workstation directly.
+
+## Open questions
+
+- Should embedding requests (`/api/embed`, `/v1/embeddings`) count as LLM
+  traffic for the lock? They do in this spec (they hold VRAM); reconsider if
+  RAG indexing starves image jobs for too long.
+- Whether to add a `POST /gpulock/release` admin endpoint to force-reset the
+  lock without restarting. Cheap to add; decide once it's been stuck once.

cmd/gpulock/main.go

@@ -0,0 +1,193 @@
+// gpulock 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 (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"os"
+	"os/signal"
+	"strings"
+	"syscall"
+	"time"
+
+	"gpu-turnstile/internal/comfy"
+	"gpu-turnstile/internal/lock"
+	"gpu-turnstile/internal/metrics"
+	"gpu-turnstile/internal/ollama"
+	"gpu-turnstile/internal/proxy"
+)
+
+// version is injected at build time via -ldflags "-X main.version=...".
+var version = "dev"
+
+type config struct {
+	listenOllama   string
+	listenComfy    string
+	ollamaURL      string
+	comfyURL       string
+	unloadTimeout  time.Duration
+	jobTimeout     time.Duration
+	llmWaitTimeout time.Duration
+	warmModel      string
+	logLevel       slog.Level
+	logJSON        bool
+}
+
+func envDuration(getenv func(string) string, name string, dst *time.Duration) error {
+	v := getenv(name)
+	if v == "" {
+		return nil
+	}
+	d, err := time.ParseDuration(v)
+	if err != nil {
+		return fmt.Errorf("%s: %w", name, err)
+	}
+	*dst = d
+	return nil
+}
+
+func loadConfig(getenv func(string) string) (config, error) {
+	cfg := config{
+		listenOllama:   ":11435",
+		listenComfy:    ":8189",
+		ollamaURL:      "http://127.0.0.1:11434",
+		comfyURL:       "http://127.0.0.1:8188",
+		unloadTimeout:  time.Minute,
+		jobTimeout:     15 * time.Minute,
+		llmWaitTimeout: 10 * time.Minute,
+		logLevel:       slog.LevelInfo,
+	}
+	for _, e := range []struct {
+		name string
+		dst  *string
+	}{
+		{"LISTEN_OLLAMA", &cfg.listenOllama},
+		{"LISTEN_COMFY", &cfg.listenComfy},
+		{"OLLAMA_URL", &cfg.ollamaURL},
+		{"COMFY_URL", &cfg.comfyURL},
+		{"WARM_MODEL", &cfg.warmModel},
+	} {
+		if v := getenv(e.name); v != "" {
+			*e.dst = v
+		}
+	}
+	for _, e := range []struct {
+		name string
+		dst  *time.Duration
+	}{
+		{"UNLOAD_TIMEOUT", &cfg.unloadTimeout},
+		{"JOB_TIMEOUT", &cfg.jobTimeout},
+		{"LLM_WAIT_TIMEOUT", &cfg.llmWaitTimeout},
+	} {
+		if err := envDuration(getenv, e.name, e.dst); err != nil {
+			return cfg, err
+		}
+	}
+	if v := getenv("LOG_LEVEL"); v != "" {
+		var level slog.Level
+		if err := level.UnmarshalText([]byte(v)); err != nil {
+			return cfg, fmt.Errorf("LOG_LEVEL: %w", err)
+		}
+		cfg.logLevel = level
+	}
+	switch strings.ToLower(getenv("LOG_FORMAT")) {
+	case "", "text":
+	case "json":
+		cfg.logJSON = true
+	default:
+		return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"")
+	}
+	return cfg, nil
+}
+
+func main() {
+	cfg, err := loadConfig(os.Getenv)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "gpulock: %v\n", err)
+		os.Exit(1)
+	}
+
+	opts := &slog.HandlerOptions{Level: cfg.logLevel}
+	var handler slog.Handler = slog.NewTextHandler(os.Stderr, opts)
+	if cfg.logJSON {
+		handler = slog.NewJSONHandler(os.Stderr, opts)
+	}
+	log := slog.New(handler)
+	slog.SetDefault(log)
+
+	log.Info("starting gpulock",
+		"version", version,
+		"listen_ollama", cfg.listenOllama,
+		"listen_comfy", cfg.listenComfy,
+		"ollama_url", cfg.ollamaURL,
+		"comfy_url", cfg.comfyURL,
+	)
+
+	ollamaClient, err := ollama.New(cfg.ollamaURL, log)
+	if err != nil {
+		log.Error("invalid configuration", "err", err)
+		os.Exit(1)
+	}
+	comfyClient, err := comfy.New(cfg.comfyURL, log)
+	if err != nil {
+		log.Error("invalid configuration", "err", err)
+		os.Exit(1)
+	}
+
+	srv, err := proxy.New(proxy.Config{
+		OllamaURL:      cfg.ollamaURL,
+		ComfyURL:       cfg.comfyURL,
+		Lock:           lock.New(log),
+		Ollama:         ollamaClient,
+		Comfy:          comfyClient,
+		Metrics:        metrics.New(),
+		Log:            log,
+		LLMWaitTimeout: cfg.llmWaitTimeout,
+		UnloadTimeout:  cfg.unloadTimeout,
+		JobTimeout:     cfg.jobTimeout,
+		WarmModel:      cfg.warmModel,
+	})
+	if err != nil {
+		log.Error("invalid configuration", "err", err)
+		os.Exit(1)
+	}
+
+	// Probe both upstreams once; failure is logged, not fatal.
+	probeCtx, probeCancel := context.WithTimeout(context.Background(), 5*time.Second)
+	if err := ollamaClient.Probe(probeCtx); err != nil {
+		log.Warn("ollama probe failed", "url", cfg.ollamaURL, "err", err)
+	}
+	if err := comfyClient.Probe(probeCtx); err != nil {
+		log.Warn("comfy probe failed", "url", cfg.comfyURL, "err", err)
+	}
+	probeCancel()
+
+	ollamaSrv := &http.Server{Addr: cfg.listenOllama, Handler: srv.OllamaHandler()}
+	comfySrv := &http.Server{Addr: cfg.listenComfy, Handler: srv.ComfyHandler()}
+
+	errCh := make(chan error, 2)
+	go func() { errCh <- ollamaSrv.ListenAndServe() }()
+	go func() { errCh <- comfySrv.ListenAndServe() }()
+
+	sigCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
+
+	select {
+	case err := <-errCh:
+		if err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Error("listener failed", "err", err)
+			os.Exit(1)
+		}
+	case <-sigCtx.Done():
+		log.Info("shutting down")
+	}
+
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer shutdownCancel()
+	ollamaSrv.Shutdown(shutdownCtx)
+	comfySrv.Shutdown(shutdownCtx)
+}

go.mod

@@ -0,0 +1,3 @@
+module gpu-turnstile
+
+go 1.23

internal/comfy/client.go

@@ -0,0 +1,138 @@
+// Package comfy is a minimal client for the ComfyUI endpoints gpulock needs
+// after a prompt has been accepted: polling /history and freeing VRAM.
+package comfy
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/url"
+	"time"
+)
+
+// Client talks to a ComfyUI upstream.
+type Client struct {
+	base string
+	hc   *http.Client
+	log  *slog.Logger
+
+	// PollInterval is how often /history/<id> is polled. Defaults to 1s;
+	// tests can shorten it.
+	PollInterval time.Duration
+}
+
+// New validates baseURL and returns a Client.
+func New(baseURL string, log *slog.Logger) (*Client, error) {
+	u, err := url.Parse(baseURL)
+	if err != nil || u.Scheme == "" || u.Host == "" {
+		return nil, fmt.Errorf("invalid ComfyUI URL %q", baseURL)
+	}
+	return &Client{
+		base:         u.String(),
+		hc:           &http.Client{Timeout: 30 * time.Second},
+		log:          log,
+		PollInterval: time.Second,
+	}, nil
+}
+
+// Probe checks that ComfyUI answers on /system_stats.
+func (c *Client) Probe(ctx context.Context) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/system_stats", nil)
+	if err != nil {
+		return err
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	io.Copy(io.Discard, resp.Body)
+	if resp.StatusCode/100 != 2 {
+		return fmt.Errorf("/system_stats returned %s", resp.Status)
+	}
+	return nil
+}
+
+type historyEntry struct {
+	Status struct {
+		Completed bool   `json:"completed"`
+		StatusStr string `json:"status_str"`
+	} `json:"status"`
+}
+
+// jobDone reports whether the history entry for promptID shows a finished
+// job (completed or errored).
+func (c *Client) jobDone(ctx context.Context, promptID string) (bool, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/history/"+promptID, nil)
+	if err != nil {
+		return false, err
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return false, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode/100 != 2 {
+		io.Copy(io.Discard, resp.Body)
+		return false, fmt.Errorf("/history/%s returned %s", promptID, resp.Status)
+	}
+	var history map[string]historyEntry
+	if err := json.NewDecoder(resp.Body).Decode(&history); err != nil {
+		return false, err
+	}
+	entry, ok := history[promptID]
+	if !ok {
+		return false, nil
+	}
+	return entry.Status.Completed || entry.Status.StatusStr == "error", nil
+}
+
+// WaitJob polls until the job for promptID is finished. The context should
+// carry the job deadline; transient polling errors are logged and retried,
+// so a returned error is always from the context.
+func (c *Client) WaitJob(ctx context.Context, promptID string) error {
+	for {
+		done, err := c.jobDone(ctx, promptID)
+		if err != nil {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			c.log.Warn("failed to poll job history", "prompt_id", promptID, "err", err)
+		}
+		if done {
+			return nil
+		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(c.PollInterval):
+		}
+	}
+}
+
+// Free asks ComfyUI to unload its models and free VRAM.
+func (c *Client) Free(ctx context.Context) error {
+	buf, err := json.Marshal(map[string]bool{"unload_models": true, "free_memory": true})
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.base+"/free", bytes.NewReader(buf))
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	io.Copy(io.Discard, resp.Body)
+	if resp.StatusCode/100 != 2 {
+		return fmt.Errorf("/free returned %s", resp.Status)
+	}
+	return nil
+}

internal/lock/lock.go

@@ -0,0 +1,157 @@
+// Package lock implements the two-mode GPU arbitration lock: any number of
+// concurrent LLM requests ("readers") or exactly one image job ("writer"),
+// with image jobs taking priority over newly arriving LLM requests.
+package lock
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+)
+
+// State is the current GPU occupancy state.
+type State string
+
+const (
+	StateIdle  State = "idle"
+	StateLLM   State = "llm"
+	StateImage State = "image"
+)
+
+type imageWaiter struct{ id uint64 }
+
+// Lock is a writer-preferring two-mode lock. The zero value is not usable;
+// construct with New.
+type Lock struct {
+	mu     sync.Mutex
+	change chan struct{} // closed and replaced on every state change
+
+	n           int  // LLM requests in flight
+	imageActive bool // an image job holds the GPU
+	imageQ      []imageWaiter
+	nextID      uint64
+
+	log *slog.Logger
+}
+
+// New returns a ready-to-use Lock. log may be nil; if set, every state
+// transition is logged at debug level.
+func New(log *slog.Logger) *Lock {
+	return &Lock{change: make(chan struct{}), log: log}
+}
+
+// broadcast wakes all waiters. Call with mu held.
+func (l *Lock) broadcast() {
+	close(l.change)
+	l.change = make(chan struct{})
+}
+
+func (l *Lock) logTransition(msg string, args ...any) {
+	if l.log != nil {
+		l.log.Debug(msg, args...)
+	}
+}
+
+// AcquireLLM blocks until no image job is active or pending, then registers
+// one in-flight LLM request. Returns ctx.Err() if the context is cancelled
+// while waiting; no state is changed in that case.
+func (l *Lock) AcquireLLM(ctx context.Context) error {
+	l.mu.Lock()
+	for l.imageActive || len(l.imageQ) > 0 {
+		ch := l.change
+		l.mu.Unlock()
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ch:
+		}
+		l.mu.Lock()
+	}
+	l.n++
+	n := l.n
+	l.mu.Unlock()
+	l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
+	return nil
+}
+
+// ReleaseLLM marks one LLM request as finished.
+func (l *Lock) ReleaseLLM() {
+	l.mu.Lock()
+	l.n--
+	n := l.n
+	if l.n == 0 {
+		l.broadcast()
+	}
+	l.mu.Unlock()
+	if n == 0 {
+		l.logTransition("lock transition", "state", StateIdle, "llm_inflight", 0)
+	}
+}
+
+// AcquireImage queues the caller FIFO behind other image jobs, blocks new
+// LLM requests immediately, and waits for in-flight LLM requests to drain
+// before granting exclusive GPU access. Returns ctx.Err() if the context is
+// cancelled while waiting; the caller is removed from the queue.
+func (l *Lock) AcquireImage(ctx context.Context) error {
+	l.mu.Lock()
+	l.nextID++
+	w := imageWaiter{id: l.nextID}
+	l.imageQ = append(l.imageQ, w)
+	pending := len(l.imageQ)
+	l.broadcast() // new LLM requests must block from now on
+	l.mu.Unlock()
+	l.logTransition("image job pending", "image_queue", pending)
+
+	for {
+		l.mu.Lock()
+		if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive {
+			l.imageQ = l.imageQ[1:]
+			l.imageActive = true
+			l.mu.Unlock()
+			l.logTransition("lock transition", "state", StateImage)
+			return nil
+		}
+		ch := l.change
+		l.mu.Unlock()
+		select {
+		case <-ctx.Done():
+			l.mu.Lock()
+			for i, q := range l.imageQ {
+				if q.id == w.id {
+					l.imageQ = append(l.imageQ[:i], l.imageQ[i+1:]...)
+					break
+				}
+			}
+			l.broadcast()
+			l.mu.Unlock()
+			l.logTransition("image job wait cancelled", "image_queue", len(l.imageQ))
+			return ctx.Err()
+		case <-ch:
+		}
+	}
+}
+
+// ReleaseImage frees the GPU after an image job.
+func (l *Lock) ReleaseImage() {
+	l.mu.Lock()
+	l.imageActive = false
+	l.broadcast()
+	l.mu.Unlock()
+	l.logTransition("lock transition", "state", StateIdle)
+}
+
+// Snapshot reports the current state, the number of in-flight LLM requests,
+// and whether an image job is active or waiting.
+func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	switch {
+	case l.imageActive:
+		state = StateImage
+	case l.n > 0:
+		state = StateLLM
+	default:
+		state = StateIdle
+	}
+	return state, l.n, l.imageActive || len(l.imageQ) > 0
+}

internal/lock/lock_test.go

@@ -0,0 +1,228 @@
+package lock
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+func waitFor(t *testing.T, ch <-chan struct{}, what string) {
+	t.Helper()
+	select {
+	case <-ch:
+	case <-time.After(2 * time.Second):
+		t.Fatalf("timed out waiting for %s", what)
+	}
+}
+
+func assertBlocked(t *testing.T, ch <-chan struct{}, what string) {
+	t.Helper()
+	select {
+	case <-ch:
+		t.Fatalf("%s should still be blocked", what)
+	case <-time.After(50 * time.Millisecond):
+	}
+}
+
+func TestLLMConcurrency(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+	for i := 0; i < 3; i++ {
+		if err := lk.AcquireLLM(ctx); err != nil {
+			t.Fatal(err)
+		}
+	}
+	state, n, pending := lk.Snapshot()
+	if state != StateLLM || n != 3 || pending {
+		t.Fatalf("got state=%s n=%d pending=%v", state, n, pending)
+	}
+	lk.ReleaseLLM()
+	lk.ReleaseLLM()
+	if state, _, _ := lk.Snapshot(); state != StateLLM {
+		t.Fatalf("state=%s, want llm", state)
+	}
+	lk.ReleaseLLM()
+	if state, n, _ := lk.Snapshot(); state != StateIdle || n != 0 {
+		t.Fatalf("got state=%s n=%d, want idle", state, n)
+	}
+}
+
+func TestImageWaitsForLLMDrainAndBlocksNewLLM(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	imageAcquired := make(chan struct{})
+	go func() {
+		if err := lk.AcquireImage(ctx); err != nil {
+			t.Error(err)
+		}
+		close(imageAcquired)
+	}()
+	assertBlocked(t, imageAcquired, "image acquire while LLMs in flight")
+
+	// While an image job is pending, new LLM requests must block even
+	// though LLM concurrency is otherwise allowed.
+	llmAcquired := make(chan struct{})
+	go func() {
+		if err := lk.AcquireLLM(ctx); err != nil {
+			t.Error(err)
+		}
+		close(llmAcquired)
+	}()
+	assertBlocked(t, llmAcquired, "LLM acquire while image pending")
+
+	lk.ReleaseLLM()
+	assertBlocked(t, imageAcquired, "image acquire with one LLM still in flight")
+
+	lk.ReleaseLLM()
+	waitFor(t, imageAcquired, "image acquire after drain")
+	assertBlocked(t, llmAcquired, "LLM acquire while image active")
+
+	state, _, pending := lk.Snapshot()
+	if state != StateImage || !pending {
+		t.Fatalf("got state=%s pending=%v", state, pending)
+	}
+
+	lk.ReleaseImage()
+	waitFor(t, llmAcquired, "LLM acquire after image release")
+	lk.ReleaseLLM()
+}
+
+func TestImageFIFO(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+
+	if err := lk.AcquireImage(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	order := make(chan int, 2)
+	for i, id := range []int{1, 2} {
+		go func() {
+			if err := lk.AcquireImage(ctx); err != nil {
+				t.Error(err)
+				return
+			}
+			order <- id
+		}()
+		// Ensure this waiter is queued before the next one starts.
+		for deadline := time.Now().Add(2 * time.Second); ; {
+			lk.mu.Lock()
+			qlen := len(lk.imageQ)
+			lk.mu.Unlock()
+			if qlen == i+1 {
+				break
+			}
+			if time.Now().After(deadline) {
+				t.Fatalf("image waiter %d never queued", id)
+			}
+			time.Sleep(time.Millisecond)
+		}
+	}
+
+	lk.ReleaseImage()
+	if got := <-order; got != 1 {
+		t.Fatalf("first image job = %d, want 1", got)
+	}
+	select {
+	case got := <-order:
+		t.Fatalf("second image job %d acquired while first active", got)
+	case <-time.After(50 * time.Millisecond):
+	}
+	lk.ReleaseImage()
+	if got := <-order; got != 2 {
+		t.Fatalf("second image job = %d, want 2", got)
+	}
+	lk.ReleaseImage()
+}
+
+func TestContextCancelRemovesLLMWaiter(t *testing.T) {
+	lk := New(nil)
+	if err := lk.AcquireImage(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan error, 1)
+	go func() { done <- lk.AcquireLLM(ctx) }()
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+	if err := <-done; err == nil {
+		t.Fatal("expected context error")
+	}
+
+	lk.ReleaseImage()
+	if state, n, _ := lk.Snapshot(); state != StateIdle || n != 0 {
+		t.Fatalf("got state=%s n=%d, want idle", state, n)
+	}
+}
+
+func TestContextCancelRemovesImageWaiter(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	wctx, cancel := context.WithCancel(ctx)
+	done := make(chan error, 1)
+	go func() { done <- lk.AcquireImage(wctx) }()
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+	if err := <-done; err == nil {
+		t.Fatal("expected context error")
+	}
+
+	lk.ReleaseLLM()
+	// The cancelled waiter must be gone: a new LLM request acquires
+	// immediately instead of blocking behind it.
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+	if _, _, pending := lk.Snapshot(); pending {
+		t.Fatal("image_pending=true after waiter was cancelled")
+	}
+	lk.ReleaseLLM()
+}
+
+// TestRace hammers the lock from many goroutines; run with -race.
+func TestRace(t *testing.T) {
+	lk := New(nil)
+	var wg sync.WaitGroup
+	for i := 0; i < 8; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			for j := 0; j < 25; j++ {
+				if i%4 == 3 {
+					if err := lk.AcquireImage(context.Background()); err != nil {
+						t.Error(err)
+						return
+					}
+					lk.Snapshot()
+					lk.ReleaseImage()
+				} else {
+					if err := lk.AcquireLLM(context.Background()); err != nil {
+						t.Error(err)
+						return
+					}
+					lk.Snapshot()
+					lk.ReleaseLLM()
+				}
+			}
+		}(i)
+	}
+	wg.Wait()
+	state, n, pending := lk.Snapshot()
+	if state != StateIdle || n != 0 || pending {
+		t.Fatalf("leaked lock state: state=%s n=%d pending=%v", state, n, pending)
+	}
+}

internal/metrics/metrics.go

@@ -0,0 +1,121 @@
+// Package metrics provides the Prometheus text exposition for gpulock
+// without any external dependencies.
+package metrics
+
+import (
+	"fmt"
+	"io"
+	"math"
+	"sync/atomic"
+)
+
+var buckets = []float64{0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 120, 300, 600}
+
+// histogram is a fixed-bucket cumulative histogram safe for concurrent use.
+type histogram struct {
+	counts []atomic.Uint64 // len(buckets)+1; last slot is +Inf
+	sum    atomic.Uint64   // math.Float64bits
+	count  atomic.Uint64
+}
+
+func newHistogram() *histogram {
+	return &histogram{counts: make([]atomic.Uint64, len(buckets)+1)}
+}
+
+func (h *histogram) observe(v float64) {
+	for i, b := range buckets {
+		if v <= b {
+			h.counts[i].Add(1)
+		}
+	}
+	h.counts[len(buckets)].Add(1)
+	h.sum.Add(math.Float64bits(v))
+	h.count.Add(1)
+}
+
+func (h *histogram) write(w io.Writer, name, label string) {
+	for i, b := range buckets {
+		fmt.Fprintf(w, "%s_bucket{%sle=\"%g\"} %d\n", name, label, b, h.counts[i].Load())
+	}
+	fmt.Fprintf(w, "%s_bucket{%sle=\"+Inf\"} %d\n", name, label, h.counts[len(buckets)].Load())
+	fmt.Fprintf(w, "%s_sum{%s} %g\n", name, trimLabel(label), math.Float64frombits(h.sum.Load()))
+	fmt.Fprintf(w, "%s_count{%s} %d\n", name, trimLabel(label), h.count.Load())
+}
+
+// trimLabel converts a `kind="llm",`-style prefix into the label set used for
+// _sum/_count lines (same labels, no trailing comma).
+func trimLabel(label string) string {
+	if len(label) > 0 && label[len(label)-1] == ',' {
+		return label[:len(label)-1]
+	}
+	return label
+}
+
+// Metrics holds all gpulock metric values.
+type Metrics struct {
+	llmWait   *histogram
+	imageWait *histogram
+	unload    *histogram
+	imageJobs atomic.Uint64
+}
+
+// New returns a zeroed Metrics.
+func New() *Metrics {
+	return &Metrics{
+		llmWait:   newHistogram(),
+		imageWait: newHistogram(),
+		unload:    newHistogram(),
+	}
+}
+
+// ObserveLockWait records how long a lock acquisition of the given kind
+// ("llm" or "image") waited.
+func (m *Metrics) ObserveLockWait(kind string, seconds float64) {
+	if kind == "image" {
+		m.imageWait.observe(seconds)
+	} else {
+		m.llmWait.observe(seconds)
+	}
+}
+
+// ObserveUnload records how long an Ollama unload took.
+func (m *Metrics) ObserveUnload(seconds float64) { m.unload.observe(seconds) }
+
+// IncImageJobs counts one accepted image job.
+func (m *Metrics) IncImageJobs() { m.imageJobs.Add(1) }
+
+// Render writes the Prometheus text exposition for the given lock snapshot.
+func (m *Metrics) Render(w io.Writer, state string, llmInflight int, imagePending bool) {
+	fmt.Fprint(w, `# HELP gpulock_state Current GPU state (1 for the active state).
+# TYPE gpulock_state gauge
+`)
+	for _, s := range []string{"idle", "llm", "image"} {
+		v := 0
+		if s == state {
+			v = 1
+		}
+		fmt.Fprintf(w, "gpulock_state{state=%q} %d\n", s, v)
+	}
+	pending := 0
+	if imagePending {
+		pending = 1
+	}
+	fmt.Fprintf(w, `# HELP gpulock_llm_inflight LLM requests currently in flight.
+# TYPE gpulock_llm_inflight gauge
+gpulock_llm_inflight %d
+# HELP gpulock_image_pending Whether an image job is active or waiting.
+# TYPE gpulock_image_pending gauge
+gpulock_image_pending %d
+# HELP gpulock_image_jobs_total Image jobs accepted by ComfyUI.
+# TYPE gpulock_image_jobs_total counter
+gpulock_image_jobs_total %d
+# HELP gpulock_lock_wait_seconds Time spent waiting to acquire the GPU lock.
+# TYPE gpulock_lock_wait_seconds histogram
+`, llmInflight, pending, m.imageJobs.Load())
+	m.llmWait.write(w, "gpulock_lock_wait_seconds", `kind="llm",`)
+	m.imageWait.write(w, "gpulock_lock_wait_seconds", `kind="image",`)
+	fmt.Fprint(w, `# HELP gpulock_unload_seconds Time spent unloading Ollama models before an image job.
+# TYPE gpulock_unload_seconds histogram
+`)
+	m.unload.write(w, "gpulock_unload_seconds", "")
+}

internal/ollama/client.go

@@ -0,0 +1,180 @@
+// Package ollama is a minimal client for the Ollama management endpoints
+// gpulock needs: listing loaded models, unloading them, and warming a model.
+package ollama
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/url"
+	"time"
+)
+
+// Client talks to an Ollama upstream.
+type Client struct {
+	base string
+	hc   *http.Client
+	log  *slog.Logger
+
+	// PollInterval is how often /api/ps is re-checked while waiting for
+	// models to unload. Defaults to 500ms; tests can shorten it.
+	PollInterval time.Duration
+}
+
+// New validates baseURL and returns a Client.
+func New(baseURL string, log *slog.Logger) (*Client, error) {
+	u, err := url.Parse(baseURL)
+	if err != nil || u.Scheme == "" || u.Host == "" {
+		return nil, fmt.Errorf("invalid Ollama URL %q", baseURL)
+	}
+	return &Client{
+		base:         u.String(),
+		hc:           &http.Client{Timeout: 30 * time.Second},
+		log:          log,
+		PollInterval: 500 * time.Millisecond,
+	}, nil
+}
+
+// Probe checks that Ollama answers on /api/version.
+func (c *Client) Probe(ctx context.Context) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/api/version", nil)
+	if err != nil {
+		return err
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	io.Copy(io.Discard, resp.Body)
+	if resp.StatusCode/100 != 2 {
+		return fmt.Errorf("/api/version returned %s", resp.Status)
+	}
+	return nil
+}
+
+type psResponse struct {
+	Models []struct {
+		Name  string `json:"name"`
+		Model string `json:"model"`
+	} `json:"models"`
+}
+
+// LoadedModels returns the names of models currently held in memory.
+func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/api/ps", nil)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode/100 != 2 {
+		io.Copy(io.Discard, resp.Body)
+		return nil, fmt.Errorf("/api/ps returned %s", resp.Status)
+	}
+	var ps psResponse
+	if err := json.NewDecoder(resp.Body).Decode(&ps); err != nil {
+		return nil, err
+	}
+	models := make([]string, 0, len(ps.Models))
+	for _, m := range ps.Models {
+		if m.Name != "" {
+			models = append(models, m.Name)
+		} else {
+			models = append(models, m.Model)
+		}
+	}
+	return models, nil
+}
+
+func (c *Client) post(ctx context.Context, path string, body any) (int, error) {
+	buf, err := json.Marshal(body)
+	if err != nil {
+		return 0, err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.base+path, bytes.NewReader(buf))
+	if err != nil {
+		return 0, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return 0, err
+	}
+	defer resp.Body.Close()
+	io.Copy(io.Discard, resp.Body)
+	return resp.StatusCode, nil
+}
+
+// unloadModel asks Ollama to evict one model. Embedding-only models reject
+// /api/generate, so fall back to /api/embed on non-2xx.
+func (c *Client) unloadModel(ctx context.Context, model string) error {
+	status, err := c.post(ctx, "/api/generate", map[string]any{"model": model, "keep_alive": 0})
+	if err == nil && status/100 == 2 {
+		return nil
+	}
+	status, err2 := c.post(ctx, "/api/embed", map[string]any{"model": model, "input": "x", "keep_alive": 0})
+	if err2 != nil {
+		return fmt.Errorf("generate: %v; embed: %v", err, err2)
+	}
+	if status/100 != 2 {
+		return fmt.Errorf("embed unload returned status %d", status)
+	}
+	return nil
+}
+
+// UnloadAll evicts every loaded model and waits until /api/ps is empty. The
+// context should carry the unload deadline; when it expires UnloadAll returns
+// ctx.Err() so the caller can degrade instead of failing. It returns the
+// elapsed wall time.
+func (c *Client) UnloadAll(ctx context.Context) (time.Duration, error) {
+	start := time.Now()
+	models, err := c.LoadedModels(ctx)
+	if err != nil {
+		return time.Since(start), fmt.Errorf("list models: %w", err)
+	}
+	for _, m := range models {
+		if err := c.unloadModel(ctx, m); err != nil {
+			c.log.Warn("failed to unload model", "model", m, "err", err)
+		}
+	}
+	for {
+		models, err := c.LoadedModels(ctx)
+		if err == nil && len(models) == 0 {
+			return time.Since(start), nil
+		}
+		if err != nil {
+			c.log.Warn("failed to poll /api/ps", "err", err)
+		}
+		select {
+		case <-ctx.Done():
+			return time.Since(start), ctx.Err()
+		case <-time.After(c.PollInterval):
+		}
+	}
+}
+
+// Warm reloads a model and pins it in memory (keep_alive=-1) so the next
+// chat request does not pay the load time.
+func (c *Client) Warm(ctx context.Context, model string) error {
+	status, err := c.post(ctx, "/api/generate", map[string]any{
+		"model":      model,
+		"prompt":     "",
+		"stream":     false,
+		"keep_alive": -1,
+	})
+	if err != nil {
+		return err
+	}
+	if status/100 != 2 {
+		return fmt.Errorf("warm generate returned status %d", status)
+	}
+	return nil
+}

internal/proxy/proxy.go

@@ -0,0 +1,285 @@
+// Package proxy contains the HTTP handlers for both gpulock listeners:
+// reverse proxies to Ollama and ComfyUI with GPU lock arbitration in front
+// of the endpoints that load models.
+package proxy
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"time"
+
+	"gpu-turnstile/internal/comfy"
+	"gpu-turnstile/internal/lock"
+	"gpu-turnstile/internal/metrics"
+	"gpu-turnstile/internal/ollama"
+)
+
+// captureLimit bounds how much of a /prompt response body is buffered while
+// looking for prompt_id. The body still passes through to the client
+// unchanged regardless of size.
+const captureLimit = 64 * 1024
+
+// Config wires a Server.
+type Config struct {
+	OllamaURL string
+	ComfyURL  string
+
+	Lock    *lock.Lock
+	Ollama  *ollama.Client
+	Comfy   *comfy.Client
+	Metrics *metrics.Metrics
+	Log     *slog.Logger
+
+	LLMWaitTimeout time.Duration
+	UnloadTimeout  time.Duration
+	JobTimeout     time.Duration
+	WarmModel      string
+}
+
+// Server serves both gpulock listeners.
+type Server struct {
+	cfg Config
+	log *slog.Logger
+
+	ollamaProxy *httputil.ReverseProxy
+	comfyProxy  *httputil.ReverseProxy
+}
+
+// New builds a Server, validating the upstream URLs.
+func New(cfg Config) (*Server, error) {
+	ollamaURL, err := url.Parse(cfg.OllamaURL)
+	if err != nil || ollamaURL.Scheme == "" || ollamaURL.Host == "" {
+		return nil, fmt.Errorf("invalid OLLAMA_URL %q", cfg.OllamaURL)
+	}
+	comfyURL, err := url.Parse(cfg.ComfyURL)
+	if err != nil || comfyURL.Scheme == "" || comfyURL.Host == "" {
+		return nil, fmt.Errorf("invalid COMFY_URL %q", cfg.ComfyURL)
+	}
+	log := cfg.Log
+	if log == nil {
+		log = slog.Default()
+	}
+	return &Server{
+		cfg:         cfg,
+		log:         log,
+		ollamaProxy: newReverseProxy(ollamaURL, log.With("upstream", "ollama")),
+		comfyProxy:  newReverseProxy(comfyURL, log.With("upstream", "comfy")),
+	}, nil
+}
+
+func newReverseProxy(target *url.URL, log *slog.Logger) *httputil.ReverseProxy {
+	return &httputil.ReverseProxy{
+		Rewrite: func(pr *httputil.ProxyRequest) {
+			pr.SetURL(target)
+			pr.SetXForwarded()
+		},
+		// Flush after every write so NDJSON/SSE streams and websocket
+		// upgrades pass through unbuffered.
+		FlushInterval: -1,
+		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
+			log.Warn("upstream error", "path", r.URL.Path, "err", err)
+			http.Error(w, "upstream unavailable", http.StatusBadGateway)
+		},
+	}
+}
+
+func (s *Server) writeHealthz(w http.ResponseWriter) {
+	state, n, pending := s.cfg.Lock.Snapshot()
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(map[string]any{
+		"state":         state,
+		"llm_inflight":  n,
+		"image_pending": pending,
+	})
+}
+
+func (s *Server) writeMetrics(w http.ResponseWriter) {
+	state, n, pending := s.cfg.Lock.Snapshot()
+	w.Header().Set("Content-Type", "text/plain; version=0.0.4")
+	s.cfg.Metrics.Render(w, string(state), n, pending)
+}
+
+// llmPaths are the Ollama endpoints that load models into VRAM and therefore
+// take the LLM lock. Everything else passes through unlocked.
+var llmPaths = map[string]bool{
+	"/api/generate":        true,
+	"/api/chat":            true,
+	"/api/embed":           true,
+	"/api/embeddings":      true,
+	"/v1/chat/completions": true,
+	"/v1/completions":      true,
+	"/v1/embeddings":       true,
+}
+
+func isLLMRequest(r *http.Request) bool {
+	return r.Method == http.MethodPost && llmPaths[r.URL.Path]
+}
+
+// OllamaHandler serves the Ollama-facing listener.
+func (s *Server) OllamaHandler() http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/healthz":
+			s.writeHealthz(w)
+			return
+		case "/metrics":
+			s.writeMetrics(w)
+			return
+		}
+		if !isLLMRequest(r) {
+			s.ollamaProxy.ServeHTTP(w, r)
+			return
+		}
+
+		start := time.Now()
+		wctx, cancel := context.WithTimeout(r.Context(), s.cfg.LLMWaitTimeout)
+		err := s.cfg.Lock.AcquireLLM(wctx)
+		cancel()
+		s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds())
+		if err != nil {
+			if errors.Is(err, context.DeadlineExceeded) && r.Context().Err() == nil {
+				http.Error(w, "GPU busy: timed out waiting for the lock", http.StatusServiceUnavailable)
+			}
+			return
+		}
+		defer s.cfg.Lock.ReleaseLLM()
+		s.ollamaProxy.ServeHTTP(w, r)
+	})
+}
+
+// ComfyHandler serves the ComfyUI-facing listener.
+func (s *Server) ComfyHandler() http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/healthz" {
+			s.writeHealthz(w)
+			return
+		}
+		if r.Method == http.MethodPost && r.URL.Path == "/prompt" {
+			s.handlePrompt(w, r)
+			return
+		}
+		s.comfyProxy.ServeHTTP(w, r)
+	})
+}
+
+// captureWriter passes the response through unchanged while recording the
+// status code and the first captureLimit bytes of the body.
+type captureWriter struct {
+	http.ResponseWriter
+	status int
+	buf    bytes.Buffer
+}
+
+func (w *captureWriter) WriteHeader(code int) {
+	w.status = code
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *captureWriter) Write(p []byte) (int, error) {
+	if w.buf.Len() < captureLimit {
+		w.buf.Write(p)
+	}
+	return w.ResponseWriter.Write(p)
+}
+
+func (w *captureWriter) Flush() {
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+func (w *captureWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
+
+// handlePrompt implements the image job flow from the spec: acquire the
+// image lock, unload Ollama, forward to ComfyUI, then track the job in the
+// background and free VRAM before releasing the lock.
+func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
+	log := s.log.With("op", "image")
+
+	start := time.Now()
+	if err := s.cfg.Lock.AcquireImage(r.Context()); err != nil {
+		if errors.Is(err, context.DeadlineExceeded) {
+			http.Error(w, "GPU busy: timed out waiting for the lock", http.StatusServiceUnavailable)
+		}
+		return
+	}
+	s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds())
+	log.Info("image lock acquired")
+
+	uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout)
+	elapsed, uerr := s.cfg.Ollama.UnloadAll(uctx)
+	ucancel()
+	s.cfg.Metrics.ObserveUnload(elapsed.Seconds())
+	switch {
+	case r.Context().Err() != nil:
+		s.cfg.Lock.ReleaseImage()
+		return
+	case uerr != nil:
+		// Degrade, don't fail the user's request on a misbehaving neighbour.
+		log.Warn("ollama unload incomplete; continuing", "err", uerr)
+	default:
+		log.Info("ollama models unloaded", "seconds", elapsed.Seconds())
+	}
+
+	cw := &captureWriter{ResponseWriter: w, status: http.StatusOK}
+	s.comfyProxy.ServeHTTP(cw, r)
+
+	var accepted struct {
+		PromptID string `json:"prompt_id"`
+	}
+	if cw.status == http.StatusOK {
+		_ = json.Unmarshal(cw.buf.Bytes(), &accepted)
+	}
+	if accepted.PromptID == "" {
+		log.Info("prompt not accepted; releasing image lock", "status", cw.status)
+		s.cfg.Lock.ReleaseImage()
+		return
+	}
+	s.cfg.Metrics.IncImageJobs()
+	go s.finishImageJob(accepted.PromptID)
+}
+
+// finishImageJob runs after the prompt has been accepted by ComfyUI: wait
+// for the job to finish, free ComfyUI's models, release the lock, and
+// optionally warm the chat model.
+func (s *Server) finishImageJob(promptID string) {
+	log := s.log.With("prompt_id", promptID)
+	log.Info("image job running")
+
+	ctx, cancel := context.WithTimeout(context.Background(), s.cfg.JobTimeout)
+	err := s.cfg.Comfy.WaitJob(ctx, promptID)
+	cancel()
+	if err != nil {
+		log.Warn("image job did not complete cleanly; releasing lock anyway", "err", err)
+	} else {
+		log.Info("image job completed")
+	}
+
+	freeCtx, freeCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	if err := s.cfg.Comfy.Free(freeCtx); err != nil {
+		log.Warn("failed to free ComfyUI models", "err", err)
+	}
+	freeCancel()
+
+	s.cfg.Lock.ReleaseImage()
+	log.Info("image lock released")
+
+	if s.cfg.WarmModel != "" {
+		if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle {
+			wctx, wcancel := context.WithTimeout(context.Background(), 2*time.Minute)
+			if err := s.cfg.Ollama.Warm(wctx, s.cfg.WarmModel); err != nil {
+				log.Warn("warm model reload failed", "model", s.cfg.WarmModel, "err", err)
+			} else {
+				log.Info("warm model reloaded", "model", s.cfg.WarmModel)
+			}
+			wcancel()
+		}
+	}
+}

internal/proxy/proxy_test.go

@@ -0,0 +1,347 @@
+package proxy
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"gpu-turnstile/internal/comfy"
+	"gpu-turnstile/internal/lock"
+	"gpu-turnstile/internal/metrics"
+	"gpu-turnstile/internal/ollama"
+)
+
+// recorder collects upstream call events in order.
+type recorder struct {
+	mu     sync.Mutex
+	events []string
+}
+
+func (r *recorder) add(e string) {
+	r.mu.Lock()
+	r.events = append(r.events, e)
+	r.mu.Unlock()
+}
+
+func (r *recorder) index(e string) int {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	for i, ev := range r.events {
+		if ev == e {
+			return i
+		}
+	}
+	return -1
+}
+
+// fakes wires up fake Ollama and ComfyUI upstreams plus the gpulock server.
+type fakes struct {
+	rec       *recorder
+	ollama    *httptest.Server
+	comfy     *httptest.Server
+	server    *httptest.Server // Ollama-facing gpulock listener
+	comfySrv  *httptest.Server // ComfyUI-facing gpulock listener
+	freeCh    chan struct{}
+	chatCh    chan struct{}
+	historyMu sync.Mutex
+	history   string
+}
+
+func newFakes(t *testing.T) *fakes {
+	t.Helper()
+	f := &fakes{rec: &recorder{}, freeCh: make(chan struct{}), chatCh: make(chan struct{})}
+
+	var psCalls int
+	var psMu sync.Mutex
+
+	f.ollama = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/api/version":
+			io.WriteString(w, `{"version":"0.0.0"}`)
+		case "/api/ps":
+			psMu.Lock()
+			psCalls++
+			c := psCalls
+			psMu.Unlock()
+			f.rec.add("ps")
+			if c == 1 {
+				io.WriteString(w, `{"models":[{"name":"chat-model"}]}`)
+			} else {
+				io.WriteString(w, `{"models":[]}`)
+			}
+		case "/api/generate":
+			f.rec.add("unload")
+			io.WriteString(w, `{}`)
+		case "/api/chat":
+			f.rec.add("chat")
+			close(f.chatCh)
+			w.Header().Set("Content-Type", "application/x-ndjson")
+			io.WriteString(w, `{"done":true}`+"\n")
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	t.Cleanup(f.ollama.Close)
+
+	f.comfy = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case r.URL.Path == "/system_stats":
+			io.WriteString(w, `{}`)
+		case r.URL.Path == "/prompt":
+			f.rec.add("prompt")
+			w.Header().Set("Content-Type", "application/json")
+			io.WriteString(w, `{"prompt_id":"p1"}`)
+		case r.URL.Path == "/history/p1":
+			f.rec.add("history")
+			f.historyMu.Lock()
+			h := f.history
+			f.historyMu.Unlock()
+			w.Header().Set("Content-Type", "application/json")
+			if h == "" {
+				io.WriteString(w, `{}`)
+			} else {
+				fmt.Fprintf(w, `{"p1":{"status":{"completed":%s,"status_str":"success"}}}`, h)
+			}
+		case r.URL.Path == "/free":
+			f.rec.add("free")
+			close(f.freeCh)
+			io.WriteString(w, `{}`)
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	t.Cleanup(f.comfy.Close)
+
+	ollamaClient, err := ollama.New(f.ollama.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ollamaClient.PollInterval = 5 * time.Millisecond
+	comfyClient, err := comfy.New(f.comfy.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	comfyClient.PollInterval = 5 * time.Millisecond
+
+	srv, err := New(Config{
+		OllamaURL:      f.ollama.URL,
+		ComfyURL:       f.comfy.URL,
+		Lock:           lock.New(nil),
+		Ollama:         ollamaClient,
+		Comfy:          comfyClient,
+		Metrics:        metrics.New(),
+		LLMWaitTimeout: 2 * time.Second,
+		UnloadTimeout:  2 * time.Second,
+		JobTimeout:     2 * time.Second,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	f.server = httptest.NewServer(srv.OllamaHandler())
+	t.Cleanup(f.server.Close)
+	f.comfySrv = httptest.NewServer(srv.ComfyHandler())
+	t.Cleanup(f.comfySrv.Close)
+	return f
+}
+
+func (f *fakes) completeJob() {
+	f.historyMu.Lock()
+	f.history = "true"
+	f.historyMu.Unlock()
+}
+
+func TestImageJobSequenceAndLLMBlocked(t *testing.T) {
+	f := newFakes(t)
+
+	// Start the image job. The response returns as soon as ComfyUI has
+	// accepted the prompt; history polling and /free run in the background.
+	resp, err := http.Post(f.comfySrv.URL+"/prompt", "application/json", strings.NewReader(`{"workflow":{}}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, _ := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	if resp.StatusCode != 200 || !strings.Contains(string(body), `"prompt_id":"p1"`) {
+		t.Fatalf("prompt response = %d %s", resp.StatusCode, body)
+	}
+
+	// While the image job is running (history not yet complete), an LLM
+	// request must be held.
+	chatDone := make(chan struct{})
+	go func() {
+		defer close(chatDone)
+		resp, err := http.Post(f.server.URL+"/api/chat", "application/json", strings.NewReader(`{}`))
+		if err == nil {
+			io.Copy(io.Discard, resp.Body)
+			resp.Body.Close()
+		}
+	}()
+
+	select {
+	case <-f.chatCh:
+		t.Fatal("chat reached Ollama while image job still running")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	// Finish the job; the background poller then calls /free, releases the
+	// lock, and the chat request goes through.
+	f.completeJob()
+	select {
+	case <-chatDone:
+	case <-time.After(3 * time.Second):
+		t.Fatal("chat request never completed")
+	}
+
+	// Assert the call sequence for one image job.
+	for _, pair := range [][2]string{
+		{"ps", "unload"},
+		{"unload", "prompt"},
+		{"prompt", "history"},
+		{"history", "free"},
+		{"free", "chat"},
+	} {
+		a, b := f.rec.index(pair[0]), f.rec.index(pair[1])
+		if a < 0 || b < 0 || a >= b {
+			f.rec.mu.Lock()
+			t.Fatalf("expected %s before %s; events: %v", pair[0], pair[1], f.rec.events)
+		}
+	}
+}
+
+func TestPromptRejectedReleasesLock(t *testing.T) {
+	f := newFakes(t)
+	// Make ComfyUI reject the prompt.
+	f.comfy.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/prompt" {
+			f.rec.add("prompt")
+			http.Error(w, "bad workflow", http.StatusBadRequest)
+			return
+		}
+		http.NotFound(w, r)
+	})
+
+	resp, err := http.Post(f.comfySrv.URL+"/prompt", "application/json", strings.NewReader(`{}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, _ := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	if resp.StatusCode != 400 || !strings.Contains(string(body), "bad workflow") {
+		t.Fatalf("prompt response = %d %s", resp.StatusCode, body)
+	}
+
+	// The lock must already be free: a chat request goes straight through.
+	resp, err = http.Post(f.server.URL+"/api/chat", "application/json", strings.NewReader(`{}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != 200 {
+		t.Fatalf("chat status = %d", resp.StatusCode)
+	}
+}
+
+func TestStreamingNotBuffered(t *testing.T) {
+	gate := make(chan struct{})
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/x-ndjson")
+		fl := w.(http.Flusher)
+		io.WriteString(w, `{"chunk":1}`+"\n")
+		fl.Flush()
+		<-gate // hold chunk 2 back until the client has seen chunk 1
+		io.WriteString(w, `{"chunk":2}`+"\n")
+		fl.Flush()
+	}))
+	defer upstream.Close()
+
+	srv, err := New(Config{
+		OllamaURL:      upstream.URL,
+		ComfyURL:       "http://127.0.0.1:1",
+		Lock:           lock.New(nil),
+		Metrics:        metrics.New(),
+		LLMWaitTimeout: time.Second,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	front := httptest.NewServer(srv.OllamaHandler())
+	defer front.Close()
+
+	resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+
+	// If gpulock buffered the stream, this read would never complete while
+	// the gate is closed.
+	line, err := bufio.NewReader(resp.Body).ReadString('\n')
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(line, `"chunk":1`) {
+		t.Fatalf("first line = %q", line)
+	}
+	close(gate)
+	rest, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(string(rest), `"chunk":2`) {
+		t.Fatalf("rest = %q", rest)
+	}
+}
+
+func TestHealthzAndMetrics(t *testing.T) {
+	f := newFakes(t)
+
+	for _, base := range []string{f.server.URL, f.comfySrv.URL} {
+		resp, err := http.Get(base + "/healthz")
+		if err != nil {
+			t.Fatal(err)
+		}
+		body, _ := io.ReadAll(resp.Body)
+		resp.Body.Close()
+		if resp.StatusCode != 200 || !strings.Contains(string(body), `"state":"idle"`) {
+			t.Fatalf("healthz = %d %s", resp.StatusCode, body)
+		}
+	}
+
+	resp, err := http.Get(f.server.URL + "/metrics")
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, _ := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	text := string(body)
+	for _, want := range []string{
+		`gpulock_state{state="idle"} 1`,
+		"gpulock_llm_inflight 0",
+		"gpulock_image_jobs_total 0",
+		`gpulock_lock_wait_seconds_bucket{kind="llm",le="+Inf"} 0`,
+		"gpulock_unload_seconds_count{} 0",
+	} {
+		if !strings.Contains(text, want) {
+			t.Fatalf("metrics missing %q", want)
+		}
+	}
+}
+
+func TestPassThroughNoLock(t *testing.T) {
+	f := newFakes(t)
+	resp, err := http.Get(f.server.URL + "/api/version")
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, _ := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	if resp.StatusCode != 200 || !strings.Contains(string(body), "0.0.0") {
+		t.Fatalf("pass-through = %d %s", resp.StatusCode, body)
+	}
+}