fix(payg): gate + charge AI document tools and AI Create sessions (#6617)

## Problem Two AI surfaces slipped through PAYG unbilled: 1. **AI document tools** — `/api/v1/ai/tools/**` (`PdfCommentAgentController`, `MathAuditorAgentController`) live in the **proprietary** module, which can't depend on `saas` and so can't carry the saas-only `@RequiresFeature`. They also lacked `@AutoJobPostMapping`, so the charge interceptor's scope gate short-circuited them **before** category resolution: **not charged, and not even entitlement-gated** — whether called directly or dispatched by the orchestrator. 2. **AI Create** — `/api/v1/ai/create` is JSON/session-based with no file input, so the multipart charge path never fired. The old per-generation charge ran through the now-dead legacy credit system, so it currently charges nothing. ## Fix - **`AiToolRoutes`** (new, saas) — single source of truth for the `/api/v1/ai/tools/**` prefix. The proprietary controllers stay untouched; the saas hot-path recognises them by path: - **`PaygChargeInterceptor`**: brings these routes into scope and bills them **AI** on a direct call. An orchestrator-dispatched call still resolves to **AUTOMATION** first (the `X-Stirling-Automation` header is checked before the path rule), so AI-tool-inside-a-workflow keeps billing as automation. - **`EntitlementGuard`**: gates them on **`AI_SUPPORT`**. - This keeps the `proprietary → saas` layering intact (no backwards dependency). - **`JobChargeService.chargeStandalone(ctx, units)`** — charges a fixed unit count for a non-file billable action, reusing the existing free-grant split + shadow row + ledger debit + `close()`→meter path on a standalone bookkeeping job (no lineage inputs, so nothing lineage-joins it). **`JobService.open(ctx, docUnits)`** opens that bare job. - **`AiCreateController.createSession`** — charges **one document per session** at creation (best-effort; entitlement is already enforced upstream by the class-level `@RequiresFeature(AI_SUPPORT)`). Follow-up edits (`outline` / `reprompt` / `draft` / `template` / `stream`) carry **no** charge — they have no charge hook, so "charge on create, follow-ups free" falls out naturally. Per the agreed scope: charge AI Create on create now; we can optimise follow-up handling later. **AI workflow categorisation (AUTOMATION vs AI) intentionally left as-is** (the orchestrator's automation header dominates by design). ## Tests - `PaygChargeInterceptorTest`: AI-tool route (no annotations) is in scope + **AI** category; same route with the automation header → **AUTOMATION**; a plain non-AI route still short-circuits. - `EntitlementGuardTest`: AI-tool route is in scope + gated on **AI_SUPPORT** (degraded team → 402; anonymous → 401 with `category: AI`). - `JobChargeServiceTest`: `chargeStandalone` charges + meters the paid portion for a subscribed team, draws the free grant (no meter) for an unsubscribed team, and rejects `BYPASSED`. `:saas:test` + `:saas:spotlessCheck` green; coverage gates met. ## Follow-ups (not in this PR) - Make AI Create follow-ups explicitly cheaper / chained if we want (currently free by absence of a hook). - Decide whether AI-tool-inside-a-workflow should bill as AI rather than AUTOMATION.
2026-07-16 19:33:11 +02:00 · 2026-06-11 20:10:44 +01:00
parent aaa2599e23
commit 37b4d24a95
9 changed files with 372 additions and 6 deletions
@@ -38,13 +38,19 @@ import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;

 import stirling.software.proprietary.security.database.repository.UserRepository;
+import stirling.software.proprietary.security.model.ApiKeyAuthenticationToken;
 import stirling.software.proprietary.security.model.User;
 import stirling.software.saas.ai.model.AiCreateSession;
 import stirling.software.saas.ai.repository.AiCreateSessionRepository;
 import stirling.software.saas.ai.service.AiCreateProxyService;
 import stirling.software.saas.ai.service.AiCreateSessionService;
 import stirling.software.saas.payg.cap.RequiresFeature;
+import stirling.software.saas.payg.charge.ChargeContext;
+import stirling.software.saas.payg.charge.JobChargeService;
+import stirling.software.saas.payg.model.BillingCategory;
 import stirling.software.saas.payg.model.FeatureGate;
+import stirling.software.saas.payg.model.JobSource;
+import stirling.software.saas.payg.model.ProcessType;
 import stirling.software.saas.service.CreditService;
 import stirling.software.saas.service.TeamCreditService;
 import stirling.software.saas.util.AuthenticationUtils;
@@ -67,6 +73,7 @@ public class AiCreateController {
    private final TeamCreditService teamCreditService;
    private final UserRepository userRepository;
    private final CreditHeaderUtils creditHeaderUtils;
+    private final JobChargeService jobChargeService;

    @PostMapping("/sessions")
    public ResponseEntity<CreateSessionResponse> createSession(
@@ -87,9 +94,45 @@ public class AiCreateController {
                session.getUserId(),
                session.getDocType(),
                session.getTemplateId());
+        chargeForCreate(session);
        return ResponseEntity.ok(new CreateSessionResponse(session.getSessionId()));
    }

+    /**
+     * Bill one document for a new AI Create session — creating a document is the charge point;
+     * follow-up edits on the same session (outline / reprompt / draft / template / stream) carry no
+     * charge. AI usage is billable, so a JWT (web) session counts the same as an API-key one.
+     *
+     * <p>Best-effort: a charge failure must not block the user's session. Entitlement is already
+     * enforced upstream — this controller is {@code @RequiresFeature(AI_SUPPORT)}, so the
+     * EntitlementGuard 402s a team with no AI allowance before we ever get here; this call only
+     * does the accounting (free-grant draw + Stripe meter).
+     */
+    private void chargeForCreate(AiCreateSession session) {
+        try {
+            Authentication auth = SecurityContextHolder.getContext().getAuthentication();
+            User user = AuthenticationUtils.getCurrentUser(auth, userRepository);
+            if (user == null || user.getTeam() == null) {
+                return;
+            }
+            JobSource source =
+                    auth instanceof ApiKeyAuthenticationToken ? JobSource.API : JobSource.WEB;
+            ChargeContext ctx =
+                    new ChargeContext(
+                            user.getId(),
+                            user.getTeam().getId(),
+                            source,
+                            ProcessType.SINGLE_TOOL,
+                            BillingCategory.AI);
+            jobChargeService.chargeStandalone(ctx, 1);
+        } catch (RuntimeException e) {
+            log.warn(
+                    "AI create session {} charge failed; session proceeds unbilled: {}",
+                    session.getSessionId(),
+                    e.getMessage());
+        }
+    }
+
    @DeleteMapping("/sessions/{sessionId}")
    public ResponseEntity<Void> deleteSession(@PathVariable String sessionId) {
        sessionService.deleteSessionForCurrentUser(sessionId);
@@ -0,0 +1,42 @@
+package stirling.software.saas.payg.cap;
+
+import org.springframework.web.servlet.HandlerMapping;
+
+import jakarta.servlet.http.HttpServletRequest;
+
+/**
+ * Central definition of the AI document-tool route namespace ({@code /api/v1/ai/tools/**}).
+ *
+ * <p>These tools (e.g. {@code PdfCommentAgentController}, {@code MathAuditorAgentController}) live
+ * in the {@code proprietary} module, which does not depend on {@code saas} and therefore cannot
+ * carry the saas-only {@link RequiresFeature} annotation. Rather than weaken the layering, the PAYG
+ * hot-path components recognise the path prefix instead:
+ *
+ * <ul>
+ *   <li>{@code PaygChargeInterceptor} brings these routes into scope and bills them as {@code
+ *       BillingCategory.AI} on a direct call (an orchestrator-dispatched call still resolves to
+ *       AUTOMATION first, via the {@code X-Stirling-Automation} header);
+ *   <li>{@code EntitlementGuard} gates them on {@link
+ *       stirling.software.saas.payg.model.FeatureGate#AI_SUPPORT}.
+ * </ul>
+ *
+ * <p>Kept as a single source of truth so the interceptor and the guard can never drift on what
+ * counts as an AI tool.
+ */
+public final class AiToolRoutes {
+
+    /** Trailing slash so it matches the tool sub-paths, not a bare {@code /api/v1/ai/tools}. */
+    public static final String PREFIX = "/api/v1/ai/tools/";
+
+    private AiToolRoutes() {}
+
+    /**
+     * True when the request resolved to an AI document-tool endpoint. Prefers the matched route
+     * pattern (context-path independent, set by Spring MVC) and falls back to the raw request URI.
+     */
+    public static boolean matches(HttpServletRequest request) {
+        Object pattern = request.getAttribute(HandlerMapping.BEST_MATCHING_PATTERN_ATTRIBUTE);
+        String path = pattern instanceof String s ? s : request.getRequestURI();
+        return path != null && path.startsWith(PREFIX);
+    }
+}
@@ -133,6 +133,51 @@ public class JobChargeService {
        return new ChargeOutcome(result.job().getId(), units, ChargeOutcome.Disposition.OPENED);
    }

+    /**
+     * Charge a fixed number of units for a billable action that isn't file/lineage-driven — e.g. an
+     * AI Create session, billed once per document at session creation. Opens a standalone
+     * bookkeeping job (no lineage inputs, so follow-up calls never lineage-join it), draws the
+     * free-grant split, and writes the shadow + ledger rows exactly as {@link #openProcess} does,
+     * then closes the job so the paid portion meters to Stripe via the same {@code afterCommit}
+     * path and idempotency key ({@code process:<jobId>:close}).
+     *
+     * <p>Each call is independent: there is no join/dedup, so two sessions charge twice (correct —
+     * each is a distinct document). The caller passes the unit count; the policy {@code
+     * minChargeUnits} floor still applies. Must not be called for {@link BillingCategory#BYPASSED}.
+     *
+     * @return the bookkeeping job id (mostly useful for tests / tracing)
+     */
+    @Transactional
+    public UUID chargeStandalone(ChargeContext ctx, int units) {
+        Objects.requireNonNull(ctx, "ctx");
+        if (ctx.billingCategory() == BillingCategory.BYPASSED) {
+            throw new IllegalArgumentException("chargeStandalone must not be called for BYPASSED");
+        }
+
+        PricingPolicy policy = policyService.getEffectivePolicy(ctx.ownerTeamId());
+        int chargeUnits = Math.max(units, policy.getMinChargeUnits());
+        int stepLimit = resolveStepLimit(policy, ctx.source());
+
+        JobContext jobCtx =
+                new JobContext(
+                        ctx.ownerUserId(),
+                        ctx.ownerTeamId(),
+                        ctx.source(),
+                        ctx.processType(),
+                        policy.getId(),
+                        stepLimit);
+        ProcessingJob job = jobService.open(jobCtx, chargeUnits);
+
+        int freeUsed = consumeFreeGrant(ctx, chargeUnits);
+        recordShadowRow(ctx, job.getId(), policy.getId(), chargeUnits, freeUsed);
+        recordLedgerDebit(ctx, job.getId(), policy.getId(), chargeUnits);
+
+        // Close immediately — nothing will lineage-join a standalone job — so the paid portion
+        // meters via the same afterCommit hook + idempotency key as a normal process completion.
+        close(job.getId());
+        return job.getId();
+    }
+
    /**
     * Draw this job's free portion from the team's one-time lifetime grant, atomically, and return
     * the units taken (0..{@code units}); the remainder is the paid portion that will be metered to
@@ -34,6 +34,7 @@ import stirling.software.common.annotations.AutoJobPostMapping;
 import stirling.software.proprietary.security.database.repository.UserRepository;
 import stirling.software.proprietary.security.model.ApiKeyAuthenticationToken;
 import stirling.software.proprietary.security.model.User;
+import stirling.software.saas.payg.cap.AiToolRoutes;
 import stirling.software.saas.payg.cap.RequiresFeature;
 import stirling.software.saas.payg.model.FeatureGate;
 import stirling.software.saas.util.AuthenticationUtils;
@@ -132,12 +133,17 @@ public class EntitlementGuard implements HandlerInterceptor {
                AnnotationUtils.findAnnotation(hm.getMethod(), RequiresFeature.class) != null
                        || AnnotationUtils.findAnnotation(hm.getBeanType(), RequiresFeature.class)
                                != null;
-        if (!hasAutoJobPostMapping && !hasRequiresFeature) {
+        // AI document tools (/api/v1/ai/tools/**) live in the proprietary module and can't carry
+        // @RequiresFeature; recognise them by path so they're gated on AI_SUPPORT — see
+        // AiToolRoutes and PaygChargeInterceptor, which classify the same routes as AI.
+        boolean aiToolRoute = AiToolRoutes.matches(request);
+        if (!hasAutoJobPostMapping && !hasRequiresFeature && !aiToolRoute) {
            skippedNoAnnotationCounter.increment();
            return true;
        }

-        FeatureGate[] required = resolveRequiredGates(hm);
+        FeatureGate[] required =
+                aiToolRoute ? new FeatureGate[] {FeatureGate.AI_SUPPORT} : resolveRequiredGates(hm);
        Authentication auth = SecurityContextHolder.getContext().getAuthentication();

        boolean anonymous = isAnonymous(auth);
@@ -38,6 +38,7 @@ import stirling.software.common.util.TempFileManager;
 import stirling.software.proprietary.security.database.repository.UserRepository;
 import stirling.software.proprietary.security.model.ApiKeyAuthenticationToken;
 import stirling.software.proprietary.security.model.User;
+import stirling.software.saas.payg.cap.AiToolRoutes;
 import stirling.software.saas.payg.cap.RequiresFeature;
 import stirling.software.saas.payg.charge.ChargeContext;
 import stirling.software.saas.payg.charge.ChargeOutcome;
@@ -199,7 +200,10 @@ public class PaygChargeInterceptor implements AsyncHandlerInterceptor {
                            || AnnotationUtils.findAnnotation(
                                            hm.getBeanType(), RequiresFeature.class)
                                    != null;
-            if (!hasAutoJobPostMapping && !hasRequiresFeature) {
+            // AI document tools (/api/v1/ai/tools/**) live in the proprietary module and can't
+            // carry @RequiresFeature, so they're recognised by path — see AiToolRoutes.
+            boolean aiToolRoute = AiToolRoutes.matches(request);
+            if (!hasAutoJobPostMapping && !hasRequiresFeature && !aiToolRoute) {
                callsShortCircuit.increment();
                return true;
            }
@@ -514,11 +518,14 @@ public class PaygChargeInterceptor implements AsyncHandlerInterceptor {
    /**
     * Resolve the {@link BillingCategory} for this request. Precedence: {@code
     * X-Stirling-Automation: true} or {@code @RequiresFeature(AUTOMATION)} → AUTOMATION;
-     * {@code @RequiresFeature(AI_SUPPORT)} → AI; API-key auth → API; otherwise BYPASSED (manual UI
-     * tool — short-circuited in {@link #preHandle}).
+     * {@code @RequiresFeature(AI_SUPPORT)} → AI; an AI document-tool route ({@link AiToolRoutes}) →
+     * AI; API-key auth → API; otherwise BYPASSED (manual UI tool — short-circuited in {@link
+     * #preHandle}).
     *
     * <p>Method-level {@code @RequiresFeature} wins over class-level. Multiple gates: AUTOMATION
-     * dominates AI within a single annotation.
+     * dominates AI within a single annotation. The AI-tool path check sits below the automation
+     * header on purpose: an AI tool dispatched inside a policy / AI workflow bills as AUTOMATION,
+     * while a direct call to it bills as AI.
     */
    private static BillingCategory determineCategory(
            HandlerMethod handler, HttpServletRequest request, Authentication auth) {
@@ -545,6 +552,11 @@ public class PaygChargeInterceptor implements AsyncHandlerInterceptor {
                return BillingCategory.AI;
            }
        }
+        // AI document tools (proprietary module, recognised by path). A direct call bills as AI; an
+        // orchestrator-dispatched call already returned AUTOMATION above via the automation header.
+        if (AiToolRoutes.matches(request)) {
+            return BillingCategory.AI;
+        }
        if (auth instanceof ApiKeyAuthenticationToken) {
            return BillingCategory.API;
        }
@@ -229,6 +229,20 @@ public class JobService {
        return new JoinOrOpenResult(saved, JoinOrOpenResult.Disposition.JOINED);
    }

+    /**
+     * Open a standalone process with no lineage inputs, for a billable action that isn't
+     * file/lineage-driven (e.g. an AI Create session). Because no input signatures are recorded,
+     * nothing downstream can lineage-join it — each such charge stands alone. {@code docUnits} is
+     * persisted so the charge service's shadow + ledger rows agree with the job.
+     */
+    @Transactional
+    public ProcessingJob open(JobContext ctx, int docUnits) {
+        Objects.requireNonNull(ctx, "ctx");
+        ProcessingJob job = openFresh(ctx, Map.of()).job();
+        job.setDocUnits(docUnits);
+        return jobRepository.save(job);
+    }
+
    private JoinOrOpenResult openFresh(
            JobContext ctx, Map<Path, Set<LineageSignature>> signaturesByInput) {
        ProcessingJob fresh = new ProcessingJob();