agent exits when drained so the launcher can wait for a clean upload

src/main/java/com/gpuflight/agent/Main.java

@@ -32,6 +32,14 @@ static void main(String[] args) throws Exception {
         Publisher publisher = PublisherFactory.create(config.publisher());
         System.out.println("[agent] Publisher: " + config.publisher().getClass().getSimpleName());
 
+        // One-shot mode for the launcher-driven `gpufl trace` upload: exit once every
+        // discovered session has fully drained, so the launcher can wait for a clean finish
+        // instead of hard-killing the agent mid-upload (which dropped late windows).
+        boolean exitWhenDrained = parseExitWhenDrained(args, System.getenv());
+        if (exitWhenDrained) {
+            System.out.println("[agent] exit-when-drained mode enabled");
+        }
+
         // In v1.2 a "source" is just a FOLDER to watch - each run writes its
         // logs under <folder>/<session_id>/<channel>.log[.gz], so sessions are
         // auto-discovered (no per-session config). Collect every folder to
@@ -94,6 +102,11 @@ static void main(String[] args) throws Exception {
         // re-watched, so this set also marks "done" for the agent's lifetime.)
         var startedSessions = java.util.concurrent.ConcurrentHashMap.<String>newKeySet();
 
+        // Live per-channel tailer tasks across ALL sessions. The exit-when-drained monitor
+        // exits the agent once this returns to 0 (after >=1 session was seen and nothing is
+        // still being written).
+        var activeTailers = new java.util.concurrent.atomic.AtomicInteger(0);
+
         // Spawn the per-channel tailer set for one discovered session.
         // Idempotent - a second call for the same (folder, session_id) is a
         // no-op, so the rescan can call it freely without re-announcing.
@@ -112,19 +125,24 @@ static void main(String[] args) throws Exception {
             var remaining = new java.util.concurrent.atomic.AtomicInteger(session.logTypes().size());
             String sid = session.sessionId();
             for (String type : session.logTypes()) {
+                activeTailers.incrementAndGet();
                 executor.submit(() -> {
-                    // Only the "system" channel carries device_metric_batch events.
-                    var dedup = "system".equals(type) ? deduplicator : null;
-                    var tailer = new LogTailer(session.folder(), session.sessionId(), type,
-                                               topicPrefix, cursorMgr, consumedFilesQueue, dedup,
-                                               streamUploadSettings);
-                    tailer.tail(publisher);
-                    // Reached only when tail() returns. Skip on interruption (agent
-                    // shutdown) so an in-flight upload is never declared complete.
-                    if (!Thread.currentThread().isInterrupted()
-                            && remaining.decrementAndGet() == 0
-                            && streamUploadSettings.enabled()) {
-                        signalSessionComplete(publisher, sid);
+                    try {
+                        // Only the "system" channel carries device_metric_batch events.
+                        var dedup = "system".equals(type) ? deduplicator : null;
+                        var tailer = new LogTailer(session.folder(), session.sessionId(), type,
+                                                   topicPrefix, cursorMgr, consumedFilesQueue, dedup,
+                                                   streamUploadSettings);
+                        tailer.tail(publisher);
+                        // Reached only when tail() returns. Skip on interruption (agent
+                        // shutdown) so an in-flight upload is never declared complete.
+                        if (!Thread.currentThread().isInterrupted()
+                                && remaining.decrementAndGet() == 0
+                                && streamUploadSettings.enabled()) {
+                            signalSessionComplete(publisher, sid);
+                        }
+                    } finally {
+                        activeTailers.decrementAndGet();
                     }
                 });
             }
@@ -188,10 +206,62 @@ static void main(String[] args) throws Exception {
             try { publisher.close(); } catch (Exception ignored) {}
         }));
 
-        // Block main thread until SIGTERM/Ctrl+C triggers the shutdown hook.
+        if (exitWhenDrained) {
+            awaitDrainThenExit(watchedFolders.keySet(), startedSessions, activeTailers);
+            System.out.println("[agent] all sessions drained - exiting");
+            executor.shutdownNow();
+            try { publisher.close(); } catch (Exception ignored) {}
+            return;
+        }
+
+        // Daemon mode: block main thread until SIGTERM/Ctrl+C triggers the shutdown hook.
         new CountDownLatch(1).await();
     }
 
+    /**
+     * Block until every discovered session has fully drained (no live tailers) and no
+     * session is still being written (no {@code .tmp/} dir), after at least one session
+     * was seen. Lets the launcher-driven one-shot {@code gpufl trace} upload wait for a
+     * clean finish instead of hard-killing the agent mid-upload. Two consecutive clean
+     * checks guard against a session discovered between ticks.
+     */
+    private static void awaitDrainThenExit(Collection<File> folders,
+                                           Set<String> startedSessions,
+                                           java.util.concurrent.atomic.AtomicInteger activeTailers)
+            throws InterruptedException {
+        int clean = 0;
+        while (true) {
+            Thread.sleep(1000);
+            boolean drained = !startedSessions.isEmpty()
+                    && activeTailers.get() == 0
+                    && !anyActiveSession(folders);
+            clean = drained ? clean + 1 : 0;
+            if (clean >= 2) return;
+        }
+    }
+
+    /** True if any session subdir under {@code folders} still has a {@code .tmp/} working
+     *  dir - i.e., a run is still writing it (the client removes .tmp once every channel
+     *  closes). Guards exit-when-drained against a run that hasn't finished yet. */
+    static boolean anyActiveSession(Collection<File> folders) {
+        for (File folder : folders) {
+            File[] subdirs = folder.listFiles(File::isDirectory);
+            if (subdirs == null) continue;
+            for (File subdir : subdirs) {
+                if (subdir.getName().startsWith(".")) continue;
+                if (new File(subdir, ".tmp").isDirectory()) return true;
+            }
+        }
+        return false;
+    }
+
+    /** Whether the launcher asked the agent to exit once its session(s) drain
+     *  (GPUFL_AGENT_EXIT_WHEN_DRAINED / --exit-when-drained). */
+    static boolean parseExitWhenDrained(String[] args, Map<String, String> env) {
+        String v = resolve(args, "exit-when-drained", "GPUFL_AGENT_EXIT_WHEN_DRAINED", null, env);
+        return v != null && !v.equalsIgnoreCase("false") && !v.equals("0");
+    }
+
     /**
      * Tell the backend that every channel of a session has finished uploading.
      * Best-effort with a few bounded retries: {@code publishSessionComplete}

src/main/java/com/gpuflight/agent/publisher/HttpPublisher.java

@@ -84,47 +84,66 @@ public boolean publishStream(String sessionId, List<String> ndjsonLines) {
         try {
             String ndjson = String.join("\n", ndjsonLines) + "\n";
             byte[] body = gzip(ndjson.getBytes(StandardCharsets.UTF_8));
-            String url = config.streamEndpoint();
-            System.out.println("[agent] HTTP stream POST starting: url=" + url
-                + " session=" + sessionId
-                + " lines=" + ndjsonLines.size()
-                + " gzipBytes=" + body.length);
+            System.out.println("[agent] HTTP stream POST starting: url=" + config.streamEndpoint()
+                + " session=" + sessionId + " lines=" + ndjsonLines.size() + " gzipBytes=" + body.length);
+            return postGzStream(sessionId, body);
+        } catch (Exception e) {
+            System.out.println("HTTP stream connection error: " + e);
+            return false;
+        }
+    }
 
+    @Override
+    public boolean publishStreamGz(String sessionId, byte[] gzBody) {
+        if (gzBody == null || gzBody.length == 0) {
+            return true;
+        }
+        System.out.println("[agent] HTTP stream POST starting (gz window): url=" + config.streamEndpoint()
+            + " session=" + sessionId + " gzipBytes=" + gzBody.length);
+        return postGzStream(sessionId, gzBody);
+    }
+
+    /**
+     * POST an already-gzipped NDJSON body to the stream endpoint and interpret the
+     * status. 2xx (accepted), 409 (already finalized on a restart) and 402 (limit)
+     * all advance; anything else is a retryable failure.
+     */
+    private boolean postGzStream(String sessionId, byte[] gzBody) {
+        try {
+            String url = config.streamEndpoint();
             var requestBuilder = HttpRequest.newBuilder()
                 .uri(URI.create(url))
                 .header("Content-Type", "application/x-ndjson")
                 .header("Content-Encoding", "gzip")
                 .header("X-GpuFlight-Session-Id", sessionId)
                 .header("X-GpuFlight-Hostname", InetAddress.getLocalHost().getHostName())
-                .POST(HttpRequest.BodyPublishers.ofByteArray(body));
+                .POST(HttpRequest.BodyPublishers.ofByteArray(gzBody));
 
             addAuthHeader(requestBuilder);
 
             var response = client.send(requestBuilder.build(), HttpResponse.BodyHandlers.ofString());
-            if (response.statusCode() >= 200 && response.statusCode() <= 299) {
-                System.out.println("[agent] HTTP stream POST accepted: status="
-                    + response.statusCode() + " session=" + sessionId);
+            int sc = response.statusCode();
+            if (sc >= 200 && sc <= 299) {
+                System.out.println("[agent] HTTP stream POST accepted: status=" + sc + " session=" + sessionId);
                 return true;
             }
-            if (response.statusCode() == 409) {
+            if (sc == 409) {
                 // Session already finalized on the backend - typically this agent
-                // re-tailing a finished session after a restart. The data is
-                // already stored and re-profiling is required to replace it, so
-                // retrying NEVER succeeds. Advance like a 2xx (mirrors the 402
-                // handling in publish()) so the tailer drains to EOF and exits
-                // instead of re-POSTing the same batch every 5s forever.
+                // re-tailing a finished session after a restart. The data is already
+                // stored and re-profiling is required to replace it, so retrying NEVER
+                // succeeds. Advance like a 2xx so the tailer finishes instead of
+                // re-POSTing forever.
                 System.out.println("[agent] HTTP stream POST: session already uploaded (409) - "
                     + "skipping, session=" + sessionId);
                 return true;
             }
-            if (response.statusCode() == 402) {
+            if (sc == 402) {
                 // GPU/workspace limit exceeded - permanent; retrying won't help.
                 System.err.println("[GPUFL] GPU limit exceeded for this workspace - "
                     + "skipping session=" + sessionId + ". " + response.body());
                 return true;
             }
-            System.out.println("HTTP stream publish failed [" + url + "]: "
-                + response.statusCode() + " - " + response.body());
+            System.out.println("HTTP stream publish failed [" + url + "]: " + sc + " - " + response.body());
             return false;
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();

src/main/java/com/gpuflight/agent/publisher/Publisher.java

@@ -14,6 +14,17 @@ default boolean publishStream(String sessionId, List<String> ndjsonLines) {
         return false;
     }
 
+    /**
+     * Upload a complete, already-gzipped NDJSON window verbatim as one chunk - the
+     * client's {@code <channel>.<index>.log.gz} sent byte-for-byte. It is exactly the
+     * gzipped-NDJSON wire format {@link #publishStream} produces, so this skips the
+     * decompress + re-batch + re-gzip round-trip. Default unsupported - the window
+     * fast-path is gated to the HTTP stream publisher.
+     */
+    default boolean publishStreamGz(String sessionId, byte[] gzBody) {
+        return false;
+    }
+
     /**
      * Signal the backend that EVERY channel of {@code sessionId} has finished
      * uploading — the agent has drained all per-channel tailers and every batch