# SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import asyncio from typing import Any import pytest from vllm_omni.experimental.ar_diffusion.session import ( ARDiffusionChunkMetadataError, ARDiffusionEventAcceptanceStatus, ARDiffusionEventQueueFullError, ARDiffusionPreparedControls, ARDiffusionSession, ARDiffusionSessionError, ARDiffusionSessionEvent, ARDiffusionSessionManager, ARDiffusionSessionStateError, ARDiffusionSessionStatus, ARDiffusionStaleEventError, ARDiffusionWorkerLifecycle, ) from vllm_omni.experimental.ar_diffusion.tick_protocol import ( ARDiffusionChunkMetadata, ARDiffusionControlInput, ARDiffusionTickRequest, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class FakeTickConsumer: def __init__(self) -> None: self.ticks: list[ARDiffusionTickRequest] = [] self.failures_remaining = 0 self.started: asyncio.Event | None = None self.release: asyncio.Event | None = None self.metadata_override: ARDiffusionChunkMetadata | None = None async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata: self.ticks.append(tick) if self.started is not None: self.started.set() if self.release is not None: await self.release.wait() if self.failures_remaining: self.failures_remaining -= 1 raise RuntimeError("fake tick failed") return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick) def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata: return output class FakeLifecycle: def __init__(self) -> None: self.resets: list[str] = [] self.closes: list[str] = [] self.close_failures_remaining = 0 async def reset_session(self, session_id: str) -> None: self.resets.append(session_id) async def close_session(self, session_id: str) -> None: self.closes.append(session_id) if self.close_failures_remaining: self.close_failures_remaining -= 1 raise RuntimeError("fake close failed") class FailingCommitReducer: def __init__(self) -> None: self.commit_calls = 0 self.reset_calls = 0 def prepare( self, *, current_controls: dict[str, ARDiffusionControlInput], events: tuple[ARDiffusionSessionEvent, ...], chunk_index: int, ) -> ARDiffusionPreparedControls: del current_controls, events, chunk_index return ARDiffusionPreparedControls(controls=()) def commit(self, prepared: ARDiffusionPreparedControls) -> None: del prepared self.commit_calls += 1 raise RuntimeError("fake reducer commit failed") def reset(self) -> None: self.reset_calls += 1 def make_session( *, max_pending_events: int = 8, consumer: FakeTickConsumer | None = None, lifecycle: FakeLifecycle | None = None, ) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]: resolved_consumer = consumer or FakeTickConsumer() resolved_lifecycle = lifecycle or FakeLifecycle() session = ARDiffusionSession( "world-1", tick_consumer=resolved_consumer, lifecycle=resolved_lifecycle, max_pending_events=max_pending_events, request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}", ) return session, resolved_consumer, resolved_lifecycle @pytest.mark.asyncio async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None: session, consumer, _ = make_session() event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn") event_one = ARDiffusionSessionEvent( event_id=1, controls=( ARDiffusionControlInput( track="camera", schema="example.camera.v1", data={"samples": [{"x": 1.0}]}, ), ), ) accepted_two = await session.accept_event(event_two) accepted_one = await session.accept_event(event_one) duplicate = await session.accept_event(event_two) assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE assert consumer.ticks == [] metadata = await session.next_chunk() assert metadata.applied_event_ids == (1, 2) assert consumer.ticks[0] == ARDiffusionTickRequest( session_id="world-1", request_id="world-1-request-0", chunk_index=0, applied_event_ids=(1, 2), prompt="turn", controls=( ARDiffusionControlInput( track="camera", schema="example.camera.v1", data={"samples": [{"x": 1.0}]}, ), ), ) assert session.pending_event_count == 0 applied_duplicate = await session.accept_event(event_two) assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE next_metadata = await session.next_chunk() assert next_metadata.chunk_index == 1 assert next_metadata.applied_event_ids == () assert consumer.ticks[1].prompt == "turn" assert consumer.ticks[1].controls == consumer.ticks[0].controls @pytest.mark.asyncio async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None: session, consumer, _ = make_session(max_pending_events=1) payload = {"samples": [{"x": 1.0}]} composite = ARDiffusionSessionEvent( event_id=5, prompt="new prompt", controls=( ARDiffusionControlInput( track="camera", schema="example.camera.v1", data=payload, ), ARDiffusionControlInput( track="action", schema="example.action.v1", data={"keys": ["left"]}, ), ), ) await session.accept_event(composite) payload["samples"][0]["x"] = 99.0 duplicate = await session.accept_event(composite) assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"): await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected")) await session.next_chunk() assert consumer.ticks[0].applied_event_ids == (5,) assert consumer.ticks[0].prompt == "new prompt" assert consumer.ticks[0].controls[0].data["keys"] == ("left",) assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},) assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]} assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]} def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None: with pytest.raises(ValueError, match="at most once"): ARDiffusionSessionEvent( event_id=1, controls=( ARDiffusionControlInput(track="camera", schema="one", data={}), ARDiffusionControlInput(track="camera", schema="two", data={}), ), ) with pytest.raises(ValueError, match="transport-safe"): ARDiffusionControlInput( track="camera", schema="example.camera.v1", data={"object": object()}, ) @pytest.mark.asyncio async def test_failed_tick_closes_worker_and_requires_reset() -> None: consumer = FakeTickConsumer() consumer.failures_remaining = 1 lifecycle = FakeLifecycle() session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle) await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me")) with pytest.raises(RuntimeError, match="fake tick failed"): await session.next_chunk() assert session.status is ARDiffusionSessionStatus.FAILED assert session.chunk_index == 0 assert session.pending_event_count == 1 assert lifecycle.closes == ["world-1"] with pytest.raises(ARDiffusionSessionStateError, match="failed"): await session.next_chunk() await session.reset() assert session.status is ARDiffusionSessionStatus.ACTIVE assert session.pending_event_count == 0 assert lifecycle.resets == ["world-1"] await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset")) metadata = await session.next_chunk() assert metadata.chunk_index == 0 assert metadata.applied_event_ids == (4,) assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)] @pytest.mark.asyncio async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None: consumer = FakeTickConsumer() consumer.started = asyncio.Event() consumer.release = asyncio.Event() session, _, _ = make_session(consumer=consumer) await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first")) first_chunk = asyncio.create_task(session.next_chunk()) await consumer.started.wait() accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second")) assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED assert consumer.ticks[0].prompt == "first" assert consumer.ticks[0].applied_event_ids == (10,) with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"): await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale")) consumer.release.set() first_metadata = await first_chunk second_metadata = await session.next_chunk() assert first_metadata.applied_event_ids == (10,) assert second_metadata.applied_event_ids == (11,) assert consumer.ticks[1].prompt == "second" @pytest.mark.asyncio async def test_metadata_mismatch_does_not_commit_snapshot() -> None: consumer = FakeTickConsumer() consumer.metadata_override = ARDiffusionChunkMetadata( session_id="another-world", request_id="wrong", chunk_index=0, ) lifecycle = FakeLifecycle() session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle) await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending")) with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"): await session.next_chunk() assert session.chunk_index == 0 assert session.pending_event_count == 1 assert session.status is ARDiffusionSessionStatus.FAILED assert lifecycle.closes == ["world-1"] @pytest.mark.asyncio async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None: consumer = FakeTickConsumer() lifecycle = FakeLifecycle() reducer = FailingCommitReducer() session = ARDiffusionSession( "world-1", tick_consumer=consumer, lifecycle=lifecycle, max_pending_events=4, request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}", control_reducer=reducer, ) await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending")) with pytest.raises(RuntimeError, match="fake reducer commit failed"): await session.next_chunk() assert reducer.commit_calls == 1 assert session.status is ARDiffusionSessionStatus.FAILED assert session.chunk_index == 0 assert session.pending_event_count == 1 assert lifecycle.closes == ["world-1"] with pytest.raises(ARDiffusionSessionStateError, match="failed"): await session.next_chunk() await session.reset() assert reducer.reset_calls == 1 assert session.status is ARDiffusionSessionStatus.ACTIVE assert session.chunk_index == 0 assert session.pending_event_count == 0 @pytest.mark.asyncio async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None: consumer = FakeTickConsumer() lifecycle = FakeLifecycle() manager = ARDiffusionSessionManager( tick_consumer=consumer, lifecycle=lifecycle, max_pending_events=4, request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}", ) session = await manager.create_session("world-reset") await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset")) await session.next_chunk() await manager.reset_session("world-reset") assert lifecycle.resets == ["world-reset"] assert session.status is ARDiffusionSessionStatus.ACTIVE assert session.chunk_index == 0 duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate")) assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE reset_tick = await session.next_chunk() assert reset_tick.chunk_index == 0 assert consumer.ticks[-1].prompt is None assert consumer.ticks[-1].controls == () await manager.close_session("world-reset") assert lifecycle.closes == ["world-reset"] assert session.status is ARDiffusionSessionStatus.CLOSED with pytest.raises(ARDiffusionSessionStateError, match="does not exist"): await manager.get_session("world-reset") disconnected = await manager.create_session("world-disconnect") await manager.disconnect("world-disconnect") assert disconnected.status is ARDiffusionSessionStatus.CLOSED assert lifecycle.closes[-1] == "world-disconnect" @pytest.mark.asyncio async def test_manager_keeps_interleaved_sessions_independent() -> None: consumer = FakeTickConsumer() manager = ARDiffusionSessionManager( tick_consumer=consumer, lifecycle=FakeLifecycle(), max_pending_events=4, request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}", ) session_a = await manager.create_session("world-a") session_b = await manager.create_session("world-b") await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0")) await session_a.next_chunk() await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0")) await session_b.next_chunk() await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1")) await session_a.next_chunk() assert [ ( tick.session_id, tick.request_id, tick.chunk_index, tick.applied_event_ids, ) for tick in consumer.ticks ] == [ ("world-a", "world-a-request-0", 0, (10,)), ("world-b", "world-b-request-0", 0, (1,)), ("world-a", "world-a-request-1", 1, (11,)), ] assert session_a.chunk_index == 2 assert session_b.chunk_index == 1 with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"): await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a")) assert ( await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1")) ).status is ARDiffusionEventAcceptanceStatus.ACCEPTED @pytest.mark.asyncio async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None: consumer = FakeTickConsumer() lifecycle = FakeLifecycle() lifecycle.close_failures_remaining = 1 manager = ARDiffusionSessionManager( tick_consumer=consumer, lifecycle=lifecycle, max_pending_events=4, ) session = await manager.create_session("world-close-retry") with pytest.raises(RuntimeError, match="fake close failed"): await manager.close_session("world-close-retry") assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED assert await manager.get_session("world-close-retry") is session with pytest.raises(ARDiffusionSessionStateError, match="already exists"): await manager.create_session("world-close-retry") with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"): await manager.reset_session("world-close-retry") await manager.close_session("world-close-retry") assert lifecycle.closes == ["world-close-retry", "world-close-retry"] assert session.status is ARDiffusionSessionStatus.CLOSED with pytest.raises(ARDiffusionSessionStateError, match="does not exist"): await manager.get_session("world-close-retry") recreated = await manager.create_session("world-close-retry") assert recreated is not session @pytest.mark.asyncio async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None: lifecycle = FakeLifecycle() lifecycle.close_failures_remaining = 1 manager = ARDiffusionSessionManager( tick_consumer=FakeTickConsumer(), lifecycle=lifecycle, max_pending_events=4, ) session = await manager.create_session("world-disconnect-retry") with pytest.raises(RuntimeError, match="fake close failed"): await manager.disconnect("world-disconnect-retry") assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED assert await manager.get_session("world-disconnect-retry") is session await manager.close_session("world-disconnect-retry") assert session.status is ARDiffusionSessionStatus.CLOSED class FakeRPCClient: def __init__(self, result: list[Any]) -> None: self.result = result self.calls: list[dict[str, Any]] = [] async def collective_rpc(self, **kwargs: Any) -> list[Any]: self.calls.append(kwargs) return self.result @pytest.mark.asyncio async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None: client = FakeRPCClient([[True]]) lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0) await lifecycle.reset_session("world-1") await lifecycle.close_session("world-1") assert client.calls == [ { "method": "reset_ar_diffusion_session", "timeout": 3.0, "args": ("world-1",), "stage_ids": [2], }, { "method": "close_ar_diffusion_session", "timeout": 3.0, "args": ("world-1",), "stage_ids": [2], }, ] @pytest.mark.asyncio async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None: lifecycle = ARDiffusionWorkerLifecycle( FakeRPCClient([{"supported": False, "error": "not an AR stage"}]), stage_ids=[1], ) with pytest.raises(ARDiffusionSessionError, match="not an AR stage"): await lifecycle.close_session("world-1")