mirror of
https://github.com/openai/codex.git
synced 2026-10-26 18:38:47 +00:00
Start Python SDK turn subscriptions at their attachment point (#44400)
## What changed Make turn event delivery depend on when each consumer attaches: - Handles returned by `thread.turn(...)` receive events from when the request is sent, including events received before the response arrives. - Other handles and explicit low-level registrations receive only events from when they join. Remove replay of completed items, token usage, and terminal events; use `thread.read(include_turns=True)` for earlier history. - Retain unread events for active subscribers and pending starts, release unused turn state, and close the returned subscription when an async start is cancelled. Preserve the existing low-level subscription when repeated `turn_start(...)` calls return the same turn. ## Testing Update coverage for late joins, request-time buffering, completion or transport failure before a start returns, async cancellation cleanup, and repeated low-level starts. Adjust external-message integration tests to allow handles to collect different items while receiving the same final response. GitOrigin-RevId: 3b7f1268df1193dee88385006ae9510101086c24
This commit is contained in:
@@ -233,6 +233,9 @@ passed to `run(...)` or `turn(...)` applies to that turn and subsequent turns.
|
||||
|
||||
## TurnHandle / AsyncTurnHandle
|
||||
|
||||
A `thread.turn(...)` handle receives events from when the call sends its request.
|
||||
Other handles start when they join; use `thread.read(include_turns=True)` for earlier history.
|
||||
|
||||
### TurnHandle
|
||||
|
||||
- `steer(input: str | Input) -> TurnSteerResponse`
|
||||
|
||||
@@ -1294,8 +1294,8 @@ def _render_thread_block(turn_fields: list[PublicFieldSpec], *, is_async: bool =
|
||||
*_approval_mode_model_arg_lines(),
|
||||
*_model_arg_lines(turn_fields),
|
||||
" )",
|
||||
f" turn = {await_prefix}{client}.turn_start(self.id, wire_input, params=params)",
|
||||
f" return {handle_type}({handle_owner}, self.id, turn.turn.id)",
|
||||
f" turn, subscription = {await_prefix}{client}._start_turn(self.id, wire_input, params=params, for_handle=True)",
|
||||
f" return {handle_type}({handle_owner}, self.id, turn.turn.id, _subscription=subscription)",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@@ -11,11 +11,7 @@ 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,
|
||||
ItemCompletedNotification,
|
||||
ThreadTokenUsageUpdatedNotification,
|
||||
)
|
||||
from .generated.v2_all import AccountLoginCompletedNotification
|
||||
from .models import JsonValue, Notification, UnknownNotification
|
||||
|
||||
ResponseQueueItem = JsonValue | BaseException
|
||||
@@ -30,29 +26,18 @@ class _TurnState:
|
||||
first_event: int = 0
|
||||
next_event: int = 0
|
||||
subscribers: dict[object, int] = field(default_factory=dict)
|
||||
subscribed: bool = False
|
||||
completed_items: dict[str, Notification] = field(default_factory=dict)
|
||||
usage: Notification | None = None
|
||||
terminal: NotificationQueueItem | None = None
|
||||
unclaimed: int = 0
|
||||
completed: bool = False
|
||||
|
||||
|
||||
class _TurnSubscription:
|
||||
"""One consumer's result snapshot and cursor over shared unread events."""
|
||||
"""One consumer's cursor over shared unread events."""
|
||||
|
||||
def __init__(self, router: MessageRouter, state: _TurnState) -> None:
|
||||
def __init__(self, router: MessageRouter, state: _TurnState, cursor: int) -> None:
|
||||
self._router = router
|
||||
self._state = state
|
||||
self._cursor = state.first_event
|
||||
self._cursor = cursor
|
||||
self._token = object()
|
||||
state.subscribers[self._token] = self._cursor
|
||||
state.subscribed = True
|
||||
self._replay: deque[NotificationQueueItem] = deque(state.completed_items.values())
|
||||
if state.usage is not None:
|
||||
self._replay.append(state.usage)
|
||||
if state.terminal is not None:
|
||||
self._replay.append(state.terminal)
|
||||
self._closed = False
|
||||
self._release = weakref.finalize(
|
||||
self, router._release_turn, weakref.ref(router), state, self._token
|
||||
@@ -60,17 +45,16 @@ class _TurnSubscription:
|
||||
|
||||
def next(self) -> Notification:
|
||||
with self._router._turn_condition:
|
||||
while not self._replay and self._cursor == self._state.next_event and not self._closed:
|
||||
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")
|
||||
if self._replay:
|
||||
item = self._replay.popleft()
|
||||
else:
|
||||
item = self._state.events[self._cursor]
|
||||
self._cursor += 1
|
||||
self._state.subscribers[self._token] = self._cursor
|
||||
self._router._prune_turn_events(self._state)
|
||||
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
|
||||
@@ -78,7 +62,6 @@ class _TurnSubscription:
|
||||
def close(self) -> None:
|
||||
with self._router._turn_condition:
|
||||
self._closed = True
|
||||
self._replay.clear()
|
||||
self._router._turn_condition.notify_all()
|
||||
self._release()
|
||||
|
||||
@@ -101,8 +84,8 @@ class MessageRouter:
|
||||
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 | None] = {}
|
||||
self._pending_turn_requests: dict[str, int] = {}
|
||||
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()
|
||||
|
||||
@@ -159,41 +142,43 @@ class MessageRouter:
|
||||
return item
|
||||
|
||||
@contextmanager
|
||||
def pending_turn(self, thread_id: str) -> Iterator[None]:
|
||||
"""Retain early completion while a turn/start response is in flight."""
|
||||
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:
|
||||
self._pending_turn_requests[thread_id] = (
|
||||
self._pending_turn_requests.get(thread_id, 0) + 1
|
||||
)
|
||||
cursors = {turn_id: state.next_event for turn_id, state in self._turn_states.items()}
|
||||
self._pending_turn_requests[thread_id] = None
|
||||
try:
|
||||
yield
|
||||
yield cursors
|
||||
finally:
|
||||
with self._lock:
|
||||
self._pending_turn_requests[thread_id] -= 1
|
||||
if self._pending_turn_requests[thread_id] == 0:
|
||||
del self._pending_turn_requests[thread_id]
|
||||
del self._pending_turn_requests[thread_id]
|
||||
for state in list(self._turn_states.values()):
|
||||
if state.thread_id == thread_id:
|
||||
if state.thread_id in (None, thread_id):
|
||||
self._prune_turn_events(state)
|
||||
self._discard_finished_turn(state)
|
||||
|
||||
def prepare_turn(self, turn_id: str, thread_id: str) -> None:
|
||||
"""Reserve the returned turn for a handle or a low-level consumer."""
|
||||
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
|
||||
state.unclaimed += 1
|
||||
self._turn_notifications.setdefault(turn_id, None)
|
||||
|
||||
def _subscribe_turn_locked(self, turn_id: str) -> _TurnSubscription:
|
||||
state = self._turn_states.setdefault(turn_id, _TurnState(turn_id))
|
||||
state.unclaimed = max(0, state.unclaimed - 1)
|
||||
return _TurnSubscription(self, state)
|
||||
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 with completed items, latest usage, and unread events."""
|
||||
"""Attach a consumer starting at the next event for this turn."""
|
||||
with self._lock:
|
||||
return self._subscribe_turn_locked(turn_id)
|
||||
state = self._turn_states.setdefault(turn_id, _TurnState(turn_id))
|
||||
return _TurnSubscription(self, state, state.next_event)
|
||||
|
||||
@staticmethod
|
||||
def _release_turn(
|
||||
@@ -202,74 +187,42 @@ class MessageRouter:
|
||||
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)
|
||||
router._discard_finished_turn(state)
|
||||
|
||||
def _prune_turn_events(self, state: _TurnState) -> None:
|
||||
if (
|
||||
not state.subscribed
|
||||
or state.unclaimed
|
||||
or self._pending_turn_requests.get(state.thread_id, 0)
|
||||
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:
|
||||
event = state.events.pop(state.first_event)
|
||||
del state.events[state.first_event]
|
||||
state.first_event += 1
|
||||
# Late joins need the completed result, not every consumed token delta
|
||||
# or intermediate usage update. Keep one snapshot entry per item.
|
||||
if isinstance(event, BaseException) or event.method == "turn/completed":
|
||||
state.terminal = event
|
||||
elif isinstance(event.payload, ItemCompletedNotification):
|
||||
item = event.payload.item
|
||||
state.completed_items[getattr(item, "root", item).id] = event
|
||||
elif isinstance(event.payload, ThreadTokenUsageUpdatedNotification):
|
||||
state.usage = event
|
||||
|
||||
def _discard_finished_turn(self, state: _TurnState) -> None:
|
||||
if (
|
||||
state.completed
|
||||
and not state.subscribers
|
||||
and not state.unclaimed
|
||||
and not self._pending_turn_requests.get(state.thread_id, 0)
|
||||
):
|
||||
self._turn_states.pop(state.id, None)
|
||||
if self._turn_notifications.get(state.id) is None:
|
||||
self._turn_notifications.pop(state.id, None)
|
||||
state.events.clear()
|
||||
state.completed_items.clear()
|
||||
state.usage = None
|
||||
state.terminal = None
|
||||
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 self._turn_notifications.get(turn_id) is None:
|
||||
self._turn_notifications[turn_id] = self._subscribe_turn_locked(turn_id)
|
||||
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 turn_id not in self._turn_notifications:
|
||||
return
|
||||
subscription = self._turn_notifications.pop(turn_id)
|
||||
if subscription is None and (state := self._turn_states.get(turn_id)) is not None:
|
||||
state.unclaimed = max(0, state.unclaimed - 1)
|
||||
self._prune_turn_events(state)
|
||||
self._discard_finished_turn(state)
|
||||
if subscription is not None:
|
||||
subscription.close()
|
||||
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:
|
||||
if turn_id not in self._turn_notifications:
|
||||
raise RuntimeError(f"turn {turn_id!r} is not registered for streaming")
|
||||
subscription = self._turn_notifications[turn_id]
|
||||
subscription = self._turn_notifications.get(turn_id)
|
||||
if subscription is None:
|
||||
subscription = self._subscribe_turn_locked(turn_id)
|
||||
self._turn_notifications[turn_id] = subscription
|
||||
raise RuntimeError(f"turn {turn_id!r} is not registered for streaming")
|
||||
return subscription.next()
|
||||
|
||||
def register_goal(self, thread_id: str) -> _GoalOperationState:
|
||||
@@ -363,10 +316,9 @@ class MessageRouter:
|
||||
state.thread_id = thread_id or state.thread_id
|
||||
state.events[state.next_event] = notification
|
||||
state.next_event += 1
|
||||
self._prune_turn_events(state)
|
||||
if notification.method == "turn/completed":
|
||||
state.completed = True
|
||||
self._discard_finished_turn(state)
|
||||
self._prune_turn_events(state)
|
||||
self._turn_condition.notify_all()
|
||||
|
||||
def fail_all(self, exc: BaseException) -> None:
|
||||
@@ -378,12 +330,13 @@ class MessageRouter:
|
||||
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
|
||||
self._prune_turn_events(state)
|
||||
state.completed = True
|
||||
self._discard_finished_turn(state)
|
||||
self._prune_turn_events(state)
|
||||
self._turn_condition.notify_all()
|
||||
goal_operations = list(self._goal_operations.values())
|
||||
self._goal_operations.clear()
|
||||
|
||||
@@ -649,8 +649,10 @@ class Thread:
|
||||
summary=summary,
|
||||
service_tier_for_turn=turn_service_tier,
|
||||
)
|
||||
turn = self._client.turn_start(self.id, wire_input, params=params)
|
||||
return TurnHandle(self._client, self.id, turn.turn.id)
|
||||
turn, subscription = self._client._start_turn(
|
||||
self.id, wire_input, params=params, for_handle=True
|
||||
)
|
||||
return TurnHandle(self._client, self.id, turn.turn.id, _subscription=subscription)
|
||||
|
||||
# END GENERATED: Thread.flat_methods
|
||||
|
||||
@@ -754,8 +756,10 @@ class AsyncThread:
|
||||
summary=summary,
|
||||
service_tier_for_turn=turn_service_tier,
|
||||
)
|
||||
turn = await self._codex._client.turn_start(self.id, wire_input, params=params)
|
||||
return AsyncTurnHandle(self._codex, self.id, turn.turn.id)
|
||||
turn, subscription = await self._codex._client._start_turn(
|
||||
self.id, wire_input, params=params, for_handle=True
|
||||
)
|
||||
return AsyncTurnHandle(self._codex, self.id, turn.turn.id, _subscription=subscription)
|
||||
|
||||
# END GENERATED: AsyncThread.flat_methods
|
||||
|
||||
@@ -782,6 +786,15 @@ class TurnHandle:
|
||||
id: str
|
||||
_subscription: _TurnSubscription = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __init__(
|
||||
self, _client: CodexClient, thread_id: str, id: str, *, _subscription=None
|
||||
) -> None:
|
||||
self._client, self.thread_id, self.id = _client, thread_id, id
|
||||
if _subscription is None:
|
||||
self.__post_init__()
|
||||
else:
|
||||
self._subscription = _subscription
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._subscription = self._client._subscribe_turn_notifications(self.id)
|
||||
|
||||
@@ -830,6 +843,13 @@ class AsyncTurnHandle:
|
||||
id: str
|
||||
_subscription: _TurnSubscription = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __init__(self, _codex: AsyncCodex, thread_id: str, id: str, *, _subscription=None) -> None:
|
||||
self._codex, self.thread_id, self.id = _codex, thread_id, id
|
||||
if _subscription is None:
|
||||
self.__post_init__()
|
||||
else:
|
||||
self._subscription = _subscription
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._subscription = self._codex._client._subscribe_turn_notifications(self.id)
|
||||
|
||||
|
||||
@@ -296,19 +296,31 @@ class AsyncCodexClient:
|
||||
params: V2TurnStartParams | JsonObject | None = None,
|
||||
) -> TurnStartResponse:
|
||||
"""Start a turn, releasing an unclaimed result if the caller is cancelled."""
|
||||
return (await self._start_turn(thread_id, input_items, params, for_handle=False))[0]
|
||||
|
||||
async def _start_turn(
|
||||
self,
|
||||
thread_id: str,
|
||||
input_items: list[JsonObject] | JsonObject | str,
|
||||
params: V2TurnStartParams | JsonObject | None,
|
||||
for_handle: bool,
|
||||
) -> tuple[TurnStartResponse, _TurnSubscription | None]:
|
||||
operation = _TURN_START_EXECUTOR.submit(
|
||||
copy_context().run, self._sync.turn_start, thread_id, input_items, params
|
||||
copy_context().run, self._sync._start_turn, thread_id, input_items, params, for_handle
|
||||
)
|
||||
try:
|
||||
return await asyncio.wrap_future(operation)
|
||||
except asyncio.CancelledError:
|
||||
|
||||
def discard_cancelled_result(completed: Future[TurnStartResponse]) -> None:
|
||||
def discard_cancelled_result(
|
||||
completed: Future[tuple[TurnStartResponse, _TurnSubscription | None]],
|
||||
) -> None:
|
||||
try:
|
||||
started = completed.result()
|
||||
_, subscription = completed.result()
|
||||
except BaseException:
|
||||
return
|
||||
self._sync._subscribe_turn_notifications(started.turn.id).close()
|
||||
if subscription is not None:
|
||||
subscription.close()
|
||||
|
||||
operation.add_done_callback(discard_cancelled_result)
|
||||
raise
|
||||
|
||||
@@ -654,6 +654,15 @@ class CodexClient:
|
||||
params: V2TurnStartParams | JsonObject | None = None,
|
||||
) -> TurnStartResponse:
|
||||
"""Start a turn and register its notification queue as early as possible."""
|
||||
return self._start_turn(thread_id, input_items, params, for_handle=False)[0]
|
||||
|
||||
def _start_turn(
|
||||
self,
|
||||
thread_id: str,
|
||||
input_items: list[JsonObject] | JsonObject | str,
|
||||
params: V2TurnStartParams | JsonObject | None,
|
||||
for_handle: bool,
|
||||
) -> tuple[TurnStartResponse, _TurnSubscription | None]:
|
||||
with self._thread_start_lock(thread_id):
|
||||
if self._router.has_goal(thread_id):
|
||||
raise InvalidRequestError(
|
||||
@@ -665,10 +674,12 @@ class CodexClient:
|
||||
"threadId": thread_id,
|
||||
"input": self._normalize_input_items(input_items),
|
||||
}
|
||||
with self._router.pending_turn(thread_id):
|
||||
with self._router.pending_turn(thread_id) as cursors:
|
||||
started = self.request("turn/start", payload, response_model=TurnStartResponse)
|
||||
self._router.prepare_turn(started.turn.id, thread_id)
|
||||
return started
|
||||
subscription = self._router.prepare_turn(
|
||||
started.turn.id, thread_id, cursors, for_handle=for_handle
|
||||
)
|
||||
return started, subscription
|
||||
|
||||
@contextmanager
|
||||
def _thread_start_lock(self, thread_id: str) -> Iterator[None]:
|
||||
|
||||
@@ -82,7 +82,10 @@ def test_external_message_joins_active_turn_with_tool_authority(tmp_path) -> Non
|
||||
harness.responses.wait_for_requests(1)
|
||||
joined = thread.turn(ExternalMessage(tool_name="notifications", content=content))
|
||||
result = consumers.submit(joined.run).result(timeout=15)
|
||||
assert original_result.result(timeout=15) == result
|
||||
first = original_result.result(timeout=15)
|
||||
assert first.final_response == result.final_response
|
||||
assert first.items[0].root.type == "userMessage"
|
||||
assert all(item.root.type != "userMessage" for item in result.items)
|
||||
assert codex._client._router._turn_states == {}
|
||||
requests = harness.responses.requests()
|
||||
|
||||
@@ -151,7 +154,7 @@ def test_async_external_message_allows_both_handles_to_consume(tmp_path) -> None
|
||||
first, second = await asyncio.wait_for(
|
||||
asyncio.gather(original_result, joined.run()), timeout=15
|
||||
)
|
||||
assert first == second
|
||||
assert (first.final_response, first.usage) == (second.final_response, second.usage)
|
||||
assert first.final_response == "Update processed"
|
||||
assert first.usage is not None
|
||||
assert codex._client._sync._router._turn_states == {}
|
||||
|
||||
@@ -617,9 +617,10 @@ def test_client_reader_routes_interleaved_turn_notifications_by_turn_id() -> Non
|
||||
)
|
||||
|
||||
|
||||
def test_turn_notification_router_buffers_events_before_registration() -> None:
|
||||
"""Early turn events should be replayed once their TurnHandle registers."""
|
||||
def test_turn_notification_router_starts_at_explicit_registration() -> None:
|
||||
"""Explicit registration receives events from when the caller attaches."""
|
||||
client = CodexClient()
|
||||
client.register_turn_notifications("turn-1")
|
||||
client._router.route_notification(
|
||||
client._coerce_notification(
|
||||
"item/agentMessage/delta",
|
||||
@@ -632,7 +633,6 @@ def test_turn_notification_router_buffers_events_before_registration() -> None:
|
||||
)
|
||||
)
|
||||
|
||||
client.register_turn_notifications("turn-1")
|
||||
event = client.next_turn_notification("turn-1")
|
||||
|
||||
assert isinstance(event.payload, AgentMessageDeltaNotification)
|
||||
|
||||
@@ -191,10 +191,10 @@ def test_turn_inputs_and_options_reach_the_client(api_type, method, content) ->
|
||||
}
|
||||
),
|
||||
)
|
||||
subscription = SimpleNamespace(next=Mock(return_value=completed), close=Mock())
|
||||
client = SimpleNamespace(
|
||||
turn_start=rpc(return_value=SimpleNamespace(turn=SimpleNamespace(id="turn-1"))),
|
||||
_subscribe_turn_notifications=Mock(
|
||||
return_value=SimpleNamespace(next=Mock(return_value=completed), close=Mock())
|
||||
_start_turn=rpc(
|
||||
return_value=(SimpleNamespace(turn=SimpleNamespace(id="turn-1")), subscription)
|
||||
),
|
||||
)
|
||||
codex = api_type.__new__(api_type)
|
||||
@@ -229,7 +229,7 @@ def test_turn_inputs_and_options_reach_the_client(api_type, method, content) ->
|
||||
"toolOutput": {"name": "notifications", "namespace": "slack", "output": content},
|
||||
}
|
||||
)
|
||||
assert _params_dict(client.turn_start.call_args.kwargs["params"]) == {
|
||||
assert _params_dict(client._start_turn.call_args.kwargs["params"]) == {
|
||||
"threadId": "thread-1",
|
||||
"input": expected_input,
|
||||
"serviceTier": "priority",
|
||||
@@ -249,7 +249,7 @@ def test_external_messages_cannot_be_mixed_with_user_input_or_sent_as_user_steer
|
||||
async_api = api_type is AsyncCodex
|
||||
rpc = AsyncMock if async_api else Mock
|
||||
client = SimpleNamespace(
|
||||
turn_start=rpc(), turn_steer=rpc(), _subscribe_turn_notifications=Mock()
|
||||
_start_turn=rpc(), turn_steer=rpc(), _subscribe_turn_notifications=Mock()
|
||||
)
|
||||
codex = api_type.__new__(api_type)
|
||||
codex._client = client
|
||||
@@ -277,7 +277,7 @@ def test_external_messages_cannot_be_mixed_with_user_input_or_sent_as_user_steer
|
||||
result = thread.turn(ExternalMessage(tool_name=" ", content="Untrusted update"))
|
||||
if async_api:
|
||||
await result
|
||||
client.turn_start.assert_not_called()
|
||||
client._start_turn.assert_not_called()
|
||||
client.turn_steer.assert_not_called()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
@@ -8,7 +8,7 @@ import pytest
|
||||
|
||||
from openai_codex import AsyncCodex
|
||||
from openai_codex._run import _collect_turn_result
|
||||
from openai_codex.api import AsyncTurnHandle, TurnHandle
|
||||
from openai_codex.api import AsyncThread, AsyncTurnHandle, Thread, TurnHandle
|
||||
from openai_codex.async_client import AsyncCodexClient
|
||||
from openai_codex.client import CodexClient
|
||||
from openai_codex.errors import TransportClosedError
|
||||
@@ -55,23 +55,24 @@ def turn_events(client, *, status="completed"):
|
||||
]
|
||||
|
||||
|
||||
def test_late_join_replays_items_already_consumed_by_original_handle():
|
||||
@pytest.mark.parametrize("consumed", [False, True])
|
||||
def test_late_join_starts_with_future_events(consumed):
|
||||
client = CodexClient()
|
||||
original = TurnHandle(client, "thread-1", "turn-1")
|
||||
events = turn_events(client)
|
||||
client._router.route_notification(events[0])
|
||||
stream = original.stream()
|
||||
first = next(stream)
|
||||
client._router.route_notification(events[1])
|
||||
usage = next(stream)
|
||||
previous = [next(stream)] if consumed else []
|
||||
joined = TurnHandle(client, "thread-1", "turn-1")
|
||||
client._router.route_notification(events[1])
|
||||
for event in events[2:]:
|
||||
client._router.route_notification(event)
|
||||
|
||||
first_result = _collect_turn_result(chain([first, usage], stream), turn_id="turn-1")
|
||||
assert joined.run() == first_result
|
||||
first_result = _collect_turn_result(chain(previous, stream), turn_id="turn-1")
|
||||
joined_result = joined.run()
|
||||
assert first_result.final_response == "done"
|
||||
assert first_result.usage.last.total_tokens == 5
|
||||
assert joined_result.final_response is None
|
||||
assert joined_result.usage == first_result.usage
|
||||
assert client._router._turn_states == {}
|
||||
|
||||
|
||||
@@ -92,7 +93,6 @@ def test_consumed_deltas_are_released_while_turn_is_active():
|
||||
client._router.route_notification(event)
|
||||
assert subscription.next() == event
|
||||
assert state.events == {}
|
||||
assert state.completed_items == {}
|
||||
assert not state.completed
|
||||
subscription.close()
|
||||
|
||||
@@ -113,32 +113,69 @@ def test_slow_subscriber_keeps_unread_deltas_until_it_consumes_them():
|
||||
slow.close()
|
||||
|
||||
|
||||
def test_completion_before_turn_start_response_is_replayed(monkeypatch):
|
||||
client = CodexClient()
|
||||
@pytest.mark.parametrize("async_api", [False, True])
|
||||
@pytest.mark.parametrize("low_level", [False, True])
|
||||
@pytest.mark.parametrize("completed", [False, True])
|
||||
def test_events_or_failure_before_turn_start_returns(monkeypatch, async_api, low_level, completed):
|
||||
codex = AsyncCodex()
|
||||
codex._initialized = True
|
||||
client = codex._client._sync if async_api else CodexClient()
|
||||
|
||||
def request_raw(method, params):
|
||||
assert method == "turn/start"
|
||||
for event in turn_events(client):
|
||||
for event in turn_events(client) if completed else []:
|
||||
client._router.route_notification(event)
|
||||
client._router.fail_all(TransportClosedError("transport failed"))
|
||||
return {"turn": {"id": "turn-1", "status": "inProgress", "items": []}}
|
||||
|
||||
monkeypatch.setattr(client, "_request_raw", request_raw)
|
||||
started = client.turn_start("thread-1", "hello")
|
||||
handle = TurnHandle(client, "thread-1", started.turn.id)
|
||||
assert handle.run().final_response == "done"
|
||||
public = codex._client if async_api else client
|
||||
thread = AsyncThread(codex, "thread-1") if async_api else Thread(client, "thread-1")
|
||||
|
||||
async def value(call):
|
||||
return await call if async_api else call
|
||||
|
||||
async def scenario():
|
||||
if low_level:
|
||||
started = await value(public.turn_start("thread-1", "hello"))
|
||||
assert (
|
||||
await value(public.wait_for_turn_completed(started.turn.id))
|
||||
).turn.id == "turn-1"
|
||||
else:
|
||||
handle = await value(thread.turn("hello"))
|
||||
assert (await value(handle.run())).final_response == "done"
|
||||
|
||||
if completed:
|
||||
asyncio.run(scenario())
|
||||
else:
|
||||
with pytest.raises(TransportClosedError, match="transport failed"):
|
||||
asyncio.run(scenario())
|
||||
assert client._router._turn_states == {}
|
||||
|
||||
|
||||
def test_pending_join_preserves_result_after_original_handle_finishes():
|
||||
def test_pending_join_starts_at_request_while_original_handle_finishes():
|
||||
client = CodexClient()
|
||||
original = TurnHandle(client, "thread-1", "turn-1")
|
||||
with client._router.pending_turn("thread-1"):
|
||||
for event in turn_events(client):
|
||||
client._router.route_notification(event)
|
||||
events = turn_events(client)
|
||||
unknown = [client._coerce_notification(name, {"turnId": "turn-1"}) for name in ("old", "live")]
|
||||
client._router.route_notification(unknown[0])
|
||||
assert original._subscription.next() is unknown[0]
|
||||
|
||||
with client._router.pending_turn("thread-1") as cursors:
|
||||
client._router.route_notification(unknown[1])
|
||||
assert original._subscription.next() is unknown[1]
|
||||
client._router.route_notification(events[0])
|
||||
client._router.route_notification(events[1])
|
||||
manual = TurnHandle(client, "thread-1", "turn-1")
|
||||
client._router.route_notification(events[2])
|
||||
result = original.run()
|
||||
client._router.prepare_turn("turn-1", "thread-1")
|
||||
joined = TurnHandle(client, "thread-1", "turn-1")
|
||||
assert joined.run() == result
|
||||
assert manual.run().usage is None
|
||||
subscription = client._router.prepare_turn("turn-1", "thread-1", cursors, for_handle=True)
|
||||
joined_handle = TurnHandle(client, "thread-1", "turn-1", _subscription=subscription)
|
||||
assert joined_handle._subscription.next() is unknown[1]
|
||||
joined = joined_handle.run()
|
||||
assert result.final_response == joined.final_response == "done"
|
||||
assert joined.usage.last.total_tokens == 5
|
||||
assert client._router._turn_states == {}
|
||||
|
||||
|
||||
@@ -220,7 +257,7 @@ def test_cancelled_turn_start_releases_result_after_response(monkeypatch):
|
||||
entered = threading.Event()
|
||||
respond = threading.Event()
|
||||
released = threading.Event()
|
||||
subscribe = client._sync._subscribe_turn_notifications
|
||||
subscribe = client._sync._router.prepare_turn
|
||||
|
||||
def request_raw(method, params):
|
||||
entered.set()
|
||||
@@ -229,8 +266,8 @@ def test_cancelled_turn_start_releases_result_after_response(monkeypatch):
|
||||
client._sync._router.route_notification(event)
|
||||
return {"turn": {"id": "turn-1", "items": [], "status": "completed"}}
|
||||
|
||||
def observe_release(turn_id):
|
||||
subscription = subscribe(turn_id)
|
||||
def observe_release(*args, **kwargs):
|
||||
subscription = subscribe(*args, **kwargs)
|
||||
close = subscription.close
|
||||
|
||||
def close_and_signal():
|
||||
@@ -241,7 +278,7 @@ def test_cancelled_turn_start_releases_result_after_response(monkeypatch):
|
||||
return subscription
|
||||
|
||||
monkeypatch.setattr(client._sync, "_request_raw", request_raw)
|
||||
monkeypatch.setattr(client._sync, "_subscribe_turn_notifications", observe_release)
|
||||
monkeypatch.setattr(client._sync._router, "prepare_turn", observe_release)
|
||||
task = asyncio.create_task(client.turn_start("thread-1", "hello"))
|
||||
try:
|
||||
assert await asyncio.to_thread(entered.wait, 2)
|
||||
@@ -270,7 +307,7 @@ def test_cancelled_queued_start_does_not_send_a_request(monkeypatch):
|
||||
occupied = executor.submit(occupy_worker)
|
||||
assert worker_started.wait(timeout=5)
|
||||
client = AsyncCodexClient()
|
||||
monkeypatch.setattr(client._sync, "turn_start", lambda *args: request_sent.set())
|
||||
monkeypatch.setattr(client._sync, "_start_turn", lambda *args, **kwargs: request_sent.set())
|
||||
|
||||
async def scenario():
|
||||
task = asyncio.create_task(client.turn_start("thread-1", "hello"))
|
||||
@@ -298,6 +335,9 @@ def test_low_level_start_keeps_implicit_registration_and_explicit_unregister(mon
|
||||
|
||||
monkeypatch.setattr(client, "_request_raw", request_raw)
|
||||
started = client.turn_start("thread-1", "hello")
|
||||
registered = client._router._turn_notifications[started.turn.id]
|
||||
assert client.turn_start("thread-1", "again").turn.id == started.turn.id
|
||||
assert client._router._turn_notifications[started.turn.id] is registered
|
||||
assert client.next_turn_notification(started.turn.id) == turn_events(client)[0]
|
||||
client.unregister_turn_notifications(started.turn.id)
|
||||
with pytest.raises(RuntimeError, match="not registered"):
|
||||
|
||||
Reference in New Issue
Block a user