warmbly/internal/api/handler/generation.go

// AI writing-assistant generation endpoint. Flow:
//  1. feature-gate the org (paid or in free trial) via CanUseWritingAssistant
//  2. atomically consume one credit (DB-enforced: no negative balance, no
//     double-charge on Idempotency-Key replay)
//  3. call the configured provider (Anthropic, falling back to OpenAI)
//  4. return {text, credits_remaining, model}
//
// On insufficient credits the consume step short-circuits with 402 BEFORE any
// provider call, so a depleted org never burns a paid completion. Because the
// debit happens before the provider call, a provider failure refunds the
// credit so the customer is not charged for a generation they never received.

package handler

import (
	"errors"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/warmbly/warmbly/internal/api/middleware"
	"github.com/warmbly/warmbly/internal/app/credits"
	"github.com/warmbly/warmbly/internal/errx"
	"github.com/warmbly/warmbly/internal/models"
	"github.com/warmbly/warmbly/internal/pkg/generation"
)

// creditsPerWrite is the credit cost of one writing-assistant call. Kept as a
// constant so pricing is in one place; tokens consumed are recorded separately
// on the ledger transaction for later cost analysis.
const creditsPerWrite = 1

// writeMaxPromptLen bounds the inbound prompt so a single request can't be used
// to drive a very large (and expensive) completion.
const writeMaxPromptLen = 8000

type generationWriteRequest struct {
	Prompt string `json:"prompt"`
	Tone   string `json:"tone"`
}

// paymentRequiredJSON emits the standard error envelope with a 402 status.
// errx has no PaymentRequired code, so this endpoint writes the 402 directly
// while keeping the same {error, message, code, request_id} shape.
func paymentRequiredJSON(c *gin.Context, message string) {
	c.JSON(http.StatusPaymentRequired, gin.H{
		"error":      "Payment Required",
		"message":    message,
		"code":       "insufficient_credits",
		"request_id": c.GetString("request_id"),
	})
}

// GenerateWriting — POST /generation/write
func (h *Handler) GenerateWriting(c *gin.Context) {
	orgID := middleware.GetOrganizationID(c)
	if orgID == nil {
		errx.JSON(c, errx.New(errx.BadRequest, "no organization selected"))
		return
	}

	var req generationWriteRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		errx.JSON(c, errx.ErrInvalid)
		return
	}
	req.Prompt = strings.TrimSpace(req.Prompt)
	if req.Prompt == "" {
		errx.JSON(c, errx.New(errx.BadRequest, "prompt is required"))
		return
	}
	if len(req.Prompt) > writeMaxPromptLen {
		errx.JSON(c, errx.New(errx.BadRequest, "prompt is too long"))
		return
	}

	// Feature gate: paid orgs and free-trial orgs may use the assistant.
	allowed, xerr := h.FeatureGateService.CanUseWritingAssistant(c.Request.Context(), *orgID)
	if xerr != nil {
		errx.JSON(c, xerr)
		return
	}
	if !allowed {
		errx.JSON(c, errx.New(errx.Forbidden, "The AI writing assistant requires an active plan or trial."))
		return
	}

	// Provider must be configured.
	if h.WritingGenerator == nil {
		errx.JSON(c, errx.New(errx.ServiceUnavailable, "AI writing assistant is not configured."))
		return
	}

	// Model routing by tier. Paid orgs get the stronger model; the active
	// provider (Anthropic or OpenAI fallback) decides the concrete model ID.
	paid, xerr := h.FeatureGateService.IsPaidOrganization(c.Request.Context(), *orgID)
	if xerr != nil {
		errx.JSON(c, xerr)
		return
	}
	model := h.WritingGenerator.ModelForTier(paid)

	// Consume one credit up front, unless this is a free/local model, which runs
	// un-metered (AI_FREE). On the metered path the DB enforces the
	// no-negative / no-replay invariants and returns 402 on a depleted balance.
	idemKey := strings.TrimSpace(c.GetHeader("Idempotency-Key"))
	local := h.WritingGenerator.IsLocal()
	// Attribute the charge to the teammate who asked for the draft.
	reqCtx := c.Request.Context()
	if actor, aerr := middleware.GetUserUUID(c); aerr == nil {
		reqCtx = models.WithCreditMeta(reqCtx, models.CreditMeta{ActorID: actor})
	}

	var remaining int
	if local {
		if bal, berr := h.CreditService.GetBalance(reqCtx, *orgID); berr == nil {
			remaining = bal
		}
	} else {
		var err error
		remaining, err = h.CreditService.Consume(
			reqCtx, *orgID, creditsPerWrite,
			"writing_assistant", model, 0, idemKey,
		)
		if err != nil {
			switch {
			case errors.Is(err, credits.ErrInsufficientCredits):
				paymentRequiredJSON(c, "You're out of AI credits. Upgrade or purchase more to keep using the writing assistant.")
			case errors.Is(err, credits.ErrCapExceeded):
				errx.JSON(c, errx.New(errx.TooManyRequests, "AI writing assistant usage limit reached, please try again later."))
			default:
				errx.JSON(c, errx.InternalError())
			}
			return
		}
	}

	// Generate. On provider failure, refund the credit so the customer is not
	// charged for a completion they never received (nothing to refund on the
	// free/local path). The refund is best-effort; a failed refund is logged via
	// the audit trail rather than surfaced.
	voice := h.orgVoice(c.Request.Context(), *orgID, req.Tone)
	result, gerr := h.WritingGenerator.GenerateWriting(c.Request.Context(), model, req.Prompt, voice)
	if gerr != nil {
		if !local {
			if bal, rerr := h.CreditService.Grant(reqCtx, *orgID, creditsPerWrite, "writing_assistant_refund"); rerr == nil {
				remaining = bal
			}
		}
		if errors.Is(gerr, generation.ErrNotConfigured) {
			errx.JSON(c, errx.New(errx.ServiceUnavailable, "AI writing assistant is not configured."))
			return
		}
		errx.JSON(c, errx.New(errx.ServiceUnavailable, "The writing assistant is temporarily unavailable. Your credit was not charged."))
		return
	}

	// Usage-based settle: price the actual tokens and charge any overage
	// beyond the flat minimum (best-effort; never fails the delivered text).
	charged := 0
	if !local {
		charged = creditsPerWrite
		if extra, serr := h.CreditService.SettleUsage(reqCtx, *orgID, creditsPerWrite, result.Model, result.TokensUsed, "writing_assistant", settleKey(idemKey)); serr == nil && extra > 0 {
			remaining -= extra
			charged += extra
		}
	}

	c.JSON(http.StatusOK, gin.H{
		"text":              result.Text,
		"credits_remaining": remaining,
		"credits_charged":   charged,
		"tokens_used":       result.TokensUsed,
		"model":             result.Model,
	})
}

// settleKey derives the usage-settle idempotency key from the call's key. An
// empty key stays empty (non-idempotent call, non-idempotent settle).
func settleKey(idemKey string) string {
	if idemKey == "" {
		return ""
	}
	return idemKey + ":usage"
}