diff --git a/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java b/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java index f72eca14d5..c4663c7698 100644 --- a/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java +++ b/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java @@ -50,6 +50,16 @@ public class InternalApiClient { "^/api/v1/(general|misc|security|convert|filter)(/[A-Za-z0-9_-]+)+$" + "|^/api/v1/ai/tools(/[A-Za-z0-9_-]+)+$"); + /** + * Marker propagated on every internal sub-step dispatch so the saas PAYG interceptor classifies + * the call as {@code BillingCategory.AUTOMATION}. By construction every {@link + * InternalApiClient#post} caller is an automation surface (pipeline executor, AI workflow, + * policy runner) running a child tool inside a parent automation flow — see the saas {@code + * PaygChargeInterceptor.determineCategory} precedence chain, where this header dominates any + * per-tool {@code @RequiresFeature} annotation. + */ + public static final String AUTOMATION_HEADER = "X-Stirling-Automation"; + private final ServletContext servletContext; private final UserServiceInterface userService; private final TempFileManager tempFileManager; @@ -96,6 +106,11 @@ public class InternalApiClient { if (apiKey != null && !apiKey.isEmpty()) { headers.add("X-API-KEY", apiKey); } + // Tag the sub-step as automation so PAYG bills it under AUTOMATION regardless of which + // tool-level @RequiresFeature annotation the dispatched controller carries (e.g. an AI-OCR + // step inside a policy run must bill as AUTOMATION, not AI). Set unconditionally because + // every caller of this dispatcher is an automation surface by design. + headers.add(AUTOMATION_HEADER, "true"); HttpEntity> entity = new HttpEntity<>(body, headers); RequestCallback requestCallback = restTemplate.httpEntityCallback(entity, Resource.class); diff --git a/app/common/src/test/java/stirling/software/common/service/InternalApiClientTest.java b/app/common/src/test/java/stirling/software/common/service/InternalApiClientTest.java index 06543e47b0..7bab141888 100644 --- a/app/common/src/test/java/stirling/software/common/service/InternalApiClientTest.java +++ b/app/common/src/test/java/stirling/software/common/service/InternalApiClientTest.java @@ -59,6 +59,53 @@ class InternalApiClientTest { servletContext, userService, tempFileManager, environment, applicationProperties); } + @Test + void postTagsRequestAsAutomation() throws Exception { + // Every InternalApiClient.post() caller is a parent automation flow dispatching a child + // tool (pipeline executor, AI workflow, policy runner). Tagging the sub-step here means + // the saas PaygChargeInterceptor classifies it as BillingCategory.AUTOMATION regardless of + // the dispatched controller's @RequiresFeature — so an AI-OCR step inside a policy run + // bills as AUTOMATION, not AI. The header value is the literal string "true" because the + // interceptor compares case-insensitively-trimmed against that token. + MultiValueMap body = new LinkedMultiValueMap<>(); + body.add("fileInput", namedResource("input.pdf", "data")); + + Path tempPath = Files.createTempFile("internal-api-automation-test", ".tmp"); + TempFile tempFile = mock(TempFile.class); + when(tempFile.getPath()).thenReturn(tempPath); + when(tempFile.getFile()).thenReturn(tempPath.toFile()); + when(tempFileManager.createManagedTempFile("internal-api")).thenReturn(tempFile); + + HttpHeaders[] captured = {null}; + + try (var ignored = + mockConstruction( + RestTemplate.class, + (rt, ctx) -> { + when(rt.httpEntityCallback(any(), eq(Resource.class))) + .thenAnswer( + inv -> { + HttpEntity entity = inv.getArgument(0); + captured[0] = entity.getHeaders(); + return (RequestCallback) req -> {}; + }); + when(rt.execute(anyString(), eq(HttpMethod.POST), any(), any())) + .thenAnswer(inv -> fakeOkResponse(inv.getArgument(3))); + })) { + + InternalApiClient mockedClient = newClient(); + mockedClient.post("/api/v1/general/merge-pdfs", body); + + assertNotNull(captured[0]); + assertEquals( + "true", + captured[0].getFirst(InternalApiClient.AUTOMATION_HEADER), + "Sub-step dispatch must carry the automation marker header"); + } finally { + Files.deleteIfExists(tempPath); + } + } + @Test void postDoesNotForceContentType() throws Exception { MultiValueMap body = new LinkedMultiValueMap<>(); diff --git a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java index 85520a239b..126ce8d473 100644 --- a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java +++ b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java @@ -40,6 +40,8 @@ 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.model.FeatureGate; import stirling.software.saas.service.CreditService; import stirling.software.saas.service.TeamCreditService; import stirling.software.saas.util.AuthenticationUtils; @@ -49,6 +51,7 @@ import stirling.software.saas.util.CreditHeaderUtils; @Profile("saas") @RequestMapping("/api/v1/ai/create") @RequiredArgsConstructor +@RequiresFeature(FeatureGate.AI_SUPPORT) @Slf4j public class AiCreateController { diff --git a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateInternalController.java b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateInternalController.java index 34c44176e5..43b7a3f059 100644 --- a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateInternalController.java +++ b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateInternalController.java @@ -23,11 +23,14 @@ import lombok.extern.slf4j.Slf4j; import stirling.software.saas.ai.model.AiCreateSession; import stirling.software.saas.ai.model.AiCreateSessionStatus; import stirling.software.saas.ai.service.AiCreateSessionService; +import stirling.software.saas.payg.cap.RequiresFeature; +import stirling.software.saas.payg.model.FeatureGate; @RestController @Profile("saas") @RequestMapping("/api/v1/ai/create/internal") @RequiredArgsConstructor +@RequiresFeature(FeatureGate.AI_SUPPORT) @Slf4j public class AiCreateInternalController { diff --git a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiProxyController.java b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiProxyController.java index 074da8707c..4dc3592b99 100644 --- a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiProxyController.java +++ b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiProxyController.java @@ -25,6 +25,8 @@ import lombok.extern.slf4j.Slf4j; import stirling.software.proprietary.security.database.repository.UserRepository; import stirling.software.proprietary.security.model.User; import stirling.software.saas.ai.service.AiProxyService; +import stirling.software.saas.payg.cap.RequiresFeature; +import stirling.software.saas.payg.model.FeatureGate; import stirling.software.saas.service.CreditService; import stirling.software.saas.service.TeamCreditService; import stirling.software.saas.util.AuthenticationUtils; @@ -33,6 +35,7 @@ import stirling.software.saas.util.CreditHeaderUtils; @RestController @Profile("saas") @RequestMapping("/api/v1/ai") +@RequiresFeature(FeatureGate.AI_SUPPORT) @Slf4j public class AiProxyController { diff --git a/app/saas/src/test/java/stirling/software/saas/payg/cap/RequiresFeatureAnnotationRolloutTest.java b/app/saas/src/test/java/stirling/software/saas/payg/cap/RequiresFeatureAnnotationRolloutTest.java new file mode 100644 index 0000000000..c2d5fcd8b7 --- /dev/null +++ b/app/saas/src/test/java/stirling/software/saas/payg/cap/RequiresFeatureAnnotationRolloutTest.java @@ -0,0 +1,52 @@ +package stirling.software.saas.payg.cap; + +import static org.assertj.core.api.Assertions.assertThat; + +import org.junit.jupiter.api.Test; +import org.springframework.core.annotation.AnnotationUtils; + +import stirling.software.saas.ai.controller.AiCreateController; +import stirling.software.saas.ai.controller.AiCreateInternalController; +import stirling.software.saas.ai.controller.AiProxyController; +import stirling.software.saas.payg.model.FeatureGate; + +/** + * Annotation-rollout guard. The saas {@code PaygChargeInterceptor} reads class-level + * {@code @RequiresFeature} via {@link AnnotationUtils#findAnnotation(Class, Class)} to decide + * whether a request bills as {@code AI}, {@code AUTOMATION}, or falls through to the auth-derived + * default. These tests pin the gate on each AI surface so the classification can't silently regress + * to {@code BYPASSED} if someone strips the annotation while refactoring. + * + *

Out of scope: {@code PipelineController} (in core) and {@code PolicyController} (in + * proprietary) — neither module can import {@code @RequiresFeature} from saas without a forbidden + * upward dependency. Their automation classification is enforced via the {@code + * X-Stirling-Automation} header set unconditionally by {@code InternalApiClient.post}; see the + * dedicated test in that module. + */ +class RequiresFeatureAnnotationRolloutTest { + + @Test + void aiCreateController_isClassifiedAsAiSupport() { + RequiresFeature ann = + AnnotationUtils.findAnnotation(AiCreateController.class, RequiresFeature.class); + assertThat(ann).isNotNull(); + assertThat(ann.value()).containsExactly(FeatureGate.AI_SUPPORT); + } + + @Test + void aiCreateInternalController_isClassifiedAsAiSupport() { + RequiresFeature ann = + AnnotationUtils.findAnnotation( + AiCreateInternalController.class, RequiresFeature.class); + assertThat(ann).isNotNull(); + assertThat(ann.value()).containsExactly(FeatureGate.AI_SUPPORT); + } + + @Test + void aiProxyController_isClassifiedAsAiSupport() { + RequiresFeature ann = + AnnotationUtils.findAnnotation(AiProxyController.class, RequiresFeature.class); + assertThat(ann).isNotNull(); + assertThat(ann.value()).containsExactly(FeatureGate.AI_SUPPORT); + } +}