mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
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.
This commit is contained in:
ConnorYoh
@@ -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
|
||||
|
||||
+8
-2
@@ -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);
|
||||
|
||||
+16
-4
@@ -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();
|
||||
|
||||
@@ -825,6 +825,94 @@ class JobChargeServiceTest {
|
||||
verify(jobService).close(jobId);
|
||||
}
|
||||
|
||||
// --- chargeStandalone() — non-file billable actions (e.g. AI Create) -----------------------
|
||||
|
||||
@Test
|
||||
void chargeStandalone_subscribedTeam_chargesAndMetersPaidPortion() {
|
||||
// AI Create-style charge: one standalone bookkeeping job, free split, ledger debit, meter.
|
||||
long teamId = 100L;
|
||||
PricingPolicy policy = stubPolicy(/*minCharge*/ 1, Map.of(JobSource.WEB, 10));
|
||||
when(policyService.getEffectivePolicy(teamId)).thenReturn(policy);
|
||||
|
||||
UUID jobId = UUID.randomUUID();
|
||||
when(jobService.open(any(JobContext.class), eq(1))).thenReturn(openJob(jobId));
|
||||
when(jobService.close(jobId)).thenReturn(openJob(jobId));
|
||||
|
||||
// Subscribed, no free grant left → the whole unit is paid and meters.
|
||||
PaygTeamExtensions ext = new PaygTeamExtensions();
|
||||
ext.setTeamId(teamId);
|
||||
ext.setStripeCustomerId("cus_x");
|
||||
ext.setPaygSubscriptionId("sub_x");
|
||||
ext.setFreeUnitsRemaining(0L);
|
||||
when(teamExtRepo.findByIdForUpdate(teamId)).thenReturn(Optional.of(ext));
|
||||
when(teamExtRepo.findById(teamId)).thenReturn(Optional.of(ext));
|
||||
when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId))
|
||||
.thenReturn(Optional.of(chargedShadowRow(jobId, teamId, 1, 0, BillingCategory.AI)));
|
||||
|
||||
ChargeContext ctx =
|
||||
new ChargeContext(
|
||||
7L, teamId, JobSource.WEB, ProcessType.SINGLE_TOOL, BillingCategory.AI);
|
||||
ArgumentCaptor<WalletLedgerEntry> ledger = ArgumentCaptor.forClass(WalletLedgerEntry.class);
|
||||
|
||||
withTransactionSynchronization(() -> service.chargeStandalone(ctx, 1));
|
||||
|
||||
verify(jobService).open(any(JobContext.class), eq(1));
|
||||
verify(jobService).close(jobId);
|
||||
verify(ledgerRepo).save(ledger.capture());
|
||||
assertThat(ledger.getValue().getEntryType()).isEqualTo(LedgerEntryType.DEBIT);
|
||||
assertThat(ledger.getValue().getAmountUnits()).isEqualTo(-1);
|
||||
assertThat(ledger.getValue().getBillingCategory()).isEqualTo(BillingCategory.AI);
|
||||
verify(shadowRepo).save(any(PaygShadowCharge.class));
|
||||
// Paid portion (1) metered to Stripe after commit, keyed by the standard process key.
|
||||
verify(meterReporter)
|
||||
.recordUsage(
|
||||
eq(teamId),
|
||||
eq("cus_x"),
|
||||
eq(1),
|
||||
eq(BillingCategory.AI),
|
||||
eq("process:" + jobId + ":close"),
|
||||
eq(jobId));
|
||||
}
|
||||
|
||||
@Test
|
||||
void chargeStandalone_freeTeamWithGrant_drawsGrantAndDoesNotMeter() {
|
||||
long teamId = 100L;
|
||||
PricingPolicy policy = stubPolicy(1, Map.of(JobSource.WEB, 10));
|
||||
when(policyService.getEffectivePolicy(teamId)).thenReturn(policy);
|
||||
|
||||
UUID jobId = UUID.randomUUID();
|
||||
when(jobService.open(any(JobContext.class), eq(1))).thenReturn(openJob(jobId));
|
||||
when(jobService.close(jobId)).thenReturn(openJob(jobId));
|
||||
|
||||
// Free grant available, no subscription → the unit is drawn from the grant, nothing meters.
|
||||
PaygTeamExtensions ext = new PaygTeamExtensions();
|
||||
ext.setTeamId(teamId);
|
||||
ext.setFreeUnitsRemaining(50L);
|
||||
when(teamExtRepo.findByIdForUpdate(teamId)).thenReturn(Optional.of(ext));
|
||||
when(teamExtRepo.findById(teamId)).thenReturn(Optional.of(ext));
|
||||
when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId))
|
||||
.thenReturn(Optional.of(chargedShadowRow(jobId, teamId, 1, 1, BillingCategory.AI)));
|
||||
|
||||
ChargeContext ctx =
|
||||
new ChargeContext(
|
||||
7L, teamId, JobSource.WEB, ProcessType.SINGLE_TOOL, BillingCategory.AI);
|
||||
|
||||
withTransactionSynchronization(() -> service.chargeStandalone(ctx, 1));
|
||||
|
||||
assertThat(ext.getFreeUnitsRemaining()).isEqualTo(49L);
|
||||
verify(meterReporter, never())
|
||||
.recordUsage(any(), any(), Mockito.anyInt(), any(), any(), any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void chargeStandalone_bypassedCategory_throws() {
|
||||
ChargeContext ctx =
|
||||
new ChargeContext(
|
||||
7L, 100L, JobSource.WEB, ProcessType.SINGLE_TOOL, BillingCategory.BYPASSED);
|
||||
assertThatThrownBy(() -> service.chargeStandalone(ctx, 1))
|
||||
.isInstanceOf(IllegalArgumentException.class);
|
||||
}
|
||||
|
||||
private static void withTransactionSynchronization(Runnable body) {
|
||||
TransactionSynchronizationManager.initSynchronization();
|
||||
try {
|
||||
|
||||
+49
@@ -149,6 +149,55 @@ class EntitlementGuardTest {
|
||||
assertThat(body.get("category").asText()).isEqualTo("AI");
|
||||
}
|
||||
|
||||
@Test
|
||||
void aiToolRoute_noAnnotation_isInScopeAndGatedOnAiSupport() throws Exception {
|
||||
// /api/v1/ai/tools/** controllers live in the proprietary module and carry no
|
||||
// @RequiresFeature; the guard recognises them by path and gates on AI_SUPPORT. A degraded
|
||||
// team is 402'd even though the handler has no annotation.
|
||||
UUID supabaseId = UUID.randomUUID();
|
||||
SecurityContextHolder.getContext().setAuthentication(jwtAuth(supabaseId));
|
||||
when(userRepository.findBySupabaseId(supabaseId))
|
||||
.thenReturn(Optional.of(userWithTeam(7L, 42L)));
|
||||
when(entitlementService.getSnapshot(42L)).thenReturn(degradedSnapshot());
|
||||
|
||||
HandlerMethod hm = handlerFor("plainEndpoint"); // no annotations
|
||||
MockHttpServletRequest req = new MockHttpServletRequest();
|
||||
req.setRequestURI("/api/v1/ai/tools/pdf-comment-agent");
|
||||
MockHttpServletResponse res = new MockHttpServletResponse();
|
||||
|
||||
boolean proceed = guard.preHandle(req, res, hm);
|
||||
|
||||
assertThat(proceed).isFalse();
|
||||
assertThat(res.getStatus()).isEqualTo(402);
|
||||
JsonNode body = json.readTree(res.getContentAsByteArray());
|
||||
assertThat(body.get("error").asText()).isEqualTo("FEATURE_DEGRADED");
|
||||
assertThat(body.get("missingGates").get(0).asText()).isEqualTo("AI_SUPPORT");
|
||||
verify(entitlementService).getSnapshot(42L);
|
||||
}
|
||||
|
||||
@Test
|
||||
void aiToolRoute_anonymous_returns401WithAiCategory() throws Exception {
|
||||
SecurityContextHolder.getContext()
|
||||
.setAuthentication(
|
||||
new AnonymousAuthenticationToken(
|
||||
"key",
|
||||
"anonymousUser",
|
||||
List.of(new SimpleGrantedAuthority("ROLE_ANONYMOUS"))));
|
||||
|
||||
HandlerMethod hm = handlerFor("plainEndpoint");
|
||||
MockHttpServletRequest req = new MockHttpServletRequest();
|
||||
req.setRequestURI("/api/v1/ai/tools/math-auditor-agent");
|
||||
MockHttpServletResponse res = new MockHttpServletResponse();
|
||||
|
||||
boolean proceed = guard.preHandle(req, res, hm);
|
||||
|
||||
assertThat(proceed).isFalse();
|
||||
assertThat(res.getStatus()).isEqualTo(401);
|
||||
JsonNode body = json.readTree(res.getContentAsByteArray());
|
||||
assertThat(body.get("error").asText()).isEqualTo("SIGNUP_REQUIRED");
|
||||
assertThat(body.get("category").asText()).isEqualTo("AI");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------------------
|
||||
// Anonymous user
|
||||
// ---------------------------------------------------------------------------------------
|
||||
|
||||
+67
@@ -612,6 +612,73 @@ class PaygChargeInterceptorTest {
|
||||
.isEqualTo(stirling.software.saas.payg.model.BillingCategory.AUTOMATION);
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandle_aiToolRoute_inScopeAndCategoryAi() throws Exception {
|
||||
// AI document tools (/api/v1/ai/tools/**) live in the proprietary module and carry no PAYG
|
||||
// annotation. The interceptor recognises them by path → in scope + AI category, so a direct
|
||||
// multipart call opens a charge. Handler has NO annotations (handlerMethodForPlain).
|
||||
authenticateWithUser(makeUser(7L, 42L));
|
||||
UUID jobId = UUID.randomUUID();
|
||||
when(chargeService.openProcess(any(), anyList()))
|
||||
.thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED));
|
||||
org.mockito.ArgumentCaptor<stirling.software.saas.payg.charge.ChargeContext> ctxCaptor =
|
||||
org.mockito.ArgumentCaptor.forClass(
|
||||
stirling.software.saas.payg.charge.ChargeContext.class);
|
||||
|
||||
MockMultipartHttpServletRequest req = newMultipart();
|
||||
req.setRequestURI("/api/v1/ai/tools/pdf-comment-agent");
|
||||
req.addFile(
|
||||
new MockMultipartFile("fileInput", "x.pdf", "application/pdf", "abc".getBytes()));
|
||||
|
||||
interceptor.preHandle(req, new MockHttpServletResponse(), handlerMethodForPlain());
|
||||
|
||||
verify(chargeService).openProcess(ctxCaptor.capture(), anyList());
|
||||
assertThat(ctxCaptor.getValue().billingCategory())
|
||||
.isEqualTo(stirling.software.saas.payg.model.BillingCategory.AI);
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandle_aiToolRoute_withAutomationHeader_isAutomation() throws Exception {
|
||||
// An AI tool dispatched inside a policy / AI workflow carries X-Stirling-Automation: true →
|
||||
// AUTOMATION wins over the AI path rule (the header is checked first).
|
||||
authenticateWithUser(makeUser(7L, 42L));
|
||||
UUID jobId = UUID.randomUUID();
|
||||
when(chargeService.openProcess(any(), anyList()))
|
||||
.thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED));
|
||||
org.mockito.ArgumentCaptor<stirling.software.saas.payg.charge.ChargeContext> ctxCaptor =
|
||||
org.mockito.ArgumentCaptor.forClass(
|
||||
stirling.software.saas.payg.charge.ChargeContext.class);
|
||||
|
||||
MockMultipartHttpServletRequest req = newMultipart();
|
||||
req.setRequestURI("/api/v1/ai/tools/pdf-comment-agent");
|
||||
req.addFile(
|
||||
new MockMultipartFile("fileInput", "x.pdf", "application/pdf", "abc".getBytes()));
|
||||
req.addHeader("X-Stirling-Automation", "true");
|
||||
|
||||
interceptor.preHandle(req, new MockHttpServletResponse(), handlerMethodForPlain());
|
||||
|
||||
verify(chargeService).openProcess(ctxCaptor.capture(), anyList());
|
||||
assertThat(ctxCaptor.getValue().billingCategory())
|
||||
.isEqualTo(stirling.software.saas.payg.model.BillingCategory.AUTOMATION);
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandle_plainRouteNoAnnotations_stillShortCircuits() throws Exception {
|
||||
// Guard against the path rule being too broad: a non-AI-tools route with no annotations
|
||||
// must
|
||||
// still short-circuit (BYPASSED path), unaffected by the AI-tools recognition.
|
||||
authenticateWithUser(makeUser(7L, 42L));
|
||||
|
||||
MockMultipartHttpServletRequest req = newMultipart(); // URI = /api/v1/security/test-tool
|
||||
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
|
||||
|
||||
boolean cont =
|
||||
interceptor.preHandle(req, new MockHttpServletResponse(), handlerMethodForPlain());
|
||||
|
||||
assertThat(cont).isTrue();
|
||||
verify(chargeService, never()).openProcess(any(), anyList());
|
||||
}
|
||||
|
||||
// --- helpers --------------------------------------------------------------------------------
|
||||
|
||||
private MockMultipartHttpServletRequest newMultipart() {
|
||||
|
||||
Reference in New Issue
Block a user