Repository: LangStream/langstream (language: python)

Issue:
Log filtering doesn't seem to work
After deploying an application, I typically look at the logs to find instances of an error log, like this:

```
bin/langstream apps logs webbot | grep -i error
```


Based on the output, I want to filter on just the worker than has the error logs. The log output seems to be prefixed with the worker ID, so I try to use that. But I get a 404 error:
```
webbot-crawler-webcrawler-source-1-0 14:02:48.066 [main] INFO  a.l.runtime.agent.AgentRunner -- Error, NOW SLEEPING
webbot-crawler-webcrawler-source-1-0 io.minio.errors.ErrorResponseException: Moved Permanently. Use region us-east-1
webbot-crawler-webcrawler-source-1-0 14:03:48.068 [main] ERROR a.l.runtime.agent.AgentRunner -- Unexpected error
webbot-crawler-webcrawler-source-1-0 io.minio.errors.ErrorResponseException: Moved Permanently. Use region us-east-1
^C
chrisbartholomew@cbartholomew-rmbp16 ~/git/langstream -  (main) $ bin/langstream apps logs webbot -f=webbot-crawler-webcrawler-source-1-0
{
  "error" : "Not Found",
  "path" : "/api/applications/cbart-gcp-useast1/webbot/logs%3Ffilter=webbot-crawler-webcrawler-source-1-0",
  "status" : 404,
  "timestamp" : "2023-08-24T14:05:15.247+00:00"
}ERROR: 404 status received
```
I have tried other combinations of IDs from the output, but I can't seem to get the filtering to work.


Discussion / hints:

Fix patch:
diff --git a/langstream-api/src/main/java/ai/langstream/api/storage/ApplicationStore.java b/langstream-api/src/main/java/ai/langstream/api/storage/ApplicationStore.java
index a0c36dcdc..8cfb70fac 100644
--- a/langstream-api/src/main/java/ai/langstream/api/storage/ApplicationStore.java
+++ b/langstream-api/src/main/java/ai/langstream/api/storage/ApplicationStore.java
@@ -56,14 +56,18 @@ class LogOptions {
         private List<String> filterReplicas;
     }
 
-    @FunctionalInterface
     interface PodLogHandler {
         void start(LogLineConsumer onLogLine);
+
+        String getPodName();
+
+        void close();
     }
 
-    @FunctionalInterface
     interface LogLineConsumer {
         boolean onLogLine(String line);
+
+        void onEnd();
     }
 
     List<PodLogHandler> logs(String tenant, String applicationId, LogOptions logOptions);
diff --git a/langstream-cli/src/main/java/ai/langstream/cli/commands/applications/GetApplicationLogsCmd.java b/langstream-cli/src/main/java/ai/langstream/cli/commands/applications/GetApplicationLogsCmd.java
index 29471a662..0b9c7267f 100644
--- a/langstream-cli/src/main/java/ai/langstream/cli/commands/applications/GetApplicationLogsCmd.java
+++ b/langstream-cli/src/main/java/ai/langstream/cli/commands/applications/GetApplicationLogsCmd.java
@@ -16,6 +16,11 @@
 package ai.langstream.cli.commands.applications;
 
 import ai.langstream.admin.client.AdminClient;
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
@@ -41,28 +46,30 @@ public class GetApplicationLogsCmd extends BaseApplicationCmd {
     public void run() {
         final String filterStr = filter == null ? "" : "?filter=" + String.join(",", filter);
         final AdminClient client = getClient();
-        client.getHttpClient()
-                .sendAsync(
-                        client.newGet(client.tenantAppPath("/" + name + "/logs" + filterStr)),
-                        HttpResponse.BodyHandlers.ofByteArrayConsumer(
-                                bytes -> {
-                                    if (bytes.isPresent()) {
-                                        command.commandLine()
-                                                .getOut()
-                                                .print(
-                                                        new String(
-                                                                bytes.get(),
-                                                                StandardCharsets.UTF_8));
-                                        command.commandLine().getOut().flush();
-                                    }
-                                }))
-                .thenApply(HttpResponse::statusCode)
-                .thenAccept(
-                        (status) -> {
-                            if (status != 200) {
-                                err("ERROR: %d status received".formatted(status));
-                            }
-                        })
-                .join();
+        final HttpRequest request =
+                client.newGet(client.tenantAppPath("/" + name + "/logs" + filterStr));
+        final HttpResponse<InputStream> response =
+                client.getHttpClient().send(request, HttpResponse.BodyHandlers.ofInputStream());
+        if (response.statusCode() != 200) {
+            throw new RuntimeException(
+                    "Failed to get application logs: "
+                            + response.statusCode()
+                            + " "
+                            + new String(response.body().readAllBytes(), StandardCharsets.UTF_8));
+        }
+
+        BufferedReader br = new BufferedReader(new InputStreamReader(response.body()));
+        String line;
+        try {
+            while ((line = br.readLine()) != null) {
+                if ("Heartbeat".equals(line)) {
+                    continue;
+                }
+                command.commandLine().getOut().println(line);
+                command.commandLine().getOut().flush();
+            }
+        } catch (IOException e) {
+            throw new RuntimeException("Unable to read logs", e);
+        }
     }
 }
diff --git a/langstream-k8s-storage/src/main/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStore.java b/langstream-k8s-storage/src/main/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStore.java
index 1f10df7f2..06024ca3a 100644
--- a/langstream-k8s-storage/src/main/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStore.java
+++ b/langstream-k8s-storage/src/main/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStore.java
@@ -43,6 +43,7 @@
 import io.fabric8.kubernetes.client.dsl.LogWatch;
 import io.fabric8.kubernetes.client.dsl.PodResource;
 import java.io.BufferedReader;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
@@ -399,48 +400,7 @@ public List<PodLogHandler> logs(String tenant, String applicationId, LogOptions
                             }
                             return true;
                         })
-                .map(
-                        pod -> {
-                            return new PodLogHandler() {
-                                @Override
-                                @SneakyThrows
-                                public void start(LogLineConsumer onLogLine) {
-                                    final PodResource podResource =
-                                            client.pods()
-                                                    .inNamespace(tenantToNamespace(tenant))
-                                                    .withName(pod);
-                                    final int replicas = extractReplicas(pod);
-                                    final String color = LOG_COLORS[replicas % LOG_COLORS.length];
-                                    try (final LogWatch watchLog =
-                                            podResource
-                                                    .inContainer("runtime")
-                                                    .withPrettyOutput()
-                                                    .withReadyWaitTimeout(Integer.MAX_VALUE)
-                                                    .watchLog(); ) {
-                                        final InputStream in = watchLog.getOutput();
-                                        try (BufferedReader br =
-                                                new BufferedReader(new InputStreamReader(in))) {
-                                            String line;
-                                            while ((line = br.readLine()) != null) {
-                                                String coloredLog =
-                                                        "\u001B["
-                                                                + color
-                                                                + "m"
-                                                                + pod
-                                                                + " "
-                                                                + line
-                                                                + "\u001B[0m\n";
-                                                final boolean shallContinue =
-                                                        onLogLine.onLogLine(coloredLog);
-                                                if (!shallContinue) {
-                                                    break;
-                                                }
-                                            }
-                                        }
-                                    }
-                                }
-                            };
-                        })
+                .map(pod -> new StreamPodLogHandler(tenant, pod))
                 .collect(Collectors.toList());
     }
 
@@ -448,4 +408,81 @@ private int extractReplicas(String pod) {
         final String[] split = pod.split("-");
         return Integer.parseInt(split[split.length - 1]);
     }
+
+    private class StreamPodLogHandler implements PodLogHandler {
+        private final String tenant;
+        private final String pod;
+        private InputStream in;
+        private volatile boolean closed;
+
+        public StreamPodLogHandler(String tenant, String pod) {
+            this.tenant = tenant;
+            this.pod = pod;
+        }
+
+        @Override
+        public String getPodName() {
+            return pod;
+        }
+
+        @Override
+        @SneakyThrows
+        public void start(LogLineConsumer onLogLine) {
+            while (true) {
+                try {
+                    boolean continueLoop = streamUntilEOF(onLogLine);
+                    if (!continueLoop) {
+                        return;
+                    }
+                    if (closed) {
+                        return;
+                    }
+                } catch (Throwable tt) {
+                    if (!closed) {
+                        log.warn("Error while streaming logs for pod {}", pod, tt);
+                    }
+                    onLogLine.onEnd();
+                }
+            }
+        }
+
+        @Override
+        public void close() {
+            closed = true;
+            try {
+                in.close();
+            } catch (Throwable ignored) {
+            }
+        }
+
+        private boolean streamUntilEOF(LogLineConsumer onLogLine) throws IOException {
+            final PodResource podResource =
+                    client.pods().inNamespace(tenantToNamespace(tenant)).withName(pod);
+            final int replicas = extractReplicas(pod);
+            final String color = LOG_COLORS[replicas % LOG_COLORS.length];
+            try (final LogWatch watchLog =
+                    podResource
+                            .inContainer("runtime")
+                            .withPrettyOutput()
+                            .withReadyWaitTimeout(Integer.MAX_VALUE)
+                            .watchLog(); ) {
+                in = watchLog.getOutput();
+                try (BufferedReader br = new BufferedReader(new InputStreamReader(in))) {
+                    String line;
+                    while ((line = br.readLine()) != null) {
+                        String coloredLog =
+                                "\u001B[%sm[%s] %s\u001B[0m\n".formatted(color, pod, line);
+                        final boolean shallContinue = onLogLine.onLogLine(coloredLog);
+                        if (!shallContinue) {
+                            return false;
+                        }
+                        if (closed) {
+                            return false;
+                        }
+                    }
+                }
+            }
+            return onLogLine.onLogLine("[%s] EOF\n".formatted(pod));
+        }
+    }
 }
diff --git a/langstream-webservice/src/main/java/ai/langstream/webservice/application/ApplicationResource.java b/langstream-webservice/src/main/java/ai/langstream/webservice/application/ApplicationResource.java
index d306bd2f2..802b41bd7 100644
--- a/langstream-webservice/src/main/java/ai/langstream/webservice/application/ApplicationResource.java
+++ b/langstream-webservice/src/main/java/ai/langstream/webservice/application/ApplicationResource.java
@@ -36,12 +36,16 @@
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.extern.slf4j.Slf4j;
 import net.lingala.zip4j.ZipFile;
+import org.apache.commons.lang3.concurrent.BasicThreadFactory;
 import org.springframework.core.io.ByteArrayResource;
 import org.springframework.core.io.Resource;
 import org.springframework.http.HttpStatus;
@@ -61,7 +65,6 @@
 import org.springframework.web.server.ResponseStatusException;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.FluxSink;
-import reactor.core.scheduler.Schedulers;
 
 @RestController
 @Tag(name = "applications")
@@ -73,7 +76,12 @@ public class ApplicationResource {
     ApplicationService applicationService;
     CodeStorageService codeStorageService;
 
-    private final ExecutorService logsThreadPool = Executors.newCachedThreadPool();
+    private final ScheduledExecutorService logsHeartbeatThreadPool =
+            Executors.newSingleThreadScheduledExecutor(
+                    new BasicThreadFactory.Builder().namingPattern("app-logs-hb-%d").build());
+    private final ExecutorService logsThreadPool =
+            Executors.newCachedThreadPool(
+                    new BasicThreadFactory.Builder().namingPattern("app-logs-%d").build());
 
     private void performAuthorization(Authentication authentication, final String tenant) {
         if (authentication == null) {
@@ -286,28 +294,59 @@ Flux<String> getApplicationLogs(
                         tenant,
                         applicationId,
                         new ApplicationStore.LogOptions(filterReplicas.orElse(null)));
-        if (podLogs.isEmpty()) {
-            return Flux.empty();
-        }
-        return Flux.create(
-                        (Consumer<FluxSink<String>>)
-                                fluxSink -> {
-                                    for (ApplicationStore.PodLogHandler podLog : podLogs) {
-                                        logsThreadPool.submit(
-                                                () -> {
-                                                    try {
-                                                        podLog.start(
-                                                                line -> {
-                                                                    fluxSink.next(line);
-                                                                    return true;
-                                                                });
-                                                    } catch (Exception e) {
-                                                        fluxSink.error(e);
+        AtomicLong lastSent = new AtomicLong(Long.MAX_VALUE);
+        final Consumer<FluxSink<String>> fluxSinkConsumer =
+                fluxSink -> {
+                    if (podLogs.isEmpty()) {
+                        fluxSink.next("No pods found\n");
+                        fluxSink.complete();
+                        return;
+                    }
+                    fluxSink.onDispose(
+                            () -> podLogs.forEach(ApplicationStore.PodLogHandler::close));
+                    logsHeartbeatThreadPool.scheduleWithFixedDelay(
+                            () -> {
+                                try {
+                                    if (lastSent.get() + TimeUnit.SECONDS.toMillis(30)
+                                            < System.currentTimeMillis()) {
+                                        fluxSink.next("Heartbeat\n");
+                                        lastSent.set(System.currentTimeMillis());
+                                    }
+                                } catch (Throwable e) {
+                                }
+                            },
+                            30,
+                            30,
+                            TimeUnit.SECONDS);
+                    for (ApplicationStore.PodLogHandler podLog : podLogs) {
+                        fluxSink.next(
+                                "Start receiving log for pod %s\n".formatted(podLog.getPodName()));
+                        logsThreadPool.submit(
+                                () -> {
+                                    try {
+                                        final ApplicationStore.LogLineConsumer logLineConsumer =
+                                                new ApplicationStore.LogLineConsumer() {
+                                                    @Override
+                                                    public boolean onLogLine(String line) {
+                                                        fluxSink.next(line);
+                                                        lastSent.set(System.currentTimeMillis());
+                                                        return true;
                                                     }
-                                                });
+
+                                                    @Override
+                                                    public void onEnd() {
+                                                        lastSent.set(System.currentTimeMillis());
+                                                        fluxSink.complete();
+                                                    }
+                                                };
+                                        podLog.start(logLineConsumer);
+                                    } catch (Exception e) {
+                                        fluxSink.error(e);
                                     }
-                                })
-                .subscribeOn(Schedulers.fromExecutor(logsThreadPool));
+                                });
+                    }
+                };
+        return Flux.create(fluxSinkConsumer);
     }
 
     @GetMapping(value = "/{tenant}/{applicationId}/code", produces = "application/zip")


Test patch:
diff --git a/langstream-k8s-storage/src/test/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStoreLogsTest.java b/langstream-k8s-storage/src/test/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStoreLogsTest.java
index fcafe7dfb..8056a5f32 100644
--- a/langstream-k8s-storage/src/test/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStoreLogsTest.java
+++ b/langstream-k8s-storage/src/test/java/ai/langstream/impl/storage/k8s/apps/KubernetesApplicationStoreLogsTest.java
@@ -82,16 +82,28 @@ void testLogs() {
         podHandlers
                 .get(0)
                 .start(
-                        line -> {
-                            assertEquals("\u001B[32mmyapp-agent111-0 hello\u001B[0m\n", line);
-                            return false;
+                        new ApplicationStore.LogLineConsumer() {
+                            @Override
+                            public boolean onLogLine(String line) {
+                                assertEquals("\u001B[32m[myapp-agent111-0] hello\u001B[0m\n", line);
+                                return false;
+                            }
+
+                            @Override
+                            public void onEnd() {}
                         });
         podHandlers
                 .get(1)
                 .start(
-                        line -> {
-                            assertEquals("\u001B[33mmyapp-agent111-1 hello\u001B[0m\n", line);
-                            return false;
+                        new ApplicationStore.LogLineConsumer() {
+                            @Override
+                            public boolean onLogLine(String line) {
+                                assertEquals("\u001B[33m[myapp-agent111-1] hello\u001B[0m\n", line);
+                                return false;
+                            }
+
+                            @Override
+                            public void onEnd() {}
                         });
 
         podHandlers =
@@ -103,9 +115,15 @@ void testLogs() {
         podHandlers
                 .get(0)
                 .start(
-                        line -> {
-                            assertEquals("\u001B[33mmyapp-agent111-1 hello\u001B[0m\n", line);
-                            return false;
+                        new ApplicationStore.LogLineConsumer() {
+                            @Override
+                            public boolean onLogLine(String line) {
+                                assertEquals("\u001B[33m[myapp-agent111-1] hello\u001B[0m\n", line);
+                                return false;
+                            }
+
+                            @Override
+                            public void onEnd() {}
                         });
     }