mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 19:33:11 +02:00
Change AI engine to execute tools in Java instead of on frontend (#6116)
# Description of Changes Redesign AI engine so that it autogenerates the `tool_models.py` file from the OpenAPI spec so the Python has access to the Java API parameters and the full list of Java tools that it can run. CI ensures that whenever someone modifies a tool endpoint that the AI enigne tool models get updated as well (the dev gets told to run `task engine:tool-models`). There's loads of advantages to having the Java be the one that actually executes the tools, rather than the frontend as it was previously set up to theoretically use: - The AI gets much better descriptions of the params from the API docs - It'll be usable headless in the future so a Java daemon could run to execute ops on files in a folder without the need for the UI to run - The Java already has all the logic it needs to execute the tools - We don't need to parse the TypeScript to find the API (which is hard because the TS wasn't designed to be computer-read to extract the API) I've also hooked up the prototype frontend to ensure it's working properly, and have built it in a way that all the tool names can be translated properly, which was always an issue with previous prototypes of this. --------- Co-authored-by: Anthony Stirling <[email protected]> Co-authored-by: EthanHealy01 <[email protected]>
This commit is contained in:
co-authored by
Anthony Stirling
EthanHealy01
parent
cc9650e7a3
commit
e5767ed58b
@@ -0,0 +1,184 @@
|
||||
package stirling.software.common.service;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.UncheckedIOException;
|
||||
import java.net.URLDecoder;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.nio.file.Files;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.core.env.Environment;
|
||||
import org.springframework.core.io.FileSystemResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.http.*;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
import org.springframework.web.client.RequestCallback;
|
||||
import org.springframework.web.client.RestTemplate;
|
||||
|
||||
import jakarta.servlet.ServletContext;
|
||||
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.model.enumeration.Role;
|
||||
import stirling.software.common.util.TempFile;
|
||||
import stirling.software.common.util.TempFileManager;
|
||||
|
||||
/**
|
||||
* Dispatches HTTP POST requests to internal Stirling API endpoints via loopback. Used by
|
||||
* PipelineProcessor and AiWorkflowService to execute tool operations programmatically without
|
||||
* leaving the JVM network stack.
|
||||
*/
|
||||
@Service
|
||||
@Slf4j
|
||||
public class InternalApiClient {
|
||||
|
||||
// Allowlist for internal dispatch. Matches a fixed namespace prefix,
|
||||
// but rejects traversal (..), URL-encoding (%), query/fragment, backslashes, and any other
|
||||
// character that could alter the resolved endpoint on the local Spring server.
|
||||
private static final Pattern ALLOWED_ENDPOINT_PATH =
|
||||
Pattern.compile("^/api/v1/(general|misc|security|convert|filter)(/[A-Za-z0-9_-]+)+$");
|
||||
|
||||
private final ServletContext servletContext;
|
||||
private final UserServiceInterface userService;
|
||||
private final TempFileManager tempFileManager;
|
||||
private final Environment environment;
|
||||
|
||||
public InternalApiClient(
|
||||
ServletContext servletContext,
|
||||
@Autowired(required = false) UserServiceInterface userService,
|
||||
TempFileManager tempFileManager,
|
||||
Environment environment) {
|
||||
this.servletContext = servletContext;
|
||||
this.userService = userService;
|
||||
this.tempFileManager = tempFileManager;
|
||||
this.environment = environment;
|
||||
}
|
||||
|
||||
/**
|
||||
* POST to an internal API endpoint. The endpointPath must start with one of the allowed
|
||||
* prefixes (e.g. {@code /api/v1/misc/compress-pdf}).
|
||||
*
|
||||
* @param endpointPath API path (e.g. {@code /api/v1/general/rotate-pdf})
|
||||
* @param body multipart form body (fileInput + parameters)
|
||||
* @return response with the result file as a {@link TempFileResource} body
|
||||
*/
|
||||
public ResponseEntity<Resource> post(String endpointPath, MultiValueMap<String, Object> body) {
|
||||
validateUrl(endpointPath);
|
||||
String url = getBaseUrl() + endpointPath;
|
||||
|
||||
RestTemplate restTemplate = new RestTemplate();
|
||||
HttpHeaders headers = new HttpHeaders();
|
||||
String apiKey = getApiKeyForUser();
|
||||
if (apiKey != null && !apiKey.isEmpty()) {
|
||||
headers.add("X-API-KEY", apiKey);
|
||||
}
|
||||
|
||||
HttpEntity<MultiValueMap<String, Object>> entity = new HttpEntity<>(body, headers);
|
||||
RequestCallback requestCallback = restTemplate.httpEntityCallback(entity, Resource.class);
|
||||
|
||||
return restTemplate.execute(
|
||||
url,
|
||||
HttpMethod.POST,
|
||||
requestCallback,
|
||||
response -> {
|
||||
try {
|
||||
TempFile tempFile = tempFileManager.createManagedTempFile("internal-api");
|
||||
Files.copy(
|
||||
response.getBody(),
|
||||
tempFile.getPath(),
|
||||
java.nio.file.StandardCopyOption.REPLACE_EXISTING);
|
||||
String filename = extractFilename(response.getHeaders());
|
||||
TempFileResource resource = new TempFileResource(tempFile, filename);
|
||||
return ResponseEntity.status(response.getStatusCode())
|
||||
.headers(response.getHeaders())
|
||||
.body(resource);
|
||||
} catch (IOException e) {
|
||||
throw new UncheckedIOException(e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract the filename from a response's {@code Content-Disposition} header. Returns {@code
|
||||
* null} if the header is missing or has no filename.
|
||||
*/
|
||||
private static String extractFilename(HttpHeaders headers) {
|
||||
String contentDisposition = headers.getFirst(HttpHeaders.CONTENT_DISPOSITION);
|
||||
if (contentDisposition == null || contentDisposition.isBlank()) {
|
||||
return null;
|
||||
}
|
||||
for (String part : contentDisposition.split(";")) {
|
||||
String trimmed = part.trim();
|
||||
if (trimmed.startsWith("filename")) {
|
||||
String[] kv = trimmed.split("=", 2);
|
||||
if (kv.length != 2) {
|
||||
continue;
|
||||
}
|
||||
String value = kv[1].trim().replace("\"", "");
|
||||
return URLDecoder.decode(value, StandardCharsets.UTF_8);
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
private String getBaseUrl() {
|
||||
// Resolve the port lazily so desktop mode (server.port=0, OS-assigned) dispatches to the
|
||||
// actual bound port. Spring publishes local.server.port once the web server is up; fall
|
||||
// back to the configured server.port for early calls (tests, non-web contexts).
|
||||
String port = environment.getProperty("local.server.port");
|
||||
if (port == null) {
|
||||
port = environment.getProperty("server.port", "8080");
|
||||
}
|
||||
return "http://localhost:" + port + servletContext.getContextPath();
|
||||
}
|
||||
|
||||
private String getApiKeyForUser() {
|
||||
if (userService == null) return "";
|
||||
String username = userService.getCurrentUsername();
|
||||
if (username != null && !username.equals("anonymousUser")) {
|
||||
return userService.getApiKeyForUser(username);
|
||||
}
|
||||
return userService.getApiKeyForUser(Role.INTERNAL_API_USER.getRoleId());
|
||||
}
|
||||
|
||||
private void validateUrl(String endpointPath) {
|
||||
if (endpointPath == null || !ALLOWED_ENDPOINT_PATH.matcher(endpointPath).matches()) {
|
||||
log.warn("Blocked internal API request to disallowed path: {}", endpointPath);
|
||||
throw new SecurityException(
|
||||
"Internal API dispatch not permitted for endpoint: " + endpointPath);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A {@link FileSystemResource} that holds a reference to its backing {@link TempFile}.
|
||||
*
|
||||
* <p>If a display filename is supplied (typically parsed from the upstream response's {@code
|
||||
* Content-Disposition} header), it is returned from {@link #getFilename()} instead of the
|
||||
* underlying temp file's path-based name.
|
||||
*/
|
||||
public static class TempFileResource extends FileSystemResource {
|
||||
private final TempFile tempFile;
|
||||
private final String displayFilename;
|
||||
|
||||
public TempFileResource(TempFile tempFile) {
|
||||
this(tempFile, null);
|
||||
}
|
||||
|
||||
public TempFileResource(TempFile tempFile, String displayFilename) {
|
||||
super(tempFile.getFile());
|
||||
this.tempFile = tempFile;
|
||||
this.displayFilename = displayFilename;
|
||||
}
|
||||
|
||||
public TempFile getTempFile() {
|
||||
return tempFile;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return displayFilename != null ? displayFilename : super.getFilename();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package stirling.software.common.service;
|
||||
|
||||
/** Provides metadata about tool endpoints for internal dispatch. */
|
||||
public interface ToolMetadataService {
|
||||
|
||||
/** Returns true if the given operation path accepts multiple input files. */
|
||||
boolean isMultiInput(String operationPath);
|
||||
|
||||
/**
|
||||
* Returns true when the endpoint's ZIP response is a transport for multiple typed results and
|
||||
* should be unpacked: multi-output endpoints (Type:SIMO / Type:MIMO) and wrapper declarations
|
||||
* such as {@code Output:ZIP-PDF} or {@code Output:IMAGE/ZIP}.
|
||||
*
|
||||
* <p>Returns false for a bare {@code Output:ZIP} (e.g. {@code get-attachments}), where the
|
||||
* archive itself is the deliverable and should be kept packed.
|
||||
*/
|
||||
boolean shouldUnpackZipResponse(String operationPath);
|
||||
}
|
||||
@@ -538,9 +538,9 @@ public final class RegexPatternUtils {
|
||||
getPattern("[^a-zA-Z0-9 ]"); // Input sanitization
|
||||
getPattern("[^a-zA-Z0-9]"); // Filename sanitization
|
||||
// API doc patterns
|
||||
getPattern("Output:(\\w+)"); // precompiled single-escaped for runtime regex \w
|
||||
getPattern("Input:(\\w+)");
|
||||
getPattern("Type:(\\w+)");
|
||||
getPattern("Output:\\s*(\\w+)");
|
||||
getPattern("Input:\\s*(\\w+)");
|
||||
getPattern("Type:\\s*(\\w+)");
|
||||
log.debug("Pre-compiled {} common regex patterns", patternCache.size());
|
||||
}
|
||||
|
||||
@@ -552,19 +552,19 @@ public final class RegexPatternUtils {
|
||||
|
||||
/* Pattern for matching Output:<TYPE> in API descriptions */
|
||||
public Pattern getApiDocOutputTypePattern() {
|
||||
return getPattern("Output:(\\w+)");
|
||||
return getPattern("Output:\\s*(\\w+)");
|
||||
}
|
||||
|
||||
/* Pattern for matching Input:<TYPE> in API descriptions */
|
||||
public Pattern getApiDocInputTypePattern() {
|
||||
return getPattern("Input:(\\w+)");
|
||||
return getPattern("Input:\\s*(\\w+)");
|
||||
}
|
||||
|
||||
/**
|
||||
* Pattern for matching Type:<CODE> in API descriptions
|
||||
*/
|
||||
public Pattern getApiDocTypePattern() {
|
||||
return getPattern("Type:(\\w+)");
|
||||
return getPattern("Type:\\s*(\\w+)");
|
||||
}
|
||||
|
||||
/* Pattern for validating file extensions (2-4 alphanumeric, case-insensitive) */
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
package stirling.software.common.util;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.OutputStream;
|
||||
import java.nio.file.Files;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.function.Consumer;
|
||||
import java.util.zip.ZipEntry;
|
||||
import java.util.zip.ZipInputStream;
|
||||
|
||||
import org.springframework.core.io.FileSystemResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
|
||||
import io.github.pixee.security.ZipSecurity;
|
||||
|
||||
import lombok.experimental.UtilityClass;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
/**
|
||||
* Helpers for detecting and extracting ZIP-formatted responses returned from Stirling API
|
||||
* endpoints. Shared between {@code PipelineProcessor} and {@code AiWorkflowService} so both callers
|
||||
* unpack ZIPs consistently (hardened against zip-slip, depth-limited, backed by managed temp
|
||||
* files).
|
||||
*/
|
||||
@Slf4j
|
||||
@UtilityClass
|
||||
public class ZipExtractionUtils {
|
||||
|
||||
private static final int MAX_UNZIP_DEPTH = 10;
|
||||
private static final byte[] ZIP_MAGIC = {0x50, 0x4B, 0x03, 0x04};
|
||||
|
||||
/**
|
||||
* Returns true if the resource starts with the standard ZIP magic bytes. CBZ files are
|
||||
* explicitly treated as non-ZIP.
|
||||
*/
|
||||
public static boolean isZip(Resource data) throws IOException {
|
||||
return isZip(data, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns true if the resource starts with the standard ZIP magic bytes. Files named with the
|
||||
* {@code .cbz} extension are excluded (handled separately by the comic viewer).
|
||||
*/
|
||||
public static boolean isZip(Resource data, String filename) throws IOException {
|
||||
if (data == null || data.contentLength() < ZIP_MAGIC.length) {
|
||||
return false;
|
||||
}
|
||||
if (filename != null && filename.toLowerCase().endsWith(".cbz")) {
|
||||
return false;
|
||||
}
|
||||
try (InputStream is = data.getInputStream()) {
|
||||
byte[] header = new byte[ZIP_MAGIC.length];
|
||||
if (is.read(header) < ZIP_MAGIC.length) {
|
||||
return false;
|
||||
}
|
||||
for (int i = 0; i < ZIP_MAGIC.length; i++) {
|
||||
if (header[i] != ZIP_MAGIC[i]) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract a ZIP resource into a flat list of resources, one per file entry. Nested ZIPs are
|
||||
* recursively extracted up to {@link #MAX_UNZIP_DEPTH}. Each entry is materialized as a
|
||||
* hardened-extracted managed temp file so downstream consumers can stream the bytes.
|
||||
*/
|
||||
public static List<Resource> extractZip(Resource zip, TempFileManager tempFileManager)
|
||||
throws IOException {
|
||||
return extractZip(zip, tempFileManager, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract a ZIP resource into a flat list of resources. Each created {@link TempFile} is also
|
||||
* passed to {@code tempFileConsumer} when non-null, giving callers the option to register the
|
||||
* temp files with an auxiliary lifecycle (e.g. {@code PipelineResult}).
|
||||
*/
|
||||
public static List<Resource> extractZip(
|
||||
Resource zip, TempFileManager tempFileManager, Consumer<TempFile> tempFileConsumer)
|
||||
throws IOException {
|
||||
return extractZipInternal(zip, tempFileManager, tempFileConsumer, 0);
|
||||
}
|
||||
|
||||
private static List<Resource> extractZipInternal(
|
||||
Resource zip,
|
||||
TempFileManager tempFileManager,
|
||||
Consumer<TempFile> tempFileConsumer,
|
||||
int depth)
|
||||
throws IOException {
|
||||
if (depth > MAX_UNZIP_DEPTH) {
|
||||
log.warn(
|
||||
"ZIP nesting depth {} exceeds limit {}, treating as file",
|
||||
depth,
|
||||
MAX_UNZIP_DEPTH);
|
||||
return List.of(zip);
|
||||
}
|
||||
log.debug("Unzipping data of length: {}", zip.contentLength());
|
||||
List<Resource> extracted = new ArrayList<>();
|
||||
try (InputStream bais = zip.getInputStream();
|
||||
ZipInputStream zis = ZipSecurity.createHardenedInputStream(bais)) {
|
||||
ZipEntry entry;
|
||||
while ((entry = zis.getNextEntry()) != null) {
|
||||
if (entry.isDirectory()) {
|
||||
continue;
|
||||
}
|
||||
TempFile tempFile = tempFileManager.createManagedTempFile("unzip");
|
||||
if (tempFileConsumer != null) {
|
||||
tempFileConsumer.accept(tempFile);
|
||||
}
|
||||
try (OutputStream os = Files.newOutputStream(tempFile.getPath())) {
|
||||
byte[] buffer = new byte[4096];
|
||||
int count;
|
||||
while ((count = zis.read(buffer)) != -1) {
|
||||
os.write(buffer, 0, count);
|
||||
}
|
||||
}
|
||||
final String filename = entry.getName();
|
||||
Resource fileResource =
|
||||
new FileSystemResource(tempFile.getFile()) {
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return filename;
|
||||
}
|
||||
};
|
||||
if (isZip(fileResource, filename)) {
|
||||
log.debug("Nested ZIP entry {} — recursing", filename);
|
||||
extracted.addAll(
|
||||
extractZipInternal(
|
||||
fileResource, tempFileManager, tempFileConsumer, depth + 1));
|
||||
} else {
|
||||
extracted.add(fileResource);
|
||||
}
|
||||
}
|
||||
}
|
||||
log.debug("Unzipping completed. {} files extracted.", extracted.size());
|
||||
return extracted;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package stirling.software.common.service;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
import static org.mockito.ArgumentMatchers.*;
|
||||
import static org.mockito.Mockito.*;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
import org.springframework.core.io.ByteArrayResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.http.*;
|
||||
import org.springframework.http.client.ClientHttpResponse;
|
||||
import org.springframework.mock.env.MockEnvironment;
|
||||
import org.springframework.util.LinkedMultiValueMap;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
import org.springframework.web.client.RequestCallback;
|
||||
import org.springframework.web.client.ResponseExtractor;
|
||||
import org.springframework.web.client.RestTemplate;
|
||||
|
||||
import jakarta.servlet.ServletContext;
|
||||
|
||||
import stirling.software.common.util.TempFile;
|
||||
import stirling.software.common.util.TempFileManager;
|
||||
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
class InternalApiClientTest {
|
||||
|
||||
@Mock ServletContext servletContext;
|
||||
@Mock UserServiceInterface userService;
|
||||
@Mock TempFileManager tempFileManager;
|
||||
|
||||
InternalApiClient client;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
lenient().when(servletContext.getContextPath()).thenReturn("");
|
||||
MockEnvironment environment = new MockEnvironment().withProperty("server.port", "8080");
|
||||
client = new InternalApiClient(servletContext, userService, tempFileManager, environment);
|
||||
}
|
||||
|
||||
@Test
|
||||
void postDoesNotForceContentType() throws Exception {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
body.add("fileInput", namedResource("input.pdf", "data"));
|
||||
|
||||
Path tempPath = Files.createTempFile("internal-api-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)));
|
||||
})) {
|
||||
|
||||
ResponseEntity<Resource> response = client.post("/api/v1/general/merge-pdfs", body);
|
||||
|
||||
assertNotNull(response);
|
||||
assertEquals(HttpStatus.OK, response.getStatusCode());
|
||||
assertNotNull(response.getBody());
|
||||
assertNull(captured[0].getContentType(), "Content-Type should not be forced");
|
||||
} finally {
|
||||
Files.deleteIfExists(tempPath);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void postRejectsDisallowedPath() {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
assertThrows(SecurityException.class, () -> client.post("/api/v1/admin/settings", body));
|
||||
}
|
||||
|
||||
@Test
|
||||
void postRejectsPathTraversal() {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
assertThrows(
|
||||
SecurityException.class,
|
||||
() -> client.post("/api/v1/misc/../../actuator/env", body));
|
||||
}
|
||||
|
||||
@Test
|
||||
void postRejectsUrlEncodedCharacters() {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
assertThrows(
|
||||
SecurityException.class, () -> client.post("/api/v1/misc/%2e%2e/actuator", body));
|
||||
}
|
||||
|
||||
@Test
|
||||
void postRejectsQueryString() {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
assertThrows(
|
||||
SecurityException.class,
|
||||
() -> client.post("/api/v1/misc/compress-pdf?redirect=evil", body));
|
||||
}
|
||||
|
||||
@Test
|
||||
void postRejectsEmptySegment() {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
assertThrows(SecurityException.class, () -> client.post("/api/v1/misc//foo", body));
|
||||
}
|
||||
|
||||
@Test
|
||||
void postRejectsTrailingSlash() {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
assertThrows(SecurityException.class, () -> client.post("/api/v1/misc/foo/", body));
|
||||
}
|
||||
|
||||
@Test
|
||||
void postRejectsNullPath() {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
assertThrows(SecurityException.class, () -> client.post(null, body));
|
||||
}
|
||||
|
||||
/** Create a ByteArrayResource with a filename (required for multipart). */
|
||||
private static Resource namedResource(String filename, String content) {
|
||||
return new ByteArrayResource(content.getBytes(StandardCharsets.UTF_8)) {
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return filename;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/** Simulate a successful HTTP response through a RestTemplate ResponseExtractor. */
|
||||
@SuppressWarnings("unchecked")
|
||||
private static ResponseEntity<Resource> fakeOkResponse(Object extractorArg) throws Exception {
|
||||
var extractor = (ResponseExtractor<ResponseEntity<Resource>>) extractorArg;
|
||||
ClientHttpResponse response = mock(ClientHttpResponse.class);
|
||||
when(response.getBody())
|
||||
.thenReturn(new ByteArrayInputStream("ok".getBytes(StandardCharsets.UTF_8)));
|
||||
HttpHeaders headers = new HttpHeaders();
|
||||
headers.add(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"out.pdf\"");
|
||||
when(response.getHeaders()).thenReturn(headers);
|
||||
lenient().when(response.getStatusCode()).thenReturn(HttpStatus.OK);
|
||||
return extractor.extractData(response);
|
||||
}
|
||||
}
|
||||
+3
-1
@@ -124,7 +124,9 @@ public class SplitPdfByChaptersController {
|
||||
@MultiFileResponse
|
||||
@Operation(
|
||||
summary = "Split PDFs by Chapters",
|
||||
description = "Splits a PDF into chapters and returns a ZIP file.")
|
||||
description =
|
||||
"Splits a PDF into chapters and returns a ZIP file. Input:PDF Output:ZIP-PDF"
|
||||
+ " Type:SISO")
|
||||
public ResponseEntity<StreamingResponseBody> splitPdf(
|
||||
@ModelAttribute SplitPdfByChaptersRequest request) throws Exception {
|
||||
MultipartFile file = request.getFileInput();
|
||||
|
||||
+19
-200
@@ -11,36 +11,26 @@ import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
import java.util.Map.Entry;
|
||||
import java.util.zip.ZipEntry;
|
||||
import java.util.zip.ZipInputStream;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.core.io.FileSystemResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.http.*;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.util.LinkedMultiValueMap;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
import org.springframework.web.client.RequestCallback;
|
||||
import org.springframework.web.client.RestTemplate;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
import io.github.pixee.security.Filenames;
|
||||
import io.github.pixee.security.ZipSecurity;
|
||||
|
||||
import jakarta.servlet.ServletContext;
|
||||
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.SPDF.SPDFApplication;
|
||||
import stirling.software.SPDF.model.PipelineConfig;
|
||||
import stirling.software.SPDF.model.PipelineOperation;
|
||||
import stirling.software.SPDF.model.PipelineResult;
|
||||
import stirling.software.SPDF.service.ApiDocService;
|
||||
import stirling.software.common.model.enumeration.Role;
|
||||
import stirling.software.common.service.UserServiceInterface;
|
||||
import stirling.software.common.util.TempFile;
|
||||
import stirling.software.common.service.InternalApiClient;
|
||||
import stirling.software.common.util.TempFileManager;
|
||||
import stirling.software.common.util.ZipExtractionUtils;
|
||||
|
||||
@Service
|
||||
@Slf4j
|
||||
@@ -48,20 +38,16 @@ public class PipelineProcessor {
|
||||
|
||||
private final ApiDocService apiDocService;
|
||||
|
||||
private final UserServiceInterface userService;
|
||||
|
||||
private final ServletContext servletContext;
|
||||
private final InternalApiClient internalApiClient;
|
||||
|
||||
private final TempFileManager tempFileManager;
|
||||
|
||||
public PipelineProcessor(
|
||||
ApiDocService apiDocService,
|
||||
@Autowired(required = false) UserServiceInterface userService,
|
||||
ServletContext servletContext,
|
||||
InternalApiClient internalApiClient,
|
||||
TempFileManager tempFileManager) {
|
||||
this.apiDocService = apiDocService;
|
||||
this.userService = userService;
|
||||
this.servletContext = servletContext;
|
||||
this.internalApiClient = internalApiClient;
|
||||
this.tempFileManager = tempFileManager;
|
||||
}
|
||||
|
||||
@@ -84,48 +70,6 @@ public class PipelineProcessor {
|
||||
return name.substring(0, underscoreIndex) + extension;
|
||||
}
|
||||
|
||||
// Allowlist of URL path prefixes permitted through the pipeline.
|
||||
private static final List<String> ALLOWED_PIPELINE_PATH_PREFIXES =
|
||||
List.of(
|
||||
"/api/v1/general/",
|
||||
"/api/v1/misc/",
|
||||
"/api/v1/security/",
|
||||
"/api/v1/convert/",
|
||||
"/api/v1/filter/");
|
||||
|
||||
private void validatePipelineUrl(String url) {
|
||||
// Strip scheme+host to get the path portion for comparison
|
||||
String path = url;
|
||||
int schemeEnd = url.indexOf("://");
|
||||
if (schemeEnd != -1) {
|
||||
int pathStart = url.indexOf('/', schemeEnd + 3);
|
||||
path = pathStart != -1 ? url.substring(pathStart) : "/";
|
||||
}
|
||||
final String pathToCheck = path;
|
||||
boolean allowed = ALLOWED_PIPELINE_PATH_PREFIXES.stream().anyMatch(pathToCheck::contains);
|
||||
if (!allowed) {
|
||||
log.warn("Blocked pipeline request to disallowed URL: {}", url);
|
||||
throw new SecurityException(
|
||||
"Pipeline operation not permitted for endpoint: " + pathToCheck);
|
||||
}
|
||||
}
|
||||
|
||||
private String getApiKeyForUser() {
|
||||
if (userService == null) return "";
|
||||
String username = userService.getCurrentUsername();
|
||||
if (username != null && !username.equals("anonymousUser")) {
|
||||
return userService.getApiKeyForUser(username);
|
||||
}
|
||||
// Scheduled/internal context — no user in security context
|
||||
return userService.getApiKeyForUser(Role.INTERNAL_API_USER.getRoleId());
|
||||
}
|
||||
|
||||
private String getBaseUrl() {
|
||||
String contextPath = servletContext.getContextPath();
|
||||
String port = SPDFApplication.getStaticPort();
|
||||
return "http://localhost:" + port + contextPath + "/";
|
||||
}
|
||||
|
||||
PipelineResult runPipelineAgainstFiles(List<Resource> outputFiles, PipelineConfig config)
|
||||
throws Exception {
|
||||
PipelineResult result = new PipelineResult();
|
||||
@@ -153,7 +97,6 @@ public class PipelineProcessor {
|
||||
"Invalid operation: " + operation + " with parameters: " + parameters);
|
||||
}
|
||||
|
||||
String url = getBaseUrl() + operation;
|
||||
List<Resource> newOutputFiles = new ArrayList<>();
|
||||
if (!isMultiInputOperation) {
|
||||
for (Resource file : outputFiles) {
|
||||
@@ -175,12 +118,15 @@ public class PipelineProcessor {
|
||||
body.add(entry.getKey(), entry.getValue());
|
||||
}
|
||||
}
|
||||
ResponseEntity<Resource> response = sendWebRequest(url, body);
|
||||
ResponseEntity<Resource> response =
|
||||
internalApiClient.post(operation, body);
|
||||
// If the operation is filter and the response body is null or empty,
|
||||
// skip
|
||||
// this
|
||||
// file
|
||||
if (response.getBody() instanceof TempFileResource tempFileResource) {
|
||||
if (response.getBody()
|
||||
instanceof
|
||||
InternalApiClient.TempFileResource tempFileResource) {
|
||||
result.addTempFile(tempFileResource.getTempFile());
|
||||
}
|
||||
|
||||
@@ -257,8 +203,9 @@ public class PipelineProcessor {
|
||||
body.add(entry.getKey(), entry.getValue());
|
||||
}
|
||||
}
|
||||
ResponseEntity<Resource> response = sendWebRequest(url, body);
|
||||
if (response.getBody() instanceof TempFileResource tempFileResource) {
|
||||
ResponseEntity<Resource> response = internalApiClient.post(operation, body);
|
||||
if (response.getBody()
|
||||
instanceof InternalApiClient.TempFileResource tempFileResource) {
|
||||
result.addTempFile(tempFileResource.getTempFile());
|
||||
}
|
||||
// Handle the response
|
||||
@@ -312,43 +259,6 @@ public class PipelineProcessor {
|
||||
return result;
|
||||
}
|
||||
|
||||
/* package */ ResponseEntity<Resource> sendWebRequest(
|
||||
String url, MultiValueMap<String, Object> body) {
|
||||
validatePipelineUrl(url);
|
||||
RestTemplate restTemplate = new RestTemplate();
|
||||
// Set up headers, including API key
|
||||
HttpHeaders headers = new HttpHeaders();
|
||||
String apiKey = getApiKeyForUser();
|
||||
if (apiKey != null && !apiKey.isEmpty()) {
|
||||
headers.add("X-API-KEY", apiKey);
|
||||
}
|
||||
|
||||
// Let the message converter set the multipart boundary/content type
|
||||
HttpEntity<MultiValueMap<String, Object>> entity = new HttpEntity<>(body, headers);
|
||||
|
||||
RequestCallback requestCallback =
|
||||
restTemplate.httpEntityCallback(entity, Resource.class /* response type hint */);
|
||||
return restTemplate.execute(
|
||||
url,
|
||||
HttpMethod.POST,
|
||||
requestCallback,
|
||||
response -> {
|
||||
try {
|
||||
TempFile tempFile = tempFileManager.createManagedTempFile("pipeline");
|
||||
Files.copy(
|
||||
response.getBody(),
|
||||
tempFile.getPath(),
|
||||
java.nio.file.StandardCopyOption.REPLACE_EXISTING);
|
||||
TempFileResource resource = new TempFileResource(tempFile);
|
||||
return ResponseEntity.status(response.getStatusCode())
|
||||
.headers(response.getHeaders())
|
||||
.body(resource);
|
||||
} catch (IOException e) {
|
||||
throw new UncheckedIOException(e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private List<Resource> processOutputFiles(
|
||||
String operation,
|
||||
ResponseEntity<Resource> response,
|
||||
@@ -367,13 +277,15 @@ public class PipelineProcessor {
|
||||
newFilename = removeTrailingNaming(extractFilename(response));
|
||||
}
|
||||
// Check if the response body is a zip file
|
||||
if (isZip(response.getBody(), newFilename)) {
|
||||
if (ZipExtractionUtils.isZip(response.getBody(), newFilename)) {
|
||||
// Unzip the file and add all the files to the new output files
|
||||
newOutputFiles.addAll(unzip(response.getBody(), result));
|
||||
newOutputFiles.addAll(
|
||||
ZipExtractionUtils.extractZip(
|
||||
response.getBody(), tempFileManager, result::addTempFile));
|
||||
} else {
|
||||
final Resource tempResource = response.getBody();
|
||||
if (tempResource instanceof TempFileResource) {
|
||||
result.addTempFile(((TempFileResource) tempResource).getTempFile());
|
||||
if (tempResource instanceof InternalApiClient.TempFileResource tfr) {
|
||||
result.addTempFile(tfr.getTempFile());
|
||||
}
|
||||
Resource outputResource =
|
||||
new FileSystemResource(tempResource.getFile()) {
|
||||
@@ -456,97 +368,4 @@ public class PipelineProcessor {
|
||||
log.info("Files successfully loaded. Starting processing...");
|
||||
return outputFiles;
|
||||
}
|
||||
|
||||
private boolean isZip(Resource data, String filename) throws IOException {
|
||||
if (data == null || data.contentLength() < 4) {
|
||||
return false;
|
||||
}
|
||||
if (filename != null) {
|
||||
String lower = filename.toLowerCase();
|
||||
if (lower.endsWith(".cbz")) {
|
||||
// Treat CBZ as non-zip for our unzipping purposes
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// Check the first four bytes of the data against the standard zip magic number
|
||||
try (InputStream is = data.getInputStream()) {
|
||||
byte[] header = new byte[4];
|
||||
if (is.read(header) < 4) {
|
||||
return false;
|
||||
}
|
||||
return header[0] == 0x50 && header[1] == 0x4B && header[2] == 0x03 && header[3] == 0x04;
|
||||
}
|
||||
}
|
||||
|
||||
private boolean isZip(Resource data) throws IOException {
|
||||
return isZip(data, null);
|
||||
}
|
||||
|
||||
private static final int MAX_UNZIP_DEPTH = 10;
|
||||
|
||||
private List<Resource> unzip(Resource data, PipelineResult result) throws IOException {
|
||||
return unzip(data, result, 0);
|
||||
}
|
||||
|
||||
private List<Resource> unzip(Resource data, PipelineResult result, int depth)
|
||||
throws IOException {
|
||||
if (depth > MAX_UNZIP_DEPTH) {
|
||||
log.warn(
|
||||
"ZIP nesting depth {} exceeds limit {}, treating as file",
|
||||
depth,
|
||||
MAX_UNZIP_DEPTH);
|
||||
return List.of(data);
|
||||
}
|
||||
log.info("Unzipping data of length: {}", data.contentLength());
|
||||
List<Resource> unzippedFiles = new ArrayList<>();
|
||||
try (InputStream bais = data.getInputStream();
|
||||
ZipInputStream zis = ZipSecurity.createHardenedInputStream(bais)) {
|
||||
ZipEntry entry;
|
||||
while ((entry = zis.getNextEntry()) != null) {
|
||||
if (entry.isDirectory()) {
|
||||
continue;
|
||||
}
|
||||
TempFile tempFile = tempFileManager.createManagedTempFile("unzip");
|
||||
result.addTempFile(tempFile);
|
||||
try (OutputStream os = Files.newOutputStream(tempFile.getPath())) {
|
||||
byte[] buffer = new byte[4096];
|
||||
int count;
|
||||
while ((count = zis.read(buffer)) != -1) {
|
||||
os.write(buffer, 0, count);
|
||||
}
|
||||
}
|
||||
final String filename = entry.getName();
|
||||
Resource fileResource =
|
||||
new FileSystemResource(tempFile.getFile()) {
|
||||
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return filename;
|
||||
}
|
||||
};
|
||||
// If the unzipped file is a zip file, unzip it
|
||||
if (isZip(fileResource, filename)) {
|
||||
log.info("File {} is a zip file. Unzipping...", filename);
|
||||
unzippedFiles.addAll(unzip(fileResource, result, depth + 1));
|
||||
} else {
|
||||
unzippedFiles.add(fileResource);
|
||||
}
|
||||
}
|
||||
}
|
||||
log.info("Unzipping completed. {} files were unzipped.", unzippedFiles.size());
|
||||
return unzippedFiles;
|
||||
}
|
||||
|
||||
private static class TempFileResource extends FileSystemResource {
|
||||
private final TempFile tempFile;
|
||||
|
||||
public TempFileResource(TempFile tempFile) {
|
||||
super(tempFile.getFile());
|
||||
this.tempFile = tempFile;
|
||||
}
|
||||
|
||||
public TempFile getTempFile() {
|
||||
return tempFile;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+2
-1
@@ -13,7 +13,8 @@ public class RotatePDFRequest extends PDFFile {
|
||||
|
||||
@Schema(
|
||||
description =
|
||||
"The angle by which to rotate the PDF file. This should be a multiple of 90.",
|
||||
"The clockwise angle by which to rotate the PDF file. Must be a multiple of"
|
||||
+ " 90.",
|
||||
type = "integer",
|
||||
requiredMode = Schema.RequiredMode.REQUIRED,
|
||||
allowableValues = {"0", "90", "180", "270"})
|
||||
|
||||
@@ -6,6 +6,7 @@ import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
import java.util.regex.Matcher;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.http.HttpEntity;
|
||||
@@ -30,7 +31,14 @@ import tools.jackson.databind.ObjectMapper;
|
||||
|
||||
@Service
|
||||
@Slf4j
|
||||
public class ApiDocService {
|
||||
public class ApiDocService implements stirling.software.common.service.ToolMetadataService {
|
||||
|
||||
// Matches a bare "Output:ZIP" declaration (i.e. ZIP is not followed by "-" or "/").
|
||||
// Bare ZIP means the archive itself is the deliverable (e.g. get-attachments), so it
|
||||
// should not be auto-unpacked. Wrapper forms like Output:ZIP-PDF or Output:IMAGE/ZIP
|
||||
// use ZIP as transport for multiple typed results and are safe to unpack.
|
||||
private static final Pattern BARE_ZIP_OUTPUT =
|
||||
Pattern.compile("Output\\s*:\\s*ZIP(?![-/])", Pattern.CASE_INSENSITIVE);
|
||||
|
||||
private final Map<String, ApiEndpoint> apiDocumentation = new HashMap<>();
|
||||
|
||||
@@ -149,6 +157,7 @@ public class ApiDocService {
|
||||
return endpoint.areParametersValid(parameters);
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean isMultiInput(String operationName) {
|
||||
if (apiDocsJsonRootNode == null || apiDocumentation.isEmpty()) {
|
||||
loadApiDocumentation();
|
||||
@@ -166,5 +175,36 @@ public class ApiDocService {
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean shouldUnpackZipResponse(String operationName) {
|
||||
if (apiDocsJsonRootNode == null || apiDocumentation.isEmpty()) {
|
||||
loadApiDocumentation();
|
||||
}
|
||||
if (!apiDocumentation.containsKey(operationName)) {
|
||||
return false;
|
||||
}
|
||||
ApiEndpoint endpoint = apiDocumentation.get(operationName);
|
||||
String description = endpoint.getDescription();
|
||||
Matcher typeMatcher =
|
||||
RegexPatternUtils.getInstance().getApiDocTypePattern().matcher(description);
|
||||
if (typeMatcher.find()) {
|
||||
String type = typeMatcher.group(1);
|
||||
// Multi-output endpoints (SIMO/MIMO) return a ZIP of their outputs.
|
||||
if (type.endsWith("MO")) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
Matcher outputMatcher =
|
||||
RegexPatternUtils.getInstance().getApiDocOutputTypePattern().matcher(description);
|
||||
if (outputMatcher.find()) {
|
||||
String output = outputMatcher.group(1).toUpperCase(Locale.ROOT);
|
||||
if (output.startsWith("ZIP")) {
|
||||
// Bare "Output:ZIP" is a single-archive deliverable, not a transport.
|
||||
return !BARE_ZIP_OUTPUT.matcher(description).find();
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// Model class for API Endpoint
|
||||
|
||||
+8
-110
@@ -4,8 +4,6 @@ import static org.junit.jupiter.api.Assertions.*;
|
||||
import static org.mockito.ArgumentMatchers.*;
|
||||
import static org.mockito.Mockito.*;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.List;
|
||||
@@ -15,24 +13,18 @@ import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.MockedConstruction;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
import org.springframework.core.io.ByteArrayResource;
|
||||
import org.springframework.core.io.FileSystemResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.http.client.ClientHttpResponse;
|
||||
import org.springframework.util.LinkedMultiValueMap;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
|
||||
import jakarta.servlet.ServletContext;
|
||||
|
||||
import stirling.software.SPDF.model.PipelineConfig;
|
||||
import stirling.software.SPDF.model.PipelineOperation;
|
||||
import stirling.software.SPDF.model.PipelineResult;
|
||||
import stirling.software.SPDF.service.ApiDocService;
|
||||
import stirling.software.common.service.UserServiceInterface;
|
||||
import stirling.software.common.service.InternalApiClient;
|
||||
import stirling.software.common.util.TempFileManager;
|
||||
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
@@ -40,9 +32,7 @@ class PipelineProcessorTest {
|
||||
|
||||
@Mock ApiDocService apiDocService;
|
||||
|
||||
@Mock UserServiceInterface userService;
|
||||
|
||||
@Mock ServletContext servletContext;
|
||||
@Mock InternalApiClient internalApiClient;
|
||||
|
||||
@Mock TempFileManager tempFileManager;
|
||||
|
||||
@@ -51,9 +41,7 @@ class PipelineProcessorTest {
|
||||
@BeforeEach
|
||||
void setUp() throws Exception {
|
||||
pipelineProcessor =
|
||||
spy(
|
||||
new PipelineProcessor(
|
||||
apiDocService, userService, servletContext, tempFileManager));
|
||||
new PipelineProcessor(apiDocService, internalApiClient, tempFileManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
@@ -65,7 +53,6 @@ class PipelineProcessorTest {
|
||||
config.setOperations(List.of(op));
|
||||
|
||||
Resource file = new MyFileByteArrayResource();
|
||||
|
||||
List<Resource> files = List.of(file);
|
||||
|
||||
when(apiDocService.isMultiInput("/api/v1/filter/filter-page-count")).thenReturn(false);
|
||||
@@ -74,13 +61,11 @@ class PipelineProcessorTest {
|
||||
when(apiDocService.isValidOperation(eq("/api/v1/filter/filter-page-count"), anyMap()))
|
||||
.thenReturn(true);
|
||||
|
||||
// Use a FileSystemResource backed by a temp file to avoid FileNotFoundException
|
||||
Path emptyTemp = Files.createTempFile("empty", ".tmp");
|
||||
Resource emptyResource = new FileSystemResource(emptyTemp.toFile());
|
||||
|
||||
doReturn(new ResponseEntity<>(emptyResource, HttpStatus.OK))
|
||||
.when(pipelineProcessor)
|
||||
.sendWebRequest(anyString(), any());
|
||||
when(internalApiClient.post(anyString(), any()))
|
||||
.thenReturn(new ResponseEntity<>(emptyResource, HttpStatus.OK));
|
||||
|
||||
PipelineResult result = pipelineProcessor.runPipelineAgainstFiles(files, config);
|
||||
|
||||
@@ -118,105 +103,18 @@ class PipelineProcessorTest {
|
||||
when(apiDocService.getExtensionTypes(anyBoolean(), anyString())).thenReturn(List.of("pdf"));
|
||||
when(apiDocService.isValidOperation(anyString(), anyMap())).thenReturn(true);
|
||||
|
||||
doReturn(new ResponseEntity<>(outputResource, HttpStatus.OK))
|
||||
.when(pipelineProcessor)
|
||||
.sendWebRequest(anyString(), any());
|
||||
when(internalApiClient.post(anyString(), any()))
|
||||
.thenReturn(new ResponseEntity<>(outputResource, HttpStatus.OK));
|
||||
|
||||
PipelineResult result = pipelineProcessor.runPipelineAgainstFiles(files, config);
|
||||
|
||||
verify(pipelineProcessor).sendWebRequest(anyString(), any());
|
||||
verify(internalApiClient).post(anyString(), any());
|
||||
|
||||
assertFalse(result.isHasErrors());
|
||||
|
||||
// Clean up
|
||||
Files.deleteIfExists(tempPath);
|
||||
}
|
||||
|
||||
@Test
|
||||
void sendWebRequestDoesNotForceContentType() throws Exception {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
body.add(
|
||||
"fileInput",
|
||||
new ByteArrayResource("data".getBytes(StandardCharsets.UTF_8)) {
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return "input.pdf";
|
||||
}
|
||||
});
|
||||
|
||||
Path tempPath = Files.createTempFile("pipeline-test", ".tmp");
|
||||
var tempFile = mock(stirling.software.common.util.TempFile.class);
|
||||
when(tempFile.getPath()).thenReturn(tempPath);
|
||||
when(tempFile.getFile()).thenReturn(tempPath.toFile());
|
||||
when(tempFileManager.createManagedTempFile("pipeline")).thenReturn(tempFile);
|
||||
|
||||
var capturedHeaders = new org.springframework.http.HttpHeaders[1];
|
||||
|
||||
try (MockedConstruction<org.springframework.web.client.RestTemplate> ignored =
|
||||
mockConstruction(
|
||||
org.springframework.web.client.RestTemplate.class,
|
||||
(mock, context) -> {
|
||||
when(mock.httpEntityCallback(any(), eq(Resource.class)))
|
||||
.thenAnswer(
|
||||
invocation -> {
|
||||
var entity = invocation.getArgument(0);
|
||||
capturedHeaders[0] =
|
||||
((org.springframework.http.HttpEntity<?>)
|
||||
entity)
|
||||
.getHeaders();
|
||||
return (org.springframework.web.client
|
||||
.RequestCallback)
|
||||
request -> {};
|
||||
});
|
||||
|
||||
when(mock.execute(
|
||||
anyString(),
|
||||
eq(org.springframework.http.HttpMethod.POST),
|
||||
any(),
|
||||
any()))
|
||||
.thenAnswer(
|
||||
invocation -> {
|
||||
@SuppressWarnings("unchecked")
|
||||
var extractor =
|
||||
(org.springframework.web.client
|
||||
.ResponseExtractor<
|
||||
ResponseEntity<Resource>>)
|
||||
invocation.getArgument(3);
|
||||
ClientHttpResponse response =
|
||||
mock(ClientHttpResponse.class);
|
||||
when(response.getBody())
|
||||
.thenReturn(
|
||||
new ByteArrayInputStream(
|
||||
"ok"
|
||||
.getBytes(
|
||||
StandardCharsets
|
||||
.UTF_8)));
|
||||
var headers =
|
||||
new org.springframework.http.HttpHeaders();
|
||||
headers.add(
|
||||
org.springframework.http.HttpHeaders
|
||||
.CONTENT_DISPOSITION,
|
||||
"attachment; filename=\"out.pdf\"");
|
||||
when(response.getHeaders()).thenReturn(headers);
|
||||
lenient()
|
||||
.when(response.getStatusCode())
|
||||
.thenReturn(HttpStatus.OK);
|
||||
return extractor.extractData(response);
|
||||
});
|
||||
})) {
|
||||
ResponseEntity<Resource> response =
|
||||
pipelineProcessor.sendWebRequest(
|
||||
"http://localhost/api/v1/general/merge-pdfs", body);
|
||||
|
||||
assertNotNull(response);
|
||||
assertEquals(HttpStatus.OK, response.getStatusCode());
|
||||
assertNotNull(response.getBody());
|
||||
assertNull(capturedHeaders[0].getContentType());
|
||||
} finally {
|
||||
Files.deleteIfExists(tempPath);
|
||||
}
|
||||
}
|
||||
|
||||
private static class MyFileByteArrayResource extends ByteArrayResource {
|
||||
public MyFileByteArrayResource() {
|
||||
super("data".getBytes());
|
||||
|
||||
@@ -9,6 +9,8 @@ import java.util.Map;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.CsvSource;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
|
||||
@@ -348,6 +350,123 @@ class ApiDocServiceTest {
|
||||
assertTrue(apiDocService.isMultiInput("/miso"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldUnpackZipResponseDetectsMultiOutputType() throws Exception {
|
||||
String json = "{\"description\": \"Output:PDF Type:SIMO\"}";
|
||||
JsonNode postNode = mapper.readTree(json);
|
||||
ApiEndpoint endpoint = new ApiEndpoint("/split", postNode);
|
||||
setApiDocumentation(Map.of("/split", endpoint));
|
||||
setApiDocsJsonRootNode();
|
||||
assertTrue(apiDocService.shouldUnpackZipResponse("/split"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldUnpackZipResponseDetectsMimoType() throws Exception {
|
||||
String json = "{\"description\": \"Output:PDF Type:MIMO\"}";
|
||||
JsonNode postNode = mapper.readTree(json);
|
||||
ApiEndpoint endpoint = new ApiEndpoint("/overlay", postNode);
|
||||
setApiDocumentation(Map.of("/overlay", endpoint));
|
||||
setApiDocsJsonRootNode();
|
||||
assertTrue(apiDocService.shouldUnpackZipResponse("/overlay"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldUnpackZipResponseDetectsZipOutputDeclaration() throws Exception {
|
||||
String json = "{\"description\": \"Output:ZIP-PDF Type:SISO\"}";
|
||||
JsonNode postNode = mapper.readTree(json);
|
||||
ApiEndpoint endpoint = new ApiEndpoint("/split-by-sections", postNode);
|
||||
setApiDocumentation(Map.of("/split-by-sections", endpoint));
|
||||
setApiDocsJsonRootNode();
|
||||
assertTrue(apiDocService.shouldUnpackZipResponse("/split-by-sections"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldUnpackZipResponseReturnsFalseForSisoPdf() throws Exception {
|
||||
String json = "{\"description\": \"Input:PDF Output:PDF Type:SISO\"}";
|
||||
JsonNode postNode = mapper.readTree(json);
|
||||
ApiEndpoint endpoint = new ApiEndpoint("/rotate", postNode);
|
||||
setApiDocumentation(Map.of("/rotate", endpoint));
|
||||
setApiDocsJsonRootNode();
|
||||
assertFalse(apiDocService.shouldUnpackZipResponse("/rotate"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldUnpackZipResponseReturnsFalseForUnknownOperation() throws Exception {
|
||||
setApiDocumentation(Map.of());
|
||||
assertFalse(apiDocService.shouldUnpackZipResponse("/unknown"));
|
||||
}
|
||||
|
||||
/**
|
||||
* Coverage test: every Stirling endpoint whose ZIP response is a transport for multiple typed
|
||||
* results (SIMO/MIMO or Output:ZIP-PDF / Output:IMAGE/ZIP etc.) must be classified as {@code
|
||||
* shouldUnpackZipResponse = true}. Descriptions below are the real
|
||||
* {@code @Operation(description=...)} strings from each controller, so if a controller is
|
||||
* renamed, tweaked or introduced without a {@code Type:} / {@code Output:ZIP-*} tag, this test
|
||||
* breaks, surfacing the bug before {@code AiWorkflowService} silently registers a multi-result
|
||||
* ZIP as a single file.
|
||||
*
|
||||
* <p>Add a new row here whenever a new unpack-eligible endpoint is introduced. Descriptions can
|
||||
* be trimmed to the part containing the relevant tags.
|
||||
*/
|
||||
@ParameterizedTest(name = "{0} → shouldUnpackZipResponse")
|
||||
@CsvSource(
|
||||
textBlock =
|
||||
"""
|
||||
/api/v1/general/split-pages, 'Split pages. Input:PDF Output:PDF Type:SIMO'
|
||||
/api/v1/general/split-pdf-by-sections, 'Split. Input:PDF Output:ZIP-PDF Type:SISO'
|
||||
/api/v1/general/split-by-size-or-count, 'Split by size. Input:PDF Output:ZIP-PDF Type:SISO'
|
||||
/api/v1/general/split-pdf-by-chapters, 'Split by chapters. Input:PDF Output:ZIP-PDF Type:SISO'
|
||||
/api/v1/general/split-for-poster-print, 'Poster split. Input: PDF Output: ZIP-PDF Type: SISO'
|
||||
/api/v1/general/overlay-pdfs, 'Overlay PDFs. Input:PDF Output:PDF Type:MIMO'
|
||||
/api/v1/misc/auto-split-pdf, 'Auto split. Input:PDF Output:ZIP-PDF Type:SISO'
|
||||
/api/v1/misc/extract-images, 'Extract images. Output:IMAGE/ZIP Type:SIMO'
|
||||
/api/v1/misc/extract-image-scans, 'Extract image scans. Input:PDF Output:IMAGE/ZIP Type:SIMO'
|
||||
""")
|
||||
void shouldUnpackZipResponseClassifiesKnownUnpackableEndpoints(
|
||||
String endpoint, String description) throws Exception {
|
||||
String json = mapper.writeValueAsString(Map.of("description", description));
|
||||
JsonNode postNode = mapper.readTree(json);
|
||||
setApiDocumentation(Map.of(endpoint, new ApiEndpoint(endpoint, postNode)));
|
||||
setApiDocsJsonRootNode();
|
||||
assertTrue(
|
||||
apiDocService.shouldUnpackZipResponse(endpoint),
|
||||
() ->
|
||||
"Expected shouldUnpackZipResponse=true for "
|
||||
+ endpoint
|
||||
+ " with description: "
|
||||
+ description);
|
||||
}
|
||||
|
||||
/**
|
||||
* Inverse coverage: endpoints whose ZIP response is the deliverable itself (or that return
|
||||
* single non-ZIP files) must not be flagged for unpacking. Catches regressions where a change
|
||||
* to the classifier accidentally widens the positive match.
|
||||
*/
|
||||
@ParameterizedTest(name = "{0} → !shouldUnpackZipResponse")
|
||||
@CsvSource(
|
||||
textBlock =
|
||||
"""
|
||||
/api/v1/general/rotate-pdf, 'Rotate. Input:PDF Output:PDF Type:SISO'
|
||||
/api/v1/general/merge-pdfs, 'Merge. Input:PDF Output:PDF Type:MISO'
|
||||
/api/v1/misc/compress-pdf, 'Compress. Input:PDF Output:PDF Type:SISO'
|
||||
/api/v1/misc/flatten, 'Flatten forms. Input:PDF Output:PDF Type:SISO'
|
||||
/api/v1/security/get-attachments, 'Extract attachments. Input:PDF Output:ZIP Type:SISO'
|
||||
""")
|
||||
void shouldUnpackZipResponseRejectsNonUnpackableEndpoints(String endpoint, String description)
|
||||
throws Exception {
|
||||
String json = mapper.writeValueAsString(Map.of("description", description));
|
||||
JsonNode postNode = mapper.readTree(json);
|
||||
setApiDocumentation(Map.of(endpoint, new ApiEndpoint(endpoint, postNode)));
|
||||
setApiDocsJsonRootNode();
|
||||
assertFalse(
|
||||
apiDocService.shouldUnpackZipResponse(endpoint),
|
||||
() ->
|
||||
"Expected shouldUnpackZipResponse=false for "
|
||||
+ endpoint
|
||||
+ " with description: "
|
||||
+ description);
|
||||
}
|
||||
|
||||
@Test
|
||||
void constructorAcceptsNullUserService() {
|
||||
ApiDocService service = new ApiDocService(mapper, servletContext, null);
|
||||
|
||||
@@ -10,6 +10,7 @@ import org.springframework.context.annotation.Configuration;
|
||||
import org.springframework.core.task.TaskDecorator;
|
||||
import org.springframework.core.task.support.TaskExecutorAdapter;
|
||||
import org.springframework.scheduling.annotation.EnableAsync;
|
||||
import org.springframework.security.concurrent.DelegatingSecurityContextExecutor;
|
||||
|
||||
@Configuration
|
||||
@EnableAsync
|
||||
@@ -49,11 +50,18 @@ public class AsyncConfig {
|
||||
return adapter;
|
||||
}
|
||||
|
||||
/**
|
||||
* AI orchestration runs on a background executor, so the incoming request's {@code
|
||||
* SecurityContext} must be propagated for downstream calls to see the authenticated user.
|
||||
* Without this, {@code JobOwnershipService} scopes job keys without a user prefix and
|
||||
* authenticated downloads fail with 403; {@code InternalApiClient} also falls back to the
|
||||
* internal-API-user key instead of the caller's.
|
||||
*/
|
||||
@Bean(name = "aiStreamExecutor")
|
||||
public Executor aiStreamExecutor() {
|
||||
TaskExecutorAdapter adapter =
|
||||
new TaskExecutorAdapter(Executors.newVirtualThreadPerTaskExecutor());
|
||||
adapter.setTaskDecorator(new MDCContextTaskDecorator());
|
||||
return adapter;
|
||||
return new DelegatingSecurityContextExecutor(adapter);
|
||||
}
|
||||
}
|
||||
|
||||
+79
-8
@@ -1,10 +1,12 @@
|
||||
package stirling.software.proprietary.controller.api;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.Executor;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Qualifier;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
@@ -25,8 +27,12 @@ import jakarta.validation.Valid;
|
||||
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.model.job.ResultFile;
|
||||
import stirling.software.common.service.JobOwnershipService;
|
||||
import stirling.software.common.service.TaskManager;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
|
||||
import stirling.software.proprietary.service.AiEngineClient;
|
||||
import stirling.software.proprietary.service.AiWorkflowService;
|
||||
|
||||
@@ -45,16 +51,30 @@ public class AiEngineController {
|
||||
private final AiWorkflowService aiWorkflowService;
|
||||
private final ObjectMapper objectMapper;
|
||||
private final Executor aiStreamExecutor;
|
||||
private final TaskManager taskManager;
|
||||
private final JobOwnershipService jobOwnershipService;
|
||||
|
||||
/**
|
||||
* SSE emitter timeout. Long enough to accommodate multi-gigabyte PDF workflows (OCR on a
|
||||
* 1000-page scan, splitting a huge PDF, etc.) without the emitter completing out from under the
|
||||
* executor. Configurable via {@code stirling.ai.streamTimeoutMs}.
|
||||
*/
|
||||
@Value("${stirling.ai.streamTimeoutMs:1800000}")
|
||||
private long streamTimeoutMs;
|
||||
|
||||
public AiEngineController(
|
||||
AiEngineClient aiEngineClient,
|
||||
AiWorkflowService aiWorkflowService,
|
||||
ObjectMapper objectMapper,
|
||||
@Qualifier("aiStreamExecutor") Executor aiStreamExecutor) {
|
||||
@Qualifier("aiStreamExecutor") Executor aiStreamExecutor,
|
||||
TaskManager taskManager,
|
||||
JobOwnershipService jobOwnershipService) {
|
||||
this.aiEngineClient = aiEngineClient;
|
||||
this.aiWorkflowService = aiWorkflowService;
|
||||
this.objectMapper = objectMapper;
|
||||
this.aiStreamExecutor = aiStreamExecutor;
|
||||
this.taskManager = taskManager;
|
||||
this.jobOwnershipService = jobOwnershipService;
|
||||
}
|
||||
|
||||
@GetMapping("/health")
|
||||
@@ -70,10 +90,14 @@ public class AiEngineController {
|
||||
@Operation(
|
||||
summary = "Run an AI workflow against a PDF",
|
||||
description =
|
||||
"Accepts a PDF upload and a user message and returns an AI workflow result")
|
||||
public ResponseEntity<AiWorkflowResponse> orchestrate(
|
||||
@Valid @ModelAttribute AiWorkflowRequest request) throws IOException {
|
||||
return ResponseEntity.ok(aiWorkflowService.orchestrate(request));
|
||||
"Accepts PDF uploads and a user message and returns an AI workflow result."
|
||||
+ " When the workflow produces files, they are registered with the job"
|
||||
+ " system and downloadable via GET /api/v1/general/files/{fileId}.")
|
||||
public AiWorkflowResponse orchestrate(@Valid @ModelAttribute AiWorkflowRequest request)
|
||||
throws IOException {
|
||||
AiWorkflowResponse result = aiWorkflowService.orchestrate(request);
|
||||
registerFileResultAsJob(result);
|
||||
return result;
|
||||
}
|
||||
|
||||
@PostMapping(value = "/orchestrate/stream", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
|
||||
@@ -83,11 +107,23 @@ public class AiEngineController {
|
||||
"Accepts a PDF upload and a user message, returns SSE events with progress"
|
||||
+ " updates followed by the final AI workflow result")
|
||||
public SseEmitter orchestrateStream(@Valid @ModelAttribute AiWorkflowRequest request) {
|
||||
SseEmitter emitter = new SseEmitter(180_000L);
|
||||
SseEmitter emitter = new SseEmitter(streamTimeoutMs);
|
||||
|
||||
emitter.onTimeout(
|
||||
() -> {
|
||||
log.warn("SSE emitter timed out for AI orchestration stream");
|
||||
// Emit an explicit error frame so the frontend reports a timeout rather than
|
||||
// silently seeing the stream end without a result.
|
||||
log.warn(
|
||||
"SSE emitter timed out for AI orchestration stream after {} ms",
|
||||
streamTimeoutMs);
|
||||
sendEvent(
|
||||
emitter,
|
||||
"error",
|
||||
Map.of(
|
||||
"message",
|
||||
"AI workflow timed out after "
|
||||
+ (streamTimeoutMs / 1000)
|
||||
+ " seconds"));
|
||||
emitter.complete();
|
||||
});
|
||||
emitter.onError(e -> log.warn("SSE emitter error for AI orchestration stream", e));
|
||||
@@ -102,15 +138,50 @@ public class AiEngineController {
|
||||
AiWorkflowResponse result =
|
||||
aiWorkflowService.orchestrate(
|
||||
request, progress -> sendEvent(emitter, "progress", progress));
|
||||
registerFileResultAsJob(result);
|
||||
sendEvent(emitter, "result", result);
|
||||
emitter.complete();
|
||||
} catch (Exception e) {
|
||||
log.error("AI orchestration stream failed", e);
|
||||
// Emit an error frame for the frontend and then complete normally. Using
|
||||
// completeWithError here as well would double-complete the emitter - the error
|
||||
// frame already conveys the failure to the client.
|
||||
sendEvent(emitter, "error", Map.of("message", e.getMessage()));
|
||||
emitter.completeWithError(e);
|
||||
emitter.complete();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Register any file results produced by the workflow with {@link TaskManager} so they are
|
||||
* downloadable via {@code GET /api/v1/general/files/{fileId}}. Uses {@code
|
||||
* setMultipleFileResults} so the fileIds we registered earlier are not mangled by TaskManager's
|
||||
* ZIP auto-extract path.
|
||||
*/
|
||||
private void registerFileResultAsJob(AiWorkflowResponse result) {
|
||||
List<AiWorkflowResultFile> files = result.getResultFiles();
|
||||
if (files == null || files.isEmpty()) {
|
||||
return;
|
||||
}
|
||||
// Scope the job key to the current user so the download endpoint's ownership check
|
||||
// passes when security is enabled. NoOpJobOwnershipService returns the UUID unchanged
|
||||
// when security is off.
|
||||
String jobKey =
|
||||
jobOwnershipService.createScopedJobKey(java.util.UUID.randomUUID().toString());
|
||||
taskManager.createTask(jobKey);
|
||||
List<ResultFile> jobFiles =
|
||||
files.stream()
|
||||
.map(
|
||||
f ->
|
||||
ResultFile.builder()
|
||||
.fileId(f.getFileId())
|
||||
.fileName(f.getFileName())
|
||||
.contentType(f.getContentType())
|
||||
.build())
|
||||
.toList();
|
||||
taskManager.setMultipleFileResults(jobKey, jobFiles);
|
||||
taskManager.setComplete(jobKey);
|
||||
}
|
||||
|
||||
private void sendEvent(SseEmitter emitter, String name, Object data) {
|
||||
try {
|
||||
emitter.send(SseEmitter.event().name(name).data(data, MediaType.APPLICATION_JSON));
|
||||
|
||||
+1
@@ -15,6 +15,7 @@ public enum AiWorkflowOutcome {
|
||||
PLAN("plan"),
|
||||
NEED_CLARIFICATION("need_clarification"),
|
||||
CANNOT_DO("cannot_do"),
|
||||
DRAFT("draft"),
|
||||
TOOL_CALL("tool_call"),
|
||||
COMPLETED("completed"),
|
||||
UNSUPPORTED_CAPABILITY("unsupported_capability"),
|
||||
|
||||
+1
@@ -8,6 +8,7 @@ public enum AiWorkflowPhase {
|
||||
ANALYZING("analyzing"),
|
||||
CALLING_ENGINE("calling_engine"),
|
||||
EXTRACTING_CONTENT("extracting_content"),
|
||||
EXECUTING_TOOL("executing_tool"),
|
||||
PROCESSING("processing");
|
||||
|
||||
private final String value;
|
||||
|
||||
+24
-1
@@ -1,15 +1,38 @@
|
||||
package stirling.software.proprietary.model.api.ai;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonInclude;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Data;
|
||||
|
||||
@Data
|
||||
@AllArgsConstructor
|
||||
@JsonInclude(JsonInclude.Include.NON_NULL)
|
||||
public class AiWorkflowProgressEvent {
|
||||
private AiWorkflowPhase phase;
|
||||
private long timestamp;
|
||||
|
||||
/** The tool endpoint path being executed, for {@link AiWorkflowPhase#EXECUTING_TOOL} events. */
|
||||
private String tool;
|
||||
|
||||
/**
|
||||
* 1-based index of the current plan step, for {@link AiWorkflowPhase#EXECUTING_TOOL} events.
|
||||
*/
|
||||
private Integer stepIndex;
|
||||
|
||||
/** Total number of plan steps, for {@link AiWorkflowPhase#EXECUTING_TOOL} events. */
|
||||
private Integer stepCount;
|
||||
|
||||
public static AiWorkflowProgressEvent of(AiWorkflowPhase phase) {
|
||||
return new AiWorkflowProgressEvent(phase, System.currentTimeMillis());
|
||||
return new AiWorkflowProgressEvent(phase, System.currentTimeMillis(), null, null, null);
|
||||
}
|
||||
|
||||
public static AiWorkflowProgressEvent executingTool(String tool, int stepIndex, int stepCount) {
|
||||
return new AiWorkflowProgressEvent(
|
||||
AiWorkflowPhase.EXECUTING_TOOL,
|
||||
System.currentTimeMillis(),
|
||||
tool,
|
||||
stepIndex,
|
||||
stepCount);
|
||||
}
|
||||
}
|
||||
|
||||
+24
@@ -44,6 +44,30 @@ public class AiWorkflowResponse {
|
||||
@Schema(description = "Structured tool steps when the workflow returns a plan")
|
||||
private List<Map<String, Object>> steps = new ArrayList<>();
|
||||
|
||||
@Schema(
|
||||
description =
|
||||
"Tool endpoint path for tool_call outcomes (e.g. /api/v1/misc/compress-pdf)")
|
||||
private String tool;
|
||||
|
||||
@Schema(description = "Tool parameters for tool_call outcomes")
|
||||
private Map<String, Object> parameters;
|
||||
|
||||
@Schema(description = "Result file ID after tool execution completes (single-file result)")
|
||||
private String fileId;
|
||||
|
||||
@Schema(description = "Result filename after tool execution completes (single-file result)")
|
||||
private String fileName;
|
||||
|
||||
@Schema(description = "Result MIME type after tool execution completes (single-file result)")
|
||||
private String contentType;
|
||||
|
||||
@Schema(
|
||||
description =
|
||||
"Result files produced by the workflow. Always populated on completed outcomes"
|
||||
+ " with at least one entry; for single-file results this mirrors"
|
||||
+ " fileId/fileName/contentType.")
|
||||
private List<AiWorkflowResultFile> resultFiles = new ArrayList<>();
|
||||
|
||||
@Schema(description = "Per-file text extraction requests from the AI engine")
|
||||
private List<AiWorkflowFileRequest> files = new ArrayList<>();
|
||||
|
||||
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
package stirling.software.proprietary.model.api.ai;
|
||||
|
||||
import io.swagger.v3.oas.annotations.media.Schema;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Data;
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
/** A single file produced by a completed AI workflow. */
|
||||
@Data
|
||||
@NoArgsConstructor
|
||||
@AllArgsConstructor
|
||||
@Schema(description = "Descriptor for a file produced by an AI workflow")
|
||||
public class AiWorkflowResultFile {
|
||||
|
||||
@Schema(description = "Stirling file ID — download with GET /api/v1/general/files/{fileId}")
|
||||
private String fileId;
|
||||
|
||||
@Schema(description = "Original filename for the file")
|
||||
private String fileName;
|
||||
|
||||
@Schema(description = "MIME type of the file", example = "application/pdf")
|
||||
private String contentType;
|
||||
}
|
||||
+218
-2
@@ -7,16 +7,33 @@ import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.apache.commons.io.FilenameUtils;
|
||||
import org.apache.pdfbox.pdmodel.PDDocument;
|
||||
import org.springframework.core.io.FileSystemResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.MediaTypeFactory;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.util.LinkedMultiValueMap;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
import io.github.pixee.security.Filenames;
|
||||
|
||||
import lombok.Data;
|
||||
import lombok.RequiredArgsConstructor;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.service.CustomPDFDocumentFactory;
|
||||
import stirling.software.common.service.FileStorage;
|
||||
import stirling.software.common.service.InternalApiClient;
|
||||
import stirling.software.common.service.ToolMetadataService;
|
||||
import stirling.software.common.util.ExceptionUtils;
|
||||
import stirling.software.common.util.TempFile;
|
||||
import stirling.software.common.util.TempFileManager;
|
||||
import stirling.software.common.util.ZipExtractionUtils;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
|
||||
@@ -24,6 +41,7 @@ import stirling.software.proprietary.model.api.ai.AiWorkflowPhase;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
|
||||
import stirling.software.proprietary.service.PdfContentExtractor.LoadedFile;
|
||||
import stirling.software.proprietary.service.PdfContentExtractor.PdfContentResult;
|
||||
import stirling.software.proprietary.service.PdfContentExtractor.WorkflowArtifact;
|
||||
@@ -39,6 +57,10 @@ public class AiWorkflowService {
|
||||
private final AiEngineClient aiEngineClient;
|
||||
private final PdfContentExtractor pdfContentExtractor;
|
||||
private final ObjectMapper objectMapper;
|
||||
private final InternalApiClient internalApiClient;
|
||||
private final FileStorage fileStorage;
|
||||
private final ToolMetadataService toolMetadataService;
|
||||
private final TempFileManager tempFileManager;
|
||||
|
||||
@FunctionalInterface
|
||||
public interface ProgressListener {
|
||||
@@ -89,12 +111,13 @@ public class AiWorkflowService {
|
||||
AiWorkflowResponse response = invokeOrchestrator(request);
|
||||
return switch (response.getOutcome()) {
|
||||
case NEED_CONTENT -> onNeedContent(response, filesByName, request, listener);
|
||||
case TOOL_CALL -> onToolCall(response, filesByName, listener);
|
||||
case PLAN -> onPlan(response, filesByName, listener);
|
||||
case ANSWER,
|
||||
NOT_FOUND,
|
||||
PLAN,
|
||||
NEED_CLARIFICATION,
|
||||
CANNOT_DO,
|
||||
TOOL_CALL,
|
||||
DRAFT,
|
||||
COMPLETED,
|
||||
UNSUPPORTED_CAPABILITY,
|
||||
CANNOT_CONTINUE ->
|
||||
@@ -174,6 +197,199 @@ public class AiWorkflowService {
|
||||
}
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
private WorkflowState onToolCall(
|
||||
AiWorkflowResponse response,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
ProgressListener listener) {
|
||||
String endpointPath = response.getTool();
|
||||
Map<String, Object> parameters = response.getParameters();
|
||||
if (endpointPath == null || endpointPath.isBlank()) {
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue("AI engine returned tool_call without a tool endpoint."));
|
||||
}
|
||||
if (parameters == null) {
|
||||
parameters = Map.of();
|
||||
}
|
||||
|
||||
try {
|
||||
List<Resource> inputFiles = toResources(filesByName);
|
||||
listener.onProgress(AiWorkflowProgressEvent.executingTool(endpointPath, 1, 1));
|
||||
List<Resource> results = executeStep(endpointPath, parameters, inputFiles);
|
||||
return new WorkflowState.Terminal(
|
||||
buildCompletedResponse(
|
||||
response.getRationale(),
|
||||
results,
|
||||
new ArrayList<>(filesByName.keySet())));
|
||||
} catch (Exception e) {
|
||||
log.error("Failed to execute tool {}: {}", endpointPath, e.getMessage(), e);
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue("Tool execution failed: " + e.getMessage()));
|
||||
}
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
private WorkflowState onPlan(
|
||||
AiWorkflowResponse response,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
ProgressListener listener) {
|
||||
List<Map<String, Object>> steps = response.getSteps();
|
||||
if (steps == null || steps.isEmpty()) {
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue("AI engine returned a plan with no steps."));
|
||||
}
|
||||
|
||||
try {
|
||||
List<Resource> currentFiles = toResources(filesByName);
|
||||
|
||||
for (int i = 0; i < steps.size(); i++) {
|
||||
Map<String, Object> step = steps.get(i);
|
||||
String endpointPath = (String) step.get("tool");
|
||||
Map<String, Object> parameters =
|
||||
step.containsKey("parameters")
|
||||
? (Map<String, Object>) step.get("parameters")
|
||||
: Map.of();
|
||||
|
||||
if (endpointPath == null || endpointPath.isBlank()) {
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue("Plan step " + (i + 1) + " has no tool endpoint."));
|
||||
}
|
||||
|
||||
listener.onProgress(
|
||||
AiWorkflowProgressEvent.executingTool(endpointPath, i + 1, steps.size()));
|
||||
currentFiles = executeStep(endpointPath, parameters, currentFiles);
|
||||
}
|
||||
|
||||
return new WorkflowState.Terminal(
|
||||
buildCompletedResponse(
|
||||
response.getSummary(),
|
||||
currentFiles,
|
||||
new ArrayList<>(filesByName.keySet())));
|
||||
} catch (Exception e) {
|
||||
log.error("Failed to execute plan: {}", e.getMessage(), e);
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue("Plan execution failed: " + e.getMessage()));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a single tool step. If the endpoint accepts multiple files, all files are sent in one
|
||||
* call. Otherwise, the endpoint is called once per file. ZIP responses are unpacked so each
|
||||
* inner file is treated as its own result (e.g. split outputs a ZIP of pages).
|
||||
*/
|
||||
private List<Resource> executeStep(
|
||||
String endpointPath, Map<String, Object> parameters, List<Resource> inputFiles)
|
||||
throws IOException {
|
||||
List<Resource> results = new ArrayList<>();
|
||||
if (toolMetadataService.isMultiInput(endpointPath)) {
|
||||
results.addAll(callEndpoint(endpointPath, parameters, inputFiles));
|
||||
} else {
|
||||
for (Resource file : inputFiles) {
|
||||
results.addAll(callEndpoint(endpointPath, parameters, List.of(file)));
|
||||
}
|
||||
}
|
||||
return results;
|
||||
}
|
||||
|
||||
/**
|
||||
* Call an endpoint and return the response body. Endpoints that are declared as ZIP-returning
|
||||
* in the API spec (multi-output, or {@code Output:ZIP-*}) are unpacked into their individual
|
||||
* entries so callers always see a flat list of result files.
|
||||
*/
|
||||
private List<Resource> callEndpoint(
|
||||
String endpointPath, Map<String, Object> parameters, List<Resource> files)
|
||||
throws IOException {
|
||||
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
|
||||
for (Resource file : files) {
|
||||
body.add("fileInput", file);
|
||||
}
|
||||
for (Map.Entry<String, Object> entry : parameters.entrySet()) {
|
||||
if (entry.getValue() instanceof List<?> list) {
|
||||
for (Object item : list) {
|
||||
body.add(entry.getKey(), item);
|
||||
}
|
||||
} else {
|
||||
body.add(entry.getKey(), entry.getValue());
|
||||
}
|
||||
}
|
||||
ResponseEntity<Resource> response = internalApiClient.post(endpointPath, body);
|
||||
if (!HttpStatus.OK.equals(response.getStatusCode()) || response.getBody() == null) {
|
||||
throw new IOException(
|
||||
"Tool returned HTTP " + response.getStatusCode() + " for " + endpointPath);
|
||||
}
|
||||
Resource resource = response.getBody();
|
||||
if (toolMetadataService.shouldUnpackZipResponse(endpointPath)) {
|
||||
return ZipExtractionUtils.extractZip(resource, tempFileManager);
|
||||
}
|
||||
return List.of(resource);
|
||||
}
|
||||
|
||||
private List<Resource> toResources(Map<String, MultipartFile> filesByName) throws IOException {
|
||||
List<Resource> resources = new ArrayList<>();
|
||||
for (MultipartFile file : filesByName.values()) {
|
||||
TempFile tempFile = tempFileManager.createManagedTempFile("ai-workflow");
|
||||
file.transferTo(tempFile.getPath());
|
||||
final String originalName = Filenames.toSimpleFileName(file.getOriginalFilename());
|
||||
resources.add(
|
||||
new FileSystemResource(tempFile.getFile()) {
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return originalName;
|
||||
}
|
||||
});
|
||||
}
|
||||
return resources;
|
||||
}
|
||||
|
||||
private AiWorkflowResponse buildCompletedResponse(
|
||||
String summary, List<Resource> resultFiles, List<String> inputFileNames)
|
||||
throws IOException {
|
||||
// Store every output file individually so each gets its own Stirling file ID and the
|
||||
// frontend can add them as independent variants without going through a zip.
|
||||
boolean preserveInputNames = inputFileNames.size() == resultFiles.size();
|
||||
List<AiWorkflowResultFile> descriptors = new ArrayList<>();
|
||||
for (int i = 0; i < resultFiles.size(); i++) {
|
||||
Resource resource = resultFiles.get(i);
|
||||
String responseName = resource.getFilename();
|
||||
String inputName = preserveInputNames ? inputFileNames.get(i) : null;
|
||||
// Prefer the input name only for 1:1 operations where the output keeps the same
|
||||
// extension (rotate, compress, etc.). For converters and other extension-changing
|
||||
// tools, the response filename from Content-Disposition is authoritative.
|
||||
String name;
|
||||
if (inputName != null
|
||||
&& FilenameUtils.getExtension(inputName)
|
||||
.equalsIgnoreCase(FilenameUtils.getExtension(responseName))) {
|
||||
name = inputName;
|
||||
} else if (responseName != null) {
|
||||
name = responseName;
|
||||
} else {
|
||||
name = "result-" + (i + 1);
|
||||
}
|
||||
String contentType =
|
||||
MediaTypeFactory.getMediaType(name)
|
||||
.orElse(MediaType.APPLICATION_OCTET_STREAM)
|
||||
.toString();
|
||||
String fileId;
|
||||
try (java.io.InputStream is = resource.getInputStream()) {
|
||||
fileId = fileStorage.storeInputStream(is, name).fileId();
|
||||
}
|
||||
descriptors.add(new AiWorkflowResultFile(fileId, name, contentType));
|
||||
}
|
||||
|
||||
AiWorkflowResponse completed = new AiWorkflowResponse();
|
||||
completed.setOutcome(AiWorkflowOutcome.COMPLETED);
|
||||
completed.setSummary(summary);
|
||||
completed.setResultFiles(descriptors);
|
||||
// Mirror the first file into the legacy single-file fields so existing clients still work.
|
||||
if (!descriptors.isEmpty()) {
|
||||
AiWorkflowResultFile first = descriptors.getFirst();
|
||||
completed.setFileId(first.getFileId());
|
||||
completed.setFileName(first.getFileName());
|
||||
completed.setContentType(first.getContentType());
|
||||
}
|
||||
return completed;
|
||||
}
|
||||
|
||||
private void validateRequest(AiWorkflowRequest request) {
|
||||
for (AiWorkflowFileInput fileInput : request.getFileInputs()) {
|
||||
if (fileInput.getFileInput().isEmpty()) {
|
||||
|
||||
+335
@@ -0,0 +1,335 @@
|
||||
package stirling.software.proprietary.service;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertNotNull;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.zip.ZipEntry;
|
||||
import java.util.zip.ZipOutputStream;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
import org.springframework.core.io.ByteArrayResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.mock.web.MockMultipartFile;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
|
||||
import stirling.software.common.model.ApplicationProperties;
|
||||
import stirling.software.common.service.CustomPDFDocumentFactory;
|
||||
import stirling.software.common.service.FileStorage;
|
||||
import stirling.software.common.service.FileStorage.StoredFile;
|
||||
import stirling.software.common.service.InternalApiClient;
|
||||
import stirling.software.common.service.ToolMetadataService;
|
||||
import stirling.software.common.util.TempFileManager;
|
||||
import stirling.software.common.util.TempFileRegistry;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
|
||||
|
||||
import tools.jackson.databind.ObjectMapper;
|
||||
import tools.jackson.databind.json.JsonMapper;
|
||||
|
||||
/**
|
||||
* Smoke tests for {@link AiWorkflowService}. Covers the TOOL_CALL and PLAN execution paths,
|
||||
* ZIP-response unpacking (split endpoints), multi-input dispatch (merge endpoints), and the 1:1
|
||||
* input-to-output filename preservation rule.
|
||||
*
|
||||
* <p>External collaborators (engine client, internal API client, tool metadata, file storage) are
|
||||
* mocked. {@link TempFileManager} is constructed with a real in-test registry so the service's
|
||||
* temp-file handling exercises real code.
|
||||
*/
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
class AiWorkflowServiceTest {
|
||||
|
||||
private static final String ROTATE_ENDPOINT = "/api/v1/general/rotate-pdf";
|
||||
private static final String SPLIT_ENDPOINT = "/api/v1/general/split-pages";
|
||||
private static final String MERGE_ENDPOINT = "/api/v1/general/merge-pdfs";
|
||||
private static final String COMPRESS_ENDPOINT = "/api/v1/misc/compress-pdf";
|
||||
|
||||
@Mock private CustomPDFDocumentFactory pdfDocumentFactory;
|
||||
@Mock private AiEngineClient aiEngineClient;
|
||||
@Mock private PdfContentExtractor pdfContentExtractor;
|
||||
@Mock private InternalApiClient internalApiClient;
|
||||
@Mock private FileStorage fileStorage;
|
||||
@Mock private ToolMetadataService toolMetadataService;
|
||||
|
||||
@TempDir Path tempDir;
|
||||
|
||||
private TempFileManager tempFileManager;
|
||||
private ObjectMapper objectMapper;
|
||||
private AiWorkflowService service;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
ApplicationProperties props = new ApplicationProperties();
|
||||
props.getSystem().getTempFileManagement().setBaseTmpDir(tempDir.toString());
|
||||
props.getSystem().getTempFileManagement().setPrefix("ai-test-");
|
||||
tempFileManager = new TempFileManager(new TempFileRegistry(), props);
|
||||
objectMapper = JsonMapper.builder().build();
|
||||
|
||||
service =
|
||||
new AiWorkflowService(
|
||||
pdfDocumentFactory,
|
||||
aiEngineClient,
|
||||
pdfContentExtractor,
|
||||
objectMapper,
|
||||
internalApiClient,
|
||||
fileStorage,
|
||||
toolMetadataService,
|
||||
tempFileManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
void toolCallSingleFilePreservesInputFilename() throws IOException {
|
||||
MockMultipartFile input = pdf("input.pdf", "original-pdf-bytes");
|
||||
stubOrchestrator(
|
||||
"""
|
||||
{"outcome":"tool_call","tool":"%s","parameters":{"angle":90},"rationale":"Rotating"}
|
||||
"""
|
||||
.formatted(ROTATE_ENDPOINT));
|
||||
when(toolMetadataService.isMultiInput(ROTATE_ENDPOINT)).thenReturn(false);
|
||||
when(toolMetadataService.shouldUnpackZipResponse(ROTATE_ENDPOINT)).thenReturn(false);
|
||||
stubEndpoint(ROTATE_ENDPOINT, pdfResource("rotated-bytes", "rotated.pdf"));
|
||||
AtomicInteger ids = stubFileStorage();
|
||||
|
||||
AiWorkflowResponse result = service.orchestrate(requestFor(input, "rotate 90"));
|
||||
|
||||
assertEquals(AiWorkflowOutcome.COMPLETED, result.getOutcome());
|
||||
assertEquals(1, result.getResultFiles().size());
|
||||
// 1:1 mapping — the single output should inherit the single input's filename.
|
||||
assertEquals("input.pdf", result.getResultFiles().get(0).getFileName());
|
||||
assertEquals("file-1", result.getResultFiles().get(0).getFileId());
|
||||
assertEquals(1, ids.get());
|
||||
verify(internalApiClient, times(1)).post(eq(ROTATE_ENDPOINT), any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void toolCallZipResponseUnpacksIntoMultipleResults() throws IOException {
|
||||
MockMultipartFile input = pdf("doc.pdf", "original");
|
||||
stubOrchestrator(
|
||||
"""
|
||||
{"outcome":"tool_call","tool":"%s","parameters":{},"rationale":"Splitting"}
|
||||
"""
|
||||
.formatted(SPLIT_ENDPOINT));
|
||||
when(toolMetadataService.isMultiInput(SPLIT_ENDPOINT)).thenReturn(false);
|
||||
when(toolMetadataService.shouldUnpackZipResponse(SPLIT_ENDPOINT)).thenReturn(true);
|
||||
stubEndpoint(
|
||||
SPLIT_ENDPOINT,
|
||||
zipResource(
|
||||
"doc.zip",
|
||||
List.of(
|
||||
new ZipEntryBytes("page-1.pdf", "page-one"),
|
||||
new ZipEntryBytes("page-2.pdf", "page-two"),
|
||||
new ZipEntryBytes("page-3.pdf", "page-three"))));
|
||||
stubFileStorage();
|
||||
|
||||
AiWorkflowResponse result = service.orchestrate(requestFor(input, "split"));
|
||||
|
||||
assertEquals(AiWorkflowOutcome.COMPLETED, result.getOutcome());
|
||||
assertEquals(3, result.getResultFiles().size());
|
||||
// Input count (1) != output count (3) so the per-entry filename is kept.
|
||||
assertEquals("page-1.pdf", result.getResultFiles().get(0).getFileName());
|
||||
assertEquals("page-2.pdf", result.getResultFiles().get(1).getFileName());
|
||||
assertEquals("page-3.pdf", result.getResultFiles().get(2).getFileName());
|
||||
}
|
||||
|
||||
@Test
|
||||
void multiInputEndpointIsCalledOnceWithAllFiles() throws IOException {
|
||||
MockMultipartFile a = pdf("a.pdf", "a-bytes");
|
||||
MockMultipartFile b = pdf("b.pdf", "b-bytes");
|
||||
stubOrchestrator(
|
||||
"""
|
||||
{"outcome":"tool_call","tool":"%s","parameters":{},"rationale":"Merging"}
|
||||
"""
|
||||
.formatted(MERGE_ENDPOINT));
|
||||
when(toolMetadataService.isMultiInput(MERGE_ENDPOINT)).thenReturn(true);
|
||||
when(toolMetadataService.shouldUnpackZipResponse(MERGE_ENDPOINT)).thenReturn(false);
|
||||
stubEndpoint(MERGE_ENDPOINT, pdfResource("merged-bytes", "merged.pdf"));
|
||||
stubFileStorage();
|
||||
|
||||
AiWorkflowResponse result =
|
||||
service.orchestrate(requestFor(new MockMultipartFile[] {a, b}, "merge these"));
|
||||
|
||||
assertEquals(AiWorkflowOutcome.COMPLETED, result.getOutcome());
|
||||
assertEquals(1, result.getResultFiles().size());
|
||||
// Two inputs but only one output → filename is not preserved from either input.
|
||||
assertEquals("merged.pdf", result.getResultFiles().get(0).getFileName());
|
||||
verify(internalApiClient, times(1)).post(eq(MERGE_ENDPOINT), any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void singleInputEndpointIsCalledOncePerFile() throws IOException {
|
||||
MockMultipartFile a = pdf("a.pdf", "a-bytes");
|
||||
MockMultipartFile b = pdf("b.pdf", "b-bytes");
|
||||
stubOrchestrator(
|
||||
"""
|
||||
{"outcome":"tool_call","tool":"%s","parameters":{"angle":90},"rationale":"Rotating"}
|
||||
"""
|
||||
.formatted(ROTATE_ENDPOINT));
|
||||
when(toolMetadataService.isMultiInput(ROTATE_ENDPOINT)).thenReturn(false);
|
||||
when(toolMetadataService.shouldUnpackZipResponse(ROTATE_ENDPOINT)).thenReturn(false);
|
||||
stubEndpoint(ROTATE_ENDPOINT, pdfResource("rotated", "rotated.pdf"));
|
||||
stubFileStorage();
|
||||
|
||||
AiWorkflowResponse result =
|
||||
service.orchestrate(requestFor(new MockMultipartFile[] {a, b}, "rotate both"));
|
||||
|
||||
assertEquals(AiWorkflowOutcome.COMPLETED, result.getOutcome());
|
||||
assertEquals(2, result.getResultFiles().size());
|
||||
// Per-file loop dispatches one call per input file.
|
||||
verify(internalApiClient, times(2)).post(eq(ROTATE_ENDPOINT), any());
|
||||
// 1:1 mapping preserves each input's filename.
|
||||
assertEquals("a.pdf", result.getResultFiles().get(0).getFileName());
|
||||
assertEquals("b.pdf", result.getResultFiles().get(1).getFileName());
|
||||
}
|
||||
|
||||
@Test
|
||||
void planExecutesStepsSequentially() throws IOException {
|
||||
MockMultipartFile input = pdf("input.pdf", "bytes");
|
||||
stubOrchestrator(
|
||||
"""
|
||||
{
|
||||
"outcome":"plan",
|
||||
"summary":"Rotate then compress",
|
||||
"steps":[
|
||||
{"tool":"%s","parameters":{"angle":90}},
|
||||
{"tool":"%s","parameters":{}}
|
||||
]
|
||||
}
|
||||
"""
|
||||
.formatted(ROTATE_ENDPOINT, COMPRESS_ENDPOINT));
|
||||
when(toolMetadataService.isMultiInput(anyString())).thenReturn(false);
|
||||
when(toolMetadataService.shouldUnpackZipResponse(anyString())).thenReturn(false);
|
||||
stubEndpoint(ROTATE_ENDPOINT, pdfResource("rotated", "rotated.pdf"));
|
||||
stubEndpoint(COMPRESS_ENDPOINT, pdfResource("compressed", "compressed.pdf"));
|
||||
stubFileStorage();
|
||||
|
||||
AiWorkflowResponse result = service.orchestrate(requestFor(input, "rotate and compress"));
|
||||
|
||||
assertEquals(AiWorkflowOutcome.COMPLETED, result.getOutcome());
|
||||
assertEquals(1, result.getResultFiles().size());
|
||||
// 1:1 input → output mapping at the plan level preserves the input's filename.
|
||||
assertEquals("input.pdf", result.getResultFiles().get(0).getFileName());
|
||||
verify(internalApiClient, times(1)).post(eq(ROTATE_ENDPOINT), any());
|
||||
verify(internalApiClient, times(1)).post(eq(COMPRESS_ENDPOINT), any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void toolCallWithoutEndpointFallsBackToCannotContinue() throws IOException {
|
||||
MockMultipartFile input = pdf("input.pdf", "bytes");
|
||||
stubOrchestrator("{\"outcome\":\"tool_call\",\"parameters\":{}}");
|
||||
|
||||
AiWorkflowResponse result = service.orchestrate(requestFor(input, "do something"));
|
||||
|
||||
assertEquals(AiWorkflowOutcome.CANNOT_CONTINUE, result.getOutcome());
|
||||
assertNotNull(result.getReason());
|
||||
verify(internalApiClient, never()).post(anyString(), any());
|
||||
}
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
private void stubOrchestrator(String responseJson) throws IOException {
|
||||
when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString())).thenReturn(responseJson);
|
||||
}
|
||||
|
||||
private void stubEndpoint(String endpoint, Resource body) {
|
||||
when(internalApiClient.post(eq(endpoint), any(MultiValueMap.class)))
|
||||
.thenReturn(ResponseEntity.ok(body));
|
||||
}
|
||||
|
||||
/**
|
||||
* Stub {@link FileStorage#storeInputStream} with sequential file IDs and an accurate byte
|
||||
* count. Returns the counter so tests can assert how many stores happened.
|
||||
*/
|
||||
private AtomicInteger stubFileStorage() throws IOException {
|
||||
AtomicInteger counter = new AtomicInteger();
|
||||
when(fileStorage.storeInputStream(any(InputStream.class), anyString()))
|
||||
.thenAnswer(
|
||||
inv -> {
|
||||
InputStream is = inv.getArgument(0);
|
||||
long size = is.readAllBytes().length;
|
||||
return new StoredFile("file-" + counter.incrementAndGet(), size);
|
||||
});
|
||||
return counter;
|
||||
}
|
||||
|
||||
private static MockMultipartFile pdf(String filename, String content) {
|
||||
return new MockMultipartFile("fileInput", filename, "application/pdf", content.getBytes());
|
||||
}
|
||||
|
||||
private static AiWorkflowRequest requestFor(MockMultipartFile file, String message) {
|
||||
return requestFor(new MockMultipartFile[] {file}, message);
|
||||
}
|
||||
|
||||
private static AiWorkflowRequest requestFor(MockMultipartFile[] files, String message) {
|
||||
AiWorkflowRequest request = new AiWorkflowRequest();
|
||||
List<AiWorkflowFileInput> inputs = new ArrayList<>();
|
||||
for (MockMultipartFile file : files) {
|
||||
AiWorkflowFileInput fileInput = new AiWorkflowFileInput();
|
||||
fileInput.setFileInput(file);
|
||||
inputs.add(fileInput);
|
||||
}
|
||||
request.setFileInputs(inputs);
|
||||
request.setUserMessage(message);
|
||||
return request;
|
||||
}
|
||||
|
||||
private static ByteArrayResource pdfResource(String content, String filename) {
|
||||
return new ByteArrayResource(content.getBytes()) {
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return filename;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private static ByteArrayResource zipResource(String filename, List<ZipEntryBytes> entries)
|
||||
throws IOException {
|
||||
ByteArrayOutputStream baos = new ByteArrayOutputStream();
|
||||
try (ZipOutputStream zos = new ZipOutputStream(baos)) {
|
||||
for (ZipEntryBytes entry : entries) {
|
||||
zos.putNextEntry(new ZipEntry(entry.name()));
|
||||
zos.write(entry.bytes());
|
||||
zos.closeEntry();
|
||||
}
|
||||
}
|
||||
byte[] zipBytes = baos.toByteArray();
|
||||
return new ByteArrayResource(zipBytes) {
|
||||
@Override
|
||||
public String getFilename() {
|
||||
return filename;
|
||||
}
|
||||
|
||||
@Override
|
||||
public InputStream getInputStream() {
|
||||
return new ByteArrayInputStream(zipBytes);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private record ZipEntryBytes(String name, byte[] bytes) {
|
||||
ZipEntryBytes(String name, String content) {
|
||||
this(name, content.getBytes());
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user