From 4557f39923ced8d8659cf9b5d08fbeb2d09a85ed Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Mon, 8 Jun 2026 12:13:06 -0700 Subject: [PATCH] Fix goal failure cleanup in Python SDK --- sdk/python/src/openai_codex/_goal.py | 34 +++++++++++-------- sdk/python/src/openai_codex/api.py | 2 ++ sdk/python/src/openai_codex/async_client.py | 4 +++ sdk/python/src/openai_codex/client.py | 24 +++++++++++++ sdk/python/tests/app_server_harness.py | 2 +- .../tests/test_app_server_goal_turns.py | 22 ++++++++---- 6 files changed, 66 insertions(+), 22 deletions(-) diff --git a/sdk/python/src/openai_codex/_goal.py b/sdk/python/src/openai_codex/_goal.py index 4e64211958..eb628bb490 100644 --- a/sdk/python/src/openai_codex/_goal.py +++ b/sdk/python/src/openai_codex/_goal.py @@ -80,9 +80,6 @@ class _GoalOperationState: elif isinstance(payload, ThreadTokenUsageUpdatedNotification): self.usage = payload.token_usage if ( - isinstance(payload, TurnCompletedNotification) - and payload.turn.status in {TurnStatus.failed, TurnStatus.interrupted} - ) or ( self.current_turn_id is None and self.completed_turn is not None and (self.cleared or _terminal_goal_status(self.status)) @@ -236,7 +233,7 @@ class _GoalStreamCursor: active: bool = False cleared: bool = False - def process(self, notification: Notification) -> tuple[list[Notification], bool]: + def process(self, notification: Notification) -> tuple[list[Notification], bool, 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") @@ -245,24 +242,27 @@ class _GoalStreamCursor: if isinstance(payload, TurnStartedNotification): self.active = True if self.started is not None: - return [], False + return [], False, False self.started = payload.turn - return [_logical_notification(notification, logical_turn_id)], False + return [_logical_notification(notification, logical_turn_id)], False, False if isinstance(payload, TurnCompletedNotification): self.active = False self.last_completed = payload if payload.turn.status in {TurnStatus.failed, TurnStatus.interrupted}: - self.state.finish() - return [self._completion(notification.method, payload)], True + return ( + [self._completion(notification.method, payload)], + True, + (payload.turn.status == TurnStatus.failed), + ) 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 - return [], False + return [self._completion(notification.method, payload)], True, False + return [], False, False events = [_logical_notification(notification, logical_turn_id)] if isinstance(payload, ThreadGoalUpdatedNotification): @@ -281,8 +281,8 @@ class _GoalStreamCursor: ): self.state.finish() events.append(self._completion("turn/completed", self.last_completed)) - return events, True - return events, False + return events, True, False + return events, False, False def _completion( self, @@ -310,6 +310,7 @@ class _GoalNotificationStream(Iterator[Notification]): state: _GoalOperationState next_notification: Callable[[], Notification] unregister: Callable[[], None] + stop_failed_goal: Callable[[], None] _cursor: _GoalStreamCursor = field(init=False) _pending: deque[Notification] = field(default_factory=deque) _closed: bool = False @@ -325,7 +326,9 @@ class _GoalNotificationStream(Iterator[Notification]): raise StopIteration try: while not self._pending: - events, completed = self._cursor.process(self.next_notification()) + events, completed, failed = self._cursor.process(self.next_notification()) + if failed: + self.stop_failed_goal() self._pending.extend(events) if completed: self._finish() @@ -356,6 +359,7 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]): state: _GoalOperationState next_notification: Callable[[], Awaitable[Notification]] unregister: Callable[[], None] + stop_failed_goal: Callable[[], Awaitable[None]] _cursor: _GoalStreamCursor = field(init=False) _pending: deque[Notification] = field(default_factory=deque) _closed: bool = False @@ -371,7 +375,9 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]): raise StopAsyncIteration try: while not self._pending: - events, completed = self._cursor.process(await self.next_notification()) + events, completed, failed = self._cursor.process(await self.next_notification()) + if failed: + await self.stop_failed_goal() self._pending.extend(events) if completed: self._finish() diff --git a/sdk/python/src/openai_codex/api.py b/sdk/python/src/openai_codex/api.py index a856c83f9c..9f37bfd2f2 100644 --- a/sdk/python/src/openai_codex/api.py +++ b/sdk/python/src/openai_codex/api.py @@ -841,6 +841,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), ) def ordinary_stream() -> Iterator[Notification]: @@ -963,6 +964,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), ) async def ordinary_stream() -> AsyncIterator[Notification]: diff --git a/sdk/python/src/openai_codex/async_client.py b/sdk/python/src/openai_codex/async_client.py index b1f2f6759c..10f4def5f5 100644 --- a/sdk/python/src/openai_codex/async_client.py +++ b/sdk/python/src/openai_codex/async_client.py @@ -227,6 +227,10 @@ 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 start_goal_operation( self, thread_id: str, diff --git a/sdk/python/src/openai_codex/client.py b/sdk/python/src/openai_codex/client.py index 4ef54108e9..d2a12bb8b1 100644 --- a/sdk/python/src/openai_codex/client.py +++ b/sdk/python/src/openai_codex/client.py @@ -501,6 +501,30 @@ 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.""" + try: + self.pause_goal(state.thread_id) + except Exception: + pass + turn_id = state.current_turn() + if turn_id is None: + return + try: + self.turn_interrupt(state.thread_id, turn_id) + except InvalidRequestError as exc: + if not exc.message.startswith("expected active turn id"): + return + next_turn_id = state.current_turn() + if next_turn_id is None or next_turn_id == turn_id: + return + try: + self.turn_interrupt(state.thread_id, next_turn_id) + except Exception: + pass + except Exception: + pass + def start_goal_operation( self, thread_id: str, diff --git a/sdk/python/tests/app_server_harness.py b/sdk/python/tests/app_server_harness.py index d609cfb6e0..2c83741f91 100644 --- a/sdk/python/tests/app_server_harness.py +++ b/sdk/python/tests/app_server_harness.py @@ -247,7 +247,7 @@ class AppServerHarness: """Write config.toml that routes model calls to the mock server.""" config_toml = self.codex_home / "config.toml" requires_openai_auth = "requires_openai_auth = true\n" if self.requires_openai_auth else "" - goals = "[features]\ngoals = true\n\n" if self.enable_goals else "" + goals = f"[features]\ngoals = {str(self.enable_goals).lower()}\n\n" config_toml.write_text( f""" model = "mock-model" diff --git a/sdk/python/tests/test_app_server_goal_turns.py b/sdk/python/tests/test_app_server_goal_turns.py index 660351a3ae..5e4f372efb 100644 --- a/sdk/python/tests/test_app_server_goal_turns.py +++ b/sdk/python/tests/test_app_server_goal_turns.py @@ -466,7 +466,7 @@ 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: - """A failed server turn should end the logical operation without another continuation.""" + """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( @@ -476,9 +476,13 @@ def test_terminal_goal_failure_stops_continuation_and_releases_routing(tmp_path) ] ) ) - harness.responses.enqueue_assistant_message( - "Recovered with an ordinary turn.", - response_id="goal-failure-follow-up", + harness.responses.enqueue_sse( + streaming_response( + "goal-failure-rollover", + "msg-goal-failure-rollover", + ["automatic ", "rollover"], + ), + delay_between_events_s=0.5, ) with Codex(config=harness.app_server_config()) as codex: @@ -486,17 +490,21 @@ 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") + harness.responses.enqueue_assistant_message( + "Recovered with an ordinary turn.", + response_id="goal-failure-follow-up", + ) follow_up = thread.run("Run after the goal failure") - harness.responses.wait_for_requests(2) + harness.responses.wait_for_requests(3) requests = harness.responses.requests() assert { "follow_up": (follow_up.status, follow_up.final_response), "request_count": len(requests), - "follow_up_input": requests[1].message_input_texts("user")[-1:], + "follow_up_input": requests[2].message_input_texts("user")[-1:], } == { "follow_up": (TurnStatus.completed, "Recovered with an ordinary turn."), - "request_count": 2, + "request_count": 3, "follow_up_input": ["Run after the goal failure"], }