[3/6] Add private Python goal-aware turn handles

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:

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),
+        )

sdk/python/src/openai_codex/api.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
 import asyncio
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import AsyncIterator, Iterator
 
 from ._approval_mode import (
     ApprovalMode as ApprovalMode,
     _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,23 +744,31 @@ 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."""
-        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)
+        if self._goal is not None:
+            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."""
@@ -766,10 +786,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,24 +807,32 @@ 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]:
+    def stream(self) -> AsyncIterator[Notification]:
         """Yield only notifications routed to this async turn handle."""
-        await self._codex._ensure_initialized()
-        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)
+        if self._goal is not None:
+            return self._goal.stream()
+
+        async def notifications() -> AsyncIterator[Notification]:
+            await self._codex._ensure_initialized()
+            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)
+
+        return notifications()
 
     async def run(self) -> TurnResult:
         """Consume the turn stream and return its completed result."""