diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 8e081bac34..ea9169e0b7 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -601,6 +601,7 @@ type ServerConfig struct { Host string `mapstructure:"host"` Port int `mapstructure:"port"` Mode string `mapstructure:"mode"` // debug/release + EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL，用于生成邮件中的外部链接 ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时（秒） IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时（秒） @@ -1459,6 +1460,9 @@ func load(allowMissingJWTSecret bool) (*Config, error) { // 环境变量支持 viper.AutomaticEnv() viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_")) + if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil { + return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err) + } // 默认值 setDefaults() @@ -1614,6 +1618,7 @@ func setDefaults() { viper.SetDefault("server.host", "0.0.0.0") viper.SetDefault("server.port", 8080) viper.SetDefault("server.mode", "release") + viper.SetDefault("server.enable_server_timing", false) viper.SetDefault("server.frontend_url", "") viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头 viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时 diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 489bc346b8..4eea3a2840 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -17,6 +17,23 @@ func resetViperWithJWTSecret(t *testing.T) { t.Setenv("JWT_SECRET", strings.Repeat("x", 32)) } +func TestLoadServerTimingConfig(t *testing.T) { + t.Run("disabled by default", func(t *testing.T) { + resetViperWithJWTSecret(t) + cfg, err := Load() + require.NoError(t, err) + require.False(t, cfg.Server.EnableServerTiming) + }) + + t.Run("enabled by exact environment variable", func(t *testing.T) { + resetViperWithJWTSecret(t) + t.Setenv("ENABLE_SERVER_TIMING", "true") + cfg, err := Load() + require.NoError(t, err) + require.True(t, cfg.Server.EnableServerTiming) + }) +} + func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) { viper.Reset() t.Setenv("JWT_SECRET", "") diff --git a/backend/internal/handler/admin/ops_ws_handler.go b/backend/internal/handler/admin/ops_ws_handler.go index 75fd7ea002..e4c42cc9c0 100644 --- a/backend/internal/handler/admin/ops_ws_handler.go +++ b/backend/internal/handler/admin/ops_ws_handler.go @@ -16,6 +16,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" @@ -323,7 +324,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) { // If realtime monitoring is disabled, prefer a successful WS upgrade followed by a clean close // with a deterministic close code. This prevents clients from spinning on 404/1006 reconnect loops. if !h.opsService.IsRealtimeMonitoringEnabled(c.Request.Context()) { - conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c)) if err != nil { c.JSON(http.StatusNotFound, gin.H{"error": "ops realtime monitoring is disabled"}) return @@ -358,7 +359,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) { defer releaseOpsWSIPSlot(clientIP) } - conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c)) if err != nil { logger.LegacyPrintf("handler.admin.ops_ws", "[OpsWS] upgrade failed: %v", err) return diff --git a/backend/internal/pkg/antigravity/client.go b/backend/internal/pkg/antigravity/client.go index e318d1cdaf..39b6d2c90c 100644 --- a/backend/internal/pkg/antigravity/client.go +++ b/backend/internal/pkg/antigravity/client.go @@ -17,6 +17,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" ) // ForbiddenError 表示上游返回 403 Forbidden @@ -279,7 +280,6 @@ func NewClient(proxyURL string) (*Client, error) { } client.Transport = transport } - return &Client{ httpClient: client, }, nil @@ -341,7 +341,7 @@ func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (* } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - resp, err := c.httpClient.Do(req) + resp, err := servertiming.Do(c.httpClient, req) if err != nil { return nil, fmt.Errorf("token 交换请求失败: %w", err) } @@ -383,7 +383,7 @@ func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenR } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - resp, err := c.httpClient.Do(req) + resp, err := servertiming.Do(c.httpClient, req) if err != nil { return nil, fmt.Errorf("token 刷新请求失败: %w", err) } @@ -414,7 +414,7 @@ func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo } req.Header.Set("Authorization", "Bearer "+accessToken) - resp, err := c.httpClient.Do(req) + resp, err := servertiming.Do(c.httpClient, req) if err != nil { return nil, fmt.Errorf("用户信息请求失败: %w", err) } @@ -465,7 +465,7 @@ func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadC req.Header.Set("Content-Type", "application/json") req.Header.Set("User-Agent", GetUserAgentForContext(ctx)) - resp, err := c.httpClient.Do(req) + resp, err := servertiming.Do(c.httpClient, req) if err != nil { lastErr = fmt.Errorf("loadCodeAssist 请求失败: %w", err) if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { @@ -544,7 +544,7 @@ func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (s req.Header.Set("Content-Type", "application/json") req.Header.Set("User-Agent", GetUserAgentForContext(ctx)) - resp, err := c.httpClient.Do(req) + resp, err := servertiming.Do(c.httpClient, req) if err != nil { lastErr = fmt.Errorf("onboardUser 请求失败: %w", err) if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { @@ -683,7 +683,7 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI req.Header.Set("Content-Type", "application/json") req.Header.Set("User-Agent", GetUserAgentForContext(ctx)) - resp, err := fetchClient.Do(req) + resp, err := servertiming.Do(fetchClient, req) if err != nil { lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err) if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { @@ -842,7 +842,7 @@ func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetU req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1") req.Host = "daily-cloudcode-pa.googleapis.com" - resp, err := c.httpClient.Do(req) + resp, err := servertiming.Do(c.httpClient, req) if err != nil { return nil, fmt.Errorf("setUserSettings 请求失败: %w", err) } @@ -885,7 +885,7 @@ func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID strin req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1") req.Host = "daily-cloudcode-pa.googleapis.com" - resp, err := c.httpClient.Do(req) + resp, err := servertiming.Do(c.httpClient, req) if err != nil { return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err) } diff --git a/backend/internal/pkg/httpclient/pool.go b/backend/internal/pkg/httpclient/pool.go index 12804cc67d..22d3c65feb 100644 --- a/backend/internal/pkg/httpclient/pool.go +++ b/backend/internal/pkg/httpclient/pool.go @@ -25,6 +25,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" ) @@ -92,6 +93,7 @@ func buildClient(opts Options) (*http.Client, error) { if opts.ValidateResolvedIP && !opts.AllowPrivateHosts { rt = newValidatedTransport(transport) } + rt = servertiming.WrapRoundTripper(rt) return &http.Client{ Transport: rt, Timeout: opts.Timeout, diff --git a/backend/internal/pkg/servertiming/collector.go b/backend/internal/pkg/servertiming/collector.go new file mode 100644 index 0000000000..553edede31 --- /dev/null +++ b/backend/internal/pkg/servertiming/collector.go @@ -0,0 +1,348 @@ +package servertiming + +import ( + "context" + "fmt" + "sort" + "strconv" + "strings" + "sync" + "time" +) + +const ( + HeaderName = "Server-Timing" + AdminUIHeader = "X-Admin-UI-Request" + MetricDatabase = "db" + MetricRedis = "redis" + dependencyPrefix = "dep_" + + maxMetricNameLength = 48 + maxIntervals = 2048 + maxHeaderLength = 4096 +) + +type contextKey struct{} + +type interval struct { + start time.Time + end time.Time +} + +type metric struct { + count int64 + intervals []interval +} + +// Collector stores request-scoped timing samples. It is safe for concurrent use. +type Collector struct { + startedAt time.Time + + mu sync.Mutex + metrics map[string]*metric + cacheStatus string +} + +// New creates a collector whose total duration starts at startedAt. +func New(startedAt time.Time) *Collector { + if startedAt.IsZero() { + startedAt = time.Now() + } + return &Collector{ + startedAt: startedAt, + metrics: make(map[string]*metric), + } +} + +// WithCollector attaches a collector to a context. +func WithCollector(ctx context.Context, collector *Collector) context.Context { + if ctx == nil { + ctx = context.Background() + } + if collector == nil { + return ctx + } + return context.WithValue(ctx, contextKey{}, collector) +} + +// FromContext returns the request timing collector, when one is active. +func FromContext(ctx context.Context) (*Collector, bool) { + if ctx == nil { + return nil, false + } + collector, ok := ctx.Value(contextKey{}).(*Collector) + return collector, ok && collector != nil +} + +// Active reports whether timing collection is enabled for this request. +func Active(ctx context.Context) bool { + _, ok := FromContext(ctx) + return ok +} + +// Record adds a completed interval and operation count to a metric. +func Record(ctx context.Context, name string, startedAt, endedAt time.Time, count int) { + collector, ok := FromContext(ctx) + if !ok { + return + } + collector.Record(name, startedAt, endedAt, count) +} + +// RecordInterval adds timing without incrementing the operation count. It is +// useful when one logical operation has multiple blocking driver calls. +func RecordInterval(ctx context.Context, name string, startedAt, endedAt time.Time) { + collector, ok := FromContext(ctx) + if !ok { + return + } + collector.record(name, startedAt, endedAt, 0) +} + +// Record adds a completed interval directly to the collector. +func (c *Collector) Record(name string, startedAt, endedAt time.Time, count int) { + if count <= 0 { + count = 1 + } + c.record(name, startedAt, endedAt, count) +} + +func (c *Collector) record(name string, startedAt, endedAt time.Time, count int) { + name = normalizeMetricName(name) + if c == nil || name == "" || startedAt.IsZero() || endedAt.Before(startedAt) { + return + } + if count < 0 { + count = 0 + } + + c.mu.Lock() + m := c.metrics[name] + if m == nil { + m = &metric{} + c.metrics[name] = m + } + m.count += int64(count) + if len(m.intervals) < maxIntervals { + m.intervals = append(m.intervals, interval{start: startedAt, end: endedAt}) + } + c.mu.Unlock() +} + +// Observe starts a metric span and returns an idempotent completion function. +func Observe(ctx context.Context, name string) func() { + collector, ok := FromContext(ctx) + name = normalizeMetricName(name) + if !ok || name == "" { + return func() {} + } + startedAt := time.Now() + var once sync.Once + return func() { + once.Do(func() { + collector.Record(name, startedAt, time.Now(), 1) + }) + } +} + +// ObserveDependency starts a named external dependency span. +func ObserveDependency(ctx context.Context, module string) func() { + return Observe(ctx, dependencyMetricName(module)) +} + +// RecordDependency records a completed external dependency interval. +func RecordDependency(ctx context.Context, module string, startedAt, endedAt time.Time) { + Record(ctx, dependencyMetricName(module), startedAt, endedAt, 1) +} + +// SetCacheStatus records the response-cache outcome for the request. +func SetCacheStatus(ctx context.Context, status string) { + collector, ok := FromContext(ctx) + if !ok { + return + } + status = normalizeCacheStatus(status) + if status == "" { + return + } + collector.mu.Lock() + collector.cacheStatus = status + collector.mu.Unlock() +} + +// HeaderValue renders a bounded, deterministic Server-Timing header. +func HeaderValue(ctx context.Context, endedAt time.Time, cacheStatus string) string { + collector, ok := FromContext(ctx) + if !ok { + return "" + } + return collector.HeaderValue(endedAt, cacheStatus) +} + +// HeaderValue renders a bounded, deterministic Server-Timing header. +func (c *Collector) HeaderValue(endedAt time.Time, cacheStatus string) string { + if c == nil { + return "" + } + if endedAt.IsZero() { + endedAt = time.Now() + } + if endedAt.Before(c.startedAt) { + endedAt = c.startedAt + } + + c.mu.Lock() + metrics := make(map[string]metric, len(c.metrics)) + allIntervals := make([]interval, 0) + dependencyIntervals := make([]interval, 0) + var dependencyCount int64 + for name, source := range c.metrics { + copied := metric{count: source.count, intervals: append([]interval(nil), source.intervals...)} + metrics[name] = copied + allIntervals = append(allIntervals, copied.intervals...) + if strings.HasPrefix(name, dependencyPrefix) { + dependencyIntervals = append(dependencyIntervals, copied.intervals...) + dependencyCount += copied.count + } + } + storedCacheStatus := c.cacheStatus + c.mu.Unlock() + + total := endedAt.Sub(c.startedAt) + blocked := unionDuration(allIntervals, c.startedAt, endedAt) + app := total - blocked + if app < 0 { + app = 0 + } + + cacheStatus = normalizeCacheStatus(cacheStatus) + if cacheStatus == "" { + cacheStatus = normalizeCacheStatus(storedCacheStatus) + } + if cacheStatus == "" { + cacheStatus = "bypass" + } + + database := metrics[MetricDatabase] + redisMetric := metrics[MetricRedis] + parts := []string{ + "total;dur=" + formatDuration(total), + "app;dur=" + formatDuration(app), + fmt.Sprintf("db;dur=%s;desc=\"queries=%d\"", formatDuration(unionDuration(database.intervals, c.startedAt, endedAt)), database.count), + fmt.Sprintf("redis;dur=%s;desc=\"commands=%d\"", formatDuration(unionDuration(redisMetric.intervals, c.startedAt, endedAt)), redisMetric.count), + "cache;desc=\"" + cacheStatus + "\"", + fmt.Sprintf("deps;dur=%s;desc=\"calls=%d\"", formatDuration(unionDuration(dependencyIntervals, c.startedAt, endedAt)), dependencyCount), + } + + dependencyNames := make([]string, 0) + for name := range metrics { + if strings.HasPrefix(name, dependencyPrefix) { + dependencyNames = append(dependencyNames, name) + } + } + sort.Strings(dependencyNames) + for _, name := range dependencyNames { + m := metrics[name] + part := fmt.Sprintf("%s;dur=%s;desc=\"calls=%d\"", name, formatDuration(unionDuration(m.intervals, c.startedAt, endedAt)), m.count) + candidate := strings.Join(append(parts, part), ", ") + if len(candidate) > maxHeaderLength { + break + } + parts = append(parts, part) + } + + return strings.Join(parts, ", ") +} + +func dependencyMetricName(module string) string { + module = normalizeMetricName(module) + module = strings.TrimPrefix(module, dependencyPrefix) + if module == "" { + module = "http" + } + return dependencyPrefix + module +} + +func normalizeMetricName(name string) string { + name = strings.ToLower(strings.TrimSpace(name)) + if name == "" { + return "" + } + var b strings.Builder + b.Grow(min(len(name), maxMetricNameLength)) + for _, r := range name { + if b.Len() >= maxMetricNameLength { + break + } + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9': + _, _ = b.WriteRune(r) + case r == '_' || r == '-': + _ = b.WriteByte('_') + } + } + return strings.Trim(b.String(), "_") +} + +func normalizeCacheStatus(status string) string { + switch strings.ToLower(strings.TrimSpace(status)) { + case "hit": + return "hit" + case "miss": + return "miss" + case "bypass": + return "bypass" + default: + return "" + } +} + +func unionDuration(intervals []interval, lowerBound, upperBound time.Time) time.Duration { + if len(intervals) == 0 || !upperBound.After(lowerBound) { + return 0 + } + normalized := make([]interval, 0, len(intervals)) + for _, item := range intervals { + start := item.start + end := item.end + if start.Before(lowerBound) { + start = lowerBound + } + if end.After(upperBound) { + end = upperBound + } + if end.After(start) { + normalized = append(normalized, interval{start: start, end: end}) + } + } + if len(normalized) == 0 { + return 0 + } + sort.Slice(normalized, func(i, j int) bool { + return normalized[i].start.Before(normalized[j].start) + }) + + currentStart := normalized[0].start + currentEnd := normalized[0].end + var total time.Duration + for _, item := range normalized[1:] { + if !item.start.After(currentEnd) { + if item.end.After(currentEnd) { + currentEnd = item.end + } + continue + } + total += currentEnd.Sub(currentStart) + currentStart = item.start + currentEnd = item.end + } + total += currentEnd.Sub(currentStart) + return total +} + +func formatDuration(value time.Duration) string { + if value < 0 { + value = 0 + } + return strconv.FormatFloat(float64(value)/float64(time.Millisecond), 'f', 1, 64) +} diff --git a/backend/internal/pkg/servertiming/collector_test.go b/backend/internal/pkg/servertiming/collector_test.go new file mode 100644 index 0000000000..1bb809f9bc --- /dev/null +++ b/backend/internal/pkg/servertiming/collector_test.go @@ -0,0 +1,129 @@ +package servertiming + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + "time" +) + +func TestCollectorHeaderValueAggregatesIntervals(t *testing.T) { + startedAt := time.Unix(100, 0) + collector := New(startedAt) + collector.Record(MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(40*time.Millisecond), 2) + collector.Record(MetricRedis, startedAt.Add(30*time.Millisecond), startedAt.Add(50*time.Millisecond), 3) + collector.Record(dependencyMetricName("openai"), startedAt.Add(70*time.Millisecond), startedAt.Add(100*time.Millisecond), 1) + collector.Record(dependencyMetricName("github"), startedAt.Add(60*time.Millisecond), startedAt.Add(90*time.Millisecond), 1) + + got := collector.HeaderValue(startedAt.Add(120*time.Millisecond), "miss") + want := `total;dur=120.0, app;dur=40.0, db;dur=30.0;desc="queries=2", redis;dur=20.0;desc="commands=3", cache;desc="miss", deps;dur=40.0;desc="calls=2", dep_github;dur=30.0;desc="calls=1", dep_openai;dur=30.0;desc="calls=1"` + if got != want { + t.Fatalf("HeaderValue() = %q, want %q", got, want) + } +} + +func TestRecordIntervalDoesNotIncrementCount(t *testing.T) { + startedAt := time.Unix(200, 0) + collector := New(startedAt) + ctx := WithCollector(context.Background(), collector) + + Record(ctx, MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(20*time.Millisecond), 1) + RecordInterval(ctx, MetricDatabase, startedAt.Add(30*time.Millisecond), startedAt.Add(40*time.Millisecond)) + + header := HeaderValue(ctx, startedAt.Add(100*time.Millisecond), "hit") + if !strings.Contains(header, `db;dur=20.0;desc="queries=1"`) { + t.Fatalf("header %q does not contain one query with both blocking intervals", header) + } + if !strings.Contains(header, "app;dur=80.0") { + t.Fatalf("header %q does not subtract the interval union from app time", header) + } +} + +func TestCollectorCacheStatusFallback(t *testing.T) { + startedAt := time.Unix(300, 0) + collector := New(startedAt) + ctx := WithCollector(context.Background(), collector) + + SetCacheStatus(ctx, " HIT ") + if got := HeaderValue(ctx, startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="hit"`) { + t.Fatalf("HeaderValue() = %q, want stored cache hit", got) + } + + other := New(startedAt) + if got := other.HeaderValue(startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="bypass"`) { + t.Fatalf("HeaderValue() = %q, want cache bypass", got) + } +} + +func TestCollectorSanitizesDependencyMetric(t *testing.T) { + startedAt := time.Unix(400, 0) + collector := New(startedAt) + ctx := WithCollector(context.Background(), collector) + RecordDependency(ctx, "GitHub API\r\nInjected;dur=999", startedAt, startedAt.Add(time.Millisecond)) + + header := HeaderValue(ctx, startedAt.Add(2*time.Millisecond), "bypass") + if strings.ContainsAny(header, "\r\n") || strings.Contains(header, ";dur=999") { + t.Fatalf("unsafe metric content reached header: %q", header) + } + if !strings.Contains(header, "dep_githubapiinjecteddur999;dur=1.0") { + t.Fatalf("sanitized dependency metric missing from header: %q", header) + } +} + +func TestCollectorBoundsHeaderLength(t *testing.T) { + startedAt := time.Unix(500, 0) + collector := New(startedAt) + for i := 0; i < 300; i++ { + collector.Record( + dependencyMetricName(fmt.Sprintf("module_%03d_with_a_deliberately_long_name", i)), + startedAt, + startedAt.Add(time.Millisecond), + 1, + ) + } + + header := collector.HeaderValue(startedAt.Add(2*time.Millisecond), "bypass") + if len(header) > maxHeaderLength { + t.Fatalf("header length = %d, want <= %d", len(header), maxHeaderLength) + } + if !strings.Contains(header, "total;dur=2.0") || !strings.Contains(header, "deps;dur=1.0") { + t.Fatalf("bounded header lost fixed metrics: %q", header) + } +} + +func TestCollectorConcurrentRecording(t *testing.T) { + startedAt := time.Now() + collector := New(startedAt) + ctx := WithCollector(context.Background(), collector) + + const workers = 25 + const recordsPerWorker = 100 + var wg sync.WaitGroup + wg.Add(workers) + for i := 0; i < workers; i++ { + go func() { + defer wg.Done() + for j := 0; j < recordsPerWorker; j++ { + Record(ctx, MetricDatabase, startedAt, startedAt.Add(time.Microsecond), 1) + } + }() + } + wg.Wait() + + header := HeaderValue(ctx, startedAt.Add(time.Millisecond), "bypass") + want := fmt.Sprintf(`queries=%d`, workers*recordsPerWorker) + if !strings.Contains(header, want) { + t.Fatalf("header %q does not contain %q", header, want) + } +} + +func TestContextHelpersHandleMissingCollector(t *testing.T) { + if Active(context.Background()) { + t.Fatal("context without collector reported active") + } + if got := HeaderValue(context.Background(), time.Now(), "hit"); got != "" { + t.Fatalf("HeaderValue() = %q without collector, want empty", got) + } +} diff --git a/backend/internal/pkg/servertiming/http.go b/backend/internal/pkg/servertiming/http.go new file mode 100644 index 0000000000..e326e24302 --- /dev/null +++ b/backend/internal/pkg/servertiming/http.go @@ -0,0 +1,104 @@ +package servertiming + +import ( + "context" + "net/http" + "strings" + "time" +) + +type dependencyModuleKey struct{} + +type timingRoundTripper struct { + base http.RoundTripper +} + +// WithDependencyModule overrides the safe module name used for an outbound call. +func WithDependencyModule(ctx context.Context, module string) context.Context { + if ctx == nil { + ctx = context.Background() + } + module = strings.TrimPrefix(normalizeMetricName(module), dependencyPrefix) + if module == "" { + return ctx + } + return context.WithValue(ctx, dependencyModuleKey{}, module) +} + +// WrapRoundTripper records outbound response-header latency for active requests. +func WrapRoundTripper(base http.RoundTripper) http.RoundTripper { + if base == nil { + base = http.DefaultTransport + } + if _, ok := base.(*timingRoundTripper); ok { + return base + } + return &timingRoundTripper{base: base} +} + +// InstrumentClient returns a shallow client copy with an instrumented transport. +func InstrumentClient(client *http.Client) *http.Client { + if client == nil { + client = &http.Client{} + } + copyClient := *client + copyClient.Transport = WrapRoundTripper(copyClient.Transport) + return ©Client +} + +// Do records response-header latency without changing the client's transport +// type. Use it for clients whose callers inspect or configure *http.Transport. +func Do(client *http.Client, req *http.Request) (*http.Response, error) { + if client == nil { + client = http.DefaultClient + } + if req == nil || !Active(req.Context()) { + return client.Do(req) + } + startedAt := time.Now() + response, err := client.Do(req) + RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now()) + return response, err +} + +func (t *timingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + if req == nil || !Active(req.Context()) { + return t.base.RoundTrip(req) + } + startedAt := time.Now() + response, err := t.base.RoundTrip(req) + RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now()) + return response, err +} + +func dependencyModule(req *http.Request) string { + if req != nil { + if module, ok := req.Context().Value(dependencyModuleKey{}).(string); ok && module != "" { + return module + } + } + if req == nil || req.URL == nil { + return "http" + } + host := strings.ToLower(req.URL.Hostname()) + switch { + case strings.Contains(host, "github"): + return "github" + case strings.Contains(host, "openai"): + return "openai" + case strings.Contains(host, "anthropic"): + return "anthropic" + case strings.Contains(host, "generativelanguage") || strings.Contains(host, "gemini"): + return "gemini" + case strings.Contains(host, "cloudcode") || strings.Contains(host, "antigravity"): + return "antigravity" + case strings.Contains(host, "googleapis") || strings.Contains(host, "google"): + return "google" + case strings.Contains(host, "amazonaws") || strings.Contains(host, "cloudflarestorage") || strings.Contains(host, "s3"): + return "s3" + case strings.Contains(host, "stripe") || strings.Contains(host, "airwallex") || strings.Contains(host, "alipay") || strings.Contains(host, "wechatpay") || strings.Contains(host, "paypal"): + return "payment" + default: + return "http" + } +} diff --git a/backend/internal/pkg/servertiming/http_test.go b/backend/internal/pkg/servertiming/http_test.go new file mode 100644 index 0000000000..d37f378414 --- /dev/null +++ b/backend/internal/pkg/servertiming/http_test.go @@ -0,0 +1,168 @@ +package servertiming + +import ( + "context" + "io" + "net/http" + "strings" + "testing" + "time" +) + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +type trackingBody struct { + read bool +} + +func (b *trackingBody) Read(_ []byte) (int, error) { + b.read = true + return 0, io.EOF +} + +func (b *trackingBody) Close() error { return nil } + +func TestWrapRoundTripperRecordsResponseHeaderLatency(t *testing.T) { + startedAt := time.Now() + collector := New(startedAt) + body := &trackingBody{} + baseCalled := false + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + baseCalled = true + return &http.Response{ + StatusCode: http.StatusOK, + Body: body, + Header: make(http.Header), + Request: req, + }, nil + }) + req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.github.com/repos/example/project", nil) + if err != nil { + t.Fatal(err) + } + + resp, err := WrapRoundTripper(base).RoundTrip(req) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if !baseCalled { + t.Fatal("base RoundTripper was not called") + } + if body.read { + t.Fatal("RoundTripper instrumentation read the response body; timing must stop at response headers") + } + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, `dep_github;dur=`) || !strings.Contains(header, `deps;dur=`) { + t.Fatalf("dependency metrics missing from header: %q", header) + } +} + +func TestWrapRoundTripperUsesContextModuleOverride(t *testing.T) { + collector := New(time.Now()) + ctx := WithDependencyModule(WithCollector(context.Background(), collector), "data-managementd") + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://private.example.test/path", nil) + if err != nil { + t.Fatal(err) + } + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil + }) + + if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil { + t.Fatal(err) + } + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, "dep_data_managementd") { + t.Fatalf("module override missing from header: %q", header) + } + if strings.Contains(header, "private.example") { + t.Fatalf("raw host leaked into header: %q", header) + } +} + +func TestWrapRoundTripperSkipsInactiveContext(t *testing.T) { + called := false + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + called = true + return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Header: make(http.Header), Request: req}, nil + }) + req, err := http.NewRequest(http.MethodGet, "https://api.openai.com/v1/models", nil) + if err != nil { + t.Fatal(err) + } + if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil { + t.Fatal(err) + } + if !called { + t.Fatal("inactive request did not reach base RoundTripper") + } +} + +func TestDoRecordsWithoutChangingTransportType(t *testing.T) { + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil + }) + client := &http.Client{Transport: base} + collector := New(time.Now()) + req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.openai.com/v1/models", nil) + if err != nil { + t.Fatal(err) + } + if _, err := Do(client, req); err != nil { + t.Fatal(err) + } + if _, ok := client.Transport.(roundTripFunc); !ok { + t.Fatalf("Do changed client transport type to %T", client.Transport) + } + if header := collector.HeaderValue(time.Now(), "bypass"); !strings.Contains(header, "dep_openai;dur=") { + t.Fatalf("dependency metric missing from header: %q", header) + } +} + +func TestDependencyModuleClassification(t *testing.T) { + tests := map[string]string{ + "https://api.github.com/repos/a/b": "github", + "https://api.openai.com/v1/models": "openai", + "https://api.anthropic.com/v1/messages": "anthropic", + "https://generativelanguage.googleapis.com/v1/models": "gemini", + "https://cloudcode-pa.googleapis.com/v1internal": "antigravity", + "https://storage.googleapis.com/bucket/object": "google", + "https://bucket.s3.amazonaws.com/object": "s3", + "https://api.stripe.com/v1/refunds": "payment", + "https://dependency.example.test/path": "http", + } + for rawURL, want := range tests { + req, err := http.NewRequest(http.MethodGet, rawURL, nil) + if err != nil { + t.Fatalf("NewRequest(%q): %v", rawURL, err) + } + if got := dependencyModule(req); got != want { + t.Errorf("dependencyModule(%q) = %q, want %q", rawURL, got, want) + } + } +} + +func TestClientInstrumentationDoesNotMutateOriginal(t *testing.T) { + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil + }) + original := &http.Client{Transport: base, Timeout: time.Second} + instrumented := InstrumentClient(original) + if instrumented == original { + t.Fatal("InstrumentClient returned the original client") + } + if _, ok := original.Transport.(roundTripFunc); !ok { + t.Fatalf("InstrumentClient mutated the original transport to %T", original.Transport) + } + if instrumented.Timeout != original.Timeout { + t.Fatal("InstrumentClient did not preserve client settings") + } + if WrapRoundTripper(instrumented.Transport) != instrumented.Transport { + t.Fatal("WrapRoundTripper wrapped an already instrumented transport twice") + } +} diff --git a/backend/internal/repository/backup_s3_store.go b/backend/internal/repository/backup_s3_store.go index 5d419f574b..2104e1e5d7 100644 --- a/backend/internal/repository/backup_s3_store.go +++ b/backend/internal/repository/backup_s3_store.go @@ -13,6 +13,7 @@ import ( "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/Wei-Shaw/sub2api/internal/service" ) @@ -63,12 +64,14 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader, return 0, fmt.Errorf("read body: %w", err) } + finish := servertiming.ObserveDependency(ctx, "s3") _, err = s.client.PutObject(ctx, &s3.PutObjectInput{ Bucket: &s.bucket, Key: &key, Body: bytes.NewReader(data), ContentType: &contentType, }) + finish() if err != nil { return 0, fmt.Errorf("S3 PutObject: %w", err) } @@ -76,10 +79,12 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader, } func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser, error) { + finish := servertiming.ObserveDependency(ctx, "s3") result, err := s.client.GetObject(ctx, &s3.GetObjectInput{ Bucket: &s.bucket, Key: &key, }) + finish() if err != nil { return nil, fmt.Errorf("S3 GetObject: %w", err) } @@ -87,10 +92,12 @@ func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser } func (s *S3BackupStore) Delete(ctx context.Context, key string) error { + finish := servertiming.ObserveDependency(ctx, "s3") _, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{ Bucket: &s.bucket, Key: &key, }) + finish() return err } @@ -107,9 +114,11 @@ func (s *S3BackupStore) PresignURL(ctx context.Context, key string, expiry time. } func (s *S3BackupStore) HeadBucket(ctx context.Context) error { + finish := servertiming.ObserveDependency(ctx, "s3") _, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{ Bucket: &s.bucket, }) + finish() if err != nil { return fmt.Errorf("S3 HeadBucket failed: %w", err) } diff --git a/backend/internal/repository/claude_oauth_service.go b/backend/internal/repository/claude_oauth_service.go index 5c5f27c86a..ec2d426ecb 100644 --- a/backend/internal/repository/claude_oauth_service.go +++ b/backend/internal/repository/claude_oauth_service.go @@ -276,5 +276,5 @@ func createReqClient(proxyURL string) (*req.Client, error) { client.SetProxyURL(trimmed) } - return client, nil + return instrumentReqClient(client), nil } diff --git a/backend/internal/repository/ent.go b/backend/internal/repository/ent.go index 64d321924d..3abb528e98 100644 --- a/backend/internal/repository/ent.go +++ b/backend/internal/repository/ent.go @@ -15,7 +15,7 @@ import ( "entgo.io/ent/dialect" entsql "entgo.io/ent/dialect/sql" - _ "github.com/lib/pq" // PostgreSQL 驱动，通过副作用导入注册驱动 + "github.com/lib/pq" ) // InitEnt 初始化 Ent ORM 客户端并返回客户端实例和底层的 *sql.DB。 @@ -48,9 +48,19 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) { // 使用 Ent 的 SQL 驱动打开 PostgreSQL 连接。 // dialect.Postgres 指定使用 PostgreSQL 方言进行 SQL 生成。 - drv, err := entsql.Open(dialect.Postgres, dsn) - if err != nil { - return nil, nil, err + var drv *entsql.Driver + if cfg.Server.EnableServerTiming { + connector, err := pq.NewConnector(dsn) + if err != nil { + return nil, nil, err + } + drv = entsql.OpenDB(dialect.Postgres, sql.OpenDB(newServerTimingConnector(connector))) + } else { + var err error + drv, err = entsql.Open(dialect.Postgres, dsn) + if err != nil { + return nil, nil, err + } } applyDBPoolSettings(drv.DB(), cfg) diff --git a/backend/internal/repository/http_upstream.go b/backend/internal/repository/http_upstream.go index bb079b0789..0d5afd500c 100644 --- a/backend/internal/repository/http_upstream.go +++ b/backend/internal/repository/http_upstream.go @@ -25,6 +25,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" @@ -186,7 +187,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i } // 执行请求 - resp, err := entry.client.Do(req) + resp, err := servertiming.Do(entry.client, req) if err != nil { s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err) // 请求失败，立即减少计数 @@ -243,7 +244,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco return nil, err } - resp, err := entry.client.Do(req) + resp, err := servertiming.Do(entry.client, req) if err != nil { atomic.AddInt64(&entry.inFlight, -1) atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano()) diff --git a/backend/internal/repository/redis.go b/backend/internal/repository/redis.go index 2b4ee4e636..0ead4644c1 100644 --- a/backend/internal/repository/redis.go +++ b/backend/internal/repository/redis.go @@ -21,7 +21,11 @@ import ( // 2. MinIdleConns: 保持最小空闲连接，减少冷启动延迟（默认 10） // 3. DialTimeout/ReadTimeout/WriteTimeout: 精确控制各阶段超时 func InitRedis(cfg *config.Config) *redis.Client { - return redis.NewClient(buildRedisOptions(cfg)) + client := redis.NewClient(buildRedisOptions(cfg)) + if cfg.Server.EnableServerTiming { + client.AddHook(serverTimingRedisHook{}) + } + return client } // buildRedisOptions 构建 Redis 连接选项 diff --git a/backend/internal/repository/req_client_pool.go b/backend/internal/repository/req_client_pool.go index 32501f7b19..95ab27ce32 100644 --- a/backend/internal/repository/req_client_pool.go +++ b/backend/internal/repository/req_client_pool.go @@ -2,11 +2,13 @@ package repository import ( "fmt" + "net/http" "strings" "sync" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/imroc/req/v3" ) @@ -57,6 +59,7 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) { if trimmed != "" { client.SetProxyURL(trimmed) } + client = instrumentReqClient(client) actual, _ := sharedReqClients.LoadOrStore(key, client) if c, ok := actual.(*req.Client); ok { @@ -65,6 +68,17 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) { return client, nil } +func instrumentReqClient(client *req.Client) *req.Client { + if client == nil { + return nil + } + client.GetTransport().WrapRoundTripFunc(func(rt http.RoundTripper) req.HttpRoundTripFunc { + timed := servertiming.WrapRoundTripper(rt) + return timed.RoundTrip + }) + return client +} + func buildReqClientKey(opts reqClientOptions) string { return fmt.Sprintf("%s|%s|%t|%t", strings.TrimSpace(opts.ProxyURL), diff --git a/backend/internal/repository/req_client_pool_test.go b/backend/internal/repository/req_client_pool_test.go index 9067d0129f..3a27841c5a 100644 --- a/backend/internal/repository/req_client_pool_test.go +++ b/backend/internal/repository/req_client_pool_test.go @@ -1,12 +1,17 @@ package repository import ( + "context" + "net/http" + "net/http/httptest" "reflect" + "strings" "sync" "testing" "time" "unsafe" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/imroc/req/v3" "github.com/stretchr/testify/require" ) @@ -118,3 +123,20 @@ func TestCreateGeminiReqClient_ForceHTTP2Disabled(t *testing.T) { require.NoError(t, err) require.Equal(t, "", forceHTTPVersion(t, client)) } + +func TestInstrumentReqClientRecordsDependency(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + defer server.Close() + + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + client := instrumentReqClient(req.C()) + response, err := client.R().SetContext(ctx).Get(server.URL) + require.NoError(t, err) + require.Equal(t, http.StatusNoContent, response.StatusCode) + + header := collector.HeaderValue(time.Now(), "bypass") + require.True(t, strings.Contains(header, "dep_http;dur="), header) +} diff --git a/backend/internal/repository/server_timing_redis.go b/backend/internal/repository/server_timing_redis.go new file mode 100644 index 0000000000..dba35450de --- /dev/null +++ b/backend/internal/repository/server_timing_redis.go @@ -0,0 +1,39 @@ +package repository + +import ( + "context" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/redis/go-redis/v9" +) + +type serverTimingRedisHook struct{} + +func (serverTimingRedisHook) DialHook(next redis.DialHook) redis.DialHook { + return next +} + +func (serverTimingRedisHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook { + return func(ctx context.Context, cmd redis.Cmder) error { + if !servertiming.Active(ctx) { + return next(ctx, cmd) + } + startedAt := time.Now() + err := next(ctx, cmd) + servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), 1) + return err + } +} + +func (serverTimingRedisHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook { + return func(ctx context.Context, cmds []redis.Cmder) error { + if !servertiming.Active(ctx) { + return next(ctx, cmds) + } + startedAt := time.Now() + err := next(ctx, cmds) + servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), len(cmds)) + return err + } +} diff --git a/backend/internal/repository/server_timing_redis_test.go b/backend/internal/repository/server_timing_redis_test.go new file mode 100644 index 0000000000..d1ae47e3b0 --- /dev/null +++ b/backend/internal/repository/server_timing_redis_test.go @@ -0,0 +1,63 @@ +package repository + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/redis/go-redis/v9" +) + +func TestServerTimingRedisHookRecordsCommands(t *testing.T) { + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + hook := serverTimingRedisHook{} + + process := hook.ProcessHook(func(context.Context, redis.Cmder) error { + time.Sleep(time.Millisecond) + return errors.New("redis failure") + }) + if err := process(ctx, redis.NewStringCmd(ctx, "get", "sensitive-key")); err == nil { + t.Fatal("ProcessHook did not return the underlying error") + } + + pipeline := hook.ProcessPipelineHook(func(context.Context, []redis.Cmder) error { + time.Sleep(time.Millisecond) + return nil + }) + commands := []redis.Cmder{ + redis.NewStringCmd(ctx, "get", "first-secret"), + redis.NewStringCmd(ctx, "get", "second-secret"), + redis.NewStatusCmd(ctx, "set", "third-secret", "value"), + } + if err := pipeline(ctx, commands); err != nil { + t.Fatal(err) + } + + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, `commands=4`) { + t.Fatalf("header %q does not report one command and a three-command pipeline", header) + } + if strings.Contains(header, "secret") || strings.Contains(header, "get") { + t.Fatalf("Redis command details leaked into header: %q", header) + } +} + +func TestServerTimingRedisHookSkipsInactiveContext(t *testing.T) { + called := false + hook := serverTimingRedisHook{} + process := hook.ProcessHook(func(context.Context, redis.Cmder) error { + called = true + return nil + }) + ctx := context.Background() + if err := process(ctx, redis.NewStringCmd(ctx, "ping")); err != nil { + t.Fatal(err) + } + if !called { + t.Fatal("inactive Redis command did not reach the next hook") + } +} diff --git a/backend/internal/repository/server_timing_sql.go b/backend/internal/repository/server_timing_sql.go new file mode 100644 index 0000000000..062663f08b --- /dev/null +++ b/backend/internal/repository/server_timing_sql.go @@ -0,0 +1,311 @@ +package repository + +import ( + "context" + "database/sql/driver" + "errors" + "io" + "reflect" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" +) + +type serverTimingConnector struct { + base driver.Connector +} + +func newServerTimingConnector(base driver.Connector) driver.Connector { + return &serverTimingConnector{base: base} +} + +func (c *serverTimingConnector) Connect(ctx context.Context) (driver.Conn, error) { + startedAt := time.Now() + conn, err := c.base.Connect(ctx) + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + if err != nil { + return nil, err + } + return &serverTimingConn{Conn: conn}, nil +} + +func (c *serverTimingConnector) Driver() driver.Driver { + return c.base.Driver() +} + +type serverTimingConn struct { + driver.Conn +} + +func (c *serverTimingConn) Prepare(query string) (driver.Stmt, error) { + stmt, err := c.Conn.Prepare(query) + if err != nil { + return nil, err + } + return &serverTimingStmt{Stmt: stmt}, nil +} + +func (c *serverTimingConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) { + startedAt := time.Now() + var ( + stmt driver.Stmt + err error + ) + if preparer, ok := c.Conn.(driver.ConnPrepareContext); ok { + stmt, err = preparer.PrepareContext(ctx, query) + } else { + stmt, err = c.Conn.Prepare(query) + } + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + if err != nil { + return nil, err + } + return &serverTimingStmt{Stmt: stmt}, nil +} + +func (c *serverTimingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) { + execer, ok := c.Conn.(driver.ExecerContext) + if !ok { + return nil, driver.ErrSkip + } + startedAt := time.Now() + result, err := execer.ExecContext(ctx, query, args) + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + return result, err +} + +func (c *serverTimingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { + queryer, ok := c.Conn.(driver.QueryerContext) + if !ok { + return nil, driver.ErrSkip + } + startedAt := time.Now() + rows, err := queryer.QueryContext(ctx, query, args) + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + if err != nil || rows == nil { + return rows, err + } + return newServerTimingRows(ctx, rows), nil +} + +func (c *serverTimingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) { + startedAt := time.Now() + var ( + tx driver.Tx + err error + ) + if beginner, ok := c.Conn.(driver.ConnBeginTx); ok { + tx, err = beginner.BeginTx(ctx, opts) + } else { + if opts.Isolation != driver.IsolationLevel(0) { + return nil, errors.New("driver does not support non-default isolation") + } + if opts.ReadOnly { + return nil, errors.New("driver does not support read-only transactions") + } + // The wrapper exposes ConnBeginTx, so it must retain database/sql's + // legacy fallback for drivers that only implement Conn.Begin. + tx, err = c.Conn.Begin() //nolint:staticcheck // Required driver compatibility fallback. + } + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + if err != nil || tx == nil { + return tx, err + } + return &serverTimingTx{Tx: tx, ctx: ctx}, nil +} + +func (c *serverTimingConn) Ping(ctx context.Context) error { + if pinger, ok := c.Conn.(driver.Pinger); ok { + startedAt := time.Now() + err := pinger.Ping(ctx) + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err + } + return nil +} + +func (c *serverTimingConn) ResetSession(ctx context.Context) error { + if resetter, ok := c.Conn.(driver.SessionResetter); ok { + startedAt := time.Now() + err := resetter.ResetSession(ctx) + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err + } + return nil +} + +func (c *serverTimingConn) IsValid() bool { + if validator, ok := c.Conn.(driver.Validator); ok { + return validator.IsValid() + } + return true +} + +func (c *serverTimingConn) CheckNamedValue(value *driver.NamedValue) error { + if checker, ok := c.Conn.(driver.NamedValueChecker); ok { + return checker.CheckNamedValue(value) + } + return driver.ErrSkip +} + +type serverTimingStmt struct { + driver.Stmt +} + +func (s *serverTimingStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) { + startedAt := time.Now() + var ( + result driver.Result + err error + ) + if execer, ok := s.Stmt.(driver.StmtExecContext); ok { + result, err = execer.ExecContext(ctx, args) + } else { + var values []driver.Value + values, err = namedValues(args) + if err == nil { + // The wrapper exposes StmtExecContext and must preserve the fallback + // database/sql would use for a legacy driver statement. + result, err = s.Stmt.Exec(values) //nolint:staticcheck // Required driver compatibility fallback. + } + } + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + return result, err +} + +func (s *serverTimingStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) { + startedAt := time.Now() + var ( + rows driver.Rows + err error + ) + if queryer, ok := s.Stmt.(driver.StmtQueryContext); ok { + rows, err = queryer.QueryContext(ctx, args) + } else { + var values []driver.Value + values, err = namedValues(args) + if err == nil { + // The wrapper exposes StmtQueryContext and must preserve the fallback + // database/sql would use for a legacy driver statement. + rows, err = s.Stmt.Query(values) //nolint:staticcheck // Required driver compatibility fallback. + } + } + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + if err != nil || rows == nil { + return rows, err + } + return newServerTimingRows(ctx, rows), nil +} + +func (s *serverTimingStmt) CheckNamedValue(value *driver.NamedValue) error { + if checker, ok := s.Stmt.(driver.NamedValueChecker); ok { + return checker.CheckNamedValue(value) + } + return driver.ErrSkip +} + +func namedValues(args []driver.NamedValue) ([]driver.Value, error) { + values := make([]driver.Value, len(args)) + for i, arg := range args { + if arg.Name != "" { + return nil, errors.New("named parameters are not supported") + } + values[i] = arg.Value + } + return values, nil +} + +type serverTimingRows struct { + driver.Rows + ctx context.Context +} + +func newServerTimingRows(ctx context.Context, rows driver.Rows) *serverTimingRows { + return &serverTimingRows{Rows: rows, ctx: ctx} +} + +func (r *serverTimingRows) Close() error { + startedAt := time.Now() + err := r.Rows.Close() + servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (r *serverTimingRows) Next(dest []driver.Value) error { + startedAt := time.Now() + err := r.Rows.Next(dest) + servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (r *serverTimingRows) HasNextResultSet() bool { + if rows, ok := r.Rows.(driver.RowsNextResultSet); ok { + return rows.HasNextResultSet() + } + return false +} + +func (r *serverTimingRows) NextResultSet() error { + rows, ok := r.Rows.(driver.RowsNextResultSet) + if !ok { + return io.EOF + } + startedAt := time.Now() + err := rows.NextResultSet() + servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (r *serverTimingRows) ColumnTypeScanType(index int) reflect.Type { + if rows, ok := r.Rows.(driver.RowsColumnTypeScanType); ok { + return rows.ColumnTypeScanType(index) + } + return reflect.TypeOf(new(any)).Elem() +} + +func (r *serverTimingRows) ColumnTypeDatabaseTypeName(index int) string { + if rows, ok := r.Rows.(driver.RowsColumnTypeDatabaseTypeName); ok { + return rows.ColumnTypeDatabaseTypeName(index) + } + return "" +} + +func (r *serverTimingRows) ColumnTypeLength(index int) (int64, bool) { + if rows, ok := r.Rows.(driver.RowsColumnTypeLength); ok { + return rows.ColumnTypeLength(index) + } + return 0, false +} + +func (r *serverTimingRows) ColumnTypeNullable(index int) (bool, bool) { + if rows, ok := r.Rows.(driver.RowsColumnTypeNullable); ok { + return rows.ColumnTypeNullable(index) + } + return false, false +} + +func (r *serverTimingRows) ColumnTypePrecisionScale(index int) (int64, int64, bool) { + if rows, ok := r.Rows.(driver.RowsColumnTypePrecisionScale); ok { + return rows.ColumnTypePrecisionScale(index) + } + return 0, 0, false +} + +type serverTimingTx struct { + driver.Tx + ctx context.Context +} + +func (t *serverTimingTx) Commit() error { + startedAt := time.Now() + err := t.Tx.Commit() + servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (t *serverTimingTx) Rollback() error { + startedAt := time.Now() + err := t.Tx.Rollback() + servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} diff --git a/backend/internal/repository/server_timing_sql_test.go b/backend/internal/repository/server_timing_sql_test.go new file mode 100644 index 0000000000..3a8bbbe03e --- /dev/null +++ b/backend/internal/repository/server_timing_sql_test.go @@ -0,0 +1,258 @@ +package repository + +import ( + "context" + "database/sql/driver" + "io" + "regexp" + "strconv" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" +) + +const fakeDriverDelay = 2 * time.Millisecond + +type timingFakeDriver struct{} + +func (timingFakeDriver) Open(string) (driver.Conn, error) { return newTimingFakeConn(), nil } + +type timingFakeConnector struct { + conn driver.Conn +} + +func (c timingFakeConnector) Connect(context.Context) (driver.Conn, error) { + time.Sleep(fakeDriverDelay) + return c.conn, nil +} + +func (timingFakeConnector) Driver() driver.Driver { return timingFakeDriver{} } + +type timingFakeConn struct{} + +func newTimingFakeConn() *timingFakeConn { return &timingFakeConn{} } + +func (c *timingFakeConn) Prepare(string) (driver.Stmt, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeStmt{}, nil +} + +func (c *timingFakeConn) PrepareContext(context.Context, string) (driver.Stmt, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeStmt{}, nil +} + +func (c *timingFakeConn) Close() error { return nil } + +func (c *timingFakeConn) Begin() (driver.Tx, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeTx{}, nil +} + +func (c *timingFakeConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeTx{}, nil +} + +func (c *timingFakeConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) { + time.Sleep(fakeDriverDelay) + return driver.RowsAffected(1), nil +} + +func (c *timingFakeConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil +} + +func (c *timingFakeConn) Ping(context.Context) error { + time.Sleep(fakeDriverDelay) + return nil +} + +func (c *timingFakeConn) ResetSession(context.Context) error { + time.Sleep(fakeDriverDelay) + return nil +} + +type timingFakeStmt struct{} + +func (s *timingFakeStmt) Close() error { return nil } +func (s *timingFakeStmt) NumInput() int { return -1 } + +func (s *timingFakeStmt) Exec([]driver.Value) (driver.Result, error) { + time.Sleep(fakeDriverDelay) + return driver.RowsAffected(1), nil +} + +func (s *timingFakeStmt) Query([]driver.Value) (driver.Rows, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil +} + +func (s *timingFakeStmt) ExecContext(context.Context, []driver.NamedValue) (driver.Result, error) { + time.Sleep(fakeDriverDelay) + return driver.RowsAffected(1), nil +} + +func (s *timingFakeStmt) QueryContext(context.Context, []driver.NamedValue) (driver.Rows, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil +} + +type timingFakeRows struct { + values [][]driver.Value + index int +} + +func (r *timingFakeRows) Columns() []string { return []string{"value"} } + +func (r *timingFakeRows) Close() error { + time.Sleep(fakeDriverDelay) + return nil +} + +func (r *timingFakeRows) Next(dest []driver.Value) error { + time.Sleep(fakeDriverDelay) + if r.index >= len(r.values) { + return io.EOF + } + copy(dest, r.values[r.index]) + r.index++ + return nil +} + +type timingFakeTx struct{} + +func (t *timingFakeTx) Commit() error { + time.Sleep(fakeDriverDelay) + return nil +} + +func (t *timingFakeTx) Rollback() error { + time.Sleep(fakeDriverDelay) + return nil +} + +func metricDuration(t *testing.T, header, metric string) float64 { + t.Helper() + re := regexp.MustCompile(`(?:^|, )` + regexp.QuoteMeta(metric) + `;dur=([0-9]+(?:\.[0-9]+)?)`) + match := re.FindStringSubmatch(header) + if len(match) != 2 { + t.Fatalf("metric %q missing from header %q", metric, header) + } + value, err := strconv.ParseFloat(match[1], 64) + if err != nil { + t.Fatalf("parse %s duration: %v", metric, err) + } + return value +} + +func TestServerTimingConnectorRecordsDriverCallsWithoutRowLifetime(t *testing.T) { + startedAt := time.Now() + collector := servertiming.New(startedAt) + ctx := servertiming.WithCollector(context.Background(), collector) + + wrapped := newServerTimingConnector(timingFakeConnector{conn: newTimingFakeConn()}) + rawConn, err := wrapped.Connect(ctx) + if err != nil { + t.Fatal(err) + } + conn, ok := rawConn.(*serverTimingConn) + if !ok { + t.Fatalf("Connect() returned %T, want *serverTimingConn", rawConn) + } + + if _, err := conn.ExecContext(ctx, "sensitive update", nil); err != nil { + t.Fatal(err) + } + rows, err := conn.QueryContext(ctx, "sensitive select", nil) + if err != nil { + t.Fatal(err) + } + values := make([]driver.Value, 1) + if err := rows.Next(values); err != nil { + t.Fatal(err) + } + + // Application work between row reads must remain app time. + time.Sleep(30 * time.Millisecond) + if err := rows.Next(values); err != io.EOF { + t.Fatalf("rows.Next() = %v, want EOF", err) + } + if err := rows.Close(); err != nil { + t.Fatal(err) + } + + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, `queries=2`) { + t.Fatalf("header %q does not report two SQL operations", header) + } + if strings.Contains(header, "sensitive") { + t.Fatalf("SQL text leaked into header: %q", header) + } + if app, db := metricDuration(t, header, "app"), metricDuration(t, header, "db"); app <= db { + t.Fatalf("row processing gap was counted as DB time: app=%.1fms db=%.1fms header=%q", app, db, header) + } +} + +func TestServerTimingPreparedStatementsAndTransactions(t *testing.T) { + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + conn := &serverTimingConn{Conn: newTimingFakeConn()} + + stmt, err := conn.PrepareContext(ctx, "prepare sensitive statement") + if err != nil { + t.Fatal(err) + } + timedStmt, ok := stmt.(*serverTimingStmt) + if !ok { + t.Fatalf("PrepareContext() returned %T, want *serverTimingStmt", stmt) + } + if _, err := timedStmt.ExecContext(ctx, nil); err != nil { + t.Fatal(err) + } + rows, err := timedStmt.QueryContext(ctx, nil) + if err != nil { + t.Fatal(err) + } + if err := rows.Close(); err != nil { + t.Fatal(err) + } + + tx, err := conn.BeginTx(ctx, driver.TxOptions{}) + if err != nil { + t.Fatal(err) + } + if err := tx.Commit(); err != nil { + t.Fatal(err) + } + if err := conn.Ping(ctx); err != nil { + t.Fatal(err) + } + if err := conn.ResetSession(ctx); err != nil { + t.Fatal(err) + } + + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, `queries=3`) { + t.Fatalf("header %q does not report prepare, exec, and query operations", header) + } + if metricDuration(t, header, "db") <= 0 { + t.Fatalf("DB duration was not recorded: %q", header) + } +} + +func TestNamedValuesRejectNamedParameters(t *testing.T) { + if _, err := namedValues([]driver.NamedValue{{Name: "secret", Value: 1}}); err == nil { + t.Fatal("namedValues accepted a named parameter") + } + values, err := namedValues([]driver.NamedValue{{Ordinal: 1, Value: "value"}}) + if err != nil { + t.Fatal(err) + } + if len(values) != 1 || values[0] != "value" { + t.Fatalf("namedValues() = %#v", values) + } +} diff --git a/backend/internal/server/middleware/cors.go b/backend/internal/server/middleware/cors.go index 03d5d025de..0283d53115 100644 --- a/backend/internal/server/middleware/cors.go +++ b/backend/internal/server/middleware/cors.go @@ -52,7 +52,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc { } allowHeaders := []string{ "Content-Type", "Content-Length", "Accept-Encoding", "X-CSRF-Token", "Authorization", - "accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", + "accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", "X-Admin-UI-Request", } // OpenAI Node SDK 会发送 x-stainless-* 请求头，需在 CORS 中显式放行。 openAIProperties := []string{ @@ -83,7 +83,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc { } c.Writer.Header().Set("Access-Control-Allow-Headers", allowHeadersValue) c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH") - c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag") + c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag, Server-Timing") c.Writer.Header().Set("Access-Control-Max-Age", "86400") } // 处理预检请求 diff --git a/backend/internal/server/middleware/cors_test.go b/backend/internal/server/middleware/cors_test.go index 6d0bea3608..6a61f696df 100644 --- a/backend/internal/server/middleware/cors_test.go +++ b/backend/internal/server/middleware/cors_test.go @@ -103,8 +103,10 @@ func TestCORS_AllowedOrigin_HasAllowHeaders(t *testing.T) { // 应设置 Allow-Headers、Allow-Methods 和 Max-Age assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"), "允许的 origin 应收到 Allow-Headers") + assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-Admin-UI-Request") assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"), "允许的 origin 应收到 Allow-Methods") + assert.Contains(t, w.Header().Get("Access-Control-Expose-Headers"), "Server-Timing") assert.Equal(t, "86400", w.Header().Get("Access-Control-Max-Age"), "允许的 origin 应收到 Max-Age=86400") assert.Equal(t, "https://allowed.example.com", w.Header().Get("Access-Control-Allow-Origin"), diff --git a/backend/internal/server/middleware/server_timing.go b/backend/internal/server/middleware/server_timing.go new file mode 100644 index 0000000000..2bb21071e0 --- /dev/null +++ b/backend/internal/server/middleware/server_timing.go @@ -0,0 +1,132 @@ +package middleware + +import ( + "net/http" + "strings" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/gin-gonic/gin" +) + +const ( + snapshotCacheHeader = "X-Snapshot-Cache" + usageCacheHeader = "X-Usage-Stats-Cache" +) + +type serverTimingResponseWriter struct { + gin.ResponseWriter + context *gin.Context + once sync.Once +} + +func (w *serverTimingResponseWriter) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} + +// ServerTiming collects timing only for requests made by the Admin web UI. +func ServerTiming(enabled bool) gin.HandlerFunc { + if !enabled { + return func(c *gin.Context) { + c.Next() + } + } + return func(c *gin.Context) { + if !isAdminUIRequest(c) || c.Request == nil { + c.Next() + return + } + + collector := servertiming.New(time.Now()) + c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector)) + writer := &serverTimingResponseWriter{ + ResponseWriter: c.Writer, + context: c, + } + c.Writer = writer + c.Next() + writer.finalize() + } +} + +func (w *serverTimingResponseWriter) WriteHeader(statusCode int) { + w.ResponseWriter.WriteHeader(statusCode) +} + +func (w *serverTimingResponseWriter) WriteHeaderNow() { + w.finalize() + w.ResponseWriter.WriteHeaderNow() +} + +func (w *serverTimingResponseWriter) Write(data []byte) (int, error) { + w.finalize() + return w.ResponseWriter.Write(data) +} + +func (w *serverTimingResponseWriter) WriteString(data string) (int, error) { + w.finalize() + return w.ResponseWriter.WriteString(data) +} + +func (w *serverTimingResponseWriter) Flush() { + w.finalize() + w.ResponseWriter.Flush() +} + +func (w *serverTimingResponseWriter) finalize() { + if w == nil { + return + } + w.once.Do(func() { + if value := ServerTimingHeaderValue(w.context); value != "" { + w.ResponseWriter.Header().Set(servertiming.HeaderName, value) + } + }) +} + +// ServerTimingHeaderValue returns a timing value only for an authenticated admin. +func ServerTimingHeaderValue(c *gin.Context) string { + if c == nil || c.Request == nil { + return "" + } + role, ok := GetUserRoleFromContext(c) + if !ok || role != "admin" { + return "" + } + return servertiming.HeaderValue(c.Request.Context(), time.Now(), responseCacheStatus(c.Writer.Header())) +} + +// ServerTimingResponseHeader builds the extra header map required by WebSocket upgrades. +func ServerTimingResponseHeader(c *gin.Context) http.Header { + value := ServerTimingHeaderValue(c) + if value == "" { + return nil + } + return http.Header{servertiming.HeaderName: []string{value}} +} + +func isAdminUIRequest(c *gin.Context) bool { + if c == nil || c.Request == nil || c.Request.URL == nil { + return false + } + if strings.TrimSpace(c.GetHeader(servertiming.AdminUIHeader)) == "1" { + return true + } + path := strings.TrimSpace(c.Request.URL.Path) + return path == "/api/v1/admin" || strings.HasPrefix(path, "/api/v1/admin/") +} + +func responseCacheStatus(header http.Header) string { + for _, name := range []string{snapshotCacheHeader, usageCacheHeader} { + switch strings.ToLower(strings.TrimSpace(header.Get(name))) { + case "hit": + return "hit" + case "miss": + return "miss" + case "bypass": + return "bypass" + } + } + return "bypass" +} diff --git a/backend/internal/server/middleware/server_timing_test.go b/backend/internal/server/middleware/server_timing_test.go new file mode 100644 index 0000000000..c064840ece --- /dev/null +++ b/backend/internal/server/middleware/server_timing_test.go @@ -0,0 +1,188 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/gin-gonic/gin" +) + +func runServerTimingRequest( + t *testing.T, + enabled bool, + path string, + marker string, + role string, + handler gin.HandlerFunc, +) *httptest.ResponseRecorder { + t.Helper() + gin.SetMode(gin.TestMode) + engine := gin.New() + engine.Use(ServerTiming(enabled)) + engine.Any("/*path", func(c *gin.Context) { + if role != "" { + c.Set(string(ContextKeyUserRole), role) + } + handler(c) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path, nil) + if marker != "" { + request.Header.Set(servertiming.AdminUIHeader, marker) + } + engine.ServeHTTP(recorder, request) + return recorder +} + +func TestServerTimingScopesAndRoleGate(t *testing.T) { + tests := []struct { + name string + enabled bool + path string + marker string + role string + wantHeader bool + }{ + {name: "disabled", enabled: false, path: "/api/v1/admin/users", role: "admin"}, + {name: "admin API path", enabled: true, path: "/api/v1/admin/users", role: "admin", wantHeader: true}, + {name: "shared API marked by admin UI", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "admin", wantHeader: true}, + {name: "non admin role", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "user"}, + {name: "unauthenticated public request", enabled: true, path: "/api/v1/settings/public", marker: "1"}, + {name: "unmarked shared API", enabled: true, path: "/api/v1/groups/available", role: "admin"}, + {name: "invalid marker", enabled: true, path: "/api/v1/groups/available", marker: "true", role: "admin"}, + {name: "admin prefix boundary", enabled: true, path: "/api/v1/administrator", role: "admin"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := runServerTimingRequest(t, tt.enabled, tt.path, tt.marker, tt.role, func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + header := recorder.Header().Get(servertiming.HeaderName) + if tt.wantHeader && header == "" { + t.Fatalf("%s header missing", servertiming.HeaderName) + } + if !tt.wantHeader && header != "" { + t.Fatalf("unexpected %s header: %q", servertiming.HeaderName, header) + } + if header != "" && (!strings.Contains(header, "total;dur=") || !strings.Contains(header, `cache;desc="bypass"`)) { + t.Fatalf("incomplete timing header: %q", header) + } + }) + } +} + +func TestServerTimingCollectorIsRequestScoped(t *testing.T) { + active := false + recorder := runServerTimingRequest(t, true, "/api/v1/keys", "1", "admin", func(c *gin.Context) { + active = servertiming.Active(c.Request.Context()) + c.Status(http.StatusNoContent) + }) + if !active { + t.Fatal("collector was not attached to marked request context") + } + if recorder.Header().Get(servertiming.HeaderName) == "" { + t.Fatal("timing header missing from status-only response") + } +} + +func TestServerTimingFinalizesBeforeEarlyCommit(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/stream", "", "admin", func(c *gin.Context) { + c.Status(http.StatusAccepted) + c.Writer.WriteHeaderNow() + }) + if got := recorder.Header().Get(servertiming.HeaderName); got == "" { + t.Fatal("timing header was not written before response commit") + } +} + +func TestServerTimingFinalizesOnFlush(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/export", "", "admin", func(c *gin.Context) { + c.Writer.Flush() + }) + if got := recorder.Header().Get(servertiming.HeaderName); got == "" { + t.Fatal("timing header was not written before stream flush") + } +} + +func TestServerTimingStatusResponses(t *testing.T) { + tests := []struct { + name string + status int + }{ + {name: "not modified", status: http.StatusNotModified}, + {name: "internal error", status: http.StatusInternalServerError}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/test", "", "admin", func(c *gin.Context) { + c.Status(tt.status) + }) + if recorder.Code != tt.status { + t.Fatalf("status = %d, want %d", recorder.Code, tt.status) + } + if got := recorder.Header().Get(servertiming.HeaderName); got == "" { + t.Fatalf("timing header missing from status %d response", tt.status) + } + }) + } +} + +func TestServerTimingResponseWriterUnwraps(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + baseWriter := c.Writer + writer := &serverTimingResponseWriter{ResponseWriter: baseWriter} + if got := writer.Unwrap(); got != baseWriter { + t.Fatalf("Unwrap() = %T, want original Gin writer", got) + } +} + +func TestServerTimingCacheOutcome(t *testing.T) { + tests := []struct { + name string + headerName string + value string + want string + }{ + {name: "snapshot hit", headerName: snapshotCacheHeader, value: "hit", want: "hit"}, + {name: "usage miss", headerName: usageCacheHeader, value: "MISS", want: "miss"}, + {name: "invalid", headerName: snapshotCacheHeader, value: "stale", want: "bypass"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/dashboard", "", "admin", func(c *gin.Context) { + c.Header(tt.headerName, tt.value) + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + want := `cache;desc="` + tt.want + `"` + if got := recorder.Header().Get(servertiming.HeaderName); !strings.Contains(got, want) { + t.Fatalf("timing header %q does not contain %q", got, want) + } + }) + } +} + +func TestServerTimingResponseHeaderForWebSocket(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ws/qps", nil) + collector := servertiming.New(time.Now()) + c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector)) + c.Set(string(ContextKeyUserRole), "admin") + + header := ServerTimingResponseHeader(c) + if header.Get(servertiming.HeaderName) == "" { + t.Fatal("WebSocket response header missing timing value") + } + + c.Set(string(ContextKeyUserRole), "user") + if got := ServerTimingResponseHeader(c); got != nil { + t.Fatalf("non-admin WebSocket received timing header: %#v", got) + } +} diff --git a/backend/internal/server/router.go b/backend/internal/server/router.go index 3d86373779..5fc70149fe 100644 --- a/backend/internal/server/router.go +++ b/backend/internal/server/router.go @@ -60,6 +60,7 @@ func SetupRouter( } return nil })) + r.Use(middleware2.ServerTiming(cfg.Server.EnableServerTiming)) // Serve embedded frontend with settings injection if available if web.HasEmbeddedFrontend() { diff --git a/backend/internal/service/channel_monitor_checker.go b/backend/internal/service/channel_monitor_checker.go index 7fb829a3cb..889b2bbed7 100644 --- a/backend/internal/service/channel_monitor_checker.go +++ b/backend/internal/service/channel_monitor_checker.go @@ -13,6 +13,7 @@ import ( "strings" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/tidwall/gjson" ) @@ -34,7 +35,7 @@ func newSSRFSafeHTTPClient(timeout time.Duration) *http.Client { TLSHandshakeTimeout: monitorTLSHandshakeTimeout, ResponseHeaderTimeout: monitorResponseHeaderTimeout, } - return &http.Client{Timeout: timeout, Transport: tr} + return &http.Client{Timeout: timeout, Transport: servertiming.WrapRoundTripper(tr)} } // CheckOptions 承载一次检测的自定义入参。 diff --git a/backend/internal/service/content_moderation.go b/backend/internal/service/content_moderation.go index 6d3b91d205..f633c8ad17 100644 --- a/backend/internal/service/content_moderation.go +++ b/backend/internal/service/content_moderation.go @@ -22,6 +22,7 @@ import ( infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" ) const ( @@ -561,7 +562,7 @@ func NewContentModerationService( userRepo: userRepo, authCacheInvalidator: authCacheInvalidator, emailService: emailService, - httpClient: &http.Client{}, + httpClient: servertiming.InstrumentClient(nil), workerCount: maxContentModerationWorkerCount, asyncQueue: make(chan contentModerationTask, maxContentModerationQueueSize), keyHealth: make(map[string]*contentModerationKeyHealth), diff --git a/backend/internal/service/payment_order.go b/backend/internal/service/payment_order.go index 04feb8002a..9b7ee08990 100644 --- a/backend/internal/service/payment_order.go +++ b/backend/internal/service/payment_order.go @@ -16,6 +16,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/payment" "github.com/Wei-Shaw/sub2api/internal/payment/provider" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/shopspring/decimal" ) @@ -445,7 +446,9 @@ func (s *PaymentService) invokeProvider(ctx context.Context, order *dbent.Paymen IsMobile: req.IsMobile, ReturnURL: providerReturnURL, }, sel, outTradeNo, payAmountStr, subject) + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") pr, err := prov.CreatePayment(ctx, providerReq) + finishProviderCall() if err != nil { slog.Error("[PaymentService] CreatePayment failed", "provider", sel.ProviderKey, "instance", sel.InstanceID, "error", err) if appErr := new(infraerrors.ApplicationError); errors.As(err, &appErr) { diff --git a/backend/internal/service/payment_order_lifecycle.go b/backend/internal/service/payment_order_lifecycle.go index 8ed18797dd..46a2e00605 100644 --- a/backend/internal/service/payment_order_lifecycle.go +++ b/backend/internal/service/payment_order_lifecycle.go @@ -13,6 +13,7 @@ import ( "github.com/Wei-Shaw/sub2api/ent/paymentorder" "github.com/Wei-Shaw/sub2api/internal/payment" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" ) // --- Cancel & Expire --- @@ -157,7 +158,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym if queryRef == "" { return "" } + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := prov.QueryOrder(ctx, queryRef) + finishProviderCall() if err != nil { slog.Warn("query upstream failed", "orderID", o.ID, "error", err) return "" @@ -199,7 +202,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym return "" } if cp, ok := prov.(payment.CancelableProvider); ok { + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") _ = cp.CancelPayment(ctx, queryRef) + finishProviderCall() } return "" } @@ -208,7 +213,9 @@ func requeryPaidOrderOnce(ctx context.Context, prov payment.Provider, queryRef s if prov == nil || strings.TrimSpace(queryRef) == "" { return nil, false } + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := prov.QueryOrder(ctx, queryRef) + finishProviderCall() if err != nil { slog.Warn("query upstream retry failed", "queryRef", queryRef, "error", err) return nil, false diff --git a/backend/internal/service/payment_refund.go b/backend/internal/service/payment_refund.go index 91822680ed..bc073a2c34 100644 --- a/backend/internal/service/payment_refund.go +++ b/backend/internal/service/payment_refund.go @@ -19,6 +19,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/payment" "github.com/Wei-Shaw/sub2api/internal/payment/provider" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" ) // --- Refund Flow --- @@ -347,12 +348,14 @@ func (s *PaymentService) gwRefund(ctx context.Context, p *RefundPlan) (*payment. }) return nil, err } + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := prov.Refund(ctx, payment.RefundRequest{ TradeNo: p.Order.PaymentTradeNo, OrderID: p.Order.OutTradeNo, Amount: formatGatewayRefundAmount(p.GatewayAmount, p.Order), Reason: p.Reason, }) + finishProviderCall() if err != nil { if resp != nil && strings.TrimSpace(resp.Status) == payment.ProviderStatusPending { return resp, nil @@ -417,12 +420,14 @@ func (s *PaymentService) QueryAndFinalizeRefund(ctx context.Context, oid int64) } pendingDetail := s.latestRefundPendingDetail(ctx, oid) + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := queryProvider.QueryRefund(ctx, payment.RefundQueryRequest{ TradeNo: o.PaymentTradeNo, OrderID: o.OutTradeNo, RefundID: pendingDetail.RefundID, Amount: formatGatewayRefundAmount(o.RefundAmount, o), }) + finishProviderCall() if err != nil { return nil, fmt.Errorf("query refund: %w", err) } diff --git a/backend/internal/service/vertex_service_account.go b/backend/internal/service/vertex_service_account.go index 256695ded5..7ccbeee43c 100644 --- a/backend/internal/service/vertex_service_account.go +++ b/backend/internal/service/vertex_service_account.go @@ -18,6 +18,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/golang-jwt/jwt/v5" ) @@ -195,7 +196,7 @@ func vertexServiceAccountProxyURL(account *Account) string { func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) { proxyURL = strings.TrimSpace(proxyURL) if proxyURL == "" { - return &http.Client{Timeout: 15 * time.Second}, nil + return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second}), nil } _, parsedProxy, err := proxyurl.Parse(proxyURL) @@ -211,7 +212,7 @@ func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) { if err := proxyutil.ConfigureTransportProxy(transport, parsedProxy); err != nil { return nil, err } - return &http.Client{Timeout: 15 * time.Second, Transport: transport}, nil + return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second, Transport: transport}), nil } func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAccountKey, proxyURL string) (string, time.Duration, error) { diff --git a/backend/internal/service/vertex_service_account_test.go b/backend/internal/service/vertex_service_account_test.go index d77a1988e9..68c756eaa2 100644 --- a/backend/internal/service/vertex_service_account_test.go +++ b/backend/internal/service/vertex_service_account_test.go @@ -13,6 +13,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) @@ -101,6 +102,24 @@ func TestVertexServiceAccountProxyURL(t *testing.T) { require.Empty(t, vertexServiceAccountProxyURL(&Account{ProxyID: &proxyID})) } +func TestVertexServiceAccountHTTPClientRecordsDependency(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + defer server.Close() + + client, err := newVertexServiceAccountHTTPClient("") + require.NoError(t, err) + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + request, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + response, err := client.Do(request) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + require.Contains(t, collector.HeaderValue(time.Now(), "bypass"), "dep_http;dur=") +} + func TestExchangeVertexServiceAccountTokenUsesProxy(t *testing.T) { privateKey, err := rsa.GenerateKey(rand.Reader, 2048) require.NoError(t, err) diff --git a/deploy/.env.example b/deploy/.env.example index f68257df9f..1d1a6ec881 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -23,6 +23,9 @@ SERVER_PORT=8080 # Server mode: release or debug SERVER_MODE=release +# Return Server-Timing for authenticated requests made by the Admin web UI +ENABLE_SERVER_TIMING=false + # Apple container image overrides (ignored by Docker Compose). Pin release tags # or digests for repeatable operator-managed deployments. APPLE_CONTAINER_SUB2API_IMAGE=weishaw/sub2api:latest diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 954263d9c4..dfc584a59a 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -20,6 +20,9 @@ server: # Mode: "debug" for development, "release" for production # 运行模式："debug" 用于开发，"release" 用于生产环境 mode: "release" + # Return Server-Timing for authenticated requests made by the Admin web UI + # 为管理端 Web 页面发出的已认证请求返回 Server-Timing + enable_server_timing: false # Frontend base URL used to generate external links in emails (e.g. password reset) # 用于生成邮件中的外部链接（例如：重置密码链接）的前端基础地址 # Example: "https://example.com" diff --git a/deploy/docker-compose.dev.yml b/deploy/docker-compose.dev.yml index 6f5b3f56f3..43f5dd3f60 100644 --- a/deploy/docker-compose.dev.yml +++ b/deploy/docker-compose.dev.yml @@ -26,6 +26,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=debug + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} - DATABASE_HOST=postgres - DATABASE_PORT=5432 diff --git a/deploy/docker-compose.local.yml b/deploy/docker-compose.local.yml index 042752e857..5fb161603b 100644 --- a/deploy/docker-compose.local.yml +++ b/deploy/docker-compose.local.yml @@ -51,6 +51,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} # ======================================================================= diff --git a/deploy/docker-compose.standalone.yml b/deploy/docker-compose.standalone.yml index 2e1d335624..40ed4751d6 100644 --- a/deploy/docker-compose.standalone.yml +++ b/deploy/docker-compose.standalone.yml @@ -37,6 +37,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} # ======================================================================= diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 22713c59aa..6aecdcfa5a 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -47,6 +47,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} # ======================================================================= diff --git a/frontend/src/api/__tests__/adminUIRequest.spec.ts b/frontend/src/api/__tests__/adminUIRequest.spec.ts new file mode 100644 index 0000000000..9064a52f1a --- /dev/null +++ b/frontend/src/api/__tests__/adminUIRequest.spec.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from 'vitest' + +import { + ADMIN_UI_REQUEST_HEADER, + shouldMarkAdminUIRequest, +} from '@/api/adminUIRequest' + +describe('Admin UI request marker', () => { + it('uses the stable request header name', () => { + expect(ADMIN_UI_REQUEST_HEADER).toBe('X-Admin-UI-Request') + }) + + it.each([ + '/admin', + '/admin/users', + '/api/v1/admin', + '/api/v1/admin/accounts?status=active', + 'https://api.example.test/api/v1/admin/dashboard', + ])('marks Admin API request %s before page navigation', (requestURL) => { + expect(shouldMarkAdminUIRequest(requestURL, '/login')).toBe(true) + }) + + it.each(['/keys', '/groups/available', '/auth/me', '/announcements'])( + 'marks shared request %s while an Admin page is active', + (requestURL) => { + expect(shouldMarkAdminUIRequest(requestURL, '/admin/dashboard')).toBe(true) + } + ) + + it.each([ + ['/keys', '/dashboard'], + ['/api/v1/administer', '/dashboard'], + ['/keys', '/administrator'], + ['', '/'], + ])('does not mark request %s on page %s', (requestURL, pagePath) => { + expect(shouldMarkAdminUIRequest(requestURL, pagePath)).toBe(false) + }) +}) diff --git a/frontend/src/api/__tests__/client.spec.ts b/frontend/src/api/__tests__/client.spec.ts index a0a05410d4..b275cca34b 100644 --- a/frontend/src/api/__tests__/client.spec.ts +++ b/frontend/src/api/__tests__/client.spec.ts @@ -12,6 +12,7 @@ describe('API Client', () => { beforeEach(async () => { localStorage.clear() + window.history.replaceState({}, '', '/') // 每次测试重新导入以获取干净的模块状态 vi.resetModules() const mod = await import('@/api/client') @@ -120,6 +121,55 @@ describe('API Client', () => { const config = adapter.mock.calls[0][0] expect(config.withCredentials).toBe(true) }) + + it('Admin API 在进入管理页面前也带 Admin UI 标记', async () => { + const adapter = vi.fn().mockResolvedValue({ + status: 200, + data: { code: 0, data: {} }, + headers: {}, + config: {}, + statusText: 'OK', + }) + apiClient.defaults.adapter = adapter + + await apiClient.get('/admin/users') + + const config = adapter.mock.calls[0][0] + expect(config.headers.get('X-Admin-UI-Request')).toBe('1') + }) + + it('管理页面调用共享 API 时带 Admin UI 标记', async () => { + window.history.replaceState({}, '', '/admin/dashboard') + const adapter = vi.fn().mockResolvedValue({ + status: 200, + data: { code: 0, data: {} }, + headers: {}, + config: {}, + statusText: 'OK', + }) + apiClient.defaults.adapter = adapter + + await apiClient.get('/groups/available') + + const config = adapter.mock.calls[0][0] + expect(config.headers.get('X-Admin-UI-Request')).toBe('1') + }) + + it('普通用户页面调用共享 API 时不带 Admin UI 标记', async () => { + const adapter = vi.fn().mockResolvedValue({ + status: 200, + data: { code: 0, data: {} }, + headers: {}, + config: {}, + statusText: 'OK', + }) + apiClient.defaults.adapter = adapter + + await apiClient.get('/groups/available') + + const config = adapter.mock.calls[0][0] + expect(config.headers.get('X-Admin-UI-Request')).toBeFalsy() + }) }) // --- 响应拦截器 --- diff --git a/frontend/src/api/adminUIRequest.ts b/frontend/src/api/adminUIRequest.ts new file mode 100644 index 0000000000..2d60e2987d --- /dev/null +++ b/frontend/src/api/adminUIRequest.ts @@ -0,0 +1,27 @@ +export const ADMIN_UI_REQUEST_HEADER = 'X-Admin-UI-Request' + +function isAdminPath(path: string): boolean { + return ( + path === '/admin' || + path.startsWith('/admin/') || + path === '/api/v1/admin' || + path.startsWith('/api/v1/admin/') + ) +} + +function requestPath(rawURL: string): string { + const value = rawURL.trim() + if (!value) return '' + try { + const origin = typeof window !== 'undefined' ? window.location.origin : 'http://localhost' + return new URL(value, origin).pathname + } catch { + return value.split(/[?#]/, 1)[0] + } +} + +export function shouldMarkAdminUIRequest(requestURL: string, pagePath?: string): boolean { + const currentPath = + pagePath ?? (typeof window !== 'undefined' ? window.location.pathname : '') + return isAdminPath(requestPath(requestURL)) || isAdminPath(currentPath) +} diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index 5df969f188..a2b4d2f650 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -6,6 +6,7 @@ import axios, { AxiosInstance, AxiosError, InternalAxiosRequestConfig, AxiosResponse } from 'axios' import type { ApiResponse } from '@/types' import { getLocale } from '@/i18n' +import { ADMIN_UI_REQUEST_HEADER, shouldMarkAdminUIRequest } from './adminUIRequest' import { getAPIBaseURL } from './url' export { buildApiUrl, buildGatewayUrl } from './url' @@ -74,6 +75,10 @@ apiClient.interceptors.request.use( config.params.timezone = getUserTimezone() } + if (config.headers && shouldMarkAdminUIRequest(String(config.url || ''))) { + config.headers[ADMIN_UI_REQUEST_HEADER] = '1' + } + return config }, (error) => { diff --git a/frontend/src/components/admin/account/AccountTestModal.vue b/frontend/src/components/admin/account/AccountTestModal.vue index 0a0e3dd9ae..0a8f853ebb 100644 --- a/frontend/src/components/admin/account/AccountTestModal.vue +++ b/frontend/src/components/admin/account/AccountTestModal.vue @@ -250,6 +250,7 @@ import TextArea from '@/components/common/TextArea.vue' import { Icon } from '@/components/icons' import { useClipboard } from '@/composables/useClipboard' import { buildApiUrl } from '@/api/client' +import { ADMIN_UI_REQUEST_HEADER } from '@/api/adminUIRequest' import { adminAPI } from '@/api/admin' import type { Account, ClaudeModel } from '@/types' @@ -438,7 +439,8 @@ const startTest = async () => { method: 'POST', headers: { Authorization: `Bearer ${localStorage.getItem('auth_token')}`, - 'Content-Type': 'application/json' + 'Content-Type': 'application/json', + [ADMIN_UI_REQUEST_HEADER]: '1' }, body: JSON.stringify(requestBody), signal: abortController.signal