feat(delegate): Add create_sub_visualizer method for custom sub-agent visualization (#1767)

Co-authored-by: openhands <[email protected]>

Author: malhotra5

openhands-sdk/openhands/sdk/conversation/visualizer/base.py

@@ -65,3 +65,26 @@ def on_event(self, event: Event) -> None:
             event: The event to visualize
         """
         pass
+
+    def create_sub_visualizer(
+        self,
+        agent_id: str,  # noqa: ARG002
+    ) -> "ConversationVisualizerBase | None":
+        """Create a visualizer for a sub-agent during delegation.
+
+        Override this method to support sub-agent visualization in multi-agent
+        delegation scenarios. The sub-visualizer will be used to display events
+        from the spawned sub-agent.
+
+        By default, returns None which means sub-agents will not have visualization.
+        Subclasses that support delegation (like DelegationVisualizer) should
+        override this method to create appropriate sub-visualizers.
+
+        Args:
+            agent_id: The identifier of the sub-agent being spawned
+
+        Returns:
+            A visualizer instance for the sub-agent, or None if sub-agent
+            visualization is not supported
+        """
+        return None

openhands-tools/openhands/tools/delegate/impl.py

@@ -9,7 +9,6 @@
 from openhands.sdk.tool.tool import ToolExecutor
 from openhands.tools.delegate.definition import DelegateObservation
 from openhands.tools.delegate.registration import get_agent_factory
-from openhands.tools.delegate.visualizer import DelegationVisualizer
 
 
 if TYPE_CHECKING:
@@ -130,14 +129,12 @@ def _spawn_agents(self, action: "DelegateAction") -> DelegateObservation:
                 factory = get_agent_factory(agent_type)
                 worker_agent = factory.factory_func(sub_agent_llm)
 
-                if isinstance(parent_visualizer, DelegationVisualizer):
-                    sub_visualizer = DelegationVisualizer(
-                        name=agent_id,
-                        highlight_regex=parent_visualizer._highlight_patterns,
-                        skip_user_messages=parent_visualizer._skip_user_messages,
-                    )
-                else:
-                    sub_visualizer = None
+                # Use parent visualizer's create_sub_visualizer method if available
+                # This allows custom visualizers (e.g., TUI-based) to create
+                # appropriate sub-visualizers for their environment
+                sub_visualizer = None
+                if parent_visualizer is not None:
+                    sub_visualizer = parent_visualizer.create_sub_visualizer(agent_id)
 
                 sub_conversation = LocalConversation(
                     agent=worker_agent,
@@ -211,13 +208,13 @@ def _delegate_tasks(self, action: "DelegateAction") -> "DelegateObservation":
             results = {}
             errors = {}
 
-            # Get the parent agent's name from the visualizer
+            # Get the parent agent's name from the visualizer if available
             parent_conversation = self.parent_conversation
             parent_name = None
             if hasattr(parent_conversation, "_visualizer"):
                 visualizer = parent_conversation._visualizer
-                if isinstance(visualizer, DelegationVisualizer):
-                    parent_name = visualizer._name
+                if visualizer is not None:
+                    parent_name = getattr(visualizer, "_name", None)
 
             def run_task(
                 agent_id: str,

openhands-tools/openhands/tools/delegate/visualizer.py

@@ -61,6 +61,24 @@ def __init__(
         )
         self._name = name
 
+    def create_sub_visualizer(self, agent_id: str) -> "DelegationVisualizer":
+        """Create a visualizer for a sub-agent during delegation.
+
+        Creates a new DelegationVisualizer instance for the sub-agent with
+        the same configuration as the parent visualizer.
+
+        Args:
+            agent_id: The identifier of the sub-agent being spawned
+
+        Returns:
+            A new DelegationVisualizer configured for the sub-agent
+        """
+        return DelegationVisualizer(
+            name=agent_id,
+            highlight_regex=self._highlight_patterns,
+            skip_user_messages=self._skip_user_messages,
+        )
+
     @staticmethod
     def _format_agent_name(name: str) -> str:
         """

tests/sdk/conversation/test_visualizer.py

@@ -504,3 +504,15 @@ def test_visualizer_conversation_state_update_event_skipped():
     block = visualizer._create_event_block(event)
     # Should return None to skip visualization
     assert block is None
+
+
+def test_default_visualizer_create_sub_visualizer_returns_none():
+    """Test that DefaultConversationVisualizer.create_sub_visualizer returns None.
+
+    This is the expected default behavior - base visualizers don't support
+    sub-agent visualization. Subclasses like DelegationVisualizer can override
+    this to provide sub-agent visualizers.
+    """
+    visualizer = DefaultConversationVisualizer()
+    result = visualizer.create_sub_visualizer("test_agent")
+    assert result is None

tests/tools/delegate/test_visualizer.py

@@ -211,3 +211,36 @@ def test_delegation_visualizer_observation_event():
     assert block is not None
     # The block contains the Rule as the first element with the title
     assert "Main Delegator Agent Observation" in str(block.renderables[0])
+
+
+def test_delegation_visualizer_create_sub_visualizer():
+    """Test create_sub_visualizer creates a new visualizer for sub-agents."""
+    parent_visualizer = DelegationVisualizer(
+        name="main_delegator",
+        highlight_regex={"test": "bold"},
+        skip_user_messages=True,
+    )
+
+    # Create sub-visualizer for a sub-agent
+    sub_visualizer = parent_visualizer.create_sub_visualizer("lodging_expert")
+
+    # Verify sub-visualizer is a DelegationVisualizer
+    assert isinstance(sub_visualizer, DelegationVisualizer)
+    # Verify sub-visualizer has the correct agent name
+    assert sub_visualizer._name == "lodging_expert"
+    # Verify settings are inherited from parent
+    assert sub_visualizer._highlight_patterns == {"test": "bold"}
+    assert sub_visualizer._skip_user_messages is True
+
+
+def test_delegation_visualizer_create_sub_visualizer_with_defaults():
+    """Test create_sub_visualizer works with default parent settings."""
+    parent_visualizer = DelegationVisualizer(name="parent")
+
+    sub_visualizer = parent_visualizer.create_sub_visualizer("child_agent")
+
+    assert isinstance(sub_visualizer, DelegationVisualizer)
+    assert sub_visualizer._name == "child_agent"
+    # Default values should be inherited
+    assert sub_visualizer._highlight_patterns is not None  # Has default patterns
+    assert sub_visualizer._skip_user_messages is False