File size: 10,383 Bytes
f15fb1d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
from __future__ import annotations

import asyncio

from app_server_harness import AppServerHarness
from app_server_helpers import (
    agent_message_texts,
    agent_message_texts_from_items,
    next_async_delta,
    next_sync_delta,
    streaming_response,
)

from openai_codex import AsyncCodex, Codex
from openai_codex.generated.v2_all import (
    AgentMessageDeltaNotification,
    TurnCompletedNotification,
    TurnStatus,
)


def test_sync_stream_routes_text_deltas_and_completion(tmp_path) -> None:
    """A sync turn stream should expose deltas, completed items, and completion."""
    with AppServerHarness(tmp_path) as harness:
        harness.responses.enqueue_sse(streaming_response("stream-1", "msg-stream-1", ["he", "llo"]))

        with Codex(config=harness.app_server_config()) as codex:
            thread = codex.thread_start()
            stream = thread.turn("stream please").stream()
            events = list(stream)
            request = harness.responses.single_request()

    assert {
        "deltas": [
            event.payload.delta
            for event in events
            if isinstance(event.payload, AgentMessageDeltaNotification)
        ],
        "agent_messages": agent_message_texts(events),
        "request_user_texts": request.message_input_texts("user")[-1:],
        "completed_statuses": [
            event.payload.turn.status
            for event in events
            if isinstance(event.payload, TurnCompletedNotification)
        ],
    } == {
        "deltas": ["he", "llo"],
        "agent_messages": ["hello"],
        "request_user_texts": ["stream please"],
        "completed_statuses": [TurnStatus.completed],
    }


def test_turn_run_returns_completed_turn(tmp_path) -> None:
    """TurnHandle.run should collect output and completion metadata."""
    with AppServerHarness(tmp_path) as harness:
        harness.responses.enqueue_assistant_message("turn complete", response_id="turn-run-1")

        with Codex(config=harness.app_server_config()) as codex:
            thread = codex.thread_start()
            turn = thread.turn("complete this turn")
            completed = turn.run()

    assert {
        "turn_id": completed.id,
        "status": completed.status,
        "agent_messages": agent_message_texts_from_items(completed.items),
        "final_response": completed.final_response,
    } == {
        "turn_id": turn.id,
        "status": TurnStatus.completed,
        "agent_messages": ["turn complete"],
        "final_response": "turn complete",
    }


def test_async_stream_routes_text_deltas_and_completion(tmp_path) -> None:
    """An async turn stream should expose the same notification sequence."""

    async def scenario() -> None:
        """Stream one async turn against the real pinned app-server."""
        with AppServerHarness(tmp_path) as harness:
            harness.responses.enqueue_sse(
                streaming_response("async-stream-1", "msg-async-stream-1", ["as", "ync"])
            )

            async with AsyncCodex(config=harness.app_server_config()) as codex:
                thread = await codex.thread_start()
                turn = await thread.turn("async stream please")
                events = [event async for event in turn.stream()]
                request = harness.responses.single_request()

        assert {
            "deltas": [
                event.payload.delta
                for event in events
                if isinstance(event.payload, AgentMessageDeltaNotification)
            ],
            "agent_messages": agent_message_texts(events),
            "request_user_texts": request.message_input_texts("user")[-1:],
            "completed_statuses": [
                event.payload.turn.status
                for event in events
                if isinstance(event.payload, TurnCompletedNotification)
            ],
        } == {
            "deltas": ["as", "ync"],
            "agent_messages": ["async"],
            "request_user_texts": ["async stream please"],
            "completed_statuses": [TurnStatus.completed],
        }

    asyncio.run(scenario())


def test_low_level_sync_stream_text_uses_real_turn_routing(tmp_path) -> None:
    """CodexClient.stream_text should stream through a real app-server turn."""
    with AppServerHarness(tmp_path) as harness:
        harness.responses.enqueue_sse(
            streaming_response("low-sync-stream", "msg-low-sync-stream", ["fir", "st"])
        )

        with Codex(config=harness.app_server_config()) as codex:
            thread = codex.thread_start()
            chunks = list(codex._client.stream_text(thread.id, "low-level sync"))  # noqa: SLF001

    assert [chunk.delta for chunk in chunks] == ["fir", "st"]


def test_low_level_async_stream_text_allows_parallel_model_list(tmp_path) -> None:
    """Async stream_text should yield without blocking another app-server request."""

    async def scenario() -> None:
        """Leave a stream open while another async request completes."""
        with AppServerHarness(tmp_path) as harness:
            harness.responses.enqueue_sse(
                streaming_response(
                    "low-async-stream",
                    "msg-low-async-stream",
                    ["one", "two", "three"],
                ),
                delay_between_events_s=0.03,
            )

            async with AsyncCodex(config=harness.app_server_config()) as codex:
                thread = await codex.thread_start()
                stream = codex._client.stream_text(  # noqa: SLF001
                    thread.id,
                    "low-level async",
                )
                first = await anext(stream)
                models_task = asyncio.create_task(codex.models())
                models = await asyncio.wait_for(models_task, timeout=1.0)
                remaining = [chunk.delta async for chunk in stream]

        assert {
            "first": first.delta,
            "remaining": remaining,
            "models_payload_has_data": isinstance(
                models.model_dump(by_alias=True, mode="json").get("data"),
                list,
            ),
        } == {
            "first": "one",
            "remaining": ["two", "three"],
            "models_payload_has_data": True,
        }

    asyncio.run(scenario())


def test_interleaved_sync_turn_streams_route_by_turn_id(tmp_path) -> None:
    """Two sync streams on one client should consume only their own notifications."""
    with AppServerHarness(tmp_path) as harness:
        harness.responses.enqueue_sse(
            streaming_response("first-stream", "msg-first", ["one-", "done"]),
            delay_between_events_s=0.01,
        )
        harness.responses.enqueue_sse(
            streaming_response("second-stream", "msg-second", ["two-", "done"]),
            delay_between_events_s=0.01,
        )

        with Codex(config=harness.app_server_config()) as codex:
            first_thread = codex.thread_start()
            second_thread = codex.thread_start()
            first_turn = first_thread.turn("first")
            second_turn = second_thread.turn("second")

            first_stream = first_turn.stream()
            second_stream = second_turn.stream()
            first_first_delta = next_sync_delta(first_stream)
            second_first_delta = next_sync_delta(second_stream)
            first_second_delta = next_sync_delta(first_stream)
            second_second_delta = next_sync_delta(second_stream)
            first_tail = list(first_stream)
            second_tail = list(second_stream)

    assert {
        "streams": sorted(
            [
                (
                    first_first_delta,
                    first_second_delta,
                    agent_message_texts(first_tail),
                ),
                (
                    second_first_delta,
                    second_second_delta,
                    agent_message_texts(second_tail),
                ),
            ]
        ),
    } == {
        "streams": [
            ("one-", "done", ["one-done"]),
            ("two-", "done", ["two-done"]),
        ],
    }


def test_interleaved_async_turn_streams_route_by_turn_id(tmp_path) -> None:
    """Two async streams on one client should consume only their own notifications."""

    async def scenario() -> None:
        """Interleave async stream consumers against one app-server process."""
        with AppServerHarness(tmp_path) as harness:
            harness.responses.enqueue_sse(
                streaming_response("async-first", "msg-async-first", ["a1", "-done"]),
                delay_between_events_s=0.01,
            )
            harness.responses.enqueue_sse(
                streaming_response("async-second", "msg-async-second", ["a2", "-done"]),
                delay_between_events_s=0.01,
            )

            async with AsyncCodex(config=harness.app_server_config()) as codex:
                first_thread = await codex.thread_start()
                second_thread = await codex.thread_start()
                first_turn = await first_thread.turn("async first")
                second_turn = await second_thread.turn("async second")

                first_stream = first_turn.stream()
                second_stream = second_turn.stream()
                first_first_delta = await next_async_delta(first_stream)
                second_first_delta = await next_async_delta(second_stream)
                first_second_delta = await next_async_delta(first_stream)
                second_second_delta = await next_async_delta(second_stream)
                first_tail = [event async for event in first_stream]
                second_tail = [event async for event in second_stream]

        assert {
            "streams": sorted(
                [
                    (
                        first_first_delta,
                        first_second_delta,
                        agent_message_texts(first_tail),
                    ),
                    (
                        second_first_delta,
                        second_second_delta,
                        agent_message_texts(second_tail),
                    ),
                ]
            ),
        } == {
            "streams": [
                ("a1", "-done", ["a1-done"]),
                ("a2", "-done", ["a2-done"]),
            ],
        }

    asyncio.run(scenario())