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/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. The DB enforces the no-negative / no-replay
	// invariants; on a depleted balance this returns 402 with no provider call.
	idemKey := strings.TrimSpace(c.GetHeader("Idempotency-Key"))
	remaining, err := h.CreditService.Consume(
		c.Request.Context(), *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. The refund is best-effort;
	// a failed refund is logged via the audit trail rather than surfaced.
	result, gerr := h.WritingGenerator.GenerateWriting(c.Request.Context(), model, req.Prompt, req.Tone)
	if gerr != nil {
		if bal, rerr := h.CreditService.Grant(c.Request.Context(), *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
	}

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