diff --git a/sdk/python/docs/api-reference.md b/sdk/python/docs/api-reference.md index 6eeea25050..5d9c160dff 100644 --- a/sdk/python/docs/api-reference.md +++ b/sdk/python/docs/api-reference.md @@ -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` diff --git a/sdk/python/scripts/update_sdk_artifacts.py b/sdk/python/scripts/update_sdk_artifacts.py index 4530a6d0bf..0c97d5156e 100755 --- a/sdk/python/scripts/update_sdk_artifacts.py +++ b/sdk/python/scripts/update_sdk_artifacts.py @@ -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) diff --git a/sdk/python/src/openai_codex/_message_router.py b/sdk/python/src/openai_codex/_message_router.py index 2d69099d4b..8ab5ff1339 100644 --- a/sdk/python/src/openai_codex/_message_router.py +++ b/sdk/python/src/openai_codex/_message_router.py @@ -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() diff --git a/sdk/python/src/openai_codex/api.py b/sdk/python/src/openai_codex/api.py index d29c8dad7f..d226744cf8 100644 --- a/sdk/python/src/openai_codex/api.py +++ b/sdk/python/src/openai_codex/api.py @@ -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) diff --git a/sdk/python/src/openai_codex/async_client.py b/sdk/python/src/openai_codex/async_client.py index 889970e038..a4c80e65c5 100644 --- a/sdk/python/src/openai_codex/async_client.py +++ b/sdk/python/src/openai_codex/async_client.py @@ -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 diff --git a/sdk/python/src/openai_codex/client.py b/sdk/python/src/openai_codex/client.py index 265bb5e9a9..5b38538c54 100644 --- a/sdk/python/src/openai_codex/client.py +++ b/sdk/python/src/openai_codex/client.py @@ -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]: diff --git a/sdk/python/tests/test_app_server_inputs.py b/sdk/python/tests/test_app_server_inputs.py index f8649a4241..26e0bfa659 100644 --- a/sdk/python/tests/test_app_server_inputs.py +++ b/sdk/python/tests/test_app_server_inputs.py @@ -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 == {} diff --git a/sdk/python/tests/test_client_rpc_methods.py b/sdk/python/tests/test_client_rpc_methods.py index f46eacad47..8b7485f0dc 100644 --- a/sdk/python/tests/test_client_rpc_methods.py +++ b/sdk/python/tests/test_client_rpc_methods.py @@ -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) diff --git a/sdk/python/tests/test_public_api_runtime_behavior.py b/sdk/python/tests/test_public_api_runtime_behavior.py index 780cff3fab..fc98f88311 100644 --- a/sdk/python/tests/test_public_api_runtime_behavior.py +++ b/sdk/python/tests/test_public_api_runtime_behavior.py @@ -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()) diff --git a/sdk/python/tests/test_turn_subscriptions.py b/sdk/python/tests/test_turn_subscriptions.py index 30937d1396..fea67c9cb0 100644 --- a/sdk/python/tests/test_turn_subscriptions.py +++ b/sdk/python/tests/test_turn_subscriptions.py @@ -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"):