payg(c4): annotate AI controllers + tag InternalApiClient sub-steps as AUTOMATION

The saas PaygChargeInterceptor classifies a request as AI / AUTOMATION /
API / BYPASSED using a fixed precedence: X-Stirling-Automation header >
@RequiresFeature(AUTOMATION) > @RequiresFeature(AI_SUPPORT) > API-key
auth > BYPASSED (manual UI).

Plug the AI surface and the automation sub-step path:

AI controllers (saas → easy import):
- AiCreateController, AiCreateInternalController, AiProxyController each
  get class-level @RequiresFeature(AI_SUPPORT). The interceptor already
  falls back to the bean type via AnnotationUtils.findAnnotation when no
  method-level annotation is present.

Automation sub-step header (InternalApiClient in common):
- Every caller of InternalApiClient.post() is a parent automation flow
  (PipelineProcessor, AiWorkflowService, PolicyExecutor) running a child
  tool via loopback HTTP. Tag every dispatch with
  X-Stirling-Automation: true so the child step bills as AUTOMATION
  regardless of its own annotation — i.e. an AI-OCR step inside a
  policy run is correctly billed AUTOMATION rather than AI.

Deviation (documented in test javadoc): PipelineController lives in
core and PolicyController lives in proprietary. Neither module can
import @RequiresFeature from saas without a forbidden upward dependency
(both build under STIRLING_FLAVOR=core/proprietary where saas is
absent). Promoting RequiresFeature + FeatureGate to common is a
non-trivial refactor (touches 18 files across multiple C-bucket
commits) and out of scope for C4. The X-Stirling-Automation header
covers sub-step dispatch via InternalApiClient; direct user-facing
calls to /api/v1/pipeline/handleData therefore bill as BYPASSED (WEB)
or API (api-key auth) rather than AUTOMATION. Annotating those
controllers is left to a future refactor that promotes the cap types
to common.

Tests:
- RequiresFeatureAnnotationRolloutTest pins the AI controller
  classifications so a refactor can't silently regress to BYPASSED.
- InternalApiClientTest.postTagsRequestAsAutomation captures the
  HttpHeaders submitted to RestTemplate and asserts the marker.
This commit is contained in:
Connor Yoh
2026-06-09 14:48:36 +01:00
parent e80e1b2190
commit 6eb41b6931
6 changed files with 123 additions and 0 deletions
@@ -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<MultiValueMap<String, Object>> entity = new HttpEntity<>(body, headers);
RequestCallback requestCallback = restTemplate.httpEntityCallback(entity, Resource.class);
@@ -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<String, Object> 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<String, Object> body = new LinkedMultiValueMap<>();
@@ -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 {
@@ -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 {
@@ -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 {
@@ -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.
*
* <p><b>Out of scope</b>: {@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);
}
}