massgen 0.1.0a3__py3-none-any.whl → 0.1.2__py3-none-any.whl

massgen/backend/chat_completions.py

@@ -31,10 +31,10 @@ from ..logger_config import log_backend_agent_message, log_stream_chunk, logger
 
 # Local imports
 from .base import FilesystemSupport, StreamChunk
-from .base_with_mcp import MCPBackend
+from .base_with_custom_tool_and_mcp import CustomToolAndMCPBackend
 
 
-class ChatCompletionsBackend(MCPBackend):
+class ChatCompletionsBackend(CustomToolAndMCPBackend):
     """Complete OpenAI-compatible Chat Completions API backend.
 
     Can be used directly with any OpenAI-compatible provider by setting provider name.
@@ -67,14 +67,14 @@ class ChatCompletionsBackend(MCPBackend):
         async for chunk in super().stream_with_tools(messages, tools, **kwargs):
             yield chunk
 
-    async def _stream_with_mcp_tools(
+    async def _stream_with_custom_and_mcp_tools(
         self,
         current_messages: List[Dict[str, Any]],
         tools: List[Dict[str, Any]],
         client,
         **kwargs,
     ) -> AsyncGenerator[StreamChunk, None]:
-        """Recursively stream MCP responses, executing function calls as needed."""
+        """Recursively stream responses, executing custom and MCP tool calls as needed."""
 
         # Build API params for this iteration
         all_params = {**self.config, **kwargs}
@@ -193,14 +193,28 @@ class ChatCompletionsBackend(MCPBackend):
 
         # Execute any captured function calls
         if captured_function_calls and response_completed:
-            # Check if any of the function calls are NOT MCP functions
-            non_mcp_functions = [call for call in captured_function_calls if call["name"] not in self._mcp_functions]
+            # Categorize function calls
+            mcp_calls = []
+            custom_calls = []
+            provider_calls = []
 
-            if non_mcp_functions:
-                logger.info(f"Non-MCP function calls detected (will be ignored in MCP execution): {[call['name'] for call in non_mcp_functions]}")
+            for call in captured_function_calls:
+                if call["name"] in self._mcp_functions:
+                    mcp_calls.append(call)
+                elif call["name"] in self._custom_tool_names:
+                    custom_calls.append(call)
+                else:
+                    provider_calls.append(call)
+
+            # If there are provider calls (non-MCP, non-custom), let API handle them
+            if provider_calls:
+                logger.info(f"Provider function calls detected: {[call['name'] for call in provider_calls]}. Ending local processing.")
+                yield StreamChunk(type="done")
+                return
 
             # Check circuit breaker status before executing MCP functions
-            if not await self._check_circuit_breaker_before_execution():
+            if mcp_calls and not await self._check_circuit_breaker_before_execution():
+                logger.warning("All MCP servers blocked by circuit breaker")
                 yield StreamChunk(
                     type="mcp_status",
                     status="mcp_blocked",
@@ -210,22 +224,30 @@ class ChatCompletionsBackend(MCPBackend):
                 yield StreamChunk(type="done")
                 return
 
-            # Execute only MCP function calls
-            mcp_functions_executed = False
+            # Initialize for execution
+            functions_executed = False
             updated_messages = current_messages.copy()
+            processed_call_ids = set()  # Track processed calls
 
-            # Check if planning mode is enabled - block MCP tool execution during planning
+            # Check if planning mode is enabled - selectively block MCP tool execution during planning
             if self.is_planning_mode_enabled():
-                logger.info("[MCP] Planning mode enabled - blocking all MCP tool execution")
-                yield StreamChunk(
-                    type="mcp_status",
-                    status="planning_mode_blocked",
-                    content="🚫 [MCP] Planning mode active - MCP tools blocked during coordination",
-                    source="planning_mode",
-                )
-                # Skip all MCP tool execution but still continue with workflow
-                yield StreamChunk(type="done")
-                return
+                blocked_tools = self.get_planning_mode_blocked_tools()
+
+                if not blocked_tools:
+                    # Empty set means block ALL MCP tools (backward compatible)
+                    logger.info("[ChatCompletions] Planning mode enabled - blocking ALL MCP tool execution")
+                    yield StreamChunk(
+                        type="mcp_status",
+                        status="planning_mode_blocked",
+                        content="🚫 [MCP] Planning mode active - all MCP tools blocked during coordination",
+                        source="planning_mode",
+                    )
+                    # Skip all MCP tool execution but still continue with workflow
+                    yield StreamChunk(type="done")
+                    return
+                else:
+                    # Selective blocking - log but continue to check each tool individually
+                    logger.info(f"[ChatCompletions] Planning mode enabled - selective blocking of {len(blocked_tools)} tools")
 
             # Create single assistant message with all tool calls
             if captured_function_calls:
@@ -252,11 +274,90 @@ class ChatCompletionsBackend(MCPBackend):
                     }
                     updated_messages.append(assistant_message)
 
-            # Execute functions and collect results
-            tool_results = []
-            for call in captured_function_calls:
+            # Execute custom tools first
+            for call in custom_calls:
+                try:
+                    # Yield custom tool call status
+                    yield StreamChunk(
+                        type="custom_tool_status",
+                        status="custom_tool_called",
+                        content=f"🔧 [Custom Tool] Calling {call['name']}...",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    # Yield custom tool arguments
+                    yield StreamChunk(
+                        type="custom_tool_status",
+                        status="function_call",
+                        content=f"Arguments for Calling {call['name']}: {call['arguments']}",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    # Execute custom tool
+                    result = await self._execute_custom_tool(call)
+
+                    # Add function result to messages
+                    function_output_msg = {
+                        "role": "tool",
+                        "tool_call_id": call["call_id"],
+                        "content": str(result),
+                    }
+                    updated_messages.append(function_output_msg)
+
+                    # Yield custom tool results
+                    yield StreamChunk(
+                        type="custom_tool_status",
+                        status="function_call_output",
+                        content=f"Results for Calling {call['name']}: {str(result)}",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    # Yield custom tool response status
+                    yield StreamChunk(
+                        type="custom_tool_status",
+                        status="custom_tool_response",
+                        content=f"✅ [Custom Tool] {call['name']} completed",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    processed_call_ids.add(call["call_id"])
+                    functions_executed = True
+                    logger.info(f"Executed custom tool: {call['name']}")
+
+                except Exception as e:
+                    logger.error(f"Error executing custom tool {call['name']}: {e}")
+                    error_msg = f"Error executing {call['name']}: {str(e)}"
+
+                    # Yield error with arguments shown
+                    yield StreamChunk(
+                        type="custom_tool_status",
+                        status="function_call",
+                        content=f"Arguments for Calling {call['name']}: {call['arguments']}",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    yield StreamChunk(
+                        type="custom_tool_status",
+                        status="custom_tool_error",
+                        content=f"❌ [Custom Tool Error] {error_msg}",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    # Add error result to messages
+                    error_output_msg = {
+                        "role": "tool",
+                        "tool_call_id": call["call_id"],
+                        "content": error_msg,
+                    }
+                    updated_messages.append(error_output_msg)
+                    processed_call_ids.add(call["call_id"])
+                    functions_executed = True
+
+            # Execute MCP function calls
+            mcp_functions_executed = False
+            for call in mcp_calls:
                 function_name = call["name"]
-                if self.is_mcp_tool_call(function_name):
+                if function_name in self._mcp_functions:
                     yield StreamChunk(
                         type="mcp_status",
                         status="mcp_tool_called",
@@ -264,7 +365,7 @@ class ChatCompletionsBackend(MCPBackend):
                         source=f"mcp_{function_name}",
                     )
 
-                    # Yield detailed MCP status as StreamChunk (similar to gemini.py)
+                    # Yield detailed MCP status as StreamChunk
                     tools_info = f" ({len(self._mcp_functions)} tools available)" if self._mcp_functions else ""
                     yield StreamChunk(
                         type="mcp_status",
@@ -275,51 +376,40 @@ class ChatCompletionsBackend(MCPBackend):
 
                     try:
                         # Execute MCP function with retry and exponential backoff
-                        (
-                            result_str,
-                            result_obj,
-                        ) = await self._execute_mcp_function_with_retry(function_name, call["arguments"])
+                        result_str, result_obj = await self._execute_mcp_function_with_retry(function_name, call["arguments"])
 
                         # Check if function failed after all retries
                         if isinstance(result_str, str) and result_str.startswith("Error:"):
                             # Log failure but still create tool response
                             logger.warning(f"MCP function {function_name} failed after retries: {result_str}")
-                            tool_results.append(
-                                {
-                                    "tool_call_id": call["call_id"],
-                                    "content": result_str,
-                                    "success": False,
-                                },
-                            )
-                        else:
-                            # Yield MCP success status as StreamChunk (similar to gemini.py)
-                            yield StreamChunk(
-                                type="mcp_status",
-                                status="mcp_tools_success",
-                                content=f"MCP tool call succeeded (call #{self._mcp_tool_calls_count})",
-                                source=f"mcp_{function_name}",
-                            )
 
-                            tool_results.append(
-                                {
-                                    "tool_call_id": call["call_id"],
-                                    "content": result_str,
-                                    "success": True,
-                                    "result_obj": result_obj,
-                                },
-                            )
+                            # Add error result to messages
+                            function_output_msg = {
+                                "role": "tool",
+                                "tool_call_id": call["call_id"],
+                                "content": result_str,
+                            }
+                            updated_messages.append(function_output_msg)
+
+                            processed_call_ids.add(call["call_id"])
+                            mcp_functions_executed = True
+                            continue
 
                     except Exception as e:
                         # Only catch unexpected non-MCP system errors
                         logger.error(f"Unexpected error in MCP function execution: {e}")
                         error_msg = f"Error executing {function_name}: {str(e)}"
-                        tool_results.append(
-                            {
-                                "tool_call_id": call["call_id"],
-                                "content": error_msg,
-                                "success": False,
-                            },
-                        )
+
+                        # Add error result to messages
+                        function_output_msg = {
+                            "role": "tool",
+                            "tool_call_id": call["call_id"],
+                            "content": error_msg,
+                        }
+                        updated_messages.append(function_output_msg)
+
+                        processed_call_ids.add(call["call_id"])
+                        mcp_functions_executed = True
                         continue
 
                     # Yield function_call status
@@ -330,60 +420,66 @@ class ChatCompletionsBackend(MCPBackend):
                         source=f"mcp_{function_name}",
                     )
 
-                    logger.info(f"Executed MCP function {function_name} (stdio/streamable-http)")
-                    mcp_functions_executed = True
-                else:
-                    # For non-MCP functions, add a dummy tool result to maintain message consistency
-                    logger.info(f"Non-MCP function {function_name} detected, creating placeholder response")
-                    tool_results.append(
-                        {
-                            "tool_call_id": call["call_id"],
-                            "content": f"Function {function_name} is not available in this MCP session.",
-                            "success": False,
-                        },
+                    # Add function output to messages and yield status chunk
+                    function_output_msg = {
+                        "role": "tool",
+                        "tool_call_id": call["call_id"],
+                        "content": str(result_str),
+                    }
+                    updated_messages.append(function_output_msg)
+
+                    # Yield function_call_output status with preview
+                    result_text = str(result_str)
+                    if hasattr(result_obj, "content") and result_obj.content:
+                        if isinstance(result_obj.content, list) and len(result_obj.content) > 0:
+                            first_item = result_obj.content[0]
+                            if hasattr(first_item, "text"):
+                                result_text = first_item.text
+
+                    yield StreamChunk(
+                        type="mcp_status",
+                        status="function_call_output",
+                        content=f"Results for Calling {function_name}: {result_text}",
+                        source=f"mcp_{function_name}",
                     )
 
-            # Add all tool response messages after the assistant message
-            for result in tool_results:
-                # Yield function_call_output status with preview
-                result_text = str(result["content"])
-                if result.get("success") and hasattr(result.get("result_obj"), "content") and result["result_obj"].content:
-                    obj = result["result_obj"]
-                    if isinstance(obj.content, list) and len(obj.content) > 0:
-                        first_item = obj.content[0]
-                        if hasattr(first_item, "text"):
-                            result_text = first_item.text
+                    logger.info(f"Executed MCP function {function_name} (stdio/streamable-http)")
+                    processed_call_ids.add(call["call_id"])
 
-                yield StreamChunk(
-                    type="mcp_status",
-                    status="function_call_output",
-                    content=f"Results for Calling {function_name}: {result_text}",
-                    source=f"mcp_{function_name}",
-                )
+                    # Yield MCP tool response status
+                    yield StreamChunk(
+                        type="mcp_status",
+                        status="mcp_tool_response",
+                        content=f"✅ [MCP Tool] {function_name} completed",
+                        source=f"mcp_{function_name}",
+                    )
 
-                function_output_msg = {
-                    "role": "tool",
-                    "tool_call_id": result["tool_call_id"],
-                    "content": result["content"],
-                }
-                updated_messages.append(function_output_msg)
+                    mcp_functions_executed = True
+                    functions_executed = True
 
-                yield StreamChunk(
-                    type="mcp_status",
-                    status="mcp_tool_response",
-                    content=f"✅ [MCP Tool] {function_name} completed",
-                    source=f"mcp_{function_name}",
-                )
+            # Ensure all captured function calls have results to prevent hanging
+            for call in captured_function_calls:
+                if call["call_id"] not in processed_call_ids:
+                    logger.warning(f"Tool call {call['call_id']} for function {call['name']} was not processed - adding error result")
+
+                    # Add missing function call and error result to messages
+                    error_output_msg = {
+                        "role": "tool",
+                        "tool_call_id": call["call_id"],
+                        "content": f"Error: Tool call {call['call_id']} for function {call['name']} was not processed. This may indicate a validation or execution error.",
+                    }
+                    updated_messages.append(error_output_msg)
+                    mcp_functions_executed = True
 
             # Trim history after function executions to bound memory usage
-            if mcp_functions_executed:
+            if functions_executed or mcp_functions_executed:
                 updated_messages = self._trim_message_history(updated_messages)
 
                 # Recursive call with updated messages
-                async for chunk in self._stream_with_mcp_tools(updated_messages, tools, client, **kwargs):
+                async for chunk in self._stream_with_custom_and_mcp_tools(updated_messages, tools, client, **kwargs):
                     yield chunk
             else:
-                # No MCP functions were executed, we're done
+                # No functions were executed, we're done
                 yield StreamChunk(type="done")
                 return
 
@@ -395,6 +491,7 @@ class ChatCompletionsBackend(MCPBackend):
                 content="✅ [MCP] Session completed",
                 source="mcp_session",
             )
+            yield StreamChunk(type="done")
             return
 
     async def _process_stream(self, stream, all_params, agent_id) -> AsyncGenerator[StreamChunk, None]:

massgen/backend/claude.py

@@ -37,10 +37,10 @@ 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_mcp import MCPBackend, UploadFileError
+from .base_with_custom_tool_and_mcp import CustomToolAndMCPBackend, UploadFileError
 
 
-class ClaudeBackend(MCPBackend):
+class ClaudeBackend(CustomToolAndMCPBackend):
     """Claude backend using Anthropic's Messages API with full multi-tool support."""
 
     def __init__(self, api_key: Optional[str] = None, **kwargs):
@@ -469,7 +469,7 @@ class ClaudeBackend(MCPBackend):
                         "Claude Files API upload left unresolved file_pending_upload marker " f"(message {msg_idx}, item {item_idx}, source {identifier}).",
                     )
 
-    async def _stream_without_mcp_tools(
+    async def _stream_without_custom_and_mcp_tools(
         self,
         messages: List[Dict[str, Any]],
         tools: List[Dict[str, Any]],
@@ -501,6 +501,8 @@ class ClaudeBackend(MCPBackend):
                     name = tool.get("function", {}).get("name") if "function" in tool else tool.get("name")
                     if name and name in self._mcp_function_names:
                         continue
+                    if name and name in self._custom_tool_names:
+                        continue
                 elif tool.get("type") == "mcp":
                     continue
                 non_mcp_tools.append(tool)
@@ -519,14 +521,14 @@ class ClaudeBackend(MCPBackend):
         async for chunk in self._process_stream(stream, all_params, agent_id):
             yield chunk
 
-    async def _stream_with_mcp_tools(
+    async def _stream_with_custom_and_mcp_tools(
         self,
         current_messages: List[Dict[str, Any]],
         tools: List[Dict[str, Any]],
         client,
         **kwargs,
     ) -> AsyncGenerator[StreamChunk, None]:
-        """Recursively stream responses, executing MCP function calls when detected."""
+        """Recursively stream responses, executing MCP and custom tool function calls when detected."""
 
         # Build API params for this iteration
         all_params = {**self.config, **kwargs}
@@ -554,6 +556,7 @@ class ClaudeBackend(MCPBackend):
         content = ""
         current_tool_uses: Dict[str, Dict[str, Any]] = {}
         mcp_tool_calls: List[Dict[str, Any]] = []
+        custom_tool_calls: List[Dict[str, Any]] = []
         response_completed = False
 
         async for event in stream:
@@ -657,8 +660,8 @@ class ClaudeBackend(MCPBackend):
                 elif event.type == "message_delta":
                     pass
                 elif event.type == "message_stop":
-                    # Identify MCP and non-MCP tool calls among current_tool_uses
-                    non_mcp_tool_calls = []
+                    # Identify MCP, custom, and non-MCP/non-custom tool calls among current_tool_uses
+                    non_mcp_non_custom_tool_calls = []
                     if current_tool_uses:
                         for tool_use in current_tool_uses.values():
                             tool_name = tool_use.get("name", "")
@@ -683,8 +686,19 @@ class ClaudeBackend(MCPBackend):
                                         },
                                     },
                                 )
+                            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,
+                                        },
+                                    },
+                                )
                             else:
-                                non_mcp_tool_calls.append(
+                                non_mcp_non_custom_tool_calls.append(
                                     {
                                         "id": tool_use["id"],
                                         "type": "function",
@@ -694,10 +708,10 @@ class ClaudeBackend(MCPBackend):
                                         },
                                     },
                                 )
-                    # Emit non-MCP tool calls for the caller to execute
-                    if non_mcp_tool_calls:
-                        log_stream_chunk("backend.claude", "tool_calls", non_mcp_tool_calls, agent_id)
-                        yield StreamChunk(type="tool_calls", tool_calls=non_mcp_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)
+                        yield StreamChunk(type="tool_calls", tool_calls=non_mcp_non_custom_tool_calls)
                     response_completed = True
                     break
             except Exception as event_error:
@@ -706,8 +720,8 @@ class ClaudeBackend(MCPBackend):
                 yield StreamChunk(type="error", error=error_msg)
                 continue
 
-        # If we captured MCP tool calls, execute them and recurse
-        if response_completed and mcp_tool_calls:
+        # If we captured MCP or custom tool calls, execute them and recurse
+        if response_completed and (mcp_tool_calls or custom_tool_calls):
             # Circuit breaker pre-execution check using base class method
             if not await self._check_circuit_breaker_before_execution():
                 yield StreamChunk(
@@ -721,11 +735,12 @@ class ClaudeBackend(MCPBackend):
 
             updated_messages = current_messages.copy()
 
-            # Build assistant message with tool_use blocks for all MCP tool calls
+            # Build assistant message with tool_use blocks for all MCP and custom tool calls
             assistant_content = []
             if content:  # Add text content if any
                 assistant_content.append({"type": "text", "text": content})
 
+            # Add tool_use blocks for MCP tools
             for tool_call in mcp_tool_calls:
                 tool_name = tool_call["function"]["name"]
                 tool_args = tool_call["function"]["arguments"]
@@ -740,10 +755,92 @@ class ClaudeBackend(MCPBackend):
                     },
                 )
 
+            # Add tool_use blocks for custom tools
+            for tool_call in custom_tool_calls:
+                tool_name = tool_call["function"]["name"]
+                tool_args = tool_call["function"]["arguments"]
+                tool_id = tool_call["id"]
+
+                assistant_content.append(
+                    {
+                        "type": "tool_use",
+                        "id": tool_id,
+                        "name": tool_name,
+                        "input": tool_args,
+                    },
+                )
+
             # Append the assistant message with tool uses
             updated_messages.append({"role": "assistant", "content": assistant_content})
 
-            # Now execute the MCP tool calls and append results
+            # First execute custom tool calls and append results
+            for tool_call in custom_tool_calls:
+                function_name = tool_call["function"]["name"]
+
+                # 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}",
+                )
+
+                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,
+                        },
+                    ],
+                }
+
+                # 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
             for tool_call in mcp_tool_calls:
                 function_name = tool_call["function"]["name"]
 
@@ -825,8 +922,8 @@ class ClaudeBackend(MCPBackend):
             # Trim updated_messages using base class method
             updated_messages = self._trim_message_history(updated_messages)
 
-            # After processing all MCP calls, recurse: async for chunk in self._stream_mcp_recursive(updated_messages, tools, client, **kwargs): yield chunk
-            async for chunk in self._stream_with_mcp_tools(updated_messages, tools, client, **kwargs):
+            # 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
         else: