Complete python gRPC requests when closing (LangStream#690)

langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/AbstractGrpcAgent.java

@@ -95,7 +95,7 @@ public void start() throws Exception {
     }
 
     @Override
-    public void setContext(AgentContext context) throws Exception {
+    public void setContext(AgentContext context) {
         this.agentContext = context;
     }
 
@@ -128,6 +128,7 @@ public void stopChannel(boolean wait) throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
+        stopBeforeRestart();
         stopChannel(true);
     }
 

langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java

@@ -100,21 +100,13 @@ public void process(List<ai.langstream.api.runner.code.Record> records, RecordSi
                                 records.size());
                         records.forEach(recordSink::emitEmptyList);
                     } else {
-                        records.forEach(
-                                e -> {
-                                    recordSink.emitError(e, stopped);
-                                });
+                        records.forEach(e -> recordSink.emitError(e, stopped));
                     }
                 }
             }
         }
     }
 
-    @Override
-    public synchronized void close() throws Exception {
-        stopBeforeRestart();
-    }
-
     private SourceRecordAndResult fromGrpc(
             ai.langstream.api.runner.code.Record sourceRecord, ProcessorResult result)
             throws IOException {
@@ -202,8 +194,8 @@ protected void stopBeforeRestart() throws Exception {
             if (request != null) {
                 try {
                     request.onCompleted();
-                } catch (IllegalStateException ignored) {
-                    log.info("Ignoring error while stopping {}", ignored + "");
+                } catch (IllegalStateException e) {
+                    log.info("Ignoring error while stopping {}", e + "");
                 }
             }
         }

langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java

@@ -144,8 +144,8 @@ protected void stopBeforeRestart() throws Exception {
             if (request != null) {
                 try {
                     request.onCompleted();
-                } catch (IllegalStateException ignored) {
-                    log.info("Ignoring error while stopping {}", ignored + "");
+                } catch (IllegalStateException e) {
+                    log.info("Ignoring error while stopping {}", e + "");
                 }
             }
         }

langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSource.java

@@ -72,7 +72,7 @@ public List<Record> read() throws Exception {
     }
 
     @Override
-    public void permanentFailure(Record record, Exception error) throws Exception {
+    public void permanentFailure(Record record, Exception error) {
         if (record instanceof GrpcAgentRecord grpcAgentRecord) {
             request.onNext(
                     SourceRequest.newBuilder()
@@ -156,8 +156,8 @@ protected void stopBeforeRestart() throws Exception {
             if (request != null) {
                 try {
                     request.onCompleted();
-                } catch (IllegalStateException ignored) {
-                    log.info("Ignoring error while stopping {}", ignored + "");
+                } catch (IllegalStateException e) {
+                    log.info("Ignoring error while stopping {}", e + "");
                 }
             }
         }

langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentProcessor.java

@@ -47,12 +47,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
+        super.close();
+        if (server != null) {
+            server.close(false);
+        }
     }
 
     @Override
     protected synchronized void stopBeforeRestart() throws Exception {
         super.stopBeforeRestart();
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 }

langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentSink.java

@@ -45,13 +45,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
         super.close();
+        if (server != null) {
+            server.close(false);
+        }
     }
 
     @Override
     protected synchronized void stopBeforeRestart() throws Exception {
         super.stopBeforeRestart();
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 }

langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentSource.java

@@ -45,13 +45,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
         super.close();
+        if (server != null) {
+            server.close(false);
+        }
     }
 
     @Override
     protected synchronized void stopBeforeRestart() throws Exception {
         super.stopBeforeRestart();
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 }

langstream-agents/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentProcessorTest.java

@@ -63,6 +63,8 @@
 public class GrpcAgentProcessorTest {
     private Server server;
     private ManagedChannel channel;
+    private GrpcAgentProcessor processor;
+    private TestAgentContext context;
     private final AtomicInteger schemaCounter = new AtomicInteger(0);
 
     private final AgentServiceGrpc.AgentServiceImplBase testProcessorService =
@@ -128,7 +130,9 @@ public void onNext(ProcessorRequest request) {
                         public void onError(Throwable throwable) {}
 
                         @Override
-                        public void onCompleted() {}
+                        public void onCompleted() {
+                            response.onCompleted();
+                        }
                     };
                 }
             };
@@ -144,6 +148,10 @@ public void setUp() throws Exception {
                         .start();
 
         channel = InProcessChannelBuilder.forName(serverName).directExecutor().build();
+        processor = new GrpcAgentProcessor(channel);
+        context = new TestAgentContext();
+        processor.setContext(context);
+        processor.start();
         schemaCounter.set(0);
     }
 
@@ -175,9 +183,6 @@ private static Stream<Arguments> primitives() {
     @ParameterizedTest
     @MethodSource("primitives")
     void testProcess(Object value, Object key, Object header) throws Exception {
-        GrpcAgentProcessor processor = new GrpcAgentProcessor(channel);
-        processor.setContext(new TestAgentContext());
-        processor.start();
         Record inputRecord =
                 SimpleRecord.builder()
                         .value(value)
@@ -188,26 +193,17 @@ void testProcess(Object value, Object key, Object header) throws Exception {
                         .build();
         assertProcessSuccessful(processor, inputRecord);
         assertProcessSuccessful(processor, inputRecord);
-        processor.close();
     }
 
     @Test
     void testEmpty() throws Exception {
-        GrpcAgentProcessor processor = new GrpcAgentProcessor(channel);
-        TestAgentContext context = new TestAgentContext();
-        processor.setContext(context);
-        processor.start();
         assertProcessSuccessful(processor, SimpleRecord.builder().build());
         assertFalse(context.failureCalled.await(1, TimeUnit.SECONDS));
-        processor.close();
     }
 
     @Test
     void testFailingRecord() throws Exception {
-        GrpcAgentProcessor processor = new GrpcAgentProcessor(channel);
         Record inputRecord = SimpleRecord.builder().origin("failing-origin").build();
-        processor.setContext(new TestAgentContext());
-        processor.start();
         CompletableFuture<Void> op = new CompletableFuture<>();
         processor.process(
                 List.of(inputRecord),
@@ -223,23 +219,17 @@ void testFailingRecord() throws Exception {
                     op.complete(null);
                 });
         op.get(5, TimeUnit.SECONDS);
-        processor.close();
     }
 
     @ParameterizedTest
     @ValueSource(
             strings = {"failing-server", "completing-server", "wrong-record-id", "wrong-schema-id"})
     void testServerError(String origin) throws Exception {
-        GrpcAgentProcessor processor = new GrpcAgentProcessor(channel);
         Record inputRecord = SimpleRecord.builder().origin(origin).build();
 
-        TestAgentContext testAgentContext = new TestAgentContext();
-        processor.setContext(testAgentContext);
-        processor.start();
         processor.process(List.of(inputRecord), result -> {});
 
-        assertTrue(testAgentContext.failureCalled.await(1, TimeUnit.SECONDS));
-        processor.close();
+        assertTrue(context.failureCalled.await(1, TimeUnit.SECONDS));
     }
 
     @Test
@@ -254,9 +244,6 @@ void testAvroAndSchema() throws Exception {
                         .endRecord();
         GenericData.Record avroRecord = new GenericData.Record(schema);
         avroRecord.put("testField", "test-string");
-        GrpcAgentProcessor processor = new GrpcAgentProcessor(channel);
-        processor.setContext(new TestAgentContext());
-        processor.start();
         Record inputRecord =
                 SimpleRecord.builder()
                         .value(avroRecord)
@@ -267,14 +254,10 @@ void testAvroAndSchema() throws Exception {
         assertProcessSuccessful(processor, inputRecord);
         assertProcessSuccessful(processor, inputRecord);
         assertEquals(1, schemaCounter.get());
-        processor.close();
     }
 
     @Test
     void testInfo() throws Exception {
-        GrpcAgentProcessor processor = new GrpcAgentProcessor(channel);
-        processor.setContext(new TestAgentContext());
-        processor.start();
         Map<String, Object> info = processor.buildAdditionalInfo();
         assertEquals("test-info-value", info.get("test-info-key"));
     }

langstream-agents/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSinkTest.java

@@ -176,7 +176,9 @@ public void onNext(SinkRequest request) {
                 public void onError(Throwable throwable) {}
 
                 @Override
-                public void onCompleted() {}
+                public void onCompleted() {
+                    responseObserver.onCompleted();
+                }
             };
         }
     }

langstream-agents/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSourceTest.java

@@ -53,7 +53,8 @@
 public class GrpcAgentSourceTest {
     private Server server;
     private ManagedChannel channel;
-
+    private GrpcAgentSource source;
+    private TestAgentContext context;
     private final TestSourceService testSourceService = new TestSourceService();
 
     @BeforeEach
@@ -67,6 +68,10 @@ public void setUp() throws Exception {
                         .start();
 
         channel = InProcessChannelBuilder.forName(serverName).directExecutor().build();
+        source = new GrpcAgentSource(channel);
+        context = new TestAgentContext();
+        source.setContext(context);
+        source.start();
     }
 
     @AfterEach
@@ -79,10 +84,6 @@ public void tearDown() throws Exception {
 
     @Test
     void testCommit() throws Exception {
-        GrpcAgentSource source = new GrpcAgentSource(channel);
-        TestAgentContext context = new TestAgentContext();
-        source.setContext(context);
-        source.start();
         List<Record> read = readRecords(source, 3);
         source.commit(List.of(read.get(0)));
         assertFalse(context.failureCalled.await(1, TimeUnit.SECONDS));
@@ -93,10 +94,6 @@ void testCommit() throws Exception {
 
     @Test
     void testSourceGrpcError() throws Exception {
-        GrpcAgentSource source = new GrpcAgentSource(channel);
-        TestAgentContext context = new TestAgentContext();
-        source.setContext(context);
-        source.start();
         List<Record> read = readRecords(source, 3);
         source.commit(List.of(read.get(1)));
         assertTrue(context.failureCalled.await(1, TimeUnit.SECONDS));
@@ -105,10 +102,6 @@ void testSourceGrpcError() throws Exception {
 
     @Test
     void testSourceGrpcCompletedUnexpectedly() throws Exception {
-        GrpcAgentSource source = new GrpcAgentSource(channel);
-        TestAgentContext context = new TestAgentContext();
-        source.setContext(context);
-        source.start();
         List<Record> read = readRecords(source, 3);
         source.commit(List.of(read.get(2)));
         assertTrue(context.failureCalled.await(1, TimeUnit.SECONDS));
@@ -117,9 +110,6 @@ void testSourceGrpcCompletedUnexpectedly() throws Exception {
 
     @Test
     void testAvroAndSchema() throws Exception {
-        GrpcAgentSource source = new GrpcAgentSource(channel);
-        source.setContext(new TestAgentContext());
-        source.start();
         List<Record> read = readRecords(source, 1);
         GenericRecord record = (GenericRecord) read.get(0).value();
         assertEquals("test-string", record.get("testField").toString());
@@ -128,9 +118,6 @@ void testAvroAndSchema() throws Exception {
 
     @Test
     void testPermanentFailure() throws Exception {
-        GrpcAgentSource source = new GrpcAgentSource(channel);
-        source.setContext(new TestAgentContext());
-        source.start();
         List<Record> read = readRecords(source, 1);
         source.permanentFailure(read.get(0), new RuntimeException("permanent-failure"));
         assertEquals(testSourceService.permanentFailure.getRecordId(), 42);
@@ -227,7 +214,9 @@ public void onNext(SourceRequest request) {
                 public void onError(Throwable throwable) {}
 
                 @Override
-                public void onCompleted() {}
+                public void onCompleted() {
+                    responseObserver.onCompleted();
+                }
             };
         }
     }

langstream-api/src/main/java/ai/langstream/api/runner/code/AgentCode.java

@@ -66,7 +66,7 @@ default void close() throws Exception {}
     /**
      * Gracefully restart the agent.
      *
-     * @throws Exception
+     * @throws Exception if an error occurs
      */
     default void restart() throws Exception {}
 }