# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Session-state semantics of the engine-resident duplex session (``DuplexEngineSession``)."""

from __future__ import annotations

import time
from collections.abc import Callable
from dataclasses import FrozenInstanceError

import pytest

from vllm_omni.engine.duplex.config import (
    DuplexCapabilities,
    DuplexConfigError,
    DuplexOverlapPolicy,
    DuplexSessionConfig,
    DuplexTurnEventType,
    DuplexTurnState,
    ResponseCreateOptions,
)
from vllm_omni.engine.duplex.contracts import DuplexFence
from vllm_omni.engine.duplex.events import TurnEvent
from vllm_omni.engine.duplex.session.engine_session import (
    RESPONSE_REQUEST_MEASUREMENT_ORIGIN,
    DuplexEngineSession,
    DuplexFenceMismatchError,
)
from vllm_omni.engine.duplex.session.playback_ledger import apply_playback_ack
from vllm_omni.metrics.stats import DUPLEX_STAGE_TABLE_EXCLUDE, OrchestratorAggregator, StageRequestStats, StageStats
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import (
    minicpmo45_native_capabilities,
)

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def _session(
    session_id: str = "duplex-test",
    config: DuplexSessionConfig | None = None,
    *,
    num_stages: int = 1,
    log_stats: bool = False,
    clock: Callable[[], float] | None = None,
) -> DuplexEngineSession:
    session = DuplexEngineSession(
        session_id=session_id,
        config=config or DuplexSessionConfig(model="test-model"),
        num_stages=num_stages,
        log_stats=log_stats,
    )
    if clock is not None:
        session._clock = clock
    return session


def _stage_stats(
    *,
    stage_id: int,
    request_id: str = "stage-req",
    num_tokens_out: int = 3,
    vllm_ttft_ms: float = 12.0,
    vllm_tpot_ms: float = 0.0,
    serving_time_to_first_output_ms: float = 0.0,
) -> StageRequestStats:
    return StageRequestStats(
        batch_id=0,
        batch_size=1,
        num_tokens_in=7,
        num_tokens_out=num_tokens_out,
        stage_gen_time_ms=120.0,
        rx_transfer_bytes=0,
        rx_decode_time_ms=0.0,
        rx_in_flight_time_ms=0.0,
        stage_stats=StageStats(),
        stage_id=stage_id,
        request_id=request_id,
        final_output_type="text",
        vllm_ttft_ms=vllm_ttft_ms,
        vllm_tpot_ms=vllm_tpot_ms,
        serving_time_to_first_output_ms=serving_time_to_first_output_ms,
    )


def test_commit_audio_input_does_not_advance_model_turn_identity():
    session = _session()

    first = session.commit_audio_input(transcript="first chunk")
    second = session.commit_audio_input(transcript="second chunk")

    assert session.input_commit_seq == 2
    assert first.input_commit_seq == 1
    assert second.input_commit_seq == 2
    assert first.turn_id == second.turn_id == session.turn_id == 0
    assert first.message["transcript"] == "first chunk"
    assert first.message["content"] == [
        {"type": "audio_url", "audio_url": {"url": "native-duplex:input-audio"}, "transcript": "first chunk"}
    ]
    assert len(session.history) == 2
    assert session.turn_state == DuplexTurnState.USER_COMMITTED


def test_commit_audio_input_accepts_explicit_turn_id():
    session = _session()

    committed = session.commit_audio_input(turn_id=4)

    assert committed.turn_id == 4
    assert session.turn_id == 0
    assert "transcript" not in committed.message


# ---- response options ----


def test_response_options_apply_to_one_response_without_mutating_session_defaults():
    session = _session(config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64))
    session.reserve_response_options(
        ResponseCreateOptions(instructions="one response", voice="override-voice", max_tokens=8)
    )

    assert session.config.instructions == "base"
    assert session.config.voice == "base-voice"
    assert session.response_config is session.config
    session.begin_response()
    assert session.config.instructions == "base"
    assert session.config.voice == "base-voice"
    assert session.response_config.instructions == "one response"
    assert session.response_config.voice == "override-voice"
    assert session.response_config.max_tokens == 8

    session.end_response()
    assert session.response_config is session.config
    assert session.config.instructions == "base"
    assert session.config.voice == "base-voice"
    assert session.config.max_tokens == 64


def test_response_options_cannot_overwrite_an_unconsumed_reservation():
    session = _session(config=DuplexSessionConfig(instructions="base"))
    session.reserve_response_options(ResponseCreateOptions(instructions="first"))

    with pytest.raises(RuntimeError, match="already reserved"):
        session.reserve_response_options(ResponseCreateOptions(instructions="second"))

    session.begin_response()
    assert session.response_config.instructions == "first"


def test_response_options_cannot_be_reserved_while_response_is_active():
    session = _session(config=DuplexSessionConfig(instructions="base"))
    session.begin_response()

    with pytest.raises(RuntimeError, match="active"):
        session.reserve_response_options(ResponseCreateOptions(instructions="too late"))


def test_discarded_response_options_do_not_apply():
    session = _session(config=DuplexSessionConfig(instructions="base"))
    session.reserve_response_options(ResponseCreateOptions(instructions="dropped"))
    session.discard_response_options()

    session.begin_response()

    assert session.response_config.instructions == "base"


# ---- response / overlap identity ----


def test_session_owns_response_and_overlap_identity():
    session = _session()

    response_id = session.begin_response()
    assert session.turn_state == DuplexTurnState.ASSISTANT_GENERATING
    session.accumulate_overlap_speech(320)
    session.accumulate_overlap_speech(180)
    session.end_response()

    assert session.active_response_id is None
    assert session.last_response_id == response_id
    assert session.overlap_speech_ms == 500
    assert session.reset_overlap_speech() == 500
    assert session.overlap_speech_ms == 0
    assert session.turn_state == DuplexTurnState.IDLE


def test_session_composes_single_owner_ledgers_with_immutable_views():
    session = _session()

    session.bind_request("req-1")
    response_id = session.begin_response(turn_id=3)
    session.mark_audio_sent(duration_ms=240)

    assert session.active_request_id == "req-1"
    assert session.active_response_id == response_id
    assert session.active_response_turn_id == 3
    assert session.playback.sent_ms == 240
    assert session.turn_state == DuplexTurnState.ASSISTANT_PLAYING
    with pytest.raises(FrozenInstanceError):
        session.playback.sent_ms = 480  # type: ignore[misc]
    assert session.playback.sent_ms == 240


def test_barge_in_advances_epoch_and_drops_uncommitted_assistant_text():
    session = _session()
    response_id = session.begin_response()
    session.bind_request("chatcmpl-duplex-test")
    session.append_assistant_text("unplayed answer")

    new_epoch = session.barge_in()

    assert response_id is not None
    assert new_epoch == 1
    assert session.epoch == 1
    assert session.accepted_fence == DuplexFence(session.session_id, epoch=1, turn_id=0)
    assert session.active_request_id is None
    assert session.active_response_id is None
    assert session.assistant_text_buffer == ()
    assert session.history == ()
    assert session.turn_state == DuplexTurnState.BARGE_IN


# ---- playback ledger ----


def test_playback_ack_tracks_committed_cursor_separately():
    session = _session()
    session.mark_audio_sent(duration_ms=10_000)

    session.acknowledge_playback(played_ms=2_000)

    assert session.playback.generated_ms == 10_000
    assert session.playback.sent_ms == 10_000
    assert session.playback.played_ms == 2_000
    assert session.playback.committed_ms == 2_000
    assert session.playback.as_dict() == {
        "generated_ms": 10_000,
        "sent_ms": 10_000,
        "played_ms": 2_000,
        "committed_ms": 2_000,
    }


def test_playback_ack_is_scoped_per_response():
    session = _session()
    first = session.begin_response()
    session.mark_audio_sent(duration_ms=3_000)
    session.end_response()
    second = session.begin_response()
    session.mark_audio_sent(duration_ms=1_000)

    session.acknowledge_playback(played_ms=2_500, response_id=first)

    assert session.playback_for_response(first).played_ms == 2_500
    assert session.playback_for_response(second).played_ms == 0
    assert session.playback.played_ms == 0


def test_history_commit_uses_audio_text_alignment_marks():
    session = _session()
    session.begin_response()

    session.append_assistant_text("hello ")
    session.mark_audio_sent(duration_ms=1_000, text_chars=6)
    session.append_assistant_text("world")
    session.mark_audio_sent(duration_ms=2_000, text_chars=11)
    assert [(mark.text_chars, mark.audio_end_ms) for mark in session.assistant_audio_text_marks] == [
        (6, 1_000),
        (11, 2_000),
    ]
    session.acknowledge_playback(played_ms=1_200, committed_ms=1_200)

    committed = session.end_response(commit_text=True)

    assert committed == {"role": "assistant", "content": "hello w"}
    assert session.history[-1] == committed
    assert session.last_assistant_full_message == {"role": "assistant", "content": "hello world"}


def test_history_commit_with_ack_only_policy_defers_unacknowledged_text():
    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
    response_id = session.begin_response()
    session.append_assistant_text("never played")
    session.mark_audio_sent(duration_ms=1_000, text_chars=12)

    committed = session.end_response(commit_text=True)

    assert committed is None
    assert session.history == ()
    assert f"item_{response_id}" in session.pending_history_item_ids


def test_history_commit_with_ack_only_policy_keeps_only_acknowledged_prefix():
    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
    session.begin_response()
    session.append_assistant_text("hello world")
    session.mark_audio_sent(duration_ms=1_000, text_chars=6)
    session.mark_audio_sent(duration_ms=2_000, text_chars=11)
    session.acknowledge_playback(played_ms=1_000, committed_ms=1_000)

    committed = session.end_response(commit_text=True)

    assert committed == {"role": "assistant", "content": "hello"}


@pytest.mark.parametrize("audio_complete", [False, True])
def test_unaligned_response_keeps_empty_turn_before_later_user_input(audio_complete):
    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
    first = session.commit_audio_input(transcript="recite a poem")
    response_id = session.begin_response()
    session.reserve_history_item(f"item_{response_id}")
    session.append_assistant_text("An answer the user has only partly heard")
    session.mark_audio_sent(
        duration_ms=10_000,
        text_requires_complete_audio=True,
        audio_complete=audio_complete,
    )

    assert session.end_response(commit_text=True) is None
    session.barge_in()
    session.clear_playback_cursor()
    second = session.commit_audio_input(transcript="stop")
    empty_answer = {"role": "assistant", "content": ""}
    assert session.history == (first.message, empty_answer, second.message)

    ack = apply_playback_ack(session, {"response_id": response_id, "played_ms": 6_981})
    assert ack[0].to_realtime()["event"]["history_committed"] is False
    assert session.history == (first.message, empty_answer, second.message)

    if audio_complete:
        ack = apply_playback_ack(session, {"response_id": response_id, "played_ms": 10_000})
        assert ack[0].to_realtime()["event"]["history_committed"] is True
        assert session.history == (
            first.message,
            {"role": "assistant", "content": "An answer the user has only partly heard"},
            second.message,
        )


@pytest.mark.parametrize(("reserve_slot", "assistant_text"), [(True, ""), (False, "unplayed answer")])
def test_unaligned_response_only_materializes_an_existing_answer_slot(reserve_slot, assistant_text):
    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
    first = session.commit_audio_input(transcript="first input")
    response_id = session.begin_response()
    if reserve_slot:
        session.reserve_history_item(f"item_{response_id}")
    session.append_assistant_text(assistant_text)
    session.mark_audio_sent(duration_ms=10_000, text_requires_complete_audio=True, audio_complete=True)

    assert session.end_response(commit_text=reserve_slot) is None
    assert session.history == (first.message,)


@pytest.mark.parametrize("operation", ["delete", "truncate"])
def test_removing_one_empty_assistant_preserves_the_other_turn_boundary(operation):
    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
    user_inputs = []
    response_ids = []
    for transcript in ("first input", "second input"):
        user_inputs.append(session.commit_audio_input(transcript=transcript).message)
        response_id = session.begin_response()
        response_ids.append(response_id)
        session.reserve_history_item(f"item_{response_id}")
        session.append_assistant_text("unplayed answer")
        session.mark_audio_sent(duration_ms=10_000, text_requires_complete_audio=True, audio_complete=True)
        session.end_response(commit_text=True)

    item_id = f"item_{response_ids[1]}"
    if operation == "delete":
        assert session.delete_history_item(item_id) is True
    else:
        # A full ACK releases the pending snapshot. Truncation then edits the
        # stored message into another empty dict before removing that item.
        ack = apply_playback_ack(session, {"response_id": response_ids[1], "played_ms": 10_000})
        assert ack[0].to_realtime()["event"]["history_committed"] is True
        assert session.truncate_history_item(item_id, audio_end_ms=0, hard=True) is True
    assert session.history == (user_inputs[0], {"role": "assistant", "content": ""}, user_inputs[1])
    ack = apply_playback_ack(session, {"response_id": response_ids[1], "played_ms": 10_000})
    assert ack[0].to_realtime()["error"]["code"] == "playback_item_not_found"
    assert session.history == (user_inputs[0], {"role": "assistant", "content": ""}, user_inputs[1])


def test_truncate_history_item_uses_response_alignment_marks():
    session = _session()
    response_id = session.begin_response()
    session.append_assistant_text("hello ")
    session.mark_audio_sent(duration_ms=1_000, text_chars=6)
    session.append_assistant_text("world")
    session.mark_audio_sent(duration_ms=2_000, text_chars=11)
    session.acknowledge_playback(played_ms=2_000)
    committed = session.end_response(commit_text=True)
    assert committed == {"role": "assistant", "content": "hello world"}
    item_id = f"item_{response_id}"
    session.register_history_item(item_id, committed)

    assert session.truncate_history_item(item_id, audio_end_ms=1_000) is True

    assert session.history[-1] == {"role": "assistant", "content": "hello"}
    assert session.history_item_ids[item_id] == {"role": "assistant", "content": "hello"}
    assert session.delete_history_item(item_id) is True
    assert session.history == ()


# ---- turn signals / model turns ----


@pytest.mark.parametrize(
    ("event_type", "turn_state"),
    [
        (DuplexTurnEventType.USER_STARTED, DuplexTurnState.USER_SPEAKING),
        (DuplexTurnEventType.USER_COMMITTED, DuplexTurnState.USER_COMMITTED),
        (DuplexTurnEventType.ASSISTANT_STARTED, DuplexTurnState.ASSISTANT_GENERATING),
        (DuplexTurnEventType.ASSISTANT_DONE, DuplexTurnState.IDLE),
        (DuplexTurnEventType.BARGE_IN, DuplexTurnState.BARGE_IN),
    ],
)
def test_signal_turn_transitions_and_returns_typed_turn_event(event_type, turn_state):
    session = _session()

    event = session.signal_turn(event_type.value)

    assert isinstance(event, TurnEvent)
    assert session.turn_state == turn_state
    wire = event.to_realtime()
    assert wire["type"] == "turn.event"
    assert wire["event"] == event_type.value
    assert wire["turn_state"] == turn_state.value


def test_signal_turn_playback_ack_updates_cursor_and_close_marks_session_closing():
    session = _session()
    session.mark_audio_sent(duration_ms=5_000)

    ack = session.signal_turn("playback_ack", {"played_ms": 1_500, "committed_ms": 1_000})
    assert ack.event == "playback_ack"
    assert session.playback.played_ms == 1_500
    assert session.playback.committed_ms == 1_000

    close = session.signal_turn(DuplexTurnEventType.CLOSE.value)
    assert close.event == "close"
    assert session.state.value == "closing"


def test_complete_model_turn_advances_turn_id_and_fence():
    session = _session()
    assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=0)

    session.complete_model_turn(0)

    assert session.turn_id == 1
    assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=1)
    assert session.accepted_fence == session.fence

    # A stale terminal for an already-completed turn does not move identity backwards.
    session.complete_model_turn(0)
    assert session.turn_id == 1
    session.complete_model_turn(3)
    assert session.turn_id == 4
    assert session.accepted_fence.turn_id == 4


# ---- fence validation ----


def test_accept_fence_rejects_stale_or_foreign_fences():
    session = _session("sid-fence")
    session.complete_model_turn(0)

    with pytest.raises(DuplexFenceMismatchError):
        session.accept_fence(DuplexFence("sid-fence", epoch=0, turn_id=0))
    with pytest.raises(DuplexFenceMismatchError):
        session.accept_fence(DuplexFence("sid-other", epoch=0, turn_id=1))

    session.accept_fence(DuplexFence("sid-fence", epoch=1, turn_id=0))
    assert session.accepted_fence == DuplexFence("sid-fence", epoch=1, turn_id=0)


def test_prepare_and_commit_append_sequence_chunks_per_turn_and_epoch():
    session = _session("sid-append")
    fence = session.fence

    first = session.commit_append(session.prepare_append(fence))
    second = session.commit_append(session.prepare_append(fence))
    assert (first.seq, first.turn_seq, first.turn_id) == (1, 1, 0)
    assert (second.seq, second.turn_seq, second.turn_id) == (2, 2, 0)

    session.complete_model_turn(0)
    third = session.commit_append(session.prepare_append(session.fence))
    assert (third.seq, third.turn_seq, third.turn_id) == (3, 1, 1)

    # Barge-in advances the epoch (turn_id is kept) and restarts the append sequence.
    session.barge_in()
    assert session.fence == DuplexFence("sid-append", epoch=1, turn_id=1)
    fourth = session.commit_append(session.prepare_append(session.fence))
    assert (fourth.seq, fourth.turn_seq, fourth.turn_id) == (1, 1, 1)
    assert session.input_seq == 1


def test_commit_append_rejects_stale_reservation():
    session = _session("sid-append-stale")
    fence = session.fence
    reservation = session.prepare_append(fence)
    session.commit_append(session.prepare_append(fence))

    with pytest.raises(RuntimeError, match="stale"):
        session.commit_append(reservation)

    with pytest.raises(DuplexFenceMismatchError):
        session.prepare_append(DuplexFence("sid-other", epoch=0, turn_id=0))


def test_cancel_fence_releases_stage_requests_and_advances_identity():
    session = _session("sid-cancel")
    cancelled = session.fence
    session.reserve_stage_request(0, "req-a", fence=cancelled)
    session.bind_stage_request(1, "req-b", fence=cancelled)
    assert session.stage_request_submitted(1, "req-b") is True
    assert session.stage_request_submitted(0, "req-a") is False
    next_fence = DuplexFence("sid-cancel", epoch=1, turn_id=0)

    stale = session.cancel_fence(cancelled, next_fence)

    assert stale == ["req-a", "req-b"]
    assert session.resource_request_ids() == []
    assert session.accepted_fence == next_fence

    with pytest.raises(DuplexFenceMismatchError):
        session.cancel_fence(next_fence, DuplexFence("sid-cancel", epoch=1, turn_id=1))
    with pytest.raises(DuplexFenceMismatchError):
        session.cancel_fence(DuplexFence("sid-other", epoch=1, turn_id=0), DuplexFence("sid-cancel", epoch=2))


def test_request_resource_keys_are_stage_id_and_request_id():
    session = _session("sid-keys")
    session.bind_stage_request(0, "req-a", fence=session.fence)
    session.bind_stage_request(1, "req-b", fence=session.fence)
    session.bind_stage_request(2, "req-tts", fence=session.fence)
    stale_keys = list(session.request_resources.keys())
    stale_ids = list(dict.fromkeys(rid for _, rid in stale_keys))
    assert stale_ids == ["req-a", "req-b", "req-tts"]
    for sid, rid in stale_keys:
        if sid < 2:
            session.request_resources.pop((sid, rid), None)
    assert session.resource_request_ids() == ["req-tts"]
    with pytest.raises(TypeError, match="unhashable"):
        dict.fromkeys(rid for _, rid in session.request_resources.items())


def test_release_resources_for_request_ids_keeps_other_ids_on_the_same_fence():
    session = _session("sid-drain-release")
    fence = session.fence
    session.bind_stage_request(2, "drain", fence=fence)
    session.bind_stage_request(3, "drain", fence=fence)
    session.bind_stage_request(0, "live", fence=fence)

    released = session.release_resources_for_request_ids(["drain"])

    assert released == ["drain"]
    assert session.resource_request_ids() == ["live"]
    assert session.release_resources_for_request_ids([]) == []


# ---- public view ----


def test_as_public_dict_exposes_identity_capabilities_and_playback():
    session = _session("duplex-public", config=DuplexSessionConfig(model="m", voice="alloy"))
    session.mark_audio_sent(duration_ms=100)

    payload = session.as_public_dict()

    assert payload["id"] == "duplex-public"
    assert payload["model"] == "m"
    assert payload["voice"] == "alloy"
    assert payload["state"] == "open"
    assert payload["turn_state"] == "assistant_playing"
    assert payload["epoch"] == 0
    assert payload["turn_id"] == 0
    assert payload["capabilities"] == DuplexCapabilities().as_dict()
    assert payload["playback"] == {"generated_ms": 100, "sent_ms": 100, "played_ms": 0, "committed_ms": 0}


# ---- DuplexSessionConfig.from_realtime ----


def test_from_realtime_rejects_unsupported_audio_formats():
    with pytest.raises(DuplexConfigError) as excinfo:
        DuplexSessionConfig.from_realtime({"input_audio_format": "mp3"})
    assert excinfo.value.code == "unsupported_audio_format"

    with pytest.raises(DuplexConfigError) as excinfo:
        DuplexSessionConfig.from_realtime({"audio": {"output": {"format": {"type": "audio/opus"}}}})
    assert excinfo.value.code == "unsupported_audio_format"


def test_from_realtime_validates_turn_detection():
    with pytest.raises(DuplexConfigError) as excinfo:
        DuplexSessionConfig.from_realtime({"turn_detection": {"type": "semantic_vad"}})
    assert excinfo.value.code == "unsupported_turn_detection"
    assert excinfo.value.param == "turn_detection"

    with pytest.raises(DuplexConfigError) as excinfo:
        DuplexSessionConfig.from_realtime({"overlap_policy": "barge_in_on_speech"})
    assert excinfo.value.code == "unsupported_turn_detection"

    server_vad = DuplexSessionConfig.from_realtime({"turn_detection": {"type": "server_vad"}})
    assert server_vad.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value
    assert server_vad.extra_body["realtime_turn_detection"]["threshold"] == 0.5

    model_owned = DuplexSessionConfig.from_realtime({"turn_detection": None})
    assert model_owned.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
    assert model_owned.extra_body["realtime_turn_detection"] is None


def test_from_realtime_maps_wire_fields_and_stores_realtime_keys_in_extra_body():
    config = DuplexSessionConfig.from_realtime(
        {
            "model": "openbmb/MiniCPM-o-4_5",
            "instructions": "be brief",
            "output_audio_format": "pcm16",
            "max_response_output_tokens": "inf",
            "audio": {"output": {"voice": "alloy", "speed": 1.25}},
            "tools": [{"type": "function", "name": "lookup"}],
            "tool_choice": "auto",
            "metadata": {"tenant": "t1"},
            "include": ["item.input_audio_transcription.logprobs"],
            "input_audio_transcription": {"model": "whisper-1"},
            "extra_body": {"custom": True},
        }
    )

    assert config.model == "openbmb/MiniCPM-o-4_5"
    assert config.instructions == "be brief"
    assert config.voice == "alloy"
    assert config.speed == 1.25
    assert config.max_tokens is None
    assert config.response_format == "pcm"
    assert config.modalities == ["text", "audio"]
    assert config.extra_body["custom"] is True
    assert config.extra_body["realtime_tools"] == [{"type": "function", "name": "lookup"}]
    assert config.extra_body["realtime_tool_choice"] == "auto"
    assert config.extra_body["realtime_metadata"] == {"tenant": "t1"}
    assert config.extra_body["realtime_include"] == ["item.input_audio_transcription.logprobs"]
    assert config.extra_body["realtime_input_audio_transcription"] == {"model": "whisper-1"}
    assert config.extra_body["realtime_audio"] == {"output": {"voice": "alloy", "speed": 1.25}}
    assert config.extra_body["realtime_output_audio_format"] == "pcm16"
    assert "extra_body" not in config.extra_body["realtime_session_payload"]
    assert config.extra_body["realtime_session_payload"]["model"] == "openbmb/MiniCPM-o-4_5"


def test_from_realtime_uses_served_model_when_payload_has_none():
    config = DuplexSessionConfig.from_realtime({}, model="served-model")

    assert config.model == "served-model"
    assert config.response_format == "pcm"
    assert config.idle_timeout_s == 300.0


def test_from_realtime_ignores_client_chosen_session_ids():
    config = DuplexSessionConfig.from_realtime({"id": "client-chosen", "session_id": "also-client-chosen"})

    assert not hasattr(config, "session_id")
    assert "id" not in config.as_dict()
    session = _session("duplex-engine-allocated", config=config)
    assert session.as_public_dict()["id"] == "duplex-engine-allocated"


# ---- DuplexSessionConfig.apply_realtime_update ----


def test_apply_realtime_update_rejects_model_change():
    config = DuplexSessionConfig(model="served-model")

    with pytest.raises(DuplexConfigError) as excinfo:
        config.apply_realtime_update({"model": "other-model"}, session_id="sid")

    assert excinfo.value.code == "model_update_unsupported"
    assert config.model == "served-model"


def test_apply_realtime_update_rejects_voice_change_after_audio_started():
    config = DuplexSessionConfig(model="m", voice="alloy")

    config.apply_realtime_update({"voice": "verse"}, audio_started=False)
    assert config.voice == "verse"
    with pytest.raises(DuplexConfigError) as excinfo:
        config.apply_realtime_update({"audio": {"output": {"voice": "alloy"}}}, audio_started=True)
    assert excinfo.value.code == "voice_update_after_audio_unsupported"
    assert config.voice == "verse"


def test_apply_realtime_update_rejects_ref_audio_change():
    config = DuplexSessionConfig(model="m")

    with pytest.raises(DuplexConfigError) as excinfo:
        config.apply_realtime_update({"ref_audio": "/tmp/voice.wav"})

    assert excinfo.value.code == "ref_audio_update_unsupported"
    assert config.ref_audio is None


def test_apply_realtime_update_patches_fields_and_realtime_keys():
    config = DuplexSessionConfig(model="m", instructions="old", extra_body={"realtime_tools": [{"name": "x"}]})

    config.apply_realtime_update(
        {
            "instructions": "new",
            "output_audio_format": "g711_ulaw",
            "max_output_tokens": 32,
            "tools": None,
            "metadata": {"k": "v"},
            "overlap_policy": "not-a-policy",
        }
    )

    assert config.instructions == "new"
    assert config.response_format == "pcm"
    assert config.max_tokens == 32
    assert "realtime_tools" not in config.extra_body
    assert config.extra_body["realtime_metadata"] == {"k": "v"}
    assert config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
    assert config.extra_body["realtime_session_payload"]["instructions"] == "new"


# ---- overlap policy / capabilities ----


def test_overlap_policy_defaults_and_invalid_values_to_listen_only():
    assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
    assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value
    assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value
    assert (
        DuplexSessionConfig._normalize_overlap_policy(" Barge_In_On_Speech ")
        == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value
    )
    assert DuplexSessionConfig.from_event({"session": {"overlap_policy": "bogus"}}).overlap_policy == "listen_only"


def test_capabilities_as_dict_reports_declared_implementation():
    caps = DuplexCapabilities().as_dict()

    assert caps["implementation_level"] == "turn_based_duplex"
    assert caps["input_modes"] == ["append_audio_chunk"]
    assert caps["supports_kv_lease"] is False
    assert caps["supports_core_kv_lease"] is False
    assert "client_event" in caps["signal_sources"]
    assert set(caps) == set(minicpmo45_native_capabilities().as_dict())


def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
    caps = minicpmo45_native_capabilities(max_sessions=2).as_dict()

    assert caps["implementation_level"] == "model_native_duplex"
    assert caps["supports_input_append"] is True
    assert caps["input_modes"] == ["append_audio_chunk"]
    assert caps["adapter_patterns"] == ["scheduler_data_plane"]
    assert caps["supports_model_internal_state"] is True
    assert caps["requires_model_runner_kv"] is True
    assert caps["requires_native_stage_role"] is True
    assert caps["supports_kv_lease"] is False
    assert caps["supports_core_kv_lease"] is False
    assert caps["supports_stage_resumption"] is True
    assert caps["supports_scheduler_native_append"] is False
    assert caps["supports_core_resumable_request"] is True
    assert caps["supports_stage_connector_handoff"] is True
    assert caps["supports_audio_truncate"] is True
    assert caps["supports_barge_in"] is True
    assert caps["target_barge_in_latency_ms"] is None
    assert caps["supports_multi_session"] is True
    assert caps["supports_multi_session_same_replica"] is True
    assert caps["supports_session_lease"] is True
    assert caps["supports_session_resume"] is True
    assert caps["session_admission_mode"] == "engine_managed"
    assert caps["stage_handoff_transport"] == "scheduler_data_plane"


def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
    caps = minicpmo45_native_capabilities(max_sessions=1).as_dict()

    assert caps["supports_multi_session"] is False
    assert caps["supports_multi_session_same_replica"] is False


def test_response_timing_binds_latest_request_start_for_model_turn():
    session = _session()
    session.mark_model_turn_request_started(0, 10.0)
    session.mark_model_turn_request_started(0, 11.0)
    session.begin_response(turn_id=0)

    first = session.mark_response_first_outputs(
        observed_at_s=11.2,
        has_text=True,
        has_audio=False,
    )
    assert first["ttft_ms"] == pytest.approx(200.0)
    assert first["measurement_origin"] == RESPONSE_REQUEST_MEASUREMENT_ORIGIN
    assert "ttfp_ms" not in first

    # After begin_response the origin is frozen at the pending start (11.0).
    # The 12.0 append must not rebind, so TTFP is 12.3-11.0 = 1300 ms, not 300.
    session.mark_model_turn_request_started(0, 12.0)
    audio = session.mark_response_first_outputs(
        observed_at_s=12.3,
        has_text=False,
        has_audio=True,
    )
    assert audio["ttft_ms"] == pytest.approx(200.0)
    assert audio["ttfp_ms"] == pytest.approx(1300.0)


def test_response_timing_attaches_metrics_only_when_newly_observed():
    session = _session()
    session.mark_model_turn_request_started(0, 10.0)
    session.begin_response(turn_id=0)

    first = session.mark_response_first_outputs(observed_at_s=10.2, has_text=True, has_audio=False)
    assert first["ttft_ms"] == pytest.approx(200.0)
    assert session.mark_response_first_outputs(observed_at_s=10.3, has_text=True, has_audio=False) == {}
    audio = session.mark_response_first_outputs(observed_at_s=10.4, has_text=False, has_audio=True)
    assert audio["ttft_ms"] == pytest.approx(200.0)
    assert audio["ttfp_ms"] == pytest.approx(400.0)
    assert session.mark_response_first_outputs(observed_at_s=10.5, has_text=True, has_audio=True) == {}


def test_response_timing_is_cleared_on_end_barge_in_and_close():
    session = _session()
    session.mark_model_turn_request_started(0, 10.0)
    session.begin_response(turn_id=0)
    assert session.mark_response_first_outputs(observed_at_s=10.2, has_text=True, has_audio=False)
    session.end_response()
    assert session.mark_response_first_outputs(observed_at_s=10.4, has_text=True, has_audio=False) == {}

    session.mark_model_turn_request_started(0, 11.0)
    session.barge_in()
    session.begin_response(turn_id=0)
    assert session.mark_response_first_outputs(observed_at_s=11.2, has_text=True, has_audio=False) == {}

    session.mark_model_turn_request_started(0, 12.0)
    session.close()
    session.begin_response(turn_id=0)
    assert session.mark_response_first_outputs(observed_at_s=12.2, has_text=True, has_audio=False) == {}


def test_complete_model_turn_drops_request_starts_for_finished_turns():
    session = _session()
    session.mark_model_turn_request_started(0, 10.0)
    session.mark_model_turn_request_started(1, 11.0)
    session.complete_model_turn(0)
    session.begin_response(turn_id=1)
    assert session.mark_response_first_outputs(
        observed_at_s=11.25,
        has_text=True,
        has_audio=False,
    )["ttft_ms"] == pytest.approx(250.0)
    session.end_response()
    session.begin_response(turn_id=0)
    assert session.mark_response_first_outputs(observed_at_s=11.5, has_text=True, has_audio=False) == {}


def test_log_stats_off_does_not_open_a_response_aggregator():
    session = _session()
    session.begin_response()
    session.observe_stage_request_stats(0, _stage_stats(stage_id=0))
    assert session._response_aggregator is None
    session.end_response()
    assert session._response_aggregator is None


def test_end_response_logs_pending_and_active_stage_request_stats(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(OrchestratorAggregator, "build_and_log_summary", lambda self: {})
    session = _session(num_stages=2, log_stats=True)
    wire_stats = _stage_stats(stage_id=0, request_id="stage0-req", num_tokens_out=3)
    session.observe_stage_request_stats(0, wire_stats)
    assert session._response_aggregator is None
    assert wire_stats.request_id == "stage0-req"

    response_id = session.begin_response()
    aggregator = session._response_aggregator
    assert aggregator is not None
    assert aggregator.num_stages == 2
    recorded = aggregator.stage_events[response_id]
    assert recorded[0].num_tokens_out == 3
    assert recorded[0].request_id == response_id
    assert wire_stats.request_id == "stage0-req"

    session.observe_stage_request_stats(1, _stage_stats(stage_id=1, num_tokens_out=5, vllm_ttft_ms=0.0))
    assert [event.stage_id for event in aggregator.stage_events[response_id]] == [0, 1]

    session.end_response()
    assert session._response_aggregator is None
    assert response_id in aggregator.e2e_done


def test_end_response_prints_one_column_per_stage(monkeypatch: pytest.MonkeyPatch):
    logged: list[OrchestratorAggregator] = []
    monkeypatch.setattr(OrchestratorAggregator, "build_and_log_summary", lambda self: logged.append(self) or {})
    session = _session(num_stages=3, log_stats=True)
    session.observe_stage_request_stats(0, _stage_stats(stage_id=0, num_tokens_out=3, vllm_ttft_ms=40.0))
    response_id = session.begin_response()
    session.observe_stage_request_stats(0, _stage_stats(stage_id=0, num_tokens_out=4, vllm_ttft_ms=12.0))
    session.observe_stage_request_stats(1, _stage_stats(stage_id=1, num_tokens_out=5, vllm_ttft_ms=20.0))
    session.observe_stage_request_stats(1, _stage_stats(stage_id=1, num_tokens_out=6, vllm_ttft_ms=0.0))
    session.observe_stage_request_stats(2, _stage_stats(stage_id=2, num_tokens_out=0, vllm_ttft_ms=0.0))
    session.end_response()

    assert len(logged) == 1
    rows = logged[0].stage_events[response_id]
    assert [event.stage_id for event in rows] == [0, 1, 2]
    assert rows[0].num_tokens_out == 7
    assert rows[0].vllm_ttft_ms == pytest.approx(40.0)
    assert rows[1].num_tokens_out == 11
    assert rows[1].vllm_ttft_ms == pytest.approx(20.0)
    assert logged[0].stage_table_exclude == DUPLEX_STAGE_TABLE_EXCLUDE


def test_end_response_omits_serving_time_to_first_output_from_stage_table(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    original = OrchestratorAggregator.build_and_log_summary
    captured: list[tuple[OrchestratorAggregator, dict[str, object]]] = []

    def _capture(self: OrchestratorAggregator) -> dict[str, object]:
        summary = original(self)
        captured.append((self, summary))
        return summary

    monkeypatch.setattr(OrchestratorAggregator, "build_and_log_summary", _capture)
    session = _session(num_stages=3, log_stats=True)
    response_id = session.begin_response()
    session.observe_stage_request_stats(0, _stage_stats(stage_id=0, serving_time_to_first_output_ms=80.631))
    session.observe_stage_request_stats(1, _stage_stats(stage_id=1, serving_time_to_first_output_ms=418.829))
    session.observe_stage_request_stats(
        2, _stage_stats(stage_id=2, vllm_ttft_ms=0.0, serving_time_to_first_output_ms=487.109)
    )
    session.end_response()

    assert len(captured) == 1
    aggregator, summary = captured[0]
    stage_table = summary.get("stage_table")
    assert isinstance(stage_table, list) and stage_table
    rows = stage_table[0]["stages"]
    assert isinstance(rows, list)
    assert [row["stage_id"] for row in rows] == [0, 1, 2]
    assert all("serving_time_to_first_output_ms" not in row for row in rows)
    assert all("vllm_ttft_ms" in row for row in rows)
    events = aggregator.stage_events[response_id]
    assert [event.serving_time_to_first_output_ms for event in events] == pytest.approx([80.631, 418.829, 487.109])


def test_http_stage_table_keeps_serving_time_to_first_output() -> None:
    agg = OrchestratorAggregator(num_stages=2, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=1)
    agg.on_stage_metrics(0, "r1", _stage_stats(stage_id=0, serving_time_to_first_output_ms=80.0))
    agg.on_stage_metrics(1, "r1", _stage_stats(stage_id=1, serving_time_to_first_output_ms=418.0))
    agg.on_finalize_request(1, "r1", req_start_ts=0.0)

    summary = agg.build_and_log_summary()
    rows = summary["stage_table"][0]["stages"]
    assert [row["serving_time_to_first_output_ms"] for row in rows] == [80.0, 418.0]


def test_logged_e2e_includes_wait_before_first_output(monkeypatch: pytest.MonkeyPatch) -> None:
    mono = {"t": 100.0}
    wall = {"t": 1_000.0}
    monkeypatch.setattr(time, "time", lambda: wall["t"])
    logged: list[OrchestratorAggregator] = []
    monkeypatch.setattr(OrchestratorAggregator, "build_and_log_summary", lambda self: logged.append(self) or {})
    session = _session(log_stats=True, clock=lambda: mono["t"])
    session.mark_model_turn_request_started(0, 100.0)

    mono["t"] = 102.0
    wall["t"] = 1_002.0
    response_id = session.begin_response(turn_id=0)
    first = session.mark_response_first_outputs(observed_at_s=102.0, has_text=True, has_audio=True)
    assert first["ttft_ms"] == pytest.approx(2000.0)

    mono["t"] = 102.1
    wall["t"] = 1_002.1
    session.end_response()

    assert logged[0].e2e_events[0].request_id == response_id
    assert logged[0].e2e_events[0].e2e_total_ms == pytest.approx(2100.0)


def test_logged_tpot_matches_client_weighted_aggregation(monkeypatch: pytest.MonkeyPatch) -> None:
    logged: list[OrchestratorAggregator] = []
    monkeypatch.setattr(OrchestratorAggregator, "build_and_log_summary", lambda self: logged.append(self) or {})
    session = _session(log_stats=True)
    response_id = session.begin_response()
    session.observe_stage_request_stats(0, _stage_stats(stage_id=0, num_tokens_out=11, vllm_tpot_ms=10.0))
    session.observe_stage_request_stats(0, _stage_stats(stage_id=0, num_tokens_out=3, vllm_tpot_ms=100.0))
    session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 11, "vllm_tpot_ms": 10.0}})
    client = session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 3, "vllm_tpot_ms": 100.0}})
    expected = client["0"]["vllm_tpot_ms"]
    session.end_response()

    assert expected == pytest.approx(25.0)
    assert logged[0].stage_events[response_id][0].vllm_tpot_ms == pytest.approx(expected)


def test_draining_response_stays_ack_admissible_after_next_begin_response() -> None:
    from vllm_omni.engine.duplex.events import ErrorEvent
    from vllm_omni.engine.duplex.session.playback_ledger import apply_playback_ack

    session = _session()
    first = session.begin_response(turn_id=1)
    session.append_assistant_text("hello")
    session.mark_audio_sent(400, text_chars=5)
    session.snapshot_active_response_for_drain()
    second = session.begin_response(turn_id=2)
    assert second != first
    assert session.playback.sent_ms == 0
    events = apply_playback_ack(
        session,
        {
            "type": "playback.ack",
            "response_id": first,
            "item_id": f"item_{first}",
            "played_ms": 100,
            "committed_ms": 100,
        },
    )
    assert not any(isinstance(event, ErrorEvent) and event.code == "playback_item_not_found" for event in events)
    session.mark_audio_sent(900, text_chars=5, response_id=first)
    assert session.playback_for_response(first).sent_ms == 900
    assert session.playback.sent_ms == 0


def test_finished_drain_keeps_sent_audio_ackable() -> None:
    from vllm_omni.engine.duplex.events import ErrorEvent
    from vllm_omni.engine.duplex.session.playback_ledger import apply_playback_ack

    session = _session()
    first = session.begin_response(turn_id=1)
    session.append_assistant_text("hello")
    session.mark_audio_sent(400, text_chars=5)
    session.snapshot_active_response_for_drain()
    second = session.begin_response(turn_id=2)
    session.release_finished_drain_response(first)
    assert first in session._conversation.assistant_response_snapshots
    assert f"item_{first}" in session._conversation.history_item_placeholders
    events = apply_playback_ack(
        session,
        {
            "type": "playback.ack",
            "response_id": first,
            "item_id": f"item_{first}",
            "played_ms": 400,
            "committed_ms": 400,
        },
    )
    assert not any(isinstance(event, ErrorEvent) for event in events)
    assert session.active_response_id == second

    silent = session.begin_response(turn_id=3)
    session.snapshot_active_response_for_drain()
    session.begin_response(turn_id=4)
    session.release_finished_drain_response(silent)
    assert silent not in session._conversation.assistant_response_snapshots
    assert f"item_{silent}" not in session._conversation.history_item_placeholders