Add streaming to Engine orchestrator (#6094)

# Description of Changes
Adds a streaming endpoint to the Java AI orchestrator
(`/api/v1/ai/orchestrate/stream` in addition to the existing
`/api/v1/ai/orchestrate`). This allows the caller to get updates of what
stage of orchestration is being run at the time so UIs can give the user
feedback.

Also contains some dubious Gradle changes to suppress errors coming from
Spotless, when it crashes in Google stuff. I'm not sure if that's
appropriate to add, feel free to ask for changes in review.

app/proprietary/src/main/java/stirling/software/proprietary/config/AsyncConfig.java

@@ -48,4 +48,12 @@ public class AsyncConfig {
         adapter.setTaskDecorator(new MDCContextTaskDecorator());
         return adapter;
     }
+
+    @Bean(name = "aiStreamExecutor")
+    public Executor aiStreamExecutor() {
+        TaskExecutorAdapter adapter =
+                new TaskExecutorAdapter(Executors.newVirtualThreadPerTaskExecutor());
+        adapter.setTaskDecorator(new MDCContextTaskDecorator());
+        return adapter;
+    }
 }

app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java

@@ -1,7 +1,10 @@
 package stirling.software.proprietary.controller.api;
 
 import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.Executor;
 
+import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
@@ -12,6 +15,7 @@ import org.springframework.web.bind.annotation.RequestBody;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.server.ResponseStatusException;
+import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 
 import io.swagger.v3.oas.annotations.Hidden;
 import io.swagger.v3.oas.annotations.Operation;
@@ -19,7 +23,6 @@ import io.swagger.v3.oas.annotations.tags.Tag;
 
 import jakarta.validation.Valid;
 
-import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
 import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
@@ -34,7 +37,6 @@ import tools.jackson.databind.ObjectMapper;
 @Slf4j
 @RestController
 @RequestMapping("/api/v1/ai")
-@RequiredArgsConstructor
 @Hidden
 @Tag(name = "AI Engine", description = "Endpoints for AI-powered PDF workflows")
 public class AiEngineController {
@@ -42,6 +44,18 @@ public class AiEngineController {
     private final AiEngineClient aiEngineClient;
     private final AiWorkflowService aiWorkflowService;
     private final ObjectMapper objectMapper;
+    private final Executor aiStreamExecutor;
+
+    public AiEngineController(
+            AiEngineClient aiEngineClient,
+            AiWorkflowService aiWorkflowService,
+            ObjectMapper objectMapper,
+            @Qualifier("aiStreamExecutor") Executor aiStreamExecutor) {
+        this.aiEngineClient = aiEngineClient;
+        this.aiWorkflowService = aiWorkflowService;
+        this.objectMapper = objectMapper;
+        this.aiStreamExecutor = aiStreamExecutor;
+    }
 
     @GetMapping("/health")
     @Operation(
@@ -62,6 +76,49 @@ public class AiEngineController {
         return ResponseEntity.ok(aiWorkflowService.orchestrate(request));
     }
 
+    @PostMapping(value = "/orchestrate/stream", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
+    @Operation(
+            summary = "Run an AI workflow with streaming progress",
+            description =
+                    "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);
+
+        emitter.onTimeout(
+                () -> {
+                    log.warn("SSE emitter timed out for AI orchestration stream");
+                    emitter.complete();
+                });
+        emitter.onError(e -> log.warn("SSE emitter error for AI orchestration stream", e));
+
+        aiStreamExecutor.execute(() -> runOrchestrationStream(request, emitter));
+
+        return emitter;
+    }
+
+    private void runOrchestrationStream(AiWorkflowRequest request, SseEmitter emitter) {
+        try {
+            AiWorkflowResponse result =
+                    aiWorkflowService.orchestrate(
+                            request, progress -> sendEvent(emitter, "progress", progress));
+            sendEvent(emitter, "result", result);
+            emitter.complete();
+        } catch (Exception e) {
+            log.error("AI orchestration stream failed", e);
+            sendEvent(emitter, "error", Map.of("message", e.getMessage()));
+            emitter.completeWithError(e);
+        }
+    }
+
+    private void sendEvent(SseEmitter emitter, String name, Object data) {
+        try {
+            emitter.send(SseEmitter.event().name(name).data(data, MediaType.APPLICATION_JSON));
+        } catch (IOException e) {
+            log.debug("Failed to send SSE event (client may have disconnected)", e);
+        }
+    }
+
     @PostMapping(value = "/pdf/edit", consumes = MediaType.APPLICATION_JSON_VALUE)
     @Operation(
             summary = "Generate a PDF edit plan",

app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowPhase.java

@@ -0,0 +1,33 @@
+package stirling.software.proprietary.model.api.ai;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonValue;
+
+/** Progress phases emitted during AI workflow orchestration. */
+public enum AiWorkflowPhase {
+    ANALYZING("analyzing"),
+    CALLING_ENGINE("calling_engine"),
+    EXTRACTING_CONTENT("extracting_content"),
+    PROCESSING("processing");
+
+    private final String value;
+
+    AiWorkflowPhase(String value) {
+        this.value = value;
+    }
+
+    @JsonValue
+    public String getValue() {
+        return value;
+    }
+
+    @JsonCreator
+    public static AiWorkflowPhase fromValue(String value) {
+        for (AiWorkflowPhase phase : values()) {
+            if (phase.value.equals(value)) {
+                return phase;
+            }
+        }
+        throw new IllegalArgumentException("Unknown AI workflow phase: " + value);
+    }
+}

app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowProgressEvent.java

@@ -0,0 +1,15 @@
+package stirling.software.proprietary.model.api.ai;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+
+@Data
+@AllArgsConstructor
+public class AiWorkflowProgressEvent {
+    private AiWorkflowPhase phase;
+    private long timestamp;
+
+    public static AiWorkflowProgressEvent of(AiWorkflowPhase phase) {
+        return new AiWorkflowProgressEvent(phase, System.currentTimeMillis());
+    }
+}

app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java

@@ -20,6 +20,8 @@ import stirling.software.common.util.ExceptionUtils;
 import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput;
 import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
 import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
+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.service.PdfContentExtractor.LoadedFile;
@@ -38,6 +40,13 @@ public class AiWorkflowService {
     private final PdfContentExtractor pdfContentExtractor;
     private final ObjectMapper objectMapper;
 
+    @FunctionalInterface
+    public interface ProgressListener {
+        void onProgress(AiWorkflowProgressEvent event);
+    }
+
+    private static final ProgressListener NOOP_LISTENER = event -> {};
+
     private sealed interface WorkflowState {
         record Pending(WorkflowTurnRequest request) implements WorkflowState {}
 
@@ -45,6 +54,11 @@ public class AiWorkflowService {
     }
 
     public AiWorkflowResponse orchestrate(AiWorkflowRequest request) throws IOException {
+        return orchestrate(request, NOOP_LISTENER);
+    }
+
+    public AiWorkflowResponse orchestrate(AiWorkflowRequest request, ProgressListener listener)
+            throws IOException {
         validateRequest(request);
 
         Map<String, MultipartFile> filesByName = new LinkedHashMap<>();
@@ -57,19 +71,24 @@ public class AiWorkflowService {
         initialRequest.setUserMessage(request.getUserMessage().trim());
         initialRequest.setFileNames(new ArrayList<>(filesByName.keySet()));
 
+        listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.ANALYZING));
+
         WorkflowState state = new WorkflowState.Pending(initialRequest);
         while (state instanceof WorkflowState.Pending pending) {
-            state = advance(pending.request(), filesByName);
+            state = advance(pending.request(), filesByName, listener);
         }
         return ((WorkflowState.Terminal) state).response();
     }
 
     private WorkflowState advance(
-            WorkflowTurnRequest request, Map<String, MultipartFile> filesByName)
+            WorkflowTurnRequest request,
+            Map<String, MultipartFile> filesByName,
+            ProgressListener listener)
             throws IOException {
+        listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.CALLING_ENGINE));
         AiWorkflowResponse response = invokeOrchestrator(request);
         return switch (response.getOutcome()) {
-            case NEED_CONTENT -> onNeedContent(response, filesByName, request);
+            case NEED_CONTENT -> onNeedContent(response, filesByName, request, listener);
             case ANSWER,
                     NOT_FOUND,
                     PLAN,
@@ -86,7 +105,8 @@ public class AiWorkflowService {
     private WorkflowState onNeedContent(
             AiWorkflowResponse response,
             Map<String, MultipartFile> filesByName,
-            WorkflowTurnRequest request)
+            WorkflowTurnRequest request,
+            ProgressListener listener)
             throws IOException {
         if (!request.getArtifacts().isEmpty()) {
             return new WorkflowState.Terminal(
@@ -119,6 +139,8 @@ public class AiWorkflowService {
                                         Collectors.toMap(
                                                 AiWorkflowFileRequest::getFileName, r -> r));
 
+        listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.EXTRACTING_CONTENT));
+
         List<LoadedFile> loadedFiles = new ArrayList<>();
         try {
             for (String fileName : fileNamesToLoad) {
@@ -133,6 +155,8 @@ public class AiWorkflowService {
                             response.getMaxPages(),
                             response.getMaxCharacters());
 
+            listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.PROCESSING));
+
             WorkflowTurnRequest nextRequest = new WorkflowTurnRequest();
             nextRequest.setUserMessage(request.getUserMessage());
             nextRequest.setFileNames(request.getFileNames());