fix(agent): recover JDBC connections after query timeouts

agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java

@@ -29,6 +29,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
     private boolean requestActive;
     private boolean leasePinnedAtRequestStart;
     private boolean sessionAffinity;
+    private boolean pooledConnectionPoisoned;
 
     @Override
     public final Connection getConnection() {
@@ -47,6 +48,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
             sessionAffinity = false;
             requestActive = false;
             leasePinnedAtRequestStart = false;
+            pooledConnectionPoisoned = false;
             loadDriver(params);
             configuredDatabase = params.getDatabase();
             connectParams = params;
@@ -256,6 +258,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
             sessionAffinity = false;
             requestActive = false;
             leasePinnedAtRequestStart = false;
+            pooledConnectionPoisoned = false;
             identifierQuote = "";
         });
     }
@@ -271,24 +274,54 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
         return poolRegistry != null;
     }
 
-    final synchronized void beginPooledRequest() throws Exception {
-        if (poolRegistry == null) {
-            return;
-        }
-        if (connectParams == null) {
-            throw new IllegalStateException("Not connected");
-        }
-        requestActive = true;
-        leasePinnedAtRequestStart = pooledLease != null;
-        if (pooledLease == null) {
-            try {
-                pooledLease = borrowPooledConnection();
-            } catch (Exception error) {
-                requestActive = false;
-                throw error;
+    final synchronized boolean quarantinePooledConnection() {
+        pooledConnectionPoisoned = true;
+        return requestActive && pooledLease != null && pooledLease.quarantine();
+    }
+
+    final void beginPooledRequest() throws Exception {
+        ConnectParams params;
+        String identity;
+        JdbcConnectionPoolRegistry registry;
+        synchronized (this) {
+            if (poolRegistry == null) {
+                return;
             }
+            if (connectParams == null || poolIdentity == null) {
+                throw new IllegalStateException("Not connected");
+            }
+            requestActive = true;
+            leasePinnedAtRequestStart = pooledLease != null;
+            if (pooledLease != null) {
+                connection = pooledLease.connection();
+                return;
+            }
+            params = connectParams;
+            identity = poolIdentity;
+            registry = poolRegistry;
+        }
+
+        JdbcConnectionPoolRegistry.Lease borrowed;
+        try {
+            borrowed = borrowPooledConnection(registry, identity, params);
+        } catch (Exception error) {
+            synchronized (this) {
+                requestActive = false;
+                leasePinnedAtRequestStart = false;
+            }
+            throw error;
+        }
+
+        synchronized (this) {
+            if (pooledConnectionPoisoned) {
+                requestActive = false;
+                leasePinnedAtRequestStart = false;
+                borrowed.evict();
+                throw new IllegalStateException("JDBC Session was quarantined while waiting for a connection");
+            }
+            pooledLease = borrowed;
+            connection = borrowed.connection();
         }
-        connection = pooledLease.connection();
     }
 
     final synchronized void finishPooledRequest(
@@ -301,6 +334,10 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
             return;
         }
         requestActive = false;
+        if (pooledConnectionPoisoned) {
+            releasePooledConnection(true);
+            return;
+        }
         if (succeeded && requiresSessionAffinity) {
             sessionAffinity = true;
             JdbcSchemaSwitcher.forget(connection);
@@ -322,7 +359,14 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
     }
 
     final synchronized void releaseIdlePooledConnection(JdbcExecutor executor) {
-        if (poolRegistry == null || requestActive || sessionAffinity || pooledLease == null) {
+        if (poolRegistry == null || requestActive || pooledLease == null) {
+            return;
+        }
+        if (pooledConnectionPoisoned) {
+            releasePooledConnection(true);
+            return;
+        }
+        if (sessionAffinity) {
             return;
         }
         if (executor.hasOpenSessions() || executor.hasActiveStatements()) {
@@ -489,18 +533,20 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
 
     private Connection openInitializedConnection(ConnectParams params) throws Exception {
         Connection opened = openConnection(params);
-        boolean initialized = false;
         try {
             afterPhysicalConnect(params, opened);
-            initialized = true;
             return opened;
-        } finally {
-            if (!initialized) {
-                try {
-                    opened.close();
-                } catch (Exception ignored) {
-                }
+        } catch (Exception error) {
+            try {
+                opened.close();
+            } catch (Exception closeError) {
+                error.addSuppressed(closeError);
+                throw AgentRpcError.resource(
+                    "close",
+                    new JdbcConnectionPoolRegistry.PhysicalConnectionStateUnknownException(error)
+                );
             }
+            throw error;
         }
     }
 
@@ -509,12 +555,24 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
         if (params == null || poolIdentity == null || poolRegistry == null) {
             throw new IllegalStateException("Not connected");
         }
-        return poolRegistry.borrow(poolIdentity, () -> openInitializedConnection(params));
+        return borrowPooledConnection(poolRegistry, poolIdentity, params);
+    }
+
+    private JdbcConnectionPoolRegistry.Lease borrowPooledConnection(
+        JdbcConnectionPoolRegistry registry,
+        String identity,
+        ConnectParams params
+    ) throws Exception {
+        return registry.borrow(
+            identity,
+            JdbcSessionRole.from(params.getSessionRole()),
+            () -> openInitializedConnection(params)
+        );
     }
 
     private void closeCurrentConnection() throws Exception {
         if (pooledLease != null) {
-            boolean evict = sessionAffinity || !preparePooledConnectionForReturn();
+            boolean evict = pooledConnectionPoisoned || sessionAffinity || !preparePooledConnectionForReturn();
             releasePooledConnection(evict);
         } else if (connection != null) {
             connection.close();

agents/common/src/main/java/com/dbx/agent/AgentRpcError.java

@@ -0,0 +1,135 @@
+package com.dbx.agent;
+
+import com.google.gson.JsonObject;
+
+import java.sql.SQLException;
+import java.sql.SQLRecoverableException;
+import java.sql.SQLTransientConnectionException;
+import java.util.Locale;
+
+final class AgentRpcError extends RuntimeException {
+    private final String category;
+    private final boolean retryable;
+    private final String disposition;
+    private final String stage;
+
+    private AgentRpcError(
+        String message,
+        String category,
+        boolean retryable,
+        String disposition,
+        String stage,
+        Throwable cause
+    ) {
+        super(message, cause);
+        this.category = category;
+        this.retryable = retryable;
+        this.disposition = disposition;
+        this.stage = stage;
+    }
+
+    static AgentRpcError resource(String stage, Throwable cause) {
+        return new AgentRpcError(
+            "Agent runtime resource limit reached",
+            "resource",
+            false,
+            "replace_runtime",
+            stage,
+            cause
+        );
+    }
+
+    static AgentRpcError backpressure(String stage, Throwable cause) {
+        return new AgentRpcError(
+            "Agent request capacity is temporarily exhausted",
+            "resource",
+            true,
+            "keep",
+            stage,
+            cause
+        );
+    }
+
+    static JsonObject toJson(Throwable error, String method, String agentSessionId) {
+        AgentRpcError classified = classify(error, method);
+        JsonObject rpcError = new JsonObject();
+        rpcError.addProperty("code", -1);
+        rpcError.addProperty("message", message(error));
+        JsonObject data = new JsonObject();
+        data.addProperty("category", classified.category);
+        data.addProperty("retryable", classified.retryable);
+        data.addProperty("sessionDisposition", classified.disposition);
+        data.addProperty("stage", classified.stage);
+        if (agentSessionId != null && !agentSessionId.trim().isEmpty()) {
+            data.addProperty("agentSessionId", agentSessionId);
+        }
+        rpcError.add("data", data);
+        return rpcError;
+    }
+
+    private static AgentRpcError classify(Throwable error, String method) {
+        AgentRpcError explicit = find(error, AgentRpcError.class);
+        if (explicit != null) {
+            return explicit;
+        }
+        SQLException sqlError = find(error, SQLException.class);
+        if (sqlError != null) {
+            String sqlState = sqlError.getSQLState();
+            String stage = stage(method);
+            boolean connectionError = "connect".equals(stage)
+                || "validate".equals(stage)
+                || sqlError instanceof SQLRecoverableException
+                || sqlError instanceof SQLTransientConnectionException
+                || (sqlState != null && sqlState.toUpperCase(Locale.ROOT).startsWith("08"));
+            boolean operationRetryable = connectionError && ("connect".equals(stage) || "validate".equals(stage));
+            String disposition = connectionError && !"connect".equals(stage) ? "quarantine" : "keep";
+            return new AgentRpcError(
+                message(error),
+                connectionError ? "connection" : "sql",
+                operationRetryable,
+                disposition,
+                stage,
+                error
+            );
+        }
+        return new AgentRpcError(message(error), "protocol", false, "keep", stage(method), error);
+    }
+
+    private static String stage(String method) {
+        if (method == null) {
+            return "request";
+        }
+        if (AgentProtocol.METHOD_CONNECT.equals(method) || AgentProtocol.METHOD_OPEN_SESSION.equals(method)) {
+            return "connect";
+        }
+        if (AgentProtocol.METHOD_VALIDATE_CONNECTION.equals(method) || AgentProtocol.METHOD_VALIDATE_SESSION.equals(method)) {
+            return "validate";
+        }
+        if (AgentProtocol.METHOD_CANCEL_SESSION.equals(method)) {
+            return "cancel";
+        }
+        if (AgentProtocol.METHOD_CLOSE_SESSION.equals(method) || AgentProtocol.METHOD_DISCONNECT.equals(method)) {
+            return "close";
+        }
+        if (AgentProtocol.METHOD_FETCH_QUERY_PAGE.equals(method)
+            || AgentProtocol.METHOD_FETCH_TABLE_READ_PAGE.equals(method)) {
+            return "fetch";
+        }
+        return "execute";
+    }
+
+    private static String message(Throwable error) {
+        return error.getMessage() == null ? error.toString() : error.getMessage();
+    }
+
+    private static <T extends Throwable> T find(Throwable error, Class<T> type) {
+        Throwable current = error;
+        while (current != null) {
+            if (type.isInstance(current)) {
+                return type.cast(current);
+            }
+            current = current.getCause();
+        }
+        return null;
+    }
+}

agents/common/src/main/java/com/dbx/agent/ConnectParams.java

@@ -22,6 +22,7 @@ public final class ConnectParams {
     private String client_key_path;
     private String gbase_server;
     private String informix_server;
+    private String sessionRole;
 
     public ConnectParams() {
         this("", 0, "", "", "", "", "", false, "", Collections.emptyList());
@@ -200,6 +201,14 @@ public final class ConnectParams {
         this.informix_server = informix_server;
     }
 
+    public String getSessionRole() {
+        return sessionRole;
+    }
+
+    public void setSessionRole(String sessionRole) {
+        this.sessionRole = sessionRole;
+    }
+
     @Override
     public boolean equals(Object other) {
         if (this == other) return true;
@@ -221,13 +230,14 @@ public final class ConnectParams {
             && Objects.equals(client_cert_path, that.client_cert_path)
             && Objects.equals(client_key_path, that.client_key_path)
             && Objects.equals(gbase_server, that.gbase_server)
-            && Objects.equals(informix_server, that.informix_server);
+            && Objects.equals(informix_server, that.informix_server)
+            && Objects.equals(sessionRole, that.sessionRole);
     }
 
     @Override
     public int hashCode() {
         return Objects.hash(host, port, database, username, password, url_params, connection_string,
-            port_explicit, mysql_compat_mode, jdbc_driver_class, jdbc_driver_paths, ssl, ca_cert_path, client_cert_path, client_key_path, gbase_server, informix_server);
+            port_explicit, mysql_compat_mode, jdbc_driver_class, jdbc_driver_paths, ssl, ca_cert_path, client_cert_path, client_key_path, gbase_server, informix_server, sessionRole);
     }
 
     @Override
@@ -249,6 +259,7 @@ public final class ConnectParams {
             + ", client_key_path=" + client_key_path
             + ", gbase_server=" + gbase_server
             + ", informix_server=" + informix_server
+            + ", sessionRole=" + sessionRole
             + ")";
     }
 }

agents/common/src/main/java/com/dbx/agent/JdbcSessionRole.java

@@ -0,0 +1,10 @@
+package com.dbx.agent;
+
+enum JdbcSessionRole {
+    WORKLOAD,
+    METADATA;
+
+    static JdbcSessionRole from(String value) {
+        return "metadata".equalsIgnoreCase(value == null ? "" : value.trim()) ? METADATA : WORKLOAD;
+    }
+}

agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java

@@ -121,6 +121,11 @@ public final class JsonRpcServer {
         }
     }
 
+    boolean quarantine() {
+        AbstractJdbcAgent jdbcAgent = pooledJdbcAgent();
+        return jdbcAgent != null && jdbcAgent.quarantinePooledConnection();
+    }
+
     void expireIdleResources() {
         jdbcExecutor.expireIdleResources();
         releaseIdlePooledConnection();

agents/common/src/main/java/com/dbx/agent/MultiSessionJsonRpcServer.java

@@ -13,22 +13,30 @@ import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 public final class MultiSessionJsonRpcServer implements AutoCloseable {
     private static final String LEGACY_SESSION_ID = "__legacy__";
     private static final int MAX_SESSIONS = 256;
+    private static final int MAX_REQUEST_THREADS = 64;
+    private static final int MAX_CLEANUP_THREADS = 16;
     private static final long MAINTENANCE_INTERVAL_MILLIS = 60_000L;
 
     private final Supplier<? extends DatabaseAgent> agentFactory;
     private final Supplier<? extends SessionRpcHandler> sessionHandlerFactory;
     private final Map<String, Session> sessions = new ConcurrentHashMap<>();
-    private final ExecutorService requests = Executors.newCachedThreadPool();
+    private final ExecutorService requests;
+    private final ExecutorService cleanup;
     private final JdbcConnectionPoolRegistry poolRegistry;
     private final Gson gson = new Gson();
     private final PrintStream protocolOutput = System.out;
@@ -38,29 +46,43 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
     private ScheduledExecutorService maintenance;
 
     public MultiSessionJsonRpcServer(Supplier<? extends DatabaseAgent> agentFactory) {
-        this(agentFactory, new JdbcConnectionPoolRegistry());
+        this(agentFactory, new JdbcConnectionPoolRegistry(), RuntimeLimits.defaults());
     }
 
     MultiSessionJsonRpcServer(
         Supplier<? extends DatabaseAgent> agentFactory,
         JdbcConnectionPoolRegistry.PoolSettings poolSettings
     ) {
-        this(agentFactory, new JdbcConnectionPoolRegistry(poolSettings));
+        this(agentFactory, new JdbcConnectionPoolRegistry(poolSettings), RuntimeLimits.defaults());
+    }
+
+    MultiSessionJsonRpcServer(
+        Supplier<? extends DatabaseAgent> agentFactory,
+        JdbcConnectionPoolRegistry.PoolSettings poolSettings,
+        RuntimeLimits runtimeLimits
+    ) {
+        this(agentFactory, new JdbcConnectionPoolRegistry(poolSettings), runtimeLimits);
     }
 
     private MultiSessionJsonRpcServer(
         Supplier<? extends DatabaseAgent> agentFactory,
-        JdbcConnectionPoolRegistry poolRegistry
+        JdbcConnectionPoolRegistry poolRegistry,
+        RuntimeLimits runtimeLimits
     ) {
         this.agentFactory = agentFactory;
         this.sessionHandlerFactory = null;
         this.poolRegistry = poolRegistry;
+        this.requests = boundedExecutor(runtimeLimits.maximumRequestThreads, "dbx-agent-request");
+        this.cleanup = boundedExecutor(runtimeLimits.maximumCleanupThreads, "dbx-agent-cleanup");
     }
 
     private MultiSessionJsonRpcServer(Supplier<? extends SessionRpcHandler> sessionHandlerFactory, boolean customHandler) {
         this.agentFactory = null;
         this.sessionHandlerFactory = sessionHandlerFactory;
         this.poolRegistry = new JdbcConnectionPoolRegistry();
+        RuntimeLimits runtimeLimits = RuntimeLimits.defaults();
+        this.requests = boundedExecutor(runtimeLimits.maximumRequestThreads, "dbx-agent-request");
+        this.cleanup = boundedExecutor(runtimeLimits.maximumCleanupThreads, "dbx-agent-cleanup");
     }
 
     /** Creates a protocol v2 server for a non-JDBC, session-scoped agent. */
@@ -82,7 +104,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
                     writeResponse(handleRequest(request));
                     return;
                 }
-                requests.submit(() -> writeResponse(handleRequest(request)));
+                executeRequest(request, this::writeResponse);
             }
         } catch (Exception e) {
             throw new RuntimeException(e);
@@ -95,6 +117,14 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
         return gson.toJson(handleRequest(JsonParser.parseString(line).getAsJsonObject()));
     }
 
+    void executeRequest(JsonObject request, Consumer<JsonObject> responseConsumer) {
+        try {
+            requests.execute(() -> responseConsumer.accept(handleRequest(request)));
+        } catch (RejectedExecutionException error) {
+            responseConsumer.accept(errorResponse(request.get("id"), AgentRpcError.backpressure("request", error)));
+        }
+    }
+
     private JsonObject handleRequest(JsonObject request) {
         JsonElement id = request.get("id");
         String method = request.get("method").getAsString();
@@ -134,10 +164,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
             }
             response.add("result", gson.toJsonTree(result));
         } catch (Throwable error) {
-            JsonObject rpcError = new JsonObject();
-            rpcError.addProperty("code", -1);
-            rpcError.addProperty("message", error.getMessage() == null ? error.toString() : error.getMessage());
-            response.add("error", rpcError);
+            response.add("error", AgentRpcError.toJson(error, method, stringOrNull(params, "agentSessionId")));
         }
         return response;
     }
@@ -165,7 +192,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
             return session.connect(params);
         } catch (Exception error) {
             sessions.remove(sessionId, session);
-            session.close();
+            session.quarantineAndClose(cleanup);
             throw error;
         }
     }
@@ -173,7 +200,13 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
     private Object closeSession(String sessionId) {
         Session session = sessions.remove(sessionId);
         if (session != null) {
-            session.close();
+            boolean replaceRuntime = session.quarantineAndClose(cleanup);
+            if (replaceRuntime) {
+                throw AgentRpcError.resource(
+                    "close",
+                    new IllegalStateException("JDBC quarantine operation limit reached")
+                );
+            }
         }
         return Collections.singletonMap("ok", true);
     }
@@ -206,7 +239,11 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
 
     private void closeAllSessions() {
         for (String sessionId : sessions.keySet()) {
-            closeSession(sessionId);
+            try {
+                closeSession(sessionId);
+            } catch (AgentRpcError ignored) {
+                // Sessions are already detached; process shutdown remains the final cleanup boundary.
+            }
         }
     }
 
@@ -263,7 +300,16 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
             }
         }
         closeAllSessions();
-        requests.shutdown();
+        requests.shutdownNow();
+        cleanup.shutdown();
+        try {
+            if (!cleanup.awaitTermination(2, TimeUnit.SECONDS)) {
+                cleanup.shutdownNow();
+            }
+        } catch (InterruptedException error) {
+            Thread.currentThread().interrupt();
+            cleanup.shutdownNow();
+        }
         poolRegistry.close();
     }
 
@@ -275,6 +321,47 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
         };
     }
 
+    private static ExecutorService boundedExecutor(int maximumThreads, String threadName) {
+        return new ThreadPoolExecutor(
+            0,
+            maximumThreads,
+            60L,
+            TimeUnit.SECONDS,
+            new SynchronousQueue<>(),
+            daemonThreadFactory(threadName),
+            new ThreadPoolExecutor.AbortPolicy()
+        );
+    }
+
+    private static JsonObject errorResponse(JsonElement id, Throwable error) {
+        JsonObject response = new JsonObject();
+        response.addProperty("jsonrpc", "2.0");
+        response.add("id", id);
+        response.add("error", AgentRpcError.toJson(error, "request", null));
+        return response;
+    }
+
+    private static String stringOrNull(JsonObject params, String key) {
+        return params.has(key) && !params.get(key).isJsonNull() ? params.get(key).getAsString() : null;
+    }
+
+    static final class RuntimeLimits {
+        private final int maximumRequestThreads;
+        private final int maximumCleanupThreads;
+
+        RuntimeLimits(int maximumRequestThreads, int maximumCleanupThreads) {
+            if (maximumRequestThreads <= 0 || maximumCleanupThreads <= 0) {
+                throw new IllegalArgumentException("Agent runtime thread limits must be positive");
+            }
+            this.maximumRequestThreads = maximumRequestThreads;
+            this.maximumCleanupThreads = maximumCleanupThreads;
+        }
+
+        private static RuntimeLimits defaults() {
+            return new RuntimeLimits(MAX_REQUEST_THREADS, MAX_CLEANUP_THREADS);
+        }
+    }
+
     private static String requiredSessionId(JsonObject params) {
         if (!params.has("agentSessionId") || params.get("agentSessionId").getAsString().trim().isEmpty()) {
             throw new IllegalArgumentException("agentSessionId is required");
@@ -293,6 +380,8 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
         private final JsonRpcServer server;
         private final SessionRpcHandler handler;
         private final ReentrantLock lock = new ReentrantLock();
+        private final AtomicReference<State> state = new AtomicReference<>(State.ACTIVE);
+        private final AtomicBoolean cleanupScheduled = new AtomicBoolean();
 
         private Session(JsonRpcServer server) {
             this.server = server;
@@ -305,8 +394,10 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
         }
 
         private Object handle(String method, JsonObject params) throws Exception {
+            requireActive();
             lock.lock();
             try {
+                requireActive();
                 return handler == null ? server.dispatchForRuntime(method, params) : handler.handle(method, params);
             } finally {
                 lock.unlock();
@@ -314,8 +405,10 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
         }
 
         private Object connect(JsonObject params) throws Exception {
+            requireActive();
             lock.lock();
             try {
+                requireActive();
                 return handler == null
                     ? server.dispatchForRuntime(AgentProtocol.METHOD_CONNECT, params)
                     : handler.connect(params);
@@ -324,9 +417,31 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
             }
         }
 
+        private boolean quarantineAndClose(ExecutorService cleanup) {
+            state.compareAndSet(State.ACTIVE, State.QUARANTINED);
+            boolean replaceRuntime = server != null && server.quarantine();
+            if (!cleanupScheduled.compareAndSet(false, true)) {
+                return replaceRuntime;
+            }
+            try {
+                cleanup.execute(this::closeWhenIdle);
+            } catch (RejectedExecutionException error) {
+                throw AgentRpcError.resource("close", error);
+            }
+            return replaceRuntime;
+        }
+
+        private void closeWhenIdle() {
+            close();
+        }
+
         private void close() {
+            state.compareAndSet(State.ACTIVE, State.QUARANTINED);
             lock.lock();
             try {
+                if (state.get() == State.CLOSED) {
+                    return;
+                }
                 if (handler != null) {
                     handler.close();
                 } else {
@@ -334,6 +449,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
                 }
             } catch (Exception ignored) {
             } finally {
+                state.set(State.CLOSED);
                 lock.unlock();
             }
         }
@@ -347,7 +463,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
         }
 
         private void expireIdleResources() {
-            if (handler != null) {
+            if (state.get() != State.ACTIVE || handler != null) {
                 return;
             }
             if (!lock.tryLock()) {
@@ -361,7 +477,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
         }
 
         private void expireIdleResources(long nowMillis, long idleTimeoutMillis) {
-            if (handler != null) {
+            if (state.get() != State.ACTIVE || handler != null) {
                 return;
             }
             if (!lock.tryLock()) {
@@ -373,5 +489,17 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
                 lock.unlock();
             }
         }
+
+        private void requireActive() {
+            if (state.get() != State.ACTIVE) {
+                throw new IllegalStateException("Agent session is quarantined");
+            }
+        }
+
+        private enum State {
+            ACTIVE,
+            QUARANTINED,
+            CLOSED
+        }
     }
 }