internal/api/middleware/response_cache.go

package middleware

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
	"github.com/tidwall/gjson"
)

type cacheCaptureWriter struct {
	gin.ResponseWriter
	statusCode int
	body       bytes.Buffer
}

func (w *cacheCaptureWriter) WriteHeader(code int) {
	w.statusCode = code
	w.ResponseWriter.WriteHeader(code)
}

func (w *cacheCaptureWriter) Write(data []byte) (int, error) {
	if w.statusCode == 0 {
		w.statusCode = http.StatusOK
	}
	w.body.Write(data)
	return w.ResponseWriter.Write(data)
}

func (w *cacheCaptureWriter) WriteString(s string) (int, error) {
	if w.statusCode == 0 {
		w.statusCode = http.StatusOK
	}
	w.body.WriteString(s)
	return w.ResponseWriter.WriteString(s)
}

func ResponseCacheMiddleware(rc cache.ResponseCache, cfg config.ResponseCacheConfig) gin.HandlerFunc {
	return func(c *gin.Context) {
		if rc == nil || !cfg.Enabled || c.Request == nil || c.Request.Method != http.MethodPost {
			c.Next()
			return
		}
		if c.FullPath() == "" {
			c.Next()
			return
		}
		contentType := c.GetHeader("Content-Type")
		if contentType != "" && !strings.Contains(strings.ToLower(contentType), "application/json") {
			c.Next()
			return
		}

		raw, err := io.ReadAll(c.Request.Body)
		if err != nil {
			c.Next()
			return
		}
		c.Request.Body = io.NopCloser(bytes.NewReader(raw))
		if len(raw) == 0 {
			c.Next()
			return
		}

		model := gjson.GetBytes(raw, "model").String()
		isStreaming := gjson.GetBytes(raw, "stream").Bool()
		if !cache.ShouldCache(cfg, model, isStreaming, 0) {
			c.Next()
			return
		}

		modelKey := c.FullPath() + "|" + model
		key := cache.ResponseCacheKey(modelKey, raw)
		if entry, ok := rc.Get(c.Request.Context(), key); ok && entry != nil {
			for k, v := range entry.Headers {
				if k != "" && v != "" {
					c.Header(k, v)
				}
			}
			c.Header("X-Cache", "HIT")
			c.Status(entry.StatusCode)
			_, _ = c.Writer.Write(entry.Body)
			c.Abort()
			return
		}

		writer := &cacheCaptureWriter{ResponseWriter: c.Writer}
		c.Writer = writer
		c.Next()

		status := writer.statusCode
		if status == 0 {
			status = c.Writer.Status()
		}
		if status < 200 || status >= 300 {
			return
		}

		body := writer.body.Bytes()
		if !cache.ShouldCache(cfg, model, isStreaming, len(body)) {
			return
		}

		headers := make(map[string]string)
		for k, values := range c.Writer.Header() {
			if len(values) > 0 {
				headers[k] = values[0]
			}
		}

		ttl := time.Duration(cfg.TTLSeconds) * time.Second
		if ttl <= 0 {
			ttl = 300 * time.Second
		}

		_ = rc.Set(c.Request.Context(), key, &cache.CachedResponse{
			StatusCode: status,
			Headers:    headers,
			Body:       bytes.Clone(body),
			CachedAt:   time.Now(),
			Model:      model,
		}, ttl)
	}
}