Pin compose example to v0.1.2

Backoff now covers all dial-phase errors (refused, timeout, DNS), TLS
handshake failures, and 5xx responses with replayable bodies; streamed
POSTs are never replayed to avoid duplicate work. First retry of an
episode logs at WARN, subsequent attempts at INFO.

LOGLEVEL (LOG_LEVEL kept as alias) now defaults to warn: startup logs
version plus every setting; INFO adds one line per incoming request and
per response with status/duration, ANSI-colored in text mode (bypasses
slog's escaping so colors render in docker compose logs); NO_COLOR or
LOG_FORMAT=json disables colors.

README.md

@@ -41,7 +41,7 @@ startup.
 | `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 |
+| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` works as an alias |
 | `LOG_FORMAT` | `text` | `json` for structured JSON logs |
 | `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
 | `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
@@ -49,6 +49,8 @@ startup.
 | `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
 | `WARM_TIMEOUT` | `2m` | Warm-model reload after an image job |
 | `SHUTDOWN_TIMEOUT` | `10s` | Graceful shutdown on SIGINT/SIGTERM |
+| `BACKOFF_INITIAL` | `1s` | First retry wait when an upstream refuses a connection |
+| `BACKOFF_MAX` | `60s` | Cap for the exponential retry backoff |
 | `PROMPT_CAPTURE_LIMIT` | `65536` | Bytes of the `/prompt` response buffered to find `prompt_id` (pass-through is unaffected) |
 
 ## Observability
@@ -58,6 +60,12 @@ startup.
   `gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`,
   `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
   (histogram, `kind="llm|image"`), `gpu_turnstile_unload_seconds`.
+- Logs: startup logs the version and every setting (visible even at the
+  default `warn` level). With `LOGLEVEL=info` or `debug`, every request
+  logs a `-->` incoming line and a `<--` response line with status and
+  duration — ANSI-colored (cyan incoming; green/yellow/red by status
+  class) in text mode, which renders in `docker compose logs` on Windows
+  Terminal. Set `NO_COLOR` to disable colors.
 
 ## Build and run
 

SPEC.md

@@ -112,13 +112,15 @@ load time. Off by default.
 | `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 |
+| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` is accepted as an alias |
 | `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
 | `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
 | `PROBE_TIMEOUT` | `5s` | startup probe of both upstreams |
 | `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
 | `WARM_TIMEOUT` | `2m` | warm-model reload after an image job |
 | `SHUTDOWN_TIMEOUT` | `10s` | graceful shutdown on SIGINT/SIGTERM |
+| `BACKOFF_INITIAL` | `1s` | first retry wait when an upstream refuses a connection |
+| `BACKOFF_MAX` | `60s` | cap for the exponential retry backoff |
 | `PROMPT_CAPTURE_LIMIT` | `65536` | bytes of the `/prompt` response buffered to find `prompt_id` (pass-through is unaffected) |
 
 Startup fails fast on unparsable values. Both upstreams are probed once at
@@ -134,7 +136,13 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
   `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
   (histogram, label `kind="llm|image"`), `gpu_turnstile_unload_seconds`.
 - Structured logs (`log/slog`, JSON when `LOG_FORMAT=json`), one line per
-  state transition and per image job phase with `prompt_id`.
+  state transition and per image job phase with `prompt_id`. Startup logs
+  the version and every setting (visible even at the default `warn`
+  level). With `LOGLEVEL=info` or `debug`, every request logs a `-->`
+  incoming line and a `<--` response line with status and duration —
+  ANSI-colored (cyan incoming; green/yellow/red by status class) in text
+  mode, which renders in `docker compose logs` on Windows Terminal. Set
+  `NO_COLOR` to disable colors.
 
 ## Edge cases to handle
 
@@ -146,6 +154,14 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
   `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.
+- Upstream unreachable while proxying (connection refused, dial timeout,
+  DNS failure, TLS handshake error): retry with exponential backoff —
+  `BACKOFF_INITIAL`, doubling per attempt, capped at `BACKOFF_MAX` — until
+  the upstream answers or the client disconnects. These are safe to retry:
+  the request never reached the upstream application. 5xx responses are
+  retried the same way, but only when the request body can be replayed
+  (GETs, or bodies with `GetBody`); streamed POSTs are never replayed to
+  avoid duplicate work such as a double-enqueued ComfyUI prompt.
 - `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.

cmd/gpu-turnstile/main.go

@@ -40,6 +40,8 @@ type config struct {
 	freeTimeout         time.Duration
 	warmTimeout         time.Duration
 	shutdownTimeout     time.Duration
+	backoffInitial      time.Duration
+	backoffMax          time.Duration
 	promptCaptureLimit  int64
 
 	warmModel string
@@ -76,9 +78,11 @@ func loadConfig(getenv func(string) string) (config, error) {
 		freeTimeout:         30 * time.Second,
 		warmTimeout:         2 * time.Minute,
 		shutdownTimeout:     10 * time.Second,
+		backoffInitial:      time.Second,
+		backoffMax:          time.Minute,
 		promptCaptureLimit:  64 * 1024,
 
-		logLevel: slog.LevelInfo,
+		logLevel: slog.LevelWarn,
 	}
 	for _, e := range []struct {
 		name string
@@ -107,6 +111,8 @@ func loadConfig(getenv func(string) string) (config, error) {
 		{"FREE_TIMEOUT", &cfg.freeTimeout},
 		{"WARM_TIMEOUT", &cfg.warmTimeout},
 		{"SHUTDOWN_TIMEOUT", &cfg.shutdownTimeout},
+		{"BACKOFF_INITIAL", &cfg.backoffInitial},
+		{"BACKOFF_MAX", &cfg.backoffMax},
 	} {
 		if err := envDuration(getenv, e.name, e.dst); err != nil {
 			return cfg, err
@@ -119,10 +125,15 @@ func loadConfig(getenv func(string) string) (config, error) {
 		}
 		cfg.promptCaptureLimit = n
 	}
-	if v := getenv("LOG_LEVEL"); v != "" {
+	// LOGLEVEL is the canonical spelling; LOG_LEVEL is kept as an alias.
+	logLevelValue := getenv("LOGLEVEL")
+	if logLevelValue == "" {
+		logLevelValue = getenv("LOG_LEVEL")
+	}
+	if logLevelValue != "" {
 		var level slog.Level
-		if err := level.UnmarshalText([]byte(v)); err != nil {
+		if err := level.UnmarshalText([]byte(logLevelValue)); err != nil {
-			return cfg, fmt.Errorf("LOG_LEVEL: %w", err)
+			return cfg, fmt.Errorf("LOGLEVEL: %w", err)
 		}
 		cfg.logLevel = level
 	}
@@ -151,12 +162,29 @@ func main() {
 	log := slog.New(handler)
 	slog.SetDefault(log)
 
-	log.Info("starting gpu-turnstile",
+	// 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(context.Background(), slog.LevelWarn, "starting gpu-turnstile",
 		"version", version,
 		"listen_ollama", cfg.listenOllama,
 		"listen_comfy", cfg.listenComfy,
 		"ollama_url", cfg.ollamaURL,
 		"comfy_url", cfg.comfyURL,
+		"unload_timeout", cfg.unloadTimeout,
+		"job_timeout", cfg.jobTimeout,
+		"llm_wait_timeout", cfg.llmWaitTimeout,
+		"unload_poll_interval", cfg.unloadPollInterval,
+		"history_poll_interval", cfg.historyPollInterval,
+		"probe_timeout", cfg.probeTimeout,
+		"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,
+		"log_level", cfg.logLevel,
+		"log_format", map[bool]string{true: "json", false: "text"}[cfg.logJSON],
 	)
 
 	ollamaClient, err := ollama.New(cfg.ollamaURL, log)
@@ -178,6 +206,7 @@ func main() {
 		Comfy:               comfyClient,
 		Metrics:             metrics.New(),
 		Log:                 log,
+		LogColor:            !cfg.logJSON && os.Getenv("NO_COLOR") == "",
 		LLMWaitTimeout:      cfg.llmWaitTimeout,
 		UnloadTimeout:       cfg.unloadTimeout,
 		JobTimeout:          cfg.jobTimeout,
@@ -185,6 +214,8 @@ func main() {
 		HistoryPollInterval: cfg.historyPollInterval,
 		FreeTimeout:         cfg.freeTimeout,
 		WarmTimeout:         cfg.warmTimeout,
+		BackoffInitial:      cfg.backoffInitial,
+		BackoffMax:          cfg.backoffMax,
 		PromptCaptureLimit:  cfg.promptCaptureLimit,
 		WarmModel:           cfg.warmModel,
 	})

compose.yaml.example

@@ -1,14 +1,20 @@
 # Example deployment for gpu-turnstile. Copy to compose.yaml and adjust.
 #
 # gpu-turnstile listens on the ports the services normally use; the actual
-# Ollama and ComfyUI instances run one port higher on the workstation.
+# Ollama and ComfyUI instances run one port higher (11435 / 8189) and must
+# bind 0.0.0.0 so the container can reach them (OLLAMA_HOST=0.0.0.0:11435,
+# ComfyUI --listen 0.0.0.0 --port 8189).
 services:
   gpu-turnstile:
-    image: git.rambossek.at/public/gpu-turnstile:v0.1.1
+    image: git.rambossek.at/public/gpu-turnstile:v0.1.2
     restart: unless-stopped
     environment:
-      OLLAMA_URL: http://<workstation-ip>:11435
+      # Services on the Docker host itself:
-      COMFY_URL: http://<workstation-ip>:8189
+      OLLAMA_URL: http://host.docker.internal:11435
+      COMFY_URL: http://host.docker.internal:8189
+      # Services on another machine: use its LAN IP instead, e.g.
+      # OLLAMA_URL: http://192.168.1.10:11435
+      # COMFY_URL: http://192.168.1.10:8189
       # UNLOAD_TIMEOUT: 60s
       # JOB_TIMEOUT: 15m
       # LLM_WAIT_TIMEOUT: 10m

internal/proxy/proxy.go

@@ -4,15 +4,22 @@
 package proxy
 
 import (
+	"bufio"
 	"bytes"
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"log/slog"
+	"net"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"os"
+	"strconv"
+	"strings"
 	"time"
 
 	"gpu-turnstile/internal/comfy"
@@ -37,10 +44,21 @@ type Config struct {
 	Metrics *metrics.Metrics
 	Log     *slog.Logger
 
+	// LogColor enables ANSI colors in per-request log lines. Ignored when
+	// the log level is above INFO (request lines are not emitted at all).
+	LogColor bool
+
 	LLMWaitTimeout time.Duration
 	UnloadTimeout  time.Duration
 	JobTimeout     time.Duration
 
+	// BackoffInitial and BackoffMax control the exponential retry backoff
+	// when an upstream refuses a connection: the wait doubles from
+	// BackoffInitial up to BackoffMax between attempts. Zero selects the
+	// defaults (1s / 60s).
+	BackoffInitial time.Duration
+	BackoffMax     time.Duration
+
 	// UnloadPollInterval and HistoryPollInterval override the clients'
 	// /api/ps and /history poll intervals when > 0.
 	UnloadPollInterval  time.Duration
@@ -57,11 +75,13 @@ type Config struct {
 
 // Server serves both gpu-turnstile listeners.
 type Server struct {
-	cfg          Config
+	cfg            Config
-	log          *slog.Logger
+	log            *slog.Logger
-	freeTimeout  time.Duration
+	freeTimeout    time.Duration
-	warmTimeout  time.Duration
+	warmTimeout    time.Duration
-	captureLimit int64
+	captureLimit   int64
+	backoffInitial time.Duration
+	backoffMax     time.Duration
 
 	ollamaProxy *httputil.ReverseProxy
 	comfyProxy  *httputil.ReverseProxy
@@ -99,19 +119,142 @@ func New(cfg Config) (*Server, error) {
 	if captureLimit <= 0 {
 		captureLimit = defaultCaptureLimit
 	}
+	backoffInitial := cfg.BackoffInitial
+	if backoffInitial <= 0 {
+		backoffInitial = time.Second
+	}
+	backoffMax := cfg.BackoffMax
+	if backoffMax <= 0 {
+		backoffMax = time.Minute
+	}
+	retry := &retryTransport{
+		base:    http.DefaultTransport,
+		initial: backoffInitial,
+		max:     backoffMax,
+		log:     log,
+	}
 	return &Server{
-		cfg:          cfg,
+		cfg:            cfg,
-		log:          log,
+		log:            log,
-		freeTimeout:  freeTimeout,
+		freeTimeout:    freeTimeout,
-		warmTimeout:  warmTimeout,
+		warmTimeout:    warmTimeout,
-		captureLimit: captureLimit,
+		captureLimit:   captureLimit,
-		ollamaProxy:  newReverseProxy(ollamaURL, log.With("upstream", "ollama")),
+		backoffInitial: backoffInitial,
-		comfyProxy:   newReverseProxy(comfyURL, log.With("upstream", "comfy")),
+		backoffMax:     backoffMax,
+		ollamaProxy:    newReverseProxy(ollamaURL, retry, log.With("upstream", "ollama")),
+		comfyProxy:     newReverseProxy(comfyURL, retry, log.With("upstream", "comfy")),
 	}, nil
 }
 
-func newReverseProxy(target *url.URL, log *slog.Logger) *httputil.ReverseProxy {
+// retryTransport retries requests whose failure means the upstream never
+// saw them — any dial-phase error (connection refused, dial timeout, DNS
+// failure), TLS handshake errors — plus 5xx responses when the request
+// body can be replayed (GETs and requests with GetBody set). The wait
+// doubles from initial up to max between attempts. The loop runs until
+// the request succeeds, fails in a non-retryable way, or the client's
+// context is cancelled.
+type retryTransport struct {
+	base    http.RoundTripper
+	initial time.Duration
+	max     time.Duration
+	log     *slog.Logger
+}
+
+// shouldRetry reports whether a RoundTrip error means the request never
+// reached the upstream application and is therefore safe to send again.
+func shouldRetry(err error) bool {
+	// Dial-phase failures: refused, timeout, unreachable, DNS (wrapped).
+	var opErr *net.OpError
+	if errors.As(err, &opErr) && opErr.Op == "dial" {
+		return true
+	}
+	var dnsErr *net.DNSError
+	if errors.As(err, &dnsErr) {
+		return true
+	}
+	// TLS handshake failures: the HTTP request was never written.
+	var recordErr tls.RecordHeaderError
+	if errors.As(err, &recordErr) {
+		return true
+	}
+	var certErr *tls.CertificateVerificationError
+	if errors.As(err, &certErr) {
+		return true
+	}
+	var alertErr tls.AlertError
+	if errors.As(err, &alertErr) {
+		return true
+	}
+	return false
+}
+
+// replayable reports whether the request body can be sent again. Bodies
+// streamed from the client (GetBody == nil) cannot, so 5xx responses to
+// POSTs are not retried: the upstream may have partially processed them,
+// and re-sending could duplicate work (e.g. a second ComfyUI prompt).
+func replayable(req *http.Request) bool {
+	return req.Body == nil || req.Body == http.NoBody || req.GetBody != nil
+}
+
+func (t *retryTransport) logAt(level slog.Level, msg string, args ...any) {
+	if t.log != nil && t.log.Enabled(context.Background(), level) {
+		t.log.Log(context.Background(), level, msg, args...)
+	}
+}
+
+func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	wait := t.initial
+	attempt := 0
+	for {
+		resp, err := t.base.RoundTrip(req)
+		switch {
+		case err != nil && !shouldRetry(err):
+			return nil, err
+		case err == nil && (resp.StatusCode < 500 || !replayable(req)):
+			return resp, nil
+		}
+
+		// Retryable failure: a transport error, or a 5xx response.
+		var reason string
+		if err != nil {
+			reason = err.Error()
+		} else {
+			reason = resp.Status
+			io.Copy(io.Discard, resp.Body)
+			resp.Body.Close()
+			if req.GetBody != nil {
+				if body, berr := req.GetBody(); berr == nil {
+					req.Body = body
+				}
+			}
+		}
+		attempt++
+		// One WARN per outage episode; subsequent attempts at INFO.
+		level := slog.LevelInfo
+		if attempt == 1 {
+			level = slog.LevelWarn
+		}
+		t.logAt(level, "upstream unavailable; retrying with backoff",
+			"path", req.URL.Path, "reason", reason, "attempt", attempt, "retry_in", wait)
+
+		select {
+		case <-req.Context().Done():
+			if err != nil {
+				return nil, err
+			}
+			return nil, req.Context().Err()
+		case <-time.After(wait):
+		}
+		wait *= 2
+		if wait > t.max {
+			wait = t.max
+		}
+	}
+}
+
+func newReverseProxy(target *url.URL, transport http.RoundTripper, log *slog.Logger) *httputil.ReverseProxy {
 	return &httputil.ReverseProxy{
+		Transport: transport,
 		Rewrite: func(pr *httputil.ProxyRequest) {
 			pr.SetURL(target)
 			pr.SetXForwarded()
@@ -142,6 +285,90 @@ func (s *Server) writeMetrics(w http.ResponseWriter) {
 	s.cfg.Metrics.Render(w, string(state), n, pending)
 }
 
+// ANSI colors for per-request log lines.
+const (
+	ansiReset  = "\x1b[0m"
+	ansiCyan   = "\x1b[36m"
+	ansiGreen  = "\x1b[32m"
+	ansiYellow = "\x1b[33m"
+	ansiRed    = "\x1b[31m"
+)
+
+// statusRecorder remembers the response status while passing everything
+// through, including streaming flushes and websocket hijacks.
+type statusRecorder struct {
+	http.ResponseWriter
+	status int
+}
+
+func (r *statusRecorder) WriteHeader(code int) {
+	r.status = code
+	r.ResponseWriter.WriteHeader(code)
+}
+
+func (r *statusRecorder) Flush() {
+	if f, ok := r.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+func (r *statusRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	h, ok := r.ResponseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, errors.New("response writer does not support hijacking")
+	}
+	return h.Hijack()
+}
+
+func (r *statusRecorder) Unwrap() http.ResponseWriter { return r.ResponseWriter }
+
+// reqLine emits one request log line. With color enabled slog cannot be
+// used: its text handler escapes the ANSI sequences, so the line is
+// written to stderr directly in the same key=value shape. Without color
+// it is a plain slog INFO line.
+func (s *Server) reqLine(log *slog.Logger, code, line string, attrs ...any) {
+	if !s.cfg.LogColor {
+		log.Info(line, attrs...)
+		return
+	}
+	var sb strings.Builder
+	sb.WriteString("time=" + time.Now().Format("2006-01-02T15:04:05.000Z07:00") + " level=INFO ")
+	sb.WriteString(code + line + ansiReset)
+	for i := 0; i+1 < len(attrs); i += 2 {
+		fmt.Fprintf(&sb, " %v=%v", attrs[i], attrs[i+1])
+	}
+	sb.WriteByte('\n')
+	os.Stderr.WriteString(sb.String())
+}
+
+// logRequests logs one line per incoming request and one per completed
+// response at INFO level, colored when enabled: cyan "-->" for incoming,
+// green/yellow/red "<--" for responses by status class. At log levels
+// above INFO it is a pass-through.
+func (s *Server) logRequests(listener string, next http.Handler) http.Handler {
+	log := s.log.With("listener", listener)
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !log.Enabled(r.Context(), slog.LevelInfo) {
+			next.ServeHTTP(w, r)
+			return
+		}
+		start := time.Now()
+		s.reqLine(log, ansiCyan, "--> "+r.Method+" "+r.URL.RequestURI(),
+			"listener", listener, "remote", r.RemoteAddr)
+		rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
+		next.ServeHTTP(rec, r)
+		code := ansiGreen
+		switch {
+		case rec.status >= 500:
+			code = ansiRed
+		case rec.status >= 400:
+			code = ansiYellow
+		}
+		s.reqLine(log, code, "<-- "+strconv.Itoa(rec.status)+" "+r.Method+" "+r.URL.RequestURI(),
+			"listener", listener, "ms", time.Since(start).Milliseconds())
+	})
+}
+
 // 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{
@@ -160,7 +387,7 @@ func isLLMRequest(r *http.Request) bool {
 
 // OllamaHandler serves the Ollama-facing listener.
 func (s *Server) OllamaHandler() http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return s.logRequests("ollama", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case "/healthz":
 			s.writeHealthz(w)
@@ -187,12 +414,12 @@ func (s *Server) OllamaHandler() http.Handler {
 		}
 		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) {
+	return s.logRequests("comfy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path == "/healthz" {
 			s.writeHealthz(w)
 			return
@@ -202,7 +429,7 @@ func (s *Server) ComfyHandler() http.Handler {
 			return
 		}
 		s.comfyProxy.ServeHTTP(w, r)
-	})
+	}))
 }
 
 // captureWriter passes the response through unchanged while recording the

internal/proxy/retry_test.go

@@ -0,0 +1,207 @@
+package proxy
+
+import (
+	"context"
+	"crypto/tls"
+	"io"
+	"net"
+	"net/http"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+)
+
+// stubTransport fails with ECONNREFUSED for the first fails requests, then
+// returns a 200 response.
+type stubTransport struct {
+	fails int
+	calls int
+}
+
+func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	s.calls++
+	if s.calls <= s.fails {
+		return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
+	}
+	return &http.Response{
+		StatusCode: 200,
+		Body:       io.NopCloser(strings.NewReader("ok")),
+		Header:     make(http.Header),
+	}, nil
+}
+
+func TestRetryTransportBackoff(t *testing.T) {
+	st := &stubTransport{fails: 3}
+	rt := &retryTransport{
+		base:    st,
+		initial: 10 * time.Millisecond,
+		max:     25 * time.Millisecond,
+	}
+
+	start := time.Now()
+	req, _ := http.NewRequest(http.MethodGet, "http://upstream/api/version", nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if st.calls != 4 {
+		t.Fatalf("calls = %d, want 4", st.calls)
+	}
+	// Waits: 10ms + 20ms + 25ms (capped) = 55ms minimum.
+	elapsed := time.Since(start)
+	if elapsed < 50*time.Millisecond {
+		t.Fatalf("elapsed = %v, want >= ~55ms of backoff", elapsed)
+	}
+	if elapsed > 5*time.Second {
+		t.Fatalf("elapsed = %v, suspiciously long", elapsed)
+	}
+}
+
+func TestRetryTransportNonRefusedErrorNotRetried(t *testing.T) {
+	rt := &retryTransport{
+		base:    &stubTransport{fails: 0},
+		initial: time.Millisecond,
+		max:     time.Millisecond,
+	}
+	req, _ := http.NewRequest(http.MethodGet, "http://upstream/", nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil || resp.StatusCode != 200 {
+		t.Fatalf("resp=%v err=%v", resp, err)
+	}
+}
+
+type alwaysRefused struct{ calls int }
+
+func (a *alwaysRefused) RoundTrip(*http.Request) (*http.Response, error) {
+	a.calls++
+	return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
+}
+
+func TestRetryTransportContextCancel(t *testing.T) {
+	ar := &alwaysRefused{}
+	rt := &retryTransport{base: ar, initial: time.Second, max: time.Second}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://upstream/", nil)
+	start := time.Now()
+	_, err := rt.RoundTrip(req)
+	if err == nil {
+		t.Fatal("expected error after context cancellation")
+	}
+	if time.Since(start) > 2*time.Second {
+		t.Fatal("retry loop did not stop on context cancellation")
+	}
+}
+
+func TestRetryViaReverseProxy(t *testing.T) {
+	// Point the proxy at a port nothing listens on; the request should be
+	// retried (not instantly 502) until the client context ends.
+	srv, err := New(Config{
+		OllamaURL:      "http://127.0.0.1:1",
+		ComfyURL:       "http://127.0.0.1:1",
+		Lock:           nil,
+		Metrics:        nil,
+		BackoffInitial: 10 * time.Millisecond,
+		BackoffMax:     20 * time.Millisecond,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	_ = srv // construction must not panic with minimal config
+
+	rt := &retryTransport{base: http.DefaultTransport, initial: 10 * time.Millisecond, max: 20 * time.Millisecond}
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:1/", nil)
+	_, err = rt.RoundTrip(req)
+	if err == nil || !shouldRetry(err) {
+		t.Fatalf("err = %v, want connection refused", err)
+	}
+}
+
+// flakyStatus returns 500 for the first fails requests, then 200.
+type flakyStatus struct {
+	fails int
+	calls int
+}
+
+func (s *flakyStatus) RoundTrip(req *http.Request) (*http.Response, error) {
+	s.calls++
+	code := 200
+	if s.calls <= s.fails {
+		code = 500
+	}
+	return &http.Response{
+		StatusCode: code,
+		Status:     http.StatusText(code),
+		Body:       io.NopCloser(strings.NewReader("")),
+		Header:     make(http.Header),
+	}, nil
+}
+
+func TestRetryTransport5xxGet(t *testing.T) {
+	fs := &flakyStatus{fails: 2}
+	rt := &retryTransport{base: fs, initial: time.Millisecond, max: 2 * time.Millisecond}
+
+	req, _ := http.NewRequest(http.MethodGet, "http://upstream/api/version", nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != 200 {
+		t.Fatalf("status = %d, want 200", resp.StatusCode)
+	}
+	if fs.calls != 3 {
+		t.Fatalf("calls = %d, want 3", fs.calls)
+	}
+}
+
+func TestRetryTransport5xxPostNotRetried(t *testing.T) {
+	fs := &flakyStatus{fails: 10}
+	rt := &retryTransport{base: fs, initial: time.Millisecond, max: time.Millisecond}
+
+	// A streamed body (no GetBody) must not be replayed after a 500.
+	req, _ := http.NewRequest(http.MethodPost, "http://upstream/prompt", io.NopCloser(strings.NewReader("{}")))
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != 500 {
+		t.Fatalf("status = %d, want 500", resp.StatusCode)
+	}
+	if fs.calls != 1 {
+		t.Fatalf("calls = %d, want 1 (no retry for streamed POST)", fs.calls)
+	}
+}
+
+func TestShouldRetryClassification(t *testing.T) {
+	cases := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"dial refused", &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}, true},
+		{"dial timeout", &net.OpError{Op: "dial", Net: "tcp", Err: timeoutErr{}}, true},
+		{"dns", &net.DNSError{Err: "no such host", IsNotFound: true}, true},
+		{"tls record", tls.RecordHeaderError{Msg: "bad"}, true},
+		{"tls alert", tls.AlertError(42), true},
+		{"read error mid-request", &net.OpError{Op: "read", Net: "tcp", Err: syscall.ECONNRESET}, false},
+		{"plain error", io.EOF, false},
+	}
+	for _, c := range cases {
+		if got := shouldRetry(c.err); got != c.want {
+			t.Errorf("%s: shouldRetry = %v, want %v", c.name, got, c.want)
+		}
+	}
+}
+
+type timeoutErr struct{}
+
+func (timeoutErr) Error() string   { return "i/o timeout" }
+func (timeoutErr) Timeout() bool   { return true }
+func (timeoutErr) Temporary() bool { return true }