class MockSlowCellGenerator(ColumnGenerator[ExpressionColumnConfig]):
    """Cell generator that simulates DGX Spark-class latency.

    Each call sleeps for ``latency`` seconds before returning. When used behind
    the throttle layer with a tight concurrency limit and the default 300s
    acquire timeout, queued tasks time out before getting a permit.

    For testing, ``timeout_after`` makes the generator raise ModelTimeoutError
    after the Nth call, reproducing what the ThrottledModelClient does when the
    throttle acquire deadline expires.
    """

    def __init__(
        self,
        *args: Any,
        latency: float = 0.0,
        timeout_after: int | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self._latency = latency
        self._timeout_after = timeout_after
        self._call_count = 0

    @staticmethod
    def get_generation_strategy() -> GenerationStrategy:
        return GenerationStrategy.CELL_BY_CELL

    def generate(self, data: dict) -> dict:
        self._call_count += 1
        if self._timeout_after is not None and self._call_count > self._timeout_after:
            raise ModelTimeoutError(
                f"Throttle acquire timed out after 300s for slow-provider/large-model [chat]"
            )
        data[self.config.name] = f"slow_{data.get('seed', '?')}"
        return data

    async def agenerate(self, data: dict) -> dict:
        self._call_count += 1
        if self._timeout_after is not None and self._call_count > self._timeout_after:
            raise ModelTimeoutError(
                f"Throttle acquire timed out after 300s for slow-provider/large-model [chat]"
            )
        await asyncio.sleep(self._latency)
        data[self.config.name] = f"slow_{data.get('seed', '?')}"
        return data


@pytest.mark.asyncio(loop_scope="session")
async def test_scheduler_slow_model_timeout_triggers_early_shutdown() -> None:
    provider = _mock_provider()
    num_records = 6

    configs = [
        SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}),
        LLMTextColumnConfig(name="answer", prompt="{{ seed }}", model_alias=MODEL_ALIAS),
    ]
    strategies = {
        "seed": GenerationStrategy.FULL_COLUMN,
        "answer": GenerationStrategy.CELL_BY_CELL,
    }

    slow_gen = MockSlowCellGenerator(
        config=_expr_config("answer"),
        resource_provider=provider,
        timeout_after=2,  # First 2 calls succeed, then every call times out
    )
    generators: dict[str, ColumnGenerator] = {
        "seed": MockSeedGenerator(config=_expr_config("seed"), resource_provider=provider),
        "answer": slow_gen,
    }

    graph = ExecutionGraph.create(configs, strategies)
    row_groups = [(0, num_records)]
    tracker = CompletionTracker.with_graph(graph, row_groups)
    buffer_manager = RowGroupBufferManager(MagicMock())

    scheduler = AsyncTaskScheduler(
        generators=generators,
        graph=graph,
        tracker=tracker,
        row_groups=row_groups,
        buffer_manager=buffer_manager,
        shutdown_error_window=4,
        shutdown_error_rate=0.5,
    )
    await asyncio.wait_for(scheduler.run(), timeout=10.0)

    assert scheduler._early_shutdown, "Early shutdown flag should be set"
    dropped = sum(1 for ri in range(num_records) if tracker.is_dropped(0, ri))
    assert dropped > 0, "Expected some rows to be dropped due to timeout errors"