vllm-omni/tests/worker/test_omni_connector_mixin.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for OmniConnectorModelRunnerMixin.

These tests use a mock connector (in-memory dict store) and do not require
GPU or vLLM runtime.
"""

from __future__ import annotations

import time
import unittest
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock, patch

import pytest
import torch

from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
    OmniKVTransferManager,
)
from vllm_omni.outputs import OmniConnectorOutput
from vllm_omni.worker.omni_connector_model_runner_mixin import (
    OmniConnectorModelRunnerMixin,
)

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

# ------------------------------------------------------------------ #
#  Mock helpers
# ------------------------------------------------------------------ #


class MockConnector:
    """In-memory connector for testing (mimics OmniConnectorBase)."""

    def __init__(self, stage_id: int = 0):
        self.stage_id = stage_id
        self._store: dict[str, Any] = {}

    def put(self, from_stage, to_stage, put_key, data):
        key = f"{from_stage}_{to_stage}_{put_key}"
        self._store[key] = data
        return True, len(str(data)), None

    def get(self, from_stage, to_stage, get_key, metadata=None):
        key = f"{from_stage}_{to_stage}_{get_key}"
        data = self._store.pop(key, None)
        if data is None:
            return None
        return data, len(str(data))

    def close(self):
        pass


def _make_model_config(
    stage_id: int = 0,
    async_chunk: bool = False,
    worker_type: str = "ar",
    custom_func: str | None = None,
) -> SimpleNamespace:
    return SimpleNamespace(
        stage_connector_config=None,
        async_chunk=async_chunk,
        worker_type=worker_type,
        custom_process_next_stage_input_func=custom_func,
    )


def _make_request(req_id: str, external_req_id: str | None = None):
    r = SimpleNamespace(
        request_id=req_id,
        external_req_id=external_req_id or req_id,
        additional_information=None,
        prompt_token_ids=[],
        num_computed_tokens=0,
    )
    return r


class MixinHost(OmniConnectorModelRunnerMixin):
    """Minimal class that mixes in the mixin for testing."""

    pass


class _FakeTPGroup:
    def __init__(self, *, world_size: int, rank_in_group: int, follower_result: Any = None):
        self.world_size = world_size
        self.rank_in_group = rank_in_group
        self.follower_result = follower_result
        self.broadcast_inputs: list[Any] = []

    def broadcast_object(self, obj: Any | None = None, src: int = 0):
        self.broadcast_inputs.append(obj)
        if self.rank_in_group == src:
            return obj
        return self.follower_result


# ------------------------------------------------------------------ #
#  Test cases
# ------------------------------------------------------------------ #


@pytest.mark.parametrize(
    ("role", "custom_func", "expected"),
    [
        ("sender", None, False),
        ("sender", "package.build_payload", True),
        ("receiver", None, True),
    ],
)
def test_init_payload_connector_ownership(role, custom_func, expected):
    model_config = _make_model_config(custom_func=custom_func)
    model_config.stage_connector_config = {
        "name": "MooncakeTransferEngineConnector",
        "extra": {"role": role},
    }

    host = MixinHost()
    connector = MockConnector()
    with (
        patch.object(host, "_create_connector", return_value=connector) as create,
        patch.object(host, "_load_custom_func", return_value=(None, None)),
    ):
        host.init_omni_connectors(model_config)

    assert (create.call_count == 1) is expected
    assert (host._omni_connector is connector) is expected
    host.shutdown_omni_connectors()


class TestMixinAsyncChunkSendRecv(unittest.TestCase):
    """Test 2: Async chunk send/recv + bg threads."""

    def test_send_chunk_passes_is_finished_and_connector(self):
        connector = MockConnector(stage_id=0)

        sender = MixinHost()
        sender.init_omni_connectors(
            model_config=_make_model_config(stage_id=0, async_chunk=True),
        )
        sender._omni_connector = connector
        sender._stage_id = 0
        sender._async_chunk = True

        seen = {}

        def mock_process(transfer_manager, pooling_output, request, is_finished=False):
            seen["connector"] = transfer_manager.connector
            seen["is_finished"] = is_finished
            return {"data": pooling_output, "finished": is_finished}

        sender._custom_process_func = mock_process

        request = _make_request("req-1", "ext-req-1")
        request.is_finished = lambda: True
        sender._send_single_request(
            {
                "stage_id": 0,
                "next_stage_id": 1,
                "request_id": "ext-req-1",
                "request": request,
                "pooling_output": {"value": 42},
            }
        )
        self.assertIs(seen["connector"], connector)
        self.assertTrue(seen["is_finished"])

        sender.shutdown_omni_connectors()

    def test_send_chunk_does_not_retry_real_type_error(self):
        connector = MockConnector(stage_id=0)

        sender = MixinHost()
        sender.init_omni_connectors(
            model_config=_make_model_config(stage_id=0, async_chunk=True),
        )
        sender._omni_connector = connector
        sender._stage_id = 0
        sender._async_chunk = True

        seen = {"calls": 0}

        def broken_process(transfer_manager, pooling_output, request, is_finished=""):
            seen["calls"] += 1
            return {"data": is_finished + "tail"}

        sender._custom_process_func = broken_process

        request = _make_request("req-1", "ext-req-1")
        request.is_finished = lambda: True
        ok = sender.send_chunk(request, pooling_output={"value": 42})
        self.assertFalse(ok)
        self.assertEqual(seen["calls"], 1)

        sender.shutdown_omni_connectors()


class TestMixinKVCacheTransfer(unittest.TestCase):
    """Test 3: KV cache delegation to OmniKVTransferManager."""

    def test_send_kv_delegates(self):
        mock_kvm = MagicMock()
        mock_kvm.handle_finished_requests_kv_transfer.return_value = ["req-1"]

        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
            kv_transfer_manager=mock_kvm,
        )

        result = host.send_kv_cache(
            finished_reqs={"req-1": {"seq_len": 10, "block_ids": [0]}},
            kv_caches=[],
            block_size=16,
            cache_dtype="float16",
        )
        self.assertEqual(result, ["req-1"])
        mock_kvm.handle_finished_requests_kv_transfer.assert_called_once()

        host.shutdown_omni_connectors()

    def test_recv_kv_delegates(self):
        mock_kvm = MagicMock()
        mock_kvm.receive_kv_cache_for_request.return_value = ({"layer_blocks": {}}, 100)

        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
            kv_transfer_manager=mock_kvm,
        )

        data, size = host.recv_kv_cache("req-1")
        self.assertIsNotNone(data)
        self.assertEqual(size, 100)
        mock_kvm.receive_kv_cache_for_request.assert_called_once()

        host.shutdown_omni_connectors()

    def test_receive_multi_kv_fetches_companions_via_mixin(self):
        mock_kvm = MagicMock()

        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
            kv_transfer_manager=mock_kvm,
        )

        host.recv_kv_cache = MagicMock(
            side_effect=[({"layer_blocks": {"k": [1]}}, 64), ({"layer_blocks": {"k": [2]}}, 32)]
        )
        seen = {}

        def collect_cfg(request_id, cfg_role_payloads):
            seen["request_id"] = request_id
            seen["cfg_role_payloads"] = cfg_role_payloads
            return {"cfg_text_kv_metadata": {"seq_len": 3}}

        req = SimpleNamespace(
            request_id="req-1",
            sampling_params=SimpleNamespace(cfg_kv_request_ids={"cfg_text": "req-1__cfg_text"}),
        )
        ok = host.receive_multi_kv_cache(req, cfg_kv_collect_func=collect_cfg)
        self.assertTrue(ok)
        host.recv_kv_cache.assert_any_call("req-1", target_device=None)
        host.recv_kv_cache.assert_any_call("req-1__cfg_text", target_device=None)
        mock_kvm.apply_kv_cache_to_request.assert_called_once_with(req, {"layer_blocks": {"k": [1]}})
        self.assertEqual(seen["request_id"], "req-1")
        self.assertEqual(
            seen["cfg_role_payloads"],
            {"cfg_text": ({"layer_blocks": {"k": [2]}}, 32)},
        )
        self.assertEqual(req.sampling_params.cfg_text_kv_metadata, {"seq_len": 3})

        host.shutdown_omni_connectors()

    def test_receive_multi_kv_skips_inactive_request(self):
        mock_kvm = MagicMock()

        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
            kv_transfer_manager=mock_kvm,
        )

        host.requests = {}
        host.recv_kv_cache = MagicMock(return_value=({"layer_blocks": {"k": [1]}}, 64))
        req = SimpleNamespace(request_id="req-1", sampling_params=None)

        ok = host.receive_multi_kv_cache(req)

        self.assertFalse(ok)
        host.recv_kv_cache.assert_not_called()
        mock_kvm.apply_kv_cache_to_request.assert_not_called()

        host.shutdown_omni_connectors()


class TestOmniConnectorOutput(unittest.TestCase):
    """Test 4: Output aggregation across transfer modes."""

    def test_output_aggregation(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
        )

        host._chunk_ready_req_ids.add("req-1")
        host._chunk_finished_req_ids.add("req-2")
        host._local_request_metadata["req-1"] = {"next_stage_prompt_len": 10}
        host._stage_recv_req_ids.add("req-3")

        output = host.get_omni_connector_output()
        self.assertIsInstance(output, OmniConnectorOutput)
        self.assertEqual(output.chunk_ready_req_ids, {"req-1"})
        self.assertEqual(output.chunk_finished_req_ids, {"req-2"})
        self.assertEqual(output.request_metadata, {"req-1": {"next_stage_prompt_len": 10}})
        self.assertEqual(output.stage_recv_req_ids, {"req-3"})

        output2 = host.get_omni_connector_output()
        self.assertEqual(output2.chunk_ready_req_ids, set())
        self.assertEqual(output2.request_metadata, {})

        host.shutdown_omni_connectors()


class TestMixinNoConnector(unittest.TestCase):
    """Edge case: mixin works gracefully without a connector."""

    def test_no_connector(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
        )
        self.assertIsNone(host._omni_connector)

        results = host.recv_full_payload_inputs(scheduler_output=None)
        self.assertIsNone(results)

        sent = host.send_full_payload_outputs(None, {"req-1": {}})
        self.assertEqual(sent, [])

        ok = host.send_chunk(_make_request("req-1"), pooling_output={})
        self.assertFalse(ok)

        output = host.get_omni_connector_output()
        self.assertIsInstance(output, OmniConnectorOutput)

        host.shutdown_omni_connectors()


class TestFinishedLoadReqsDrain(unittest.TestCase):
    """Test A1 fix: get_omni_connector_output drains _finished_load_reqs."""

    def test_finished_load_reqs_flow_to_chunk_ready(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
        )

        host._finished_load_reqs.add("req-1")
        host._finished_load_reqs.add("req-2")

        output = host.get_omni_connector_output()
        self.assertIn("req-1", output.chunk_ready_req_ids)
        self.assertIn("req-2", output.chunk_ready_req_ids)

        self.assertEqual(len(host._finished_load_reqs), 0)
        self.assertEqual(len(host._chunk_ready_req_ids), 0)

        host.shutdown_omni_connectors()


class TestLoadCustomFuncSelection(unittest.TestCase):
    def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self):
        incompatible_paths = [
            "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav",
            "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit",
            "vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow",
            "vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion",
        ]

        for func_path in incompatible_paths:
            selected_path, func = MixinHost._load_custom_func(
                SimpleNamespace(
                    async_chunk=False,
                    custom_process_input_func=func_path,
                    custom_process_next_stage_input_func=None,
                )
            )
            assert selected_path != func_path
            assert func is None or MixinHost._is_connector_payload_builder(func)


class TestFullPayloadSendWithCustomFunc(unittest.TestCase):
    """Test B4: send_full_payload_outputs with full_payload_mode custom process func."""

    def test_full_payload_send_passes_is_finished_and_connector(self):
        seen = {}

        def full_payload_func(transfer_manager, pooling_output, request, is_finished=False):
            seen["connector"] = transfer_manager.connector
            seen["is_finished"] = is_finished
            seen["data"] = pooling_output
            seen["rid"] = request.request_id if request else None
            return {"processed": True, "finished": is_finished}

        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
        )
        host._omni_connector = MockConnector(stage_id=0)
        host._stage_id = 0
        host._custom_process_func = full_payload_func

        req = _make_request("req-1")
        req.is_finished = lambda: True
        sent = host.send_full_payload_outputs(
            scheduler_output=None,
            outputs={"req-1": ({"raw": 100}, req)},
        )
        self.assertEqual(sent, ["req-1"])
        self.assertEqual(
            seen,
            {
                "connector": host._omni_connector,
                "is_finished": True,
                "data": {"raw": 100},
                "rid": "req-1",
            },
        )

        host.shutdown_omni_connectors()

    def test_accumulate_and_flush(self):
        call_log = []

        def full_payload_func(transfer_manager, pooling_output, request):
            call_log.append(request.request_id if request else None)
            return {"processed": True}

        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
        )
        host._omni_connector = MockConnector(stage_id=0)
        host._stage_id = 0
        host._custom_process_func = full_payload_func

        req = _make_request("req-1")
        host.accumulate_full_payload_output("req-1", {"raw": 42}, req)
        self.assertEqual(len(host._pending_full_payload_send), 1)

        host.flush_full_payload_outputs({"req-1"})
        self.assertEqual(len(host._pending_full_payload_send), 0)
        self.assertEqual(len(call_log), 1)
        self.assertEqual(call_log[0], "req-1")

        time.sleep(0.1)
        host.shutdown_omni_connectors()


class TestKVSentReqIdsAccumulation(unittest.TestCase):
    """Test that kv_sent_req_ids accumulates results from send_kv_cache."""

    def test_kv_sent_accumulation(self):
        mock_kvm = MagicMock()
        mock_kvm.handle_finished_requests_kv_transfer.return_value = ["req-1", "req-2"]

        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(),
            kv_transfer_manager=mock_kvm,
        )

        host.send_kv_cache(
            finished_reqs={"req-1": {}, "req-2": {}},
            kv_caches=[],
            block_size=16,
            cache_dtype="float16",
        )

        output = host.get_omni_connector_output()
        self.assertIn("req-1", output.kv_sent_req_ids)
        self.assertIn("req-2", output.kv_sent_req_ids)

        output2 = host.get_omni_connector_output()
        self.assertEqual(output2.kv_sent_req_ids, [])

        host.shutdown_omni_connectors()


class TestChunkStreamCompletedGuard(unittest.TestCase):
    """Test that register_chunk_recv is skipped after finish sentinel.

    This validates the fix for the race condition where the scheduling
    coordinator re-registers a request for chunk polling after its
    upstream chunk stream has already finished (is_finished sentinel
    received), causing the bg recv thread to poll for a non-existent
    shared-memory segment (e.g. ``_0_7`` when only 7 chunks 0–6 exist).
    """

    def _make_host(self, stage_id: int = 1) -> MixinHost:
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=stage_id, async_chunk=True),
        )
        host._omni_connector = MockConnector(stage_id=stage_id)
        host._stage_id = stage_id
        host._async_chunk = True
        return host

    def test_register_blocked_after_finish_sentinel(self):
        """register_chunk_recv must be a no-op after the finish sentinel."""
        host = self._make_host(stage_id=1)

        req = _make_request("req-1", "ext-req-1")

        # Simulate the bg thread having received the finish sentinel:
        with host._lock:
            host._chunk_stream_completed.add("req-1")

        # Now try to re-register — this mimics the coordinator asking
        # the model runner to poll for the next (non-existent) chunk.
        host.register_chunk_recv(req)

        # The request must NOT appear in _pending_load_reqs
        self.assertNotIn(
            "req-1",
            host._pending_load_reqs,
            "register_chunk_recv should skip requests whose chunk stream is already complete",
        )

        host.shutdown_omni_connectors()

    def test_register_allowed_before_finish(self):
        """register_chunk_recv works normally before finish sentinel."""
        host = self._make_host(stage_id=1)
        req = _make_request("req-1", "ext-req-1")

        host.register_chunk_recv(req)
        self.assertIn(
            "req-1",
            host._pending_load_reqs,
            "register_chunk_recv should add request to pending when stream is not yet complete",
        )

        host.shutdown_omni_connectors()

    def test_finish_sentinel_populates_completed_set(self):
        """Receiving is_finished=True adds to _chunk_stream_completed."""
        host = self._make_host(stage_id=1)

        # Simulate _poll_single_request receiving is_finished=True
        req_id = "req-1"
        with host._lock:
            host._chunk_finished_req_ids.add(req_id)
            host._chunk_stream_completed.add(req_id)
            host._local_stage_payload_cache[req_id] = {"finished": True}
            host._local_request_metadata[req_id] = {}
            host._finished_load_reqs.add(req_id)
            host._pending_load_reqs.pop(req_id, None)

        self.assertIn(req_id, host._chunk_stream_completed)

        # Subsequent register_chunk_recv should be blocked
        req = _make_request(req_id, f"ext-{req_id}")
        host.register_chunk_recv(req)
        self.assertNotIn(req_id, host._pending_load_reqs)

        host.shutdown_omni_connectors()

    def test_stage_0_always_skipped(self):
        """Stage-0 has no upstream, register_chunk_recv is always no-op."""
        host = self._make_host(stage_id=0)
        host._stage_id = 0

        req = _make_request("req-1")
        host.register_chunk_recv(req)
        self.assertNotIn("req-1", host._pending_load_reqs)

        host.shutdown_omni_connectors()

    def test_full_payload_recv_guard_still_works(self):
        """Pre-existing guard: staged full-payload results prevent registration."""
        host = self._make_host(stage_id=1)

        with host._lock:
            host._stage_recv_req_ids.add("req-1")

        req = _make_request("req-1", "ext-req-1")
        host.register_chunk_recv(req)
        self.assertNotIn("req-1", host._pending_load_reqs)

        host.shutdown_omni_connectors()


class TestCleanupFinishedRequest(unittest.TestCase):
    """Test cleanup_finished_request frees per-request mixin state."""

    def _make_host(self, stage_id: int = 1) -> MixinHost:
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=stage_id, async_chunk=True),
        )
        host._omni_connector = MockConnector(stage_id=stage_id)
        host._stage_id = stage_id
        host._async_chunk = True
        return host

    def test_cleanup_removes_all_state(self):
        """cleanup_finished_request removes all tracking dicts/sets."""
        host = self._make_host(stage_id=1)
        req_id = "req-1"
        ext_id = "ext-req-1"

        # Simulate state accumulated during a request's lifetime
        host._request_ids_mapping[req_id] = ext_id
        host._put_req_chunk[ext_id] = 5
        host._get_req_chunk[req_id] = 3
        host._send_side_request_payload[ext_id] = {"some": "data"}
        host._code_prompt_token_ids[ext_id] = [[1, 2, 3]]
        host._cached_ic[ext_id] = 16
        host._chunk_stream_completed.add(req_id)
        host._stage_recv_req_ids.add(req_id)
        host._local_stage_payload_cache[req_id] = {"engine_inputs": {}}
        host._local_request_metadata[req_id] = {"prompt_len": 10}

        # Cleanup
        host.cleanup_finished_request(req_id)

        # All state should be gone
        self.assertNotIn(req_id, host._request_ids_mapping)
        self.assertNotIn(ext_id, host._put_req_chunk)
        self.assertNotIn(req_id, host._get_req_chunk)
        self.assertNotIn(ext_id, host._send_side_request_payload)
        self.assertNotIn(ext_id, host._code_prompt_token_ids)
        self.assertNotIn(ext_id, host._cached_ic)
        self.assertNotIn(req_id, host._chunk_stream_completed)
        self.assertNotIn(req_id, host._stage_recv_req_ids)
        self.assertNotIn(req_id, host._local_stage_payload_cache)
        self.assertNotIn(req_id, host._local_request_metadata)

        host.shutdown_omni_connectors()

    def test_cleanup_removes_per_cycle_ready_state(self):
        """cleanup_finished_request clears ready/finished carry-over for req-id reuse."""
        host = self._make_host(stage_id=1)
        req_id = "req-1"

        host._pending_load_reqs[req_id] = _make_request(req_id, "ext-req-1")
        host._finished_load_reqs.add(req_id)
        host._chunk_ready_req_ids.add(req_id)
        host._chunk_finished_req_ids.add(req_id)

        host.cleanup_finished_request(req_id)

        self.assertNotIn(req_id, host._pending_load_reqs)
        self.assertNotIn(req_id, host._finished_load_reqs)
        self.assertNotIn(req_id, host._chunk_ready_req_ids)
        self.assertNotIn(req_id, host._chunk_finished_req_ids)

        host.shutdown_omni_connectors()

    def test_cleanup_without_mapping(self):
        """cleanup works for Stage-0 where _request_ids_mapping isn't set."""
        host = self._make_host(stage_id=0)
        host._stage_id = 0
        req_id = "req-1"

        # Stage-0 uses req_id directly (no ext_id mapping)
        host._put_req_chunk[req_id] = 3
        host._get_req_chunk[req_id] = 0
        host._cached_ic[req_id] = 4

        host.cleanup_finished_request(req_id)

        self.assertNotIn(req_id, host._put_req_chunk)
        self.assertNotIn(req_id, host._get_req_chunk)
        self.assertNotIn(req_id, host._cached_ic)

        host.shutdown_omni_connectors()

    def test_deferred_cleanup_removes_cached_ic(self):
        host = self._make_host(stage_id=1)
        req_id = "req-1"
        ext_id = "ext-req-1"

        host._request_ids_mapping[req_id] = ext_id
        host._pending_save_counts[ext_id] = 1
        host._cached_ic[ext_id] = 8

        host.cleanup_finished_request(req_id)

        self.assertIn(ext_id, host._deferred_send_cleanup)
        self.assertIn(ext_id, host._cached_ic)

        host._decrement_pending_save_count(ext_id)

        self.assertNotIn(ext_id, host._deferred_send_cleanup)
        self.assertNotIn(ext_id, host._cached_ic)

        host.shutdown_omni_connectors()

    def test_prune_inactive_requests_cleans_stale_state_but_keeps_active(self):
        """Inactive request IDs should be pruned without touching active ones."""
        host = self._make_host(stage_id=1)
        active_req_id = "req-active"
        stale_req_id = "req-stale"
        stale_ext_id = "ext-stale"

        host._request_ids_mapping[active_req_id] = "ext-active"
        host._request_ids_mapping[stale_req_id] = stale_ext_id
        host._put_req_chunk[stale_ext_id] = 2
        host._get_req_chunk[stale_req_id] = 1
        host._finished_load_reqs.add(stale_req_id)
        host._chunk_ready_req_ids.update({active_req_id, stale_req_id})
        host._chunk_finished_req_ids.add(stale_req_id)
        host._chunk_stream_completed.add(stale_req_id)
        host._stage_recv_req_ids.add(active_req_id)
        host._send_side_request_payload[stale_ext_id] = {"stale": True}
        host._code_prompt_token_ids[stale_ext_id] = [[1, 2, 3]]

        pruned = host.prune_inactive_requests({active_req_id})

        self.assertEqual(pruned, {stale_req_id})
        self.assertIn(active_req_id, host._request_ids_mapping)
        self.assertIn(active_req_id, host._chunk_ready_req_ids)
        self.assertIn(active_req_id, host._stage_recv_req_ids)
        self.assertNotIn(stale_req_id, host._request_ids_mapping)
        self.assertNotIn(stale_ext_id, host._put_req_chunk)
        self.assertNotIn(stale_req_id, host._get_req_chunk)
        self.assertNotIn(stale_req_id, host._pending_load_reqs)
        self.assertNotIn(stale_req_id, host._finished_load_reqs)
        self.assertNotIn(stale_req_id, host._chunk_ready_req_ids)
        self.assertNotIn(stale_req_id, host._chunk_finished_req_ids)
        self.assertNotIn(stale_req_id, host._chunk_stream_completed)
        self.assertNotIn(stale_req_id, host._stage_recv_req_ids)
        self.assertNotIn(stale_ext_id, host._send_side_request_payload)
        self.assertNotIn(stale_ext_id, host._code_prompt_token_ids)

        host.shutdown_omni_connectors()

    def test_prune_inactive_requests_keeps_recently_received_full_payload_state(self):
        """Late bg-thread receives must survive until the scheduler catches up."""
        host = self._make_host(stage_id=1)
        req_id = "req-recv-race"
        ext_id = "ext-recv-race"

        host._request_ids_mapping[req_id] = ext_id
        host._put_req_chunk[ext_id] = 1
        host._local_stage_payload_cache[req_id] = {"engine_inputs": {"ids": [1, 2, 3]}}
        host._local_request_metadata[req_id] = {"next_stage_prompt_len": 3}
        host._stage_recv_req_ids.add(req_id)

        pruned = host.prune_inactive_requests(set())

        self.assertEqual(pruned, set())
        self.assertIn(req_id, host._request_ids_mapping)
        self.assertIn(req_id, host._local_stage_payload_cache)
        self.assertIn(req_id, host._local_request_metadata)
        self.assertIn(req_id, host._stage_recv_req_ids)
        self.assertIn(ext_id, host._put_req_chunk)

        # Once the scheduler has consumed the wake-up and the request really
        # disappears from all protected sets, prune should clean it up.
        host._stage_recv_req_ids.clear()
        host._local_stage_payload_cache.clear()
        host._local_request_metadata.clear()

        pruned = host.prune_inactive_requests(set())

        self.assertEqual(pruned, {req_id})
        self.assertNotIn(req_id, host._request_ids_mapping)
        self.assertNotIn(ext_id, host._put_req_chunk)

        host.shutdown_omni_connectors()


class TestSendChunkCachesMapping(unittest.TestCase):
    """Test that send_chunk caches internal→external req ID mapping."""

    def test_send_chunk_populates_request_ids_mapping(self):
        """send_chunk should cache the internal→external mapping."""
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=0, async_chunk=True),
        )
        host._omni_connector = MockConnector(stage_id=0)
        host._stage_id = 0
        host._async_chunk = True

        def mock_process(transfer_manager, pooling_output, request):
            return {"data": "test", "finished": False}

        host._custom_process_func = mock_process

        request = _make_request("internal-1", "external-1")
        host.send_chunk(request, pooling_output={"v": 1})

        # The mapping should be cached
        self.assertEqual(
            host._request_ids_mapping.get("internal-1"),
            "external-1",
        )

        time.sleep(0.1)
        host.shutdown_omni_connectors()


class TestLocalPayloadCacheLifecycle(unittest.TestCase):
    """Unit tests for the local payload cache API (RFC §2.4)."""

    def _make_host(self) -> MixinHost:
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=0),
        )
        host._omni_connector = MockConnector(stage_id=0)
        host._stage_id = 0
        return host

    def test_put_get_pop(self):
        host = self._make_host()
        payload = {"engine_inputs": {"ids": [1, 2, 3]}}
        host.put_local_stage_payload("r1", payload)

        self.assertEqual(host.get_local_stage_payload("r1"), payload)
        popped = host.pop_local_stage_payload("r1")
        self.assertEqual(popped, payload)
        self.assertIsNone(host.get_local_stage_payload("r1"))
        host.shutdown_omni_connectors()

    def test_recv_full_payload_inputs_populates_local_cache(self):
        host = self._make_host()
        host._omni_connector = MockConnector(stage_id=0)
        host._stage_id = 0

        # Simulate a full payload already staged by the bg recv path
        with host._lock:
            host._local_stage_payload_cache["r1"] = {"tok": [10]}
            host._stage_recv_req_ids.add("r1")

        host.recv_full_payload_inputs(scheduler_output=None)
        self.assertEqual(host.get_local_stage_payload("r1"), {"tok": [10]})
        host.shutdown_omni_connectors()

    def test_rank0_only_polls_connector_for_tp_full_payload(self):
        host = self._make_host()
        host._omni_connector = MagicMock()
        host._stage_id = 2
        host._local_rank = 0
        host._request_ids_mapping["r1"] = "ext-r1"
        host._get_req_chunk["r1"] = 0
        payload = {"tok": [10], "finished": torch.tensor(True)}
        connector_result = (payload, 123)
        host._omni_connector.get.return_value = connector_result
        tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)

        with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
            made_progress = host._poll_single_request("r1")

        self.assertTrue(made_progress)
        host._omni_connector.get.assert_called_once_with("1", "2", "ext-r1_1_0")
        self.assertEqual(tp_group.broadcast_inputs, [])
        self.assertEqual(host.get_local_stage_payload("r1"), payload)
        self.assertIn("r1", host._full_payload_pending_broadcast_req_ids)
        self.assertNotIn("r1", host._stage_recv_req_ids)
        self.assertIsNone(host.get_local_request_metadata("r1"))
        host.shutdown_omni_connectors()

    def test_tp_follower_skips_connector_poll_for_full_payload(self):
        host = self._make_host()
        host._omni_connector = MagicMock()
        host._stage_id = 2
        host._local_rank = 1
        host._request_ids_mapping["r1"] = "ext-r1"
        host._get_req_chunk["r1"] = 0
        tp_group = _FakeTPGroup(world_size=2, rank_in_group=1)

        with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
            made_progress = host._poll_single_request("r1")

        self.assertFalse(made_progress)
        host._omni_connector.get.assert_not_called()
        self.assertEqual(tp_group.broadcast_inputs, [])
        self.assertNotIn("r1", host._local_stage_payload_cache)
        host.shutdown_omni_connectors()

    def test_recv_full_payload_inputs_broadcasts_tp_leader_results_to_followers(self):
        host = self._make_host()
        host._omni_connector = MagicMock()
        host._stage_id = 2
        host._local_rank = 1
        host._pending_load_reqs["r1"] = object()
        payload = {"tok": [10], "finished": torch.tensor(True)}
        tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result={"r1": payload})

        with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
            results = host.recv_full_payload_inputs(scheduler_output=None)

        self.assertEqual(results, {"r1": payload})
        self.assertEqual(host.get_local_stage_payload("r1"), payload)
        self.assertEqual(host.get_local_request_metadata("r1"), {})
        self.assertEqual(host._stage_recv_req_ids, {"r1"})
        self.assertNotIn("r1", host._pending_load_reqs)
        self.assertEqual(tp_group.broadcast_inputs, [None])
        host.shutdown_omni_connectors()


class TestTPAsyncChunkFanout(unittest.TestCase):
    def _make_host(self, rank: int) -> MixinHost:
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
        )
        host._omni_connector = MagicMock()
        host._stage_id = 2
        host._async_chunk = True
        host._model_mode = "gen"
        host._local_rank = rank
        host._request_ids_mapping["r1"] = "ext-r1"
        host._get_req_chunk["r1"] = 0
        return host

    def test_rank0_only_polls_connector_for_tp_async_chunk(self):
        host = self._make_host(rank=0)
        payload = {
            "codes": {"audio": [10, 11]},
            "meta": {"left_context_size": 0, "finished": torch.tensor(False)},
        }
        host._omni_connector.get.return_value = (payload, 123)
        tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)

        with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
            made_progress = host._poll_single_request("r1")

        self.assertTrue(made_progress)
        host._omni_connector.get.assert_called_once_with("1", "2", "ext-r1_1_0")
        self.assertEqual(host.get_local_stage_payload("r1"), payload)
        self.assertIn("r1", host._finished_load_reqs)
        self.assertIn("r1", host._async_chunk_updated_req_ids)
        self.assertEqual(tp_group.broadcast_inputs, [])
        host.shutdown_omni_connectors()

    def test_tp_follower_skips_connector_poll_for_async_chunk(self):
        host = self._make_host(rank=1)
        tp_group = _FakeTPGroup(world_size=2, rank_in_group=1)

        with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
            made_progress = host._poll_single_request("r1")

        self.assertFalse(made_progress)
        host._omni_connector.get.assert_not_called()
        self.assertIsNone(host.get_local_stage_payload("r1"))
        self.assertEqual(tp_group.broadcast_inputs, [])
        host.shutdown_omni_connectors()

    def test_get_output_broadcasts_tp_async_chunk_payloads_to_followers(self):
        host = self._make_host(rank=1)
        host._pending_load_reqs["r1"] = object()
        payload = {
            "code_predictor_codes": [10, 11],
            "left_context_size": 0,
            "finished": torch.tensor(True),
        }
        packet = {
            "staged_payloads": {"r1": payload},
            "request_metadata": {"r1": {"code_predictor_codes": [10, 11], "left_context_size": 0}},
            "newly_finished": {"r1"},
            "chunk_finished": {"r1"},
        }
        tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result=packet)

        with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
            output = host.get_omni_connector_output()

        self.assertEqual(output.chunk_ready_req_ids, {"r1"})
        self.assertEqual(output.chunk_finished_req_ids, {"r1"})
        self.assertEqual(
            output.request_metadata,
            {"r1": {"code_predictor_codes": [10, 11], "left_context_size": 0}},
        )
        self.assertEqual(host.get_local_stage_payload("r1"), payload)
        self.assertNotIn("r1", host._pending_load_reqs)
        self.assertIn("r1", host._chunk_stream_completed)
        self.assertEqual(tp_group.broadcast_inputs, [None])
        host.shutdown_omni_connectors()


class TestKVTransferLifecycle(unittest.TestCase):
    """Unit tests for KV transfer lifecycle methods."""

    def _make_host(self) -> MixinHost:
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=0),
        )
        return host

    def test_mark_drain_ack_complete(self):
        host = self._make_host()
        self.assertFalse(host.has_pending_kv_work())

        host.mark_kv_transfer("r1", seq_len=100, block_ids=[0, 1, 2])
        self.assertTrue(host.has_pending_kv_work())
        self.assertTrue(host.is_kv_transfer_triggered("r1"))

        # Drain moves pending → active
        pending = host.drain_pending_kv_transfers()
        self.assertEqual(pending, {"r1": {"seq_len": 100, "block_ids": [0, 1, 2]}})
        self.assertIn("r1", host._kv_active_transfers)
        self.assertTrue(host.has_pending_kv_work())

        # Ack moves active → completed
        host.ack_kv_transfers(["r1"])
        self.assertNotIn("r1", host._kv_active_transfers)
        self.assertIn("r1", host._kv_completed_transfers)

        # Drain completed
        completed = host.drain_completed_kv_transfers()
        self.assertEqual(completed, {"r1"})
        self.assertFalse(host.has_pending_kv_work())
        host.shutdown_omni_connectors()

    def test_mark_dedup(self):
        host = self._make_host()
        host.mark_kv_transfer("r1", seq_len=100, block_ids=[0])
        host.mark_kv_transfer("r1", seq_len=200, block_ids=[0, 1])
        # Second mark is a no-op
        self.assertEqual(host._kv_pending_transfers["r1"]["seq_len"], 100)
        host.shutdown_omni_connectors()

    def test_cleanup_removes_kv_state(self):
        host = self._make_host()
        host.mark_kv_transfer("r1", seq_len=50, block_ids=[0])
        host.drain_pending_kv_transfers()
        host.cleanup_finished_request("r1")
        self.assertFalse(host.is_kv_transfer_triggered("r1"))
        self.assertNotIn("r1", host._kv_active_transfers)
        self.assertFalse(host.has_pending_kv_work())
        host.shutdown_omni_connectors()


class TestAsyncPayloadLifecycle(unittest.TestCase):
    """Regression tests for async payload delivery lifecycle."""

    def test_send_side_request_payload_not_cleared_before_payload_is_consumable(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
        )
        host._request_ids_mapping["r1"] = "r1"
        payload = {
            "embed": {"decode": torch.ones(1, 2)},
            "ids": {"output": [1]},
            "meta": {
                "finished": torch.tensor(False),
                "override_keys": [["embed", "decode"], ["ids", "output"]],
            },
        }

        host._accumulate_payload("r1", dict(payload))
        with host._lock:
            host._finished_load_reqs.add("r1")

        host.get_omni_connector_output()
        self.assertIn("r1", host._send_side_request_payload)
        host.shutdown_omni_connectors()

    def test_payload_consumable_ignores_token_horizon_only_updates(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
        )
        payload = {
            "ids": {"output": [1, 2, 3]},
            "embed": {"decode_token_start": 2, "decode_token_end": 3},
            "meta": {
                "finished": torch.tensor(False),
                "override_keys": [
                    ["ids", "output"],
                    ["embed", "decode_token_start"],
                    ["embed", "decode_token_end"],
                ],
            },
        }
        self.assertFalse(host._payload_is_consumable(payload))
        host.shutdown_omni_connectors()

    def test_payload_consumable_accepts_decode_embeddings(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
        )
        payload = {
            "ids": {"output": [1, 2, 3]},
            "embed": {"decode": torch.ones(1, 2)},
            "meta": {"finished": torch.tensor(False)},
        }
        self.assertTrue(host._payload_is_consumable(payload))
        host.shutdown_omni_connectors()

    def test_ar_metadata_only_followup_chunk_does_not_rewake_request(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
        )
        host._omni_connector = MagicMock()
        host._stage_id = 1
        host._async_chunk = True
        host._model_mode = "ar"
        host._request_ids_mapping["r1"] = "ext-r1"
        host._get_req_chunk["r1"] = 0

        host._omni_connector.get.side_effect = [
            (
                {
                    "embed": {"decode": torch.ones(1, 2)},
                    "meta": {"finished": torch.tensor(False)},
                },
                1,
            ),
            (
                {
                    "meta": {"next_stage_prompt_len": 7, "finished": torch.tensor(False)},
                },
                1,
            ),
        ]

        host._poll_single_request("r1")
        output1 = host.get_omni_connector_output()
        self.assertEqual(output1.chunk_ready_req_ids, {"r1"})

        host._poll_single_request("r1")
        output2 = host.get_omni_connector_output()
        self.assertEqual(output2.chunk_ready_req_ids, set())
        self.assertEqual(output2.request_metadata, {"r1": {"next_stage_prompt_len": 7}})

        host.shutdown_omni_connectors()

    def test_non_ar_recv_does_not_overwrite_unconsumed_staged_chunk(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
        )
        host._omni_connector = MagicMock()
        host._stage_id = 2
        host._async_chunk = True
        host._model_mode = "gen"
        host._request_ids_mapping["r1"] = "ext-r1"
        host._get_req_chunk["r1"] = 1
        host._local_stage_payload_cache["r1"] = {
            "code_predictor_codes": [1, 2, 3],
            "left_context_size": 0,
            "finished": torch.tensor(False),
        }

        made_progress = host._poll_single_request("r1")

        self.assertFalse(made_progress)
        host._omni_connector.get.assert_not_called()
        self.assertEqual(host._get_req_chunk["r1"], 1)

        host.shutdown_omni_connectors()

    def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self):
        host = MixinHost()
        host.init_omni_connectors(
            model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
        )
        host._omni_connector = MagicMock()
        host._stage_id = 2
        host._async_chunk = True
        host._model_mode = "gen"
        host._request_ids_mapping["r1"] = "ext-r1"
        host._get_req_chunk["r1"] = 1
        host._local_request_metadata["r1"] = {
            "code_predictor_codes": [10, 11, 12],
            "left_context_size": 0,
        }
        host._finished_load_reqs.add("r1")

        made_progress = host._poll_single_request("r1")

        self.assertFalse(made_progress)
        host._omni_connector.get.assert_not_called()
        self.assertEqual(host._get_req_chunk["r1"], 1)

        output = host.get_omni_connector_output()
        self.assertEqual(output.request_metadata["r1"]["code_predictor_codes"], [10, 11, 12])
        self.assertEqual(output.chunk_ready_req_ids, {"r1"})

        host._omni_connector.get.return_value = (
            {
                "codes": {"audio": [20, 21, 22]},
                "meta": {"left_context_size": 0, "finished": torch.tensor(False)},
            },
            1,
        )
        made_progress = host._poll_single_request("r1")

        self.assertTrue(made_progress)
        host._omni_connector.get.assert_called_once()
        self.assertEqual(host._get_req_chunk["r1"], 2)

        host.shutdown_omni_connectors()


class TestRankAwareKVRouting(unittest.TestCase):
    def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost:
        host = MixinHost()
        host.init_omni_connectors(model_config=_make_model_config(stage_id=1))
        host._from_tp = from_tp
        host._to_tp = to_tp
        host._local_rank = local_rank
        return host

    def test_recv_keys_use_remote_rank_as_from_rank(self):
        host = self._make_host(from_tp=4, to_tp=2, local_rank=1)
        self.assertEqual(
            host.get_rank_aware_kv_keys("req", from_stage=0),
            ["req_0_0_2_1", "req_0_0_3_1"],
        )
        host.shutdown_omni_connectors()

    def test_single_rank_hooks_use_default_manager_key(self):
        host = self._make_host(from_tp=1, to_tp=1, local_rank=0)
        expected = ["omni_0_to_1_kv_cache_req"]

        self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
        self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
        host.shutdown_omni_connectors()

    def test_init_uses_effective_manager_tp2_topology(self):
        model_config = _make_model_config(stage_id=0)
        model_config.omni_kv_config = {
            "need_send_cache": True,
            "omni_from_stage": "0",
            "omni_to_stage": "1",
            "rank_mapping": {"from_tp": 2, "to_tp": 2},
        }
        with patch(
            "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank",
            return_value=1,
        ):
            manager = OmniKVTransferManager.from_model_config(model_config)

        host = MixinHost()
        host.init_omni_connectors(model_config=model_config, kv_transfer_manager=manager)

        expected = ["req_0_0_1_1"]
        self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
        self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
        self.assertEqual(host.get_kv_rank_mapping()["from_tp"], 2)
        self.assertEqual(host.get_kv_rank_mapping()["to_tp"], 2)
        host.shutdown_omni_connectors()

    def test_send_keys_route_from_rank_gt_to_rank(self):
        host = self._make_host(from_tp=4, to_tp=2, local_rank=3)
        self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0), ["req_0_0_3_1"])
        host.shutdown_omni_connectors()

    def test_invalid_recv_rank_mapping_raises(self):
        host = self._make_host(from_tp=3, to_tp=2, local_rank=1)
        with self.assertRaises(ValueError):
            host.get_rank_aware_kv_keys("req", from_stage=0)
        host.shutdown_omni_connectors()

    def test_invalid_send_rank_mapping_raises(self):
        host = self._make_host(from_tp=3, to_tp=2, local_rank=1)
        with self.assertRaises(ValueError):
            host.get_rank_aware_kv_send_keys("req", from_stage=0)
        host.shutdown_omni_connectors()

    def test_merge_rank_sharded_payloads_concatenates_head_dimension(self):
        host = self._make_host(from_tp=4, to_tp=2, local_rank=0)
        payloads = [
            {"layer_blocks": {"key_cache": [torch.ones(2, 1, 3)], "value_cache": [torch.ones(2, 1, 3)]}},
            {"layer_blocks": {"key_cache": [torch.full((2, 1, 3), 2.0)], "value_cache": [torch.full((2, 1, 3), 2.0)]}},
        ]
        merged = host._merge_rank_sharded_kv_payloads(payloads)
        self.assertEqual(tuple(merged["layer_blocks"]["key_cache"][0].shape), (2, 2, 3))
        self.assertTrue(torch.equal(merged["layer_blocks"]["key_cache"][0][:, 0], torch.ones(2, 3)))
        self.assertTrue(torch.equal(merged["layer_blocks"]["key_cache"][0][:, 1], torch.full((2, 3), 2.0)))
        host.shutdown_omni_connectors()

    def test_slice_rank_sharded_payload_splits_head_dimension(self):
        host = self._make_host(from_tp=2, to_tp=4, local_rank=1)
        payload = {
            "layer_blocks": {
                "key_cache": [torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)],
                "value_cache": [torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)],
            },
            "metadata": {},
        }
        sliced = host._slice_rank_sharded_kv_payload(payload)
        self.assertEqual(tuple(sliced["layer_blocks"]["key_cache"][0].shape), (2, 2, 3))
        expected = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)[:, 2:4, :]
        self.assertTrue(torch.equal(sliced["layer_blocks"]["key_cache"][0], expected))
        host.shutdown_omni_connectors()


class TestAttachOmniConnectorOutput(unittest.TestCase):
    def test_wraps_empty_model_runner_output_when_signals_exist(self):
        from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT

        host = MixinHost()
        host.get_omni_connector_output = lambda: OmniConnectorOutput(chunk_ready_req_ids={"req-1"})

        wrapped = host.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)

        self.assertIsNot(wrapped, EMPTY_MODEL_RUNNER_OUTPUT)
        self.assertEqual(wrapped.omni_connector_output.chunk_ready_req_ids, {"req-1"})


class TestConnectorConfigValidation(unittest.TestCase):
    def test_invalid_connector_name_raises(self):
        host = MixinHost()
        model_config = _make_model_config(stage_id=1)
        model_config.stage_connector_config = {"name": "   "}

        with self.assertRaisesRegex(RuntimeError, "missing connector name"):
            host.init_omni_connectors(model_config=model_config)


class _FailingConnector:
    """Connector whose put() fails a configurable number of times."""

    def __init__(self, fail_count: int = 1, raise_on_fail: bool = False):
        self._fail_count = fail_count
        self._raise_on_fail = raise_on_fail
        self.attempt = 0

    def put(self, from_stage, to_stage, put_key, data):
        self.attempt += 1
        if self.attempt <= self._fail_count:
            if self._raise_on_fail:
                raise ConnectionError("transient connector error")
            return False, 0, None
        return True, len(str(data)), None

    def get(self, *a, **kw):
        return None

    def close(self):
        pass


class TestSendRetry(unittest.TestCase):
    """Tests for P1-2: failed connector sends must be retried."""

    def _make_sender(self, connector):
        sender = MixinHost()
        sender.init_omni_connectors(
            model_config=_make_model_config(stage_id=0, async_chunk=True),
        )
        sender._omni_connector = connector
        sender._stage_id = 0
        sender._async_chunk = True
        return sender

    def _make_task(self, req_id="r1"):
        return {
            "stage_id": 0,
            "next_stage_id": 1,
            "request_id": req_id,
            "data": {"payload": "test"},
        }

    def test_send_single_request_returns_false_on_put_failure(self):
        connector = _FailingConnector(fail_count=999)
        sender = self._make_sender(connector)

        result = sender._send_single_request(self._make_task())
        self.assertFalse(result)
        sender.shutdown_omni_connectors()

    def test_send_single_request_does_not_decrement_on_failure(self):
        connector = _FailingConnector(fail_count=999)
        sender = self._make_sender(connector)
        sender._pending_save_counts["r1"] = 1

        sender._send_single_request(self._make_task())
        self.assertEqual(sender._pending_save_counts.get("r1"), 1, "pending count must NOT be decremented on failure")
        sender.shutdown_omni_connectors()

    def test_send_single_request_decrements_on_success(self):
        connector = MockConnector(stage_id=0)
        sender = self._make_sender(connector)
        sender._pending_save_counts["r1"] = 1

        result = sender._send_single_request(self._make_task())
        self.assertTrue(result)
        self.assertNotIn("r1", sender._pending_save_counts, "pending count should be zero/removed on success")
        sender.shutdown_omni_connectors()

    def test_requeue_or_drop_requeues_on_first_failure(self):
        sender = self._make_sender(MockConnector(stage_id=0))
        task = self._make_task()

        sender._requeue_or_drop_failed_send(task)

        self.assertEqual(task.get("_retry_count"), 1)
        with sender._lock:
            dq = sender._pending_save_reqs.get("r1")
        self.assertIsNotNone(dq)
        self.assertEqual(len(dq), 1)
        sender.shutdown_omni_connectors()

    def test_requeue_or_drop_drops_after_max_retries(self):
        sender = self._make_sender(MockConnector(stage_id=0))
        sender._pending_save_counts["r1"] = 1
        task = self._make_task()
        task["_retry_count"] = sender._MAX_SEND_RETRIES  # already at max

        sender._requeue_or_drop_failed_send(task)

        with sender._lock:
            dq = sender._pending_save_reqs.get("r1")
        self.assertTrue(dq is None or len(dq) == 0, "task should NOT be re-enqueued after max retries")
        self.assertNotIn("r1", sender._pending_save_counts, "pending count should be cleaned up on final drop")
        sender.shutdown_omni_connectors()

    def test_save_loop_retries_on_exception(self):
        """Integration: _save_loop retries a task when put() raises."""
        from collections import deque

        connector = _FailingConnector(fail_count=1, raise_on_fail=True)
        sender = self._make_sender(connector)
        task = self._make_task()

        with sender._lock:
            sender._pending_save_reqs["r1"] = deque([task])
        sender._pending_save_counts["r1"] = 1

        sender._stop_event.clear()

        def run_one_loop():
            sender._save_loop()

        sender._stop_event.set()  # will exit after one iteration
        # Run manually instead of threading
        # Simulate: pop task, send fails, requeue
        popped_task = None
        with sender._lock:
            dq = sender._pending_save_reqs.get("r1")
            if dq:
                popped_task = dq.popleft()
                if not dq:
                    del sender._pending_save_reqs["r1"]

        if popped_task is not None:
            success = False
            try:
                success = sender._send_single_request(popped_task)
            except Exception:
                pass
            if not success:
                sender._requeue_or_drop_failed_send(popped_task)

        # After first failure, task should be re-enqueued
        with sender._lock:
            dq = sender._pending_save_reqs.get("r1")
        self.assertIsNotNone(dq)
        self.assertEqual(len(dq), 1)
        requeued = dq[0]
        self.assertEqual(requeued.get("_retry_count"), 1)

        # Second attempt should succeed (connector now returns True)
        success = sender._send_single_request(requeued)
        self.assertTrue(success)
        sender.shutdown_omni_connectors()


if __name__ == "__main__":
    unittest.main()