diff --git a/tests/lib/streaming/agents/test_streams.py b/tests/lib/streaming/agents/test_streams.py index 0e77bcabff..b1495bfc65 100644 --- a/tests/lib/streaming/agents/test_streams.py +++ b/tests/lib/streaming/agents/test_streams.py @@ -481,7 +481,8 @@ def handler(arguments: dict[str, object]) -> str: repeated_call = {"event_id": "original-call", **call()} payloads = [turn_event("created"), repeated_call] - for index in range(5000): + delta_count = 5000 + for index in range(delta_count): delta_event: dict[str, object] = { "type": "agent.session.turn.output_text.delta", "event_id": f"delta-{index}", @@ -492,10 +493,12 @@ def handler(arguments: dict[str, object]) -> str: "content_index": 0, "delta": str(index), } - payloads.extend([delta_event, delta_event]) + payloads.append(delta_event) + payloads.append(payloads[-1]) # Repeat the last delta after sustained cache eviction. payloads.extend([repeated_call, turn_event("completed"), idle()]) server.body = EventBody(payloads) deltas: list[str] = [] + tool_event_deliveries = 0 if isinstance(sdk, AsyncOpenAI): async with sdk.beta.agents.sessions.stream( "session_test", input="Hello", tool_handlers={"search": handler} @@ -503,6 +506,8 @@ def handler(arguments: dict[str, object]) -> str: async for event in stream: if event.type == "agent.session.turn.output_text.delta": deltas.append(event.delta) + elif event.event_id == repeated_call["event_id"]: + tool_event_deliveries += 1 assert len(stream._state.event_ids) <= 1024 assert len(stream._state.recent_events) <= 1024 else: @@ -512,9 +517,12 @@ def handler(arguments: dict[str, object]) -> str: for event in sync_stream: if event.type == "agent.session.turn.output_text.delta": deltas.append(event.delta) + elif event.event_id == repeated_call["event_id"]: + tool_event_deliveries += 1 assert len(sync_stream._state.event_ids) <= 1024 assert len(sync_stream._state.recent_events) <= 1024 - assert deltas == [str(index) for index in range(5000)] + assert deltas == [str(index) for index in range(delta_count)] + assert tool_event_deliveries == 2 assert calls == [{"query": "test"}] assert len(server.inputs()) == 2 assert server.body.closed