rony-chat-bot/internal/server/health.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"
)

// HealthResponse is the shape returned by /api/health.
//
// Status is "healthy" when every component is up, "degraded" when non-critical
// components are down, "unhealthy" when the bot cannot serve. The HTTP code
// is 200 for healthy, 503 for unhealthy. Degraded returns 200 (the bot can
// still answer, just without RAG or with no LLM).
type HealthResponse struct {
	Status     string                     `json:"status"`
	Version    string                     `json:"version"`
	CheckedAt  string                     `json:"checked_at"`
	Components map[string]ComponentHealth `json:"components"`
}

// ComponentHealth reports the state of one dependency. Latency is a
// human-readable duration ("12ms"); Error is set only when Status="down".
type ComponentHealth struct {
	Status  string         `json:"status"`            // "up" | "down"
	Latency string         `json:"latency,omitempty"`  // e.g. "12ms"
	Error   string         `json:"error,omitempty"`
	Details map[string]any `json:"details,omitempty"`
}

// healthCheckTimeout caps each component probe. The whole /api/health call
// finishes in roughly this duration even if a dependency is hung.
const healthCheckTimeout = 2 * time.Second

func (h *Handlers) Health(w http.ResponseWriter, r *http.Request) {
	// `?deep=true` adds an FTS5 row count to the store probe; same latency
	// budget. Cheap enough that we don't bother splitting into two endpoints.
	deep := r.URL.Query().Get("deep") == "true"

	// Pad the context so even a slow store doesn't kill the request before
	// the LLM probe finishes.
	ctx, cancel := context.WithTimeout(r.Context(), healthCheckTimeout+500*time.Millisecond)
	defer cancel()

	// Run both probes in parallel so a slow LLM doesn't delay the store
	// check (and vice versa).
	type probe struct {
		name string
		fn   func(context.Context) ComponentHealth
	}
	results := make(map[string]ComponentHealth, 2)
	var mu sync.Mutex
	var wg sync.WaitGroup
	for _, p := range []probe{
		{"llm", h.checkLLM},
		{"store", func(ctx context.Context) ComponentHealth { return h.checkStore(ctx, deep) }},
	} {
		wg.Add(1)
		p := p
		go func() {
			defer wg.Done()
			r := p.fn(ctx)
			mu.Lock()
			results[p.name] = r
			mu.Unlock()
		}()
	}
	wg.Wait()

	resp := HealthResponse{
		Status:     "healthy",
		Version:    h.version,
		CheckedAt:  time.Now().UTC().Format(time.RFC3339),
		Components: results,
	}

	// Decide overall status. LLM is critical; store is non-critical
	// (RAG degrades to "I don't have that information" when it can't be read).
	httpStatus := http.StatusOK
	if results["llm"].Status != "up" {
		resp.Status = "unhealthy"
		httpStatus = http.StatusServiceUnavailable
	} else if results["store"].Status != "up" {
		resp.Status = "degraded"
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(httpStatus)
	_ = json.NewEncoder(w).Encode(resp)
}

// checkLLM pings the provider's health endpoint with a short timeout.
func (h *Handlers) checkLLM(parent context.Context) ComponentHealth {
	p := h.cfg.DefaultProvider()
	healthURL := deriveHealthURL(p.Endpoint, p.Type)
	if healthURL == "" {
		return ComponentHealth{
			Status:  "down",
			Error:   "no health endpoint known for provider type " + p.Type,
			Details: map[string]any{"provider": p.Type, "model": p.Model},
		}
	}

	ctx, cancel := context.WithTimeout(parent, healthCheckTimeout)
	defer cancel()

	start := time.Now()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, healthURL, nil)
	if err != nil {
		return ComponentHealth{Status: "down", Error: err.Error(), Latency: time.Since(start).String()}
	}
	resp, err := http.DefaultClient.Do(req)
	latency := time.Since(start)
	if err != nil {
		return ComponentHealth{Status: "down", Error: err.Error(), Latency: latency.String()}
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
		return ComponentHealth{
			Status:  "up",
			Latency: latency.String(),
			Details: map[string]any{"provider": p.Type, "model": p.Model, "url": healthURL},
		}
	}
	return ComponentHealth{
		Status:  "down",
		Latency: latency.String(),
		Error:   fmt.Sprintf("HTTP %d", resp.StatusCode),
		Details: map[string]any{"url": healthURL},
	}
}

// deriveHealthURL maps a provider's chat-completions URL to its health URL.
// Returns "" when the provider has no usable health probe.
func deriveHealthURL(endpoint, providerType string) string {
	if endpoint == "" {
		return ""
	}
	// Strip trailing /v1, /chat/completions, etc., to get the base.
	base := endpoint
	for _, suffix := range []string{"/v1/chat/completions", "/chat/completions", "/v1"} {
		if strings.HasSuffix(base, suffix) {
			base = strings.TrimSuffix(base, suffix)
			break
		}
	}
	switch providerType {
	case "llamacpp", "ollama":
		// Both expose /health at the root.
		return base + "/health"
	case "openai":
		// OpenAI doesn't have /health, but /models is auth-light and stable.
		return base + "/models"
	default:
		// anthropic and unknown: no cheap probe without a real request.
		return ""
	}
}

// checkStore pings the SQLite store with a trivial query. In deep mode it
// also counts indexed chunks to confirm the FTS5 schema is alive.
func (h *Handlers) checkStore(parent context.Context, deep bool) ComponentHealth {
	if h.store == nil {
		return ComponentHealth{Status: "down", Error: "store not configured"}
	}
	ctx, cancel := context.WithTimeout(parent, healthCheckTimeout)
	defer cancel()

	start := time.Now()
	var n int
	err := h.store.DB().QueryRowContext(ctx, "SELECT 1").Scan(&n)
	latency := time.Since(start)
	if err != nil {
		return ComponentHealth{Status: "down", Error: err.Error(), Latency: latency.String()}
	}

	details := map[string]any{}
	if deep {
		_ = h.store.DB().QueryRowContext(ctx,
			"SELECT count(*) FROM portfolio_chunks").Scan(&n)
		details["chunks"] = n
	}
	return ComponentHealth{Status: "up", Latency: latency.String(), Details: details}
}