diff --git a/sdk/python/src/openai_codex/_goal.py b/sdk/python/src/openai_codex/_goal.py index eb628bb490..72148003dd 100644 --- a/sdk/python/src/openai_codex/_goal.py +++ b/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 diff --git a/sdk/python/src/openai_codex/api.py b/sdk/python/src/openai_codex/api.py index 2971f8d1ac..57ba84330d 100644 --- a/sdk/python/src/openai_codex/api.py +++ b/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() diff --git a/sdk/python/src/openai_codex/async_client.py b/sdk/python/src/openai_codex/async_client.py index 2aecca864c..38d621356a 100644 --- a/sdk/python/src/openai_codex/async_client.py +++ b/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) diff --git a/sdk/python/src/openai_codex/client.py b/sdk/python/src/openai_codex/client.py index d2a12bb8b1..9d3622c166 100644 --- a/sdk/python/src/openai_codex/client.py +++ b/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 diff --git a/sdk/python/tests/app_server_harness.py b/sdk/python/tests/app_server_harness.py index 2c83741f91..daba753d82 100644 --- a/sdk/python/tests/app_server_harness.py +++ b/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}, }, } diff --git a/sdk/python/tests/test_app_server_goal_turns.py b/sdk/python/tests/test_app_server_goal_turns.py index 7ce8f6d1fd..3d45a9f735 100644 --- a/sdk/python/tests/test_app_server_goal_turns.py +++ b/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."""