from __future__ import annotations import queue import threading import weakref from collections import deque from contextlib import contextmanager from dataclasses import dataclass, field from typing import Iterator from ._goal import _GoalOperationState from .errors import CodexError, TransportClosedError, map_jsonrpc_error from .generated.notification_registry import notification_turn_id from .generated.v2_all import AccountLoginCompletedNotification from .models import JsonValue, Notification, UnknownNotification ResponseQueueItem = JsonValue | BaseException NotificationQueueItem = Notification | BaseException @dataclass class _TurnState: id: str thread_id: str | None = None events: dict[int, NotificationQueueItem] = field(default_factory=dict) first_event: int = 0 next_event: int = 0 subscribers: dict[object, int] = field(default_factory=dict) completed: bool = False class _TurnSubscription: """One consumer's cursor over shared unread events.""" def __init__(self, router: MessageRouter, state: _TurnState, cursor: int) -> None: self._router = router self._state = state self._cursor = cursor self._token = object() state.subscribers[self._token] = self._cursor self._closed = False self._release = weakref.finalize( self, router._release_turn, weakref.ref(router), state, self._token ) def next(self) -> Notification: with self._router._turn_condition: while self._cursor == self._state.next_event and not self._closed: if self._state.completed: raise TransportClosedError("Turn is no longer streaming") self._router._turn_condition.wait() if self._closed: raise TransportClosedError("Turn subscription closed") item = self._state.events[self._cursor] self._cursor += 1 self._state.subscribers[self._token] = self._cursor self._router._prune_turn_events(self._state) if isinstance(item, BaseException): raise item return item def close(self) -> None: with self._router._turn_condition: self._closed = True self._router._turn_condition.notify_all() self._release() class MessageRouter: """Route reader-thread messages to the SDK operation waiting for them. The app-server stdio transport is a single ordered stream, so only the reader thread should consume stdout. This router keeps the rest of the SDK from competing for that stream by giving each in-flight JSON-RPC request its own queue and each turn consumer its own event cursor. """ def __init__(self) -> None: """Create empty response, turn, and global notification queues.""" # GC can release abandoned subscriptions during another routing operation. self._lock = threading.RLock() self._response_waiters: dict[str, queue.Queue[ResponseQueueItem]] = {} self._login_notifications: dict[str, queue.Queue[NotificationQueueItem]] = {} self._pending_login_notifications: dict[str, deque[Notification]] = {} self._turn_condition = threading.Condition(self._lock) self._turn_states: dict[str, _TurnState] = {} self._turn_notifications: dict[str, _TurnSubscription] = {} self._pending_turn_requests: dict[str, BaseException | None] = {} 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]: """Register a one-shot queue for a JSON-RPC response id.""" waiter: queue.Queue[ResponseQueueItem] = queue.Queue(maxsize=1) with self._lock: self._response_waiters[request_id] = waiter return waiter def discard_response_waiter(self, request_id: str) -> None: """Remove a response waiter when the request could not be written.""" with self._lock: self._response_waiters.pop(request_id, None) def next_global_notification(self) -> Notification: """Block until the next notification that is not scoped to a turn.""" item = self._global_notifications.get() if isinstance(item, BaseException): raise item return item def register_login(self, login_id: str) -> None: """Register a queue for one interactive login attempt.""" login_queue: queue.Queue[NotificationQueueItem] = queue.Queue() with self._lock: if login_id in self._login_notifications: return pending = self._pending_login_notifications.pop(login_id, deque()) self._login_notifications[login_id] = login_queue for notification in pending: login_queue.put(notification) def unregister_login(self, login_id: str) -> None: """Stop routing future notifications for one login attempt.""" with self._lock: self._login_notifications.pop(login_id, None) def next_login_notification(self, login_id: str) -> Notification: """Block until the next notification for a registered login attempt.""" with self._lock: login_queue = self._login_notifications.get(login_id) if login_queue is None: raise RuntimeError(f"login {login_id!r} is not registered for waiting") item = login_queue.get() if isinstance(item, BaseException): raise item return item @contextmanager def pending_turn(self, thread_id: str) -> Iterator[dict[str, int]]: """Buffer events from the point a turn/start request is sent.""" with self._lock: cursors = {turn_id: state.next_event for turn_id, state in self._turn_states.items()} self._pending_turn_requests[thread_id] = None try: yield cursors finally: with self._lock: del self._pending_turn_requests[thread_id] for state in list(self._turn_states.values()): if state.thread_id in (None, thread_id): self._prune_turn_events(state) def prepare_turn( self, turn_id: str, thread_id: str, cursors: dict[str, int], *, for_handle: bool ) -> _TurnSubscription | None: """Attach the requesting handle or the single low-level consumer.""" with self._lock: state = self._turn_states.setdefault(turn_id, _TurnState(turn_id, thread_id)) state.thread_id = thread_id if not for_handle and turn_id in self._turn_notifications: return None if not state.completed and (err := self._pending_turn_requests[thread_id]) is not None: state.events[state.next_event] = err state.next_event += 1 state.completed = True subscription = _TurnSubscription(self, state, cursors.get(turn_id, 0)) if not for_handle: self._turn_notifications[turn_id] = subscription return subscription def subscribe_turn(self, turn_id: str) -> _TurnSubscription: """Attach a consumer starting at the next event for this turn.""" with self._lock: state = self._turn_states.setdefault(turn_id, _TurnState(turn_id)) return _TurnSubscription(self, state, state.next_event) @staticmethod def _release_turn( router_ref: weakref.ReferenceType[MessageRouter], state: _TurnState, token: object ) -> None: router = router_ref() if router is not None: with router._lock: default = router._turn_notifications.get(state.id) if default is not None and default._token is token: del router._turn_notifications[state.id] state.subscribers.pop(token, None) router._prune_turn_events(state) def _prune_turn_events(self, state: _TurnState) -> None: if state.thread_id in self._pending_turn_requests or ( state.thread_id is None and self._pending_turn_requests ): return consumed = min(state.subscribers.values(), default=state.next_event) while state.first_event < consumed: del state.events[state.first_event] state.first_event += 1 if not state.subscribers and self._turn_states.get(state.id) is state: del self._turn_states[state.id] def register_turn(self, turn_id: str) -> None: """Register the default consumer used by the low-level client API.""" with self._lock: if turn_id not in self._turn_notifications: self._turn_notifications[turn_id] = self.subscribe_turn(turn_id) def unregister_turn(self, turn_id: str) -> None: """Close only the low-level consumer, leaving other handles subscribed.""" with self._lock: if subscription := self._turn_notifications.get(turn_id): subscription.close() def next_turn_notification(self, turn_id: str) -> Notification: """Block until the next event for the default low-level consumer.""" with self._lock: subscription = self._turn_notifications.get(turn_id) if subscription is None: raise RuntimeError(f"turn {turn_id!r} is not registered for streaming") return subscription.next() def register_goal(self, thread_id: str) -> _GoalOperationState: """Register one thread-scoped logical goal operation before it starts.""" state = _GoalOperationState(thread_id=thread_id) state.activate_turn_routing() return self._register_goal(state) def reserve_goal(self, thread_id: str) -> _GoalOperationState: """Reserve a thread route without accepting physical turns yet.""" return self._register_goal(_GoalOperationState(thread_id=thread_id)) def _register_goal(self, state: _GoalOperationState) -> _GoalOperationState: with self._lock: if state.thread_id in self._goal_operations: raise RuntimeError( f"thread {state.thread_id!r} already has an active goal operation" ) self._goal_operations[state.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 has_goal(self, thread_id: str) -> bool: """Return whether a logical goal operation owns this thread route.""" with self._lock: return thread_id in self._goal_operations def route_response(self, msg: dict[str, JsonValue]) -> None: """Deliver a JSON-RPC response or error to its request waiter.""" request_id = msg.get("id") with self._lock: waiter = self._response_waiters.pop(str(request_id), None) if waiter is None: return if "error" in msg: err = msg["error"] if isinstance(err, dict): waiter.put( map_jsonrpc_error( int(err.get("code", -32000)), str(err.get("message", "unknown")), err.get("data"), ) ) else: waiter.put(CodexError("Malformed JSON-RPC error response")) return waiter.put(msg.get("result")) def route_notification(self, notification: Notification) -> None: """Deliver a notification to a turn queue or the global queue.""" login_id = self._notification_login_id(notification) if login_id is not None: with self._lock: login_queue = self._login_notifications.get(login_id) if login_queue is None: self._pending_login_notifications.setdefault(login_id, deque()).append( notification ) return login_queue.put(notification) 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/") ): if 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 with self._turn_condition: state = self._turn_states.setdefault(turn_id, _TurnState(turn_id, thread_id)) state.thread_id = thread_id or state.thread_id state.events[state.next_event] = notification state.next_event += 1 if notification.method == "turn/completed": state.completed = True self._prune_turn_events(state) self._turn_condition.notify_all() def fail_all(self, exc: BaseException) -> None: """Wake every blocked waiter when the reader thread exits.""" with self._lock: response_waiters = list(self._response_waiters.values()) self._response_waiters.clear() login_queues = list(self._login_notifications.values()) self._login_notifications.clear() self._pending_login_notifications.clear() for thread_id in self._pending_turn_requests: self._pending_turn_requests[thread_id] = exc for state in list(self._turn_states.values()): state.events[state.next_event] = exc state.next_event += 1 state.completed = True self._prune_turn_events(state) self._turn_condition.notify_all() 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: waiter.put(exc) for login_queue in login_queues: login_queue.put(exc) for goal_operation in goal_operations: goal_operation.fail(exc) self._global_notifications.put(exc) def _notification_turn_id(self, notification: Notification) -> str | None: """Extract routing ids from generated metadata or raw unknown payloads.""" payload = notification.payload if isinstance(payload, UnknownNotification): raw_turn_id = payload.params.get("turnId") if isinstance(raw_turn_id, str): return raw_turn_id raw_turn = payload.params.get("turn") if isinstance(raw_turn, dict): raw_nested_turn_id = raw_turn.get("id") if isinstance(raw_nested_turn_id, str): 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