// Package usage records per-call cost-tracking rows for the imagen CLI // to mai.imagen_usage on Supabase. The writer is best-effort by design — // the calling adapter logs failures and proceeds, because the image // itself has already landed on disk by the time we record. package usage import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "net/url" "os" "strings" "time" "mgit.msbls.de/m/ImaGen/internal/backend" ) // Default REST schema is the mai schema where mai.imagen_usage lives. const supabaseSchema = "mai" // SupabaseSink writes rows via PostgREST. It uses Accept-Profile/ // Content-Profile headers to target the mai schema instead of public. type SupabaseSink struct { URL string // SUPABASE_URL — e.g. https://msup.msbls.de APIKey string // SUPABASE_SERVICE_KEY HTTP *http.Client } // NewSupabaseSinkFromEnv reads SUPABASE_URL and SUPABASE_SERVICE_KEY // (falling back to MAI_SUPABASE_KEY) and returns a sink ready to use. // Returns nil + ok=false if the env vars are not configured — the CLI // uses that to skip cost-tracking gracefully. func NewSupabaseSinkFromEnv() (*SupabaseSink, bool) { u := strings.TrimRight(os.Getenv("SUPABASE_URL"), "/") if u == "" { return nil, false } key := os.Getenv("SUPABASE_SERVICE_KEY") if key == "" { key = os.Getenv("MAI_SUPABASE_KEY") } if key == "" { return nil, false } return &SupabaseSink{ URL: u, APIKey: key, HTTP: &http.Client{Timeout: 10 * time.Second}, }, true } type supabaseRow struct { Backend string `json:"backend"` Model string `json:"model"` Seed *int64 `json:"seed,omitempty"` PromptHash string `json:"prompt_hash"` LatencyMs int `json:"latency_ms"` CostUSDEstimate *float64 `json:"cost_usd_estimate,omitempty"` Caller string `json:"caller,omitempty"` } // Record inserts one row into mai.imagen_usage. func (s *SupabaseSink) Record(ctx context.Context, row backend.UsageRow) error { body, err := json.Marshal(supabaseRow{ Backend: row.Backend, Model: row.Model, Seed: row.Seed, PromptHash: row.PromptHash, LatencyMs: row.LatencyMs, CostUSDEstimate: row.CostUSDEstimate, Caller: row.Caller, }) if err != nil { return fmt.Errorf("usage: marshal: %w", err) } endpoint := s.URL + "/rest/v1/imagen_usage" req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) if err != nil { return err } req.Header.Set("apikey", s.APIKey) req.Header.Set("Authorization", "Bearer "+s.APIKey) req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept-Profile", supabaseSchema) req.Header.Set("Content-Profile", supabaseSchema) req.Header.Set("Prefer", "return=minimal") resp, err := s.HTTP.Do(req) if err != nil { return fmt.Errorf("usage: POST: %w", err) } defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { respBody, _ := io.ReadAll(resp.Body) return fmt.Errorf("usage: POST %d: %s", resp.StatusCode, snip(respBody)) } return nil } // Row is the read-side row shape (only the fields the CLI needs). type Row struct { CreatedAt time.Time `json:"created_at"` Backend string `json:"backend"` Model string `json:"model"` Seed *int64 `json:"seed"` PromptHash string `json:"prompt_hash"` LatencyMs *int `json:"latency_ms"` CostUSDEstimate *float64 `json:"cost_usd_estimate"` Caller *string `json:"caller"` } // Query returns rows from mai.imagen_usage filtered by created_at >= since. // Pass zero time to fetch the full table (capped server-side by PostgREST // — we set a hard 5000-row limit here too). func (s *SupabaseSink) Query(ctx context.Context, since time.Time) ([]Row, error) { q := url.Values{} q.Set("select", "created_at,backend,model,seed,prompt_hash,latency_ms,cost_usd_estimate,caller") q.Set("order", "created_at.desc") q.Set("limit", "5000") if !since.IsZero() { q.Set("created_at", "gte."+since.UTC().Format(time.RFC3339)) } endpoint := s.URL + "/rest/v1/imagen_usage?" + q.Encode() req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { return nil, err } req.Header.Set("apikey", s.APIKey) req.Header.Set("Authorization", "Bearer "+s.APIKey) req.Header.Set("Accept-Profile", supabaseSchema) resp, err := s.HTTP.Do(req) if err != nil { return nil, fmt.Errorf("usage: GET: %w", err) } defer resp.Body.Close() body, _ := io.ReadAll(resp.Body) if resp.StatusCode < 200 || resp.StatusCode >= 300 { return nil, fmt.Errorf("usage: GET %d: %s", resp.StatusCode, snip(body)) } var rows []Row if err := json.Unmarshal(body, &rows); err != nil { return nil, fmt.Errorf("usage: parse rows: %w (body: %s)", err, snip(body)) } return rows, nil } func snip(b []byte) string { const max = 500 s := strings.TrimSpace(string(b)) if len(s) > max { s = s[:max] + "..." } return s }