ImaGen/internal/usage/usage.go

// 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
}