Sample tito_session_mismatch instead of running it on every record fetch

miles/rollout/session/sessions.py

@@ -2,6 +2,7 @@
 import logging
 import time
 import uuid
+import zlib
 
 import orjson
 from fastapi import Request
@@ -112,13 +113,20 @@ async def get_session(session_id: str):
                 raise SessionNotFoundError(f"session not found: session_id={session_id}")
             return GetSessionResponse(session_id=session_id, records=[], metadata={})
         metadata = {}
-        try:
-            mismatch = registry.compute_session_mismatch(session)
-        except TokenizationError:
-            logger.exception("Failed to compute tito_session_mismatch for session %s", session_id)
-            mismatch = None
-        if mismatch is not None:
-            metadata["tito_session_mismatch"] = mismatch
+        # Re-tokenizing the full trajectory to check tito_session_mismatch is a
+        # diagnostic canary, not a training-correctness guard (train tokens come
+        # from the inference response). It runs synchronously here, so at long
+        # context it blocks the event loop and starves record retrieval. Sample
+        # a deterministic fraction by session_id to keep the signal cheap.
+        sample_rate = getattr(args, "tito_session_mismatch_sample_rate", 0.0)
+        if sample_rate > 0.0 and (zlib.crc32(session_id.encode()) % 10000) < sample_rate * 10000:
+            try:
+                mismatch = registry.compute_session_mismatch(session)
+            except TokenizationError:
+                logger.exception("Failed to compute tito_session_mismatch for session %s", session_id)
+                mismatch = None
+            if mismatch is not None:
+                metadata["tito_session_mismatch"] = mismatch
         metadata["accumulated_token_ids"] = session.token_ids
         metadata["max_trim_tokens"] = registry.tito_tokenizer.max_trim_tokens
         return GetSessionResponse(

miles/utils/arguments.py

@@ -1623,6 +1623,17 @@ def add_session_arguments(parser):
                 choices=[t.value for t in TITOTokenizerType],
                 help="TITO tokenizer type for pretokenized prefix reuse. Controls how token IDs are computed for messages appended after the pretokenized prefix in multi-turn agentic sessions.",
             )
+            parser.add_argument(
+                "--tito-session-mismatch-sample-rate",
+                type=float,
+                default=0.0,
+                help="Fraction of sessions (0.0-1.0) for which the session server re-tokenizes the full "
+                "trajectory to compute tito_session_mismatch. This is a diagnostic canary for chat-template/"
+                "tokenizer drift; it does not affect training tokens (those come from the inference response). "
+                "It runs synchronously on the record-fetch path, so at long context it blocks the event loop and "
+                "starves record retrieval. Default 0.0 (off). Set to 0.02 for a cheap drift signal, 1.0 for legacy "
+                "always-on behaviour.",
+            )
             parser.add_argument(
                 "--tito-allowed-append-roles",
                 nargs="+",