feat: Add server-specific files to the request for application with schema while proxying the completion request

server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java

@@ -12,6 +12,7 @@
 import com.epam.aidial.core.server.data.ApiKeyData;
 import com.epam.aidial.core.server.data.ErrorData;
 import com.epam.aidial.core.server.function.BaseRequestFunction;
+import com.epam.aidial.core.server.function.CollectRequestApplicationFilesFn;
 import com.epam.aidial.core.server.function.CollectRequestAttachmentsFn;
 import com.epam.aidial.core.server.function.CollectRequestDataFn;
 import com.epam.aidial.core.server.function.CollectResponseAttachmentsFn;
@@ -73,7 +74,8 @@ public DeploymentPostController(Proxy proxy, ProxyContext context) {
                 new ApplyDefaultDeploymentSettingsFn(proxy, context),
                 new EnhanceAssistantRequestFn(proxy, context),
                 new EnhanceModelRequestFn(proxy, context),
-                new AppendApplicationPropertiesFn(proxy, context));
+                new AppendApplicationPropertiesFn(proxy, context),
+                new CollectRequestApplicationFilesFn(proxy, context));
     }
 
     public Future<?> handle(String deploymentId, String deploymentApi) {
@@ -267,6 +269,7 @@ void handleRequestBody(Buffer requestBody) {
             if (ProxyUtil.processChain(tree, enhancementFunctions)) {
                 context.setRequestBody(Buffer.buffer(ProxyUtil.MAPPER.writeValueAsBytes(tree)));
             }
+            proxy.getApiKeyStore().assignPerRequestApiKey(context.getProxyApiKeyData());
         } catch (Throwable e) {
             if (e instanceof HttpException httpException) {
                 respond(httpException.getStatus(), httpException.getMessage());

server/src/main/java/com/epam/aidial/core/server/controller/ResourceController.java

@@ -12,11 +12,11 @@
 import com.epam.aidial.core.server.service.ApplicationService;
 import com.epam.aidial.core.server.service.PermissionDeniedException;
 import com.epam.aidial.core.server.service.ResourceNotFoundException;
-import com.epam.aidial.core.server.service.ShareService;
 import com.epam.aidial.core.server.util.ApplicationTypeSchemaProcessingException;
 import com.epam.aidial.core.server.util.ApplicationTypeSchemaUtils;
 import com.epam.aidial.core.server.util.ProxyUtil;
 import com.epam.aidial.core.server.util.ResourceDescriptorFactory;
+import com.epam.aidial.core.server.validation.ApplicationTypeResourceException;
 import com.epam.aidial.core.server.validation.ApplicationTypeSchemaValidationException;
 import com.epam.aidial.core.storage.data.MetadataBase;
 import com.epam.aidial.core.storage.data.ResourceItemMetadata;
@@ -37,6 +37,7 @@
 import java.util.List;
 
 import static com.epam.aidial.core.storage.http.HttpStatus.BAD_REQUEST;
+import static com.epam.aidial.core.storage.http.HttpStatus.FORBIDDEN;
 import static com.epam.aidial.core.storage.http.HttpStatus.INTERNAL_SERVER_ERROR;
 
 @Slf4j
@@ -130,11 +131,9 @@ private Future<?> getResource(ResourceDescriptor descriptor, boolean hasWriteAcc
         Future<Pair<ResourceItemMetadata, String>> responseFuture = (descriptor.getType() == ResourceTypes.APPLICATION)
                 ? getApplicationData(descriptor, hasWriteAccess, etagHeader) : getResourceData(descriptor, etagHeader);
 
-        responseFuture.onSuccess(pair -> {
-                    context.putHeader(HttpHeaders.ETAG, pair.getKey().getEtag())
-                            .exposeHeaders()
-                            .respond(HttpStatus.OK, pair.getValue());
-                })
+        responseFuture.onSuccess(pair -> context.putHeader(HttpHeaders.ETAG, pair.getKey().getEtag())
+                .exposeHeaders()
+                .respond(HttpStatus.OK, pair.getValue()))
                 .onFailure(error -> handleError(descriptor, error));
 
         return Future.succeededFuture();
@@ -177,7 +176,9 @@ private void validateCustomApplication(Application application) {
                         throw new HttpException(BAD_REQUEST, "No read access to file: " + file.getUrl());
                     });
         } catch (ValidationException | IllegalArgumentException | ApplicationTypeSchemaValidationException e) {
-            throw new HttpException(BAD_REQUEST, " Custom application validation failed", e);
+            throw new HttpException(BAD_REQUEST, "Custom application validation failed", e);
+        } catch (ApplicationTypeResourceException e) {
+            throw new HttpException(FORBIDDEN, "Failed to access application resource " + e.getResourceUri(), e);
         } catch (ApplicationTypeSchemaProcessingException e) {
             throw new HttpException(INTERNAL_SERVER_ERROR, "Custom application processing exception", e);
         }
@@ -232,11 +233,9 @@ private Future<?> putResource(ResourceDescriptor descriptor) {
             });
         }
 
-        responseFuture.onSuccess((metadata) -> {
-                    context.putHeader(HttpHeaders.ETAG, metadata.getEtag())
-                            .exposeHeaders()
-                            .respond(HttpStatus.OK, metadata);
-                })
+        responseFuture.onSuccess((metadata) -> context.putHeader(HttpHeaders.ETAG, metadata.getEtag())
+                .exposeHeaders()
+                .respond(HttpStatus.OK, metadata))
                 .onFailure(error -> handleError(descriptor, error));
 
         return Future.succeededFuture();

server/src/main/java/com/epam/aidial/core/server/function/CollectRequestApplicationFilesFn.java

@@ -0,0 +1,67 @@
+package com.epam.aidial.core.server.function;
+
+import com.epam.aidial.core.config.Application;
+import com.epam.aidial.core.config.Deployment;
+import com.epam.aidial.core.server.Proxy;
+import com.epam.aidial.core.server.ProxyContext;
+import com.epam.aidial.core.server.data.ApiKeyData;
+import com.epam.aidial.core.server.data.AutoSharedData;
+import com.epam.aidial.core.server.security.AccessService;
+import com.epam.aidial.core.server.util.ApplicationTypeSchemaUtils;
+import com.epam.aidial.core.server.validation.ApplicationTypeResourceException;
+import com.epam.aidial.core.storage.data.ResourceAccessType;
+import com.epam.aidial.core.storage.http.HttpException;
+import com.epam.aidial.core.storage.http.HttpStatus;
+import com.epam.aidial.core.storage.resource.ResourceDescriptor;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.List;
+
+
+@Slf4j
+public class CollectRequestApplicationFilesFn extends BaseRequestFunction<ObjectNode> {
+    public CollectRequestApplicationFilesFn(Proxy proxy, ProxyContext context) {
+        super(proxy, context);
+    }
+
+    @Override
+    public Boolean apply(ObjectNode tree) {
+        try {
+            Deployment deployment = context.getDeployment();
+            if (!(deployment instanceof Application application && application.getApplicationTypeSchemaId() != null)) {
+                return false;
+            }
+            if (application.getApplicationProperties() == null) {
+                throw new HttpException(HttpStatus.INTERNAL_SERVER_ERROR, "Typed application's properties not set");
+            }
+            List<ResourceDescriptor> resources = ApplicationTypeSchemaUtils.getServerFiles(context.getConfig(), application, proxy.getEncryptionService(),
+                    proxy.getResourceService());
+            appendFilesToProxyApiKeyData(resources);
+            return false;
+        } catch (HttpException ex) {
+            throw ex;
+        } catch (ApplicationTypeResourceException ex) {
+            throw new HttpException(HttpStatus.FORBIDDEN, ex.getMessage());
+        } catch (Exception e) {
+            throw new HttpException(HttpStatus.INTERNAL_SERVER_ERROR, e.getMessage());
+        }
+    }
+
+    private void appendFilesToProxyApiKeyData(List<ResourceDescriptor> resources) {
+        ApiKeyData apiKeyData = context.getProxyApiKeyData();
+        for (ResourceDescriptor resource : resources) {
+            String resourceUrl = resource.getUrl();
+            AccessService accessService = proxy.getAccessService();
+            if (accessService.hasReadAccess(resource, context)) {
+                if (resource.isFolder()) {
+                    apiKeyData.getAttachedFolders().put(resourceUrl, new AutoSharedData(ResourceAccessType.READ_ONLY));
+                } else {
+                    apiKeyData.getAttachedFiles().put(resourceUrl, new AutoSharedData(ResourceAccessType.READ_ONLY));
+                }
+            } else {
+                throw new HttpException(HttpStatus.FORBIDDEN, "Access denied to the file %s".formatted(resourceUrl));
+            }
+        }
+    }
+}

server/src/main/java/com/epam/aidial/core/server/function/CollectRequestAttachmentsFn.java

@@ -28,9 +28,6 @@ public CollectRequestAttachmentsFn(Proxy proxy, ProxyContext context) {
     @Override
     public Boolean apply(ObjectNode tree) {
         ProxyUtil.collectAttachedFilesFromRequest(tree, this::processAttachedFile);
-        // assign api key data after processing attachments
-        ApiKeyData destApiKeyData = context.getProxyApiKeyData();
-        proxy.getApiKeyStore().assignPerRequestApiKey(destApiKeyData);
         return false;
     }
 

server/src/main/java/com/epam/aidial/core/server/util/ApplicationTypeSchemaUtils.java

@@ -4,6 +4,7 @@
 import com.epam.aidial.core.config.Config;
 import com.epam.aidial.core.server.ProxyContext;
 import com.epam.aidial.core.server.security.EncryptionService;
+import com.epam.aidial.core.server.validation.ApplicationTypeResourceException;
 import com.epam.aidial.core.server.validation.ApplicationTypeSchemaValidationException;
 import com.epam.aidial.core.server.validation.DialFileKeyword;
 import com.epam.aidial.core.server.validation.DialMetaKeyword;
@@ -154,8 +155,19 @@ public static void replaceCustomAppFiles(Application application, Map<String, St
         application.setApplicationProperties(customPropertiesMap);
     }
 
+    public static List<ResourceDescriptor> getServerFiles(Config config, Application application, EncryptionService encryptionService,
+                                                          ResourceService resourceService) {
+        return getFiles(config, application, encryptionService, resourceService, ListCollector.FileCollectorType.ONLY_SERVER_FILES);
+    }
+
+    public static List<ResourceDescriptor> getFiles(Config config, Application application, EncryptionService encryptionService,
+                                                    ResourceService resourceService) {
+        return getFiles(config, application, encryptionService, resourceService, ListCollector.FileCollectorType.ALL_FILES);
+    }
+
     @SuppressWarnings("unchecked")
-    public static List<ResourceDescriptor> getFiles(Config config, Application application, EncryptionService encryptionService, ResourceService resourceService) {
+    private static List<ResourceDescriptor> getFiles(Config config, Application application, EncryptionService encryptionService,
+                                                    ResourceService resourceService, ListCollector.FileCollectorType collectorName) {
         try {
             String customApplicationSchema = getCustomApplicationSchemaOrThrow(config, application);
             if (customApplicationSchema == null) {
@@ -169,7 +181,7 @@ public static List<ResourceDescriptor> getFiles(Config config, Application appli
             if (!validationResult.isEmpty()) {
                 throw new ApplicationTypeSchemaValidationException("Failed to validate custom app against the schema", validationResult);
             }
-            ListCollector<String> propsCollector = (ListCollector<String>) collectorContext.getCollectorMap().get("file");
+            ListCollector<String> propsCollector = (ListCollector<String>) collectorContext.getCollectorMap().get(collectorName.getValue());
             if (propsCollector == null) {
                 return Collections.emptyList();
             }
@@ -178,15 +190,15 @@ public static List<ResourceDescriptor> getFiles(Config config, Application appli
                 try {
                     ResourceDescriptor descriptor = ResourceDescriptorFactory.fromAnyUrl(item, encryptionService);
                     if (!descriptor.isFolder() && !resourceService.hasResource(descriptor)) {
-                        throw new ApplicationTypeSchemaValidationException("Resource listed as dependent to the application not found or inaccessible: " + item);
+                        throw new ApplicationTypeResourceException("Resource listed as dependent to the application not found or inaccessible", item);
                     }
                     result.add(descriptor);
                 } catch (IllegalArgumentException e) {
-                    throw new ApplicationTypeSchemaValidationException("Failed to get resource descriptor for url: " + item, e);
+                    throw new ApplicationTypeResourceException("Failed to get resource descriptor for url", item, e);
                 }
             }
             return result;
-        } catch (ApplicationTypeSchemaValidationException e) {
+        } catch (ApplicationTypeSchemaValidationException | ApplicationTypeResourceException e) {
             throw e;
         } catch (Exception e) {
             throw new ApplicationTypeSchemaProcessingException("Failed to obtain list of files attached to the custom app", e);

server/src/main/java/com/epam/aidial/core/server/validation/ApplicationTypeResourceException.java

@@ -0,0 +1,26 @@
+package com.epam.aidial.core.server.validation;
+
+import lombok.Getter;
+
+
+/**
+ * Exception thrown when there is an issue with a resource associated with an application type.
+ * This exception is typically used when a resource listed as dependent to the application
+ * is not found, inaccessible, or there is a failure in obtaining the resource descriptor.
+ */
+@Getter
+public class ApplicationTypeResourceException extends RuntimeException {
+
+    private final String resourceUri;
+
+    public ApplicationTypeResourceException(String message, String resourceUri, Throwable cause) {
+        super(message, cause);
+        this.resourceUri = resourceUri;
+    }
+
+    public ApplicationTypeResourceException(String message, String resourceUri) {
+        super(message);
+        this.resourceUri = resourceUri;
+    }
+
+}

server/src/main/java/com/epam/aidial/core/server/validation/ApplicationTypeSchemaValidationException.java

@@ -14,10 +14,6 @@ public ApplicationTypeSchemaValidationException(String message, Set<ValidationMe
         this.validationMessages = validationMessages;
     }
 
-    public ApplicationTypeSchemaValidationException(String message, Throwable cause) {
-        super(message, cause);
-    }
-
     public ApplicationTypeSchemaValidationException(String message) {
         super(message);
     }

server/src/main/java/com/epam/aidial/core/server/validation/DialFileKeyword.java

@@ -33,22 +33,31 @@ private static class DialFileCollectorValidator extends BaseJsonValidator {
         private static final ErrorMessageType ERROR_MESSAGE_TYPE = () -> "dial:file";
 
         private final Boolean value;
+        private final Boolean isServerProp;
 
         public DialFileCollectorValidator(SchemaLocation schemaLocation, JsonNodePath evaluationPath, JsonNode schemaNode,
                                           JsonSchema parentSchema, Keyword keyword,
                                           ValidationContext validationContext, boolean suppressSubSchemaRetrieval) {
             super(schemaLocation, evaluationPath, schemaNode, parentSchema, ERROR_MESSAGE_TYPE, keyword, validationContext, suppressSubSchemaRetrieval);
             this.value = schemaNode.booleanValue();
+            JsonNode metaNode = parentSchema.getSchemaNode().get("dial:meta");
+            JsonNode propertyKindNode = (metaNode != null) ? metaNode.get("dial:propertyKind") : null;
+            this.isServerProp = (propertyKindNode != null) && propertyKindNode.asText().equalsIgnoreCase("server");
         }
 
         @Override
         @SuppressWarnings("unchecked")
         public Set<ValidationMessage> validate(ExecutionContext executionContext, JsonNode jsonNode, JsonNode jsonNode1, JsonNodePath jsonNodePath) {
             if (value) {
                 CollectorContext collectorContext = executionContext.getCollectorContext();
-                ListCollector<String> serverPropsCollector = (ListCollector<String>) collectorContext.getCollectorMap()
-                        .computeIfAbsent("file", k -> new ListCollector<String>());
-                serverPropsCollector.combine(List.of(jsonNode.asText()));
+                ListCollector<String> fileCollector = (ListCollector<String>) collectorContext.getCollectorMap()
+                        .computeIfAbsent(ListCollector.FileCollectorType.ALL_FILES.getValue(), k -> new ListCollector<String>());
+                fileCollector.combine(List.of(jsonNode.asText()));
+                if (isServerProp) {
+                    ListCollector<String> serverFileCollector = (ListCollector<String>) collectorContext.getCollectorMap()
+                            .computeIfAbsent(ListCollector.FileCollectorType.ONLY_SERVER_FILES.getValue(), k -> new ListCollector<String>());
+                    serverFileCollector.combine(List.of(jsonNode.asText()));
+                }
             }
             return Set.of();
         }

server/src/main/java/com/epam/aidial/core/server/validation/ListCollector.java

@@ -1,11 +1,25 @@
 package com.epam.aidial.core.server.validation;
 
 import com.networknt.schema.Collector;
+import lombok.Getter;
 
 import java.util.ArrayList;
 import java.util.List;
 
 public class ListCollector<T> implements Collector<List<T>> {
+
+    @Getter
+    public enum FileCollectorType {
+        ALL_FILES("file"),
+        ONLY_SERVER_FILES("server_file");
+
+        private final String value;
+
+        FileCollectorType(String value) {
+            this.value = value;
+        }
+    }
+
     private final List<T> references = new ArrayList<>();
 
     @Override

server/src/test/java/com/epam/aidial/core/server/ResourceApiTest.java

@@ -403,7 +403,7 @@ void testApplicationWithTypeSchemaCreation_Failed_FailAccessFile() {
                        "description": "My application description"
                   }
                 """);
-        Assertions.assertEquals(400, response.status());
+        Assertions.assertEquals(403, response.status());
     }
 
     @Test