package common import ( "net/http" "strings" commonpkg "github.com/QuantumNous/new-api/common" "github.com/gin-gonic/gin" ) const ( ContextKeyRelayChatID = "relay_chat_id" ContextKeyRelayUpstreamID = "relay_upstream_id" ) func SetRelayChatID(c *gin.Context, chatID string) { if c == nil { return } chatID = strings.TrimSpace(chatID) if chatID == "" { return } c.Set(ContextKeyRelayChatID, chatID) } func GetRelayChatID(c *gin.Context) string { if c == nil { return "" } return strings.TrimSpace(c.GetString(ContextKeyRelayChatID)) } func SetRelayUpstreamID(c *gin.Context, upstreamID string) { if c == nil { return } upstreamID = strings.TrimSpace(upstreamID) if upstreamID == "" { return } c.Set(ContextKeyRelayUpstreamID, upstreamID) } func GetRelayUpstreamID(c *gin.Context) string { if c == nil { return "" } return strings.TrimSpace(c.GetString(ContextKeyRelayUpstreamID)) } func ExtractTopLevelChatID(c *gin.Context) (string, error) { if c == nil { return "", nil } bodyStorage, err := commonpkg.GetBodyStorage(c) if err != nil { return "", err } bodyBytes, err := bodyStorage.Bytes() if err != nil || len(bodyBytes) == 0 { return "", err } var payload map[string]interface{} if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil { return "", nil } return strings.TrimSpace(commonpkg.Interface2String(payload["chat_id"])), nil } func CaptureUpstreamIDFromHTTPResponse(info *RelayInfo, resp *http.Response) { if info == nil || resp == nil { return } for _, key := range []string{"x-request-id", "request-id"} { value := strings.TrimSpace(resp.Header.Get(key)) if value != "" { info.UpstreamID = value return } } }