From a01f41fa9c03499c46e0f491ff124e5fd2615635 Mon Sep 17 00:00:00 2001
From: yixi ning <ningyixi0@gmail.com>
Date: Sat, 17 Oct 2026 08:28:08 +0000
Subject: [PATCH] fix(server): tolerate overflowing logprobs in triplet queries

---
 agentlightning/server/routes/events.py |  2 +-
 tests/server/test_endpoints.py         | 33 ++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/agentlightning/server/routes/events.py b/agentlightning/server/routes/events.py
index 5afad7974..71ff02d1f 100644
--- a/agentlightning/server/routes/events.py
+++ b/agentlightning/server/routes/events.py
@@ -114,7 +114,7 @@ def _extract_choice_log_probs(choice: dict[str, Any]) -> list[float] | None:
     for v in raw:
         try:
             f = float(v)
-        except (TypeError, ValueError):
+        except (TypeError, ValueError, OverflowError):
             return None
         if not math.isfinite(f):
             return None
diff --git a/tests/server/test_endpoints.py b/tests/server/test_endpoints.py
index 55ad9f549..695ba81c4 100644
--- a/tests/server/test_endpoints.py
+++ b/tests/server/test_endpoints.py
@@ -166,6 +166,39 @@ def test_delete_rollout(client: TestClient, auth_headers: dict[str, str]):
     assert client.delete(f"/api/rollouts/{rollout_id}", headers=auth_headers).status_code == 204
 
 
+@pytest.mark.parametrize(
+    "logprobs",
+    [
+        {"content": [{"logprob": -0.5}, {"logprob": 10**400}]},
+        {"token_logprobs": [-0.5, 10**400]},
+    ],
+    ids=["chat", "completions"],
+)
+def test_triplet_events_ignore_overflowing_logprobs(client: TestClient, auth_headers: dict[str, str], logprobs: dict):
+    rollout_id = _rollout(client, auth_headers)["rollout_id"]
+    posted = client.post(
+        f"/api/rollouts/{rollout_id}/attempt/0/events",
+        json={
+            "event_type": "model_request",
+            "data": {
+                "response": {
+                    "prompt_token_ids": [1, 2],
+                    "choices": [{"token_ids": [3, 4], "logprobs": logprobs}],
+                },
+            },
+        },
+        headers=auth_headers,
+    )
+    assert posted.status_code == 200
+
+    triplets = client.get(f"/api/rollouts/{rollout_id}/events?format=triplet", headers=auth_headers)
+    assert triplets.status_code == 200
+    data = triplets.json()[0]["data"]
+    assert data["prompt_token_ids"] == [1, 2]
+    assert data["response_token_ids"] == [3, 4]
+    assert data["response_log_probs"] is None
+
+
 def test_enqueue_with_client_rollout_id_is_idempotent(client: TestClient, auth_headers: dict[str, str]):
     rid = "fixed-rollout-id-123"
     created = client.post(