Keep monotonic_ulid() monotonic when the clock does not move forward

llm/utils.py

@@ -722,8 +722,9 @@ def monotonic_ulid() -> ULID:
     other ULID returned by this function inside the same process.
 
     It works the same way the reference JavaScript `monotonicFactory` does:
-    * If the current call happens in the same millisecond as the previous
-        one, the 80-bit randomness part is incremented by exactly one.
+    * If the current call reads a millisecond that is not later than the
+        previous one, the 80-bit randomness part is incremented by exactly
+        one and the previous timestamp is reused.
     * As soon as the system clock moves forward, a brand-new ULID with
         cryptographically secure randomness is generated.
     * If more than 2**80 ULIDs are requested within a single millisecond
@@ -742,8 +743,10 @@ def monotonic_ulid() -> ULID:
         # Decode timestamp from the last ULID we handed out
         last_ms = int.from_bytes(_last[:TIMESTAMP_LEN], "big")
 
-        # If the millisecond is the same, increment the randomness
-        if now_ms == last_ms:
+        # If the clock did not move forward, increment the randomness and keep
+        # the previous timestamp. `time.time_ns()` is read outside the lock and
+        # the wall clock can step backwards, so now_ms may be below last_ms.
+        if now_ms <= last_ms:
             rand_int = int.from_bytes(_last[TIMESTAMP_LEN:], "big") + 1
             if rand_int >= 1 << (RANDOMNESS_LEN * 8):
                 raise OverflowError(

tests/test_utils.py

@@ -1,7 +1,10 @@
 import json
+import threading
+from types import SimpleNamespace
 
 import pytest
 
+import llm.utils
 from llm import Toolbox, get_key
 from llm.utils import (
     extract_fenced_code_block,
@@ -510,6 +513,46 @@ def test_monotonic_ulids():
     assert ulids == sorted(ulids)
 
 
+def _pin_clock(monkeypatch, time_ns):
+    monkeypatch.setattr(llm.utils, "_last", None)
+    monkeypatch.setattr(llm.utils, "time", SimpleNamespace(time_ns=time_ns))
+
+
+def test_monotonic_ulids_when_clock_steps_backwards(monkeypatch):
+    readings = iter([2_000_000_000, 1_999_000_000])
+    _pin_clock(monkeypatch, lambda: next(readings))
+    first = monotonic_ulid()
+    second = monotonic_ulid()
+    assert second > first
+    # The earlier timestamp is discarded rather than encoded into the ULID
+    assert second.timestamp == first.timestamp
+
+
+def test_monotonic_ulids_across_threads(monkeypatch):
+    # time_ns() is read outside the lock, so the thread holding the earlier
+    # reading can reach the lock after the thread holding the later one
+    readings = {"early": 3_000_000_000, "late": 3_001_000_000}
+    _pin_clock(monkeypatch, lambda: readings[threading.current_thread().name])
+
+    late_is_done = threading.Event()
+    generated = {}
+
+    def generate():
+        name = threading.current_thread().name
+        if name == "early":
+            assert late_is_done.wait(timeout=10)
+        generated[name] = monotonic_ulid()
+        late_is_done.set()
+
+    threads = [threading.Thread(target=generate, name=name) for name in readings]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=10)
+
+    assert generated["early"] > generated["late"]
+
+
 def test_toolbox_config_capture():
     """Test that Toolbox captures __init__ parameters in _config"""