fix streamed response corruption in generation logs

arena/app.py

@@ -79,6 +79,7 @@
 	_leaderboard_summary,
 	_message_text_content,
 	_stats_footer,
+	_status_message,
 	_streaming_outputs,
 	_targeted_chatbot_value_updates,
 	_upsert_assistant_message,
@@ -810,12 +811,11 @@ def _apply_stream_chunk(
 		errored_slots.add(slot)
 		round_state["slot_logs"][slot]["status"] = "error"
 		round_state["slot_logs"][slot]["error"] = str(chunk.get("error") or "Unknown error")
-		error_prefix = "\n" if assistant_message_indices[slot] is not None else ""
-		assistant_message_indices[slot] = _upsert_assistant_message(
-			history=histories[slot],
-			message_index=assistant_message_indices[slot],
-			content=f"{error_prefix}[Error] {round_state['slot_logs'][slot]['error']}",
-			append=True,
+		histories[slot].append(
+			_status_message(
+				f"[Error] {round_state['slot_logs'][slot]['error']}",
+				"Generation Error",
+			)
 		)
 		reasoning_index = reasoning_message_indices[slot]
 		if reasoning_index is not None:
@@ -944,9 +944,7 @@ def _apply_stream_chunk(
 	if not changed_history and chunk.get("event") not in {"complete", "error"}:
 		return None
 
-	_finalize_round_state_logs(
-		round_state, histories, assistant_message_indices, reasoning_message_indices
-	)
+	_finalize_round_state_logs(round_state, histories, reasoning_message_indices)
 	return slot
 
 
@@ -958,9 +956,7 @@ def _finalize_generation_state(
 	completed_slots: set[int],
 	errored_slots: set[int],
 ) -> None:
-	_finalize_round_state_logs(
-		round_state, histories, assistant_message_indices, reasoning_message_indices
-	)
+	_finalize_round_state_logs(round_state, histories, reasoning_message_indices)
 	round_state["completed_slots"] = sorted(completed_slots)
 	round_state["errored_slots"] = sorted(errored_slots)
 	round_state["generation_completed_at"] = datetime.now(timezone.utc).isoformat()
@@ -1035,15 +1031,8 @@ async def stream_all_models(
 	for slot, request_metadata in enumerate(reasoning_requests):
 		reasoning_warning = request_metadata.get("reasoning_warning")
 		if isinstance(reasoning_warning, str) and reasoning_warning:
-			histories[slot].append(
-				{
-					"role": "assistant",
-					"content": f"[Warning] {reasoning_warning}",
-				}
-			)
-	_finalize_round_state_logs(
-		round_state, histories, assistant_message_indices, reasoning_message_indices
-	)
+			histories[slot].append(_status_message(f"[Warning] {reasoning_warning}", "Warning"))
+	_finalize_round_state_logs(round_state, histories, reasoning_message_indices)
 
 	yield _streaming_outputs(
 		user_input="",
@@ -1159,10 +1148,10 @@ async def stream_all_models(
 			round_state["slot_logs"][slot]["status"] = "error"
 			round_state["slot_logs"][slot]["error"] = GENERATION_INTERRUPTED_MESSAGE
 			histories[slot].append(
-				{
-					"role": "assistant",
-					"content": f"[Error] {GENERATION_INTERRUPTED_MESSAGE}",
-				}
+				_status_message(
+					f"[Error] {GENERATION_INTERRUPTED_MESSAGE}",
+					"Generation Error",
+				)
 			)
 			reasoning_index = reasoning_message_indices[slot]
 			if reasoning_index is not None:

arena/ui/display.py

@@ -69,10 +69,17 @@ def _serialize_history(history: list[Any]) -> list[dict[str, Any]]:
 	return [_serialize_message(message) for message in history]
 
 
+def _status_message(content: str, title: str) -> gr.ChatMessage:
+	return gr.ChatMessage(
+		role="assistant",
+		content=content,
+		metadata={"title": title, "status": "done"},
+	)
+
+
 def _finalize_round_state_logs(
 	round_state: dict[str, Any],
 	histories: list[list[Any]],
-	assistant_message_indices: list[int | None],
 	reasoning_message_indices: list[int | None],
 ) -> None:
 	slot_logs = round_state.get("slot_logs")
@@ -83,9 +90,6 @@ def _finalize_round_state_logs(
 		history = histories[slot]
 		slot_log = slot_logs[slot]
 		slot_log["message_history"] = _serialize_history(history)
-		assistant_index = assistant_message_indices[slot]
-		if isinstance(assistant_index, int) and 0 <= assistant_index < len(history):
-			slot_log["final_response"] = _message_text_content(history[assistant_index]).strip()
 		reasoning_index = reasoning_message_indices[slot]
 		if isinstance(reasoning_index, int) and 0 <= reasoning_index < len(history):
 			slot_log["reasoning_trace"] = _message_text_content(history[reasoning_index]).strip()

tests/integration/test_generation_flow.py

@@ -215,6 +215,15 @@ def test_apply_stream_chunk_finalizes_reasoning_on_error() -> None:
 			"usage": {"completion_tokens_details": {"reasoning_tokens": 3}},
 		},
 	)
+	app_module._apply_stream_chunk(
+		round_state,
+		histories,
+		assistant_message_indices,
+		reasoning_message_indices,
+		completed_slots,
+		errored_slots,
+		{"slot": 1, "delta": "Partial answer"},
+	)
 	slot = app_module._apply_stream_chunk(
 		round_state,
 		histories,
@@ -229,9 +238,13 @@ def test_apply_stream_chunk_finalizes_reasoning_on_error() -> None:
 	assert 1 in errored_slots
 	assert round_state["slot_logs"][1]["status"] == "error"
 	assert "provider failure" in str(round_state["slot_logs"][1]["error"])
+	assert round_state["slot_logs"][1]["final_response"] == "Partial answer"
 	assert isinstance(histories[1][1], gr.ChatMessage)
 	assert histories[1][1].metadata["status"] == "done"
-	assert histories[1][2]["content"] == "[Error] provider failure"
+	assert histories[1][2]["content"] == "Partial answer"
+	assert isinstance(histories[1][3], gr.ChatMessage)
+	assert histories[1][3].content == "[Error] provider failure"
+	assert histories[1][3].metadata["title"] == "Generation Error"
 
 
 def test_apply_stream_chunk_complete_records_usage_and_stats_footer() -> None:
@@ -666,6 +679,10 @@ async def test_stream_all_models_warns_when_selected_reasoning_is_unsupported(
 	assert final_round_state["slot_logs"][1]["message_history"][1]["content"].startswith(
 		"[Warning] Reasoning effort 'high' was selected"
 	)
+	assert final_round_state["slot_logs"][1]["message_history"][1]["metadata"] == {
+		"title": "Warning",
+		"status": "done",
+	}
 
 
 @pytest.mark.anyio
@@ -737,6 +754,13 @@ async def test_stream_all_models_recovers_from_background_stream_failure(monkeyp
 	assert final_round_state["errored_slots"] == [0, 1, 2]
 	assert final_round_state["slot_logs"][0]["final_response"] == "first chunk"
 	assert final_round_state["slot_logs"][0]["status"] == "error"
+	error_messages = [
+		message
+		for message in final_round_state["slot_logs"][0]["message_history"]
+		if message.get("metadata", {}).get("title") == "Generation Error"
+	]
+	assert len(error_messages) == 1
+	assert error_messages[0]["content"].startswith("[Error] Generation stopped")
 	assert (
 		final_round_state["slot_logs"][0]["error"]
 		== "Generation stopped before this round could finish."

tests/unit/ui/test_display.py

@@ -10,6 +10,7 @@
 	_reasoning_trace_title,
 	_serialize_history,
 	_stats_footer,
+	_status_message,
 	_upsert_assistant_message,
 	_upsert_reasoning_message,
 )
@@ -54,12 +55,20 @@ def test_message_text_content_and_serialize_history_handle_supported_types() ->
 	]
 
 
+def test_status_message_is_a_distinct_completed_chat_message() -> None:
+	message = _status_message("[Error] provider failure", "Generation Error")
+
+	assert message.role == "assistant"
+	assert message.content == "[Error] provider failure"
+	assert message.metadata == {"title": "Generation Error", "status": "done"}
+
+
 def test_finalize_round_state_logs_captures_histories_and_outputs() -> None:
 	round_state = {
 		"slot_logs": [
-			{"selection_slot": 0},
+			{"selection_slot": 0, "final_response": "Accumulated answer"},
 			{"selection_slot": 1},
-			{"selection_slot": 2},
+			{"selection_slot": 2, "final_response": "Another accumulated answer"},
 		]
 	}
 	histories = [
@@ -71,11 +80,11 @@ def test_finalize_round_state_logs_captures_histories_and_outputs() -> None:
 		[{"role": "assistant", "content": "Another answer"}],
 	]
 
-	_finalize_round_state_logs(round_state, histories, [0, None, 0], [1, None, None])
+	_finalize_round_state_logs(round_state, histories, [1, None, None])
 
-	assert round_state["slot_logs"][0]["final_response"] == "Final answer"
+	assert round_state["slot_logs"][0]["final_response"] == "Accumulated answer"
 	assert round_state["slot_logs"][0]["reasoning_trace"] == "Reasoning trail"
-	assert round_state["slot_logs"][2]["final_response"] == "Another answer"
+	assert round_state["slot_logs"][2]["final_response"] == "Another accumulated answer"
 	assert round_state["slot_logs"][1]["message_history"] == []