From 8a8b3e2009de69d86600937bf64f230ceb35a11c Mon Sep 17 00:00:00 2001 From: Rio Yu <52408936+rioyu123@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:08:00 +0800 Subject: [PATCH 1/2] fix(server): tolerate malformed responses in triplet export Signed-off-by: Rio Yu <52408936+rioyu123@users.noreply.github.com> --- agentlightning/server/routes/events.py | 17 ++++--- tests/server/test_endpoints.py | 69 ++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 6 deletions(-) diff --git a/agentlightning/server/routes/events.py b/agentlightning/server/routes/events.py index 291023ef0..28f83dd9c 100644 --- a/agentlightning/server/routes/events.py +++ b/agentlightning/server/routes/events.py @@ -113,23 +113,28 @@ 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) srv = data.get("server", {}) + if not isinstance(srv, dict): + srv = {} trimmed = { "prompt_token_ids": prompt_token_ids, "response_token_ids": response_token_ids, diff --git a/tests/server/test_endpoints.py b/tests/server/test_endpoints.py index 81bf82f70..3140e8d81 100644 --- a/tests/server/test_endpoints.py +++ b/tests/server/test_endpoints.py @@ -286,6 +286,75 @@ 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-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", From 58c0a4df0e2dad1024b7c84de64ada1bef4b1601 Mon Sep 17 00:00:00 2001 From: Rio Yu <52408936+rioyu123@users.noreply.github.com> Date: Sun, 30 Aug 2026 00:27:17 +0800 Subject: [PATCH 2/2] fix(server): normalize malformed response token IDs Signed-off-by: Rio Yu <52408936+rioyu123@users.noreply.github.com> --- agentlightning/server/routes/events.py | 3 +++ tests/server/test_endpoints.py | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/agentlightning/server/routes/events.py b/agentlightning/server/routes/events.py index 28f83dd9c..56b2e37f9 100644 --- a/agentlightning/server/routes/events.py +++ b/agentlightning/server/routes/events.py @@ -132,6 +132,9 @@ def _trim_model_request(data: dict[str, Any]) -> dict[str, Any]: 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 = {} diff --git a/tests/server/test_endpoints.py b/tests/server/test_endpoints.py index 3140e8d81..192b93656 100644 --- a/tests/server/test_endpoints.py +++ b/tests/server/test_endpoints.py @@ -292,6 +292,13 @@ def post_model_request(prompt_token_ids: object, response_token_ids: list[int]) 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},