Route Python SDK turn notifications by ID

sdk/python/docs/api-reference.md

@@ -2,7 +2,7 @@
 
 Public surface of `codex_app_server` for app-server v2.
 
-This SDK surface is experimental. The current implementation intentionally allows only one active turn consumer (`Thread.run()`, `TurnHandle.stream()`, or `TurnHandle.run()`) per client instance at a time.
+This SDK surface is experimental. Turn streams are routed by turn ID so one client can consume multiple active turns concurrently.
 
 ## Package Entry
 
@@ -137,8 +137,8 @@ Use `turn(...)` when you need low-level turn control (`stream()`, `steer()`,
 
 Behavior notes:
 
-- `stream()` and `run()` are exclusive per client instance in the current experimental build
-- starting a second turn consumer on the same `Codex` instance raises `RuntimeError`
+- `stream()` and `run()` consume only notifications for their own turn ID
+- one `Codex` instance can stream multiple active turns concurrently
 
 ### AsyncTurnHandle
 
@@ -149,8 +149,8 @@ Behavior notes:
 
 Behavior notes:
 
-- `stream()` and `run()` are exclusive per client instance in the current experimental build
-- starting a second turn consumer on the same `AsyncCodex` instance raises `RuntimeError`
+- `stream()` and `run()` consume only notifications for their own turn ID
+- one `AsyncCodex` instance can stream multiple active turns concurrently
 
 ## Inputs
 

sdk/python/docs/getting-started.md

@@ -45,7 +45,7 @@ What happened:
 - `thread.run("...")` started a turn, consumed events until completion, and returned the final assistant response plus collected items and usage.
 - `result.final_response` is `None` when no final-answer or phase-less assistant message item completes for the turn.
 - use `thread.turn(...)` when you need a `TurnHandle` for streaming, steering, interrupting, or turn IDs/status
-- one client can have only one active turn consumer (`thread.run(...)`, `TurnHandle.stream()`, or `TurnHandle.run()`) at a time in the current experimental build
+- one client can consume multiple active turns concurrently; turn streams are routed by turn ID
 
 ## 3) Continue the same thread (multi-turn)
 

sdk/python/scripts/update_sdk_artifacts.py

@@ -585,6 +585,43 @@ def _notification_specs() -> list[tuple[str, str]]:
     return specs
 
 
+def _notification_turn_id_specs(
+    specs: list[tuple[str, str]],
+) -> tuple[list[str], list[str]]:
+    server_notifications = json.loads(
+        (schema_root_dir() / "ServerNotification.json").read_text()
+    )
+    definitions = server_notifications.get("definitions", {})
+    if not isinstance(definitions, dict):
+        return ([], [])
+
+    direct: list[str] = []
+    nested: list[str] = []
+    for _, class_name in specs:
+        definition = definitions.get(class_name)
+        if not isinstance(definition, dict):
+            continue
+        props = definition.get("properties", {})
+        if not isinstance(props, dict):
+            continue
+        if "turnId" in props:
+            direct.append(class_name)
+            continue
+        turn = props.get("turn")
+        if isinstance(turn, dict) and turn.get("$ref") == "#/definitions/Turn":
+            nested.append(class_name)
+
+    return (sorted(set(direct)), sorted(set(nested)))
+
+
+def _type_tuple_source(class_names: list[str]) -> str:
+    if not class_names:
+        return "()"
+    if len(class_names) == 1:
+        return f"({class_names[0]},)"
+    return "(\n" + "".join(f"    {class_name},\n" for class_name in class_names) + ")"
+
+
 def generate_notification_registry() -> None:
     out = (
         sdk_root()
@@ -595,6 +632,7 @@ def generate_notification_registry() -> None:
     )
     specs = _notification_specs()
     class_names = sorted({class_name for _, class_name in specs})
+    direct_turn_id_types, nested_turn_types = _notification_turn_id_specs(specs)
 
     lines = [
         "# Auto-generated by scripts/update_sdk_artifacts.py",
@@ -616,7 +654,26 @@ def generate_notification_registry() -> None:
     )
     for method, class_name in specs:
         lines.append(f'    "{method}": {class_name},')
-    lines.extend(["}", ""])
+    lines.extend(
+        [
+            "}",
+            "",
+            "DIRECT_TURN_ID_NOTIFICATION_TYPES: tuple[type[BaseModel], ...] = "
+            f"{_type_tuple_source(direct_turn_id_types)}",
+            "",
+            "NESTED_TURN_NOTIFICATION_TYPES: tuple[type[BaseModel], ...] = "
+            f"{_type_tuple_source(nested_turn_types)}",
+            "",
+            "",
+            "def notification_turn_id(payload: BaseModel) -> str | None:",
+            "    if isinstance(payload, DIRECT_TURN_ID_NOTIFICATION_TYPES):",
+            "        return payload.turn_id if isinstance(payload.turn_id, str) else None",
+            "    if isinstance(payload, NESTED_TURN_NOTIFICATION_TYPES):",
+            "        return payload.turn.id",
+            "    return None",
+            "",
+        ]
+    )
 
     out.write_text("\n".join(lines))
 

sdk/python/src/codex_app_server/_message_router.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import queue
+import threading
+from collections import deque
+
+from .errors import AppServerError, map_jsonrpc_error
+from .generated.notification_registry import notification_turn_id
+from .models import JsonValue, Notification, UnknownNotification
+
+ResponseQueueItem = JsonValue | BaseException
+NotificationQueueItem = Notification | BaseException
+
+
+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.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._response_waiters: dict[str, queue.Queue[ResponseQueueItem]] = {}
+        self._turn_notifications: dict[str, queue.Queue[NotificationQueueItem]] = {}
+        self._pending_turn_notifications: dict[str, deque[Notification]] = {}
+        self._global_notifications: queue.Queue[NotificationQueueItem] = queue.Queue()
+
+    def create_response_waiter(self, request_id: str) -> queue.Queue[ResponseQueueItem]:
+        """Register a one-shot queue for a JSON-RPC response id."""
+
+        waiter: queue.Queue[ResponseQueueItem] = queue.Queue(maxsize=1)
+        with self._lock:
+            self._response_waiters[request_id] = waiter
+        return waiter
+
+    def discard_response_waiter(self, request_id: str) -> None:
+        """Remove a response waiter when the request could not be written."""
+
+        with self._lock:
+            self._response_waiters.pop(request_id, None)
+
+    def next_global_notification(self) -> Notification:
+        """Block until the next notification that is not scoped to a turn."""
+
+        item = self._global_notifications.get()
+        if isinstance(item, BaseException):
+            raise item
+        return item
+
+    def register_turn(self, turn_id: str) -> None:
+        """Register a queue for a turn stream and replay early events."""
+
+        turn_queue: queue.Queue[NotificationQueueItem] = queue.Queue()
+        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)
+
+    def unregister_turn(self, turn_id: str) -> None:
+        """Stop routing future turn events to the stream queue."""
+
+        with self._lock:
+            self._turn_notifications.pop(turn_id, None)
+
+    def next_turn_notification(self, turn_id: str) -> Notification:
+        """Block until the next notification for a registered turn."""
+
+        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
+
+    def route_response(self, msg: dict[str, JsonValue]) -> None:
+        """Deliver a JSON-RPC response or error to its request waiter."""
+
+        request_id = msg.get("id")
+        with self._lock:
+            waiter = self._response_waiters.pop(str(request_id), None)
+        if waiter is None:
+            return
+
+        if "error" in msg:
+            err = msg["error"]
+            if isinstance(err, dict):
+                waiter.put(
+                    map_jsonrpc_error(
+                        int(err.get("code", -32000)),
+                        str(err.get("message", "unknown")),
+                        err.get("data"),
+                    )
+                )
+            else:
+                waiter.put(AppServerError("Malformed JSON-RPC error response"))
+            return
+
+        waiter.put(msg.get("result"))
+
+    def route_notification(self, notification: Notification) -> None:
+        """Deliver a notification to a turn queue or the global queue."""
+
+        turn_id = self._notification_turn_id(notification)
+        if turn_id is None:
+            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)
+
+    def fail_all(self, exc: BaseException) -> None:
+        """Wake every blocked waiter when the reader thread exits."""
+
+        with self._lock:
+            response_waiters = list(self._response_waiters.values())
+            self._response_waiters.clear()
+            turn_queues = list(self._turn_notifications.values())
+            self._pending_turn_notifications.clear()
+        # Put the same transport failure into every queue so no SDK call blocks
+        # forever waiting for a response that cannot arrive.
+        for waiter in response_waiters:
+            waiter.put(exc)
+        for turn_queue in turn_queues:
+            turn_queue.put(exc)
+        self._global_notifications.put(exc)
+
+    def _notification_turn_id(self, notification: Notification) -> str | None:
+        payload = notification.payload
+        if isinstance(payload, UnknownNotification):
+            raw_turn_id = payload.params.get("turnId")
+            if isinstance(raw_turn_id, str):
+                return raw_turn_id
+            raw_turn = payload.params.get("turn")
+            if isinstance(raw_turn, dict):
+                raw_nested_turn_id = raw_turn.get("id")
+                if isinstance(raw_nested_turn_id, str):
+                    return raw_nested_turn_id
+            return None
+        return notification_turn_id(payload)

sdk/python/src/codex_app_server/api.py

@@ -38,14 +38,14 @@
 )
 from .models import InitializeResponse, JsonObject, Notification, ServerInfo
 from ._inputs import (
-    ImageInput,
+    ImageInput as ImageInput,
     Input,
-    InputItem,
-    LocalImageInput,
-    MentionInput,
+    InputItem as InputItem,
+    LocalImageInput as LocalImageInput,
+    MentionInput as MentionInput,
     RunInput,
-    SkillInput,
-    TextInput,
+    SkillInput as SkillInput,
+    TextInput as TextInput,
     _normalize_run_input,
     _to_wire_input,
 )
@@ -274,6 +274,7 @@ def thread_archive(self, thread_id: str) -> ThreadArchiveResponse:
     def thread_unarchive(self, thread_id: str) -> Thread:
         unarchived = self._client.thread_unarchive(thread_id)
         return Thread(self._client, unarchived.thread.id)
+
     # END GENERATED: Codex.flat_methods
 
     def models(self, *, include_hidden: bool = False) -> ModelListResponse:
@@ -476,6 +477,7 @@ async def thread_unarchive(self, thread_id: str) -> AsyncThread:
         await self._ensure_initialized()
         unarchived = await self._client.thread_unarchive(thread_id)
         return AsyncThread(self, unarchived.thread.id)
+
     # END GENERATED: AsyncCodex.flat_methods
 
     async def models(self, *, include_hidden: bool = False) -> ModelListResponse:
@@ -555,6 +557,7 @@ def turn(
         )
         turn = self._client.turn_start(self.id, wire_input, params=params)
         return TurnHandle(self._client, self.id, turn.turn.id)
+
     # END GENERATED: Thread.flat_methods
 
     def read(self, *, include_turns: bool = False) -> ThreadReadResponse:
@@ -644,6 +647,7 @@ async def turn(
             params=params,
         )
         return AsyncTurnHandle(self._codex, self.id, turn.turn.id)
+
     # END GENERATED: AsyncThread.flat_methods
 
     async def read(self, *, include_turns: bool = False) -> ThreadReadResponse:
@@ -674,11 +678,10 @@ def interrupt(self) -> TurnInterruptResponse:
         return self._client.turn_interrupt(self.thread_id, self.id)
 
     def stream(self) -> Iterator[Notification]:
-        # TODO: replace this client-wide experimental guard with per-turn event demux.
-        self._client.acquire_turn_consumer(self.id)
+        self._client.register_turn_notifications(self.id)
         try:
             while True:
-                event = self._client.next_notification()
+                event = self._client.next_turn_notification(self.id)
                 yield event
                 if (
                     event.method == "turn/completed"
@@ -687,7 +690,7 @@ def stream(self) -> Iterator[Notification]:
                 ):
                     break
         finally:
-            self._client.release_turn_consumer(self.id)
+            self._client.unregister_turn_notifications(self.id)
 
     def run(self) -> AppServerTurn:
         completed: TurnCompletedNotification | None = None
@@ -728,11 +731,10 @@ async def interrupt(self) -> TurnInterruptResponse:
 
     async def stream(self) -> AsyncIterator[Notification]:
         await self._codex._ensure_initialized()
-        # TODO: replace this client-wide experimental guard with per-turn event demux.
-        self._codex._client.acquire_turn_consumer(self.id)
+        self._codex._client.register_turn_notifications(self.id)
         try:
             while True:
-                event = await self._codex._client.next_notification()
+                event = await self._codex._client.next_turn_notification(self.id)
                 yield event
                 if (
                     event.method == "turn/completed"
@@ -741,7 +743,7 @@ async def stream(self) -> AsyncIterator[Notification]:
                 ):
                     break
         finally:
-            self._codex._client.release_turn_consumer(self.id)
+            self._codex._client.unregister_turn_notifications(self.id)
 
     async def run(self) -> AppServerTurn:
         completed: TurnCompletedNotification | None = None