fix(router): decode ATOM KV events with map schema

infera/router/kv_event/events.py

@@ -132,11 +132,11 @@ class SglangKVEventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc
 def batch_type_for_engine(engine) -> type:
     """KVEventBatch schema matching a worker's engine wire format.
 
-    vLLM emits tagged-MAP events; SGLang (and anything else, the historical
-    default) emits tagged-ARRAY events.
+    vLLM and Infera's ATOM event hook emit tagged-MAP events; SGLang (the
+    historical default) emits tagged-ARRAY events.
     """
     from infera.common.worker_pool import EngineType
 
-    if engine == EngineType.VLLM:
+    if engine in (EngineType.VLLM, EngineType.ATOM):
         return KVEventBatch
     return SglangKVEventBatch

tests/unit/router/test_kv_event_client.py

@@ -42,8 +42,10 @@
     AllBlocksCleared,
     BlockRemoved,
     BlockStored,
+    KVEventBatch,
     SglangBlockStored,
     SglangKVEventBatch,
+    batch_type_for_engine,
 )
 from infera.router.kv_event.hasher import ROUTER_SEED, hash_chunk
 
@@ -335,6 +337,32 @@ def test_on_worker_removed_unknown_id_is_noop():
 # ----------------------------------------------------------------------
 
 
+def test_atom_decoder_matches_infera_atom_hook_wire_format():
+    """ATOM's Infera hook publishes the same tagged-map schema as vLLM.
+
+    Selecting SGLang's tagged-array decoder for an ATOM worker rejects every
+    event and leaves KV-aware routing with an empty cache view.
+    """
+    payload = msgspec.msgpack.encode(
+        KVEventBatch(
+            ts=1.0,
+            events=[
+                _make_stored(
+                    block_hashes=[111],
+                    parent_block_hash=None,
+                    token_ids=[1, 2, 3, 4],
+                )
+            ],
+        )
+    )
+
+    decoded = msgspec.msgpack.decode(payload, type=batch_type_for_engine(EngineType.ATOM))
+
+    assert isinstance(decoded, KVEventBatch)
+    assert isinstance(decoded.events[0], BlockStored)
+    assert decoded.events[0].block_hashes == [111]
+
+
 @pytest.mark.asyncio
 async def test_real_zmq_publisher_drives_cache_view():
     """End-to-end: a real ZMQ PUB on `tcp://127.0.0.1:<port>` emits a