fix: NPE in streaming chat-completions function chain #1963

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

@@ -26,6 +26,10 @@ public CollectChatCompletionUsageFn(Proxy proxy, ProxyContext context) {
 
     @Override
     public Future<JsonNode> apply(JsonNode tree) {
+        if (tree == null) {
+            return Future.succeededFuture(tree);
+        }
+
         usage = MergeChunks.merge(usage, tree.get("usage"));
         if (tree.get("service_tier") != null) {
             serviceTier = tree.get("service_tier");

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

@@ -38,7 +38,7 @@ public Future<JsonNode> apply(JsonNode tree) {
                 processAttachedFile(attachment, permittedAttachments);
             }
             if (permittedAttachments.isEmpty()) {
-                return Future.succeededFuture();
+                return Future.succeededFuture(tree);
             }
             return proxy.getTaskExecutor().submit(() -> {
                 proxy.getApiKeyStore().updatePerRequestApiKey(perRequestKey, json -> updateAutoSharedAttachments(json, permittedAttachments, perRequestKey));

server/src/main/java/com/epam/aidial/core/server/vertx/stream/BufferingReadStream.java

@@ -267,10 +267,10 @@ public void onComplete() {
         @SneakyThrows
         private Future<Void> handle(SseEvent event) {
             Future<JsonNode> result;
-            if (functions.isEmpty() || skipEvent(event)) {
+            String data = event.getData();
+            if (functions.isEmpty() || skipEvent(event) || data == null || data.isBlank()) {
                 result = Future.succeededFuture();
             } else {
-                String data = event.getData();
                 try {
                     JsonNode tree = ProxyUtil.MAPPER.readTree(data);
                     result = Future.succeededFuture(tree);

server/src/test/java/com/epam/aidial/core/server/function/ChatCompletionFunctionChainTest.java

@@ -0,0 +1,75 @@
+package com.epam.aidial.core.server.function;
+
+import com.epam.aidial.core.server.ProxyContext;
+import com.epam.aidial.core.server.data.ApiKeyData;
+import com.epam.aidial.core.server.util.ProxyUtil;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import io.vertx.core.Future;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doCallRealMethod;
+import static org.mockito.Mockito.when;
+
+/**
+ * Exercises the real function chain wired by DeploymentPostController for streaming
+ * {@code /chat/completions} responses (StripUsagePerModelFn -> CollectResponseChatCompletionAttachmentsFn
+ * -> CollectChatCompletionUsageFn), composed the same way {@code BufferingReadStream.BaseEventListener}
+ * chains it via {@code Future.compose}. Regression test for a null tree silently produced by one
+ * function and blindly dereferenced by the next.
+ */
+@ExtendWith(MockitoExtension.class)
+class ChatCompletionFunctionChainTest {
+
+    @Mock
+    private ProxyContext context;
+
+    @Test
+    public void testChainToleratesChunksWithNoAttachmentsAndCollectsUsage() throws JsonProcessingException {
+        ApiKeyData apiKeyData = new ApiKeyData();
+        apiKeyData.setPerRequestKey("per-request-key");
+        when(context.getApiKeyData()).thenReturn(apiKeyData);
+        when(context.isStreamingRequest()).thenReturn(true);
+        doCallRealMethod().when(context).setPricingUsageNode(any());
+        doCallRealMethod().when(context).getPricingUsageNode();
+
+        List<BaseResponseFunction> functions = List.of(
+                new StripUsagePerModelFn(null, context),
+                new CollectResponseChatCompletionAttachmentsFn(null, context),
+                new CollectChatCompletionUsageFn(null, context));
+
+        Future<JsonNode> contentChunk = process(functions, """
+                { "choices": [{ "index": 0, "delta": { "content": "hi" } }] }
+                """);
+        assertTrue(contentChunk.succeeded(), () -> "chain failed: " + contentChunk.cause());
+
+        Future<JsonNode> usageChunk = process(functions, """
+                {
+                  "choices": [{ "index": 0, "delta": {} }],
+                  "usage": { "completion_tokens": 33, "prompt_tokens": 19, "total_tokens": 52 }
+                }
+                """);
+        assertTrue(usageChunk.succeeded(), () -> "chain failed: " + usageChunk.cause());
+
+        JsonNode pricingUsageNode = context.getPricingUsageNode();
+        assertNotNull(pricingUsageNode);
+        assertEquals(52, pricingUsageNode.path("usage").path("total_tokens").asLong());
+    }
+
+    private static Future<JsonNode> process(List<BaseResponseFunction> functions, String json) throws JsonProcessingException {
+        Future<JsonNode> result = Future.succeededFuture(ProxyUtil.MAPPER.readTree(json));
+        for (BaseResponseFunction fn : functions) {
+            result = result.compose(fn);
+        }
+        return result;
+    }
+}

server/src/test/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFnTest.java

@@ -11,6 +11,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doCallRealMethod;
 
@@ -67,6 +68,15 @@ public void testPassesEventThroughUnchanged() throws JsonProcessingException {
         assertEquals(input, output);
     }
 
+    @Test
+    public void testTolerantOfNullTree() {
+        CollectChatCompletionUsageFn fn = new CollectChatCompletionUsageFn(null, context);
+
+        JsonNode output = fn.apply(null).result();
+
+        assertNull(output);
+    }
+
     private static JsonNode tree(String json) throws JsonProcessingException {
         return ProxyUtil.MAPPER.readTree(json);
     }

server/src/test/java/com/epam/aidial/core/server/function/CollectResponseChatCompletionAttachmentsFnTest.java

@@ -1,8 +1,10 @@
 package com.epam.aidial.core.server.function;
 
 import com.epam.aidial.core.server.ProxyContext;
+import com.epam.aidial.core.server.data.ApiKeyData;
 import com.epam.aidial.core.server.util.ProxyUtil;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -12,6 +14,7 @@
 import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.mockito.Mockito.when;
 
 
@@ -218,4 +221,21 @@ public void testCollectAttachmentsFromResponse_ChatStreamingResponse() throws Js
                 "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/file2.txt"), files);
 
     }
+
+    @Test
+    public void testApplyPassesTreeThroughUnchangedWhenNoAttachments() throws JsonProcessingException {
+        ApiKeyData apiKeyData = new ApiKeyData();
+        apiKeyData.setPerRequestKey("per-request-key");
+        when(context.getApiKeyData()).thenReturn(apiKeyData);
+
+        JsonNode tree = ProxyUtil.MAPPER.readTree("""
+                {
+                  "choices": [{ "index": 0, "delta": { "content": "hi" } }]
+                }
+                """);
+
+        JsonNode result = new CollectResponseChatCompletionAttachmentsFn(null, context).apply(tree).result();
+
+        assertSame(tree, result);
+    }
 }