diff --git a/app/common/src/main/java/stirling/software/SPDF/config/EndpointConfiguration.java b/app/common/src/main/java/stirling/software/SPDF/config/EndpointConfiguration.java index 9d3b92cd5..a154ed2a7 100644 --- a/app/common/src/main/java/stirling/software/SPDF/config/EndpointConfiguration.java +++ b/app/common/src/main/java/stirling/software/SPDF/config/EndpointConfiguration.java @@ -68,6 +68,36 @@ public class EndpointConfiguration { return endpoint.startsWith("/") ? endpoint.substring(1) : endpoint; } + /** + * Translate a full request URI like {@code /api/v1/general/remove-pages} into the endpoint key + * used by this configuration ({@code remove-pages}). Convert endpoints are a special case - + * {@code /api/v1/convert/pdf/img} is registered as {@code pdf-to-img}. Returns {@code null} if + * the URI is not an {@code /api/v1//} path. + */ + public static String endpointKeyForUri(String uri) { + if (uri == null || !uri.contains("/api/v1")) { + return null; + } + String[] parts = uri.split("/"); + if (parts.length <= 4) { + return null; + } + if ("convert".equals(parts[3]) && parts.length > 5) { + return parts[4] + "-to-" + parts[5]; + } + return parts[4]; + } + + /** + * Convenience wrapper around {@link #isEndpointEnabled(String)} that accepts a full request URI + * and translates it to the endpoint key. Falls back to treating the URI itself as a key for + * non-{@code /api/v1/...} paths so callers can pass arbitrary URIs. + */ + public boolean isEndpointEnabledForUri(String uri) { + String key = endpointKeyForUri(uri); + return isEndpointEnabled(key != null ? key : uri); + } + public void enableEndpoint(String endpoint) { String normalized = normalizeEndpoint(endpoint); endpointStatuses.put(normalized, true); diff --git a/app/common/src/test/java/stirling/software/SPDF/config/EndpointConfigurationTest.java b/app/common/src/test/java/stirling/software/SPDF/config/EndpointConfigurationTest.java new file mode 100644 index 000000000..6c827ffc2 --- /dev/null +++ b/app/common/src/test/java/stirling/software/SPDF/config/EndpointConfigurationTest.java @@ -0,0 +1,39 @@ +package stirling.software.SPDF.config; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +import org.junit.jupiter.api.Test; + +class EndpointConfigurationTest { + + @Test + void endpointKeyForUriExtractsSimpleKebab() { + assertEquals( + "remove-pages", + EndpointConfiguration.endpointKeyForUri("/api/v1/general/remove-pages")); + assertEquals( + "compress-pdf", + EndpointConfiguration.endpointKeyForUri("/api/v1/misc/compress-pdf")); + assertEquals( + "add-watermark", + EndpointConfiguration.endpointKeyForUri("/api/v1/security/add-watermark")); + } + + @Test + void endpointKeyForUriComposesConvertEndpoints() { + assertEquals( + "pdf-to-img", EndpointConfiguration.endpointKeyForUri("/api/v1/convert/pdf/img")); + assertEquals( + "pdf-to-word", EndpointConfiguration.endpointKeyForUri("/api/v1/convert/pdf/word")); + assertEquals( + "html-to-pdf", EndpointConfiguration.endpointKeyForUri("/api/v1/convert/html/pdf")); + } + + @Test + void endpointKeyForUriReturnsNullForNonApiPaths() { + assertNull(EndpointConfiguration.endpointKeyForUri(null)); + assertNull(EndpointConfiguration.endpointKeyForUri("/some-page")); + assertNull(EndpointConfiguration.endpointKeyForUri("/api/v1/general")); + } +} diff --git a/app/core/src/main/java/stirling/software/SPDF/config/EndpointInterceptor.java b/app/core/src/main/java/stirling/software/SPDF/config/EndpointInterceptor.java index 52fb42e07..ce96bbd17 100644 --- a/app/core/src/main/java/stirling/software/SPDF/config/EndpointInterceptor.java +++ b/app/core/src/main/java/stirling/software/SPDF/config/EndpointInterceptor.java @@ -21,28 +21,7 @@ public class EndpointInterceptor implements HandlerInterceptor { HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception { String requestURI = request.getRequestURI(); - boolean isEnabled; - - // Extract the specific endpoint name (e.g: /api/v1/general/remove-pages -> remove-pages) - if (requestURI.contains("/api/v1") && requestURI.split("/").length > 4) { - - String[] requestURIParts = requestURI.split("/"); - String requestEndpoint; - - // Endpoint: /api/v1/convert/pdf/img becomes pdf-to-img - if ("convert".equals(requestURIParts[3]) && requestURIParts.length > 5) { - requestEndpoint = requestURIParts[4] + "-to-" + requestURIParts[5]; - } else { - requestEndpoint = requestURIParts[4]; - } - - log.debug("Request endpoint: {}", requestEndpoint); - isEnabled = endpointConfiguration.isEndpointEnabled(requestEndpoint); - log.debug("Is endpoint enabled: {}", isEnabled); - } else { - isEnabled = endpointConfiguration.isEndpointEnabled(requestURI); - } - + boolean isEnabled = endpointConfiguration.isEndpointEnabledForUri(requestURI); if (!isEnabled) { response.sendError(HttpServletResponse.SC_FORBIDDEN, "This endpoint is disabled"); return false; diff --git a/app/core/src/test/java/stirling/software/SPDF/config/EndpointInterceptorTest.java b/app/core/src/test/java/stirling/software/SPDF/config/EndpointInterceptorTest.java index be97d7e65..b064a169e 100644 --- a/app/core/src/test/java/stirling/software/SPDF/config/EndpointInterceptorTest.java +++ b/app/core/src/test/java/stirling/software/SPDF/config/EndpointInterceptorTest.java @@ -27,60 +27,19 @@ class EndpointInterceptorTest { } @Test - void preHandleAllowsEnabledApiEndpoint() throws Exception { + void preHandleAllowsEnabledEndpoint() throws Exception { when(request.getRequestURI()).thenReturn("/api/v1/general/remove-pages"); - when(endpointConfiguration.isEndpointEnabled("remove-pages")).thenReturn(true); + when(endpointConfiguration.isEndpointEnabledForUri("/api/v1/general/remove-pages")) + .thenReturn(true); assertTrue(interceptor.preHandle(request, response, new Object())); } @Test - void preHandleBlocksDisabledApiEndpoint() throws Exception { + void preHandleBlocksDisabledEndpoint() throws Exception { when(request.getRequestURI()).thenReturn("/api/v1/general/remove-pages"); - when(endpointConfiguration.isEndpointEnabled("remove-pages")).thenReturn(false); + when(endpointConfiguration.isEndpointEnabledForUri("/api/v1/general/remove-pages")) + .thenReturn(false); assertFalse(interceptor.preHandle(request, response, new Object())); verify(response).sendError(HttpServletResponse.SC_FORBIDDEN, "This endpoint is disabled"); } - - @Test - void preHandleExtractsConvertEndpointCorrectly() throws Exception { - when(request.getRequestURI()).thenReturn("/api/v1/convert/pdf/img"); - when(endpointConfiguration.isEndpointEnabled("pdf-to-img")).thenReturn(true); - assertTrue(interceptor.preHandle(request, response, new Object())); - } - - @Test - void preHandleBlocksDisabledConvertEndpoint() throws Exception { - when(request.getRequestURI()).thenReturn("/api/v1/convert/pdf/img"); - when(endpointConfiguration.isEndpointEnabled("pdf-to-img")).thenReturn(false); - assertFalse(interceptor.preHandle(request, response, new Object())); - } - - @Test - void preHandleUsesFullUriForNonApiPaths() throws Exception { - when(request.getRequestURI()).thenReturn("/some-page"); - when(endpointConfiguration.isEndpointEnabled("/some-page")).thenReturn(true); - assertTrue(interceptor.preHandle(request, response, new Object())); - } - - @Test - void preHandleBlocksDisabledNonApiPath() throws Exception { - when(request.getRequestURI()).thenReturn("/some-page"); - when(endpointConfiguration.isEndpointEnabled("/some-page")).thenReturn(false); - assertFalse(interceptor.preHandle(request, response, new Object())); - } - - @Test - void preHandleUsesFullUriForShortApiPath() throws Exception { - // URI with /api/v1 but not enough segments (split length <= 4) - when(request.getRequestURI()).thenReturn("/api/v1/general"); - when(endpointConfiguration.isEndpointEnabled("/api/v1/general")).thenReturn(true); - assertTrue(interceptor.preHandle(request, response, new Object())); - } - - @Test - void preHandleExtractsNonConvertApiEndpoint() throws Exception { - when(request.getRequestURI()).thenReturn("/api/v1/security/add-watermark"); - when(endpointConfiguration.isEndpointEnabled("add-watermark")).thenReturn(true); - assertTrue(interceptor.preHandle(request, response, new Object())); - } } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java index 279f40881..068314903 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java @@ -34,11 +34,14 @@ 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.AiEngineEndpointResolver; import stirling.software.proprietary.service.AiWorkflowService; import tools.jackson.core.JacksonException; import tools.jackson.databind.JsonNode; import tools.jackson.databind.ObjectMapper; +import tools.jackson.databind.node.ArrayNode; +import tools.jackson.databind.node.ObjectNode; @Slf4j @RestController @@ -53,6 +56,7 @@ public class AiEngineController { private final Executor aiStreamExecutor; private final TaskManager taskManager; private final JobOwnershipService jobOwnershipService; + private final AiEngineEndpointResolver endpointResolver; /** * SSE emitter timeout. Long enough to accommodate multi-gigabyte PDF workflows (OCR on a @@ -68,13 +72,15 @@ public class AiEngineController { ObjectMapper objectMapper, @Qualifier("aiStreamExecutor") Executor aiStreamExecutor, TaskManager taskManager, - JobOwnershipService jobOwnershipService) { + JobOwnershipService jobOwnershipService, + AiEngineEndpointResolver endpointResolver) { this.aiEngineClient = aiEngineClient; this.aiWorkflowService = aiWorkflowService; this.objectMapper = objectMapper; this.aiStreamExecutor = aiStreamExecutor; this.taskManager = taskManager; this.jobOwnershipService = jobOwnershipService; + this.endpointResolver = endpointResolver; } @GetMapping("/health") @@ -197,17 +203,36 @@ public class AiEngineController { "Sends a user message to the PDF edit agent which returns a structured plan" + " of tool operations to perform") public ResponseEntity pdfEdit(@RequestBody String requestBody) throws IOException { - validateJson(requestBody); - String response = aiEngineClient.post("/api/v1/pdf/edit", requestBody); + JsonNode parsed = parseJson(requestBody); + if (!parsed.isObject()) { + throw new ResponseStatusException( + HttpStatus.BAD_REQUEST, "Request body must be a JSON object"); + } + String forwardedBody = withEnabledEndpoints((ObjectNode) parsed); + String response = aiEngineClient.post("/api/v1/pdf/edit", forwardedBody); return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response); } - private void validateJson(String body) { + private JsonNode parseJson(String body) { try { - objectMapper.readValue(body, JsonNode.class); + return objectMapper.readValue(body, JsonNode.class); } catch (JacksonException e) { throw new ResponseStatusException( HttpStatus.BAD_REQUEST, "Request body is not valid JSON"); } } + + /** + * Always overwrite {@code enabled_endpoints} with the server's view of which endpoints are + * usable. The engine must not trust a client-supplied list - the gate is owned by the Java + * EndpointConfiguration. Values are full URL paths (e.g. {@code /api/v1/misc/compress-pdf}) + * that the engine matches against its {@code ToolEndpoint} enum, silently dropping any it + * doesn't recognise (which lets the two sides drift in either direction without breaking). + */ + private String withEnabledEndpoints(ObjectNode body) { + ArrayNode enabled = objectMapper.createArrayNode(); + endpointResolver.getEnabledEndpointUrls().forEach(enabled::add); + body.set("enabled_endpoints", enabled); + return body.toString(); + } } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineEndpointResolver.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineEndpointResolver.java new file mode 100644 index 000000000..ee2ce63eb --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineEndpointResolver.java @@ -0,0 +1,87 @@ +package stirling.software.proprietary.service; + +import java.lang.reflect.Method; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.TreeSet; + +import org.springframework.context.ApplicationContext; +import org.springframework.context.event.ContextRefreshedEvent; +import org.springframework.context.event.EventListener; +import org.springframework.stereotype.Service; +import org.springframework.web.servlet.mvc.method.RequestMappingInfo; +import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping; + +import lombok.extern.slf4j.Slf4j; + +import stirling.software.SPDF.config.EndpointConfiguration; + +/** + * Discovers every {@code /api/v1/...} request mapping in the application and exposes the subset + * that {@link EndpointConfiguration} reports as currently enabled. The AI engine receives this list + * as-is and silently drops anything it doesn't recognise, so we don't try to predict what the + * engine considers a tool - we just emit what's enabled here. + */ +@Slf4j +@Service +public class AiEngineEndpointResolver { + + private static final String API_PREFIX = "/api/v1/"; + + private final ApplicationContext applicationContext; + private final EndpointConfiguration endpointConfiguration; + // Written once on the Spring startup thread during ContextRefreshedEvent, read on HTTP + // request threads. Spring's lifecycle establishes happens-before (the servlet container + // and its worker threads are started after refresh completes), so no volatile is needed. + private Set apiUrls = Set.of(); + + public AiEngineEndpointResolver( + ApplicationContext applicationContext, EndpointConfiguration endpointConfiguration) { + this.applicationContext = applicationContext; + this.endpointConfiguration = endpointConfiguration; + } + + @EventListener(ContextRefreshedEvent.class) + public void discoverApiUrls() { + Set discovered = new TreeSet<>(); + for (RequestMappingHandlerMapping mapping : + applicationContext.getBeansOfType(RequestMappingHandlerMapping.class).values()) { + for (RequestMappingInfo info : mapping.getHandlerMethods().keySet()) { + for (String pattern : extractPatterns(info)) { + if (pattern.startsWith(API_PREFIX)) { + discovered.add(pattern); + } + } + } + } + apiUrls = Set.copyOf(discovered); + log.debug("Discovered {} /api/v1/ endpoint URLs for AI engine filtering", apiUrls.size()); + } + + public List getEnabledEndpointUrls() { + return apiUrls.stream() + .filter(endpointConfiguration::isEndpointEnabledForUri) + .sorted() + .toList(); + } + + private static Set extractPatterns(RequestMappingInfo info) { + try { + Method getDirectPaths = info.getClass().getMethod("getDirectPaths"); + Object result = getDirectPaths.invoke(info); + if (result instanceof Set set) { + Set patterns = new HashSet<>(); + for (Object value : set) { + if (value instanceof String s) { + patterns.add(s); + } + } + return patterns; + } + } catch (Exception e) { + log.trace("getDirectPaths unavailable on RequestMappingInfo", e); + } + return Set.of(); + } +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java index a8d923053..c77da89ce 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java @@ -71,6 +71,7 @@ public class AiWorkflowService { private final ToolMetadataService toolMetadataService; private final TempFileManager tempFileManager; private final FileIdStrategy fileIdStrategy; + private final AiEngineEndpointResolver endpointResolver; @FunctionalInterface public interface ProgressListener { @@ -125,6 +126,7 @@ public class AiWorkflowService { request.getConversationHistory() == null ? new ArrayList<>() : new ArrayList<>(request.getConversationHistory())); + initialRequest.setEnabledEndpoints(endpointResolver.getEnabledEndpointUrls()); listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.ANALYZING)); @@ -219,6 +221,7 @@ public class AiWorkflowService { nextRequest.setConversationHistory(request.getConversationHistory()); nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults)); nextRequest.setResumeWith(response.getResumeWith()); + nextRequest.setEnabledEndpoints(request.getEnabledEndpoints()); return new WorkflowState.Pending(nextRequest); } finally { for (LoadedFile lf : loadedFiles) { @@ -624,5 +627,6 @@ public class AiWorkflowService { private List conversationHistory = new ArrayList<>(); private List artifacts = new ArrayList<>(); private String resumeWith; + private List enabledEndpoints = new ArrayList<>(); } } diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java index 620c62e5d..b0ffbf19d 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java @@ -79,6 +79,7 @@ class AiWorkflowServiceTest { @Mock private FileStorage fileStorage; @Mock private ToolMetadataService toolMetadataService; @Mock private FileIdStrategy fileIdStrategy; + @Mock private AiEngineEndpointResolver endpointResolver; @TempDir Path tempDir; @@ -110,7 +111,9 @@ class AiWorkflowServiceTest { fileStorage, toolMetadataService, tempFileManager, - fileIdStrategy); + fileIdStrategy, + endpointResolver); + when(endpointResolver.getEnabledEndpointUrls()).thenReturn(List.of()); } @Test diff --git a/engine/src/stirling/agents/pdf_edit.py b/engine/src/stirling/agents/pdf_edit.py index aecea50aa..4dac512c1 100644 --- a/engine/src/stirling/agents/pdf_edit.py +++ b/engine/src/stirling/agents/pdf_edit.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from collections.abc import Iterable from typing import Literal, overload from pydantic import Field @@ -140,7 +141,6 @@ class PdfEditParameterSelector: class PdfEditAgent: def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime - self.supported_operations = list(OPERATIONS) self.parameter_selector = PdfEditParameterSelector(runtime) async def orchestrate(self, request: OrchestratorRequest) -> PdfEditResponse: @@ -156,6 +156,7 @@ class PdfEditAgent: files=request.files, conversation_history=request.conversation_history, page_text=extracted_text.files if extracted_text is not None else [], + enabled_endpoints=request.enabled_endpoints, ) ) @@ -165,19 +166,37 @@ class PdfEditAgent: async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse: ... async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse: logger.info( - "[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s msg=%r", + "[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s enabled=%s msg=%r", [file.name for file in request.files], has_page_text(request.page_text), allow_need_content, + request.enabled_endpoints, request.user_message, ) - selection = await self._select_plan(request, allow_need_content=allow_need_content) + supported_operations = self._get_supported_operations(request) + unavailable_operations = self._get_unavailable_operations(supported_operations) + if not supported_operations: + return EditCannotDoResponse(reason="No PDF edit operations are available on this server.") + selection = await self._select_plan( + request, supported_operations, unavailable_operations, allow_need_content=allow_need_content + ) if isinstance(selection, EditClarificationRequest | EditCannotDoResponse): logger.info("[pdf-edit] selection -> %s: %s", selection.outcome, Pretty(selection)) return selection if isinstance(selection, NeedContentResponse): logger.info("[pdf-edit] selection -> need_content: %s", selection.reason) return self._fill_need_content_defaults(selection, request) + enabled = set(supported_operations) + unsupported = [op for op in selection.operations if op not in enabled] + if unsupported: + logger.warning("[pdf-edit] plan referenced unavailable operations: %s", [op.name for op in unsupported]) + return EditCannotDoResponse( + reason=( + "The following operations are not available on this server " + "(either disabled by the administrator or not installed): " + + ", ".join(op.name for op in unsupported) + ) + ) logger.info("[pdf-edit] plan: %s", [op.name for op in selection.operations]) steps: list[ToolOperationStep] = [] for operation_index, operation_id in enumerate(selection.operations): @@ -202,18 +221,39 @@ class PdfEditAgent: async def _select_plan( self, request: PdfEditRequest, + supported_operations: Iterable[ToolEndpoint], + unavailable_operations: Iterable[ToolEndpoint], + *, allow_need_content: bool = True, ) -> PdfEditPlanOutput: can_request_content = allow_need_content and not has_page_text(request.page_text) - agent = self._build_selection_agent(allow_need_content=can_request_content) - return await agent.select(self._build_selection_prompt(request)) + agent = self._build_selection_agent( + supported_operations, unavailable_operations, allow_need_content=can_request_content + ) + return await agent.select(self._build_selection_prompt(request, supported_operations, unavailable_operations)) - def _build_selection_agent(self, *, allow_need_content: bool) -> PdfEditSelectionAgent: + def _build_selection_agent( + self, + supported_operations: Iterable[ToolEndpoint], + unavailable_operations: Iterable[ToolEndpoint], + *, + allow_need_content: bool, + ) -> PdfEditSelectionAgent: + unavailable_clause = ( + f" The following operations exist on this server but are NOT currently available " + f"(disabled by the administrator or not installed in this build) and must NOT appear " + f"in any plan: {self._get_operations_prompt(unavailable_operations)}. " + "If the user asks for one of these, return cannot_do, name the operation, and explain " + "that it exists but isn't available on this server." + if unavailable_operations + else "" + ) return PdfEditSelectionAgent( self.runtime, base_system_prompt=( "Plan PDF edit requests. " - f"Supported operations are: {self._supported_operations_prompt()}. " + f"Supported operations are: {self._get_operations_prompt(supported_operations)}." + f"{unavailable_clause} " "Return an ordered list of one or more supported operations for the plan. " "Do not produce operation parameters in this stage. " "Return need_clarification when the request is genuinely ambiguous. " @@ -224,17 +264,38 @@ class PdfEditAgent: allow_need_content=allow_need_content, ) - def _build_selection_prompt(self, request: PdfEditRequest) -> str: + def _build_selection_prompt( + self, + request: PdfEditRequest, + supported_operations: Iterable[ToolEndpoint], + unavailable_operations: Iterable[ToolEndpoint], + ) -> str: + unavailable_line = ( + "Unavailable operations (exist but not currently usable): " + f"{self._get_operations_prompt(unavailable_operations)}\n" + if unavailable_operations + else "" + ) return ( f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n" f"User request: {request.user_message}\n" f"Files: {format_file_names(request.files)}\n" - f"Supported operations: {self._supported_operations_prompt()}\n" + f"Supported operations: {self._get_operations_prompt(supported_operations)}\n" + f"{unavailable_line}" f"Extracted page text:\n{format_page_text(request.page_text)}" ) - def _supported_operations_prompt(self) -> str: - return ", ".join(f"{op.name} ({op.value})" for op in self.supported_operations) + def _get_supported_operations(self, request: PdfEditRequest) -> Iterable[ToolEndpoint]: + return request.enabled_endpoints + + @staticmethod + def _get_unavailable_operations(supported_operations: Iterable[ToolEndpoint]) -> Iterable[ToolEndpoint]: + supported_set = set(supported_operations) + return [op for op in OPERATIONS if op not in supported_set] + + @staticmethod + def _get_operations_prompt(operations: Iterable[ToolEndpoint]) -> str: + return ", ".join(f"{op.name} ({op.value})" for op in operations) def _fill_need_content_defaults( self, diff --git a/engine/src/stirling/contracts/common.py b/engine/src/stirling/contracts/common.py index 21e618ab1..73df8b5b2 100644 --- a/engine/src/stirling/contracts/common.py +++ b/engine/src/stirling/contracts/common.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Iterable from enum import StrEnum from typing import Literal, assert_never @@ -197,3 +198,14 @@ class ToolOperationStep(ApiModel): actual_type = type(self.parameters).__name__ raise ValueError(f"Parameters for tool {self.tool} must be {expected_type.__name__}, got {actual_type}.") return self + + +def drop_unknown_tool_endpoints(value: Iterable[str | ToolEndpoint]) -> list[ToolEndpoint]: + """Coerce inbound endpoint identifiers into `ToolEndpoint` members, dropping unknowns. + + Java sends the full set of endpoints it considers enabled. The engine and the Java + backend may have version drift in either direction, so we silently drop anything we + don't recognise rather than failing the request. Anything dropped simply doesn't + appear as a supported tool to the planner. + """ + return [ToolEndpoint(item) for item in value if item in ToolEndpoint] diff --git a/engine/src/stirling/contracts/orchestrator.py b/engine/src/stirling/contracts/orchestrator.py index dc7c1f152..9972d043b 100644 --- a/engine/src/stirling/contracts/orchestrator.py +++ b/engine/src/stirling/contracts/orchestrator.py @@ -2,9 +2,9 @@ from __future__ import annotations from typing import Annotated, Literal -from pydantic import Field +from pydantic import BeforeValidator, Field -from stirling.models import ApiModel +from stirling.models import ApiModel, ToolEndpoint from .agent_drafts import AgentDraftResponse from .common import ( @@ -17,6 +17,7 @@ from .common import ( SupportedCapability, ToolReportArtifact, WorkflowOutcome, + drop_unknown_tool_endpoints, ) from .execution import NextExecutionAction from .pdf_edit import PdfEditTerminalResponse @@ -37,6 +38,10 @@ class OrchestratorRequest(ApiModel): conversation_history: list[ConversationMessage] = Field(default_factory=list) artifacts: list[WorkflowArtifact] = Field(default_factory=list) resume_with: SupportedCapability | None = None + # See `PdfEditRequest.enabled_endpoints`. + enabled_endpoints: Annotated[list[ToolEndpoint], BeforeValidator(drop_unknown_tool_endpoints)] = Field( + default_factory=list + ) class UnsupportedCapabilityResponse(ApiModel): diff --git a/engine/src/stirling/contracts/pdf_edit.py b/engine/src/stirling/contracts/pdf_edit.py index be79bf563..35665ce64 100644 --- a/engine/src/stirling/contracts/pdf_edit.py +++ b/engine/src/stirling/contracts/pdf_edit.py @@ -2,9 +2,9 @@ from __future__ import annotations from typing import Annotated, Literal -from pydantic import Field +from pydantic import BeforeValidator, Field -from stirling.models import ApiModel +from stirling.models import ApiModel, ToolEndpoint from .common import ( AiFile, @@ -14,6 +14,7 @@ from .common import ( SupportedCapability, ToolOperationStep, WorkflowOutcome, + drop_unknown_tool_endpoints, ) @@ -22,6 +23,13 @@ class PdfEditRequest(ApiModel): files: list[AiFile] = Field(default_factory=list) conversation_history: list[ConversationMessage] = Field(default_factory=list) page_text: list[ExtractedFileText] = Field(default_factory=list) + # The set of endpoints the Java backend considers usable. Unknown URLs are silently + # dropped so the engine and Java can drift in either direction without breaking + # validation. An empty list means no operations are available - the planner will + # return `cannot_do`. + enabled_endpoints: Annotated[list[ToolEndpoint], BeforeValidator(drop_unknown_tool_endpoints)] = Field( + default_factory=list + ) class EditPlanResponse(ApiModel): diff --git a/engine/tests/test_pdf_edit_agent.py b/engine/tests/test_pdf_edit_agent.py index dd9bf7f51..a40b671db 100644 --- a/engine/tests/test_pdf_edit_agent.py +++ b/engine/tests/test_pdf_edit_agent.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Iterable from dataclasses import dataclass import pytest @@ -20,7 +21,7 @@ from stirling.contracts import ( SupportedCapability, ToolOperationStep, ) -from stirling.models import FileId +from stirling.models import OPERATIONS, FileId from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint from stirling.services.runtime import AppRuntime @@ -69,9 +70,19 @@ class StubPdfEditAgent(PdfEditAgent): if parameter_selector is not None: self.parameter_selector = parameter_selector + def _get_supported_operations(self, request: PdfEditRequest) -> Iterable[ToolEndpoint]: + # Tests construct requests without `enabled_endpoints`; pretend everything is enabled + # unless the test explicitly supplies an enabled set. + if request.enabled_endpoints: + return request.enabled_endpoints + return OPERATIONS + async def _select_plan( self, request: PdfEditRequest, + supported_operations: Iterable[ToolEndpoint], + unavailable_operations: Iterable[ToolEndpoint], + *, allow_need_content: bool = True, ) -> PdfEditPlanOutput: return self.selection @@ -207,14 +218,20 @@ async def test_pdf_edit_agent_builds_selection_agent_matching_content_availabili agent = PdfEditAgent(runtime) captured: list[bool] = [] - def record(*, allow_need_content: bool) -> PdfEditSelectionAgent: + def record( + supported_operations: Iterable[ToolEndpoint], + unavailable_operations: Iterable[ToolEndpoint], + *, + allow_need_content: bool, + ) -> PdfEditSelectionAgent: captured.append(allow_need_content) raise _StopSelectionError() agent._build_selection_agent = record + supported = list(OPERATIONS) with pytest.raises(_StopSelectionError): - await agent._select_plan(PdfEditRequest(user_message="Rotate.")) + await agent._select_plan(PdfEditRequest(user_message="Rotate."), supported, []) with pytest.raises(_StopSelectionError): await agent._select_plan( PdfEditRequest( @@ -225,10 +242,12 @@ async def test_pdf_edit_agent_builds_selection_agent_matching_content_availabili pages=[PdfTextSelection(page_number=1, text="content")], ) ], - ) + ), + supported, + [], ) with pytest.raises(_StopSelectionError): - await agent._select_plan(PdfEditRequest(user_message="Rotate."), allow_need_content=False) + await agent._select_plan(PdfEditRequest(user_message="Rotate."), supported, [], allow_need_content=False) assert captured == [True, False, False] @@ -282,3 +301,95 @@ async def test_pdf_edit_agent_passes_page_text_to_parameter_selector(runtime: Ap ) assert parameter_selector.calls[0].request.page_text == page_text + + +def test_pdf_edit_request_drops_unknown_enabled_urls() -> None: + request = PdfEditRequest.model_validate( + { + "user_message": "ignore me", + "enabled_endpoints": [ + ToolEndpoint.COMPRESS_PDF.value, + "/api/v1/not-a-real/endpoint", + ], + } + ) + + assert request.enabled_endpoints == [ToolEndpoint.COMPRESS_PDF] + + +@pytest.mark.anyio +async def test_pdf_edit_agent_supported_operations_defaults_to_empty( + runtime: AppRuntime, +) -> None: + agent = PdfEditAgent(runtime) + supported = agent._get_supported_operations(PdfEditRequest(user_message="hi")) + + assert list(supported) == [] + + +@pytest.mark.anyio +async def test_pdf_edit_agent_supported_operations_uses_provided_list( + runtime: AppRuntime, +) -> None: + agent = PdfEditAgent(runtime) + request = PdfEditRequest( + user_message="Compress this PDF.", + enabled_endpoints=[ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF], + ) + + supported = agent._get_supported_operations(request) + + assert list(supported) == [ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF] + + +@pytest.mark.anyio +async def test_pdf_edit_agent_returns_cannot_do_when_no_operations_enabled( + runtime: AppRuntime, +) -> None: + agent = PdfEditAgent(runtime) + response = await agent.handle(PdfEditRequest(user_message="Do anything.", enabled_endpoints=[])) + + assert isinstance(response, EditCannotDoResponse) + + +def test_pdf_edit_selection_prompt_includes_unavailable_operations(runtime: AppRuntime) -> None: + agent = PdfEditAgent(runtime) + request = PdfEditRequest( + user_message="Run OCR.", + enabled_endpoints=[ToolEndpoint.FLATTEN], + ) + supported = agent._get_supported_operations(request) + unavailable = agent._get_unavailable_operations(supported) + + prompt = agent._build_selection_prompt(request, supported, unavailable) + + assert "Unavailable operations" in prompt + assert "OCR_PDF" in prompt + assert ToolEndpoint.OCR_PDF.value in prompt + + +@pytest.mark.anyio +async def test_pdf_edit_agent_rejects_plan_referencing_unavailable_operations( + runtime: AppRuntime, +) -> None: + parameter_selector = RecordingParameterSelector() + agent = StubPdfEditAgent( + runtime, + PdfEditPlanSelection( + operations=[ToolEndpoint.COMPRESS_PDF], + summary="Compress.", + ), + parameter_selector=parameter_selector, + ) + + response = await agent.handle( + PdfEditRequest( + user_message="Compress this PDF.", + enabled_endpoints=[ToolEndpoint.FLATTEN], + ) + ) + + assert isinstance(response, EditCannotDoResponse) + assert "not available" in response.reason + assert "COMPRESS_PDF" in response.reason + assert parameter_selector.calls == []