ai_site/gateway/internal/proxy/director.go

package proxy

import (
	"context"
	"errors"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"
)

type Director struct {
	Platform *url.URL
	AI       *url.URL
	Timeout  time.Duration
}

func New(platformURL, aiURL string, timeoutSec int) (*Director, error) {
	p, err := url.Parse(platformURL)
	if err != nil {
		return nil, err
	}
	a, err := url.Parse(aiURL)
	if err != nil {
		return nil, err
	}
	if timeoutSec <= 0 {
		timeoutSec = 180
	}
	return &Director{Platform: p, AI: a, Timeout: time.Duration(timeoutSec) * time.Second}, nil
}

// eofRetryTransport：空闲连接被 platform 关掉时首包 EOF，自动换新连接重试一次。
type eofRetryTransport struct {
	base http.RoundTripper
}

func (t eofRetryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	base := t.base
	if base == nil {
		base = http.DefaultTransport
	}
	resp, err := base.RoundTrip(req)
	if err == nil || req.Context().Err() != nil {
		return resp, err
	}
	if !isTransientUpstreamEOF(err) {
		return resp, err
	}
	// Body 可能已部分读；仅在尚未发出/可读失败时安全重试 GET/有 GetBody 的请求
	if req.Body != nil && req.GetBody == nil && req.ContentLength != 0 {
		return resp, err
	}
	if req.GetBody != nil {
		body, gerr := req.GetBody()
		if gerr != nil {
			return resp, err
		}
		req = req.Clone(req.Context())
		req.Body = body
	}
	return base.RoundTrip(req)
}

func isTransientUpstreamEOF(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
		return true
	}
	msg := strings.ToLower(err.Error())
	return strings.Contains(msg, "eof") ||
		strings.Contains(msg, "connection reset") ||
		strings.Contains(msg, "broken pipe") ||
		strings.Contains(msg, "server closed idle connection")
}

func (d *Director) Handler() http.HandlerFunc {
	baseTransport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   10 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		MaxIdleConns:          100,
		MaxIdleConnsPerHost:   8,
		IdleConnTimeout:       30 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		ResponseHeaderTimeout: d.Timeout,
		// Windows 上 platform 重启后空闲连接易 EOF；关掉 keep-alive 比偶发 502 更稳
		DisableKeepAlives: true,
	}
	transport := eofRetryTransport{base: baseTransport}

	return func(w http.ResponseWriter, r *http.Request) {
		target, path := d.route(r.URL.Path)
		timeout := d.Timeout
		// dbsync agent push/pull：Windows SQLite 冷打开可超过默认超时
		if isDBSyncHeavyPath(path) && timeout < 180*time.Second {
			timeout = 180 * time.Second
		}

		proxy := httputil.NewSingleHostReverseProxy(target)
		proxy.Transport = transport
		proxy.FlushInterval = 100 * time.Millisecond
		proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {
			rw.Header().Set("Content-Type", "application/json")
			msg := escape(err.Error())
			code := http.StatusBadGateway
			// 客户端/网关超时 → 503，便于 agent 按 retryable 重试（与 platform 语义一致）
			if errors.Is(err, context.DeadlineExceeded) || errors.Is(req.Context().Err(), context.DeadlineExceeded) ||
				strings.Contains(strings.ToLower(err.Error()), "deadline exceeded") ||
				strings.Contains(strings.ToLower(err.Error()), "context canceled") {
				code = http.StatusServiceUnavailable
			}
			rw.WriteHeader(code)
			_, _ = rw.Write([]byte(
				`{"code":` + strconv.Itoa(code) + `,"message":"upstream unavailable: ` + msg + `","retryable":true}`,
			))
		}
		proxy.ModifyResponse = func(resp *http.Response) error {
			for _, h := range []string{
				"Access-Control-Allow-Origin",
				"Access-Control-Allow-Credentials",
				"Access-Control-Allow-Headers",
				"Access-Control-Allow-Methods",
				"Access-Control-Expose-Headers",
				"Access-Control-Max-Age",
			} {
				resp.Header.Del(h)
			}
			resp.Header.Set("X-Gateway", "aijianzhan-gateway")
			return nil
		}

		r.URL.Path = path
		r.Host = target.Host
		r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))
		r.Header.Set("X-Forwarded-Proto", "http")
		if r.Header.Get("X-Request-Id") == "" {
			r.Header.Set("X-Request-Id", newRequestID())
		}
		reqID := r.Header.Get("X-Request-Id")
		started := time.Now()
		rec := &statusRecorder{ResponseWriter: w, status: 200}

		ctx, cancel := context.WithTimeout(r.Context(), timeout)
		defer cancel()
		proxy.ServeHTTP(rec, r.WithContext(ctx))
		if isDBSyncHeavyPath(path) || rec.status >= 500 {
			log.Printf("gateway %s %s status=%d dur=%s req=%s", r.Method, path, rec.status, time.Since(started), reqID)
		}
	}
}

type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (s *statusRecorder) WriteHeader(code int) {
	s.status = code
	s.ResponseWriter.WriteHeader(code)
}

func isDBSyncHeavyPath(path string) bool {
	if strings.Contains(path, "/agent/sync/") {
		return true
	}
	if !strings.Contains(path, "/admin/sync/") {
		return false
	}
	return strings.Contains(path, "/push") ||
		strings.HasSuffix(path, "/reconcile") ||
		strings.HasSuffix(path, "/ingest") ||
		strings.HasSuffix(path, "/prepare") ||
		strings.HasSuffix(path, "/start") ||
		strings.HasSuffix(path, "/pull") ||
		strings.HasSuffix(path, "/bootstrap")
}

func (d *Director) route(path string) (*url.URL, string) {
	if strings.HasPrefix(path, "/ai/") || path == "/ai" {
		next := strings.TrimPrefix(path, "/ai")
		if next == "" {
			next = "/"
		}
		return d.AI, next
	}
	if path == "/api/v1/apps/generate" || strings.HasPrefix(path, "/api/v1/apps/generate?") {
		return d.AI, path
	}
	if strings.HasPrefix(path, "/api/v1/apps:generate") || strings.Contains(path, "/apps%3Agenerate") {
		return d.AI, "/api/v1/apps/generate"
	}
	return d.Platform, path
}

func escape(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `"`, `\"`)
	if len(s) > 200 {
		s = s[:200]
	}
	return s
}

func newRequestID() string {
	return strings.ReplaceAll(time.Now().UTC().Format("20060102T150405.000000000"), ".", "")
}

// Health 聚合探测
func (d *Director) Health(w http.ResponseWriter, _ *http.Request) {
	client := &http.Client{Timeout: 2 * time.Second}
	out := []struct {
		Name string `json:"name"`
		OK   bool   `json:"ok"`
	}{
		probe(client, "platform", d.Platform.String()+"/api/v1/auth/login"),
		probe(client, "ai", d.AI.String()+"/health"),
	}
	allOK := true
	for _, s := range out {
		if !s.OK {
			allOK = false
			break
		}
	}
	code := http.StatusOK
	if !allOK {
		code = http.StatusServiceUnavailable
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_, _ = io.WriteString(w, `{"gateway":"ok","upstreams":[`)
	for i, s := range out {
		if i > 0 {
			_, _ = io.WriteString(w, ",")
		}
		ok := "false"
		if s.OK {
			ok = "true"
		}
		_, _ = io.WriteString(w, `{"name":"`+s.Name+`","ok":`+ok+`}`)
	}
	_, _ = io.WriteString(w, `]}`)
}

func probe(client *http.Client, name, rawURL string) struct {
	Name string `json:"name"`
	OK   bool   `json:"ok"`
} {
	req, _ := http.NewRequest(http.MethodGet, rawURL, nil)
	if name == "platform" {
		req, _ = http.NewRequest(http.MethodPost, rawURL, strings.NewReader(`{}`))
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := client.Do(req)
	if err != nil {
		return struct {
			Name string `json:"name"`
			OK   bool   `json:"ok"`
		}{Name: name, OK: false}
	}
	defer resp.Body.Close()
	return struct {
		Name string `json:"name"`
		OK   bool   `json:"ok"`
	}{Name: name, OK: true}
}