massgen 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

massgen/backend/claude.py

@@ -37,7 +37,12 @@ from ..formatter import ClaudeFormatter
 from ..logger_config import log_backend_agent_message, log_stream_chunk, logger
 from ..mcp_tools.backend_utils import MCPErrorHandler
 from .base import FilesystemSupport, StreamChunk
-from .base_with_custom_tool_and_mcp import CustomToolAndMCPBackend, UploadFileError
+from .base_with_custom_tool_and_mcp import (
+    CustomToolAndMCPBackend,
+    CustomToolChunk,
+    ToolExecutionConfig,
+    UploadFileError,
+)
 
 
 class ClaudeBackend(CustomToolAndMCPBackend):
@@ -521,6 +526,87 @@ class ClaudeBackend(CustomToolAndMCPBackend):
         async for chunk in self._process_stream(stream, all_params, agent_id):
             yield chunk
 
+    def _append_tool_result_message(
+        self,
+        updated_messages: List[Dict[str, Any]],
+        call: Dict[str, Any],
+        result: Any,
+        tool_type: str,
+    ) -> None:
+        """Append tool result to messages in Claude format.
+
+        Args:
+            updated_messages: Message list to append to
+            call: Tool call dictionary with call_id, name, arguments
+            result: Tool execution result
+            tool_type: "custom" or "mcp"
+
+        Note:
+            Claude uses tool_result format with tool_use_id.
+        """
+        tool_result_msg = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "tool_result",
+                    "tool_use_id": call.get("call_id", "") or call.get("id", ""),
+                    "content": str(result),
+                },
+            ],
+        }
+        updated_messages.append(tool_result_msg)
+
+    def _append_tool_error_message(
+        self,
+        updated_messages: List[Dict[str, Any]],
+        call: Dict[str, Any],
+        error_msg: str,
+        tool_type: str,
+    ) -> None:
+        """Append tool error to messages in Claude format.
+
+        Args:
+            updated_messages: Message list to append to
+            call: Tool call dictionary with call_id, name, arguments
+            error_msg: Error message string
+            tool_type: "custom" or "mcp"
+
+        Note:
+            Claude uses tool_result format with tool_use_id for errors too.
+        """
+        error_result_msg = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "tool_result",
+                    "tool_use_id": call.get("call_id", "") or call.get("id", ""),
+                    "content": error_msg,
+                },
+            ],
+        }
+        updated_messages.append(error_result_msg)
+
+    async def _execute_custom_tool(self, call: Dict[str, Any]) -> AsyncGenerator[CustomToolChunk, None]:
+        """Execute custom tool with streaming support - async generator for base class.
+
+        This method is called by _execute_tool_with_logging and yields CustomToolChunk
+        objects for intermediate streaming output. The base class detects the async
+        generator and streams intermediate results to users in real-time.
+
+        Args:
+            call: Tool call dictionary with name and arguments
+
+        Yields:
+            CustomToolChunk objects with streaming data
+
+        Note:
+            - Intermediate chunks (completed=False) are streamed to users in real-time
+            - Final chunk (completed=True) contains the accumulated result for message history
+            - The base class automatically handles extracting and displaying intermediate chunks
+        """
+        async for chunk in self.stream_custom_tool_execution(call):
+            yield chunk
+
     async def _stream_with_custom_and_mcp_tools(
         self,
         current_messages: List[Dict[str, Any]],
@@ -660,8 +746,9 @@ class ClaudeBackend(CustomToolAndMCPBackend):
                 elif event.type == "message_delta":
                     pass
                 elif event.type == "message_stop":
-                    # Identify MCP, custom, and non-MCP/non-custom tool calls among current_tool_uses
-                    non_mcp_non_custom_tool_calls = []
+                    captured_calls = []
+                    tool_use_by_name: Dict[str, Dict[str, Any]] = {}
+
                     if current_tool_uses:
                         for tool_use in current_tool_uses.values():
                             tool_name = tool_use.get("name", "")
@@ -675,39 +762,68 @@ class ClaudeBackend(CustomToolAndMCPBackend):
                             except json.JSONDecodeError:
                                 parsed_input = {"raw_input": tool_input}
 
-                            if self.is_mcp_tool_call(tool_name):
-                                mcp_tool_calls.append(
-                                    {
-                                        "id": tool_use["id"],
-                                        "type": "function",
-                                        "function": {
-                                            "name": tool_name,
-                                            "arguments": parsed_input,
-                                        },
+                            captured_calls.append(
+                                {
+                                    "name": tool_name,
+                                    "arguments": json.dumps(parsed_input) if isinstance(parsed_input, dict) else str(parsed_input),
+                                    "call_id": tool_use["id"],
+                                },
+                            )
+
+                            # Store tool_use info for reconstruction
+                            tool_use_by_name[tool_use["id"]] = {
+                                "id": tool_use["id"],
+                                "parsed_input": parsed_input,
+                            }
+
+                    # Use helper to categorize tool calls
+                    if captured_calls:
+                        categorized_mcp, categorized_custom, categorized_provider = self._categorize_tool_calls(captured_calls)
+
+                        # Reconstruct Claude-specific format for each category
+                        for call in categorized_mcp:
+                            tool_info = tool_use_by_name[call["call_id"]]
+                            mcp_tool_calls.append(
+                                {
+                                    "id": tool_info["id"],
+                                    "type": "function",
+                                    "function": {
+                                        "name": call["name"],
+                                        "arguments": tool_info["parsed_input"],
                                     },
-                                )
-                            elif self.is_custom_tool_call(tool_name):
-                                custom_tool_calls.append(
-                                    {
-                                        "id": tool_use["id"],
-                                        "type": "function",
-                                        "function": {
-                                            "name": tool_name,
-                                            "arguments": parsed_input,
-                                        },
+                                },
+                            )
+
+                        for call in categorized_custom:
+                            tool_info = tool_use_by_name[call["call_id"]]
+                            custom_tool_calls.append(
+                                {
+                                    "id": tool_info["id"],
+                                    "type": "function",
+                                    "function": {
+                                        "name": call["name"],
+                                        "arguments": tool_info["parsed_input"],
                                     },
-                                )
-                            else:
-                                non_mcp_non_custom_tool_calls.append(
-                                    {
-                                        "id": tool_use["id"],
-                                        "type": "function",
-                                        "function": {
-                                            "name": tool_name,
-                                            "arguments": parsed_input,
-                                        },
+                                },
+                            )
+
+                        # Build non-MCP/non-custom tool calls (including workflow tools)
+                        non_mcp_non_custom_tool_calls = []
+                        for call in categorized_provider:
+                            tool_info = tool_use_by_name[call["call_id"]]
+                            non_mcp_non_custom_tool_calls.append(
+                                {
+                                    "id": tool_info["id"],
+                                    "type": "function",
+                                    "function": {
+                                        "name": call["name"],
+                                        "arguments": tool_info["parsed_input"],
                                     },
-                                )
+                                },
+                            )
+                    else:
+                        non_mcp_non_custom_tool_calls = []
+
                     # Emit non-MCP/non-custom tool calls for the caller to execute
                     if non_mcp_non_custom_tool_calls:
                         log_stream_chunk("backend.claude", "tool_calls", non_mcp_non_custom_tool_calls, agent_id)
@@ -773,156 +889,72 @@ class ClaudeBackend(CustomToolAndMCPBackend):
             # Append the assistant message with tool uses
             updated_messages.append({"role": "assistant", "content": assistant_content})
 
-            # First execute custom tool calls and append results
-            for tool_call in custom_tool_calls:
-                function_name = tool_call["function"]["name"]
+            # Configuration for custom tool execution
+            CUSTOM_TOOL_CONFIG = ToolExecutionConfig(
+                tool_type="custom",
+                chunk_type="custom_tool_status",
+                emoji_prefix="🔧 [Custom Tool]",
+                success_emoji="✅ [Custom Tool]",
+                error_emoji="❌ [Custom Tool Error]",
+                source_prefix="custom_",
+                status_called="custom_tool_called",
+                status_response="custom_tool_response",
+                status_error="custom_tool_error",
+                execution_callback=self._execute_custom_tool,
+            )
 
-                # Yield custom tool call status
-                yield StreamChunk(
-                    type="custom_tool_status",
-                    status="custom_tool_called",
-                    content=f"🔧 [Custom Tool] Calling {function_name}...",
-                    source=f"custom_{function_name}",
-                )
+            # Configuration for MCP tool execution
+            MCP_TOOL_CONFIG = ToolExecutionConfig(
+                tool_type="mcp",
+                chunk_type="mcp_status",
+                emoji_prefix="🔧 [MCP Tool]",
+                success_emoji="✅ [MCP Tool]",
+                error_emoji="❌ [MCP Tool Error]",
+                source_prefix="mcp_",
+                status_called="mcp_tool_called",
+                status_response="mcp_tool_response",
+                status_error="mcp_tool_error",
+                execution_callback=self._execute_mcp_function_with_retry,
+            )
 
-                try:
-                    # Execute custom function
-                    result_str = await self._execute_custom_tool(
-                        {
-                            "name": function_name,
-                            "arguments": json.dumps(tool_call["function"]["arguments"])
-                            if isinstance(tool_call["function"].get("arguments"), (dict, list))
-                            else tool_call["function"].get("arguments", "{}"),
-                            "call_id": tool_call["id"],
-                        },
-                    )
-                    if not result_str or result_str.startswith("Error:"):
-                        logger.warning(f"Custom function {function_name} failed: {result_str or 'unknown error'}")
-                        result_str = result_str or "Tool execution failed"
-                except Exception as e:
-                    logger.error(f"Unexpected error in custom function execution: {e}")
-                    result_str = f"Error executing custom tool: {str(e)}"
-
-                # Build tool result message
-                tool_result_msg = {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "tool_result",
-                            "tool_use_id": tool_call["id"],
-                            "content": result_str,
-                        },
-                    ],
+            def normalize_tool_call(tool_call: Dict[str, Any]) -> Dict[str, Any]:
+                """Convert Claude tool call format to unified format."""
+                return {
+                    "name": tool_call["function"]["name"],
+                    "arguments": json.dumps(tool_call["function"]["arguments"]) if isinstance(tool_call["function"].get("arguments"), (dict, list)) else tool_call["function"].get("arguments", "{}"),
+                    "call_id": tool_call["id"],  # Normalize "id" to "call_id"
                 }
 
-                # Append to updated_messages
-                updated_messages.append(tool_result_msg)
-
-                yield StreamChunk(
-                    type="custom_tool_status",
-                    status="function_call",
-                    content=f"Arguments for Calling {function_name}: {json.dumps(tool_call['function'].get('arguments', {}))}",
-                    source=f"custom_{function_name}",
-                )
-
-                yield StreamChunk(
-                    type="custom_tool_status",
-                    status="function_call_output",
-                    content=f"Results for Calling {function_name}: {result_str}",
-                    source=f"custom_{function_name}",
-                )
-
-                logger.info(f"Executed custom function {function_name}")
-                yield StreamChunk(
-                    type="custom_tool_status",
-                    status="custom_tool_response",
-                    content=f"✅ [Custom Tool] {function_name} completed",
-                    source=f"custom_{function_name}",
-                )
-
-            # Then execute MCP tool calls and append results
+            # Execute custom tools using unified method
+            for tool_call in custom_tool_calls:
+                # Normalize Claude tool call format to unified format
+                normalized_call = normalize_tool_call(tool_call)
+
+                # Use unified execution method
+                async for chunk in self._execute_tool_with_logging(
+                    normalized_call,
+                    CUSTOM_TOOL_CONFIG,
+                    updated_messages,
+                    set(),
+                ):
+                    yield chunk
+
+            # Execute MCP tools using unified method
             for tool_call in mcp_tool_calls:
-                function_name = tool_call["function"]["name"]
+                # Normalize Claude tool call format to unified format
+                normalized_call = normalize_tool_call(tool_call)
+
+                # Use unified execution method
+                async for chunk in self._execute_tool_with_logging(
+                    normalized_call,
+                    MCP_TOOL_CONFIG,
+                    updated_messages,
+                    set(),
+                ):
+                    yield chunk
 
-                # Yield MCP tool call status
-                yield StreamChunk(
-                    type="mcp_status",
-                    status="mcp_tool_called",
-                    content=f"🔧 [MCP Tool] Calling {function_name}...",
-                    source=f"mcp_{function_name}",
-                )
-
-                try:
-                    # Execute MCP function
-                    args_json = json.dumps(tool_call["function"]["arguments"]) if isinstance(tool_call["function"].get("arguments"), (dict, list)) else tool_call["function"].get("arguments", "{}")
-                    result_list = await self._execute_mcp_function_with_retry(function_name, args_json)
-                    if not result_list or (isinstance(result_list[0], str) and result_list[0].startswith("Error:")):
-                        logger.warning(f"MCP function {function_name} failed after retries: {result_list[0] if result_list else 'unknown error'}")
-                        continue
-                    result_str = result_list[0]
-                    result_obj = result_list[1] if len(result_list) > 1 else None
-                except Exception as e:
-                    logger.error(f"Unexpected error in MCP function execution: {e}")
-                    continue
-
-                # Build tool result message: { "role":"user", "content":[{ "type":"tool_result", "tool_use_id": tool_call["id"], "content": result_str }] }
-                tool_result_msg = {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "tool_result",
-                            "tool_use_id": tool_call["id"],
-                            "content": result_str,
-                        },
-                    ],
-                }
-
-                # Append to updated_messages
-                updated_messages.append(tool_result_msg)
-
-                yield StreamChunk(
-                    type="mcp_status",
-                    status="function_call",
-                    content=f"Arguments for Calling {function_name}: {json.dumps(tool_call['function'].get('arguments', {}))}",
-                    source=f"mcp_{function_name}",
-                )
-
-                # If result_obj might be structured, try to display summary
-                result_display = None
-                try:
-                    if hasattr(result_obj, "content") and result_obj.content:
-                        part = result_obj.content[0]
-                        if hasattr(part, "text"):
-                            result_display = str(part.text)
-                except Exception:
-                    result_display = None
-                if result_display:
-                    yield StreamChunk(
-                        type="mcp_status",
-                        status="function_call_output",
-                        content=f"Results for Calling {function_name}: {result_display}",
-                        source=f"mcp_{function_name}",
-                    )
-                else:
-                    yield StreamChunk(
-                        type="mcp_status",
-                        status="function_call_output",
-                        content=f"Results for Calling {function_name}: {result_str}",
-                        source=f"mcp_{function_name}",
-                    )
-
-                logger.info(f"Executed MCP function {function_name} (stdio/streamable-http)")
-                yield StreamChunk(
-                    type="mcp_status",
-                    status="mcp_tool_response",
-                    content=f"✅ [MCP Tool] {function_name} completed",
-                    source=f"mcp_{function_name}",
-                )
-
-            # Trim updated_messages using base class method
             updated_messages = self._trim_message_history(updated_messages)
 
-            # After processing all tool calls, recurse
             async for chunk in self._stream_with_custom_and_mcp_tools(updated_messages, tools, client, **kwargs):
                 yield chunk
             return
@@ -1167,10 +1199,8 @@ class ClaudeBackend(CustomToolAndMCPBackend):
             content=f"\n⚠️  {user_message} ({error}); continuing without MCP tools\n",
         )
 
-        # Build non-MCP configuration and stream fallback
         fallback_params = dict(api_params)
 
-        # Remove any MCP tools from the tools list
         if "tools" in fallback_params and self._mcp_functions:
             mcp_names = set(self._mcp_functions.keys())
             non_mcp_tools = []
@@ -1195,7 +1225,7 @@ class ClaudeBackend(CustomToolAndMCPBackend):
         function_name: str,
         arguments_json: str,
         max_retries: int = 3,
-    ) -> List[str | Any]:
+    ) -> Tuple[str, Any]:
         """Execute MCP function with Claude-specific formatting."""
         # Use parent class method which returns tuple
         result_str, result_obj = await super()._execute_mcp_function_with_retry(
@@ -1204,10 +1234,9 @@ class ClaudeBackend(CustomToolAndMCPBackend):
             max_retries,
         )
 
-        # Convert to list format expected by Claude streaming
         if result_str.startswith("Error:"):
-            return [result_str]
-        return [result_str, result_obj]
+            return (result_str, {"error": result_str})
+        return (result_str, result_obj)
 
     def create_tool_result_message(self, tool_call: Dict[str, Any], result_content: str) -> Dict[str, Any]:
         """Create tool result message in Claude's expected format."""