Harden Python goal cancellation

sdk/python/src/openai_codex/_goal.py

@@ -1,3 +1,4 @@
+import asyncio
 import queue
 import threading
 import time
@@ -233,7 +234,7 @@ class _GoalStreamCursor:
     active: bool = False
     cleared: bool = False
 
-    def process(self, notification: Notification) -> tuple[list[Notification], bool, bool]:
+    def process(self, notification: Notification) -> tuple[list[Notification], bool]:
         logical_turn_id = self.state.logical_turn_id
         if logical_turn_id is None:
             raise RuntimeError("goal operation has not been bound to a logical turn id")
@@ -242,27 +243,23 @@ class _GoalStreamCursor:
         if isinstance(payload, TurnStartedNotification):
             self.active = True
             if self.started is not None:
-                return [], False, False
+                return [], False
             self.started = payload.turn
-            return [_logical_notification(notification, logical_turn_id)], False, False
+            return [_logical_notification(notification, logical_turn_id)], False
 
         if isinstance(payload, TurnCompletedNotification):
             self.active = False
             self.last_completed = payload
             if payload.turn.status in {TurnStatus.failed, TurnStatus.interrupted}:
-                return (
-                    [self._completion(notification.method, payload)],
-                    True,
-                    (payload.turn.status == TurnStatus.failed),
-                )
+                return [self._completion(notification.method, payload)], True
             if self.status is None and not self.cleared:
                 raise RuntimeError(
                     "the connected Codex runtime did not activate goal mode for this turn"
                 )
             if self.cleared or _terminal_goal_status(self.status):
                 self.state.finish()
-                return [self._completion(notification.method, payload)], True, False
-            return [], False, False
+                return [self._completion(notification.method, payload)], True
+            return [], False
 
         events = [_logical_notification(notification, logical_turn_id)]
         if isinstance(payload, ThreadGoalUpdatedNotification):
@@ -281,8 +278,8 @@ class _GoalStreamCursor:
         ):
             self.state.finish()
             events.append(self._completion("turn/completed", self.last_completed))
-            return events, True, False
-        return events, False, False
+            return events, True
+        return events, False
 
     def _completion(
         self,
@@ -310,7 +307,7 @@ class _GoalNotificationStream(Iterator[Notification]):
     state: _GoalOperationState
     next_notification: Callable[[], Notification]
     unregister: Callable[[], None]
-    stop_failed_goal: Callable[[], None]
+    cancel_goal: Callable[[], None]
     _cursor: _GoalStreamCursor = field(init=False)
     _pending: deque[Notification] = field(default_factory=deque)
     _closed: bool = False
@@ -326,9 +323,7 @@ class _GoalNotificationStream(Iterator[Notification]):
             raise StopIteration
         try:
             while not self._pending:
-                events, completed, failed = self._cursor.process(self.next_notification())
-                if failed:
-                    self.stop_failed_goal()
+                events, completed = self._cursor.process(self.next_notification())
                 self._pending.extend(events)
                 if completed:
                     self._finish()
@@ -336,6 +331,10 @@ class _GoalNotificationStream(Iterator[Notification]):
         except _GoalStreamClosed:
             self.close()
             raise StopIteration from None
+        except KeyboardInterrupt:
+            self.cancel_goal()
+            self.close()
+            raise
         except BaseException:
             self.close()
             raise
@@ -359,7 +358,7 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]):
     state: _GoalOperationState
     next_notification: Callable[[], Awaitable[Notification]]
     unregister: Callable[[], None]
-    stop_failed_goal: Callable[[], Awaitable[None]]
+    cancel_goal: Callable[[], Awaitable[None]]
     _cursor: _GoalStreamCursor = field(init=False)
     _pending: deque[Notification] = field(default_factory=deque)
     _closed: bool = False
@@ -375,9 +374,7 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]):
             raise StopAsyncIteration
         try:
             while not self._pending:
-                events, completed, failed = self._cursor.process(await self.next_notification())
-                if failed:
-                    await self.stop_failed_goal()
+                events, completed = self._cursor.process(await self.next_notification())
                 self._pending.extend(events)
                 if completed:
                     self._finish()
@@ -385,6 +382,10 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]):
         except _GoalStreamClosed:
             await self.aclose()
             raise StopAsyncIteration from None
+        except asyncio.CancelledError:
+            await self.cancel_goal()
+            await self.aclose()
+            raise
         except BaseException:
             await self.aclose()
             raise

sdk/python/src/openai_codex/api.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import asyncio
-import re
 from dataclasses import dataclass
 from typing import AsyncIterator, Iterator
 
@@ -45,7 +44,7 @@ from ._run import (
 )
 from ._sandbox import Sandbox as Sandbox, _sandbox_mode, _sandbox_policy
 from .async_client import AsyncCodexClient
-from .client import CodexClient, CodexConfig
+from .client import CodexClient, CodexConfig, _active_turn_id_from_error
 from .errors import InvalidRequestError
 from .generated.v2_all import (
     ApiKeyLoginAccountParams,
@@ -92,11 +91,6 @@ def _normalize_goal_objective(objective: str) -> str:
     return objective
 
 
-def _active_turn_id_from_error(exc: InvalidRequestError) -> str | None:
-    match = re.search(r" but found `?([^`]+)`?$", exc.message)
-    return match.group(1) if match is not None else None
-
-
 def _inactive_turn_error() -> InvalidRequestError:
     return InvalidRequestError(-32600, "no active turn to steer")
 
@@ -845,7 +839,7 @@ class TurnHandle:
                 self._goal,
                 lambda: self._client.next_goal_notification(self._goal),
                 lambda: self._client.unregister_goal_operation(self._goal),
-                lambda: self._client.stop_failed_goal(self._goal),
+                lambda: self._client.cancel_goal_operation(self._goal),
             )
 
         def ordinary_stream() -> Iterator[Notification]:
@@ -870,6 +864,10 @@ class TurnHandle:
         stream = self.stream()
         try:
             return _collect_turn_result(stream, turn_id=self.id)
+        except KeyboardInterrupt:
+            if self._goal is not None and not self._goal.is_finished():
+                self._client.cancel_goal_operation(self._goal)
+            raise
         finally:
             stream.close()
 
@@ -968,7 +966,7 @@ class AsyncTurnHandle:
                 self._goal,
                 next_goal_notification,
                 lambda: self._codex._client.unregister_goal_operation(self._goal),
-                lambda: self._codex._client.stop_failed_goal(self._goal),
+                lambda: self._codex._client.cancel_goal_operation(self._goal),
             )
 
         async def ordinary_stream() -> AsyncIterator[Notification]:
@@ -994,5 +992,9 @@ class AsyncTurnHandle:
         stream = self.stream()
         try:
             return await _collect_async_turn_result(stream, turn_id=self.id)
+        except asyncio.CancelledError:
+            if self._goal is not None and not self._goal.is_finished():
+                await self._codex._client.cancel_goal_operation(self._goal)
+            raise
         finally:
             await stream.aclose()

sdk/python/src/openai_codex/async_client.py

@@ -228,9 +228,9 @@ class AsyncCodexClient:
         """Pause the active goal through the wrapped sync client."""
         return await self._call_sync(self._sync.pause_goal, thread_id)
 
-    async def stop_failed_goal(self, state: _GoalOperationState) -> None:
-        """Stop continuation work after a physical goal turn fails."""
-        await self._call_sync(self._sync.stop_failed_goal, state)
+    async def cancel_goal_operation(self, state: _GoalOperationState) -> None:
+        """Stop continuation work after a logical goal operation is cancelled."""
+        await self._call_sync(self._sync.cancel_goal_operation, state)
 
     async def start_goal_operation(
         self,
@@ -259,7 +259,7 @@ class AsyncCodexClient:
 
                 def stop_cancelled_goal() -> None:
                     try:
-                        self._sync.stop_failed_goal(state)
+                        self._sync.cancel_goal_operation(state)
                     finally:
                         state.finish()
                         self._sync.unregister_goal_operation(state)

sdk/python/src/openai_codex/client.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 import subprocess
 import threading
 import uuid
@@ -65,6 +66,11 @@ RUNTIME_PKG_NAME = "openai-codex-cli-bin"
 _GOAL_START_TIMEOUT_S = 30.0
 
 
+def _active_turn_id_from_error(exc: InvalidRequestError) -> str | None:
+    match = re.search(r" but found `?([^`]+)`?$", exc.message)
+    return match.group(1) if match is not None else None
+
+
 def _params_dict(
     params: (
         V2ThreadStartParams
@@ -501,8 +507,8 @@ class CodexClient:
         """Pause the active goal used by a logical goal turn."""
         return self.thread_goal_set(thread_id, status=ThreadGoalStatus.paused)
 
-    def stop_failed_goal(self, state: _GoalOperationState) -> None:
-        """Best-effort cleanup after a physical goal turn fails."""
+    def cancel_goal_operation(self, state: _GoalOperationState) -> None:
+        """Best-effort cleanup after a logical goal operation is cancelled."""
         try:
             self.pause_goal(state.thread_id)
         except Exception:
@@ -515,7 +521,7 @@ class CodexClient:
         except InvalidRequestError as exc:
             if not exc.message.startswith("expected active turn id"):
                 return
-            next_turn_id = state.current_turn()
+            next_turn_id = _active_turn_id_from_error(exc) or state.current_turn()
             if next_turn_id is None or next_turn_id == turn_id:
                 return
             try:
@@ -560,12 +566,9 @@ class CodexClient:
                     f"{int(_GOAL_START_TIMEOUT_S)} seconds"
                 )
             return state, turn_id
-        except BaseException:
-            if activated:
-                try:
-                    self.pause_goal(thread_id)
-                except Exception:
-                    pass
+        except BaseException as exc:
+            if activated or not isinstance(exc, InvalidRequestError):
+                self.cancel_goal_operation(state)
             state.finish()
             self.unregister_goal_operation(state)
             raise

sdk/python/tests/app_server_harness.py

@@ -448,12 +448,12 @@ def ev_function_call(call_id: str, name: str, arguments: str) -> Json:
     }
 
 
-def ev_failed(response_id: str, message: str) -> Json:
+def ev_failed(response_id: str, message: str, *, code: str = "server_error") -> Json:
     """Return a failed model response event."""
     return {
         "type": "response.failed",
         "response": {
             "id": response_id,
-            "error": {"code": "server_error", "message": message},
+            "error": {"code": code, "message": message},
         },
     }

sdk/python/tests/test_app_server_goal_turns.py

@@ -1,4 +1,6 @@
+import _thread
 import asyncio
+import threading
 import time
 
 import pytest
@@ -18,6 +20,7 @@ from app_server_helpers import (
 )
 
 from openai_codex import AsyncCodex, Codex
+from openai_codex._goal import _GoalOperationState
 from openai_codex.errors import InvalidRequestError, TransportClosedError
 from openai_codex.generated.notification_registry import notification_turn_id
 from openai_codex.generated.v2_all import (
@@ -467,14 +470,72 @@ def test_goal_interrupt_pauses_continuation_and_leaves_thread_usable(tmp_path) -
     }
 
 
-def test_terminal_goal_failure_stops_continuation_and_releases_routing(tmp_path) -> None:
+def test_cancel_goal_retries_the_server_reported_active_turn(tmp_path) -> None:
+    """Cancellation should survive notification lag while a continuation rolls over."""
+    with AppServerHarness(tmp_path, enable_goals=True) as harness:
+        _enqueue_interruptible_goal(
+            harness,
+            "cancel-rollover",
+            initial_text="Initial work complete.",
+            continuation_chunks=["continuation ", "still ", "running ", "now"],
+            follow_up_text="Rollover follow-up complete.",
+        )
+
+        with Codex(config=harness.app_server_config()) as codex:
+            thread = codex.thread_start()
+            turn = thread.start_goal("Cancel while continuation routing catches up")
+            harness.responses.wait_for_requests(2)
+
+            stale_state = _GoalOperationState(
+                thread_id=thread.id,
+                current_turn_id=turn.id,
+            )
+            codex._client.cancel_goal_operation(stale_state)
+            interrupted = turn.run()
+            follow_up = thread.run("Continue after rollover cancellation")
+            requests = harness.responses.wait_for_requests(3)
+            persisted = codex._client.request(
+                "thread/goal/get",
+                {"threadId": thread.id},
+                response_model=ThreadGoalGetResponse,
+            ).goal
+
+    assert {
+        "goal_status": persisted.status if persisted else None,
+        "interrupted": (interrupted.id, interrupted.status),
+        "follow_up": (follow_up.status, follow_up.final_response),
+        "request_count": len(requests),
+    } == {
+        "goal_status": ThreadGoalStatus.paused,
+        "interrupted": (turn.id, TurnStatus.interrupted),
+        "follow_up": (TurnStatus.completed, "Rollover follow-up complete."),
+        "request_count": 3,
+    }
+
+
+@pytest.mark.parametrize(
+    ("error_code", "expected_status"),
+    [
+        ("server_error", ThreadGoalStatus.blocked),
+        ("insufficient_quota", ThreadGoalStatus.usage_limited),
+    ],
+)
+def test_terminal_goal_failure_preserves_status_and_releases_routing(
+    tmp_path,
+    error_code,
+    expected_status,
+) -> None:
     """A failed server turn should stop rollover work and leave the thread usable."""
     with AppServerHarness(tmp_path, enable_goals=True) as harness:
         harness.responses.enqueue_sse(
             sse(
                 [
                     ev_response_created("goal-terminal-failure"),
-                    ev_failed("goal-terminal-failure", "goal model failed"),
+                    ev_failed(
+                        "goal-terminal-failure",
+                        "goal model failed",
+                        code=error_code,
+                    ),
                 ]
             )
         )
@@ -483,6 +544,11 @@ def test_terminal_goal_failure_stops_continuation_and_releases_routing(tmp_path)
             with pytest.raises(RuntimeError, match="goal model failed"):
                 thread.run_goal("Fail this goal turn")
 
+            persisted = codex._client.request(
+                "thread/goal/get",
+                {"threadId": thread.id},
+                response_model=ThreadGoalGetResponse,
+            ).goal
             harness.responses.enqueue_assistant_message(
                 "Recovered with an ordinary turn.",
                 response_id="goal-failure-follow-up",
@@ -492,16 +558,88 @@ def test_terminal_goal_failure_stops_continuation_and_releases_routing(tmp_path)
             requests = harness.responses.requests()
 
     assert {
+        "goal_status": persisted.status if persisted else None,
         "follow_up": (follow_up.status, follow_up.final_response),
         "request_count": len(requests),
         "follow_up_input": requests[1].message_input_texts("user")[-1:],
     } == {
+        "goal_status": expected_status,
         "follow_up": (TurnStatus.completed, "Recovered with an ordinary turn."),
         "request_count": 2,
         "follow_up_input": ["Run after the goal failure"],
     }
 
 
+def test_sync_goal_run_cancellation_stops_work_and_releases_routing(tmp_path) -> None:
+    """Keyboard interruption should pause and interrupt a running logical goal."""
+    with AppServerHarness(tmp_path, enable_goals=True) as harness:
+        harness.responses.enqueue_sse(
+            streaming_response(
+                "cancelled-sync-run",
+                "msg-cancelled-sync-run",
+                ["cancelled ", "sync ", "goal"],
+            ),
+            delay_between_events_s=0.5,
+        )
+        harness.responses.enqueue_assistant_message(
+            "Sync follow-up complete.",
+            response_id="cancelled-sync-run-follow-up",
+        )
+
+        with Codex(config=harness.app_server_config()) as codex:
+            thread = codex.thread_start()
+            turn = thread.start_goal("Cancel this running sync goal")
+            goal_state = turn._goal
+            assert goal_state is not None
+
+            def interrupt_when_collecting() -> None:
+                harness.responses.wait_for_requests(1)
+                deadline = time.monotonic() + 5
+                while not goal_state._notifications.empty():
+                    if time.monotonic() >= deadline:
+                        raise AssertionError("goal stream did not begin collecting")
+                    time.sleep(0.01)
+                _thread.interrupt_main()
+
+            interrupter = threading.Thread(target=interrupt_when_collecting, daemon=True)
+            interrupter.start()
+            try:
+                with pytest.raises(KeyboardInterrupt):
+                    turn.run()
+            finally:
+                interrupter.join(timeout=5)
+
+            deadline = time.monotonic() + 5
+            while True:
+                current = thread.read()
+                if isinstance(current.thread.status.root, IdleThreadStatus):
+                    break
+                if time.monotonic() >= deadline:
+                    raise AssertionError("cancelled sync goal turn did not stop")
+                time.sleep(0.01)
+
+            persisted = codex._client.request(
+                "thread/goal/get",
+                {"threadId": thread.id},
+                response_model=ThreadGoalGetResponse,
+            ).goal
+            registered_goals = dict(codex._client._router._goal_operations)
+            follow_up = thread.run("Continue after cancelling the sync goal")
+            requests = harness.responses.wait_for_requests(2)
+
+    assert {
+        "goal_status": persisted.status if persisted else None,
+        "follow_up": (follow_up.status, follow_up.final_response),
+        "request_count": len(requests),
+        "registered_goals": registered_goals,
+    } == {
+        "goal_status": ThreadGoalStatus.paused,
+        "follow_up": (TurnStatus.completed, "Sync follow-up complete."),
+        "request_count": 2,
+        "registered_goals": {},
+    }
+
+
 def test_closing_goal_stream_releases_real_process_routing(tmp_path) -> None:
     """Closing a stream should release SDK routing without pausing the persisted goal."""
     with AppServerHarness(tmp_path, enable_goals=True) as harness:
@@ -779,6 +917,67 @@ def test_async_goal_interrupts_an_active_continuation(tmp_path) -> None:
     asyncio.run(scenario())
 
 
+def test_async_goal_run_cancellation_stops_work_and_releases_routing(tmp_path) -> None:
+    """Task cancellation should pause and interrupt a running logical goal."""
+
+    async def scenario() -> None:
+        with AppServerHarness(tmp_path, enable_goals=True) as harness:
+            harness.responses.enqueue_sse(
+                streaming_response(
+                    "cancelled-async-run",
+                    "msg-cancelled-async-run",
+                    ["cancelled ", "async ", "goal"],
+                ),
+                delay_between_events_s=0.5,
+            )
+            harness.responses.enqueue_assistant_message(
+                "Async run follow-up complete.",
+                response_id="cancelled-async-run-follow-up",
+            )
+
+            async with AsyncCodex(config=harness.app_server_config()) as codex:
+                thread = await codex.thread_start()
+                turn = await thread.start_goal("Cancel this running async goal")
+                running = asyncio.create_task(turn.run())
+                await asyncio.sleep(0)
+                await asyncio.to_thread(harness.responses.wait_for_requests, 1)
+                running.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await running
+
+                deadline = time.monotonic() + 5
+                while True:
+                    current = await thread.read()
+                    if isinstance(current.thread.status.root, IdleThreadStatus):
+                        break
+                    if time.monotonic() >= deadline:
+                        raise AssertionError("cancelled async goal turn did not stop")
+                    await asyncio.sleep(0.01)
+
+                persisted = await codex._client.request(
+                    "thread/goal/get",
+                    {"threadId": thread.id},
+                    response_model=ThreadGoalGetResponse,
+                )
+                registered_goals = dict(codex._client._sync._router._goal_operations)
+                follow_up = await thread.run("Continue after cancelling the async goal")
+                requests = await asyncio.to_thread(harness.responses.wait_for_requests, 2)
+
+        assert {
+            "goal_status": persisted.goal.status if persisted.goal else None,
+            "follow_up": (follow_up.status, follow_up.final_response),
+            "request_count": len(requests),
+            "registered_goals": registered_goals,
+        } == {
+            "goal_status": ThreadGoalStatus.paused,
+            "follow_up": (TurnStatus.completed, "Async run follow-up complete."),
+            "request_count": 2,
+            "registered_goals": {},
+        }
+
+    asyncio.run(scenario())
+
+
 def test_async_goal_start_cancellation_interrupts_work_and_releases_routing(tmp_path) -> None:
     """Cancelling startup should pause and interrupt the physical goal turn."""