From a15c04cac13211054ad3aa32956337bddd09ac6c Mon Sep 17 00:00:00 2001 From: Amirhossein Khalili Date: Fri, 21 Aug 2026 01:48:47 +0330 Subject: [PATCH] fix(growth): harden AI jobs and cost controls --- internal/api/growth.go | 27 ++++- internal/api/growth_test.go | 42 +++++++ internal/api/server.go | 25 +++-- internal/growthreport/avalai.go | 103 ++++++++++++------ internal/growthreport/avalai_test.go | 28 +++++ internal/growthreport/coordinator.go | 12 +- internal/growthreport/coordinator_test.go | 2 +- internal/growthreport/model.go | 3 + internal/growthreport/orchestrator.go | 59 ++++++---- internal/growthreport/repository.go | 32 ++++-- .../repository_integration_test.go | 81 ++++++++++++++ internal/growthreport/service_test.go | 18 ++- migrations/000002_growth_analyst.sql | 12 ++ 13 files changed, 366 insertions(+), 78 deletions(-) create mode 100644 internal/growthreport/repository_integration_test.go diff --git a/internal/api/growth.go b/internal/api/growth.go index b4aa6ae..4ab6e28 100644 --- a/internal/api/growth.go +++ b/internal/api/growth.go @@ -131,6 +131,17 @@ func (s *Server) growthReportEvents(w http.ResponseWriter, r *http.Request) { return } merchantID := sessionFrom(r.Context()).MerchantID + slots := s.growthSSESlots + if slots == nil { + slots = defaultGrowthSSESlots + } + select { + case slots <- struct{}{}: + defer func() { <-slots }() + default: + s.problem(w, r, http.StatusTooManyRequests, "RATE_LIMITED", "Too many growth report streams") + return + } record, err := s.growth.Get(r.Context(), merchantID, id) if errors.Is(err, growthreport.ErrNotFound) { s.problem(w, r, 404, "RESOURCE_NOT_FOUND", "Resource not found") @@ -148,7 +159,17 @@ func (s *Server) growthReportEvents(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") w.Header().Set("Cache-Control", "no-cache, no-transform") w.Header().Set("X-Accel-Buffering", "no") - ticker := time.NewTicker(2 * time.Second) + maxDuration := s.growthSSEMaxDuration + if maxDuration <= 0 { + maxDuration = defaultGrowthSSEMaxDuration + } + pollInterval := s.growthSSEPollInterval + if pollInterval <= 0 { + pollInterval = defaultGrowthSSEPollInterval + } + streamContext, cancel := context.WithTimeout(r.Context(), maxDuration) + defer cancel() + ticker := time.NewTicker(pollInterval) defer ticker.Stop() heartbeat := time.NewTicker(15 * time.Second) defer heartbeat.Stop() @@ -165,10 +186,10 @@ func (s *Server) growthReportEvents(w http.ResponseWriter, r *http.Request) { return } select { - case <-r.Context().Done(): + case <-streamContext.Done(): return case <-ticker.C: - record, err = s.growth.Get(r.Context(), merchantID, id) + record, err = s.growth.Get(streamContext, merchantID, id) if err != nil { return } diff --git a/internal/api/growth_test.go b/internal/api/growth_test.go index 6479d07..0103e0f 100644 --- a/internal/api/growth_test.go +++ b/internal/api/growth_test.go @@ -127,3 +127,45 @@ func TestGrowthSSEAuthorizesBeforeSettingStreamHeaders(t *testing.T) { t.Fatalf("status=%d content-type=%q", response.Code, response.Header().Get("Content-Type")) } } + +func TestGrowthSSERejectsWhenGlobalCapacityIsFull(t *testing.T) { + t.Parallel() + slots := make(chan struct{}, 1) + slots <- struct{}{} + reads := 0 + service := fakeGrowthService{get: func(int64, string) (growthreport.Record, error) { + reads++ + return growthreport.Record{Report: growthreport.Report{ReportID: "growth_0123456789abcdef01234567", Status: growthreport.StatusQueued}}, nil + }} + server := &Server{logger: slog.New(slog.NewTextHandler(io.Discard, nil)), growth: service, growthSSESlots: slots} + request := httptest.NewRequest("GET", "/", nil) + request.SetPathValue("reportId", "growth_0123456789abcdef01234567") + request = request.WithContext(context.WithValue(request.Context(), sessionKey{}, identity.Session{MerchantID: 7})) + response := httptest.NewRecorder() + server.growthReportEvents(response, request) + if response.Code != http.StatusTooManyRequests || reads != 0 { + t.Fatalf("status=%d database reads=%d", response.Code, reads) + } +} + +func TestGrowthSSEEndsAtConfiguredDeadline(t *testing.T) { + t.Parallel() + reads := 0 + service := fakeGrowthService{get: func(int64, string) (growthreport.Record, error) { + reads++ + return growthreport.Record{Report: growthreport.Report{ReportID: "growth_0123456789abcdef01234567", Status: growthreport.StatusQueued, UpdatedAt: time.Unix(int64(reads), 0)}}, nil + }} + server := &Server{logger: slog.New(slog.NewTextHandler(io.Discard, nil)), growth: service, growthSSESlots: make(chan struct{}, 1), growthSSEMaxDuration: 5 * time.Millisecond, growthSSEPollInterval: time.Millisecond} + request := httptest.NewRequest("GET", "/", nil) + request.SetPathValue("reportId", "growth_0123456789abcdef01234567") + request = request.WithContext(context.WithValue(request.Context(), sessionKey{}, identity.Session{MerchantID: 7})) + response := httptest.NewRecorder() + started := time.Now() + server.growthReportEvents(response, request) + if elapsed := time.Since(started); elapsed > 250*time.Millisecond { + t.Fatalf("stream remained open for %s", elapsed) + } + if reads < 1 || reads > 10 { + t.Fatalf("database reads=%d", reads) + } +} diff --git a/internal/api/server.go b/internal/api/server.go index 0eae150..c065fbf 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -23,22 +23,33 @@ import ( const maxExportBytes = 50 << 20 +const ( + defaultGrowthSSECapacity = 64 + defaultGrowthSSEMaxDuration = 5 * time.Minute + defaultGrowthSSEPollInterval = 2 * time.Second +) + +var defaultGrowthSSESlots = make(chan struct{}, defaultGrowthSSECapacity) + var errExportBytesExceeded = errors.New("export exceeds byte limit") type Server struct { - identity *identity.Service - analytics *analytics.Service - growth growthReportService - logger *slog.Logger - frontendOrigin string - cookieSecure bool + identity *identity.Service + analytics *analytics.Service + growth growthReportService + logger *slog.Logger + frontendOrigin string + cookieSecure bool + growthSSESlots chan struct{} + growthSSEMaxDuration time.Duration + growthSSEPollInterval time.Duration } type sessionKey struct{} type tokenKey struct{} func NewServer(identityService *identity.Service, analyticsService *analytics.Service, logger *slog.Logger, frontendOrigin string, cookieSecure bool, growthService ...growthReportService) *Server { - server := &Server{identity: identityService, analytics: analyticsService, logger: logger, frontendOrigin: frontendOrigin, cookieSecure: cookieSecure} + server := &Server{identity: identityService, analytics: analyticsService, logger: logger, frontendOrigin: frontendOrigin, cookieSecure: cookieSecure, growthSSESlots: defaultGrowthSSESlots, growthSSEMaxDuration: defaultGrowthSSEMaxDuration, growthSSEPollInterval: defaultGrowthSSEPollInterval} if len(growthService) > 0 { server.growth = growthService[0] } diff --git a/internal/growthreport/avalai.go b/internal/growthreport/avalai.go index 45fb490..17e8ed0 100644 --- a/internal/growthreport/avalai.go +++ b/internal/growthreport/avalai.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "math/big" "net/http" "strings" "time" @@ -60,17 +61,19 @@ func (c *AvalAIClient) GenerateWithTools(ctx context.Context, evidence json.RawM "text": map[string]any{"format": map[string]any{"type": "json_schema", "name": "growth_copy", "strict": true, "schema": aiCopySchema()}}, } var last error + var total Usage for attempt := 0; attempt < 2; attempt++ { copy, usage, transient, err := c.workflow(ctx, base, evidence, known, executor) + total = mergeUsage(total, usage) if err == nil { - return copy, usage, nil + return copy, total, nil } last = err if !transient || ctx.Err() != nil { break } } - return AICopy{}, Usage{}, last + return AICopy{}, total, last } type responseEnvelope struct { @@ -99,9 +102,9 @@ type outputItem struct { func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, evidence json.RawMessage, known map[string]struct{}, executor GrowthAnalyticsTools) (AICopy, Usage, bool, error) { firstBody := cloneRequest(base) firstBody["input"] = string(evidence) - first, latency, transient, err := c.do(ctx, firstBody) + first, usage, transient, err := c.do(ctx, firstBody) if err != nil { - return AICopy{}, Usage{}, transient, err + return AICopy{}, usage, transient, err } calls := make([]outputItem, 0, 4) for _, raw := range first.Output { @@ -110,9 +113,6 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden calls = append(calls, item) } } - usage := Usage{InputTokens: first.Usage.InputTokens, OutputTokens: first.Usage.OutputTokens, LatencyMS: latency} - usage.ProviderRequestID = first.ProviderRequestID - usage.EstimatedCost = costString(first.EstimatedCost) final := first if len(calls) > 0 { if executor == nil { @@ -141,19 +141,11 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden } secondBody := cloneRequest(base) secondBody["input"] = input - second, secondLatency, secondTransient, err := c.do(ctx, secondBody) + second, secondUsage, secondTransient, err := c.do(ctx, secondBody) + usage = mergeUsage(usage, secondUsage) if err != nil { return AICopy{}, usage, secondTransient, err } - usage.InputTokens += second.Usage.InputTokens - usage.OutputTokens += second.Usage.OutputTokens - usage.LatencyMS += secondLatency - if second.ProviderRequestID != "" { - usage.ProviderRequestID = second.ProviderRequestID - } - if cost := costString(second.EstimatedCost); cost != "" { - usage.EstimatedCost = cost - } final = second for _, raw := range second.Output { var item outputItem @@ -169,44 +161,76 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden return copy, usage, false, nil } -func (c *AvalAIClient) do(parent context.Context, bodyValue map[string]any) (responseEnvelope, int64, bool, error) { +func (c *AvalAIClient) do(parent context.Context, bodyValue map[string]any) (responseEnvelope, Usage, bool, error) { body, err := json.Marshal(bodyValue) if err != nil { - return responseEnvelope{}, 0, false, fmt.Errorf("encoding AvalAI request: %w", err) + return responseEnvelope{}, Usage{}, false, fmt.Errorf("encoding AvalAI request: %w", err) } ctx, cancel := context.WithTimeout(parent, c.config.Timeout) defer cancel() req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.config.BaseURL, "/")+"/v1/responses", bytes.NewReader(body)) if err != nil { - return responseEnvelope{}, 0, false, fmt.Errorf("creating AvalAI request: %w", err) + return responseEnvelope{}, Usage{}, false, fmt.Errorf("creating AvalAI request: %w", err) } req.Header.Set("Authorization", "Bearer "+c.config.APIKey) req.Header.Set("Content-Type", "application/json") started := time.Now() resp, err := c.client.Do(req) if err != nil { - return responseEnvelope{}, 0, true, fmt.Errorf("calling AvalAI: %w", err) + return responseEnvelope{}, Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds()}, true, fmt.Errorf("calling AvalAI: %w", err) } defer resp.Body.Close() responseBody, err := io.ReadAll(io.LimitReader(resp.Body, maxAvalAIResponseBytes+1)) if err != nil { - return responseEnvelope{}, 0, true, fmt.Errorf("reading AvalAI response: %w", err) + return responseEnvelope{}, Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds(), ProviderRequestID: resp.Header.Get("X-Request-ID")}, true, fmt.Errorf("reading AvalAI response: %w", err) } if len(responseBody) > maxAvalAIResponseBytes { - return responseEnvelope{}, 0, false, errors.New("avalai: response too large") - } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return responseEnvelope{}, 0, resp.StatusCode == 429 || resp.StatusCode >= 500, fmt.Errorf("avalai: HTTP %d", resp.StatusCode) + return responseEnvelope{}, Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds(), ProviderRequestID: resp.Header.Get("X-Request-ID")}, false, errors.New("avalai: response too large") } var result responseEnvelope - if err := json.Unmarshal(responseBody, &result); err != nil { - return responseEnvelope{}, 0, false, fmt.Errorf("decoding AvalAI response: %w", err) + decodeErr := json.Unmarshal(responseBody, &result) + result.ProviderRequestID = resp.Header.Get("X-Request-ID") + usage := Usage{InputTokens: result.Usage.InputTokens, OutputTokens: result.Usage.OutputTokens, LatencyMS: time.Since(started).Milliseconds(), ProviderRequests: 1, ProviderRequestID: result.ProviderRequestID, EstimatedCost: costString(result.EstimatedCost)} + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return result, usage, resp.StatusCode == 429 || resp.StatusCode >= 500, fmt.Errorf("avalai: HTTP %d", resp.StatusCode) + } + if decodeErr != nil { + return responseEnvelope{}, usage, false, fmt.Errorf("decoding AvalAI response: %w", decodeErr) } if result.Status != "completed" { - return responseEnvelope{}, 0, false, fmt.Errorf("avalai: incomplete response status %q", result.Status) + return result, usage, false, fmt.Errorf("avalai: incomplete response status %q", result.Status) } - result.ProviderRequestID = resp.Header.Get("X-Request-ID") - return result, time.Since(started).Milliseconds(), false, nil + return result, usage, false, nil +} + +func mergeUsage(total, next Usage) Usage { + total.InputTokens += next.InputTokens + total.OutputTokens += next.OutputTokens + total.LatencyMS += next.LatencyMS + total.ProviderRequests += next.ProviderRequests + total.ReservedTokens += next.ReservedTokens + if next.ProviderRequestID != "" { + total.ProviderRequestID = next.ProviderRequestID + } + total.EstimatedCost = addDecimalStrings(total.EstimatedCost, next.EstimatedCost) + return total +} + +func addDecimalStrings(left, right string) string { + if left == "" { + return right + } + if right == "" { + return left + } + a, okA := new(big.Rat).SetString(left) + b, okB := new(big.Rat).SetString(right) + if !okA || !okB { + return left + } + text := new(big.Rat).Add(a, b).FloatString(12) + text = strings.TrimRight(strings.TrimRight(text, "0"), ".") + return text } func costString(raw json.RawMessage) string { @@ -287,6 +311,9 @@ func ValidateAICopy(copy AICopy, known map[string]struct{}) error { return errors.New("avalai: numeric claim rejected") } } + if containsQuantityWord(field) { + return errors.New("avalai: written numeric claim rejected") + } lower := strings.ToLower(field) for _, forbidden := range []string{"payer", "card", "رمز", "شماره کارت", "sql", "merchant_id"} { if strings.Contains(lower, forbidden) { @@ -306,6 +333,20 @@ func ValidateAICopy(copy AICopy, known map[string]struct{}) error { return nil } +func containsQuantityWord(value string) bool { + quantityWords := map[string]struct{}{ + "zero": {}, "one": {}, "two": {}, "three": {}, "four": {}, "five": {}, "six": {}, "seven": {}, "eight": {}, "nine": {}, "ten": {}, "eleven": {}, "twelve": {}, "thirteen": {}, "fourteen": {}, "fifteen": {}, "sixteen": {}, "seventeen": {}, "eighteen": {}, "nineteen": {}, "twenty": {}, "thirty": {}, "forty": {}, "fifty": {}, "sixty": {}, "seventy": {}, "eighty": {}, "ninety": {}, "hundred": {}, "thousand": {}, "million": {}, "billion": {}, "percent": {}, "percentage": {}, "double": {}, "triple": {}, "half": {}, + "صفر": {}, "یک": {}, "دو": {}, "سه": {}, "چهار": {}, "پنج": {}, "شش": {}, "هفت": {}, "هشت": {}, "نه": {}, "ده": {}, "یازده": {}, "دوازده": {}, "سیزده": {}, "چهارده": {}, "پانزده": {}, "شانزده": {}, "هفده": {}, "هجده": {}, "نوزده": {}, "بیست": {}, "سی": {}, "چهل": {}, "پنجاه": {}, "شصت": {}, "هفتاد": {}, "هشتاد": {}, "نود": {}, "صد": {}, "هزار": {}, "میلیون": {}, "میلیارد": {}, "درصد": {}, "برابر": {}, "نصف": {}, + } + for _, word := range strings.Fields(strings.ToLower(value)) { + word = strings.Trim(word, ".,;:!?،؛؟٪%()[]{}«»\"'") + if _, found := quantityWords[word]; found { + return true + } + } + return false +} + func growthTools() []map[string]any { names := []string{"get_growth_context", "get_event_evidence", "get_peer_context", "get_campaign_playbook"} tools := make([]map[string]any, 0, len(names)) diff --git a/internal/growthreport/avalai_test.go b/internal/growthreport/avalai_test.go index 6812522..02724fb 100644 --- a/internal/growthreport/avalai_test.go +++ b/internal/growthreport/avalai_test.go @@ -3,6 +3,7 @@ package growthreport import ( "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "strings" @@ -10,6 +11,31 @@ import ( "time" ) +func TestAvalAICountsFailedRetryRequestsAndUsage(t *testing.T) { + t.Parallel() + requests := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requests++ + w.Header().Set("X-Request-ID", fmt.Sprintf("request-%d", requests)) + if requests == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = w.Write([]byte(`{"usage":{"input_tokens":7,"output_tokens":3},"estimated_cost":"0.01"}`)) + return + } + _, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{\"briefs\":[{\"headline\":\"فرصت مناسب\",\"interpretation\":\"هم‌راستا با شواهد است.\",\"offerSuggestion\":\"پیشنهاد آماده کنید.\",\"messageSuggestion\":\"پیام روشن ارسال کنید.\",\"preparationAction\":\"مسیر پرداخت را بررسی کنید.\",\"limitation\":\"نتیجه قطعی نیست.\",\"evidenceIds\":[\"event_teacher\"]}]}"}]}],"usage":{"input_tokens":11,"output_tokens":5},"estimated_cost":"0.02"}`)) + })) + defer server.Close() + + client := NewAvalAIClient(AvalAIConfig{BaseURL: server.URL, APIKey: "test", Model: "test", Timeout: time.Second}, server.Client()) + _, usage, err := client.Generate(context.Background(), json.RawMessage(`{"evidence":"event_teacher"}`), map[string]struct{}{"event_teacher": {}}) + if err != nil { + t.Fatal(err) + } + if requests != 2 || usage.ProviderRequests != 2 || usage.InputTokens != 18 || usage.OutputTokens != 8 || usage.EstimatedCost != "0.03" { + t.Fatalf("requests=%d usage=%+v", requests, usage) + } +} + func TestValidateAICopyRejectsNumbersAndUnknownEvidence(t *testing.T) { t.Parallel() tests := []struct { @@ -19,6 +45,8 @@ func TestValidateAICopyRejectsNumbersAndUnknownEvidence(t *testing.T) { }{ {name: "grounded Persian copy", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسبی پیش رو دارید", Interpretation: "این مناسبت با الگوی موفق گذشته هم‌راستا است.", EvidenceIDs: []string{"event_teacher"}}}}, want: true}, {name: "numeric claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "رشد ۲۰ درصدی", EvidenceIDs: []string{"event_teacher"}}}}, want: false}, + {name: "Persian number-word claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "رشد بیست درصدی", EvidenceIDs: []string{"event_teacher"}}}}, want: false}, + {name: "English number-word claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "growth by twenty percent", EvidenceIDs: []string{"event_teacher"}}}}, want: false}, {name: "unknown evidence", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسب", EvidenceIDs: []string{"unknown"}}}}, want: false}, } known := map[string]struct{}{"event_teacher": {}} diff --git a/internal/growthreport/coordinator.go b/internal/growthreport/coordinator.go index 6461afa..1fc7873 100644 --- a/internal/growthreport/coordinator.go +++ b/internal/growthreport/coordinator.go @@ -37,17 +37,17 @@ func (c *Coordinator) AllowCreate(ctx context.Context, userID int64, now time.Ti func (c *Coordinator) AllowUncached(ctx context.Context, merchantID int64, now time.Time) error { return c.increment(ctx, fmt.Sprintf("growth:scope:%d:%s", merchantID, tehranDate(now)), c.config.UncachedPerDay, 36*time.Hour, now) } -func (c *Coordinator) ReserveBudget(ctx context.Context, tokens int64, now time.Time) error { +func (c *Coordinator) ReserveBudget(ctx context.Context, requests, tokens int64, now time.Time) error { if c.config.DailyRequests < 1 || c.config.DailyTokens < 1 { return ErrBudgetExhausted } - if c.redis == nil { + if c.redis == nil || requests < 1 || tokens < 1 || requests > c.config.DailyRequests || tokens > c.config.DailyTokens { return ErrBudgetExhausted } - if err := c.incrementRedis(ctx, "growth:budget:req:"+tehranDate(now), 1, c.config.DailyRequests, 36*time.Hour); err != nil { - return ErrBudgetExhausted - } - if err := c.incrementRedis(ctx, "growth:budget:tok:"+tehranDate(now), tokens, c.config.DailyTokens, 36*time.Hour); err != nil { + script := redis.NewScript(`local req=tonumber(redis.call('GET',KEYS[1]) or '0');local tok=tonumber(redis.call('GET',KEYS[2]) or '0');local addReq=tonumber(ARGV[1]);local addTok=tonumber(ARGV[2]);if req+addReq>tonumber(ARGV[4]) or tok+addTok>tonumber(ARGV[5]) then return 0 end;redis.call('INCRBY',KEYS[1],addReq);redis.call('INCRBY',KEYS[2],addTok);redis.call('PEXPIRE',KEYS[1],ARGV[3]);redis.call('PEXPIRE',KEYS[2],ARGV[3]);return 1`) + date := tehranDate(now) + allowed, err := script.Run(ctx, c.redis, []string{"growth:budget:req:" + date, "growth:budget:tok:" + date}, requests, tokens, (36 * time.Hour).Milliseconds(), c.config.DailyRequests, c.config.DailyTokens).Int() + if err != nil || allowed != 1 { return ErrBudgetExhausted } return nil diff --git a/internal/growthreport/coordinator_test.go b/internal/growthreport/coordinator_test.go index 9bd4e16..1e176b4 100644 --- a/internal/growthreport/coordinator_test.go +++ b/internal/growthreport/coordinator_test.go @@ -20,7 +20,7 @@ func TestCoordinatorFallsBackToBoundedProcessCounters(t *testing.T) { if err := coordinator.AllowCreate(context.Background(), 7, now); !errors.Is(err, ErrRateLimited) { t.Fatalf("third request error=%v", err) } - if err := coordinator.ReserveBudget(context.Background(), 1600, now); !errors.Is(err, ErrBudgetExhausted) { + if err := coordinator.ReserveBudget(context.Background(), MaxProviderRequests, 1600, now); !errors.Is(err, ErrBudgetExhausted) { t.Fatalf("budget without Redis error=%v", err) } } diff --git a/internal/growthreport/model.go b/internal/growthreport/model.go index 8802772..4373dad 100644 --- a/internal/growthreport/model.go +++ b/internal/growthreport/model.go @@ -41,6 +41,8 @@ type Usage struct { InputTokens int64 `json:"inputTokens"` OutputTokens int64 `json:"outputTokens"` LatencyMS int64 `json:"latencyMs"` + ProviderRequests int64 `json:"providerRequests"` + ReservedTokens int64 `json:"reservedTokens"` ProviderRequestID string `json:"providerRequestId,omitempty"` EstimatedCost string `json:"estimatedCost,omitempty"` } @@ -77,6 +79,7 @@ type Record struct { Usage Usage ManualRefresh bool TehranRefreshDay *time.Time + ClaimToken string } type CreateRequest struct { diff --git a/internal/growthreport/orchestrator.go b/internal/growthreport/orchestrator.go index dac1490..0737c40 100644 --- a/internal/growthreport/orchestrator.go +++ b/internal/growthreport/orchestrator.go @@ -15,6 +15,8 @@ import ( var ErrDisabled = errors.New("growth report: disabled") +const MaxProviderRequests int64 = 4 + type ServiceConfig struct { Enabled bool Model, PromptVersion, SchemaVersion string @@ -42,13 +44,15 @@ func (s *Service) Create(ctx context.Context, userID, merchantID int64, request if err != nil { return Record{}, false, err } - fingerprint := growth.Fingerprint(growth.FingerprintInput{MerchantID: merchantID, TerminalID: request.TerminalID, DatasetWatermark: watermark, ForecastStart: start, HorizonDays: 90, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ModelConfigVersion: growth.DefaultModelConfigVersion}) + fingerprint := s.fingerprint(merchantID, request.TerminalID, watermark, start) var refreshDay *time.Time if request.Refresh { day := tehranMidnight(now) refreshDay = &day } - record, cached, err := s.repository.CreateOrReuse(ctx, CreateRecord{ID: newReportID(), MerchantID: merchantID, TerminalID: request.TerminalID, Fingerprint: fingerprint, DatasetWatermark: watermark, ForecastStartDate: &start, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ManualRefresh: request.Refresh, RefreshDay: refreshDay, IdempotencyKey: idempotencyKey, RequestHash: RequestHash(merchantID, request.TerminalID, request.Refresh)}) + record, cached, err := s.repository.CreateOrReuse(ctx, CreateRecord{ID: newReportID(), MerchantID: merchantID, TerminalID: request.TerminalID, Fingerprint: fingerprint, DatasetWatermark: watermark, ForecastStartDate: &start, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ManualRefresh: request.Refresh, RefreshDay: refreshDay, IdempotencyKey: idempotencyKey, RequestHash: RequestHash(merchantID, request.TerminalID, request.Refresh)}, func() error { + return s.coordinator.AllowUncached(ctx, merchantID, now) + }) if err != nil { if errors.Is(err, ErrRefreshLimited) { next, quotaErr := s.repository.NextRefreshAt(ctx, merchantID, now) @@ -59,12 +63,6 @@ func (s *Service) Create(ctx context.Context, userID, merchantID int64, request } return Record{}, false, err } - if !cached { - if err := s.coordinator.AllowUncached(ctx, merchantID, now); err != nil { - _ = s.repository.Complete(ctx, record.ReportID, StatusFailed, json.RawMessage(`{}`), nil, Usage{}, "UNCACHED_SCOPE_LIMIT") - return Record{}, false, err - } - } record.Cached = cached record.NextRefreshAt, err = s.repository.NextRefreshAt(ctx, merchantID, now) return record, cached, err @@ -80,11 +78,21 @@ func (s *Service) Latest(ctx context.Context, merchantID int64, terminalID strin record, err := s.repository.Latest(ctx, merchantID, terminalID) if err == nil { record.Cached = true + watermark, start, metadataErr := s.repository.DatasetMetadata(ctx) + if metadataErr != nil { + record.Stale = true + } else { + record.Stale = record.Fingerprint != s.fingerprint(merchantID, terminalID, watermark, start) + } record.NextRefreshAt, _ = s.repository.NextRefreshAt(ctx, merchantID, s.now()) } return record, err } +func (s *Service) fingerprint(merchantID int64, terminalID, watermark string, start time.Time) string { + return growth.Fingerprint(growth.FingerprintInput{MerchantID: merchantID, TerminalID: terminalID, DatasetWatermark: watermark, ForecastStart: start, HorizonDays: 90, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ModelConfigVersion: growth.DefaultModelConfigVersion}) +} + type Processor struct { repository *Repository builder Builder @@ -108,10 +116,10 @@ func (p *Processor) ProcessNext(ctx context.Context) error { } analysis, err := p.builder.Build(ctx, record.MerchantID, terminal) if err != nil { - _ = p.repository.Complete(ctx, record.ReportID, StatusFailed, json.RawMessage(`{}`), nil, Usage{}, "ANALYTICS_FAILED") + _ = p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, StatusFailed, json.RawMessage(`{}`), nil, Usage{}, "ANALYTICS_FAILED") return err } - if err := p.repository.StoreEvidence(ctx, record.ReportID, analysis.Encoded); err != nil { + if err := p.repository.StoreEvidence(ctx, record.ReportID, record.ClaimToken, analysis.Encoded); err != nil { return err } copy := analysis.Fallback @@ -119,30 +127,41 @@ func (p *Processor) ProcessNext(ctx context.Context) error { failure := "AI_DISABLED_OR_UNAVAILABLE" if analysis.ShouldCallAI && p.aiEnabled && p.avalai != nil { reservedTokens := EstimatedTokenReservation(analysis.Encoded) - if err := p.coordinator.ReserveBudget(ctx, reservedTokens, p.now()); err == nil { - release, enterErr := p.coordinator.Enter(ctx) - if enterErr == nil { + release, enterErr := p.coordinator.Enter(ctx) + if enterErr == nil { + defer release() + if err := p.coordinator.ReserveBudget(ctx, MaxProviderRequests, reservedTokens, p.now()); err == nil { generated, usage, aiErr := p.avalai.GenerateWithTools(ctx, analysis.Encoded, analysis.KnownEvidence, NewEvidenceTools(analysis)) - release() + usage.ReservedTokens = reservedTokens if aiErr == nil { - slog.Info("growth AI generation complete", "report_id", record.ReportID, "input_tokens", usage.InputTokens, "output_tokens", usage.OutputTokens, "latency_ms", usage.LatencyMS) - return p.repository.Complete(ctx, record.ReportID, StatusReady, analysis.Encoded, &generated, usage, "") + slog.Info("growth AI generation complete", "report_id", record.ReportID, "provider_requests", usage.ProviderRequests, "reserved_tokens", usage.ReservedTokens, "input_tokens", usage.InputTokens, "output_tokens", usage.OutputTokens, "latency_ms", usage.LatencyMS) + return p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, StatusReady, analysis.Encoded, &generated, usage, "") } - slog.Warn("growth AI output rejected", "report_id", record.ReportID, "reason", "AI_REJECTED") + slog.Warn("growth AI output rejected", "report_id", record.ReportID, "reason", "AI_REJECTED", "provider_requests", usage.ProviderRequests, "reserved_tokens", usage.ReservedTokens, "input_tokens", usage.InputTokens, "output_tokens", usage.OutputTokens, "latency_ms", usage.LatencyMS) failure = "AI_REJECTED" + return p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, status, analysis.Encoded, ©, usage, failure) + } else { + failure = "AI_BUDGET_EXHAUSTED" + slog.Warn("growth AI skipped", "report_id", record.ReportID, "reason", failure) } } else { - failure = "AI_BUDGET_EXHAUSTED" + failure = "AI_CONCURRENCY_UNAVAILABLE" slog.Warn("growth AI skipped", "report_id", record.ReportID, "reason", failure) } } if !analysis.ShouldCallAI || !p.aiEnabled { slog.Info("growth deterministic fallback used", "report_id", record.ReportID, "reason", failure) } - return p.repository.Complete(ctx, record.ReportID, status, analysis.Encoded, ©, Usage{}, failure) + return p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, status, analysis.Encoded, ©, Usage{}, failure) } -func EstimatedTokenReservation(evidence json.RawMessage) int64 { return int64(2*len(evidence) + 3200) } +// EstimatedTokenReservation reserves for two complete attempts, each with two +// provider turns, before any network call. Bytes are deliberately counted as +// tokens and the fixed allowance covers tool schemas, instructions, tool +// outputs, and every turn's 1,600-token output ceiling. +func EstimatedTokenReservation(evidence json.RawMessage) int64 { + return int64(12*len(evidence) + 32000) +} func newReportID() string { data := make([]byte, 12) diff --git a/internal/growthreport/repository.go b/internal/growthreport/repository.go index b135b5d..0635bc7 100644 --- a/internal/growthreport/repository.go +++ b/internal/growthreport/repository.go @@ -20,6 +20,8 @@ var ( type Repository struct{ pool *pgxpool.Pool } +const reportLeaseDuration = 3 * time.Minute + func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} } type CreateRecord struct { @@ -40,7 +42,7 @@ type CreateRecord struct { RequestHash string } -func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord) (Record, bool, error) { +func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord, beforeInsert func() error) (Record, bool, error) { tx, err := r.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable}) if err != nil { return Record{}, false, fmt.Errorf("beginning report transaction: %w", err) @@ -96,6 +98,11 @@ func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord) (Record return Record{}, false, ErrRefreshLimited } } + if beforeInsert != nil { + if err := beforeInsert(); err != nil { + return Record{}, false, err + } + } var revision int if err := tx.QueryRow(ctx, `SELECT coalesce(max(revision),0)+1 FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text`, in.MerchantID, nilIfEmpty(in.TerminalID)).Scan(&revision); err != nil { return Record{}, false, fmt.Errorf("calculating growth report revision: %w", err) @@ -139,6 +146,9 @@ func (r *Repository) Claim(ctx context.Context) (Record, error) { return Record{}, fmt.Errorf("beginning report claim: %w", err) } defer func() { _ = tx.Rollback(ctx) }() + if _, err := tx.Exec(ctx, `UPDATE growth_reports SET status='QUEUED',lease_token=NULL,lease_expires_at=NULL,updated_at=now(),failure_code='WORKER_LEASE_EXPIRED' WHERE status IN ('ANALYZING','GENERATING') AND (lease_expires_at IS NULL OR lease_expires_at $4),updated_at=now() WHERE id=$1 AND lease_token=$2 AND status='ANALYZING'`, reportID, claimToken, evidence, int(reportLeaseDuration/time.Second)) if err != nil { return fmt.Errorf("storing deterministic evidence: %w", err) } @@ -169,7 +182,7 @@ func (r *Repository) StoreEvidence(ctx context.Context, reportID string, evidenc return nil } -func (r *Repository) Complete(ctx context.Context, reportID string, status Status, evidence json.RawMessage, copy *AICopy, usage Usage, failureCode string) error { +func (r *Repository) CompleteClaimed(ctx context.Context, reportID, claimToken string, status Status, evidence json.RawMessage, copy *AICopy, usage Usage, failureCode string) error { if status != StatusReady && status != StatusPartial && status != StatusFailed { return errors.New("growth report: invalid terminal status") } @@ -177,14 +190,17 @@ func (r *Repository) Complete(ctx context.Context, reportID string, status Statu if copy != nil { ai = copy } - _, err := r.pool.Exec(ctx, `UPDATE growth_reports SET status=$2,deterministic_evidence=$3,ai_response=$4,failure_code=$5,input_tokens=$6,output_tokens=$7,latency_ms=$8,provider_request_id=$9,estimated_cost=$10,updated_at=now(),completed_at=now() WHERE id=$1`, reportID, status, evidence, ai, nilIfEmpty(failureCode), usage.InputTokens, usage.OutputTokens, usage.LatencyMS, nilIfEmpty(usage.ProviderRequestID), nilIfEmpty(usage.EstimatedCost)) + command, err := r.pool.Exec(ctx, `UPDATE growth_reports SET status=$3,deterministic_evidence=$4,ai_response=$5,failure_code=$6,input_tokens=$7,output_tokens=$8,latency_ms=$9,provider_requests=$10,reserved_tokens=$11,provider_request_id=$12,estimated_cost=$13,lease_token=NULL,lease_expires_at=NULL,updated_at=now(),completed_at=now() WHERE id=$1 AND lease_token=$2`, reportID, claimToken, status, evidence, ai, nilIfEmpty(failureCode), usage.InputTokens, usage.OutputTokens, usage.LatencyMS, usage.ProviderRequests, usage.ReservedTokens, nilIfEmpty(usage.ProviderRequestID), nilIfEmpty(usage.EstimatedCost)) if err != nil { return fmt.Errorf("completing growth report: %w", err) } + if command.RowsAffected() != 1 { + return ErrNotFound + } return nil } -const reportSelect = `SELECT id,revision,merchant_id,terminal_id,terminal_source_key,fingerprint,status,horizon_days,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,deterministic_evidence,ai_response,failure_code,input_tokens,output_tokens,latency_ms,provider_request_id,estimated_cost,manual_refresh,tehran_refresh_day,created_at,updated_at FROM growth_reports` +const reportSelect = `SELECT id,revision,merchant_id,terminal_id,terminal_source_key,fingerprint,status,horizon_days,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,deterministic_evidence,ai_response,failure_code,input_tokens,output_tokens,latency_ms,provider_requests,reserved_tokens,provider_request_id,estimated_cost,manual_refresh,tehran_refresh_day,created_at,updated_at FROM growth_reports` type rowScanner interface{ Scan(...any) error } type queryer interface { @@ -205,7 +221,7 @@ func scanRecord(row rowScanner) (Record, error) { var ai []byte var terminalSource *string var providerRequestID, estimatedCost *string - err := row.Scan(&r.ReportID, &r.Revision, &r.MerchantID, &r.TerminalDBID, &terminalSource, &r.Fingerprint, &status, &r.HorizonDays, &r.ForecastStartDate, &r.DatasetWatermark, &r.CatalogVersion, &r.AnalyticsVersion, &r.PromptVersion, &r.SchemaVersion, &r.Model, &evidence, &ai, &r.FailureCode, &r.Usage.InputTokens, &r.Usage.OutputTokens, &r.Usage.LatencyMS, &providerRequestID, &estimatedCost, &r.ManualRefresh, &r.TehranRefreshDay, &r.CreatedAt, &r.UpdatedAt) + err := row.Scan(&r.ReportID, &r.Revision, &r.MerchantID, &r.TerminalDBID, &terminalSource, &r.Fingerprint, &status, &r.HorizonDays, &r.ForecastStartDate, &r.DatasetWatermark, &r.CatalogVersion, &r.AnalyticsVersion, &r.PromptVersion, &r.SchemaVersion, &r.Model, &evidence, &ai, &r.FailureCode, &r.Usage.InputTokens, &r.Usage.OutputTokens, &r.Usage.LatencyMS, &r.Usage.ProviderRequests, &r.Usage.ReservedTokens, &providerRequestID, &estimatedCost, &r.ManualRefresh, &r.TehranRefreshDay, &r.CreatedAt, &r.UpdatedAt) if errors.Is(err, pgx.ErrNoRows) { return Record{}, ErrNotFound } diff --git a/internal/growthreport/repository_integration_test.go b/internal/growthreport/repository_integration_test.go new file mode 100644 index 0000000..ba8fd1c --- /dev/null +++ b/internal/growthreport/repository_integration_test.go @@ -0,0 +1,81 @@ +package growthreport + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestRepositoryRecoversExpiredActiveLeases(t *testing.T) { + pool, merchantID := growthRepositoryTestPool(t) + ctx := context.Background() + now := time.Now().UTC() + for index, status := range []Status{StatusAnalyzing, StatusGenerating} { + id := fmt.Sprintf("growth_%024x", index+1) + _, err := pool.Exec(ctx, `INSERT INTO growth_reports(id,merchant_id,fingerprint,status,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,lease_token,lease_expires_at,created_at,updated_at) VALUES($1,$2,$3,$4,current_date,'test','test','test','test','test','test',$5,$6,$7,$7)`, id, merchantID, id, status, "expired-token", now.Add(-time.Minute), now.Add(time.Duration(index)*time.Second)) + if err != nil { + t.Fatal(err) + } + } + + repository := NewRepository(pool) + first, err := repository.Claim(ctx) + if err != nil { + t.Fatal(err) + } + if first.Status != StatusAnalyzing || first.ClaimToken == "" || first.ClaimToken == "expired-token" { + t.Fatalf("recovered claim=%+v", first) + } + if err := repository.CompleteClaimed(ctx, first.ReportID, "expired-token", StatusPartial, json.RawMessage(`{}`), nil, Usage{}, "stale"); !errors.Is(err, ErrNotFound) { + t.Fatalf("stale completion error=%v", err) + } + second, err := repository.Claim(ctx) + if err != nil || second.ReportID == first.ReportID || second.ClaimToken == "" { + t.Fatalf("second claim=%+v error=%v", second, err) + } +} + +func TestRepositoryChecksUncachedAllowanceBeforeInsert(t *testing.T) { + pool, merchantID := growthRepositoryTestPool(t) + repository := NewRepository(pool) + start := time.Now().UTC() + input := CreateRecord{ID: "growth_ffffffffffffffffffffffff", MerchantID: merchantID, Fingerprint: "blocked", DatasetWatermark: "test", ForecastStartDate: &start, CatalogVersion: "test", AnalyticsVersion: "test", PromptVersion: "test", SchemaVersion: "test", Model: "test"} + _, _, err := repository.CreateOrReuse(context.Background(), input, func() error { return ErrRateLimited }) + if !errors.Is(err, ErrRateLimited) { + t.Fatalf("CreateOrReuse error=%v", err) + } + var count int + if err := pool.QueryRow(context.Background(), `SELECT count(*) FROM growth_reports WHERE id=$1`, input.ID).Scan(&count); err != nil || count != 0 { + t.Fatalf("queued reports=%d error=%v", count, err) + } +} + +func growthRepositoryTestPool(t *testing.T) (*pgxpool.Pool, int64) { + t.Helper() + databaseURL := os.Getenv("TEST_DATABASE_URL") + if databaseURL == "" { + t.Skip("TEST_DATABASE_URL is not configured") + } + pool, err := pgxpool.New(context.Background(), databaseURL) + if err != nil { + t.Fatal(err) + } + sourceKey := fmt.Sprintf("growth-test-%d", time.Now().UnixNano()) + var merchantID int64 + if err := pool.QueryRow(context.Background(), `INSERT INTO merchants(source_key,category_id,category_title,display_name) VALUES($1,1,'test','test') RETURNING id`, sourceKey).Scan(&merchantID); err != nil { + pool.Close() + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), `DELETE FROM growth_reports WHERE merchant_id=$1`, merchantID) + _, _ = pool.Exec(context.Background(), `DELETE FROM merchants WHERE id=$1`, merchantID) + pool.Close() + }) + return pool, merchantID +} diff --git a/internal/growthreport/service_test.go b/internal/growthreport/service_test.go index 60ddd46..d7f46f5 100644 --- a/internal/growthreport/service_test.go +++ b/internal/growthreport/service_test.go @@ -27,11 +27,25 @@ func TestNextTehranRefreshUsesCalendarBoundary(t *testing.T) { } } -func TestEstimatedTokenReservationCoversTwoTurnsAndGrowsWithEvidence(t *testing.T) { +func TestEstimatedTokenReservationCoversEveryRetryAndTurn(t *testing.T) { t.Parallel() small := EstimatedTokenReservation([]byte("{}")) large := EstimatedTokenReservation(make([]byte, 1000)) - if small < 3200 || large <= small || large != 5200 { + if small < 4*1600 || large <= small || large != 44000 { t.Fatalf("reservations small=%d large=%d", small, large) } } + +func TestServiceFingerprintChangesWithDatasetAndConfiguration(t *testing.T) { + t.Parallel() + start := time.Date(2026, 8, 21, 0, 0, 0, 0, time.UTC) + service := &Service{config: ServiceConfig{Model: "model-a", PromptVersion: "prompt-a", SchemaVersion: "schema-a"}} + current := service.fingerprint(7, "T-1", "watermark-a", start) + if current == service.fingerprint(7, "T-1", "watermark-b", start) { + t.Fatal("dataset watermark did not invalidate fingerprint") + } + service.config.PromptVersion = "prompt-b" + if current == service.fingerprint(7, "T-1", "watermark-a", start) { + t.Fatal("prompt version did not invalidate fingerprint") + } +} diff --git a/migrations/000002_growth_analyst.sql b/migrations/000002_growth_analyst.sql index 6ea5c94..8944182 100644 --- a/migrations/000002_growth_analyst.sql +++ b/migrations/000002_growth_analyst.sql @@ -20,8 +20,12 @@ CREATE TABLE IF NOT EXISTS growth_reports ( input_tokens bigint NOT NULL DEFAULT 0 CHECK (input_tokens >= 0), output_tokens bigint NOT NULL DEFAULT 0 CHECK (output_tokens >= 0), latency_ms bigint NOT NULL DEFAULT 0 CHECK (latency_ms >= 0), + provider_requests bigint NOT NULL DEFAULT 0 CHECK (provider_requests >= 0), + reserved_tokens bigint NOT NULL DEFAULT 0 CHECK (reserved_tokens >= 0), provider_request_id text, estimated_cost text, + lease_token text, + lease_expires_at timestamptz, manual_refresh boolean NOT NULL DEFAULT false, tehran_refresh_day date, created_at timestamptz NOT NULL DEFAULT now(), @@ -33,6 +37,14 @@ ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS provider_request_id text; -- statement-break ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS estimated_cost text; -- statement-break +ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS provider_requests bigint NOT NULL DEFAULT 0 CHECK (provider_requests >= 0); +-- statement-break +ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS reserved_tokens bigint NOT NULL DEFAULT 0 CHECK (reserved_tokens >= 0); +-- statement-break +ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS lease_token text; +-- statement-break +ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS lease_expires_at timestamptz; +-- statement-break ALTER TABLE growth_reports ALTER COLUMN deterministic_evidence DROP NOT NULL; -- statement-break CREATE UNIQUE INDEX IF NOT EXISTS growth_reports_current_fingerprint_idx