From ba7a6867139a3e6abfe782e499e6ff6cc85a8bb0 Mon Sep 17 00:00:00 2001
From: Victor Skvortsov <vds003@gmail.com>
Date: Tue, 9 Jun 2026 16:33:36 +0500
Subject: [PATCH] Fix dropping ssh connections to non-provisioned terminating
 instances

---
 .../background/pipeline_tasks/instances/termination.py      | 3 +--
 src/dstack/_internal/server/services/runner/pool.py         | 6 ++++++
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py
index a4bf6d329..c4cce13b2 100644
--- a/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py
+++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py
@@ -12,7 +12,6 @@
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services.instances import get_instance_provisioning_data
 from dstack._internal.server.services.runner.pool import (
-    InstanceConnectionKey,
     instance_connection_pool,
 )
 from dstack._internal.utils.common import get_current_datetime, run_async
@@ -82,7 +81,7 @@ async def terminate_instance(instance_model: InstanceModel) -> ProcessResult:
                 )
 
     if job_provisioning_data is not None:
-        instance_connection_pool.drop(InstanceConnectionKey.from_jpd(job_provisioning_data))
+        instance_connection_pool.drop_by_jpd(job_provisioning_data)
 
     result.instance_update_map["deleted"] = True
     result.instance_update_map["deleted_at"] = NOW_PLACEHOLDER
diff --git a/src/dstack/_internal/server/services/runner/pool.py b/src/dstack/_internal/server/services/runner/pool.py
index b91d1d312..e3a012967 100644
--- a/src/dstack/_internal/server/services/runner/pool.py
+++ b/src/dstack/_internal/server/services/runner/pool.py
@@ -130,6 +130,12 @@ def drop(self, key: InstanceConnectionKey) -> None:
             except Exception:
                 logger.exception("Failed to close instance connection %s", key)
 
+    def drop_by_jpd(self, jpd: JobProvisioningData, jrd: Optional[JobRuntimeData] = None):
+        if jpd.hostname is None or jpd.ssh_port is None:
+            return
+        key = InstanceConnectionKey.from_jpd(jpd, jrd)
+        self.drop(key)
+
     def startup_cleanup(self) -> None:
         """
         Removes connection dirs left by a previous server process (e.g. after SIGKILL).