From 5dde7280ec4f2005dbb41a6368b915f9a8f5e8cd Mon Sep 17 00:00:00 2001
From: Sean Tang <171081544+seant-aws@users.noreply.github.com>
Date: Tue, 14 Jul 2026 17:44:15 +0000
Subject: [PATCH] fix: poll through CANCELING state to deliver terminal
 callback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The _poll_for_completion thread exited the polling loop as soon as
the session state transitioned away from RUNNING.  When cancel_action()
is called, the pyo3 binding immediately sets the snapshot state to
CANCELING — causing the poller to exit before the subprocess has been
killed and before the Rust callback delivers the terminal CANCELED
ActionStatus.

The result: the worker agent never receives the CANCELED callback,
leaving the action stuck in CANCELING forever from the service's
perspective.

Fix: treat CANCELING as an in-progress state (like RUNNING) and
continue polling.  The poller now exits only when the state reaches
READY or READY_ENDING, at which point action_status contains the
correct terminal state.

Signed-off-by: Sean Tang <seant-aws@users.noreply.github.com>
Signed-off-by: Sean Tang <171081544+seant-aws@users.noreply.github.com>
---
 src/openjd/sessions/_v1/_session.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/openjd/sessions/_v1/_session.py b/src/openjd/sessions/_v1/_session.py
index c0301726..fb1aa563 100644
--- a/src/openjd/sessions/_v1/_session.py
+++ b/src/openjd/sessions/_v1/_session.py
@@ -186,7 +186,7 @@ def observable(s: ActionStatus) -> tuple:
             while True:
                 state = self._rust_session.state
                 status = self._rust_session.action_status
-                if state == SessionState.RUNNING:
+                if state in (SessionState.RUNNING, SessionState.CANCELING):
                     if status is not None and self._callback:
                         if not reported_running:
                             # Defensive: the synchronous initial callback
@@ -201,7 +201,7 @@ def observable(s: ActionStatus) -> tuple:
                             self._callback(self._session_id, status)
                     time.sleep(0.05)
                     continue
-                # Not RUNNING anymore — action is done.
+                # Not RUNNING/CANCELING anymore — action is done.
                 # If we never saw RUNNING (e.g. Rust finished before we got here),
                 # still report the RUNNING transition first, so the agent state
                 # machine sees Start → End in order.