vojo/apps/ai-bot/failover.go

package main

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"time"
)

// failover.go is the local-first routing decorator over the LLMClient seam: prefer
// the (free) home backend while it is reachable, fall back to the cloud transparently
// when it is not. The cascade keeps calling b.llm.Complete and never learns which leg
// answered — LLMResponse.Model carries that to billing/telemetry.
//
// Health is decided by an ACTIVE background probe plus a request-path circuit
// breaker, both feeding one atomic healthy flag:
//
//   - unhealthy → requests go straight to the cloud (no per-request timeout tax on a
//     dead tunnel);
//   - healthy → try local under its own short deadline (LOCAL_LLM_TIMEOUT_SECONDS —
//     a half-alive tunnel must fail fast enough to leave the shared RequestBudget
//     room for the cloud leg), and on a transport/5xx/timeout failure trip the
//     breaker and retry the SAME request on the cloud;
//   - recovery is probe-driven ONLY (anti-flapping): a tripped breaker needs
//     tripProbesToClose consecutive probe successes to close, so a tunnel that
//     accepts TCP but drops requests can't flap traffic between legs.
//
// A terminal 4xx from the local leg (except 429) does NOT fall back: that is a
// config/request error (a wrong LOCAL_LLM_MODEL tag, a malformed body) and the cloud
// would only mask it. 429 (a busy single-slot GPU) falls back like a timeout.
// Per-request fallback logs at DEBUG; only state transitions and a both-legs failure
// are louder.

const (
	// localProbeTimeout bounds one health probe. The probe is a bare GET /models —
	// answered from ollama's process memory, no inference — so anything slower than
	// this means the tunnel, not the model, is struggling.
	localProbeTimeout = 3 * time.Second
	// tripProbesToClose is how many consecutive probe successes re-close a breaker
	// tripped by a REQUEST failure. The probe can pass while real requests fail (a
	// cold model being loaded, a half-alive tunnel), so one success is not proof;
	// two spaced an interval apart also rate-limits the cold-load abort loop (an
	// aborted request cancels ollama's model load, keeping every next request cold).
	tripProbesToClose = 2
	// neverOnlineWarnAfter: consecutive probe failures — without the local leg EVER
	// having come online since boot — that trigger a one-time WARN. A wrong-valued
	// LOCAL_LLM_BASE_URL (say, missing the /v1 suffix) passes the boot presence
	// check but 404s forever, and per-request fallbacks log only at DEBUG — so
	// without this the enabled feature would be silently inoperative, serving 100%
	// of traffic from the paid cloud (the exact outcome the config fail-fast exists
	// to prevent).
	neverOnlineWarnAfter = 3
)

type failoverLLMClient struct {
	primary  LLMClient // local backend (free, preferred)
	fallback LLMClient // cloud (paid, always available)
	log      *slog.Logger

	legTimeout time.Duration // per-request deadline for the local leg
	probeURL   string        // <LOCAL_LLM_BASE_URL>/models
	probeKey   string        // optional bearer for the probe (usually empty)
	interval   time.Duration // background probe period
	httpc      *http.Client

	// Breaker state. healthy starts false: until the first probe passes, traffic is
	// served by the cloud, so a boot with the home box offline never eats the probe
	// timeout on user requests.
	mu          sync.Mutex
	healthy     bool
	consecOK    int  // consecutive probe successes so far
	needOK      int  // successes required to close (1 normally, tripProbesToClose after a trip)
	consecFail  int  // consecutive probe failures (the never-online WARN counter)
	everHealthy bool // the leg has been online at least once since boot
	warnedDown  bool // the one-time never-online WARN fired
}

// NewFailoverLLMClient wraps primary (local) and fallback (cloud) into one LLMClient
// and starts the background health probe, which lives until ctx is cancelled (the
// process signal context — same lifetime as the bot).
func NewFailoverLLMClient(ctx context.Context, primary, fallback LLMClient, cfg *Config, logger *slog.Logger) LLMClient {
	c := &failoverLLMClient{
		primary:    primary,
		fallback:   fallback,
		log:        logger,
		legTimeout: cfg.LocalLLMTimeout,
		probeURL:   cfg.LocalLLMBaseURL + "/models",
		probeKey:   cfg.LocalLLMAPIKey,
		interval:   cfg.LocalLLMHealthInterval,
		httpc:      &http.Client{},
		needOK:     1,
	}
	// One INFO at construction, so the absence of the follow-up "online" line is a
	// documented signal (see the never-online WARN in noteProbe).
	logger.Info("local llm failover enabled; cloud serves until the first successful probe", "probe_url", c.probeURL)
	go c.probeLoop(ctx)
	return c
}

func (c *failoverLLMClient) Complete(ctx context.Context, req LLMRequest) (*LLMResponse, error) {
	if !c.isHealthy() {
		c.log.DebugContext(ctx, "local llm unhealthy; serving from cloud")
		return c.fallback.Complete(ctx, req)
	}

	lctx, cancel := context.WithTimeout(ctx, c.legTimeout)
	resp, err := c.primary.Complete(lctx, req)
	cancel()
	if err == nil {
		if strings.TrimSpace(resp.Text) != "" {
			return resp, nil
		}
		// A 2xx with EMPTY content from the local leg (thinking ate the whole token
		// budget, a content quirk of the abliterated model). On the plain cloud path
		// an empty completion is kept — it was billed — but the local one is free, so
		// retry the turn on the cloud instead of shipping the ⚠️-react empty path.
		// The backend is up and answering: not a breaker trip.
		c.log.DebugContext(ctx, "local llm returned empty content; retrying on cloud")
		return c.fallback.Complete(ctx, req)
	}
	if ctx.Err() != nil {
		// The WHOLE request budget is gone, not just the local leg — a cloud call
		// would die the same way, so don't burn a paid attempt; and don't trip the
		// breaker over our own cancellation.
		return nil, err
	}
	var se *httpStatusError
	if errors.As(err, &se) && se.status >= 400 && se.status < 500 && se.status != http.StatusTooManyRequests {
		return nil, err // terminal 4xx: a local config error — fail loud, never mask with the cloud
	}

	c.trip()
	c.log.DebugContext(ctx, "local llm failed; falling back to cloud", "err", err)
	resp, ferr := c.fallback.Complete(ctx, req)
	if ferr != nil {
		// Both legs down — this is the only per-request failure loud enough for WARN
		// (the caller will still ERROR+react; this line preserves the local cause).
		c.log.WarnContext(ctx, "local llm and cloud fallback both failed", "local_err", err, "cloud_err", ferr)
		return nil, ferr
	}
	return resp, nil
}

func (c *failoverLLMClient) isHealthy() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.healthy
}

// trip opens the breaker after an in-request local failure. Recovery is probe-driven
// only: requests never re-try the local leg until the prober has re-closed it.
func (c *failoverLLMClient) trip() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.healthy {
		c.log.Info("local llm tripped offline by a request failure; serving from cloud until probes recover")
	}
	c.healthy = false
	c.consecOK = 0
	c.needOK = tripProbesToClose
}

// noteProbe folds one probe result into the breaker. A failed probe opens it (the
// tunnel is observably down — don't wait for a user request to find out) but does
// not raise needOK: plain down/up recovers on the next success, only a request-path
// trip demands the longer streak.
func (c *failoverLLMClient) noteProbe(ok bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if !ok {
		if c.healthy {
			c.log.Info("local llm probe failed; serving from cloud")
		}
		c.healthy = false
		c.consecOK = 0
		c.consecFail++
		// Never came online since boot → the feature is silently inoperative (every
		// request quietly bills the cloud, visible only at DEBUG). Say so ONCE, loud.
		if !c.everHealthy && !c.warnedDown && c.consecFail >= neverOnlineWarnAfter {
			c.warnedDown = true
			c.log.Warn("local llm has never come online since boot; ALL traffic is served by the paid cloud — check LOCAL_LLM_BASE_URL and the tunnel", "probe_url", c.probeURL)
		}
		return
	}
	c.consecFail = 0
	c.consecOK++
	if !c.healthy && c.consecOK >= c.needOK {
		c.healthy = true
		c.everHealthy = true
		c.needOK = 1
		c.log.Info("local llm online; serving from local backend")
	}
}

func (c *failoverLLMClient) probeLoop(ctx context.Context) {
	t := time.NewTicker(c.interval)
	defer t.Stop()
	for {
		c.noteProbe(c.probe(ctx)) // first probe runs immediately, not an interval late
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

// probe GETs the OpenAI models listing on the local base URL — cheap (no inference)
// and it proves the whole path (SSH tunnel → relay → server) end to end, which is
// exactly what dies when the home box goes away. It does NOT prove the model is
// resident in VRAM; keeping it warm is OLLAMA_KEEP_ALIVE=-1's job (a cold model
// answers slowly, which the request-path breaker handles).
func (c *failoverLLMClient) probe(ctx context.Context) bool {
	pctx, cancel := context.WithTimeout(ctx, localProbeTimeout)
	defer cancel()
	req, err := http.NewRequestWithContext(pctx, http.MethodGet, c.probeURL, nil)
	if err != nil {
		return false
	}
	if c.probeKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.probeKey)
	}
	resp, err := c.httpc.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10)) // drain for connection reuse
	return resp.StatusCode >= 200 && resp.StatusCode < 300
}