mirror of
https://github.com/openai/codex.git
synced 2026-05-29 15:30:22 +00:00
Route Python SDK turn notifications by ID (#21778)
## Why The Python SDK previously protected the stdio transport with a single active turn-consumer guard. That avoided competing reads from stdout, but it also meant one `Codex`/`AsyncCodex` client could not stream multiple active turns at the same time. Notifications could also arrive before the caller received a `TurnHandle` and registered for streaming, so the SDK needed an explicit routing layer instead of letting individual API calls read directly from the shared transport. ## What Changed - Added a private `MessageRouter` that owns per-request response queues, per-turn notification queues, pending turn-notification replay, and global notification delivery behind a single stdout reader thread. - Generated typed notification routing metadata so turn IDs come from known payload shapes instead of router-side attribute guessing, with explicit fallback handling for unknown notification payloads. - Updated sync and async turn streaming so `TurnHandle.stream()`/`run()` and `stream_text()` consume only notifications for their own turn ID, while `AsyncAppServerClient` no longer serializes all transport calls behind one async lock. - Cleared pending turn-notification buffers when unregistered turns complete so never-consumed turn handles do not leave stale queues behind. - Removed the internal stream-until helper now that turn completion waiting can register directly with routed turn notifications. - Updated Python SDK docs and focused tests for concurrent transport calls, interleaved turn routing, buffered early notifications, unknown notification routing, async delegation, and routed turn completion behavior. ## Validation - `uv run --extra dev ruff format scripts/update_sdk_artifacts.py src/codex_app_server/_message_router.py src/codex_app_server/client.py src/codex_app_server/generated/notification_registry.py tests/test_client_rpc_methods.py tests/test_public_api_runtime_behavior.py tests/test_async_client_behavior.py` - `uv run --extra dev ruff check scripts/update_sdk_artifacts.py src/codex_app_server/_message_router.py src/codex_app_server/client.py src/codex_app_server/generated/notification_registry.py tests/test_client_rpc_methods.py tests/test_public_api_runtime_behavior.py tests/test_async_client_behavior.py` - `uv run --extra dev pytest tests/test_client_rpc_methods.py tests/test_public_api_runtime_behavior.py tests/test_async_client_behavior.py` - `git diff --check` --------- Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Ahmed Ibrahim
@@ -2,11 +2,17 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
from codex_app_server.async_client import AsyncAppServerClient
|
||||
from codex_app_server.generated.v2_all import (
|
||||
AgentMessageDeltaNotification,
|
||||
TurnCompletedNotification,
|
||||
)
|
||||
from codex_app_server.models import Notification, UnknownNotification
|
||||
|
||||
|
||||
def test_async_client_serializes_transport_calls() -> None:
|
||||
def test_async_client_allows_concurrent_transport_calls() -> None:
|
||||
async def scenario() -> int:
|
||||
client = AsyncAppServerClient()
|
||||
active = 0
|
||||
@@ -24,10 +30,10 @@ def test_async_client_serializes_transport_calls() -> None:
|
||||
await asyncio.gather(client.model_list(), client.model_list())
|
||||
return max_active
|
||||
|
||||
assert asyncio.run(scenario()) == 1
|
||||
assert asyncio.run(scenario()) == 2
|
||||
|
||||
|
||||
def test_async_stream_text_is_incremental_and_blocks_parallel_calls() -> None:
|
||||
def test_async_stream_text_is_incremental_without_blocking_parallel_calls() -> None:
|
||||
async def scenario() -> tuple[str, list[str], bool]:
|
||||
client = AsyncAppServerClient()
|
||||
|
||||
@@ -46,19 +52,155 @@ def test_async_stream_text_is_incremental_and_blocks_parallel_calls() -> None:
|
||||
stream = client.stream_text("thread-1", "hello")
|
||||
first = await anext(stream)
|
||||
|
||||
blocked_before_stream_done = False
|
||||
competing_call = asyncio.create_task(client.model_list())
|
||||
await asyncio.sleep(0.01)
|
||||
blocked_before_stream_done = not competing_call.done()
|
||||
competing_call_done_before_stream_done = competing_call.done()
|
||||
|
||||
remaining: list[str] = []
|
||||
async for item in stream:
|
||||
remaining.append(item)
|
||||
|
||||
await competing_call
|
||||
return first, remaining, blocked_before_stream_done
|
||||
return first, remaining, competing_call_done_before_stream_done
|
||||
|
||||
first, remaining, blocked = asyncio.run(scenario())
|
||||
first, remaining, was_unblocked = asyncio.run(scenario())
|
||||
assert first == "first"
|
||||
assert remaining == ["second", "third"]
|
||||
assert blocked
|
||||
assert was_unblocked
|
||||
|
||||
|
||||
def test_async_client_turn_notification_methods_delegate_to_sync_client() -> None:
|
||||
async def scenario() -> tuple[list[tuple[str, str]], Notification, str]:
|
||||
client = AsyncAppServerClient()
|
||||
event = Notification(
|
||||
method="unknown/direct",
|
||||
payload=UnknownNotification(params={"turnId": "turn-1"}),
|
||||
)
|
||||
completed = TurnCompletedNotification.model_validate(
|
||||
{
|
||||
"threadId": "thread-1",
|
||||
"turn": {"id": "turn-1", "items": [], "status": "completed"},
|
||||
}
|
||||
)
|
||||
calls: list[tuple[str, str]] = []
|
||||
|
||||
def fake_register(turn_id: str) -> None:
|
||||
calls.append(("register", turn_id))
|
||||
|
||||
def fake_unregister(turn_id: str) -> None:
|
||||
calls.append(("unregister", turn_id))
|
||||
|
||||
def fake_next(turn_id: str) -> Notification:
|
||||
calls.append(("next", turn_id))
|
||||
return event
|
||||
|
||||
def fake_wait(turn_id: str) -> TurnCompletedNotification:
|
||||
calls.append(("wait", turn_id))
|
||||
return completed
|
||||
|
||||
client._sync.register_turn_notifications = fake_register # type: ignore[method-assign]
|
||||
client._sync.unregister_turn_notifications = fake_unregister # type: ignore[method-assign]
|
||||
client._sync.next_turn_notification = fake_next # type: ignore[method-assign]
|
||||
client._sync.wait_for_turn_completed = fake_wait # type: ignore[method-assign]
|
||||
|
||||
client.register_turn_notifications("turn-1")
|
||||
next_event = await client.next_turn_notification("turn-1")
|
||||
completed_event = await client.wait_for_turn_completed("turn-1")
|
||||
client.unregister_turn_notifications("turn-1")
|
||||
|
||||
return calls, next_event, completed_event.turn.id
|
||||
|
||||
calls, next_event, completed_turn_id = asyncio.run(scenario())
|
||||
|
||||
assert (
|
||||
calls,
|
||||
next_event,
|
||||
completed_turn_id,
|
||||
) == (
|
||||
[
|
||||
("register", "turn-1"),
|
||||
("next", "turn-1"),
|
||||
("wait", "turn-1"),
|
||||
("unregister", "turn-1"),
|
||||
],
|
||||
Notification(
|
||||
method="unknown/direct",
|
||||
payload=UnknownNotification(params={"turnId": "turn-1"}),
|
||||
),
|
||||
"turn-1",
|
||||
)
|
||||
|
||||
|
||||
def test_async_stream_text_uses_sync_turn_routing() -> None:
|
||||
async def scenario() -> tuple[list[tuple[str, str]], list[str]]:
|
||||
client = AsyncAppServerClient()
|
||||
notifications = [
|
||||
Notification(
|
||||
method="item/agentMessage/delta",
|
||||
payload=AgentMessageDeltaNotification.model_validate(
|
||||
{
|
||||
"delta": "first",
|
||||
"itemId": "item-1",
|
||||
"threadId": "thread-1",
|
||||
"turnId": "turn-1",
|
||||
}
|
||||
),
|
||||
),
|
||||
Notification(
|
||||
method="item/agentMessage/delta",
|
||||
payload=AgentMessageDeltaNotification.model_validate(
|
||||
{
|
||||
"delta": "second",
|
||||
"itemId": "item-2",
|
||||
"threadId": "thread-1",
|
||||
"turnId": "turn-1",
|
||||
}
|
||||
),
|
||||
),
|
||||
Notification(
|
||||
method="turn/completed",
|
||||
payload=TurnCompletedNotification.model_validate(
|
||||
{
|
||||
"threadId": "thread-1",
|
||||
"turn": {"id": "turn-1", "items": [], "status": "completed"},
|
||||
}
|
||||
),
|
||||
),
|
||||
]
|
||||
calls: list[tuple[str, str]] = []
|
||||
|
||||
def fake_turn_start(thread_id: str, text: str, *, params=None): # type: ignore[no-untyped-def]
|
||||
calls.append(("turn_start", thread_id))
|
||||
return SimpleNamespace(turn=SimpleNamespace(id="turn-1"))
|
||||
|
||||
def fake_register(turn_id: str) -> None:
|
||||
calls.append(("register", turn_id))
|
||||
|
||||
def fake_next(turn_id: str) -> Notification:
|
||||
calls.append(("next", turn_id))
|
||||
return notifications.pop(0)
|
||||
|
||||
def fake_unregister(turn_id: str) -> None:
|
||||
calls.append(("unregister", turn_id))
|
||||
|
||||
client._sync.turn_start = fake_turn_start # type: ignore[method-assign]
|
||||
client._sync.register_turn_notifications = fake_register # type: ignore[method-assign]
|
||||
client._sync.next_turn_notification = fake_next # type: ignore[method-assign]
|
||||
client._sync.unregister_turn_notifications = fake_unregister # type: ignore[method-assign]
|
||||
|
||||
chunks = [chunk async for chunk in client.stream_text("thread-1", "hello")]
|
||||
return calls, [chunk.delta for chunk in chunks]
|
||||
|
||||
calls, deltas = asyncio.run(scenario())
|
||||
|
||||
assert (calls, deltas) == (
|
||||
[
|
||||
("turn_start", "thread-1"),
|
||||
("register", "turn-1"),
|
||||
("next", "turn-1"),
|
||||
("next", "turn-1"),
|
||||
("next", "turn-1"),
|
||||
("unregister", "turn-1"),
|
||||
],
|
||||
["first", "second"],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user