Add rate limiting

internal/server/health.go

@@ -0,0 +1,87 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"time"
+)
+
+// HealthStatus represents the health check response.
+type HealthStatus struct {
+	Status    string            `json:"status"`
+	Timestamp int64             `json:"timestamp"`
+	Checks    map[string]string `json:"checks,omitempty"`
+}
+
+// handleHealth returns a basic health check endpoint.
+// This is suitable for Kubernetes liveness probes.
+func (s *GatewayServer) handleHealth(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet {
+		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	status := HealthStatus{
+		Status:    "healthy",
+		Timestamp: time.Now().Unix(),
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
+	_ = json.NewEncoder(w).Encode(status)
+}
+
+// handleReady returns a readiness check that verifies dependencies.
+// This is suitable for Kubernetes readiness probes and load balancer health checks.
+func (s *GatewayServer) handleReady(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet {
+		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	checks := make(map[string]string)
+	allHealthy := true
+
+	// Check conversation store connectivity
+	ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
+	defer cancel()
+
+	// Test conversation store by attempting a simple operation
+	testID := "health_check_test"
+	_, err := s.convs.Get(testID)
+	if err != nil {
+		checks["conversation_store"] = "unhealthy: " + err.Error()
+		allHealthy = false
+	} else {
+		checks["conversation_store"] = "healthy"
+	}
+
+	// Check if at least one provider is configured
+	models := s.registry.Models()
+	if len(models) == 0 {
+		checks["providers"] = "unhealthy: no providers configured"
+		allHealthy = false
+	} else {
+		checks["providers"] = "healthy"
+	}
+
+	_ = ctx // Use context if needed
+
+	status := HealthStatus{
+		Timestamp: time.Now().Unix(),
+		Checks:    checks,
+	}
+
+	if allHealthy {
+		status.Status = "ready"
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+	} else {
+		status.Status = "not_ready"
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusServiceUnavailable)
+	}
+
+	_ = json.NewEncoder(w).Encode(status)
+}

internal/server/health_test.go

@@ -0,0 +1,150 @@
+package server
+
+import (
+	"encoding/json"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+)
+
+func TestHealthEndpoint(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	registry := newMockRegistry()
+	convStore := newMockConversationStore()
+
+	server := New(registry, convStore, logger)
+
+	tests := []struct {
+		name           string
+		method         string
+		expectedStatus int
+	}{
+		{
+			name:           "GET returns healthy status",
+			method:         http.MethodGet,
+			expectedStatus: http.StatusOK,
+		},
+		{
+			name:           "POST returns method not allowed",
+			method:         http.MethodPost,
+			expectedStatus: http.StatusMethodNotAllowed,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := httptest.NewRequest(tt.method, "/health", nil)
+			w := httptest.NewRecorder()
+
+			server.handleHealth(w, req)
+
+			if w.Code != tt.expectedStatus {
+				t.Errorf("expected status %d, got %d", tt.expectedStatus, w.Code)
+			}
+
+			if tt.expectedStatus == http.StatusOK {
+				var status HealthStatus
+				if err := json.NewDecoder(w.Body).Decode(&status); err != nil {
+					t.Fatalf("failed to decode response: %v", err)
+				}
+
+				if status.Status != "healthy" {
+					t.Errorf("expected status 'healthy', got %q", status.Status)
+				}
+
+				if status.Timestamp == 0 {
+					t.Error("expected non-zero timestamp")
+				}
+			}
+		})
+	}
+}
+
+func TestReadyEndpoint(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+
+	tests := []struct {
+		name           string
+		setupRegistry  func() *mockRegistry
+		convStore      *mockConversationStore
+		expectedStatus int
+		expectedReady  bool
+	}{
+		{
+			name: "returns ready when all checks pass",
+			setupRegistry: func() *mockRegistry {
+				reg := newMockRegistry()
+				reg.addModel("test-model", "test-provider")
+				return reg
+			},
+			convStore:      newMockConversationStore(),
+			expectedStatus: http.StatusOK,
+			expectedReady:  true,
+		},
+		{
+			name: "returns not ready when no providers configured",
+			setupRegistry: func() *mockRegistry {
+				return newMockRegistry()
+			},
+			convStore:      newMockConversationStore(),
+			expectedStatus: http.StatusServiceUnavailable,
+			expectedReady:  false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := New(tt.setupRegistry(), tt.convStore, logger)
+
+			req := httptest.NewRequest(http.MethodGet, "/ready", nil)
+			w := httptest.NewRecorder()
+
+			server.handleReady(w, req)
+
+			if w.Code != tt.expectedStatus {
+				t.Errorf("expected status %d, got %d", tt.expectedStatus, w.Code)
+			}
+
+			var status HealthStatus
+			if err := json.NewDecoder(w.Body).Decode(&status); err != nil {
+				t.Fatalf("failed to decode response: %v", err)
+			}
+
+			if tt.expectedReady {
+				if status.Status != "ready" {
+					t.Errorf("expected status 'ready', got %q", status.Status)
+				}
+			} else {
+				if status.Status != "not_ready" {
+					t.Errorf("expected status 'not_ready', got %q", status.Status)
+				}
+			}
+
+			if status.Timestamp == 0 {
+				t.Error("expected non-zero timestamp")
+			}
+
+			if status.Checks == nil {
+				t.Error("expected checks map to be present")
+			}
+		})
+	}
+}
+
+func TestReadyEndpointMethodNotAllowed(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	registry := newMockRegistry()
+	convStore := newMockConversationStore()
+	server := New(registry, convStore, logger)
+
+	req := httptest.NewRequest(http.MethodPost, "/ready", nil)
+	w := httptest.NewRecorder()
+
+	server.handleReady(w, req)
+
+	if w.Code != http.StatusMethodNotAllowed {
+		t.Errorf("expected status %d, got %d", http.StatusMethodNotAllowed, w.Code)
+	}
+}

internal/server/server.go

@@ -44,6 +44,8 @@ func New(registry ProviderRegistry, convs conversation.Store, logger *slog.Logge
 func (s *GatewayServer) RegisterRoutes(mux *http.ServeMux) {
 	mux.HandleFunc("/v1/responses", s.handleResponses)
 	mux.HandleFunc("/v1/models", s.handleModels)
+	mux.HandleFunc("/health", s.handleHealth)
+	mux.HandleFunc("/ready", s.handleReady)
 }
 
 func (s *GatewayServer) handleModels(w http.ResponseWriter, r *http.Request) {