From 627fcb19dab32ec1f07d7e9ffa0780483120a7e0 Mon Sep 17 00:00:00 2001
From: Jason Morcos <Jason-Morcos@users.noreply.github.com>
Date: Sat, 15 Aug 2026 11:44:24 -0700
Subject: [PATCH] fix(protocol): keep completed sessions on late cancel

---
 smartthings_local/protocol/dtls_session.py |  3 +-
 tests/test_session_interruption.py         | 40 ++++++++++++++++++++--
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/smartthings_local/protocol/dtls_session.py b/smartthings_local/protocol/dtls_session.py
index 2da687c..d2e1499 100644
--- a/smartthings_local/protocol/dtls_session.py
+++ b/smartthings_local/protocol/dtls_session.py
@@ -382,6 +382,7 @@ class DtlsCoapSession:
         io_failed = False
         cancelled = False
         interrupted = False
+        completed = False
         try:
             try:
                 completed = _drive_dtls_handshake(
@@ -403,7 +404,7 @@ class DtlsCoapSession:
         finally:
             if wake_subscription is not None:
                 interrupted = cancel._unsubscribe(*wake_subscription)
-        if cancelled or interrupted:
+        if cancelled or (interrupted and not completed):
             sock.close()
             raise SessionClosedError()
         if backend_failed:
diff --git a/tests/test_session_interruption.py b/tests/test_session_interruption.py
index 519d673..232349b 100644
--- a/tests/test_session_interruption.py
+++ b/tests/test_session_interruption.py
@@ -220,9 +220,44 @@ def test_cancel_wakes_blocked_connect_without_poll_latency(monkeypatch):
         peer.close()
 
 
-def test_cancel_wins_race_with_reported_handshake_success(monkeypatch):
+def test_reported_handshake_success_wins_cancel_during_unsubscribe(
+    monkeypatch,
+):
+    class CancelDuringUnsubscribe(ConnectCancellation):
+        def _unsubscribe(self, reader, writer):
+            self.set()
+            return super()._unsubscribe(reader, writer)
+
+    cancel = CancelDuringUnsubscribe()
+    connection = _Connection(succeed=True)
+    data_socket, peer = socket.socketpair()
+    endpoint = _install_connection(monkeypatch, connection, data_socket)
+    session = _session()
+
+    try:
+        session.connect(cancel=cancel)
+        assert cancel.is_set()
+        assert session.sock is data_socket
+        assert session.conn is connection
+        assert session.endpoint is endpoint
+        assert session.dest == endpoint.sockaddr
+        assert not cancel._writers
+    finally:
+        session.close()
+        peer.close()
+
+
+def test_cancel_during_backend_failure_does_not_read_unset_completion(
+    monkeypatch,
+):
     cancel = ConnectCancellation()
-    connection = _Connection(on_success=cancel.set, succeed=True)
+    connection = _Connection()
+
+    def fail_after_cancel():
+        cancel.set()
+        raise SSL.Error("synthetic backend failure")
+
+    connection.do_handshake = fail_after_cancel
     data_socket, peer = socket.socketpair()
     _install_connection(monkeypatch, connection, data_socket)
     session = _session()
@@ -233,6 +268,7 @@ def test_cancel_wins_race_with_reported_handshake_success(monkeypatch):
         assert data_socket.fileno() == -1
         assert session.sock is None
         assert session.conn is None
+        assert not cancel._writers
     finally:
         peer.close()