diff --git a/sdk/python/README.md b/sdk/python/README.md index 5fbb145c70..fccc92b011 100644 --- a/sdk/python/README.md +++ b/sdk/python/README.md @@ -67,6 +67,7 @@ Use Python's standard `help(openai_codex)`, `help(Codex)`, or - [Getting started](https://github.com/openai/codex/blob/main/sdk/python/docs/getting-started.md) - [API reference](https://github.com/openai/codex/blob/main/sdk/python/docs/api-reference.md) +- [Untrusted external messages](https://github.com/openai/codex/blob/main/sdk/python/docs/api-reference.md#externalmessage) - [FAQ](https://github.com/openai/codex/blob/main/sdk/python/docs/faq.md) - [Examples](https://github.com/openai/codex/blob/main/sdk/python/examples/README.md) diff --git a/sdk/python/docs/api-reference.md b/sdk/python/docs/api-reference.md index 6442daaf73..6eeea25050 100644 --- a/sdk/python/docs/api-reference.md +++ b/sdk/python/docs/api-reference.md @@ -31,6 +31,7 @@ from openai_codex import ( LocalImageInput, SkillInput, MentionInput, + ExternalMessage, ) from openai_codex.types import ( Account, @@ -201,7 +202,7 @@ These options have the same behavior on sync and async `run(...)` and `turn(...) | `turn_service_tier: str | None = None` | Overrides the tier for a newly started turn only. `None` inherits the thread setting; `"default"` selects standard speed. Does not change the thread default and is ignored when input joins an active turn. | | `source: str | None = None` | Labels the caller that initiated a new turn, such as `"review_ui"`. This is metadata; it does not schedule work or grant authority. Ignored when input joins an active turn. | -`turn_service_tier`, `source`, and explicit `include_turns` +`ExternalMessage`, `turn_service_tier`, `source`, and explicit `include_turns` on resume/fork require Codex CLI 0.151.0 or newer. The SDK raises `CodexError` before sending these options to an older runtime, which would otherwise ignore them. Published SDK releases install a matching runtime automatically; when @@ -267,7 +268,7 @@ Behavior notes: InputItem = TextInput | ImageInput | LocalImageInput | SkillInput | MentionInput Input = list[InputItem] | InputItem -RunInput = Input | str +RunInput = Input | str | ExternalMessage ``` Use `ImageInput` with a base64-encoded `data:image/...` URL. HTTP and HTTPS image URLs are @@ -276,6 +277,58 @@ deprecated; download remote images and pass their local paths with `LocalImageIn Use a plain `str` as shorthand for `TextInput(...)` anywhere a turn input is accepted: `thread.run("...")`, `thread.turn("...")`, and `turn.steer("...")`. +### ExternalMessage + +`ExternalMessage` supplies **untrusted content** from another agent, tool, or +application. Content reaches the model with tool-level authority, below user +and developer instructions. It does not establish user authorization or +approval. Keep the thread's sandbox and approval policies appropriate for the +work the user has authorized. + +```python +from openai_codex import ExternalMessage + +message = ExternalMessage( + tool_name="notifications", + namespace="slack", + content="Deployment notification: the staging checks failed.", +) +result = thread.run(message) +``` + +| Field | Meaning | +| --- | --- | +| `tool_name: str` | Required, nonempty name of the tool or application delivering the message. | +| `content` | Required text, or a sequence of structured content dictionaries or generated `FunctionCallOutputContentItem` models. Structured image content requires inline data URLs. | +| `namespace: str | None = None` | Optional namespace for the tool name. | + +Pass one `ExternalMessage` as the complete input to `run(...)` or `turn(...)`. +It starts a turn when the thread is idle or joins an active regular turn. It +appears in saved history and item notifications as a `functionCallOutput` +item, retaining tool authority. No preceding tool call or call ID is required. +Tool names and namespaces identify the source; they are not proof of its +identity or permission to act. + +When a message joins an active turn, both handles can stream or collect the +result independently. A joining handle receives previously completed items and +the latest usage, followed by live notifications. Consumed transient events such +as token deltas are discarded. Both handles collect the complete result, and +closing one stream leaves the other active. + +The async calls use the same object: + +```python +result = await async_thread.run(message) +``` + +Use `await async_thread.turn(message)` to collect a handle for streaming and +interruption. An `ExternalMessage` cannot be mixed into a user-input list. +`TurnHandle.steer(...)` accepts user input; deliver an external message to an +active turn through `thread.turn(message)`. + +See the [external message examples](../examples/16_external_message) for a user +request followed by an external notification. + ## Public Types The SDK wrappers return and accept public Codex protocol models wherever possible: diff --git a/sdk/python/docs/faq.md b/sdk/python/docs/faq.md index 7886e335e8..36a2c519b8 100644 --- a/sdk/python/docs/faq.md +++ b/sdk/python/docs/faq.md @@ -37,6 +37,36 @@ Choose `run()` for most apps. Choose `stream()` for progress UIs, custom timeout If your app is not already async, stay with `Codex`. +## How do I pass untrusted external content? + +Use `ExternalMessage` for messages from other agents, tools, or applications: + +```python +from openai_codex import ExternalMessage + +result = thread.run(ExternalMessage( + tool_name="notifications", + namespace="slack", + content="Deployment notification: the staging checks failed.", +)) +``` + +The content has tool-level authority, below user and developer instructions. +It does not authorize actions or approve requests. Establish the user's task +separately and keep the thread's sandbox and approval policies in place. +Plain strings and `TextInput` represent user input. + +An external message starts a turn or joins an active regular turn and is +preserved in history. Pass it as the entire input to `thread.run(...)` or +`thread.turn(...)`; the async methods accept the same object. See the +[API reference](api-reference.md#externalmessage) and +[runnable example](../examples/16_external_message). + +External messages and the new `include_turns`, `turn_service_tier`, and `source` +options require CLI 0.151.0 or newer. If a custom executable is too old, the SDK +raises `CodexError` before sending the request. Upgrade that executable or use +the runtime installed with a matching SDK release. + ## Does `include_turns=False` remove the conversation's context? No. On `thread_resume(...)` and `thread_fork(...)`, it only skips loading turn diff --git a/sdk/python/docs/getting-started.md b/sdk/python/docs/getting-started.md index 9ea561a6d5..c1de2e8b62 100644 --- a/sdk/python/docs/getting-started.md +++ b/sdk/python/docs/getting-started.md @@ -70,6 +70,11 @@ with Codex() as codex: Use `Thread.turn(...)` when you need a `TurnHandle` for streaming, steering, or interrupting an active turn. +For **untrusted content** from another agent, tool, or application, pass an +[`ExternalMessage`](api-reference.md#externalmessage). It retains tool-level +authority and does not establish user authorization or approval. Plain strings +and `TextInput` represent user input. + ## 4. Choose Sandbox Access Use one enum for the initial thread and later turn overrides: diff --git a/sdk/python/examples/16_external_message/async.py b/sdk/python/examples/16_external_message/async.py new file mode 100644 index 0000000000..afc8625b3b --- /dev/null +++ b/sdk/python/examples/16_external_message/async.py @@ -0,0 +1,40 @@ +"""Process an untrusted notification within a task authorized by the user.""" + +import asyncio +import sys +from pathlib import Path + +_EXAMPLES_ROOT = Path(__file__).resolve().parents[1] +if str(_EXAMPLES_ROOT) not in sys.path: + sys.path.insert(0, str(_EXAMPLES_ROOT)) + +from _bootstrap import ensure_local_sdk_src, runtime_config + +ensure_local_sdk_src() + +from openai_codex import AsyncCodex, ExternalMessage, Sandbox + + +async def main() -> None: + async with AsyncCodex(config=runtime_config()) as codex: + thread = await codex.thread_start(sandbox=Sandbox.read_only) + await thread.run( + "When deployment notifications arrive, summarize their status and suggest " + "what I should check. Do not change files or deploy anything." + ) + + # External content has tool authority; it does not supply user permission. + result = await thread.run( + ExternalMessage( + tool_name="notifications", + namespace="slack", + content="Staging deployment failed: the health check returned HTTP 503.", + ), + source="slack_notification", + ) + print("status:", result.status) + print("text:", result.final_response) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/sdk/python/examples/16_external_message/sync.py b/sdk/python/examples/16_external_message/sync.py new file mode 100644 index 0000000000..260dad31e5 --- /dev/null +++ b/sdk/python/examples/16_external_message/sync.py @@ -0,0 +1,33 @@ +"""Process an untrusted notification within a task authorized by the user.""" + +import sys +from pathlib import Path + +_EXAMPLES_ROOT = Path(__file__).resolve().parents[1] +if str(_EXAMPLES_ROOT) not in sys.path: + sys.path.insert(0, str(_EXAMPLES_ROOT)) + +from _bootstrap import ensure_local_sdk_src, runtime_config + +ensure_local_sdk_src() + +from openai_codex import Codex, ExternalMessage, Sandbox + +with Codex(config=runtime_config()) as codex: + thread = codex.thread_start(sandbox=Sandbox.read_only) + thread.run( + "When deployment notifications arrive, summarize their status and suggest " + "what I should check. Do not change files or deploy anything." + ) + + # External content has tool authority; it does not supply user permission. + result = thread.run( + ExternalMessage( + tool_name="notifications", + namespace="slack", + content="Staging deployment failed: the health check returned HTTP 503.", + ), + source="slack_notification", + ) + print("status:", result.status) + print("text:", result.final_response) diff --git a/sdk/python/examples/README.md b/sdk/python/examples/README.md index 79be5890ed..137f74277d 100644 --- a/sdk/python/examples/README.md +++ b/sdk/python/examples/README.md @@ -11,6 +11,10 @@ and `openai_codex.types`. Examples use plain strings for text-only turns and typed input objects for multimodal or structured input lists. +Use `ExternalMessage` for untrusted content from another agent, tool, or +application. It retains tool-level authority and does not grant user +authorization or approval; example 16 establishes the user's task first. + ## Prerequisites - Python `>=3.10` @@ -89,3 +93,5 @@ python examples/01_quickstart_constructor/async.py - separate `steer()` and `interrupt()` demos with concise summaries - `15_login_and_account/` - browser-login handle lifecycle, cancellation, and account inspection +- `16_external_message/` + - process an untrusted external notification within a user-authorized task diff --git a/sdk/python/scripts/update_sdk_artifacts.py b/sdk/python/scripts/update_sdk_artifacts.py index 2b26e8e900..4530a6d0bf 100755 --- a/sdk/python/scripts/update_sdk_artifacts.py +++ b/sdk/python/scripts/update_sdk_artifacts.py @@ -1259,7 +1259,8 @@ def _render_thread_block(turn_fields: list[PublicFieldSpec], *, is_async: bool = " ) -> TurnResult:", ' """Run a complete turn and collect its final result.', "", - " Accepts the same input and options as turn().", + " Accepts the same input and options as turn(), including ExternalMessage", + " for untrusted external content with tool-level authority.", ' """', f" turn = {await_prefix}self.turn(", " input,", @@ -1275,18 +1276,21 @@ def _render_thread_block(turn_fields: list[PublicFieldSpec], *, is_async: bool = *_approval_mode_override_signature_lines(), *_kw_signature_lines(turn_fields), f" ) -> {handle_type}:", - ' """Start a turn and return a handle for streaming or control.', + ' """Start a turn or join an active regular turn and return its handle.', "", + " ExternalMessage supplies untrusted content with tool-level authority;", + " it does not establish user authorization or approval.", " turn_service_tier applies only to this new turn; service_tier updates", " the thread default. source labels what initiated a new turn and grants", " no authority. Both turn_service_tier and source are ignored when joining.", ' """', - " wire_input = _to_wire_input(_normalize_run_input(input))", + " wire_input, tool_output = _to_wire_turn_input(input)", *([" await self._codex._ensure_initialized()"] if is_async else []), _approval_mode_assignment_line("_approval_mode_override_settings"), " params = TurnStartParams(", " thread_id=self.id,", " input=wire_input,", + " tool_output=tool_output,", *_approval_mode_model_arg_lines(), *_model_arg_lines(turn_fields), " )", diff --git a/sdk/python/src/openai_codex/__init__.py b/sdk/python/src/openai_codex/__init__.py index 2829958dc1..d6784855aa 100644 --- a/sdk/python/src/openai_codex/__init__.py +++ b/sdk/python/src/openai_codex/__init__.py @@ -23,6 +23,7 @@ from .api import ( ChatgptLoginHandle, Codex, DeviceCodeLoginHandle, + ExternalMessage, ImageInput, Input, InputItem, @@ -72,6 +73,7 @@ __all__ = [ "Input", "InputItem", "RunInput", + "ExternalMessage", "TextInput", "ImageInput", "LocalImageInput", diff --git a/sdk/python/src/openai_codex/_inputs.py b/sdk/python/src/openai_codex/_inputs.py index 0a1fbe1100..d047fc3157 100644 --- a/sdk/python/src/openai_codex/_inputs.py +++ b/sdk/python/src/openai_codex/_inputs.py @@ -1,7 +1,9 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass +from .generated.v2_all import FunctionCallOutputContentItem, TurnToolOutput from .models import JsonObject @@ -42,9 +44,27 @@ class MentionInput: path: str +@dataclass(slots=True) +class ExternalMessage: + """Untrusted content supplied by another agent, tool, or application. + + Content has tool-level authority, below user and developer instructions. It + does not establish user authorization or approval. Pass this as the whole + input to ``thread.run()`` or ``thread.turn()`` to start a turn or join an + active regular turn. ``tool_name`` identifies the tool delivering it. + + ``content`` accepts text or Responses-compatible function-output content + items. Structured items can be dictionaries; no generated wrapper is needed. + """ + + tool_name: str + content: str | Sequence[JsonObject | FunctionCallOutputContentItem] + namespace: str | None = None + + InputItem = TextInput | ImageInput | LocalImageInput | SkillInput | MentionInput Input = list[InputItem] | InputItem -RunInput = Input | str +RunInput = Input | str | ExternalMessage def _to_wire_item(item: InputItem) -> JsonObject: @@ -67,7 +87,17 @@ def _to_wire_input(input: Input) -> list[JsonObject]: return [_to_wire_item(input)] -def _normalize_run_input(input: RunInput) -> Input: +def _normalize_run_input(input: Input | str) -> Input: if isinstance(input, str): return TextInput(input) return input + + +def _to_wire_turn_input(input: RunInput) -> tuple[list[JsonObject], TurnToolOutput | None]: + if isinstance(input, ExternalMessage): + if not isinstance(input.tool_name, str) or not input.tool_name.strip(): + raise ValueError("ExternalMessage.tool_name must be a nonempty string") + return [], TurnToolOutput.model_validate( + {"name": input.tool_name, "namespace": input.namespace, "output": input.content} + ) + return _to_wire_input(_normalize_run_input(input)), None diff --git a/sdk/python/src/openai_codex/_message_router.py b/sdk/python/src/openai_codex/_message_router.py index c979c8c8dd..2d69099d4b 100644 --- a/sdk/python/src/openai_codex/_message_router.py +++ b/sdk/python/src/openai_codex/_message_router.py @@ -2,35 +2,107 @@ 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, map_jsonrpc_error +from .errors import CodexError, TransportClosedError, map_jsonrpc_error from .generated.notification_registry import notification_turn_id -from .generated.v2_all import AccountLoginCompletedNotification +from .generated.v2_all import ( + AccountLoginCompletedNotification, + ItemCompletedNotification, + ThreadTokenUsageUpdatedNotification, +) 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) + 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.""" + + def __init__(self, router: MessageRouter, state: _TurnState) -> None: + self._router = router + self._state = state + self._cursor = state.first_event + 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 + ) + + def next(self) -> Notification: + with self._router._turn_condition: + while not self._replay and self._cursor == self._state.next_event and not self._closed: + 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) + if isinstance(item, BaseException): + raise item + return item + + def close(self) -> None: + with self._router._turn_condition: + self._closed = True + self._replay.clear() + 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 - and active turn stream its own queue. + its own queue and each turn consumer its own event cursor. """ def __init__(self) -> None: """Create empty response, turn, and global notification queues.""" - self._lock = threading.Lock() + # 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_notifications: dict[str, queue.Queue[NotificationQueueItem]] = {} - self._pending_turn_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._goal_operations: dict[str, _GoalOperationState] = {} self._global_notifications: queue.Queue[NotificationQueueItem] = queue.Queue() @@ -86,37 +158,119 @@ class MessageRouter: raise item return item - def register_turn(self, turn_id: str) -> None: - """Register a queue for a turn stream and replay early events.""" - - turn_queue: queue.Queue[NotificationQueueItem] = queue.Queue() + @contextmanager + def pending_turn(self, thread_id: str) -> Iterator[None]: + """Retain early completion while a turn/start response is in flight.""" with self._lock: - if turn_id in self._turn_notifications: - return - # A turn can emit events immediately after turn/start, before the - # caller receives the TurnHandle and starts streaming. - pending = self._pending_turn_notifications.pop(turn_id, deque()) - self._turn_notifications[turn_id] = turn_queue - for notification in pending: - turn_queue.put(notification) + self._pending_turn_requests[thread_id] = ( + self._pending_turn_requests.get(thread_id, 0) + 1 + ) + try: + yield + 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] + for state in list(self._turn_states.values()): + if state.thread_id == 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.""" + 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) + + def subscribe_turn(self, turn_id: str) -> _TurnSubscription: + """Attach a consumer with completed items, latest usage, and unread events.""" + with self._lock: + return self._subscribe_turn_locked(turn_id) + + @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: + 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) + ): + return + consumed = min(state.subscribers.values(), default=state.next_event) + while state.first_event < consumed: + event = state.events.pop(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 + + 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) def unregister_turn(self, turn_id: str) -> None: - """Stop routing future turn events to the stream queue.""" - + """Close only the low-level consumer, leaving other handles subscribed.""" with self._lock: - self._turn_notifications.pop(turn_id, None) + 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() def next_turn_notification(self, turn_id: str) -> Notification: - """Block until the next notification for a registered turn.""" - + """Block until the next event for the default low-level consumer.""" with self._lock: - turn_queue = self._turn_notifications.get(turn_id) - if turn_queue is None: - raise RuntimeError(f"turn {turn_id!r} is not registered for streaming") - item = turn_queue.get() - if isinstance(item, BaseException): - raise item - return item + 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] + if subscription is None: + subscription = self._subscribe_turn_locked(turn_id) + self._turn_notifications[turn_id] = subscription + return subscription.next() def register_goal(self, thread_id: str) -> _GoalOperationState: """Register one thread-scoped logical goal operation before it starts.""" @@ -204,15 +358,16 @@ class MessageRouter: self._global_notifications.put(notification) return - with self._lock: - turn_queue = self._turn_notifications.get(turn_id) - if turn_queue is None: - if notification.method == "turn/completed": - self._pending_turn_notifications.pop(turn_id, None) - return - self._pending_turn_notifications.setdefault(turn_id, deque()).append(notification) - return - turn_queue.put(notification) + 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 + self._prune_turn_events(state) + if notification.method == "turn/completed": + state.completed = True + self._discard_finished_turn(state) + self._turn_condition.notify_all() def fail_all(self, exc: BaseException) -> None: """Wake every blocked waiter when the reader thread exits.""" @@ -223,8 +378,13 @@ class MessageRouter: login_queues = list(self._login_notifications.values()) self._login_notifications.clear() self._pending_login_notifications.clear() - turn_queues = list(self._turn_notifications.values()) - self._pending_turn_notifications.clear() + 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._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 @@ -233,8 +393,6 @@ class MessageRouter: waiter.put(exc) for login_queue in login_queues: 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) diff --git a/sdk/python/src/openai_codex/api.py b/sdk/python/src/openai_codex/api.py index d68d2c8771..d29c8dad7f 100644 --- a/sdk/python/src/openai_codex/api.py +++ b/sdk/python/src/openai_codex/api.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import AsyncIterator, Iterator from ._approval_mode import ( @@ -11,6 +11,7 @@ from ._approval_mode import ( ) from ._initialize_metadata import validate_initialize_metadata from ._inputs import ( + ExternalMessage as ExternalMessage, ImageInput as ImageInput, Input as Input, InputItem as InputItem, @@ -21,6 +22,7 @@ from ._inputs import ( TextInput as TextInput, _normalize_run_input, _to_wire_input, + _to_wire_turn_input, ) from ._login import ( AsyncChatgptLoginHandle, @@ -32,6 +34,7 @@ from ._login import ( start_chatgpt_login, start_device_code_login, ) +from ._message_router import _TurnSubscription from ._run import ( TurnResult, _collect_async_turn_result, @@ -584,7 +587,8 @@ class Thread: ) -> TurnResult: """Run a complete turn and collect its final result. - Accepts the same input and options as turn(). + Accepts the same input and options as turn(), including ExternalMessage + for untrusted external content with tool-level authority. """ turn = self.turn( input, @@ -618,17 +622,20 @@ class Thread: summary: ReasoningSummary | None = None, turn_service_tier: str | None = None, ) -> TurnHandle: - """Start a turn and return a handle for streaming or control. + """Start a turn or join an active regular turn and return its handle. + ExternalMessage supplies untrusted content with tool-level authority; + it does not establish user authorization or approval. turn_service_tier applies only to this new turn; service_tier updates the thread default. source labels what initiated a new turn and grants no authority. Both turn_service_tier and source are ignored when joining. """ - wire_input = _to_wire_input(_normalize_run_input(input)) + wire_input, tool_output = _to_wire_turn_input(input) approval_policy, approvals_reviewer = _approval_mode_override_settings(approval_mode) params = TurnStartParams( thread_id=self.id, input=wire_input, + tool_output=tool_output, approval_policy=approval_policy, approvals_reviewer=approvals_reviewer, cwd=cwd, @@ -684,7 +691,8 @@ class AsyncThread: ) -> TurnResult: """Run a complete turn and collect its final result. - Accepts the same input and options as turn(). + Accepts the same input and options as turn(), including ExternalMessage + for untrusted external content with tool-level authority. """ turn = await self.turn( input, @@ -718,18 +726,21 @@ class AsyncThread: summary: ReasoningSummary | None = None, turn_service_tier: str | None = None, ) -> AsyncTurnHandle: - """Start a turn and return a handle for streaming or control. + """Start a turn or join an active regular turn and return its handle. + ExternalMessage supplies untrusted content with tool-level authority; + it does not establish user authorization or approval. turn_service_tier applies only to this new turn; service_tier updates the thread default. source labels what initiated a new turn and grants no authority. Both turn_service_tier and source are ignored when joining. """ - wire_input = _to_wire_input(_normalize_run_input(input)) + wire_input, tool_output = _to_wire_turn_input(input) await self._codex._ensure_initialized() approval_policy, approvals_reviewer = _approval_mode_override_settings(approval_mode) params = TurnStartParams( thread_id=self.id, input=wire_input, + tool_output=tool_output, approval_policy=approval_policy, approvals_reviewer=approvals_reviewer, cwd=cwd, @@ -769,9 +780,13 @@ class TurnHandle: _client: CodexClient thread_id: str id: str + _subscription: _TurnSubscription = field(init=False, repr=False, compare=False) - def steer(self, input: RunInput) -> TurnSteerResponse: - """Send additional input to this active turn.""" + def __post_init__(self) -> None: + self._subscription = self._client._subscribe_turn_notifications(self.id) + + def steer(self, input: Input | str) -> TurnSteerResponse: + """Send additional user input to this active turn.""" return self._client.turn_steer( self.thread_id, self.id, @@ -784,10 +799,9 @@ class TurnHandle: 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) + event = self._subscription.next() yield event if ( event.method == "turn/completed" @@ -796,7 +810,7 @@ class TurnHandle: ): break finally: - self._client.unregister_turn_notifications(self.id) + self._subscription.close() def run(self) -> TurnResult: """Consume the turn stream and return its completed result.""" @@ -814,9 +828,13 @@ class AsyncTurnHandle: _codex: AsyncCodex thread_id: str id: str + _subscription: _TurnSubscription = field(init=False, repr=False, compare=False) - async def steer(self, input: RunInput) -> TurnSteerResponse: - """Send additional input to this active turn.""" + def __post_init__(self) -> None: + self._subscription = self._codex._client._subscribe_turn_notifications(self.id) + + async def steer(self, input: Input | str) -> TurnSteerResponse: + """Send additional user input to this active turn.""" await self._codex._ensure_initialized() return await self._codex._client.turn_steer( self.thread_id, @@ -832,10 +850,9 @@ class AsyncTurnHandle: async 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) + event = await asyncio.to_thread(self._subscription.next) yield event if ( event.method == "turn/completed" @@ -844,7 +861,7 @@ class AsyncTurnHandle: ): break finally: - self._codex._client.unregister_turn_notifications(self.id) + self._subscription.close() async def run(self) -> TurnResult: """Consume the turn stream and return its completed result.""" diff --git a/sdk/python/src/openai_codex/async_client.py b/sdk/python/src/openai_codex/async_client.py index 07e7d0d405..889970e038 100644 --- a/sdk/python/src/openai_codex/async_client.py +++ b/sdk/python/src/openai_codex/async_client.py @@ -3,12 +3,14 @@ from __future__ import annotations import asyncio import threading from collections.abc import Iterator -from concurrent.futures import Future +from concurrent.futures import Future, ThreadPoolExecutor +from contextvars import copy_context from typing import AsyncIterator, Callable, ParamSpec, TypeVar from pydantic import BaseModel from ._goal import _GoalOperationState +from ._message_router import _TurnSubscription from .client import CodexClient, CodexConfig from .generated.v2_all import ( AccountLoginCompletedNotification, @@ -48,6 +50,9 @@ ModelT = TypeVar("ModelT", bound=BaseModel) ParamsT = ParamSpec("ParamsT") ReturnT = TypeVar("ReturnT") +# Bound workers while allowing cancellation cleanup to outlive an asyncio waiter. +_TURN_START_EXECUTOR = ThreadPoolExecutor(thread_name_prefix="codex-turn-start") + class AsyncCodexClient: """Async wrapper around CodexClient using thread offloading.""" @@ -97,6 +102,9 @@ class AsyncCodexClient: """Initialize the Codex session.""" return await self._call_sync(self._sync.initialize) + def _subscribe_turn_notifications(self, turn_id: str) -> _TurnSubscription: + return self._sync._subscribe_turn_notifications(turn_id) + def register_turn_notifications(self, turn_id: str) -> None: """Register a turn notification queue on the wrapped sync client.""" self._sync.register_turn_notifications(turn_id) @@ -287,8 +295,23 @@ class AsyncCodexClient: input_items: list[JsonObject] | JsonObject | str, params: V2TurnStartParams | JsonObject | None = None, ) -> TurnStartResponse: - """Start a turn using the wrapped sync client.""" - return await self._call_sync(self._sync.turn_start, thread_id, input_items, params) + """Start a turn, releasing an unclaimed result if the caller is cancelled.""" + operation = _TURN_START_EXECUTOR.submit( + copy_context().run, self._sync.turn_start, thread_id, input_items, params + ) + try: + return await asyncio.wrap_future(operation) + except asyncio.CancelledError: + + def discard_cancelled_result(completed: Future[TurnStartResponse]) -> None: + try: + started = completed.result() + except BaseException: + return + self._sync._subscribe_turn_notifications(started.turn.id).close() + + operation.add_done_callback(discard_cancelled_result) + raise async def turn_interrupt(self, thread_id: str, turn_id: str) -> TurnInterruptResponse: """Interrupt a turn using the wrapped sync client.""" diff --git a/sdk/python/src/openai_codex/client.py b/sdk/python/src/openai_codex/client.py index b8513a5caa..265bb5e9a9 100644 --- a/sdk/python/src/openai_codex/client.py +++ b/sdk/python/src/openai_codex/client.py @@ -15,7 +15,7 @@ from pydantic import BaseModel from ._goal import _GoalOperationState from ._initialize_metadata import _split_user_agent -from ._message_router import MessageRouter +from ._message_router import MessageRouter, _TurnSubscription from ._runtime_requirements import CheckoutCapabilities, require_runtime_version from ._version import __version__ as SDK_VERSION from .errors import CodexError, InvalidRequestError, TransportClosedError @@ -332,7 +332,7 @@ class CodexClient: response_model: type[ModelT], ) -> ModelT: runtime_fields = { - "turn/start": ("turnTrigger", "serviceTierForTurn"), + "turn/start": ("toolOutput", "turnTrigger", "serviceTierForTurn"), "thread/resume": ("excludeTurns",), "thread/fork": ("excludeTurns",), } @@ -407,6 +407,9 @@ class CodexClient: """Return the next routed notification for the requested login id.""" return self._router.next_login_notification(login_id) + def _subscribe_turn_notifications(self, turn_id: str) -> _TurnSubscription: + return self._router.subscribe_turn(turn_id) + def register_turn_notifications(self, turn_id: str) -> None: """Start routing notifications for one turn into its dedicated queue.""" self._router.register_turn(turn_id) @@ -662,9 +665,10 @@ class CodexClient: "threadId": thread_id, "input": self._normalize_input_items(input_items), } - started = self.request("turn/start", payload, response_model=TurnStartResponse) - self.register_turn_notifications(started.turn.id) - return started + with self._router.pending_turn(thread_id): + started = self.request("turn/start", payload, response_model=TurnStartResponse) + self._router.prepare_turn(started.turn.id, thread_id) + return started @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 29b8fc7726..f8649a4241 100644 --- a/sdk/python/tests/test_app_server_inputs.py +++ b/sdk/python/tests/test_app_server_inputs.py @@ -1,11 +1,179 @@ from __future__ import annotations +import asyncio import base64 +from concurrent.futures import ThreadPoolExecutor from app_server_harness import AppServerHarness -from app_server_helpers import TINY_PNG_BYTES +from app_server_helpers import TINY_PNG_BYTES, streaming_response -from openai_codex import Codex, ImageInput, LocalImageInput, SkillInput, TextInput +from openai_codex import ( + AsyncCodex, + Codex, + ExternalMessage, + ImageInput, + LocalImageInput, + SkillInput, + TextInput, +) + + +def _external_items(request) -> list[dict]: + """Select model-visible external content without generated item identifiers.""" + return [ + {key: value for key, value in item.items() if key != "id"} + for item in request.input() + if item.get("type") == "function_call_output" + ] + + +def test_external_message_preserves_tool_authority_through_resume(tmp_path) -> None: + content = "External update: deployment completed." + expected = { + "type": "function_call_output", + "name": "notifications", + "namespace": "slack", + "output": content, + } + with AppServerHarness(tmp_path) as harness: + harness.responses.enqueue_assistant_message("Update received", response_id="external") + harness.responses.enqueue_assistant_message("Still available", response_id="resumed") + with Codex(config=harness.app_server_config()) as codex: + thread = codex.thread_start() + result = thread.run( + ExternalMessage(tool_name="notifications", namespace="slack", content=content) + ) + external_item = next( + item for item in result.items if item.root.type == "functionCallOutput" + ) + with Codex(config=harness.app_server_config()) as codex: + resumed = codex.thread_resume(thread.id, include_turns=False) + history = resumed.read(include_turns=True) + assert external_item in history.thread.turns[0].items + resumed.run("Summarize the external update.") + requests = harness.responses.requests() + + assert result.final_response == "Update received" + assert [_external_items(request) for request in requests] == [[expected], [expected]] + assert [ + text + for request in requests + for role in ("user", "developer") + for text in request.message_input_texts(role) + if content in text + ] == [] + + +def test_external_message_joins_active_turn_with_tool_authority(tmp_path) -> None: + content = "External update while the agent is running." + with AppServerHarness(tmp_path) as harness: + harness.responses.enqueue_sse( + streaming_response("external-first", "msg-first", ["Working"]), + delay_between_events_s=0.2, + ) + harness.responses.enqueue_assistant_message( + "Update processed", response_id="external-second" + ) + with ThreadPoolExecutor(max_workers=2) as consumers: + with Codex(config=harness.app_server_config()) as codex: + thread = codex.thread_start() + original = thread.turn("Monitor deployment updates.") + original_result = consumers.submit(original.run) + 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 + assert codex._client._router._turn_states == {} + requests = harness.responses.requests() + + assert result.usage is not None + assert (joined.id, result.final_response) == (original.id, "Update processed") + assert [_external_items(request) for request in requests] == [ + [], + [ + { + "type": "function_call_output", + "name": "notifications", + "output": content, + } + ], + ] + assert [request.message_input_texts("user")[-1] for request in requests] == [ + "Monitor deployment updates.", + "Monitor deployment updates.", + ] + + +def test_async_external_message_reaches_model_with_tool_authority(tmp_path) -> None: + async def scenario() -> None: + with AppServerHarness(tmp_path) as harness: + harness.responses.enqueue_assistant_message( + "Async update received", response_id="external-async" + ) + async with AsyncCodex(config=harness.app_server_config()) as codex: + thread = await codex.thread_start() + result = await thread.run( + ExternalMessage(tool_name="notifications", content="External async update") + ) + request = harness.responses.single_request() + + assert result.final_response == "Async update received" + assert _external_items(request) == [ + { + "type": "function_call_output", + "name": "notifications", + "output": "External async update", + } + ] + assert "External async update" not in request.message_input_texts("user") + + asyncio.run(scenario()) + + +def test_async_external_message_allows_both_handles_to_consume(tmp_path) -> None: + async def scenario() -> None: + with AppServerHarness(tmp_path) as harness: + harness.responses.enqueue_sse( + streaming_response("async-first", "msg-first", ["Working"]), + delay_between_events_s=0.2, + ) + harness.responses.enqueue_assistant_message( + "Update processed", response_id="async-second" + ) + async with AsyncCodex(config=harness.app_server_config()) as codex: + thread = await codex.thread_start() + original = await thread.turn("Monitor deployment updates.") + original_result = asyncio.create_task(original.run()) + await asyncio.to_thread(harness.responses.wait_for_requests, 1) + joined = await thread.turn( + ExternalMessage(tool_name="notifications", content="Update") + ) + first, second = await asyncio.wait_for( + asyncio.gather(original_result, joined.run()), timeout=15 + ) + assert first == second + assert first.final_response == "Update processed" + assert first.usage is not None + assert codex._client._sync._router._turn_states == {} + + asyncio.run(scenario()) + + +def test_external_message_uses_core_tool_output_truncation(tmp_path) -> None: + content = "External observation. " * 500 + with AppServerHarness(tmp_path) as harness: + harness.responses.enqueue_assistant_message( + "Context received", response_id="truncated-external" + ) + with Codex(config=harness.app_server_config()) as codex: + thread = codex.thread_start(config={"tool_output_token_limit": 32}) + thread.run(ExternalMessage(tool_name="notifications", content=content)) + request = harness.responses.single_request() + + [item] = _external_items(request) + assert (item["name"], item["type"]) == ("notifications", "function_call_output") + assert len(item["output"]) < len(content) + assert "truncated" in item["output"].lower() def test_data_url_image_input_reaches_responses_api( diff --git a/sdk/python/tests/test_client_rpc_methods.py b/sdk/python/tests/test_client_rpc_methods.py index 3f9035ee30..f46eacad47 100644 --- a/sdk/python/tests/test_client_rpc_methods.py +++ b/sdk/python/tests/test_client_rpc_methods.py @@ -84,6 +84,7 @@ def _initialized_client( @pytest.mark.parametrize( ("method", "params"), [ + ("turn/start", {"input": [], "toolOutput": {"name": "delegate", "output": "Investigate"}}), ("turn/start", {"input": [], "turnTrigger": "automation"}), ("turn/start", {"input": [], "serviceTierForTurn": "default"}), ("thread/resume", {"threadId": "thread-1", "excludeTurns": False}), @@ -117,7 +118,7 @@ def test_new_options_accept_supported_runtime_metadata( monkeypatch: pytest.MonkeyPatch, metadata: JsonObject ) -> None: client, requests = _initialized_client(monkeypatch, metadata) - params = {"input": [], "turnTrigger": "automation"} + params = {"input": [], "toolOutput": {"name": "delegate", "output": "Investigate"}} client.request("turn/start", params, response_model=InitializeResponse) @@ -665,7 +666,7 @@ def test_turn_notification_router_clears_unregistered_turn_when_completed() -> N ) ) - assert client._router._pending_turn_notifications == {} + assert client._router._turn_states == {} def test_turn_notification_router_routes_unknown_turn_notifications() -> None: diff --git a/sdk/python/tests/test_public_api_runtime_behavior.py b/sdk/python/tests/test_public_api_runtime_behavior.py index a948a86178..780cff3fab 100644 --- a/sdk/python/tests/test_public_api_runtime_behavior.py +++ b/sdk/python/tests/test_public_api_runtime_behavior.py @@ -13,7 +13,9 @@ from openai_codex.api import ( ApprovalMode, AsyncCodex, Codex, + ExternalMessage, Sandbox, + TextInput, ) from openai_codex.client import _params_dict from openai_codex.generated.v2_all import TurnCompletedNotification, TurnStartParams @@ -171,8 +173,11 @@ def test_include_turns_preserves_omission_and_inverts_explicit_values( @pytest.mark.parametrize("api_type", [Codex, AsyncCodex]) @pytest.mark.parametrize("method", ["run", "turn"]) -def test_turn_inputs_and_options_reach_the_client(api_type, method) -> None: - """Turn options reach the client through both sync and async entry points.""" +@pytest.mark.parametrize( + "content", [None, "External update", [{"type": "input_text", "text": "External update"}]] +) +def test_turn_inputs_and_options_reach_the_client(api_type, method, content) -> None: + """User and external inputs preserve distinct wire representations for every entry point.""" async def scenario() -> None: async_api = api_type is AsyncCodex @@ -188,9 +193,9 @@ def test_turn_inputs_and_options_reach_the_client(api_type, method) -> None: ) client = SimpleNamespace( turn_start=rpc(return_value=SimpleNamespace(turn=SimpleNamespace(id="turn-1"))), - register_turn_notifications=Mock(), - unregister_turn_notifications=Mock(), - next_turn_notification=rpc(return_value=completed), + _subscribe_turn_notifications=Mock( + return_value=SimpleNamespace(next=Mock(return_value=completed), close=Mock()) + ), ) codex = api_type.__new__(api_type) codex._client = client @@ -200,7 +205,11 @@ def test_turn_inputs_and_options_reach_the_client(api_type, method) -> None: if async_api else public_api_module.Thread(client, "thread-1") ) - input = "Continue." + input = ( + "Continue." + if content is None + else ExternalMessage(tool_name="notifications", namespace="slack", content=content) + ) turn = getattr(thread, method)( input, service_tier="priority", @@ -210,18 +219,70 @@ def test_turn_inputs_and_options_reach_the_client(api_type, method) -> None: if async_api: turn = await turn assert turn.id == "turn-1" - expected_input = [{"type": "text", "text": "Continue.", "text_elements": []}] + expected_input = ( + [{"type": "text", "text": "Continue.", "text_elements": []}] if content is None else [] + ) + expected_tool_output = ( + {} + if content is None + else { + "toolOutput": {"name": "notifications", "namespace": "slack", "output": content}, + } + ) assert _params_dict(client.turn_start.call_args.kwargs["params"]) == { "threadId": "thread-1", "input": expected_input, "serviceTier": "priority", "serviceTierForTurn": "default", "turnTrigger": "automation", + **expected_tool_output, } asyncio.run(scenario()) +@pytest.mark.parametrize("api_type", [Codex, AsyncCodex]) +def test_external_messages_cannot_be_mixed_with_user_input_or_sent_as_user_steering( + api_type, +) -> None: + async def scenario() -> None: + 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() + ) + codex = api_type.__new__(api_type) + codex._client = client + codex._initialized = True + thread = ( + public_api_module.AsyncThread(codex, "thread-1") + if async_api + else public_api_module.Thread(client, "thread-1") + ) + handle = ( + public_api_module.AsyncTurnHandle(codex, "thread-1", "turn-1") + if async_api + else public_api_module.TurnHandle(client, "thread-1", "turn-1") + ) + external = ExternalMessage(tool_name="notifications", content="Untrusted update") + for operation, input in ( + (thread.turn, [TextInput("User request"), external]), + (handle.steer, external), + ): + with pytest.raises(TypeError): + result = operation(input) + if async_api: + await result + with pytest.raises(ValueError, match="tool_name"): + result = thread.turn(ExternalMessage(tool_name=" ", content="Untrusted update")) + if async_api: + await result + client.turn_start.assert_not_called() + client.turn_steer.assert_not_called() + + asyncio.run(scenario()) + + def _approval_mode_turn_params(approval_mode: ApprovalMode) -> TurnStartParams: """Build real generated turn params from one public approval mode.""" approval_policy, approvals_reviewer = public_api_module._approval_mode_settings(approval_mode) diff --git a/sdk/python/tests/test_public_api_signatures.py b/sdk/python/tests/test_public_api_signatures.py index f9349196da..33e6edf55e 100644 --- a/sdk/python/tests/test_public_api_signatures.py +++ b/sdk/python/tests/test_public_api_signatures.py @@ -14,6 +14,7 @@ from openai_codex import ( AsyncTurnHandle, Codex, CodexConfig, + ExternalMessage, Sandbox, Thread, TurnHandle, @@ -46,6 +47,7 @@ EXPECTED_ROOT_EXPORTS = [ "Input", "InputItem", "RunInput", + "ExternalMessage", "TextInput", "ImageInput", "LocalImageInput", @@ -187,7 +189,7 @@ def test_turn_input_methods_accept_string_shortcut() -> None: assert { fn: inspect.signature(fn).parameters["input"].annotation for fn in (TurnHandle.steer, AsyncTurnHandle.steer) - } == dict.fromkeys((TurnHandle.steer, AsyncTurnHandle.steer), "RunInput") + } == dict.fromkeys((TurnHandle.steer, AsyncTurnHandle.steer), "Input | str") def test_root_exports_approval_mode() -> None: @@ -228,6 +230,7 @@ def test_curated_public_api_has_builtin_help_documentation() -> None: "TurnHandle": TurnHandle, "AsyncTurnHandle": AsyncTurnHandle, "TurnResult": TurnResult, + "ExternalMessage": ExternalMessage, "Sandbox": Sandbox, "thread_start": Codex.thread_start, "thread_resume": Codex.thread_resume, diff --git a/sdk/python/tests/test_turn_subscriptions.py b/sdk/python/tests/test_turn_subscriptions.py new file mode 100644 index 0000000000..30937d1396 --- /dev/null +++ b/sdk/python/tests/test_turn_subscriptions.py @@ -0,0 +1,305 @@ +import asyncio +import gc +import threading +from concurrent.futures import ThreadPoolExecutor +from itertools import chain + +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.async_client import AsyncCodexClient +from openai_codex.client import CodexClient +from openai_codex.errors import TransportClosedError + + +def turn_events(client, *, status="completed"): + scope = {"threadId": "thread-1", "turnId": "turn-1"} + usage = { + "inputTokens": 2, + "cachedInputTokens": 0, + "outputTokens": 3, + "reasoningOutputTokens": 0, + "totalTokens": 5, + } + return [ + client._coerce_notification( + "item/completed", + { + **scope, + "completedAtMs": 1, + "item": { + "id": "message", + "type": "agentMessage", + "text": "done", + "phase": "final_answer", + }, + }, + ), + client._coerce_notification( + "thread/tokenUsage/updated", {**scope, "tokenUsage": {"last": usage, "total": usage}} + ), + client._coerce_notification( + "turn/completed", + { + "threadId": "thread-1", + "turn": { + "id": "turn-1", + "items": [], + "status": status, + "error": {"message": "model failed"} if status == "failed" else None, + }, + }, + ), + ] + + +def test_late_join_replays_items_already_consumed_by_original_handle(): + 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) + joined = TurnHandle(client, "thread-1", "turn-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 + assert first_result.final_response == "done" + assert first_result.usage.last.total_tokens == 5 + assert client._router._turn_states == {} + + +def test_consumed_deltas_are_released_while_turn_is_active(): + client = CodexClient() + subscription = client._subscribe_turn_notifications("turn-1") + state = client._router._turn_states["turn-1"] + for index in range(1000): + event = client._coerce_notification( + "item/agentMessage/delta", + { + "threadId": "thread-1", + "turnId": "turn-1", + "itemId": "message", + "delta": str(index), + }, + ) + client._router.route_notification(event) + assert subscription.next() == event + assert state.events == {} + assert state.completed_items == {} + assert not state.completed + subscription.close() + + +def test_slow_subscriber_keeps_unread_deltas_until_it_consumes_them(): + client = CodexClient() + fast = client._subscribe_turn_notifications("turn-1") + slow = client._subscribe_turn_notifications("turn-1") + event = client._coerce_notification( + "item/agentMessage/delta", + {"threadId": "thread-1", "turnId": "turn-1", "itemId": "message", "delta": "hello"}, + ) + client._router.route_notification(event) + assert fast.next() == event + assert slow.next() == event + assert client._router._turn_states["turn-1"].events == {} + fast.close() + slow.close() + + +def test_completion_before_turn_start_response_is_replayed(monkeypatch): + client = CodexClient() + + def request_raw(method, params): + assert method == "turn/start" + for event in turn_events(client): + client._router.route_notification(event) + 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" + assert client._router._turn_states == {} + + +def test_pending_join_preserves_result_after_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) + result = original.run() + client._router.prepare_turn("turn-1", "thread-1") + joined = TurnHandle(client, "thread-1", "turn-1") + assert joined.run() == result + assert client._router._turn_states == {} + + +def test_failed_start_releases_early_completed_state(): + client = CodexClient() + with pytest.raises(ValueError, match="request failed"): + with client._router.pending_turn("thread-1"): + for event in turn_events(client): + client._router.route_notification(event) + raise ValueError("request failed") + assert client._router._turn_states == {} + assert client._router._pending_turn_requests == {} + + +def test_closing_one_stream_leaves_other_subscriber_intact(): + client = CodexClient() + original = TurnHandle(client, "thread-1", "turn-1") + joined = TurnHandle(client, "thread-1", "turn-1") + events = turn_events(client) + client._router.route_notification(events[0]) + stream = original.stream() + next(stream) + stream.close() + for event in events[1:]: + client._router.route_notification(event) + assert joined.run().final_response == "done" + assert client._router._turn_states == {} + + +@pytest.mark.parametrize("failure", ["transport", "model"]) +def test_both_handles_observe_failure_and_release_state(failure): + client = CodexClient() + handles = [TurnHandle(client, "thread-1", "turn-1") for _ in range(2)] + if failure == "transport": + client._router.fail_all(TransportClosedError("transport failed")) + else: + for event in turn_events(client, status="failed"): + client._router.route_notification(event) + for handle in handles: + with pytest.raises((TransportClosedError, RuntimeError), match=f"{failure} failed"): + handle.run() + assert client._router._turn_states == {} + + +def test_abandoned_handle_releases_completed_history(): + client = CodexClient() + handle = TurnHandle(client, "thread-1", "turn-1") + for event in turn_events(client): + client._router.route_notification(event) + with client._router._lock: + del handle + gc.collect() + assert client._router._turn_states == {} + + +def test_cancelled_async_consumer_leaves_other_handle_intact(): + async def scenario(): + codex = AsyncCodex() + codex._initialized = True + client = codex._client._sync + original = AsyncTurnHandle(codex, "thread-1", "turn-1") + joined = AsyncTurnHandle(codex, "thread-1", "turn-1") + task = asyncio.create_task(original.run()) + await asyncio.sleep(0) # Let the stream enter its wait before cancelling it. + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + for event in turn_events(client): + client._router.route_notification(event) + assert (await asyncio.wait_for(joined.run(), timeout=2)).final_response == "done" + assert client._router._turn_states == {} + + asyncio.run(scenario()) + + +def test_cancelled_turn_start_releases_result_after_response(monkeypatch): + async def scenario(): + client = AsyncCodexClient() + entered = threading.Event() + respond = threading.Event() + released = threading.Event() + subscribe = client._sync._subscribe_turn_notifications + + def request_raw(method, params): + entered.set() + assert respond.wait(timeout=2) + for event in turn_events(client._sync): + client._sync._router.route_notification(event) + return {"turn": {"id": "turn-1", "items": [], "status": "completed"}} + + def observe_release(turn_id): + subscription = subscribe(turn_id) + close = subscription.close + + def close_and_signal(): + close() + released.set() + + subscription.close = close_and_signal + return subscription + + monkeypatch.setattr(client._sync, "_request_raw", request_raw) + monkeypatch.setattr(client._sync, "_subscribe_turn_notifications", observe_release) + task = asyncio.create_task(client.turn_start("thread-1", "hello")) + try: + assert await asyncio.to_thread(entered.wait, 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + respond.set() + assert await asyncio.to_thread(released.wait, 2) + assert client._sync._router._turn_states == {} + + asyncio.run(scenario()) + + +def test_cancelled_queued_start_does_not_send_a_request(monkeypatch): + with ThreadPoolExecutor(max_workers=1) as executor: + monkeypatch.setattr("openai_codex.async_client._TURN_START_EXECUTOR", executor) + release_worker = threading.Event() + worker_started = threading.Event() + request_sent = threading.Event() + + def occupy_worker(): + worker_started.set() + assert release_worker.wait(timeout=5) + + occupied = executor.submit(occupy_worker) + assert worker_started.wait(timeout=5) + client = AsyncCodexClient() + monkeypatch.setattr(client._sync, "turn_start", lambda *args: request_sent.set()) + + async def scenario(): + task = asyncio.create_task(client.turn_start("thread-1", "hello")) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + try: + asyncio.run(scenario()) + finally: + release_worker.set() + occupied.result(timeout=5) + executor.submit(lambda: None).result(timeout=5) + assert not request_sent.is_set() + + +def test_low_level_start_keeps_implicit_registration_and_explicit_unregister(monkeypatch): + client = CodexClient() + + def request_raw(method, params): + for event in turn_events(client): + client._router.route_notification(event) + return {"turn": {"id": "turn-1", "status": "completed", "items": []}} + + monkeypatch.setattr(client, "_request_raw", request_raw) + started = client.turn_start("thread-1", "hello") + 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"): + client.next_turn_notification(started.turn.id) + assert client._router._turn_states == {}