test(ci): cover decode integration branches

Exercise worker snapshot imports, proxy appends, cancellation, concurrency, and memory-drain paths so the platform-neutral gate remains deterministic at 100% coverage.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

tests/inference_engine/distributed/test_mlx_ring.py

@@ -37,6 +37,25 @@ def test_probe_never_raises_and_is_structured():
         assert env.world_size == 0
 
 
+def test_probe_reports_import_failure_on_every_platform(monkeypatch):
+    def fail_import(_name):
+        raise ImportError("synthetic missing mlx")
+
+    monkeypatch.setattr(
+        "inference_engine.distributed.mlx_ring.importlib.import_module",
+        fail_import,
+    )
+    env = probe_ring_environment()
+    assert not env.is_available
+    assert env.backend == ""
+    assert env.rank == 0
+    assert env.world_size == 0
+    assert env.failure_reason == (
+        "mlx.core.distributed import failed: "
+        "ImportError: synthetic missing mlx"
+    )
+
+
 @pytest.mark.skipif(
     platform.machine() == "arm64",
     reason="Linux-gate branch: asserts the mlx-absent probe result",

tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py

@@ -283,6 +283,82 @@ def test_successful_local_import_suffix_and_on_reuse(monkeypatch):
     hook.close()
 
 
+def test_worker_proxy_import_uses_returned_snapshot_state():
+    compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+    hook = DistributedPrefillCacheHook(
+        PrefixCacheStore(
+            compatibility,
+            max_bytes=1024,
+            node_id="head",
+        ),
+        compression=CompressionCodec.NONE,
+    )
+    block_hash = chained_block_hashes([1, 2], compatibility)[0]
+
+    class WorkerProxy(_Verifier):
+        def import_snapshot(self, payload, imported_compatibility):
+            assert payload == b"worker-snapshot"
+            assert imported_compatibility == compatibility
+            return {
+                "next_global_position": 2,
+                "cached_token_ids": [1, 2],
+                "block_hash": block_hash.hex(),
+            }
+
+    verifier = WorkerProxy()
+    reused = hook._try_import(
+        verifier,
+        [1, 2],
+        _Hit(
+            "local",
+            "lease",
+            1,
+            2,
+            len(b"worker-snapshot"),
+            b"worker-snapshot",
+            block_hash=block_hash,
+        ),
+    )
+    assert reused == 2
+    assert verifier.cached_token_sequence == [1, 2]
+    assert verifier.next_global_position == 2
+    assert hook.stats.local_hits == 1
+    hook.close()
+
+
+def test_worker_proxy_append_path_skips_snapshot_publication(monkeypatch):
+    compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+    hook = DistributedPrefillCacheHook(
+        PrefixCacheStore(
+            compatibility,
+            max_bytes=1024,
+            node_id="head",
+        ),
+    )
+    published = []
+    monkeypatch.setattr(
+        hook,
+        "_publish_boundary",
+        lambda *args: published.append(args),
+    )
+
+    class WorkerProxy(_Verifier):
+        is_decode_worker_proxy = True
+
+        def append_accepted_tokens(self, tokens):
+            self.cached_token_sequence.extend(tokens)
+            self.next_global_position += len(tokens)
+
+    verifier = WorkerProxy()
+    verifier.prefill([1, 2])
+    hashes = chained_block_hashes([1, 2, 3, 4], compatibility)
+    hook._compute_and_publish(verifier, [1, 2, 3, 4], hashes, reused=2)
+    assert verifier.cached_token_sequence == [1, 2, 3, 4]
+    assert hook.stats.tokens_computed == 2
+    assert published == []
+    hook.close()
+
+
 def test_import_budget_and_remote_worker_failures(monkeypatch):
     compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
     store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head")

tests/inference_engine/server/test_grpc_app.py

@@ -31,6 +31,7 @@
     DEFAULT_BIND_ADDRESS,
     GrpcServerConfig,
     RuntimeServiceServicer,
+    _ThreadedEventStream,
     create_grpc_server,
 )
 from inference_engine.server.proto_gen.kakeya.v1 import (
@@ -883,6 +884,232 @@ def generate(self, session_id, *, max_tokens, **kw):
         await server.stop(grace=0.1)
 
 
+async def test_threaded_event_stream_stops_putting_after_cancel():
+    cancel_event = threading.Event()
+    stream = _ThreadedEventStream((), cancel_event)
+    stream._queue.put(("occupied", None))
+    cancel_event.set()
+    stream._put(("ignored", None))
+    assert stream._queue.get_nowait() == ("occupied", None)
+    assert stream._queue.empty()
+
+
+async def test_threaded_event_stream_breaks_after_cancelled_event():
+    cancel_event = threading.Event()
+
+    class CancelOnFirst:
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            cancel_event.set()
+            return "first"
+
+    stream = _ThreadedEventStream(CancelOnFirst(), cancel_event)
+    stream._run()
+    assert stream._queue.get_nowait() == ("first", None)
+
+
+class _DirectAbort(Exception):
+    pass
+
+
+class _DirectContext:
+    def __init__(self):
+        self.code = None
+        self.detail = ""
+        self.callback = None
+
+    async def abort(self, code, detail):
+        self.code = code
+        self.detail = detail
+        raise _DirectAbort
+
+    def cancelled(self):
+        return False
+
+    def add_done_callback(self, callback):
+        self.callback = callback
+
+
+async def test_watch_context_removes_cancelled_session():
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(store)
+    context = _DirectContext()
+    cancel_event = threading.Event()
+    servicer._watch_context(context, cancel_event, session.session_id)
+    assert context.callback is not None
+    context.callback(type("Done", (), {"cancelled": lambda self: True})())
+    assert cancel_event.is_set()
+    assert store.active_count == 0
+
+
+async def test_create_session_rejects_memory_drain():
+    store = SessionStore(capacity=1)
+    governor = type("Governor", (), {"draining": True})()
+    servicer = RuntimeServiceServicer(store, memory_governor=governor)
+    context = _DirectContext()
+    with pytest.raises(_DirectAbort):
+        await servicer.CreateSession(runtime_pb2.CreateSessionRequest(), context)
+    assert context.code == grpc.StatusCode.RESOURCE_EXHAUSTED
+    assert store.active_count == 0
+
+
+async def test_append_rejects_concurrent_operation():
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    coordinator = type(
+        "Coordinator",
+        (),
+        {"append_tokens": lambda self, session_id, token_ids: 0},
+    )()
+    servicer = RuntimeServiceServicer(store, append_coordinator=coordinator)
+    assert servicer._acquire_operation(session.session_id)
+    context = _DirectContext()
+    with pytest.raises(_DirectAbort):
+        await servicer.AppendTokens(
+            runtime_pb2.AppendTokensRequest(
+                session_id=session.session_id,
+                token_ids=[1],
+            ),
+            context,
+        )
+    assert context.code == grpc.StatusCode.ABORTED
+    servicer._release_operation(session.session_id)
+
+
+async def test_append_passes_cancel_event_and_resets_liveness():
+    observed = {}
+
+    class Coordinator:
+        def append_tokens(self, session_id, token_ids, cancel_event):
+            observed["session_id"] = session_id
+            observed["tokens"] = token_ids
+            observed["cancel_event"] = cancel_event
+            return len(token_ids)
+
+    class Liveness:
+        def update(self, phase):
+            observed["phase"] = phase
+
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(
+        store,
+        append_coordinator=Coordinator(),
+        liveness=Liveness(),
+    )
+    response = await servicer.AppendTokens(
+        runtime_pb2.AppendTokensRequest(
+            session_id=session.session_id,
+            token_ids=[4, 5],
+        ),
+        _DirectContext(),
+    )
+    assert response.history_length == 2
+    assert observed["session_id"] == session.session_id
+    assert observed["tokens"] == [4, 5]
+    assert isinstance(observed["cancel_event"], threading.Event)
+    assert observed["phase"] == "idle"
+
+
+async def test_append_asyncio_cancellation_removes_session():
+    class Coordinator:
+        def append_tokens(self, session_id, token_ids):
+            raise asyncio.CancelledError
+
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(store, append_coordinator=Coordinator())
+    with pytest.raises(asyncio.CancelledError):
+        await servicer.AppendTokens(
+            runtime_pb2.AppendTokensRequest(
+                session_id=session.session_id,
+                token_ids=[1],
+            ),
+            _DirectContext(),
+        )
+    assert store.active_count == 0
+
+
+async def test_append_worker_cancellation_maps_cancelled():
+    from inference_engine.session import OperationCancelledError
+
+    class Coordinator:
+        def append_tokens(self, session_id, token_ids):
+            raise OperationCancelledError("worker cancelled")
+
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(store, append_coordinator=Coordinator())
+    context = _DirectContext()
+    with pytest.raises(_DirectAbort):
+        await servicer.AppendTokens(
+            runtime_pb2.AppendTokensRequest(
+                session_id=session.session_id,
+                token_ids=[1],
+            ),
+            context,
+        )
+    assert context.code == grpc.StatusCode.CANCELLED
+    assert store.active_count == 0
+
+
+async def test_generate_rejects_concurrent_operation():
+    class Coordinator:
+        def generate(self, session_id, **kwargs):
+            return iter(())
+
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(
+        store,
+        generation_coordinator=Coordinator(),
+    )
+    assert servicer._acquire_operation(session.session_id)
+    context = _DirectContext()
+    with pytest.raises(_DirectAbort):
+        async for _ in servicer.Generate(
+            runtime_pb2.GenerateRequest(
+                session_id=session.session_id,
+                max_tokens=1,
+            ),
+            context,
+        ):
+            pass
+    assert context.code == grpc.StatusCode.ABORTED
+    servicer._release_operation(session.session_id)
+
+
+async def test_generate_busy_error_maps_aborted():
+    from inference_engine.session import SessionGenerationBusyError
+
+    class Coordinator:
+        def generate(self, session_id, **kwargs):
+            if False:
+                yield
+            raise SessionGenerationBusyError(session_id)
+
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(
+        store,
+        generation_coordinator=Coordinator(),
+    )
+    context = _DirectContext()
+    with pytest.raises(_DirectAbort):
+        async for _ in servicer.Generate(
+            runtime_pb2.GenerateRequest(
+                session_id=session.session_id,
+                max_tokens=1,
+            ),
+            context,
+        ):
+            pass
+    assert context.code == grpc.StatusCode.ABORTED
+
+
 async def test_factory_wires_prefill_cache_service():
     from inference_engine.distributed.capability import CacheCompatibility
     from inference_engine.distributed.prefill_cache import PrefixCacheStore