From 3f702084d20b3c3a3ea916f3110e99b22bda60b3 Mon Sep 17 00:00:00 2001 From: Chia Date: Thu, 6 Aug 2026 15:58:57 +1200 Subject: feat: complete commercial developer workflows Add tenant-safe usage observability, prepaid billing controls, API key lifecycle management, Embeddings metering, configurable billing alerts, and resilient provider health propagation. Harden Stripe failure handling, migrations, readiness, and the authenticated control-plane UI with end-to-end verification evidence. --- internal/providerhealth/prober.go | 171 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 171 insertions(+) create mode 100644 internal/providerhealth/prober.go (limited to 'internal/providerhealth/prober.go') diff --git a/internal/providerhealth/prober.go b/internal/providerhealth/prober.go new file mode 100644 index 0000000..820fd04 --- /dev/null +++ b/internal/providerhealth/prober.go @@ -0,0 +1,171 @@ +package providerhealth + +import ( + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "strings" + "time" + + "aigw/internal/catalog" + "aigw/internal/domain" +) + +type ProbeMetrics interface { + ProviderProbe(success bool) +} + +type ProbeOptions struct { + Enabled bool + Interval time.Duration + Timeout time.Duration + Catalog *catalog.Catalog + Tracker *Tracker + Metrics ProbeMetrics + Logger *slog.Logger + Client *http.Client +} + +type Prober struct { + enabled bool + interval time.Duration + timeout time.Duration + catalog *catalog.Catalog + tracker *Tracker + metrics ProbeMetrics + logger *slog.Logger + client *http.Client +} + +type probeTarget struct { + provider domain.Provider + keys []RouteKey +} + +func NewProber(options ProbeOptions) *Prober { + if options.Interval <= 0 { + options.Interval = 30 * time.Second + } + if options.Timeout <= 0 { + options.Timeout = 5 * time.Second + } + if options.Logger == nil { + options.Logger = slog.Default() + } + if options.Client == nil { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.Proxy = http.ProxyFromEnvironment + options.Client = &http.Client{Transport: transport, Timeout: options.Timeout} + } + return &Prober{enabled: options.Enabled, interval: options.Interval, timeout: options.Timeout, + catalog: options.Catalog, tracker: options.Tracker, metrics: options.Metrics, logger: options.Logger, client: options.Client} +} + +func (p *Prober) Run(ctx context.Context) { + if p == nil || !p.enabled || p.catalog == nil || p.tracker == nil { + return + } + p.ProbeOnce(ctx) + ticker := time.NewTicker(p.interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + p.ProbeOnce(ctx) + } + } +} + +func (p *Prober) ProbeOnce(ctx context.Context) { + for _, target := range p.targets() { + if ctx.Err() != nil { + return + } + p.probe(ctx, target) + } +} + +func (p *Prober) targets() []probeTarget { + providerIndexes := make(map[string]int) + seen := make(map[RouteKey]struct{}) + result := make([]probeTarget, 0) + for _, model := range p.catalog.AllModels() { + for _, route := range model.Routes { + key := RouteKey{ModelID: model.ID, ProviderID: route.Provider.ID, WireAPI: route.Provider.EffectiveWireAPI()} + if key.ProviderID == "" { + continue + } + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + index, exists := providerIndexes[route.Provider.ID] + if !exists { + index = len(result) + providerIndexes[route.Provider.ID] = index + result = append(result, probeTarget{provider: route.Provider}) + } + result[index].keys = append(result[index].keys, key) + } + } + return result +} + +func (p *Prober) probe(parent context.Context, target probeTarget) { + ctx, cancel := context.WithTimeout(parent, p.timeout) + defer cancel() + request, err := http.NewRequestWithContext(ctx, http.MethodGet, providerModelsURL(target.provider), nil) + if err != nil { + p.observe(target, 0, 0, false) + return + } + request.Header.Set("Accept", "application/json") + request.Header.Set("User-Agent", "aigw-health/0.1") + if target.provider.Protocol == domain.ProtocolAnthropic { + request.Header.Set("x-api-key", target.provider.APIKey) + request.Header.Set("anthropic-version", "2023-06-01") + } else { + request.Header.Set("Authorization", "Bearer "+target.provider.APIKey) + } + started := time.Now() + response, err := p.client.Do(request) + latency := time.Since(started) + if err != nil { + p.observe(target, 0, latency, false) + p.logger.Warn("provider_probe_failed", "provider_id", target.provider.ID, "error", err) + return + } + body, readErr := io.ReadAll(io.LimitReader(response.Body, 1<<20)) + _ = response.Body.Close() + success := response.StatusCode >= 200 && response.StatusCode < 300 && readErr == nil && + strings.HasPrefix(strings.ToLower(response.Header.Get("Content-Type")), "application/json") && validModelsEnvelope(body) + p.observe(target, response.StatusCode, latency, success) + if !success { + p.logger.Warn("provider_probe_failed", "provider_id", target.provider.ID, "status_code", response.StatusCode) + } +} + +func (p *Prober) observe(target probeTarget, statusCode int, latency time.Duration, success bool) { + for _, key := range target.keys { + p.tracker.Observe(key, Observation{StatusCode: statusCode, Latency: latency, Failed: !success, Active: true}) + } + if p.metrics != nil { + p.metrics.ProviderProbe(success) + } +} + +func validModelsEnvelope(body []byte) bool { + var envelope struct { + Data []json.RawMessage `json:"data"` + } + return json.Unmarshal(body, &envelope) == nil && envelope.Data != nil +} + +func providerModelsURL(provider domain.Provider) string { + baseURL := strings.TrimRight(provider.BaseURL, "/") + return baseURL + "/models" +} -- cgit v1.2.3