new-api/middleware/relay_capture_integration_test.go

package middleware

import (
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func setupCaptureTestRouter() *gin.Engine {
	gin.SetMode(gin.TestMode)
	return gin.New()
}

const captureTestDir = "./data/relay-capture"

type relayCaptureTestRecord struct {
	RequestID  string `json:"request_id"`
	CapturedAt string `json:"captured_at"`
	DurationMs int64  `json:"duration_ms"`
	UserID     int64  `json:"user_id"`
	Request    struct {
		Method  string              `json:"method"`
		Path    string              `json:"path"`
		Query   string              `json:"query"`
		Headers map[string][]string `json:"headers"`
		Body    string              `json:"body"`
	} `json:"request"`
	Response struct {
		StatusCode int                 `json:"status_code"`
		Headers    map[string][]string `json:"headers"`
		Body       string              `json:"body"`
		BodyChunks []string            `json:"body_chunks"`
	} `json:"response"`
	CaptureErrors []string `json:"capture_errors"`
}

func readCaptureRecord(t *testing.T, requestID string) relayCaptureTestRecord {
	t.Helper()
	today := time.Now().Format("2006-01-02")
	path := filepath.Join(captureTestDir, today, requestID+".json")
	content, err := os.ReadFile(path)
	require.NoError(t, err)
	require.NotContains(t, string(content), "=== REQUEST")
	require.NotContains(t, string(content), "=== RESPONSE")
	require.NotContains(t, string(content), "=== END")

	var record relayCaptureTestRecord
	require.NoError(t, common.Unmarshal(content, &record))
	return record
}

func readLatestCaptureRecord(t *testing.T) relayCaptureTestRecord {
	t.Helper()
	today := time.Now().Format("2006-01-02")
	dateDir := filepath.Join(captureTestDir, today)
	files, err := os.ReadDir(dateDir)
	require.NoError(t, err)
	require.NotEmpty(t, files, "should have capture json files")
	latest := filepath.Join(dateDir, files[len(files)-1].Name())
	require.True(t, strings.HasSuffix(latest, ".json"))

	content, err := os.ReadFile(latest)
	require.NoError(t, err)
	var record relayCaptureTestRecord
	require.NoError(t, common.Unmarshal(content, &record))
	return record
}

func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	requestID := fmt.Sprintf("test-non-stream-%d", time.Now().UnixNano())
	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		c.Set(common.RequestIdKey, requestID)
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"id":      "chatcmpl-test",
			"object":  "chat.completion",
			"choices": []gin.H{{"message": gin.H{"role": "assistant", "content": "Hello!"}}},
		})
	})

	body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}`
	req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body)))
	req.Header.Set("Authorization", "Bearer sk-test")
	req.Header.Set("Content-Type", "application/json")

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)

	record := readCaptureRecord(t, requestID)

	assert.Equal(t, requestID, record.RequestID)
	assert.NotEmpty(t, record.CapturedAt)
	assert.GreaterOrEqual(t, record.DurationMs, int64(0))
	assert.Equal(t, int64(42), record.UserID)
	assert.Equal(t, "POST", record.Request.Method)
	assert.Equal(t, "/v1/chat/completions", record.Request.Path)
	assert.Equal(t, body, record.Request.Body)
	assert.Equal(t, []string{"[REDACTED]"}, record.Request.Headers["Authorization"])
	assert.Equal(t, []string{"application/json"}, record.Request.Headers["Content-Type"])
	assert.Equal(t, 200, record.Response.StatusCode)
	assert.Contains(t, record.Response.Headers["Content-Type"][0], "application/json")
	assert.Contains(t, record.Response.Body, "chatcmpl-test")
	assert.NotEmpty(t, record.Response.BodyChunks)
	assert.Empty(t, record.CaptureErrors)
}

func TestRelayCapture_Streaming_EnabledUser(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	requestID := fmt.Sprintf("test-stream-%d", time.Now().UnixNano())
	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		c.Set(common.RequestIdKey, requestID)
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.Header("Content-Type", "text/event-stream")
		c.Header("Cache-Control", "no-cache")
		c.Header("Connection", "keep-alive")

		flusher := c.Writer.(http.Flusher)

		chunks := []string{
			"data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n",
			"data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"world!\"}}]}\n\n",
			"data: [DONE]\n\n",
		}

		for _, chunk := range chunks {
			_, _ = c.Writer.Write([]byte(chunk))
			flusher.Flush()
			time.Sleep(5 * time.Millisecond)
		}
	})

	body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"stream":true}`
	req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body)))
	req.Header.Set("Authorization", "Bearer sk-test")
	req.Header.Set("Content-Type", "application/json")

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)

	record := readCaptureRecord(t, requestID)

	assert.Equal(t, requestID, record.RequestID)
	assert.Equal(t, "POST", record.Request.Method)
	assert.Equal(t, "/v1/chat/completions", record.Request.Path)
	assert.Equal(t, body, record.Request.Body)
	assert.Equal(t, 200, record.Response.StatusCode)
	assert.Contains(t, record.Response.Body, "Hello ")
	assert.Contains(t, record.Response.Body, "world!")
	assert.Contains(t, record.Response.Body, "[DONE]")
	require.Len(t, record.Response.BodyChunks, 3)
	assert.Contains(t, record.Response.BodyChunks[0], "Hello ")
	assert.Contains(t, record.Response.BodyChunks[1], "world!")
	assert.Contains(t, record.Response.BodyChunks[2], "[DONE]")
	assert.Empty(t, record.CaptureErrors)
}

func TestRelayCapture_RedactsSensitiveRequestHeaders(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	requestID := fmt.Sprintf("test-redaction-%d", time.Now().UnixNano())
	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		c.Set(common.RequestIdKey, requestID)
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.JSON(200, gin.H{"result": "ok"})
	})

	req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`)))
	req.Header.Set("Authorization", "Bearer sk-test")
	req.Header.Set("X-Api-Key", "x-secret")
	req.Header.Set("Cookie", "session=secret")

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)

	today := time.Now().Format("2006-01-02")
	path := filepath.Join(captureTestDir, today, requestID+".json")
	content, err := os.ReadFile(path)
	require.NoError(t, err)
	contentStr := string(content)
	assert.Contains(t, contentStr, "[REDACTED]")
	assert.NotContains(t, contentStr, "Bearer sk-test")
	assert.NotContains(t, contentStr, "x-secret")
	assert.NotContains(t, contentStr, "session=secret")
}

func TestRelayCapture_DisabledUser(t *testing.T) {
	clearCaptureCache()

	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 99)
		c.Set(common.RequestIdKey, fmt.Sprintf("test-disabled-%d", time.Now().UnixNano()))
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.JSON(200, gin.H{"result": "ok"})
	})

	req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`)))
	req.Header.Set("Authorization", "Bearer sk-test")

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)

	today := time.Now().Format("2006-01-02")
	dateDir := filepath.Join(captureTestDir, today)
	if _, err := os.Stat(dateDir); err == nil {
		files, _ := os.ReadDir(dateDir)
		for _, f := range files {
			assert.NotContains(t, f.Name(), "test-disabled", "disabled user should not generate capture files")
		}
	}
}

func TestRelayCapture_LargeStream(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	requestID := fmt.Sprintf("test-large-%d", time.Now().UnixNano())
	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		c.Set(common.RequestIdKey, requestID)
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.Header("Content-Type", "text/event-stream")
		flusher := c.Writer.(http.Flusher)

		for i := 0; i < 50; i++ {
			chunk := fmt.Sprintf("data: {\"id\":\"%d\",\"choices\":[{\"delta\":{\"content\":\"chunk%d \"}}]}\n\n", i, i)
			_, _ = c.Writer.Write([]byte(chunk))
			flusher.Flush()
		}
		_, _ = c.Writer.Write([]byte("data: [DONE]\n\n"))
		flusher.Flush()
	})

	req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4","stream":true}`)))
	req.Header.Set("Authorization", "Bearer sk-test")

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)

	record := readCaptureRecord(t, requestID)
	assert.Contains(t, record.Response.Body, "chunk0")
	assert.Contains(t, record.Response.Body, "chunk49")
	assert.Contains(t, record.Response.Body, "[DONE]")
	assert.Len(t, record.Response.BodyChunks, 51)
}

func TestRelayCapture_GetRequest(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	requestID := fmt.Sprintf("test-get-%d", time.Now().UnixNano())
	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		c.Set(common.RequestIdKey, requestID)
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.GET("/v1/models", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"data": []gin.H{{"id": "gpt-4", "object": "model"}},
		})
	})

	req, _ := http.NewRequest("GET", "/v1/models", nil)
	req.Header.Set("Authorization", "Bearer sk-test")

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)

	record := readCaptureRecord(t, requestID)
	assert.Equal(t, "GET", record.Request.Method)
	assert.Equal(t, "/v1/models", record.Request.Path)
	assert.Equal(t, "", record.Request.Body)
	assert.Equal(t, 200, record.Response.StatusCode)
	assert.Contains(t, record.Response.Body, "gpt-4")
}

func TestRelayCapture_ZeroUserID(t *testing.T) {
	clearCaptureCache()

	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set(common.RequestIdKey, fmt.Sprintf("test-zero-%d", time.Now().UnixNano()))
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.GET("/v1/models", func(c *gin.Context) {
		c.JSON(200, gin.H{"data": []gin.H{}})
	})

	req, _ := http.NewRequest("GET", "/v1/models", nil)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)
}

func TestRelayCapture_EmptyRequestIDUsesFallbackFilename(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.GET("/v1/models", func(c *gin.Context) {
		c.JSON(200, gin.H{"data": []gin.H{}})
	})

	req, _ := http.NewRequest("GET", "/v1/models", nil)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)

	today := time.Now().Format("2006-01-02")
	dateDir := filepath.Join(captureTestDir, today)
	files, err := os.ReadDir(dateDir)
	require.NoError(t, err)
	require.Len(t, files, 1)
	assert.True(t, strings.HasPrefix(files[0].Name(), "capture-"))
	assert.True(t, strings.HasSuffix(files[0].Name(), ".json"))

	record := readLatestCaptureRecord(t)
	assert.True(t, strings.HasPrefix(record.RequestID, "capture-"))
	assert.Equal(t, "GET", record.Request.Method)
	assert.Equal(t, "/v1/models", record.Request.Path)
}

func TestRelayCapture_ResponseUnchanged(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	router := setupCaptureTestRouter()
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		c.Set(common.RequestIdKey, fmt.Sprintf("test-unchanged-%d", time.Now().UnixNano()))
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"id":      "chatcmpl-123",
			"choices": []gin.H{{"message": gin.H{"content": "exact response"}}},
		})
	})

	req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`)))
	req.Header.Set("Content-Type", "application/json")

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	assert.Equal(t, 200, rec.Code)
	assert.Contains(t, rec.Body.String(), "exact response")
	assert.Contains(t, rec.Body.String(), "chatcmpl-123")
}

func TestRelayCapture_ConcurrentRequests(t *testing.T) {
	clearCaptureCache()
	t.Cleanup(func() { os.RemoveAll(captureTestDir) })

	SetCaptureEnabled(42, true)
	defer SetCaptureEnabled(42, false)

	router := setupCaptureTestRouter()
	var reqCounter int64
	router.Use(func(c *gin.Context) {
		c.Set("id", 42)
		n := atomic.AddInt64(&reqCounter, 1)
		c.Set(common.RequestIdKey, fmt.Sprintf("concurrent-%d-%d", n, time.Now().UnixNano()))
		c.Next()
	})
	router.Use(RelayCaptureMiddleware())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	done := make(chan bool, 10)
	for i := 0; i < 10; i++ {
		go func(idx int) {
			req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(fmt.Sprintf(`{"req":%d}`, idx))))
			req.Header.Set("Content-Type", "application/json")
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)
			assert.Equal(t, 200, rec.Code)
			done <- true
		}(i)
	}

	for i := 0; i < 10; i++ {
		select {
		case <-done:
		case <-time.After(10 * time.Second):
			t.Fatal("timeout waiting for concurrent request")
		}
	}

	time.Sleep(100 * time.Millisecond)
	today := time.Now().Format("2006-01-02")
	dateDir := filepath.Join(captureTestDir, today)
	files, err := os.ReadDir(dateDir)
	require.NoError(t, err)
	assert.Equal(t, 10, len(files), "should have 10 capture json files")

	for _, f := range files {
		require.True(t, strings.HasSuffix(f.Name(), ".json"))
		content, err := os.ReadFile(filepath.Join(dateDir, f.Name()))
		require.NoError(t, err)
		var record relayCaptureTestRecord
		require.NoError(t, common.Unmarshal(content, &record))
		assert.Equal(t, "POST", record.Request.Method)
		assert.Equal(t, "/v1/chat/completions", record.Request.Path)
		assert.Equal(t, 200, record.Response.StatusCode)
		assert.Contains(t, record.Response.Body, "ok")
	}
}