diff --git a/cli/main.py b/cli/main.py index 7e5ff3fb..60d5ba83 100644 --- a/cli/main.py +++ b/cli/main.py @@ -93,6 +93,7 @@ class MessageBuffer: self.agent_status = {} self.current_agent = None self.report_sections = {} + self.report_finalizers = {} self.selected_analysts = [] self._last_message_id = None @@ -119,9 +120,11 @@ class MessageBuffer: # Build report_sections dynamically self.report_sections = {} + self.report_finalizers = {} for section, (analyst_key, _) in self.REPORT_SECTIONS.items(): if analyst_key is None or analyst_key in self.selected_analysts: self.report_sections[section] = None + self.report_finalizers[section] = self.REPORT_SECTIONS[section][1] # Reset other state self.current_report = None @@ -145,8 +148,8 @@ class MessageBuffer: count = 0 for section, content in self.report_sections.items(): if content is not None: - section_info = self.REPORT_SECTIONS.get(section) - if section_info and self.agent_status.get(section_info[1]) == "completed": + finalizing_agent = self.report_finalizers.get(section) + if finalizing_agent and self.agent_status.get(finalizing_agent) == "completed": count += 1 return count diff --git a/tests/cli/test_main.py b/tests/cli/test_main.py index dfe8fe01..bee14a29 100644 --- a/tests/cli/test_main.py +++ b/tests/cli/test_main.py @@ -1,5 +1,5 @@ import pytest -from cli.main import extract_content_string +from cli.main import MessageBuffer, extract_content_string def test_extract_content_string_empty(): assert extract_content_string(None) is None @@ -25,3 +25,28 @@ def test_extract_content_string_dict(): def test_extract_content_string_list(): assert extract_content_string([{"type": "text", "text": "hello"}, " world"]) == "hello world" + + +def test_message_buffer_completed_reports_count_tracks_active_sections_only(): + buffer = MessageBuffer() + buffer.init_for_analysis(["market", "news"]) + + buffer.report_sections["market_report"] = "done" + buffer.report_sections["news_report"] = "done" + buffer.report_sections["investment_plan"] = "done" + buffer.agent_status["Market Analyst"] = "completed" + buffer.agent_status["News Analyst"] = "pending" + buffer.agent_status["Research Manager"] = "completed" + + assert buffer.get_completed_reports_count() == 2 + + +def test_message_buffer_completed_reports_count_ignores_unknown_section(): + buffer = MessageBuffer() + buffer.init_for_analysis(["market"]) + + buffer.report_sections["market_report"] = "done" + buffer.report_sections["unexpected_report"] = "done" + buffer.agent_status["Market Analyst"] = "completed" + + assert buffer.get_completed_reports_count() == 1