修复vertexai模型无法通过配置的http代理正确访问的问题 (#844)

Author: OPSTime

providers/vertexai/base.go

@@ -33,16 +33,20 @@ type VertexAIProviderFactory struct{}
 
 // 创建 VertexAIProvider
 func (f VertexAIProviderFactory) Create(channel *model.Channel) base.ProviderInterface {
+	proxyAddr := ""
+	if channel.Proxy != nil {
+		proxyAddr = *channel.Proxy
+	}
+
 	vertexAIProvider := &VertexAIProvider{
 		BaseProvider: base.BaseProvider{
 			Config:    getConfig(),
 			Channel:   channel,
-			Requester: requester.NewHTTPRequester(*channel.Proxy, nil),
+			Requester: requester.NewHTTPRequester(proxyAddr, nil),
 		},
 	}
 
 	getKeyConfig(vertexAIProvider)
-
 	return vertexAIProvider
 }
 
@@ -109,16 +113,19 @@ func (p *VertexAIProvider) GetToken() (string, error) {
 		return "", fmt.Errorf("failed to unmarshal credentials: %w", err)
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
 	proxyAddr := ""
 	if p.Channel.Proxy != nil && *p.Channel.Proxy != "" {
 		proxyAddr = *p.Channel.Proxy
+		logger.SysLog(fmt.Sprintf("Vertex AI proxy: %s", proxyAddr))
 	}
 
+	// 尝试使用gRPC客户端获取token
 	client, err := credentials.NewIamCredentialsClient(ctx, option.WithCredentialsJSON([]byte(p.Channel.Key)), option.WithGRPCDialOption(grpc.WithContextDialer(customDialer(proxyAddr))))
 	if err != nil {
+		logger.SysError(fmt.Sprintf("Failed to create IAM credentials client: %v", err))
 		return "", fmt.Errorf("failed to create IAM credentials client: %w", err)
 	}
 	defer client.Close()
@@ -172,8 +179,6 @@ func errorHandle(vertexaiError *VertexaiError) *types.OpenAIError {
 		return nil
 	}
 
-	logger.SysError(fmt.Sprintf("VertexAI error: %s", vertexaiError.Error.Message))
-
 	return &types.OpenAIError{
 		Message: "VertexAI错误",
 		Type:    "gemini_error",
@@ -183,24 +188,84 @@ func errorHandle(vertexaiError *VertexaiError) *types.OpenAIError {
 }
 
 func customDialer(proxyAddr string) func(context.Context, string) (net.Conn, error) {
-
 	return func(ctx context.Context, addr string) (net.Conn, error) {
+		// 创建统一的dialer配置
+		dialer := &net.Dialer{
+			Timeout:   20 * time.Second,
+			KeepAlive: 30 * time.Second,
+		}
+
+		// 无代理直接连接
 		if proxyAddr == "" {
-			return net.Dial("tcp", addr)
+			return dialer.DialContext(ctx, "tcp", addr)
 		}
 
 		proxyURL, err := url.Parse(proxyAddr)
 		if err != nil {
 			return nil, fmt.Errorf("error parsing proxy address: %w", err)
 		}
 
-		dialer := &net.Dialer{}
-
-		dialerProxy, err := proxy.FromURL(proxyURL, dialer)
-		if err != nil {
-			return nil, fmt.Errorf("failed to create HTTP dialer: %v", err)
+		// 根据代理类型选择连接方式
+		switch proxyURL.Scheme {
+		case "http":
+			return connectViaHTTPProxy(ctx, proxyURL, addr)
+		case "https":
+			logger.SysError("Warning: HTTPS proxy not compatible with gRPC, using direct connection")
+			return dialer.DialContext(ctx, "tcp", addr)
+		case "socks5", "socks5h":
+			return connectViaSOCKS5Proxy(ctx, dialer, proxyURL, addr)
+		default:
+			return nil, fmt.Errorf("unsupported proxy scheme: %s", proxyURL.Scheme)
 		}
+	}
+}
+
+// connectViaHTTPProxy 通过HTTP代理建立连接
+func connectViaHTTPProxy(ctx context.Context, proxyURL *url.URL, targetAddr string) (net.Conn, error) {
+	dialer := &net.Dialer{
+		Timeout:   20 * time.Second,
+		KeepAlive: 30 * time.Second,
+	}
+
+	proxyConn, err := dialer.DialContext(ctx, "tcp", proxyURL.Host)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to HTTP proxy: %w", err)
+	}
+
+	// 发送HTTP CONNECT请求
+	connectReq := fmt.Sprintf("CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", targetAddr, targetAddr)
+	if _, err = proxyConn.Write([]byte(connectReq)); err != nil {
+		proxyConn.Close()
+		return nil, fmt.Errorf("failed to send CONNECT request: %w", err)
+	}
 
-		return dialerProxy.Dial("tcp", addr)
+	// 读取代理响应
+	response := make([]byte, 1024)
+	n, err := proxyConn.Read(response)
+	if err != nil {
+		proxyConn.Close()
+		return nil, fmt.Errorf("failed to read proxy response: %w", err)
+	}
+
+	responseStr := string(response[:n])
+	if !strings.Contains(responseStr, "200 Connection established") && !strings.Contains(responseStr, "200 OK") {
+		proxyConn.Close()
+		return nil, fmt.Errorf("HTTP proxy CONNECT failed: %s", responseStr)
 	}
+
+	return proxyConn, nil
+}
+
+// connectViaSOCKS5Proxy 通过SOCKS5代理建立连接
+func connectViaSOCKS5Proxy(ctx context.Context, dialer *net.Dialer, proxyURL *url.URL, addr string) (net.Conn, error) {
+	dialerProxy, err := proxy.FromURL(proxyURL, dialer)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create SOCKS5 proxy dialer: %v", err)
+	}
+
+	if contextDialer, ok := dialerProxy.(proxy.ContextDialer); ok {
+		return contextDialer.DialContext(ctx, "tcp", addr)
+	}
+
+	return dialerProxy.Dial("tcp", addr)
 }