From 3fbee02b960b42a0afff6ac6035834ee2df83bc6 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 9 Jun 2026 13:37:58 -0700 Subject: [PATCH 1/2] [3/6] Add private Python goal-aware turn handles --- sdk/python/src/openai_codex/_goal.py | 7 +- sdk/python/src/openai_codex/_goal_handle.py | 182 ++++++++++++++++++++ sdk/python/src/openai_codex/api.py | 37 +++- 3 files changed, 224 insertions(+), 2 deletions(-) create mode 100644 sdk/python/src/openai_codex/_goal_handle.py diff --git a/sdk/python/src/openai_codex/_goal.py b/sdk/python/src/openai_codex/_goal.py index 37a63fc28f4c..6bb3c32ebcca 100644 --- a/sdk/python/src/openai_codex/_goal.py +++ b/sdk/python/src/openai_codex/_goal.py @@ -174,7 +174,12 @@ def resolve_active_turn(self, expected: str, active: str) -> None: """Adopt a server-reported active id when routed state is still stale.""" with self._condition: if self.current_turn_id in {None, expected}: - self.current_turn_id = active + if self.completed_turn is not None and self.completed_turn.id == active: + self.current_turn_id = None + if self.cleared or _terminal_goal_status(self.status): + self._finished = True + else: + self.current_turn_id = active self._condition.notify_all() def turn_for_interrupt(self) -> str | None: diff --git a/sdk/python/src/openai_codex/_goal_handle.py b/sdk/python/src/openai_codex/_goal_handle.py new file mode 100644 index 000000000000..d2a1b3e2c6ce --- /dev/null +++ b/sdk/python/src/openai_codex/_goal_handle.py @@ -0,0 +1,182 @@ +import asyncio +from dataclasses import dataclass + +from ._goal import ( + _AsyncGoalNotificationStream, + _GoalNotificationStream, + _GoalOperationState, +) +from ._inputs import RunInput, _normalize_run_input, _to_wire_input +from .async_client import AsyncCodexClient +from .client import CodexClient, _active_turn_id_from_error +from .errors import InvalidRequestError +from .generated.v2_all import TurnInterruptResponse, TurnSteerResponse + + +def _inactive_turn_error() -> InvalidRequestError: + return InvalidRequestError(-32600, "no active turn to steer") + + +def _inactive_interrupt_error() -> InvalidRequestError: + return InvalidRequestError(-32600, "no active turn to interrupt") + + +@dataclass(slots=True) +class _GoalTurnHandleAdapter: + """Implement synchronous turn controls for one logical goal operation.""" + + client: CodexClient + state: _GoalOperationState + thread_id: str + logical_turn_id: str + + def steer(self, input: RunInput) -> TurnSteerResponse: + wire_input = _to_wire_input(_normalize_run_input(input)) + turn_id = self.state.active_turn() + if turn_id is None: + raise _inactive_turn_error() + try: + response = self.client.turn_steer(self.thread_id, turn_id, wire_input) + except InvalidRequestError as exc: + if not ( + exc.message == "no active turn to steer" + or exc.message.startswith("expected active turn id") + ): + raise + next_turn_id = _active_turn_id_from_error(exc) + if next_turn_id is None: + next_turn_id = self.state.active_turn(after=turn_id) + if next_turn_id is None: + raise _inactive_turn_error() from exc + response = self.client.turn_steer(self.thread_id, next_turn_id, wire_input) + self.state.resolve_active_turn(turn_id, next_turn_id) + return response.model_copy(update={"turn_id": self.logical_turn_id}) + + def interrupt(self) -> TurnInterruptResponse: + if not self.state.begin_interrupt(): + raise _inactive_interrupt_error() + try: + self.client.pause_goal(self.thread_id) + turn_id = self.state.turn_for_interrupt() + if turn_id is None: + response = TurnInterruptResponse() + else: + try: + response = self.client.turn_interrupt(self.thread_id, turn_id) + except InvalidRequestError as exc: + if exc.message == "no active turn to interrupt": + response = TurnInterruptResponse() + elif exc.message.startswith("expected active turn id"): + next_turn_id = _active_turn_id_from_error(exc) or self.state.current_turn() + if next_turn_id is None or next_turn_id == turn_id: + response = TurnInterruptResponse() + else: + try: + response = self.client.turn_interrupt( + self.thread_id, + next_turn_id, + ) + except InvalidRequestError as retry_exc: + if retry_exc.message != "no active turn to interrupt": + raise + response = TurnInterruptResponse() + self.state.resolve_active_turn(turn_id, next_turn_id) + else: + raise + self.state.confirm_interrupt() + return response + except BaseException: + self.state.cancel_interrupt() + raise + + def stream(self) -> _GoalNotificationStream: + return _GoalNotificationStream( + self.state, + lambda: self.client.next_goal_notification(self.state), + lambda: self.client.unregister_goal_operation(self.state), + lambda: self.client.cancel_goal_operation(self.state), + ) + + +@dataclass(slots=True) +class _AsyncGoalTurnHandleAdapter: + """Implement asynchronous turn controls for one logical goal operation.""" + + client: AsyncCodexClient + state: _GoalOperationState + thread_id: str + logical_turn_id: str + + async def steer(self, input: RunInput) -> TurnSteerResponse: + wire_input = _to_wire_input(_normalize_run_input(input)) + turn_id = await asyncio.to_thread(self.state.active_turn) + if turn_id is None: + raise _inactive_turn_error() + try: + response = await self.client.turn_steer(self.thread_id, turn_id, wire_input) + except InvalidRequestError as exc: + if not ( + exc.message == "no active turn to steer" + or exc.message.startswith("expected active turn id") + ): + raise + next_turn_id = _active_turn_id_from_error(exc) + if next_turn_id is None: + next_turn_id = await asyncio.to_thread( + self.state.active_turn, + after=turn_id, + ) + if next_turn_id is None: + raise _inactive_turn_error() from exc + response = await self.client.turn_steer( + self.thread_id, + next_turn_id, + wire_input, + ) + self.state.resolve_active_turn(turn_id, next_turn_id) + return response.model_copy(update={"turn_id": self.logical_turn_id}) + + async def interrupt(self) -> TurnInterruptResponse: + if not self.state.begin_interrupt(): + raise _inactive_interrupt_error() + try: + await self.client.pause_goal(self.thread_id) + turn_id = self.state.turn_for_interrupt() + if turn_id is None: + response = TurnInterruptResponse() + else: + try: + response = await self.client.turn_interrupt(self.thread_id, turn_id) + except InvalidRequestError as exc: + if exc.message == "no active turn to interrupt": + response = TurnInterruptResponse() + elif exc.message.startswith("expected active turn id"): + next_turn_id = _active_turn_id_from_error(exc) or self.state.current_turn() + if next_turn_id is None or next_turn_id == turn_id: + response = TurnInterruptResponse() + else: + try: + response = await self.client.turn_interrupt( + self.thread_id, + next_turn_id, + ) + except InvalidRequestError as retry_exc: + if retry_exc.message != "no active turn to interrupt": + raise + response = TurnInterruptResponse() + self.state.resolve_active_turn(turn_id, next_turn_id) + else: + raise + self.state.confirm_interrupt() + return response + except BaseException: + self.state.cancel_interrupt() + raise + + def stream(self) -> _AsyncGoalNotificationStream: + return _AsyncGoalNotificationStream( + self.state, + lambda: self.client.next_goal_notification(self.state), + lambda: self.client.unregister_goal_operation(self.state), + lambda: self.client.cancel_goal_operation(self.state), + ) diff --git a/sdk/python/src/openai_codex/api.py b/sdk/python/src/openai_codex/api.py index 6fc9a8243d63..18baa93e022b 100644 --- a/sdk/python/src/openai_codex/api.py +++ b/sdk/python/src/openai_codex/api.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import AsyncIterator, Iterator from ._approval_mode import ( @@ -9,6 +9,10 @@ _approval_mode_override_settings, _approval_mode_settings, ) +from ._goal_handle import ( + _AsyncGoalTurnHandleAdapter, + _GoalTurnHandleAdapter, +) from ._initialize_metadata import validate_initialize_metadata from ._inputs import ( ImageInput as ImageInput, @@ -721,9 +725,17 @@ class TurnHandle: _client: CodexClient thread_id: str id: str + _goal: _GoalTurnHandleAdapter | None = field( + default=None, + init=False, + repr=False, + compare=False, + ) def steer(self, input: RunInput) -> TurnSteerResponse: """Send additional input to this active turn.""" + if self._goal is not None: + return self._goal.steer(input) return self._client.turn_steer( self.thread_id, self.id, @@ -732,10 +744,15 @@ def steer(self, input: RunInput) -> TurnSteerResponse: def interrupt(self) -> TurnInterruptResponse: """Request interruption of this active turn.""" + if self._goal is not None: + return self._goal.interrupt() return self._client.turn_interrupt(self.thread_id, self.id) def stream(self) -> Iterator[Notification]: """Yield only notifications routed to this turn handle.""" + if self._goal is not None: + yield from self._goal.stream() + return self._client.register_turn_notifications(self.id) try: while True: @@ -766,10 +783,18 @@ class AsyncTurnHandle: _codex: AsyncCodex thread_id: str id: str + _goal: _AsyncGoalTurnHandleAdapter | None = field( + default=None, + init=False, + repr=False, + compare=False, + ) async def steer(self, input: RunInput) -> TurnSteerResponse: """Send additional input to this active turn.""" await self._codex._ensure_initialized() + if self._goal is not None: + return await self._goal.steer(input) return await self._codex._client.turn_steer( self.thread_id, self.id, @@ -779,11 +804,21 @@ async def steer(self, input: RunInput) -> TurnSteerResponse: async def interrupt(self) -> TurnInterruptResponse: """Request interruption of this active turn.""" await self._codex._ensure_initialized() + if self._goal is not None: + return await self._goal.interrupt() return await self._codex._client.turn_interrupt(self.thread_id, self.id) async def stream(self) -> AsyncIterator[Notification]: """Yield only notifications routed to this async turn handle.""" await self._codex._ensure_initialized() + if self._goal is not None: + stream = self._goal.stream() + try: + async for event in stream: + yield event + finally: + await stream.aclose() + return self._codex._client.register_turn_notifications(self.id) try: while True: From 206478d338ebc5763a9a92e150bc3e15a28c914f Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 9 Jun 2026 16:29:57 -0700 Subject: [PATCH 2/2] Release goal routing when streams close --- sdk/python/src/openai_codex/api.py | 69 +++++++++++++++--------------- 1 file changed, 35 insertions(+), 34 deletions(-) diff --git a/sdk/python/src/openai_codex/api.py b/sdk/python/src/openai_codex/api.py index 18baa93e022b..c77ba3acf8cf 100644 --- a/sdk/python/src/openai_codex/api.py +++ b/sdk/python/src/openai_codex/api.py @@ -751,21 +751,24 @@ def interrupt(self) -> TurnInterruptResponse: def stream(self) -> Iterator[Notification]: """Yield only notifications routed to this turn handle.""" if self._goal is not None: - yield from self._goal.stream() - return - self._client.register_turn_notifications(self.id) - try: - while True: - event = self._client.next_turn_notification(self.id) - yield event - if ( - event.method == "turn/completed" - and isinstance(event.payload, TurnCompletedNotification) - and event.payload.turn.id == self.id - ): - break - finally: - self._client.unregister_turn_notifications(self.id) + return self._goal.stream() + + def notifications() -> Iterator[Notification]: + self._client.register_turn_notifications(self.id) + try: + while True: + event = self._client.next_turn_notification(self.id) + yield event + if ( + event.method == "turn/completed" + and isinstance(event.payload, TurnCompletedNotification) + and event.payload.turn.id == self.id + ): + break + finally: + self._client.unregister_turn_notifications(self.id) + + return notifications() def run(self) -> TurnResult: """Consume the turn stream and return its completed result.""" @@ -808,30 +811,28 @@ async def interrupt(self) -> TurnInterruptResponse: return await self._goal.interrupt() return await self._codex._client.turn_interrupt(self.thread_id, self.id) - async def stream(self) -> AsyncIterator[Notification]: + def stream(self) -> AsyncIterator[Notification]: """Yield only notifications routed to this async turn handle.""" - await self._codex._ensure_initialized() if self._goal is not None: - stream = self._goal.stream() + return self._goal.stream() + + async def notifications() -> AsyncIterator[Notification]: + await self._codex._ensure_initialized() + self._codex._client.register_turn_notifications(self.id) try: - async for event in stream: + while True: + event = await self._codex._client.next_turn_notification(self.id) yield event + if ( + event.method == "turn/completed" + and isinstance(event.payload, TurnCompletedNotification) + and event.payload.turn.id == self.id + ): + break finally: - await stream.aclose() - return - self._codex._client.register_turn_notifications(self.id) - try: - while True: - event = await self._codex._client.next_turn_notification(self.id) - yield event - if ( - event.method == "turn/completed" - and isinstance(event.payload, TurnCompletedNotification) - and event.payload.turn.id == self.id - ): - break - finally: - self._codex._client.unregister_turn_notifications(self.id) + self._codex._client.unregister_turn_notifications(self.id) + + return notifications() async def run(self) -> TurnResult: """Consume the turn stream and return its completed result."""