codex/sdk/python/tests/test_async_client_behavior.py

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_allows_concurrent_transport_calls() -> None:
    async def scenario() -> int:
        client = AsyncAppServerClient()
        active = 0
        max_active = 0

        def fake_model_list(include_hidden: bool = False) -> bool:
            nonlocal active, max_active
            active += 1
            max_active = max(max_active, active)
            time.sleep(0.05)
            active -= 1
            return include_hidden

        client._sync.model_list = fake_model_list  # type: ignore[method-assign]
        await asyncio.gather(client.model_list(), client.model_list())
        return max_active

    assert asyncio.run(scenario()) == 2


def test_async_stream_text_is_incremental_without_blocking_parallel_calls() -> None:
    async def scenario() -> tuple[str, list[str], bool]:
        client = AsyncAppServerClient()

        def fake_stream_text(thread_id: str, text: str, params=None):  # type: ignore[no-untyped-def]
            yield "first"
            time.sleep(0.03)
            yield "second"
            yield "third"

        def fake_model_list(include_hidden: bool = False) -> str:
            return "done"

        client._sync.stream_text = fake_stream_text  # type: ignore[method-assign]
        client._sync.model_list = fake_model_list  # type: ignore[method-assign]

        stream = client.stream_text("thread-1", "hello")
        first = await anext(stream)

        competing_call = asyncio.create_task(client.model_list())
        await asyncio.sleep(0.01)
        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, competing_call_done_before_stream_done

    first, remaining, was_unblocked = asyncio.run(scenario())
    assert first == "first"
    assert remaining == ["second", "third"]
    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"],
    )