kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/worker/runtimes/claude_code/tool_mapper.py

@@ -0,0 +1,403 @@
+"""
+Tool mapping module for Claude Code runtime.
+
+This module handles the translation of generic skill types to specific
+Claude Code SDK tool names, with validation and error handling.
+
+BUG FIX #6: Added tool name validation before adding to allowed_tools.
+Issue #4 Fix: Externalized tool mappings to configuration file.
+"""
+
+from typing import List, Tuple, Set, Dict, Optional
+import structlog
+import yaml
+import os
+from pathlib import Path
+
+logger = structlog.get_logger(__name__)
+
+
+# Known builtin Claude Code tools from SDK documentation
+# This registry is the source of truth for validating builtin tool existence
+CLAUDE_CODE_BUILTIN_TOOLS: Set[str] = {
+    "Read",
+    "Write",
+    "Edit",
+    "Glob",
+    "Grep",
+    "Bash",
+    "BashOutput",
+    "KillShell",
+    "WebFetch",
+    "WebSearch",
+    "Task",
+    "NotebookEdit",
+    "TodoWrite",
+    "ExitPlanMode",
+    "AskUserQuestion",
+    "Skill",
+    "SlashCommand",
+    # MCP resource tools (both naming conventions)
+    "ListMcpResources",
+    "ReadMcpResource",
+    "mcp__list_resources",
+    "mcp__read_resource",
+}
+
+# Backward compatibility alias
+KNOWN_BUILTIN_TOOLS: Set[str] = CLAUDE_CODE_BUILTIN_TOOLS
+
+# DEPRECATED: Hardcoded mapping (kept for backward compatibility)
+# Use ToolMappingRegistry instead
+SKILL_TO_TOOL_MAPPING = {
+    "shell": ["Bash", "BashOutput", "KillShell"],
+    "file_system": ["Read", "Write", "Edit", "Glob", "Grep"],
+    "web": ["WebFetch", "WebSearch"],
+    "docker": ["Bash"],  # Docker commands via Bash
+    "kubernetes": ["Bash"],  # kubectl via Bash
+    "git": ["Bash"],  # git commands via Bash
+    "task": ["Task"],  # Subagent tasks
+    "notebook": ["NotebookEdit"],
+    "planning": ["TodoWrite", "ExitPlanMode"],
+}
+
+
+class ToolMappingRegistry:
+    """
+    Issue #4 Fix: Dynamic tool mapping registry with configuration file support.
+
+    This class loads tool mappings from a YAML configuration file instead of
+    hardcoding them, making the system more maintainable and extensible.
+    """
+
+    def __init__(self, config_path: Optional[Path] = None):
+        """
+        Initialize the tool mapping registry.
+
+        Args:
+            config_path: Optional path to configuration file.
+                        Defaults to tool_mapping_config.yaml in this directory.
+        """
+        self.config_path = config_path or self._get_default_config_path()
+        self.mappings: Dict[str, List[str]] = {}
+        self.default_tools: List[str] = []
+        self.overrides: Dict[str, Dict[str, List[str]]] = {}
+        self._load_config()
+
+    def _get_default_config_path(self) -> Path:
+        """Get path to default configuration file."""
+        return Path(__file__).parent / "tool_mapping_config.yaml"
+
+    def _load_config(self):
+        """Load mappings from configuration file."""
+        try:
+            # Check for environment variable override
+            env_config_path = os.environ.get("CLAUDE_CODE_TOOL_MAPPING_CONFIG")
+            if env_config_path:
+                config_path = Path(env_config_path)
+                logger.info(
+                    "using_tool_mapping_config_from_env",
+                    path=str(config_path),
+                )
+            else:
+                config_path = self.config_path
+
+            if not config_path.exists():
+                logger.warning(
+                    "tool_mapping_config_not_found",
+                    path=str(config_path),
+                    fallback="using hardcoded mappings",
+                )
+                self._load_fallback_mappings()
+                return
+
+            with open(config_path) as f:
+                config = yaml.safe_load(f)
+
+            self.mappings = config.get("skill_type_mappings", {})
+            self.default_tools = config.get("default_tools", ["Read", "Write", "Bash"])
+            self.overrides = config.get("overrides", {})
+
+            logger.info(
+                "tool_mapping_config_loaded",
+                path=str(config_path),
+                skill_types=list(self.mappings.keys()),
+                default_tools=self.default_tools,
+            )
+
+        except Exception as e:
+            logger.error(
+                "failed_to_load_tool_mapping_config",
+                path=str(self.config_path),
+                error=str(e),
+                fallback="using hardcoded mappings",
+                exc_info=True,
+            )
+            self._load_fallback_mappings()
+
+    def _load_fallback_mappings(self):
+        """Load fallback hardcoded mappings."""
+        self.mappings = SKILL_TO_TOOL_MAPPING.copy()
+        self.default_tools = ["Read", "Write", "Bash"]
+        logger.debug("using_fallback_hardcoded_mappings")
+
+    def get_tools_for_skill_type(self, skill_type: str) -> List[str]:
+        """
+        Get Claude Code tools for a skill type.
+
+        Args:
+            skill_type: Generic skill type (e.g., "file_system", "shell")
+
+        Returns:
+            List of Claude Code tool names
+        """
+        # Check for environment variable override first
+        env_key = f"CLAUDE_CODE_MAPPING_{skill_type.upper()}"
+        env_override = os.environ.get(env_key)
+        if env_override:
+            tools = [t.strip() for t in env_override.split(",")]
+            logger.debug(
+                "using_env_override_for_skill_type",
+                skill_type=skill_type,
+                tools=tools,
+                env_key=env_key,
+            )
+            return tools
+
+        # Check configuration
+        tools = self.mappings.get(skill_type, [])
+        if not tools:
+            logger.warning(
+                "no_mapping_for_skill_type",
+                skill_type=skill_type,
+                available_types=list(self.mappings.keys()),
+            )
+
+        return tools
+
+    def get_default_tools(self) -> List[str]:
+        """Get default tools when no specific skill types are provided."""
+        # Check for environment variable override
+        env_default = os.environ.get("CLAUDE_CODE_DEFAULT_TOOLS")
+        if env_default:
+            tools = [t.strip() for t in env_default.split(",")]
+            logger.debug("using_env_override_for_default_tools", tools=tools)
+            return tools
+
+        return self.default_tools
+
+    def validate_mapping(self, skill_type: str, tools: List[str]) -> bool:
+        """
+        Validate that all tools in a mapping exist in CLAUDE_CODE_BUILTIN_TOOLS.
+
+        Args:
+            skill_type: Skill type being validated
+            tools: List of tool names
+
+        Returns:
+            True if all tools are valid, False otherwise
+        """
+        invalid_tools = [t for t in tools if t not in CLAUDE_CODE_BUILTIN_TOOLS]
+        if invalid_tools:
+            logger.error(
+                "invalid_tools_in_mapping",
+                skill_type=skill_type,
+                invalid_tools=invalid_tools,
+                valid_tools=sorted(CLAUDE_CODE_BUILTIN_TOOLS),
+            )
+            return False
+        return True
+
+    def reload_config(self):
+        """Reload configuration from file."""
+        logger.info("reloading_tool_mapping_config")
+        self._load_config()
+
+
+# Global registry instance
+_tool_mapping_registry: Optional[ToolMappingRegistry] = None
+
+
+def get_tool_mapping_registry() -> ToolMappingRegistry:
+    """
+    Get the global tool mapping registry instance.
+
+    Returns:
+        ToolMappingRegistry instance
+    """
+    global _tool_mapping_registry
+    if _tool_mapping_registry is None:
+        _tool_mapping_registry = ToolMappingRegistry()
+    return _tool_mapping_registry
+
+
+def map_skills_to_tools(skills: list) -> List[str]:
+    """
+    Map skills to Claude Code tool names.
+
+    This function translates generic skill types to the specific
+    tool names that Claude Code SDK understands.
+
+    Issue #4 Fix: Now uses ToolMappingRegistry instead of hardcoded mappings.
+
+    Args:
+        skills: List of skill objects
+
+    Returns:
+        List of Claude Code tool names (deduplicated)
+    """
+    registry = get_tool_mapping_registry()
+    tools = []
+
+    for skill in skills:
+        # Get skill type
+        skill_type = None
+        if hasattr(skill, "type"):
+            skill_type = skill.type
+        elif isinstance(skill, dict):
+            skill_type = skill.get("type")
+
+        # Map to Claude Code tools using registry
+        if skill_type:
+            mapped_tools = registry.get_tools_for_skill_type(skill_type)
+            if mapped_tools:
+                tools.extend(mapped_tools)
+
+    # Deduplicate and add default tools if none specified
+    if tools:
+        unique_tools = list(set(tools))
+    else:
+        unique_tools = registry.get_default_tools()
+
+    logger.info(
+        "mapped_skills_to_claude_code_tools",
+        skill_count=len(skills),
+        tool_count=len(unique_tools),
+        tools=unique_tools,
+    )
+
+    return unique_tools
+
+
+def validate_builtin_tools(tool_names: List[str]) -> Tuple[List[str], List[str]]:
+    """
+    Validate that builtin tool names exist in Claude Code SDK.
+
+    Issue #2 Fix: This function ensures Claude Code builtin tools are validated
+    before being used, preventing runtime errors.
+
+    Args:
+        tool_names: List of builtin tool names to validate
+
+    Returns:
+        Tuple of (valid_tools, invalid_tools)
+
+    Raises:
+        ValueError: If any builtin tool doesn't exist in CLAUDE_CODE_BUILTIN_TOOLS
+    """
+    valid = []
+    invalid = []
+
+    for tool_name in tool_names:
+        if tool_name in CLAUDE_CODE_BUILTIN_TOOLS:
+            valid.append(tool_name)
+        else:
+            invalid.append(tool_name)
+
+    if invalid:
+        error_msg = (
+            f"Invalid builtin tool names detected: {invalid}. "
+            f"These tools don't exist in Claude Code SDK. "
+            f"Valid builtin tools are: {sorted(CLAUDE_CODE_BUILTIN_TOOLS)}"
+        )
+        logger.error(
+            "invalid_builtin_tools_detected",
+            invalid_tools=invalid,
+            valid_tools=valid,
+            available_tools=sorted(CLAUDE_CODE_BUILTIN_TOOLS),
+        )
+        raise ValueError(error_msg)
+
+    logger.debug(
+        "builtin_tools_validated",
+        tool_count=len(valid),
+        tools=valid,
+    )
+
+    return valid, invalid
+
+
+def validate_tool_names(tool_names: List[str]) -> Tuple[List[str], List[str]]:
+    """
+    Validate tool names before adding to allowed_tools.
+
+    BUG FIX #6: This function ensures only valid tool names are used.
+
+    Args:
+        tool_names: List of tool names to validate
+
+    Returns:
+        Tuple of (valid_tools, invalid_tools)
+    """
+    valid = []
+    invalid = []
+
+    for tool_name in tool_names:
+        # Builtin tools or MCP tools (mcp__ prefix) are valid
+        if tool_name in KNOWN_BUILTIN_TOOLS or tool_name.startswith("mcp__"):
+            valid.append(tool_name)
+        else:
+            invalid.append(tool_name)
+
+    if invalid:
+        logger.warning(
+            "invalid_tool_names_filtered",
+            invalid_tools=invalid,
+            valid_count=len(valid),
+            invalid_count=len(invalid),
+            message="These tools will be filtered out as they're not recognized",
+        )
+
+    return valid, invalid
+
+
+def sanitize_tool_name(name: str) -> str:
+    """
+    Sanitize a tool name to meet universal LLM provider requirements.
+
+    Uses the universal validator to ensure compatibility across all providers.
+
+    Args:
+        name: Tool name to sanitize
+
+    Returns:
+        Sanitized tool name that works for all LLM providers
+    """
+    from control_plane_api.worker.utils.tool_validation import sanitize_tool_name as universal_sanitize
+    return universal_sanitize(name)
+
+
+def construct_mcp_tool_name(server_name: str, tool_name: str = None) -> str:
+    """
+    Construct a full MCP tool name following the naming convention.
+
+    Convention: mcp__<server_name>__<tool_name>
+
+    Args:
+        server_name: Name of the MCP server
+        tool_name: Optional tool name. If omitted, only server name is used.
+
+    Returns:
+        Full MCP tool name (sanitized)
+    """
+    sanitized_server = sanitize_tool_name(server_name)
+
+    if tool_name:
+        sanitized_tool = sanitize_tool_name(tool_name)
+        # Avoid duplication: mcp__run_ado_test not mcp__run_ado_test__run_ado_test
+        if sanitized_tool == sanitized_server:
+            return f"mcp__{sanitized_server}"
+        else:
+            return f"mcp__{sanitized_server}__{sanitized_tool}"
+    else:
+        return f"mcp__{sanitized_server}"

control_plane_api/worker/runtimes/claude_code/utils.py

@@ -0,0 +1,149 @@
+"""
+Utility functions for Claude Code runtime.
+
+This module provides helper functions for prompt building, session management,
+and SDK version checking.
+"""
+
+from typing import Optional
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+def build_prompt_with_history(context: any) -> str:
+    """
+    Build prompt with conversation history.
+
+    Since ClaudeSDKClient maintains session continuity, we include
+    the conversation history as context in the prompt.
+
+    Args:
+        context: RuntimeExecutionContext
+
+    Returns:
+        Prompt string with history context
+    """
+    if not context.conversation_history:
+        return context.prompt
+
+    # Build context from history
+    history_context = "Previous conversation:\n"
+    for msg in context.conversation_history[-10:]:  # Last 10 messages
+        role = msg.get("role", "unknown")
+        content = msg.get("content", "")
+        if len(content) > 200:
+            history_context += f"{role.capitalize()}: {content[:200]}...\n"
+        else:
+            history_context += f"{role.capitalize()}: {content}\n"
+
+    return f"{history_context}\n\nCurrent request:\n{context.prompt}"
+
+
+def get_sdk_version() -> str:
+    """
+    Get Claude Code SDK version.
+
+    Returns:
+        SDK version string or "unknown"
+    """
+    try:
+        import claude_agent_sdk
+
+        return getattr(claude_agent_sdk, "__version__", "unknown")
+    except Exception:
+        return "unknown"
+
+
+def extract_usage_from_result_message(message: any) -> dict:
+    """
+    Extract usage metrics from ResultMessage.
+
+    Args:
+        message: ResultMessage object
+
+    Returns:
+        Usage dict with token counts
+    """
+    if not hasattr(message, "usage") or not message.usage:
+        logger.warning("result_message_has_no_usage")
+        return {}
+
+    # ResultMessage.usage is a dict, not an object - use .get() instead of getattr()
+    usage_dict = message.usage
+
+    # Handle both dict and object usage formats
+    if isinstance(usage_dict, dict):
+        input_tokens = usage_dict.get("input_tokens", 0)
+        output_tokens = usage_dict.get("output_tokens", 0)
+        cache_read_tokens = usage_dict.get("cache_read_input_tokens", 0)
+        cache_creation_tokens = usage_dict.get("cache_creation_input_tokens", 0)
+    else:
+        # Fallback for object-style usage
+        input_tokens = getattr(usage_dict, "input_tokens", 0)
+        output_tokens = getattr(usage_dict, "output_tokens", 0)
+        cache_read_tokens = getattr(usage_dict, "cache_read_input_tokens", 0)
+        cache_creation_tokens = getattr(usage_dict, "cache_creation_input_tokens", 0)
+
+    # Use Anthropic field names for consistency with analytics
+    usage = {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": input_tokens + output_tokens,
+        "cache_read_tokens": cache_read_tokens,
+        "cache_creation_tokens": cache_creation_tokens,
+    }
+
+    logger.info(
+        "usage_extracted_from_claude_code",
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=usage["total_tokens"],
+    )
+
+    return usage
+
+
+def extract_session_id_from_result_message(
+    message: any, execution_id: str
+) -> Optional[str]:
+    """
+    Extract and validate session_id from ResultMessage.
+
+    BUG FIX #4: Validates session_id before returning.
+
+    Args:
+        message: ResultMessage object
+        execution_id: Execution ID for logging
+
+    Returns:
+        Valid session_id or None
+    """
+    session_id = getattr(message, "session_id", None)
+
+    if not session_id:
+        logger.warning(
+            "no_session_id_in_result_message",
+            execution_id=execution_id[:8],
+            message="Multi-turn conversations may not work without session_id",
+        )
+        return None
+
+    # Validate format
+    if not isinstance(session_id, str) or len(session_id) < 10:
+        logger.warning(
+            "invalid_session_id_format_in_result_message",
+            execution_id=execution_id[:8],
+            session_id_type=type(session_id).__name__,
+            session_id_length=len(session_id) if isinstance(session_id, str) else 0,
+        )
+        return None
+
+    logger.info(
+        "session_id_captured_for_conversation_continuity",
+        execution_id=execution_id[:8],
+        session_id_prefix=session_id[:16],
+        message="This session_id will enable multi-turn conversations",
+    )
+
+    return session_id