From b220c7cb45a6984390b3633b5f6ae67021d2fd26 Mon Sep 17 00:00:00 2001 From: Patrick Lacey Date: Thu, 6 Aug 2026 00:10:07 -0400 Subject: [PATCH 1/2] fix(logging): preserve streamed final responses --- arena/app.py | 12 +++--------- arena/ui/display.py | 4 ---- tests/integration/test_generation_flow.py | 12 +++++++++++- tests/unit/ui/test_display.py | 10 +++++----- 4 files changed, 19 insertions(+), 19 deletions(-) diff --git a/arena/app.py b/arena/app.py index 06ee40b..8fc27f2 100644 --- a/arena/app.py +++ b/arena/app.py @@ -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() @@ -1041,9 +1037,7 @@ async def stream_all_models( "content": f"[Warning] {reasoning_warning}", } ) - _finalize_round_state_logs( - round_state, histories, assistant_message_indices, reasoning_message_indices - ) + _finalize_round_state_logs(round_state, histories, reasoning_message_indices) yield _streaming_outputs( user_input="", diff --git a/arena/ui/display.py b/arena/ui/display.py index d03fc56..41c3dcc 100644 --- a/arena/ui/display.py +++ b/arena/ui/display.py @@ -72,7 +72,6 @@ def _serialize_history(history: list[Any]) -> list[dict[str, Any]]: 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 +82,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() diff --git a/tests/integration/test_generation_flow.py b/tests/integration/test_generation_flow.py index ef83fdb..77989f4 100644 --- a/tests/integration/test_generation_flow.py +++ b/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,10 @@ 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\n[Error] provider failure" def test_apply_stream_chunk_complete_records_usage_and_stats_footer() -> None: diff --git a/tests/unit/ui/test_display.py b/tests/unit/ui/test_display.py index 210fb74..d87c1f3 100644 --- a/tests/unit/ui/test_display.py +++ b/tests/unit/ui/test_display.py @@ -57,9 +57,9 @@ def test_message_text_content_and_serialize_history_handle_supported_types() -> 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 +71,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"] == [] From 7d30e0ca03887e95d8059a7cb31562555929c66d Mon Sep 17 00:00:00 2001 From: Patrick Lacey Date: Thu, 6 Aug 2026 00:12:12 -0400 Subject: [PATCH 2/2] fix(ui): separate generation status messages --- arena/app.py | 27 +++++++++-------------- arena/ui/display.py | 8 +++++++ tests/integration/test_generation_flow.py | 16 +++++++++++++- tests/unit/ui/test_display.py | 9 ++++++++ 4 files changed, 43 insertions(+), 17 deletions(-) diff --git a/arena/app.py b/arena/app.py index 8fc27f2..b34c4ef 100644 --- a/arena/app.py +++ b/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: @@ -1031,12 +1031,7 @@ 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}", - } - ) + histories[slot].append(_status_message(f"[Warning] {reasoning_warning}", "Warning")) _finalize_round_state_logs(round_state, histories, reasoning_message_indices) yield _streaming_outputs( @@ -1153,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: diff --git a/arena/ui/display.py b/arena/ui/display.py index 41c3dcc..4306da0 100644 --- a/arena/ui/display.py +++ b/arena/ui/display.py @@ -69,6 +69,14 @@ 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]], diff --git a/tests/integration/test_generation_flow.py b/tests/integration/test_generation_flow.py index 77989f4..5cf4166 100644 --- a/tests/integration/test_generation_flow.py +++ b/tests/integration/test_generation_flow.py @@ -241,7 +241,10 @@ def test_apply_stream_chunk_finalizes_reasoning_on_error() -> None: 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"] == "Partial answer\n[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: @@ -676,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 @@ -747,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." diff --git a/tests/unit/ui/test_display.py b/tests/unit/ui/test_display.py index d87c1f3..b4c5172 100644 --- a/tests/unit/ui/test_display.py +++ b/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,6 +55,14 @@ 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": [