fix(agent): keep detached probes reconnectable; add Client.connectAndExitProbe

btrace-agent/src/main/java/io/btrace/agent/Client.java

@@ -288,6 +288,11 @@ static Collection<String> listProbes() {
     return probes;
   }
 
+  /** Returns {@code true} once {@link #onExit(int)} has started removing this probe. */
+  final boolean isShuttingDown() {
+    return shuttingDown;
+  }
+
   synchronized void onExit(int exitCode) {
     if (!shuttingDown) {
       shuttingDown = true;

btrace-agent/src/main/java/io/btrace/agent/RemoteClient.java

@@ -103,6 +103,8 @@ public Boolean apply(Command value) {
       AtomicReferenceFieldUpdater.newUpdater(RemoteClient.class, Socket.class, "sock");
   private final AtomicReferenceFieldUpdater<RemoteClient, WireProtocol> protocolUpdater =
       AtomicReferenceFieldUpdater.newUpdater(RemoteClient.class, WireProtocol.class, "protocol");
+  // Keeps sock and protocol a consistent pair across reconnect() and releaseTransport().
+  private final Object transportLock = new Object();
 
   private final CircularBuffer<Command> delayedCommands = new CircularBuffer<>(5000);
   private final AtomicReference<TerminalHandshake> terminalHandshake = new AtomicReference<>();
@@ -279,13 +281,20 @@ private void initClient() {
             () -> {
               try {
                 BTraceRuntime.enter();
-                while (true) {
+                while (!isShuttingDown()) {
+                  WireProtocol input;
+                  Socket inputSocket;
+                  synchronized (transportLock) {
+                    input = protocol;
+                    inputSocket = sock;
+                  }
                   try {
-                    if (protocol == null) {
+                    if (input == null) {
+                      // Detached: wait for reconnect() to install a new transport.
                       LockSupport.parkNanos(500_000_000L); // sleep 500ms
                       continue;
                     }
-                    Command cmd = protocol.read();
+                    Command cmd = input.read();
                     switch (cmd.getType()) {
                       case Command.EXIT:
                         {
@@ -341,6 +350,12 @@ private void initClient() {
                       if (log.isDebugEnabled()) {
                         log.debug("client command stream closed: {}", exp.toString());
                       }
+                      if (terminalHandshake.get() == null) {
+                        // The probe outlives its client: keep it reconnectable and keep reading
+                        // once a client reconnects.
+                        releaseTransport(input, inputSocket);
+                        continue;
+                      }
                     } else {
                       log.debug("Error while processing BTrace command", exp);
                     }
@@ -548,10 +563,38 @@ protected void closeAll() throws IOException {
     }
   }
 
+  /**
+   * Drops a client transport whose stream has ended, leaving the probe running and reconnectable. A
+   * transport that reconnect() has already replaced is closed without touching the new one.
+   */
+  private void releaseTransport(WireProtocol input, Socket inputSocket) {
+    synchronized (transportLock) {
+      if (protocol == input) {
+        protocol = null;
+        sock = null;
+        disconnected = true;
+      }
+    }
+    try {
+      input.close();
+    } catch (IOException ignore) {
+      // best effort
+    }
+    if (inputSocket != null) {
+      try {
+        inputSocket.close();
+      } catch (IOException ignore) {
+        // best effort
+      }
+    }
+  }
+
   void reconnect(WireProtocol protocol, Socket socket) throws IOException {
-    this.sock = socket;
-    this.protocol = protocol;
-    this.disconnected = false;
+    synchronized (transportLock) {
+      this.sock = socket;
+      this.protocol = protocol;
+      this.disconnected = false;
+    }
     onCommand(Command.NULL);
   }
 }

btrace-agent/src/test/java/io/btrace/agent/RemoteClientReconnectTest.java

@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) 2026, Jaroslav Bachorik <j.bachorik@btrace.io>.
+ * All rights reserved.
+ *
+ * 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
+ *
+ *     https://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 io.btrace.agent;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.btrace.core.ArgsMap;
+import io.btrace.core.BTraceRuntime;
+import io.btrace.core.SharedSettings;
+import io.btrace.core.comm.BinaryWireProtocol;
+import io.btrace.core.comm.ExitCommand;
+import io.btrace.core.comm.WireProtocol;
+import java.lang.reflect.Proxy;
+import java.net.InetAddress;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.jupiter.api.Test;
+
+class RemoteClientReconnectTest {
+  private static final int EXIT_CODE = 41;
+
+  @Test
+  void reconnectedClientCommandsReachTheProbeAfterThePeerWentAway() throws Exception {
+    CountDownLatch exitRequested = new CountDownLatch(1);
+    AtomicInteger requestedExitCode = new AtomicInteger(Integer.MIN_VALUE);
+    BTraceRuntime.Impl runtime =
+        runtimeProxy(
+            exitCode -> {
+              requestedExitCode.set(exitCode);
+              exitRequested.countDown();
+            });
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getLoopbackAddress())) {
+      Future<RemoteClient> accepted =
+          executor.submit(
+              () -> {
+                Socket socket = server.accept();
+                ClientContext context =
+                    new ClientContext(null, null, new ArgsMap(), new SharedSettings());
+                return RemoteClient.createForTerminalTest(
+                    context, protocol(socket), socket, runtime, exitCode -> {});
+              });
+      RemoteClient remote;
+      try (Socket first = new Socket(server.getInetAddress(), server.getLocalPort())) {
+        remote = accepted.get(5, TimeUnit.SECONDS);
+      }
+      // The client side closed its socket, as a detaching client does.
+      awaitTrue(remote::isDisconnected, "agent did not notice the client going away");
+
+      Future<Socket> reconnecting = executor.submit(server::accept);
+      try (Socket second = new Socket(server.getInetAddress(), server.getLocalPort())) {
+        Socket agentSide = reconnecting.get(5, TimeUnit.SECONDS);
+        remote.reconnect(protocol(agentSide), agentSide);
+        assertFalse(remote.isDisconnected());
+
+        WireProtocol client = protocol(second);
+        client.write(new ExitCommand(EXIT_CODE));
+        client.flush();
+
+        assertTrue(
+            exitRequested.await(5, TimeUnit.SECONDS),
+            "EXIT from the reconnected client never reached the probe");
+        assertEquals(EXIT_CODE, requestedExitCode.get());
+      }
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+
+  private static WireProtocol protocol(Socket socket) throws java.io.IOException {
+    return new BinaryWireProtocol(socket.getInputStream(), socket.getOutputStream());
+  }
+
+  private static void awaitTrue(java.util.function.BooleanSupplier condition, String message)
+      throws InterruptedException {
+    long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+    while (!condition.getAsBoolean()) {
+      if (System.nanoTime() > deadline) {
+        throw new AssertionError(message);
+      }
+      Thread.sleep(20);
+    }
+  }
+
+  private static BTraceRuntime.Impl runtimeProxy(final ExitRecorder recorder) {
+    return (BTraceRuntime.Impl)
+        Proxy.newProxyInstance(
+            RemoteClientReconnectTest.class.getClassLoader(),
+            new Class<?>[] {BTraceRuntime.Impl.class},
+            (proxy, method, args) -> {
+              if ("handleExit".equals(method.getName())) {
+                recorder.record(((Integer) args[0]).intValue());
+              }
+              Class<?> type = method.getReturnType();
+              if (type == Boolean.TYPE) return Boolean.FALSE;
+              if (type == Integer.TYPE) return Integer.valueOf(0);
+              if (type == Long.TYPE) return Long.valueOf(0L);
+              return null;
+            });
+  }
+
+  private interface ExitRecorder {
+    void record(int exitCode);
+  }
+}

btrace-client/src/main/java/io/btrace/client/Client.java

@@ -301,9 +301,8 @@ public byte[] compile(String fileName, String classPath) {
    * Scans the extensions directory and returns a classpath string with all API JARs.
    *
    * <p>Looks in the following locations (first match wins): - System property 'btrace.libs'
-   * (assumed to point to BTrace libs directory); scans sibling 'extensions/' - Derived from
-   * java.class.path entry containing 'btrace-client' (for dist layouts); scans sibling
-   * 'extensions/'
+   * (assumed to point to BTrace libs directory); scans sibling 'extensions/' - Derived from a
+   * btrace-client JAR on java.class.path (for legacy dist layouts); scans sibling 'extensions/'
    */
   private String getExtensionApiClasspath() {
     try {
@@ -318,19 +317,12 @@ private String getExtensionApiClasspath() {
         }
       }
 
-      // 2) Fall back to locating btrace-client in the classpath
-      String classPath = System.getProperty("java.class.path");
-      for (String entry : classPath.split(File.pathSeparator)) {
-        if (entry.contains("btrace-client")) {
-          File clientPath = new File(entry);
-          File libsDir = clientPath.getParentFile();
-          if (libsDir != null) {
-            File btraceHome = libsDir.getParentFile();
-            String cp = scanExtensionsDir(new File(btraceHome, "extensions"));
-            if (!cp.isEmpty()) {
-              return cp;
-            }
-          }
+      // 2) Fall back to locating the btrace-client JAR in the classpath
+      File btraceHome = btraceHomeFromClassPath(System.getProperty("java.class.path"));
+      if (btraceHome != null) {
+        String cp = scanExtensionsDir(new File(btraceHome, "extensions"));
+        if (!cp.isEmpty()) {
+          return cp;
         }
       }
       // 3) Fall back to extensions embedded in the agent JAR itself. A Maven or jbang user has no
@@ -348,6 +340,31 @@ private String getExtensionApiClasspath() {
     return "";
   }
 
+  /**
+   * Returns the BTrace home of a {@code <home>/libs/btrace-client[-<version>].jar} entry on {@code
+   * classPath}. Only the entry's file name is matched, so a directory that merely contains
+   * "btrace-client" in its path, such as a checkout or worktree, is never taken for the client.
+   *
+   * @return the home directory, or {@code null} when no such JAR is on the classpath
+   */
+  static File btraceHomeFromClassPath(String classPath) {
+    if (classPath == null) {
+      return null;
+    }
+    for (String entry : classPath.split(File.pathSeparator)) {
+      File clientJar = new File(entry);
+      String name = clientJar.getName();
+      boolean isClientJar =
+          name.endsWith(".jar")
+              && (name.equals("btrace-client.jar") || name.startsWith("btrace-client-"));
+      File libsDir = clientJar.getAbsoluteFile().getParentFile();
+      if (isClientJar && libsDir != null && libsDir.getParentFile() != null) {
+        return libsDir.getParentFile();
+      }
+    }
+    return null;
+  }
+
   /**
    * Returns the agent JAR as a compile classpath entry when it declares embedded extensions.
    *
@@ -1130,6 +1147,22 @@ void connectAndListFailedExtensions(String host, CommandListener listener) throw
     }
   }
 
+  /**
+   * Reconnects to a probe that a client detached from and stops it, as {@code -r <probe-id> exit}
+   * does. Returns once the agent confirms the exit; the connection is always closed.
+   *
+   * @param host the agent host
+   * @param probeId a probe id as reported by {@link #connectAndListProbes}
+   * @throws IOException if the agent is unavailable or does not know {@code probeId}
+   */
+  public void connectAndExitProbe(String host, String probeId) throws IOException {
+    try {
+      reconnect(host, probeId, cmd -> {}, new String[] {"exit", null});
+    } finally {
+      closeQuietly();
+    }
+  }
+
   void reconnect(String host, String resumeProbe, CommandListener listener, String[] command)
       throws IOException {
     if (sock != null) {