[1/6] Add Python goal routing foundation

sdk/python/src/openai_codex/_goal.py

@@ -0,0 +1,156 @@
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+
+from .generated.v2_all import (
+    ThreadGoalClearedNotification,
+    ThreadGoalStatus,
+    ThreadGoalUpdatedNotification,
+    Turn,
+    TurnCompletedNotification,
+    TurnStartedNotification,
+)
+from .models import Notification
+
+
+class _GoalStreamClosed(Exception):
+    """Wake a notification reader after its logical stream closes."""
+
+
+def _terminal_goal_status(status: ThreadGoalStatus | None) -> bool:
+    return status in {
+        ThreadGoalStatus.paused,
+        ThreadGoalStatus.blocked,
+        ThreadGoalStatus.usage_limited,
+        ThreadGoalStatus.budget_limited,
+        ThreadGoalStatus.complete,
+    }
+
+
+@dataclass(slots=True)
+class _GoalOperationState:
+    """Private state for one goal operation exposed as a logical turn."""
+
+    thread_id: str
+    logical_turn_id: str | None = None
+    current_turn_id: str | None = None
+    status: ThreadGoalStatus | None = None
+    started_turn: Turn | None = None
+    completed_turn: Turn | None = None
+    interrupted: bool = False
+    interrupt_requested: bool = False
+    cleared: bool = False
+    _condition: threading.Condition = field(default_factory=threading.Condition)
+    _notifications: queue.Queue[Notification | BaseException] = field(default_factory=queue.Queue)
+    _failure: BaseException | None = None
+    _finished: bool = False
+
+    def observe(self, notification: Notification) -> None:
+        payload = notification.payload
+        with self._condition:
+            if isinstance(payload, TurnStartedNotification):
+                if self.logical_turn_id is None:
+                    self.logical_turn_id = payload.turn.id
+                self.current_turn_id = payload.turn.id
+                if self.started_turn is None:
+                    self.started_turn = payload.turn
+            elif isinstance(payload, TurnCompletedNotification):
+                self.completed_turn = payload.turn
+                if self.current_turn_id == payload.turn.id:
+                    self.current_turn_id = None
+            elif isinstance(payload, ThreadGoalUpdatedNotification):
+                self.status = payload.goal.status
+                if self.status == ThreadGoalStatus.active:
+                    self.cleared = False
+            elif isinstance(payload, ThreadGoalClearedNotification):
+                self.cleared = True
+            if (
+                self.current_turn_id is None
+                and self.completed_turn is not None
+                and (self.cleared or _terminal_goal_status(self.status))
+            ):
+                self._finished = True
+            self._condition.notify_all()
+        self._notifications.put(notification)
+
+    def wait_for_start(self, timeout: float) -> str | None:
+        """Wait for the runtime-generated first turn without consuming its event."""
+        deadline = time.monotonic() + timeout
+        with self._condition:
+            while self.started_turn is None or self.logical_turn_id is None:
+                if self._failure is not None:
+                    raise self._failure
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    return None
+                self._condition.wait(remaining)
+            return self.logical_turn_id
+
+    def fail(self, exc: BaseException) -> None:
+        with self._condition:
+            self._failure = exc
+            self._condition.notify_all()
+        self._notifications.put(exc)
+
+    def next_notification(self) -> Notification:
+        item = self._notifications.get()
+        if isinstance(item, BaseException):
+            raise item
+        return item
+
+    def finish(self) -> None:
+        """Mark the logical operation inactive and wake waiting controls."""
+        with self._condition:
+            self._finished = True
+            self.current_turn_id = None
+            self._condition.notify_all()
+
+    def is_finished(self) -> bool:
+        with self._condition:
+            return self._finished
+
+    def begin_interrupt(self) -> bool:
+        with self._condition:
+            if self._finished:
+                return False
+            self.interrupt_requested = True
+            return True
+
+    def confirm_interrupt(self) -> None:
+        with self._condition:
+            self.interrupted = True
+            self.interrupt_requested = False
+
+    def cancel_interrupt(self) -> None:
+        with self._condition:
+            self.interrupt_requested = False
+
+    def explicit_interrupt(self, status: ThreadGoalStatus | None) -> bool:
+        with self._condition:
+            return self.interrupted or (
+                self.interrupt_requested and status == ThreadGoalStatus.paused
+            )
+
+    def active_turn(self, *, after: str | None = None) -> str | None:
+        """Wait for the current turn, or return None once the goal has ended."""
+        with self._condition:
+            while True:
+                if self._failure is not None:
+                    raise self._failure
+                if self._finished:
+                    return None
+                if self.current_turn_id is not None and self.current_turn_id != after:
+                    return self.current_turn_id
+                if self.cleared or _terminal_goal_status(self.status):
+                    return None
+                self._condition.wait()
+
+    def current_turn(self) -> str | None:
+        """Return the current physical turn without waiting for rollover."""
+        with self._condition:
+            return self.current_turn_id
+
+    def wake_notification_reader(self) -> None:
+        """Release a reader blocked after its stream has been closed."""
+        self._notifications.put(_GoalStreamClosed())

sdk/python/src/openai_codex/_message_router.py

@@ -4,6 +4,7 @@
 import threading
 from collections import deque
 
+from ._goal import _GoalOperationState
 from .errors import CodexError, map_jsonrpc_error
 from .generated.notification_registry import notification_turn_id
 from .generated.v2_all import AccountLoginCompletedNotification
@@ -30,6 +31,7 @@ def __init__(self) -> None:
         self._pending_login_notifications: dict[str, deque[Notification]] = {}
         self._turn_notifications: dict[str, queue.Queue[NotificationQueueItem]] = {}
         self._pending_turn_notifications: dict[str, deque[Notification]] = {}
+        self._goal_operations: dict[str, _GoalOperationState] = {}
         self._global_notifications: queue.Queue[NotificationQueueItem] = queue.Queue()
 
     def create_response_waiter(self, request_id: str) -> queue.Queue[ResponseQueueItem]:
@@ -116,6 +118,21 @@ def next_turn_notification(self, turn_id: str) -> Notification:
             raise item
         return item
 
+    def register_goal(self, thread_id: str) -> _GoalOperationState:
+        """Register one thread-scoped logical goal operation before it starts."""
+        state = _GoalOperationState(thread_id=thread_id)
+        with self._lock:
+            if thread_id in self._goal_operations:
+                raise RuntimeError(f"thread {thread_id!r} already has an active goal operation")
+            self._goal_operations[thread_id] = state
+        return state
+
+    def unregister_goal(self, state: _GoalOperationState) -> None:
+        """Stop routing notifications to a completed logical goal operation."""
+        with self._lock:
+            if self._goal_operations.get(state.thread_id) is state:
+                self._goal_operations.pop(state.thread_id)
+
     def route_response(self, msg: dict[str, JsonValue]) -> None:
         """Deliver a JSON-RPC response or error to its request waiter."""
 
@@ -157,6 +174,17 @@ def route_notification(self, notification: Notification) -> None:
             return
 
         turn_id = self._notification_turn_id(notification)
+        thread_id = self._notification_thread_id(notification)
+        if thread_id is not None:
+            with self._lock:
+                goal_state = self._goal_operations.get(thread_id)
+            if goal_state is not None and (
+                turn_id is not None or notification.method.startswith("thread/goal/")
+            ):
+                goal_state.observe(notification)
+                if goal_state.is_finished():
+                    self.unregister_goal(goal_state)
+                return
         if turn_id is None:
             self._global_notifications.put(notification)
             return
@@ -182,6 +210,8 @@ def fail_all(self, exc: BaseException) -> None:
             self._pending_login_notifications.clear()
             turn_queues = list(self._turn_notifications.values())
             self._pending_turn_notifications.clear()
+            goal_operations = list(self._goal_operations.values())
+            self._goal_operations.clear()
         # Put the same transport failure into every queue so no SDK call blocks
         # forever waiting for a response that cannot arrive.
         for waiter in response_waiters:
@@ -190,24 +220,12 @@ def fail_all(self, exc: BaseException) -> None:
             login_queue.put(exc)
         for turn_queue in turn_queues:
             turn_queue.put(exc)
+        for goal_operation in goal_operations:
+            goal_operation.fail(exc)
         self._global_notifications.put(exc)
 
-    def _notification_login_id(self, notification: Notification) -> str | None:
-        """Extract the login attempt id from completion notifications."""
-        if notification.method != "account/login/completed":
-            return None
-
-        payload = notification.payload
-        if isinstance(payload, AccountLoginCompletedNotification):
-            return payload.login_id
-        if isinstance(payload, UnknownNotification):
-            raw_login_id = payload.params.get("loginId")
-            if isinstance(raw_login_id, str):
-                return raw_login_id
-        return None
-
     def _notification_turn_id(self, notification: Notification) -> str | None:
-        """Extract routing ids from known generated payloads or raw unknown payloads."""
+        """Extract routing ids from generated metadata or raw unknown payloads."""
         payload = notification.payload
         if isinstance(payload, UnknownNotification):
             raw_turn_id = payload.params.get("turnId")
@@ -220,3 +238,26 @@ def _notification_turn_id(self, notification: Notification) -> str | None:
                     return raw_nested_turn_id
             return None
         return notification_turn_id(payload)
+
+    def _notification_thread_id(self, notification: Notification) -> str | None:
+        """Extract thread ids from typed payloads or raw unknown payloads."""
+        payload = notification.payload
+        if isinstance(payload, UnknownNotification):
+            raw_thread_id = payload.params.get("threadId")
+            return raw_thread_id if isinstance(raw_thread_id, str) else None
+        thread_id = getattr(payload, "thread_id", None)
+        return thread_id if isinstance(thread_id, str) else None
+
+    def _notification_login_id(self, notification: Notification) -> str | None:
+        """Extract the login attempt id from completion notifications."""
+        if notification.method != "account/login/completed":
+            return None
+
+        payload = notification.payload
+        if isinstance(payload, AccountLoginCompletedNotification):
+            return payload.login_id
+        if isinstance(payload, UnknownNotification):
+            raw_login_id = payload.params.get("loginId")
+            if isinstance(raw_login_id, str):
+                return raw_login_id
+        return None

sdk/python/src/openai_codex/async_client.py

@@ -6,6 +6,7 @@
 
 from pydantic import BaseModel
 
+from ._goal import _GoalOperationState
 from .client import CodexClient, CodexConfig
 from .generated.v2_all import (
     AccountLoginCompletedNotification,
@@ -21,6 +22,9 @@
     ThreadCompactStartResponse,
     ThreadForkParams as V2ThreadForkParams,
     ThreadForkResponse,
+    ThreadGoalClearResponse,
+    ThreadGoalSetResponse,
+    ThreadGoalStatus,
     ThreadListParams as V2ThreadListParams,
     ThreadListResponse,
     ThreadReadResponse,
@@ -107,6 +111,14 @@ def unregister_turn_notifications(self, turn_id: str) -> None:
         """Unregister a turn notification queue on the wrapped sync client."""
         self._sync.unregister_turn_notifications(turn_id)
 
+    def register_goal_operation(self, thread_id: str) -> _GoalOperationState:
+        """Register a logical goal route on the wrapped sync client."""
+        return self._sync.register_goal_operation(thread_id)
+
+    def unregister_goal_operation(self, state: _GoalOperationState) -> None:
+        """Release one logical goal route."""
+        self._sync.unregister_goal_operation(state)
+
     async def request(
         self,
         method: str,
@@ -192,6 +204,29 @@ async def thread_compact(self, thread_id: str) -> ThreadCompactStartResponse:
         """Start thread compaction using the wrapped sync client."""
         return await self._call_sync(self._sync.thread_compact, thread_id)
 
+    async def thread_goal_clear(self, thread_id: str) -> ThreadGoalClearResponse:
+        """Clear the persisted goal through the wrapped sync client."""
+        return await self._call_sync(self._sync.thread_goal_clear, thread_id)
+
+    async def thread_goal_set(
+        self,
+        thread_id: str,
+        *,
+        objective: str | None = None,
+        status: ThreadGoalStatus | None = None,
+    ) -> ThreadGoalSetResponse:
+        """Create or update a persisted goal through the wrapped sync client."""
+        return await self._call_sync(
+            self._sync.thread_goal_set,
+            thread_id,
+            objective=objective,
+            status=status,
+        )
+
+    async def pause_goal(self, thread_id: str) -> ThreadGoalSetResponse:
+        """Pause the active goal through the wrapped sync client."""
+        return await self._call_sync(self._sync.pause_goal, thread_id)
+
     async def turn_start(
         self,
         thread_id: str,
@@ -256,6 +291,10 @@ async def next_turn_notification(self, turn_id: str) -> Notification:
         """Wait for the next notification routed to one turn."""
         return await self._call_sync(self._sync.next_turn_notification, turn_id)
 
+    async def next_goal_notification(self, state: _GoalOperationState) -> Notification:
+        """Wait for the next notification in a logical goal turn."""
+        return await self._call_sync(self._sync.next_goal_notification, state)
+
     async def wait_for_login_completed(
         self,
         login_id: str,