fix: use the closed window in the eof response (#221)

Signed-off-by: Vaibhav Tiwari <vaibhav.tiwari33@gmail.com>

Author: vaibhavtiwari33

src/main/java/io/numaproj/numaflow/accumulator/AccumulatorActor.java

@@ -34,6 +34,7 @@ public Receive createReceive() {
         return ReceiveBuilder
                 .create()
                 .match(HandlerDatum.class, this::invokeHandler)
+                .match(AccumulatorOuterClass.KeyedWindow.class, this::handleCloseWindow)
                 .match(String.class, this::sendEOF)
                 .build();
     }
@@ -42,17 +43,29 @@ private void invokeHandler(HandlerDatum handlerDatum) {
         this.accumulator.processMessage(handlerDatum, outputStream);
     }
 
+    // CLOSE: echo the exact close window (including slot)
+    private void handleCloseWindow(AccumulatorOuterClass.KeyedWindow closeWindow) {
+        sendEOFResponse(closeWindow);
+    }
+
+    // Fallback: the input stream completed without a CLOSE (broadcast EOF). Keep prior
+    // behavior — echo the OPEN window (start/end/keys).
     private void sendEOF(String EOF) {
+        sendEOFResponse(AccumulatorOuterClass.KeyedWindow
+                .newBuilder()
+                .setStart(this.keyedWindow.getStart())
+                .setEnd(this.keyedWindow.getEnd())
+                .addAllKeys(this.keyedWindow.getKeysList())
+                .build());
+    }
+
+    private void sendEOFResponse(AccumulatorOuterClass.KeyedWindow eofWindow) {
         // invoke handleEndOfStream to materialize the messages received so far.
         this.accumulator.handleEndOfStream(outputStream);
 
         AccumulatorOuterClass.AccumulatorResponse eofResponse = AccumulatorOuterClass.AccumulatorResponse
                 .newBuilder()
-                .setWindow(AccumulatorOuterClass.KeyedWindow
-                        .newBuilder()
-                        .setStart(this.keyedWindow.getStart())
-                        .setEnd(this.keyedWindow.getEnd())
-                        .addAllKeys(this.keyedWindow.getKeysList()))
+                .setWindow(eofWindow)
                 .setEOF(true)
                 .build();
 

src/main/java/io/numaproj/numaflow/accumulator/AccumulatorSupervisorActor.java

@@ -106,7 +106,8 @@ private void invokeActor(AccumulatorOuterClass.AccumulatorRequest request) {
                 break;
             }
             case CLOSE: {
-                actorsMap.get(uniqueId).tell(Constants.EOF, getSelf());
+                // Send the CLOSE window to the child actor so it can echo it in the EOF response
+                actorsMap.get(uniqueId).tell(request.getOperation().getKeyedWindow(), getSelf());
                 actorsMap.remove(uniqueId);
                 break;
             }

src/test/java/io/numaproj/numaflow/accumulator/ServerTest.java

@@ -1,6 +1,7 @@
 package io.numaproj.numaflow.accumulator;
 
 import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
 import io.grpc.ManagedChannel;
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.inprocess.InProcessServerBuilder;
@@ -124,6 +125,95 @@ public void testAccumulatorSingleKey() {
         }
     }
 
+    @Test
+    public void testAccumulatorEOFEchoesCloseWindow() {
+        List<String> keys = List.of("test-accumulator");
+
+        AccumulatorOuterClass.KeyedWindow openWindow = AccumulatorOuterClass.KeyedWindow
+                .newBuilder()
+                .setStart(Timestamp.newBuilder().setSeconds(0).build())
+                .setEnd(Timestamp.newBuilder().setSeconds(60).build())
+                .setSlot("slot-0")
+                .addAllKeys(keys)
+                .build();
+
+        AccumulatorOuterClass.AccumulatorRequest openRequest = AccumulatorOuterClass.AccumulatorRequest
+                .newBuilder()
+                .setPayload(AccumulatorOuterClass.Payload
+                        .newBuilder()
+                        .setValue(ByteString.copyFromUtf8("test-payload"))
+                        .addAllKeys(keys)
+                        .build())
+                .setOperation(AccumulatorOuterClass.AccumulatorRequest.WindowOperation
+                        .newBuilder()
+                        .setEvent(AccumulatorOuterClass.AccumulatorRequest.WindowOperation.Event.OPEN)
+                        .setKeyedWindow(openWindow)
+                        .build())
+                .build();
+
+        AccumulatorOuterClass.AccumulatorRequest appendRequest = AccumulatorOuterClass.AccumulatorRequest
+                .newBuilder()
+                .setPayload(AccumulatorOuterClass.Payload
+                        .newBuilder()
+                        .setValue(ByteString.copyFromUtf8("test-payload"))
+                        .addAllKeys(keys)
+                        .build())
+                .setOperation(AccumulatorOuterClass.AccumulatorRequest.WindowOperation
+                        .newBuilder()
+                        .setEvent(AccumulatorOuterClass.AccumulatorRequest.WindowOperation.Event.APPEND)
+                        .setKeyedWindow(openWindow)
+                        .build())
+                .build();
+
+        // CLOSE carries a distinct window that must be echoed verbatim in the EOF response.
+        AccumulatorOuterClass.KeyedWindow closeWindow = AccumulatorOuterClass.KeyedWindow
+                .newBuilder()
+                .setStart(Timestamp.newBuilder().setSeconds(1000).build())
+                .setEnd(Timestamp.newBuilder().setSeconds(2000).build())
+                .setSlot("slot-7")
+                .addAllKeys(keys)
+                .build();
+        AccumulatorOuterClass.AccumulatorRequest closeRequest = AccumulatorOuterClass.AccumulatorRequest
+                .newBuilder()
+                .setOperation(AccumulatorOuterClass.AccumulatorRequest.WindowOperation
+                        .newBuilder()
+                        .setEvent(AccumulatorOuterClass.AccumulatorRequest.WindowOperation.Event.CLOSE)
+                        .setKeyedWindow(closeWindow)
+                        .build())
+                .build();
+
+        // 2 data responses + 1 EOF response.
+        AccumulatorStreamObserver responseObserver = new AccumulatorStreamObserver(3);
+
+        var stub = AccumulatorGrpc.newStub(inProcessChannel);
+        var requestStreamObserver = stub.accumulateFn(responseObserver);
+
+        requestStreamObserver.onNext(openRequest);
+        requestStreamObserver.onNext(appendRequest);
+        requestStreamObserver.onNext(closeRequest);
+        requestStreamObserver.onCompleted();
+
+        try {
+            responseObserver.done.get();
+        } catch (InterruptedException | ExecutionException e) {
+            fail("Error while waiting for response" + e.getMessage());
+        }
+
+        List<AccumulatorOuterClass.AccumulatorResponse> responses = responseObserver.getResponses();
+        assertEquals(3, responses.size());
+
+        AccumulatorOuterClass.AccumulatorResponse eof = null;
+        for (AccumulatorOuterClass.AccumulatorResponse response : responses) {
+            if (response.getEOF()) {
+                eof = response;
+            }
+        }
+        assertEquals(1000, eof.getWindow().getStart().getSeconds());
+        assertEquals(2000, eof.getWindow().getEnd().getSeconds());
+        assertEquals("slot-7", eof.getWindow().getSlot());
+        assertEquals(keys, eof.getWindow().getKeysList());
+    }
+
     private static class TestAccumFn extends Accumulator {
         @Override
         public void processMessage(Datum datum, OutputStreamObserver outputStream) {