fix(server): tolerate malformed responses in triplet export

agentlightning/server/routes/events.py

@@ -113,23 +113,31 @@ def _trim_model_request(data: dict[str, Any]) -> dict[str, Any]:
     if isinstance(resp, dict):
         prompt_token_ids = resp.get("prompt_token_ids", [])
         choices = resp.get("choices", [])
-        if choices:
+        if isinstance(choices, list) and choices and isinstance(choices[0], dict):
+            choice = choices[0]
             if not prompt_token_ids:
-                prompt_token_ids = choices[0].get("prompt_token_ids", [])
-            response_token_ids = choices[0].get("token_ids", [])
-            response_log_probs = _extract_choice_log_probs(choices[0])
+                prompt_token_ids = choice.get("prompt_token_ids", [])
+            response_token_ids = choice.get("token_ids", [])
+            response_log_probs = _extract_choice_log_probs(choice)
     elif isinstance(resp, list):
         # Legacy: raw SSE chunks (pre-assembly format, backward compat).
         for chunk in resp:
+            if not isinstance(chunk, dict):
+                continue
             if not prompt_token_ids and chunk.get("prompt_token_ids"):
                 prompt_token_ids = chunk["prompt_token_ids"]
             choices = chunk.get("choices", [])
-            if choices:
+            if isinstance(choices, list) and choices and isinstance(choices[0], dict):
                 tids = choices[0].get("token_ids")
-                if tids:
+                if isinstance(tids, list):
                     response_token_ids.extend(tids)
 
+    if not isinstance(response_token_ids, list):
+        response_token_ids = []
+
     srv = data.get("server", {})
+    if not isinstance(srv, dict):
+        srv = {}
     trimmed = {
         "prompt_token_ids": prompt_token_ids,
         "response_token_ids": response_token_ids,

tests/server/test_endpoints.py

@@ -286,6 +286,82 @@ def post_model_request(prompt_token_ids: object, response_token_ids: list[int])
     assert [event["data"]["response_token_ids"] for event in model_requests] == [[10], [20], [50]]
 
 
+@pytest.mark.parametrize(
+    ("response_data", "server", "expected_prompt_ids", "expected_response_ids"),
+    [
+        pytest.param({"choices": [None]}, {"model": MODEL_NAME, "version": 3}, [], [], id="null-choice"),
+        pytest.param([None], {"model": MODEL_NAME, "version": 3}, [], [], id="null-legacy-chunk"),
+        pytest.param({"choices": "invalid"}, {"model": MODEL_NAME, "version": 3}, [], [], id="invalid-choices"),
+        pytest.param(
+            {"choices": [{"token_ids": 2}]},
+            {"model": MODEL_NAME, "version": 3},
+            [],
+            [],
+            id="invalid-non-stream-token-ids",
+        ),
+        pytest.param(
+            [{"choices": [{"token_ids": 2}]}],
+            {"model": MODEL_NAME, "version": 3},
+            [],
+            [],
+            id="invalid-legacy-token-ids",
+        ),
+        pytest.param({}, None, [], [], id="null-server"),
+        pytest.param(
+            [
+                None,
+                {"prompt_token_ids": [1], "choices": [None]},
+                {"choices": [{"token_ids": [2]}]},
+            ],
+            {"model": MODEL_NAME, "version": 3},
+            [1],
+            [2],
+            id="mixed-legacy-chunks",
+        ),
+    ],
+)
+def test_triplet_events_tolerate_malformed_response_shapes(
+    client: TestClient,
+    auth_headers: dict[str, str],
+    response_data: object,
+    server: object,
+    expected_prompt_ids: list[int],
+    expected_response_ids: list[int],
+):
+    rollout = _rollout(client, auth_headers)
+    rollout_id = rollout["rollout_id"]
+    posted = client.post(
+        f"/api/rollouts/{rollout_id}/attempt/0/events",
+        json={
+            "event_type": "model_request",
+            "data": {
+                "response": response_data,
+                "server": server,
+                "http_status": 502,
+                "status": "error",
+            },
+        },
+        headers=auth_headers,
+    )
+    assert posted.status_code == 200
+
+    response = client.get(
+        f"/api/rollouts/{rollout_id}/events",
+        params={"event_type": "model_request", "format": "triplet"},
+        headers=auth_headers,
+    )
+
+    assert response.status_code == 200
+    data = response.json()[0]["data"]
+    assert data["prompt_token_ids"] == expected_prompt_ids
+    assert data["response_token_ids"] == expected_response_ids
+    assert data["response_log_probs"] is None
+    assert data["http_status"] == 502
+    assert data["status"] == "error"
+    expected_server = server if isinstance(server, dict) else {"model": None, "version": None}
+    assert data["server"] == expected_server
+
+
 def test_model_endpoints(client: TestClient, auth_headers: dict[str, str]):
     created = client.post(
         "/api/models",