fix: Core sets parent_deployment to a wrong value #530 (#604)

server/src/main/java/com/epam/aidial/core/server/log/GfLogStore.java

@@ -14,6 +14,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.annotations.VisibleForTesting;
 import io.netty.buffer.ByteBufInputStream;
 import io.vertx.core.Vertx;
 import io.vertx.core.buffer.Buffer;
@@ -105,13 +106,10 @@ private void append(ProxyContext context, LogEntry entry) throws JsonProcessingE
         append(entry, context.getDeployment().getName(), true);
         append(entry, "\"", false);
 
-        String sourceDeployment = context.getSourceDeployment();
-        if (sourceDeployment != null) {
-            String initialDeployment = context.getInitialDeployment();
+        String parentDeployment = getParentDeployment(context);
+        if (parentDeployment != null) {
             append(entry, ",\"parent_deployment\":\"", false);
-            // if deployment(callee) is configured with interceptors the return initial deployment(which triggers interceptors)
-            // otherwise return source deployment(caller)
-            append(entry, initialDeployment != null ? initialDeployment : sourceDeployment, true);
+            append(entry, parentDeployment, true);
             append(entry, "\"", false);
         }
 
@@ -341,4 +339,24 @@ static boolean isStreamingResponse(@Nullable Buffer response) {
         }
         return j == dataToken.length();
     }
+
+    @VisibleForTesting
+    static String getParentDeployment(ProxyContext context) {
+        List<String> interceptors = context.getInterceptors();
+        if (interceptors == null) {
+            return context.getSourceDeployment();
+        }
+        // skip interceptors and return the deployment which called the current one
+        List<String> executionPath = context.getExecutionPath();
+        int i = executionPath.size() - 2;
+        for (int j = interceptors.size() - 1; i >= 0 && j >= 0; i--, j--) {
+            String deployment = executionPath.get(i);
+            String interceptor = interceptors.get(j);
+            if (!deployment.equals(interceptor)) {
+                log.warn("Can't find parent deployment because interceptor path doesn't match: expected - {}, actual - {}", interceptor, deployment);
+                return null;
+            }
+        }
+        return i < 0 ? null : executionPath.get(i);
+    }
 }

server/src/test/java/com/epam/aidial/core/server/log/GfLogStoreTest.java

@@ -1,12 +1,18 @@
 package com.epam.aidial.core.server.log;
 
+import com.epam.aidial.core.server.ProxyContext;
 import io.vertx.core.buffer.Buffer;
 import org.junit.jupiter.api.Test;
 
+import java.util.List;
+
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 @SuppressWarnings("checkstyle:LineLength")
 public class GfLogStoreTest {
@@ -192,4 +198,58 @@ public void testAssembleStreamingResponse2() {
 
 
     }
+
+    @Test
+    public void testGetParentDeployment_NoInterceptors() {
+        ProxyContext context = mock(ProxyContext.class);
+        // app calls model without interceptors
+        when(context.getInterceptors()).thenReturn(null);
+        when(context.getSourceDeployment()).thenReturn("app");
+
+        String result = GfLogStore.getParentDeployment(context);
+
+        assertEquals("app", result);
+    }
+
+    @Test
+    public void testGetParentDeployment_DeploymentWithInterceptors1() {
+        ProxyContext context = mock(ProxyContext.class);
+        // app calls model with interceptors
+        List<String> interceptors = List.of("interceptor1", "interceptor2");
+        when(context.getInterceptors()).thenReturn(interceptors);
+        List<String> executionPath = List.of("app", "interceptor1", "interceptor2", "model");
+        when(context.getExecutionPath()).thenReturn(executionPath);
+
+        String result = GfLogStore.getParentDeployment(context);
+
+        assertEquals("app", result);
+    }
+
+    @Test
+    public void testGetParentDeployment_DeploymentWithInterceptors2() {
+        ProxyContext context = mock(ProxyContext.class);
+        // chat calls model with interceptors
+        List<String> interceptors = List.of("interceptor1", "interceptor2");
+        when(context.getInterceptors()).thenReturn(interceptors);
+        List<String> executionPath = List.of("interceptor1", "interceptor2", "model");
+        when(context.getExecutionPath()).thenReturn(executionPath);
+
+        String result = GfLogStore.getParentDeployment(context);
+
+        assertNull(result);
+    }
+
+    @Test
+    public void testGetParentDeployment_InterceptorPathMismatch() {
+        ProxyContext context = mock(ProxyContext.class);
+        // app calls model with interceptors but interceptor1 calls some dep1 in the middle using the same per request key
+        List<String> interceptors = List.of("interceptor1", "interceptor2");
+        when(context.getInterceptors()).thenReturn(interceptors);
+        List<String> executionPath = List.of("app", "interceptor1", "dep1", "interceptor2", "model");
+        when(context.getExecutionPath()).thenReturn(executionPath);
+
+        String result = GfLogStore.getParentDeployment(context);
+
+        assertNull(result);
+    }
 }