From f3f1817aea69c135d660e552e22d32dbb08dfc5e Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Mon, 23 Dec 2024 20:48:10 +0800
Subject: [PATCH] feat: Add request start time context key and update
 middleware

- Introduced a new constant `ContextKeyRequestStartTime` to store the request start time in the context, enhancing request tracking.
- Updated the `Distribute` middleware to set the request start time in the context using the new constant.
- Modified the `GenRelayInfo` function to retrieve the request start time from the context, ensuring accurate timing information is used in relay operations.
---
 constant/context_key.go    |  4 ++++
 middleware/distributor.go  |  2 ++
 relay/common/relay_info.go | 13 +++++++------
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/constant/context_key.go b/constant/context_key.go
index 3f2495e1..b02f2d43 100644
--- a/constant/context_key.go
+++ b/constant/context_key.go
@@ -1 +1,5 @@
 package constant
+
+const (
+	ContextKeyRequestStartTime = "request_start_time"
+)
diff --git a/middleware/distributor.go b/middleware/distributor.go
index 022dd0af..0af6f82e 100644
--- a/middleware/distributor.go
+++ b/middleware/distributor.go
@@ -12,6 +12,7 @@ import (
 	"one-api/service"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
 )
@@ -112,6 +113,7 @@ func Distribute() func(c *gin.Context) {
 				}
 			}
 		}
+		c.Set(constant.ContextKeyRequestStartTime, time.Now())
 		SetupContextForSelectedChannel(c, channel, modelRequest.Model)
 		c.Next()
 	}
diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go
index 3bfc2ef6..f650eb95 100644
--- a/relay/common/relay_info.go
+++ b/relay/common/relay_info.go
@@ -2,8 +2,9 @@ package common
 
 import (
 	"one-api/common"
+	"one-api/constant"
 	"one-api/dto"
-	"one-api/relay/constant"
+	relayconstant "one-api/relay/constant"
 	"strings"
 	"time"
 
@@ -66,13 +67,13 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 	userId := c.GetInt("id")
 	group := c.GetString("group")
 	tokenUnlimited := c.GetBool("token_unlimited_quota")
-	startTime := time.Now()
+	startTime := c.GetTime(constant.ContextKeyRequestStartTime)
 	// firstResponseTime = time.Now() - 1 second
 
-	apiType, _ := constant.ChannelType2APIType(channelType)
+	apiType, _ := relayconstant.ChannelType2APIType(channelType)
 
 	info := &RelayInfo{
-		RelayMode:         constant.Path2RelayMode(c.Request.URL.Path),
+		RelayMode:         relayconstant.Path2RelayMode(c.Request.URL.Path),
 		BaseUrl:           c.GetString("base_url"),
 		RequestURLPath:    c.Request.URL.String(),
 		ChannelType:       channelType,
@@ -158,10 +159,10 @@ func GenTaskRelayInfo(c *gin.Context) *TaskRelayInfo {
 	group := c.GetString("group")
 	startTime := time.Now()
 
-	apiType, _ := constant.ChannelType2APIType(channelType)
+	apiType, _ := relayconstant.ChannelType2APIType(channelType)
 
 	info := &TaskRelayInfo{
-		RelayMode:      constant.Path2RelayMode(c.Request.URL.Path),
+		RelayMode:      relayconstant.Path2RelayMode(c.Request.URL.Path),
 		BaseUrl:        c.GetString("base_url"),
 		RequestURLPath: c.Request.URL.String(),
 		ChannelType:    channelType,