diff --git a/sdk/python/src/openai_codex/_goal.py b/sdk/python/src/openai_codex/_goal.py index 4f76e0cc72..0bd8e8870a 100644 --- a/sdk/python/src/openai_codex/_goal.py +++ b/sdk/python/src/openai_codex/_goal.py @@ -8,6 +8,7 @@ from typing import AsyncIterator, Awaitable, Callable, Iterator from .generated.notification_registry import notification_turn_id from .generated.v2_all import ( + CodexErrorInfoValue, ItemCompletedNotification, ThreadGoalClearedNotification, ThreadGoalStatus, @@ -37,6 +38,13 @@ def _terminal_goal_status(status: ThreadGoalStatus | None) -> bool: } +def _failed_goal_status(turn: Turn) -> ThreadGoalStatus: + error_info = turn.error.codex_error_info if turn.error is not None else None + if error_info is not None and error_info.root == CodexErrorInfoValue.usage_limit_exceeded: + return ThreadGoalStatus.usage_limited + return ThreadGoalStatus.blocked + + @dataclass(slots=True) class _GoalOperationState: """Private state for one goal operation exposed as a logical turn.""" @@ -313,6 +321,7 @@ class _GoalNotificationStream(Iterator[Notification]): next_notification: Callable[[], Notification] unregister: Callable[[], None] cancel_goal: Callable[[], None] + finish_failed_goal: Callable[[ThreadGoalStatus], None] _cursor: _GoalStreamCursor = field(init=False) _pending: deque[Notification] = field(default_factory=deque) _closed: bool = False @@ -328,7 +337,15 @@ class _GoalNotificationStream(Iterator[Notification]): raise StopIteration try: while not self._pending: - events, completed = self._cursor.process(self.next_notification()) + notification = self.next_notification() + events, completed = self._cursor.process(notification) + payload = notification.payload + if ( + not completed + and isinstance(payload, TurnCompletedNotification) + and payload.turn.status == TurnStatus.failed + ): + self.finish_failed_goal(_failed_goal_status(payload.turn)) self._pending.extend(events) if completed: self._finish() @@ -364,6 +381,7 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]): next_notification: Callable[[], Awaitable[Notification]] unregister: Callable[[], None] cancel_goal: Callable[[], Awaitable[None]] + finish_failed_goal: Callable[[ThreadGoalStatus], Awaitable[None]] _cursor: _GoalStreamCursor = field(init=False) _pending: deque[Notification] = field(default_factory=deque) _closed: bool = False @@ -379,7 +397,15 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]): raise StopAsyncIteration try: while not self._pending: - events, completed = self._cursor.process(await self.next_notification()) + notification = await self.next_notification() + events, completed = self._cursor.process(notification) + payload = notification.payload + if ( + not completed + and isinstance(payload, TurnCompletedNotification) + and payload.turn.status == TurnStatus.failed + ): + await self.finish_failed_goal(_failed_goal_status(payload.turn)) 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 57ba84330d..d9353f5f3e 100644 --- a/sdk/python/src/openai_codex/api.py +++ b/sdk/python/src/openai_codex/api.py @@ -840,6 +840,7 @@ class TurnHandle: lambda: self._client.next_goal_notification(self._goal), lambda: self._client.unregister_goal_operation(self._goal), lambda: self._client.cancel_goal_operation(self._goal), + lambda status: self._client.finish_failed_goal(self._goal, status), ) def ordinary_stream() -> Iterator[Notification]: @@ -967,6 +968,7 @@ class AsyncTurnHandle: next_goal_notification, lambda: self._codex._client.unregister_goal_operation(self._goal), lambda: self._codex._client.cancel_goal_operation(self._goal), + lambda status: self._codex._client.finish_failed_goal(self._goal, status), ) 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 38d621356a..d00a7e9fda 100644 --- a/sdk/python/src/openai_codex/async_client.py +++ b/sdk/python/src/openai_codex/async_client.py @@ -228,6 +228,14 @@ class AsyncCodexClient: """Pause the active goal through the wrapped sync client.""" return await self._call_sync(self._sync.pause_goal, thread_id) + async def finish_failed_goal( + self, + state: _GoalOperationState, + status: ThreadGoalStatus, + ) -> None: + """Persist the terminal status for a failed physical goal turn.""" + await self._call_sync(self._sync.finish_failed_goal, state, status) + 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) diff --git a/sdk/python/src/openai_codex/client.py b/sdk/python/src/openai_codex/client.py index 9d3622c166..1d059ca8a2 100644 --- a/sdk/python/src/openai_codex/client.py +++ b/sdk/python/src/openai_codex/client.py @@ -507,6 +507,14 @@ class CodexClient: """Pause the active goal used by a logical goal turn.""" return self.thread_goal_set(thread_id, status=ThreadGoalStatus.paused) + def finish_failed_goal( + self, + state: _GoalOperationState, + status: ThreadGoalStatus, + ) -> None: + """Persist the terminal status for a failed physical goal turn.""" + self.thread_goal_set(state.thread_id, status=status) + def cancel_goal_operation(self, state: _GoalOperationState) -> None: """Best-effort cleanup after a logical goal operation is cancelled.""" try: