feat: 添加过载保护、背压和熔断机制

internal/server/handlers.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"strings"
 	"time"
@@ -10,10 +11,12 @@ import (
 	"github.com/edgeai/gateway/internal/adapter"
 	"github.com/edgeai/gateway/internal/auth"
 	"github.com/edgeai/gateway/internal/config"
+	ctxasm "github.com/edgeai/gateway/internal/context"
 	"github.com/edgeai/gateway/internal/handler"
 	"github.com/edgeai/gateway/internal/middleware"
 	"github.com/edgeai/gateway/internal/observability"
 	"github.com/edgeai/gateway/internal/router"
+	"github.com/edgeai/gateway/internal/session"
 	"github.com/edgeai/gateway/internal/task"
 	"github.com/edgeai/gateway/pkg/api"
 	"github.com/google/uuid"
@@ -70,9 +73,26 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
 		priority = int(task.PriorityNormal) // downgrade to P2 if not allowed P0
 	}
 
+	// Backpressure check: reject low-priority requests under high load
+	s.backpressure.Update(s.scheduler.RunningCount(), s.scheduler.QueueLength())
+	if !s.backpressure.ShouldAccept(priority) {
+		reason := s.backpressure.RejectReason(priority)
+		s.metrics.IncRequest("backpressure_rejected")
+		handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrQueueFull, reason, requestID))
+		return
+	}
+
+	// Circuit breaker check
+	if !s.breaker.AllowRequest() {
+		s.metrics.IncRequest("circuit_open")
+		handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrModelUnavailable, "circuit breaker open, please retry later", requestID))
+		return
+	}
+
 	// Create task
 	taskID := uuid.New().String()
 	tk := task.NewTask(taskID, requestID, identity.AppID, identity.TenantID, req.Model, task.TaskPriority(priority), req.Stream)
+	tk.SessionID = req.SessionID
 
 	// Submit to scheduler
 	if err := s.scheduler.Submit(tk); err != nil {
@@ -84,6 +104,7 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
 	s.metrics.SetQueueLength(s.scheduler.QueueLength())
 
 	// Wait for task to be dequeued
+	queueStart := time.Now()
 	ctx, cancel := context.WithTimeout(r.Context(), time.Duration(s.cfg.Timeouts.DefaultQueueMs)*time.Millisecond)
 	defer cancel()
 
@@ -91,19 +112,37 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
 	if err != nil {
 		s.scheduler.Complete(tk.ID)
 		s.metrics.IncRequest("queue_timeout")
+		tk.Transition(task.StateTimedOut)
 		handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrQueueTimeout, "queue timeout", requestID))
 		return
 	}
 
+	// Record queue time
+	dequeued.QueueMs = int(time.Since(queueStart) / time.Millisecond)
+
 	// Transition to RUNNING
 	dequeued.Transition(task.StateRunning)
 	s.metrics.SetRunningTasks(s.scheduler.RunningCount())
 
-	// Build adapter request
+	// Build adapter request — assemble context if session_id provided
+	messages := req.Messages
+	if req.SessionID != "" {
+		sess, err := s.sessions.Get(req.SessionID)
+		if err == nil && sess != nil {
+			// Load session history and assemble context
+			history := s.loadSessionHistory(sess)
+			if len(history) > 0 {
+				assembler := ctxasm.NewAssembler(&s.cfg.Context)
+				result := assembler.Assemble(history, req.Messages, target.ContextWindow, target.MaxOutputTokens, req.ContextPolicy)
+				messages = result.Messages
+			}
+		}
+	}
+
 	adapterReq := &adapter.ChatRequest{
 		RequestID:   requestID,
 		Model:       target.ActualModel,
-		Messages:    req.Messages,
+		Messages:    messages,
 		MaxTokens:   target.MaxOutputTokens,
 		Temperature: req.Temperature,
 		TopP:        req.TopP,
@@ -131,6 +170,7 @@ func (s *Server) handleStreaming(w http.ResponseWriter, r *http.Request, adapter
 	}
 
 	tk.Transition(task.StateStreaming)
+	inferenceStart := time.Now()
 
 	ch, err := adapterInst.ChatCompletionStream(r.Context(), req)
 	if err != nil {
@@ -146,11 +186,16 @@ func (s *Server) handleStreaming(w http.ResponseWriter, r *http.Request, adapter
 		s.logger.Error("streaming error", observability.F().Event("stream_error").TaskID(tk.ID).Reason(err.Error()))
 		tk.Transition(task.StateFailed)
 		s.metrics.IncTask("failed")
+		s.breaker.RecordError()
 	} else {
 		tk.Transition(task.StateCompleted)
 		s.metrics.IncTask("completed")
+		s.breaker.RecordSuccess()
 	}
 
+	tk.InferenceMs = int(time.Since(inferenceStart) / time.Millisecond)
+	tk.TotalMs = int(time.Since(tk.CreatedAt) / time.Millisecond)
+
 	s.metrics.AddTokens(inputTokens, outputTokens)
 	s.scheduler.Complete(tk.ID)
 	s.metrics.SetRunningTasks(s.scheduler.RunningCount())
@@ -158,18 +203,46 @@ func (s *Server) handleStreaming(w http.ResponseWriter, r *http.Request, adapter
 	s.metrics.IncRequest("stream_ok")
 }
 
-func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adapterInst adapter.ModelAdapter, req *adapter.ChatRequest, tk *task.Task, requestID string, _ *router.ModelTarget, logicalModel string) {
+func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adapterInst adapter.ModelAdapter, req *adapter.ChatRequest, tk *task.Task, requestID string, target *router.ModelTarget, logicalModel string) {
 	ctx, cancel := context.WithTimeout(r.Context(), time.Duration(s.cfg.Timeouts.DefaultInferenceMs)*time.Millisecond)
 	defer cancel()
 
+	inferenceStart := time.Now()
+
 	resp, err := adapterInst.ChatCompletion(ctx, req)
 	if err != nil {
+		// Try overload fallback strategies
+		fallbackResp, fbErr := s.tryOverloadFallback(r.Context(), req, target, logicalModel)
+		if fbErr == nil && fallbackResp != nil {
+			tk.Degraded = true
+			tk.InferenceMs = int(time.Since(inferenceStart) / time.Millisecond)
+			tk.TotalMs = int(time.Since(tk.CreatedAt) / time.Millisecond)
+			s.scheduler.Complete(tk.ID)
+			tk.Transition(task.StateCompleted)
+			s.metrics.IncTask("completed")
+			s.metrics.IncRequest("ok")
+			s.metrics.IncRequest("overload_fallback")
+			s.metrics.AddTokens(fallbackResp.InputTokens, fallbackResp.OutputTokens)
+			s.breaker.RecordSuccess()
+			s.scheduler.Complete(tk.ID)
+			s.metrics.SetRunningTasks(s.scheduler.RunningCount())
+			s.metrics.SetQueueLength(s.scheduler.QueueLength())
+
+			chatResp := handler.BuildChatResponse(requestID, tk.ID, logicalModel, fallbackResp)
+			chatResp.Degraded = true
+			chatResp.Timing = s.buildTiming(tk)
+			handler.WriteJSON(w, http.StatusOK, chatResp)
+			return
+		}
+
 		s.scheduler.Complete(tk.ID)
 		tk.Transition(task.StateFailed)
 		s.metrics.IncTask("failed")
 		s.metrics.IncRequest("error")
+		s.breaker.RecordError()
 
 		if strings.Contains(err.Error(), "timeout") || ctx.Err() != nil {
+			tk.Transition(task.StateTimedOut)
 			handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInferenceTimeout, "inference timeout", requestID))
 		} else {
 			handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrModelUnavailable, err.Error(), requestID))
@@ -177,18 +250,61 @@ func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adap
 		return
 	}
 
+	tk.InferenceMs = int(time.Since(inferenceStart) / time.Millisecond)
+	tk.TotalMs = int(time.Since(tk.CreatedAt) / time.Millisecond)
+
 	tk.Transition(task.StateCompleted)
 	s.metrics.IncTask("completed")
 	s.metrics.IncRequest("ok")
 	s.metrics.AddTokens(resp.InputTokens, resp.OutputTokens)
+	s.breaker.RecordSuccess()
 	s.scheduler.Complete(tk.ID)
 	s.metrics.SetRunningTasks(s.scheduler.RunningCount())
 	s.metrics.SetQueueLength(s.scheduler.QueueLength())
 
 	chatResp := handler.BuildChatResponse(requestID, tk.ID, logicalModel, resp)
+	chatResp.Timing = s.buildTiming(tk)
 	handler.WriteJSON(w, http.StatusOK, chatResp)
 }
 
+// tryOverloadFallback attempts fallback strategies when the primary model fails.
+func (s *Server) tryOverloadFallback(ctx context.Context, req *adapter.ChatRequest, target *router.ModelTarget, logicalModel string) (*adapter.ChatResponse, error) {
+	exclude := map[string]bool{target.Endpoint: true}
+
+	result := s.overload.Resolve(logicalModel, exclude)
+	if result.Rejected {
+		return nil, fmt.Errorf("overload: %s", result.Reason)
+	}
+
+	fallbackAdapter, err := s.registry.Get(result.Target.Provider)
+	if err != nil {
+		return nil, fmt.Errorf("overload: adapter not found: %w", err)
+	}
+
+	fallbackReq := &adapter.ChatRequest{
+		RequestID:   req.RequestID,
+		Model:       result.Target.ActualModel,
+		Messages:    req.Messages,
+		MaxTokens:   result.Target.MaxOutputTokens,
+		Temperature: req.Temperature,
+		TopP:        req.TopP,
+		Stream:      false,
+		CancelCh:    req.CancelCh,
+	}
+
+	s.logger.Info("overload fallback",
+		observability.F().
+			Event("overload_fallback").
+			Set("strategy", router.StrategyName(result.Strategy)).
+			Set("from", logicalModel).
+			Set("to", result.Target.LogicalModel))
+
+	fbCtx, fbCancel := context.WithTimeout(ctx, time.Duration(s.cfg.Timeouts.DefaultInferenceMs)*time.Millisecond)
+	defer fbCancel()
+
+	return fallbackAdapter.ChatCompletion(fbCtx, fallbackReq)
+}
+
 func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodGet {
 		handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
@@ -282,3 +398,87 @@ func (s *Server) handleSessionByID(w http.ResponseWriter, r *http.Request) {
 		handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
 	}
 }
+
+// loadSessionHistory returns the message history from a session.
+func (s *Server) loadSessionHistory(sess *session.Session) []api.Message {
+	if sess == nil {
+		return nil
+	}
+	return sess.Messages
+}
+
+// buildTiming constructs Timing metadata from a task.
+func (s *Server) buildTiming(tk *task.Task) *api.Timing {
+	return &api.Timing{
+		QueueMs:      tk.QueueMs,
+		FirstTokenMs: tk.FirstTokenMs,
+		InferenceMs:  tk.InferenceMs,
+		TotalMs:      tk.TotalMs,
+	}
+}
+
+func (s *Server) handleTasks(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet {
+		handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
+		return
+	}
+	handler.WriteJSON(w, http.StatusOK, map[string]any{"tasks": []any{}})
+}
+
+func (s *Server) handleTaskByID(w http.ResponseWriter, r *http.Request) {
+	requestID := middleware.GetRequestID(r.Context())
+	taskID := strings.TrimPrefix(r.URL.Path, "/v1/tasks/")
+
+	switch r.Method {
+	case http.MethodGet:
+		tk, ok := s.scheduler.GetTask(taskID)
+		if !ok {
+			handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "task not found", requestID))
+			return
+		}
+		resp := map[string]any{
+			"task_id":       tk.ID,
+			"request_id":    tk.RequestID,
+			"session_id":    tk.SessionID,
+			"state":         string(tk.GetState()),
+			"priority":      int(tk.Priority),
+			"logical_model": tk.LogicalModel,
+			"stream":        tk.Stream,
+			"created_at":    tk.CreatedAt.Format(time.RFC3339),
+			"degraded":      tk.Degraded,
+		}
+		if tk.StartedAt != nil {
+			resp["started_at"] = tk.StartedAt.Format(time.RFC3339)
+		}
+		if tk.CompletedAt != nil {
+			resp["completed_at"] = tk.CompletedAt.Format(time.RFC3339)
+		}
+		if tk.CancelReason != "" {
+			resp["cancel_reason"] = tk.CancelReason
+		}
+		if tk.ErrorMessage != "" {
+			resp["error_message"] = tk.ErrorMessage
+		}
+		resp["timing"] = s.buildTiming(tk)
+		handler.WriteJSON(w, http.StatusOK, resp)
+
+	case http.MethodDelete:
+		tk, ok := s.scheduler.GetTask(taskID)
+		if !ok {
+			handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "task not found", requestID))
+			return
+		}
+		if tk.IsTerminal() {
+			handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "task already in terminal state", requestID))
+			return
+		}
+		tk.Cancel("client requested cancellation")
+		s.scheduler.Complete(tk.ID)
+		s.metrics.SetRunningTasks(s.scheduler.RunningCount())
+		s.metrics.SetQueueLength(s.scheduler.QueueLength())
+		handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "cancelled"})
+
+	default:
+		handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
+	}
+}

internal/server/server.go

@@ -22,15 +22,18 @@ import (
 
 // Server is the main HTTP server for the AI gateway.
 type Server struct {
-	cfg       *config.Config
-	logger    *observability.Logger
-	metrics   *observability.Metrics
-	HTTPSrv   *http.Server
-	auth      *auth.Authenticator
-	registry  *adapter.Registry
-	modelMap  *router.LogicalModelMapping
-	scheduler *scheduler.Scheduler
-	sessions  *session.Store
+	cfg          *config.Config
+	logger       *observability.Logger
+	metrics      *observability.Metrics
+	HTTPSrv      *http.Server
+	auth         *auth.Authenticator
+	registry     *adapter.Registry
+	modelMap     *router.LogicalModelMapping
+	scheduler    *scheduler.Scheduler
+	sessions     *session.Store
+	breaker      *scheduler.CircuitBreaker
+	backpressure *scheduler.BackpressureManager
+	overload     *router.OverloadResolver
 }
 
 // New creates a new Server instance with all components wired.
@@ -60,6 +63,9 @@ func New(cfg *config.Config, logger *observability.Logger) (*Server, error) {
 	// Initialize logical model mapping
 	modelMap := router.NewLogicalModelMapping(cfg)
 
+	// Initialize overload resolver
+	overloadResolver := router.NewOverloadResolver(&cfg.Routing, modelMap)
+
 	// Register adapters for each unique endpoint
 	registered := make(map[string]bool)
 	for _, mc := range cfg.Models {
@@ -78,18 +84,39 @@ func New(cfg *config.Config, logger *observability.Logger) (*Server, error) {
 	// Initialize scheduler
 	sched := scheduler.NewScheduler(&cfg.Scheduler, logger)
 
+	// Initialize circuit breaker
+	breaker := scheduler.NewCircuitBreaker(
+		cfg.CircuitBreaker.ErrorRateThreshold,
+		cfg.CircuitBreaker.MinRequests,
+		cfg.CircuitBreaker.WindowSeconds,
+		cfg.CircuitBreaker.OpenDurationSeconds,
+		cfg.CircuitBreaker.HalfOpenMaxRequests,
+	)
+
+	// Initialize backpressure manager
+	bp := scheduler.NewBackpressureManager(
+		cfg.Scheduler.MaxRunningTasks,
+		cfg.Scheduler.MaxQueuedTasks,
+		cfg.Backpressure.Level1Threshold,
+		cfg.Backpressure.Level2Threshold,
+		cfg.Backpressure.Level3Threshold,
+	)
+
 	// Initialize metrics
 	metrics := observability.NewMetrics()
 
 	s := &Server{
-		cfg:       cfg,
-		logger:    logger,
-		metrics:   metrics,
-		auth:      authenticator,
-		registry:  registry,
-		modelMap:  modelMap,
-		scheduler: sched,
-		sessions:  sessionStore,
+		cfg:          cfg,
+		logger:       logger,
+		metrics:      metrics,
+		auth:         authenticator,
+		registry:     registry,
+		modelMap:     modelMap,
+		scheduler:    sched,
+		sessions:     sessionStore,
+		breaker:      breaker,
+		backpressure: bp,
+		overload:     overloadResolver,
 	}
 
 	mux := http.NewServeMux()
@@ -128,6 +155,10 @@ func (s *Server) registerRoutes(mux *http.ServeMux) {
 	// Session management
 	mux.HandleFunc("/v1/sessions", s.handleSessions)
 	mux.HandleFunc("/v1/sessions/", s.handleSessionByID)
+
+	// Task management
+	mux.HandleFunc("/v1/tasks", s.handleTasks)
+	mux.HandleFunc("/v1/tasks/", s.handleTaskByID)
 }
 
 // Authenticator returns the authenticator instance (for testing/management).