Make poll intervals and operational timeouts configurable

New env vars: UNLOAD_POLL_INTERVAL, HISTORY_POLL_INTERVAL, PROBE_TIMEOUT,
FREE_TIMEOUT, WARM_TIMEOUT, SHUTDOWN_TIMEOUT, PROMPT_CAPTURE_LIMIT.
Defaults unchanged; invalid values fail fast at startup.

internal/proxy/proxy.go

@@ -21,10 +21,10 @@ import (
 	"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
+// defaultCaptureLimit 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 defaultCaptureLimit = 64 * 1024
 
 // Config wires a Server.
 type Config struct {
@@ -40,13 +40,28 @@ type Config struct {
 	LLMWaitTimeout time.Duration
 	UnloadTimeout  time.Duration
 	JobTimeout     time.Duration
-	WarmModel      string
+
+	// UnloadPollInterval and HistoryPollInterval override the clients'
+	// /api/ps and /history poll intervals when > 0.
+	UnloadPollInterval  time.Duration
+	HistoryPollInterval time.Duration
+	// FreeTimeout and WarmTimeout bound the /free call and the warm-model
+	// reload; zero selects the defaults.
+	FreeTimeout time.Duration
+	WarmTimeout time.Duration
+	// PromptCaptureLimit overrides defaultCaptureLimit when > 0.
+	PromptCaptureLimit int64
+
+	WarmModel string
 }
 
 // Server serves both gpu-turnstile listeners.
 type Server struct {
-	cfg Config
-	log *slog.Logger
+	cfg          Config
+	log          *slog.Logger
+	freeTimeout  time.Duration
+	warmTimeout  time.Duration
+	captureLimit int64
 
 	ollamaProxy *httputil.ReverseProxy
 	comfyProxy  *httputil.ReverseProxy
@@ -66,11 +81,32 @@ func New(cfg Config) (*Server, error) {
 	if log == nil {
 		log = slog.Default()
 	}
+	if cfg.UnloadPollInterval > 0 {
+		cfg.Ollama.PollInterval = cfg.UnloadPollInterval
+	}
+	if cfg.HistoryPollInterval > 0 {
+		cfg.Comfy.PollInterval = cfg.HistoryPollInterval
+	}
+	freeTimeout := cfg.FreeTimeout
+	if freeTimeout <= 0 {
+		freeTimeout = 30 * time.Second
+	}
+	warmTimeout := cfg.WarmTimeout
+	if warmTimeout <= 0 {
+		warmTimeout = 2 * time.Minute
+	}
+	captureLimit := cfg.PromptCaptureLimit
+	if captureLimit <= 0 {
+		captureLimit = defaultCaptureLimit
+	}
 	return &Server{
-		cfg:         cfg,
-		log:         log,
-		ollamaProxy: newReverseProxy(ollamaURL, log.With("upstream", "ollama")),
-		comfyProxy:  newReverseProxy(comfyURL, log.With("upstream", "comfy")),
+		cfg:          cfg,
+		log:          log,
+		freeTimeout:  freeTimeout,
+		warmTimeout:  warmTimeout,
+		captureLimit: captureLimit,
+		ollamaProxy:  newReverseProxy(ollamaURL, log.With("upstream", "ollama")),
+		comfyProxy:   newReverseProxy(comfyURL, log.With("upstream", "comfy")),
 	}, nil
 }
 
@@ -170,11 +206,12 @@ func (s *Server) ComfyHandler() http.Handler {
 }
 
 // captureWriter passes the response through unchanged while recording the
-// status code and the first captureLimit bytes of the body.
+// status code and the first limit bytes of the body.
 type captureWriter struct {
 	http.ResponseWriter
 	status int
 	buf    bytes.Buffer
+	limit  int64
 }
 
 func (w *captureWriter) WriteHeader(code int) {
@@ -183,7 +220,7 @@ func (w *captureWriter) WriteHeader(code int) {
 }
 
 func (w *captureWriter) Write(p []byte) (int, error) {
-	if w.buf.Len() < captureLimit {
+	if int64(w.buf.Len()) < w.limit {
 		w.buf.Write(p)
 	}
 	return w.ResponseWriter.Write(p)
@@ -228,7 +265,7 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
 		log.Info("ollama models unloaded", "seconds", elapsed.Seconds())
 	}
 
-	cw := &captureWriter{ResponseWriter: w, status: http.StatusOK}
+	cw := &captureWriter{ResponseWriter: w, status: http.StatusOK, limit: s.captureLimit}
 	s.comfyProxy.ServeHTTP(cw, r)
 
 	var accepted struct {
@@ -262,7 +299,7 @@ func (s *Server) finishImageJob(promptID string) {
 		log.Info("image job completed")
 	}
 
-	freeCtx, freeCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	freeCtx, freeCancel := context.WithTimeout(context.Background(), s.freeTimeout)
 	if err := s.cfg.Comfy.Free(freeCtx); err != nil {
 		log.Warn("failed to free ComfyUI models", "err", err)
 	}
@@ -273,7 +310,7 @@ func (s *Server) finishImageJob(promptID string) {
 
 	if s.cfg.WarmModel != "" {
 		if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle {
-			wctx, wcancel := context.WithTimeout(context.Background(), 2*time.Minute)
+			wctx, wcancel := context.WithTimeout(context.Background(), s.warmTimeout)
 			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 {