[python] Add support for hot reload of python code in docker mode

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

@@ -101,11 +101,26 @@ protected Map<String, Object> buildAdditionalInfo() {
         }
     }
 
+    protected synchronized void stop() throws Exception {
+        stopChannel(false);
+    }
+
+    public void stopChannel(boolean wait) throws Exception {
+        ManagedChannel currentChannel;
+        synchronized (this) {
+            currentChannel = channel;
+        }
+        if (currentChannel != null) {
+            ManagedChannel shutdown = currentChannel.shutdown();
+            if (wait) {
+                shutdown.awaitTermination(10, TimeUnit.SECONDS);
+            }
+        }
+    }
+
     @Override
     public synchronized void close() throws Exception {
-        if (channel != null) {
-            channel.shutdown();
-        }
+        stopChannel(true);
     }
 
     protected Object fromGrpc(Value value) throws IOException {

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

@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import lombok.extern.slf4j.Slf4j;
 
@@ -40,6 +41,8 @@ public class GrpcAgentProcessor extends AbstractGrpcAgent implements AgentProces
 
     private final StreamObserver<ProcessorResponse> responseObserver = getResponseObserver();
 
+    private final AtomicBoolean restarting = new AtomicBoolean(false);
+
     private record RecordAndSink(
             ai.langstream.api.runner.code.Record sourceRecord, RecordSink sink) {}
 
@@ -58,6 +61,7 @@ public synchronized void onNewSchemaToSend(Schema schema) {
 
     @Override
     public void start() throws Exception {
+        restarting.set(false);
         super.start();
         request = AgentServiceGrpc.newStub(channel).withWaitForReady().process(responseObserver);
     }
@@ -80,16 +84,21 @@ public synchronized void process(
             }
         }
         if (requestBuilder.getRecordsCount() > 0) {
-            request.onNext(requestBuilder.build());
+            try {
+                request.onNext(requestBuilder.build());
+            } catch (IllegalStateException stopped) {
+                if (restarting.get()) {
+                    log.info("Ignoring error during restart {}", stopped + "");
+                } else {
+                    throw stopped;
+                }
+            }
         }
     }
 
     @Override
     public synchronized void close() throws Exception {
-        if (request != null) {
-            request.onCompleted();
-        }
-        super.close();
+        stop();
     }
 
     private SourceRecordAndResult fromGrpc(
@@ -150,17 +159,36 @@ public void onNext(ProcessorResponse response) {
 
             @Override
             public void onError(Throwable throwable) {
-                agentContext.criticalFailure(
-                        new RuntimeException(
-                                "gRPC server sent error: %s".formatted(throwable.getMessage()),
-                                throwable));
+                if (!restarting.get()) {
+                    agentContext.criticalFailure(
+                            new RuntimeException(
+                                    "gRPC server sent error: %s".formatted(throwable.getMessage()),
+                                    throwable));
+                } else {
+                    log.info("Ignoring error during restart {}", throwable + "");
+                }
             }
 
             @Override
             public void onCompleted() {
-                agentContext.criticalFailure(
-                        new RuntimeException("gRPC server completed the stream unexpectedly"));
+                if (!restarting.get()) {
+                    agentContext.criticalFailure(
+                            new RuntimeException("gRPC server completed the stream unexpectedly"));
+                } else {
+                    log.info("Ignoring error server stop during restart");
+                }
             }
         };
     }
+
+    protected void stop() throws Exception {
+        log.info("Restarting...");
+        restarting.set(true);
+        synchronized (this) {
+            if (request != null) {
+                request.onCompleted();
+            }
+        }
+        super.stop();
+    }
 }

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

@@ -39,7 +39,19 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close();
-        super.close();
+        if (server != null) server.close(false);
+    }
+
+    @Override
+    protected synchronized void stop() throws Exception {
+        if (server != null) server.close(true);
+    }
+
+    @Override
+    public void restart() throws Exception {
+        super.stop();
+        stop();
+        start();
+        super.start();
     }
 }

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

@@ -39,7 +39,20 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close();
+        if (server != null) server.close(false);
         super.close();
     }
+
+    @Override
+    protected synchronized void stop() throws Exception {
+        if (server != null) server.close(true);
+    }
+
+    @Override
+    public void restart() throws Exception {
+        super.stop();
+        stop();
+        start();
+        super.start();
+    }
 }

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

@@ -39,7 +39,20 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close();
+        if (server != null) server.close(false);
         super.close();
     }
+
+    @Override
+    protected synchronized void stop() throws Exception {
+        if (server != null) server.close(true);
+    }
+
+    @Override
+    public void restart() throws Exception {
+        super.stop();
+        stop();
+        start();
+        super.start();
+    }
 }

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

@@ -117,14 +117,16 @@ private AgentContextConfiguration computeAgentContextConfiguration() {
         return agentContextConfiguration;
     }
 
-    public void close() throws Exception {
+    public void close(boolean ignoreErrors) throws Exception {
         if (pythonProcess != null) {
             pythonProcess.destroy();
             int exitCode = pythonProcess.waitFor();
             log.info("Python process exited with code {}", exitCode);
 
-            if (exitCode != 0) {
-                throw new RuntimeException("Python code exited with code " + exitCode);
+            if (!ignoreErrors) {
+                if (exitCode != 0) {
+                    throw new RuntimeException("Python code exited with code " + exitCode);
+                }
             }
         }
     }

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

@@ -18,10 +18,14 @@
 import ai.langstream.api.runtime.ComponentType;
 import java.util.List;
 import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** Body of the agent */
 public interface AgentCode extends AutoCloseable {
 
+    static final Logger log = LoggerFactory.getLogger(AgentCode.class);
+
     String agentId();
 
     /**
@@ -58,4 +62,13 @@ default void close() throws Exception {}
      * @return information about the agent
      */
     List<AgentStatusResponse> getAgentStatus();
+
+    /**
+     * Gracefully restart the agent.
+     *
+     * @throws Exception
+     */
+    default void restart() throws Exception {
+        log.info("Restart is not supported for agent type {}", agentType());
+    }
 }

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

@@ -134,6 +134,11 @@ public void close() throws Exception {
                 executeWithContextClassloader(AgentCode::close);
             }
 
+            @Override
+            public void restart() throws Exception {
+                executeWithContextClassloader(AgentCode::restart);
+            }
+
             @Override
             public ComponentType componentType() {
                 return callNoExceptionWithContextClassloader(AgentCode::componentType);
@@ -192,6 +197,11 @@ public void close() throws Exception {
                 executeWithContextClassloader(AgentCode::close);
             }
 
+            @Override
+            public void restart() throws Exception {
+                executeWithContextClassloader(AgentCode::restart);
+            }
+
             @Override
             public ComponentType componentType() {
                 return callNoExceptionWithContextClassloader(AgentCode::componentType);
@@ -249,6 +259,11 @@ public void close() throws Exception {
                 executeWithContextClassloader(AgentCode::close);
             }
 
+            @Override
+            public void restart() throws Exception {
+                executeWithContextClassloader(AgentCode::restart);
+            }
+
             @Override
             public ComponentType componentType() {
                 return callNoExceptionWithContextClassloader(AgentCode::componentType);

langstream-cli/src/main/java/ai/langstream/cli/commands/docker/ApplicationWatcher.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ai.langstream.cli.commands.docker;
+
+import static java.nio.file.StandardWatchEventKinds.ENTRY_CREATE;
+import static java.nio.file.StandardWatchEventKinds.ENTRY_DELETE;
+import static java.nio.file.StandardWatchEventKinds.ENTRY_MODIFY;
+import static java.nio.file.StandardWatchEventKinds.OVERFLOW;
+
+import java.nio.file.Path;
+import java.nio.file.WatchEvent;
+import java.nio.file.WatchKey;
+import java.nio.file.WatchService;
+import java.util.function.Consumer;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class ApplicationWatcher {
+
+    public static void watchApplication(
+            Path codeDirectory, Consumer<String> changedFiles, WatchService watcher)
+            throws Exception {
+
+        Thread watchThread =
+                new Thread(
+                        () -> {
+                            try {
+                                watchFiles(watcher, codeDirectory, changedFiles);
+                            } catch (Throwable e) {
+                                e.printStackTrace();
+                            }
+                        });
+
+        watchThread.start();
+    }
+
+    private static void watchFiles(WatchService watcher, Path dir, Consumer<String> changedFiles)
+            throws Exception {
+        WatchKey register = dir.register(watcher, ENTRY_CREATE, ENTRY_DELETE, ENTRY_MODIFY);
+
+        log.info("Watching files in {}, key {}", dir, register);
+        for (; ; ) {
+
+            // wait for key to be signaled
+            WatchKey key;
+            try {
+                key = watcher.take();
+            } catch (InterruptedException x) {
+                return;
+            }
+
+            for (WatchEvent<?> event : key.pollEvents()) {
+                WatchEvent.Kind<?> kind = event.kind();
+
+                // This key is registered only
+                // for ENTRY_CREATE events,
+                // but an OVERFLOW event can
+                // occur regardless if events
+                // are lost or discarded.
+                if (kind == OVERFLOW) {
+                    continue;
+                }
+
+                // The filename is the
+                // context of the event.
+                WatchEvent<Path> ev = (WatchEvent<Path>) event;
+                Path filename = ev.context();
+
+                // Verify that the new
+                //  file is a text file.
+                try {
+                    // Resolve the filename against the directory.
+                    // If the filename is "test" and the directory is "foo",
+                    // the resolved name is "test/foo".
+                    Path child = dir.resolve(filename);
+                    changedFiles.accept(filename.toAbsolutePath().toString());
+                } catch (Exception x) {
+                    log.error("Error while watching files", x);
+                    continue;
+                }
+            }
+
+            // Reset the key -- this step is critical if you want to
+            // receive further watch events.  If the key is no longer valid,
+            // the directory is inaccessible so exit the loop.
+            boolean valid = key.reset();
+            if (!valid) {
+                log.info("Key is not valid, exiting");
+                break;
+            }
+        }
+    }
+}