fix: AG2 single-agent coordination issues

[db-ol]

Summary

• Fix AttributeError when using AG2 adapter in single-agent mode (workflow_tools not defined)
• Add coordination stage handling for AG2 single-agent execution (INITIAL_ANSWER, ENFORCEMENT, PRESENTATION)
• Fix vote rejection for single agents by accepting votes when agent has already produced an answer

Test plan

• Tested with ag2_coder.yaml config
• Verified code execution works (factorial, prime number checks)
• Verified coordination completes successfully with single agent
• Manual test: uv run python -m massgen.cli --automation --config massgen/configs/ag2/ag2_coder.yaml "Create a Python function that calculates factorial of 5 and print the result"
• Expected output: Code executes, returns "Factorial of 5 is: 120", coordination completes

Changes

massgen/adapters/ag2_adapter.py

• Added _execute_single_agent_with_coordination() method to handle MassGen coordination stages
• Guarded unregister_tools_for_agent call with if self.is_group_chat

massgen/orchestrator.py

• Added exception for single-agent vote handling to accept votes when agent has its own answer

Summary by CodeRabbit

• Bug Fixes

  • Resolved stalling where single-agent votes on its own answer were ignored; such votes are now accepted so progression continues.
  • Fixed tool cleanup so workflow tools are unregistered only in group-chat mode, preventing premature cleanup in single-agent flows.

• New Features

  • Enhanced single-agent execution with coordination-stage handling to improve streaming behavior and state logging.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Warning

Rate limit exceeded

@db-ol has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 10 minutes and 59 seconds before requesting another review.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

📝 Walkthrough

Walkthrough

Added a coordination-aware single-agent streaming execution path in AG2Adapter that handles INITIAL_ANSWER, ENFORCEMENT, and PRESENTATION stages; updated streaming cleanup to only unregister workflow tools in group-chat mode. Orchestrator vote handling now accepts a single-agent vote when restart_pending exists but the agent already has an answer.

Changes

Cohort / File(s)	Summary
Coordination-based single-agent execution massgen/adapters/ag2_adapter.py	Added _execute_single_agent_with_coordination() to handle coordination stages (store INITIAL_ANSWER, synthesize self-vote during ENFORCEMENT, normal PRESENTATION), integrated it into streaming execution for single-agent flows, added json and uuid imports, and limited workflow tool cleanup to group-chat mode.
Orchestrator single-agent vote handling massgen/orchestrator.py	Special-cased vote handling: when there is exactly one agent, restart_pending is true, and that agent already has an answer, clear restart_pending and accept the vote (log acceptance) instead of ignoring it. No public signatures changed.

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant AG2Adapter
    participant Orchestrator
    participant Tools

    Client->>AG2Adapter: start streaming single-agent run (messages, tools)
    AG2Adapter->>AG2Adapter: _execute_single_agent_with_coordination()
    AG2Adapter->>AG2Adapter: INITIAL_ANSWER -> store last_answer
    AG2Adapter->>Tools: ENFORCEMENT -> synthesize vote tool call (self)
    Tools->>Orchestrator: submit vote
    Orchestrator->>Orchestrator: if single-agent & restart_pending & has_answer -> clear restart_pending, accept vote
    Orchestrator->>AG2Adapter: vote accepted / proceed
    AG2Adapter->>Client: PRESENTATION -> stream final chunks

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Documentation Updated	⚠️ Warning	PR lacks design documentation in docs/dev_notes/ for the complex coordination stage handling architecture despite having method docstrings and inline comments.	Add a design document explaining the coordination stage handling approach, its rationale, and orchestrator interactions.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: fixing AG2 single-agent coordination issues, which aligns with all three key fixes in the changeset.
Description check	✅ Passed	The description is well-structured with summary, changes, and test plan sections. However, it lacks several required template sections including Type of change checklist, Pre-commit status, and a How to Test section with bash commands.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Capabilities Registry Check	✅ Passed	The PR does not introduce any backend or model changes. Modifications are exclusively focused on AG2 adapter coordination handling and orchestrator vote logic adjustments.
Config Parameter Sync	✅ Passed	PR does not introduce new YAML configuration parameters; coordination stages are hardcoded constants, not user-configurable parameters.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@massgen/adapters/ag2_adapter.py`:
- Around line 330-352: The docstring is out of sync with the implementation:
when coordination_stage == CoordinationStage.ENFORCEMENT the code synthesizes a
"vote" tool call rather than a "new_answer" and never uses the stored
_last_single_agent_answer; update the docstring near coordination_stage handling
to describe that ENFORCEMENT synthesizes a "vote" tool call for a single agent
(including the reason text) and either (A) change the implementation to
synthesize a "new_answer" tool call that includes _last_single_agent_answer if
that behavior is intended, or (B) keep the current "vote" behavior and
remove/clarify the stored _last_single_agent_answer usage (or reference it in
the vote payload) so code and docstring match (refer to coordination_stage,
CoordinationStage.ENFORCEMENT, _last_single_agent_answer, and the
vote/new_answer tool call).

🧹 Nitpick comments (2)

massgen/adapters/ag2_adapter.py (2)

306-308: Unused tools parameter.

The tools parameter is declared but never used in this method. Either prefix it with underscore to indicate it's intentionally unused, or remove it if not needed.

Suggested fix

     async def _execute_single_agent_with_coordination(
-        self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]]
+        self, messages: List[Dict[str, Any]], _tools: List[Dict[str, Any]]
     ) -> AsyncGenerator[StreamChunk, None]:

---

326-326: Consider initializing _last_single_agent_answer in __init__.

The attribute _last_single_agent_answer is dynamically created (line 335) and checked with hasattr here. For better clarity and IDE support, consider initializing it in __init__ or _setup_single_agent.

Suggested initialization in _setup_single_agent

     def _setup_single_agent(self):
         """Set up a single AG2 agent."""
         self.agent = setup_agent_from_config(self.agent_config)
-
         self.is_group_chat = False
+        self._last_single_agent_answer: str | None = None

Then update the check:

-            {"stage": str(self.coordination_stage), "has_stored_answer": hasattr(self, "_last_single_agent_answer")},
+            {"stage": str(self.coordination_stage), "has_stored_answer": self._last_single_agent_answer is not None},