new-api/test/e2e/mock_capture_server.go

// +build ignore

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/middleware"
	"github.com/gin-gonic/gin"
)

func main() {
	gin.SetMode(gin.ReleaseMode)
	middleware.SetCaptureEnabled(1, true)

	router := gin.New()
	router.Use(gin.Recovery())

	// 模拟 TokenAuth
	router.Use(func(c *gin.Context) {
		auth := c.GetHeader("Authorization")
		if auth == "" || auth == "Bearer disabled" {
			c.JSON(401, gin.H{"error": gin.H{"message": "Unauthorized"}})
			c.Abort()
			return
		}
		c.Set("id", 1)
		c.Next()
	})

	// 模拟 RequestId
	router.Use(func(c *gin.Context) {
		reqID := c.GetHeader(common.RequestIdKey)
		if reqID == "" {
			reqID = fmt.Sprintf("req-%d", time.Now().UnixNano())
		}
		c.Set(common.RequestIdKey, reqID)
		c.Next()
	})

	router.Use(middleware.RelayCaptureMiddleware())

	// ═══════════════════════════════════════════════════════
	// OpenAI Chat Completions API
	// ═══════════════════════════════════════════════════════
	router.POST("/v1/chat/completions", func(c *gin.Context) {
		stream := c.Query("stream") == "true"
		if stream {
			c.Header("Content-Type", "text/event-stream")
			f := c.Writer.(http.Flusher)
			for _, s := range []string{
				"data: {\"id\":\"chatcmpl-001\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n",
				"data: {\"id\":\"chatcmpl-001\",\"choices\":[{\"delta\":{\"content\":\"World!\"}}]}\n\n",
				"data: [DONE]\n\n",
			} {
				c.Writer.Write([]byte(s))
				f.Flush()
			}
			return
		}
		c.JSON(200, gin.H{
			"id":      "chatcmpl-001",
			"object":  "chat.completion",
			"model":   "gpt-4",
			"choices": []gin.H{{"index": 0, "message": gin.H{"role": "assistant", "content": "Hello from Chat Completions!"}}},
			"usage":   gin.H{"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10},
		})
	})

	// ═══════════════════════════════════════════════════════
	// OpenAI Responses API (new)
	// ═══════════════════════════════════════════════════════
	router.POST("/v1/responses", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"id":     "resp_abc123",
			"object": "response",
			"status": "completed",
			"output": []gin.H{{
				"type": "message",
				"role": "assistant",
				"content": []gin.H{{
					"type": "output_text",
					"text": "Hello from Responses API!",
				}},
			}},
			"usage": gin.H{"input_tokens": 8, "output_tokens": 4, "total_tokens": 12},
		})
	})

	// ═══════════════════════════════════════════════════════
	// Claude Messages API (Anthropic)
	// ═══════════════════════════════════════════════════════
	router.POST("/v1/messages", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"id":         "msg_01AbCdEfGhIjKlMnOpQrStUv",
			"type":       "message",
			"role":       "assistant",
			"model":      "claude-sonnet-4-6",
			"content":    []gin.H{{"type": "text", "text": "Hello! How can I assist you today?"}},
			"stop_reason": "end_turn",
			"usage":      gin.H{"input_tokens": 8, "output_tokens": 12},
		})
	})

	// ═══════════════════════════════════════════════════════
	// Gemini Native API (/v1beta/models/*)
	// ═══════════════════════════════════════════════════════
	router.POST("/v1beta/models/*path", func(c *gin.Context) {
		path := c.Param("path")
		c.JSON(200, gin.H{
			"candidates": []gin.H{{
				"content": gin.H{
					"role":  "model",
					"parts": []gin.H{{"text": fmt.Sprintf("Gemini response for: /%s", path)}},
				},
				"finishReason": "STOP",
			}},
			"usageMetadata": gin.H{"promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15},
		})
	})

	// ═══════════════════════════════════════════════════════
	// Gemini via /v1/models/* (main route)
	// ═══════════════════════════════════════════════════════
	router.POST("/v1/models/*path", func(c *gin.Context) {
		path := c.Param("path")
		c.JSON(200, gin.H{
			"candidates": []gin.H{{
				"content": gin.H{
					"role":  "model",
					"parts": []gin.H{{"text": fmt.Sprintf("Gemini(v1) response for: /%s", path)}},
				},
				"finishReason": "STOP",
			}},
		})
	})

	// ═══════════════════════════════════════════════════════
	// Models list
	// ═══════════════════════════════════════════════════════
	router.GET("/v1/models", func(c *gin.Context) {
		c.JSON(200, gin.H{"data": []gin.H{{"id": "gpt-4"}, {"id": "claude-3"}, {"id": "gemini-2.0-flash"}}})
	})

	// ═══════════════════════════════════════════════════════
	// Error endpoint for testing error status codes
	// ═══════════════════════════════════════════════════════
	router.POST("/v1/error", func(c *gin.Context) {
		code := 500
		if c := c.Query("code"); c != "" {
			fmt.Sscanf(c, "%d", &code)
		}
		c.JSON(code, gin.H{"error": gin.H{"message": "test error", "type": "test_error"}})
	})

	// ═══════════════════════════════════════════════════════
	// Capture file management endpoints
	// ═══════════════════════════════════════════════════════
	router.GET("/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	router.GET("/capture-files", func(c *gin.Context) {
		today := time.Now().Format("2006-01-02")
		dir := fmt.Sprintf("./data/relay-capture/%s", today)
		entries, err := os.ReadDir(dir)
		if err != nil {
			c.JSON(200, gin.H{"files": []string{}, "error": err.Error()})
			return
		}
		var files []string
		for _, e := range entries {
			files = append(files, e.Name())
		}
		c.JSON(200, gin.H{"dir": dir, "files": files})
	})

	router.GET("/capture-files/:name", func(c *gin.Context) {
		today := time.Now().Format("2006-01-02")
		path := fmt.Sprintf("./data/relay-capture/%s/%s", today, c.Param("name"))
		data, err := os.ReadFile(path)
		if err != nil {
			c.JSON(404, gin.H{"error": err.Error()})
			return
		}
		var record map[string]interface{}
		json.Unmarshal(data, &record)
		c.JSON(200, record)
	})

	router.DELETE("/capture-files", func(c *gin.Context) {
		os.RemoveAll("./data/relay-capture")
		c.JSON(200, gin.H{"status": "cleaned"})
	})

	port := "18080"
	if p := os.Getenv("PORT"); p != "" {
		port = p
	}
	log.Printf("Mock Relay Server on :%s (chat|responses|messages|gemini)", port)
	http.ListenAndServe(":"+port, router)
}