new-api/middleware/muse_request_context_test.go

package middleware

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

type museRequestContextResponse struct {
	MuseUserID    string `json:"muse_user_id"`
	MuseWorkID    string `json:"muse_work_id"`
	MuseScene     string `json:"muse_scene"`
	MuseRequestID string `json:"muse_request_id"`
	TraceID       string `json:"trace_id"`
}

func setupMuseRequestContextTestDB(t *testing.T) *gorm.DB {
	t.Helper()

	common.UsingSQLite = true
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	common.RedisEnabled = false

	db, err := gorm.Open(sqlite.Open("file:muse_request_context_test?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	model.DB = db
	model.LOG_DB = db
	require.NoError(t, db.AutoMigrate(&model.User{}, &model.Log{}))
	require.NoError(t, db.Create(&model.User{
		Id:       1,
		Username: "alice",
		Password: "hashed",
		Status:   common.UserStatusEnabled,
	}).Error)

	t.Cleanup(func() {
		sqlDB, err := db.DB()
		if err == nil {
			_ = sqlDB.Close()
		}
	})
	return db
}

func TestMuseRequestContextExtractsHeaders(t *testing.T) {
	gin.SetMode(gin.TestMode)

	req := httptest.NewRequest(
		http.MethodPost,
		"/v1/chat/completions",
		strings.NewReader(`{"metadata":{"work_id":"42","scene":"suggestion_generation","muse_user_id":"u-1"}}`),
	)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set(museHeaderUserID, "u-1")
	req.Header.Set(museHeaderWorkID, "42")
	req.Header.Set(museHeaderRequest, "req-1")
	req.Header.Set(museHeaderTraceID, "trace-1")

	recorder := httptest.NewRecorder()
	router := gin.New()
	router.Use(MuseRequestContext())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.JSON(http.StatusOK, museRequestContextResponse{
			MuseUserID:    common.GetContextKeyString(c, constant.ContextKeyMuseUserID),
			MuseWorkID:    common.GetContextKeyString(c, constant.ContextKeyMuseWorkID),
			MuseScene:     common.GetContextKeyString(c, constant.ContextKeyMuseScene),
			MuseRequestID: common.GetContextKeyString(c, constant.ContextKeyMuseRequestID),
			TraceID:       common.GetContextKeyString(c, constant.ContextKeyMuseTraceID),
		})
	})

	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)

	var response museRequestContextResponse
	require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
	require.Equal(t, "u-1", response.MuseUserID)
	require.Equal(t, "42", response.MuseWorkID)
	require.Equal(t, "suggestion_generation", response.MuseScene)
	require.Equal(t, "trace-1", response.TraceID)
	require.Equal(t, "req-1", response.MuseRequestID)
}

func TestMuseRequestContextFallsBackToMetadata(t *testing.T) {
	gin.SetMode(gin.TestMode)

	req := httptest.NewRequest(
		http.MethodPost,
		"/v1/chat/completions",
		strings.NewReader(`{"metadata":{"muse_user_id":"u-2","work_id":42,"scene":"suggestion_generation","request_id":"req-meta","trace_id":"trace-meta"}}`),
	)
	req.Header.Set("Content-Type", "application/json")

	recorder := httptest.NewRecorder()
	router := gin.New()
	router.Use(MuseRequestContext())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		c.JSON(http.StatusOK, museRequestContextResponse{
			MuseUserID:    common.GetContextKeyString(c, constant.ContextKeyMuseUserID),
			MuseWorkID:    common.GetContextKeyString(c, constant.ContextKeyMuseWorkID),
			MuseScene:     common.GetContextKeyString(c, constant.ContextKeyMuseScene),
			MuseRequestID: common.GetContextKeyString(c, constant.ContextKeyMuseRequestID),
			TraceID:       common.GetContextKeyString(c, constant.ContextKeyMuseTraceID),
		})
	})

	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)

	var response museRequestContextResponse
	require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
	require.Equal(t, "u-2", response.MuseUserID)
	require.Equal(t, "42", response.MuseWorkID)
	require.Equal(t, "suggestion_generation", response.MuseScene)
	require.Equal(t, "trace-meta", response.TraceID)
	require.Equal(t, "req-meta", response.MuseRequestID)
}

func TestMuseRequestContextRecordConsumeLogIncludesMuseFields(t *testing.T) {
	gin.SetMode(gin.TestMode)
	setupMuseRequestContextTestDB(t)

	req := httptest.NewRequest(
		http.MethodPost,
		"/v1/chat/completions",
		strings.NewReader(`{"metadata":{"scene":"suggestion_generation","request_id":"req-log","trace_id":"trace-log"}}`),
	)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set(museHeaderUserID, "u-log")
	req.Header.Set(museHeaderWorkID, "77")
	recorder := httptest.NewRecorder()

	router := gin.New()
	router.Use(MuseRequestContext())
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		model.RecordConsumeLog(c, 1, model.RecordConsumeLogParams{
			ModelName: "gpt-4o",
			Content:   "ok",
			Group:     "default",
			Other:     map[string]interface{}{"existing": true},
		})
		var logEntry model.Log
		require.NoError(t, model.LOG_DB.Last(&logEntry).Error)
		c.JSON(http.StatusOK, gin.H{
			"request_id": logEntry.RequestId,
			"other":      logEntry.Other,
		})
	})

	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)
	var response struct {
		RequestID string `json:"request_id"`
		Other     string `json:"other"`
	}
	require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
	require.Empty(t, response.RequestID)
	require.Contains(t, response.Other, `"muse_user_id":"u-log"`)
	require.Contains(t, response.Other, `"muse_work_id":"77"`)
	require.Contains(t, response.Other, `"muse_request_id":"req-log"`)
	require.Contains(t, response.Other, `"muse_scene":"suggestion_generation"`)
	require.Contains(t, response.Other, `"muse_trace_id":"trace-log"`)
}