377 lines
15 KiB
Python
377 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Claude Code Subagent Management System
|
|
|
|
This script provides comprehensive management for Claude Code subagents,
|
|
including initialization, coordination, and workflow execution.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SubagentStatus(Enum):
|
|
AVAILABLE = "available"
|
|
BUSY = "busy"
|
|
ERROR = "error"
|
|
OFFLINE = "offline"
|
|
|
|
|
|
@dataclass
|
|
class SubagentInstance:
|
|
name: str
|
|
config: dict[str, Any]
|
|
status: SubagentStatus = SubagentStatus.AVAILABLE
|
|
current_task: str | None = None
|
|
load_factor: float = 0.0
|
|
|
|
|
|
class SubagentManager:
|
|
"""Advanced subagent management and coordination system"""
|
|
|
|
def __init__(self, registry_path: str = ".claude-code/subagents/registry.json"):
|
|
self.registry_path = Path(registry_path)
|
|
self.subagents: dict[str, SubagentInstance] = {}
|
|
self.collaboration_patterns: dict[str, list[str]] = {}
|
|
self.active_workflows: dict[str, dict] = {}
|
|
|
|
self._load_registry()
|
|
self._initialize_subagents()
|
|
|
|
def _load_registry(self):
|
|
"""Load subagent registry configuration"""
|
|
try:
|
|
with self.registry_path.open() as f:
|
|
self.registry = json.load(f)
|
|
|
|
self.collaboration_patterns = self.registry.get("interaction_patterns", {})
|
|
logger.info(f"Loaded registry with {len(self.registry.get('subagents', {}))} subagents")
|
|
|
|
except FileNotFoundError:
|
|
logger.error(f"Registry file not found: {self.registry_path}")
|
|
self.registry = {"subagents": {}}
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Invalid JSON in registry: {e}")
|
|
self.registry = {"subagents": {}}
|
|
|
|
def _initialize_subagents(self):
|
|
"""Initialize all subagents from registry"""
|
|
for _category, subagents in self.registry.get("subagents", {}).items():
|
|
for name, config in subagents.items():
|
|
try:
|
|
subagent_config = self._load_subagent_config(config["file"])
|
|
self.subagents[name] = SubagentInstance(name=name, config=subagent_config)
|
|
logger.info(f"Initialized subagent: {name}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize subagent {name}: {e}")
|
|
|
|
def _load_subagent_config(self, config_file: str) -> dict[str, Any]:
|
|
"""Load individual subagent configuration"""
|
|
config_path = self.registry_path.parent / config_file
|
|
|
|
try:
|
|
with config_path.open() as f:
|
|
return json.load(f)
|
|
except FileNotFoundError:
|
|
logger.warning(f"Config file not found: {config_path}")
|
|
return {}
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Invalid JSON in config {config_path}: {e}")
|
|
return {}
|
|
|
|
def get_subagent_capabilities(self, subagent_name: str) -> list[str]:
|
|
"""Get capabilities of a specific subagent"""
|
|
if subagent_name in self.subagents:
|
|
return self.subagents[subagent_name].config.get("capabilities", [])
|
|
return []
|
|
|
|
def find_subagents_by_capability(self, capability: str) -> list[str]:
|
|
"""Find all subagents with a specific capability"""
|
|
matching_subagents = []
|
|
for name, instance in self.subagents.items():
|
|
if capability in instance.config.get("capabilities", []):
|
|
matching_subagents.append(name)
|
|
return matching_subagents
|
|
|
|
def get_collaboration_pattern(self, pattern_name: str) -> list[str]:
|
|
"""Get predefined collaboration pattern"""
|
|
return self.collaboration_patterns.get(pattern_name, [])
|
|
|
|
def analyze_task_requirements(self, task_description: str) -> dict[str, Any]:
|
|
"""Analyze task and determine required capabilities and subagents"""
|
|
# This would use NLP/ML in a real implementation
|
|
# For now, provide a rule-based analysis
|
|
|
|
capabilities_needed = []
|
|
priority_subagents = []
|
|
|
|
# Code quality tasks
|
|
if any(keyword in task_description.lower() for keyword in ["lint", "format", "quality", "ruff", "type"]):
|
|
capabilities_needed.extend(["static-code-analysis", "type-checking"])
|
|
priority_subagents.append("python-quality-analyst")
|
|
|
|
# Testing tasks
|
|
if any(keyword in task_description.lower() for keyword in ["test", "bdd", "behave", "coverage"]):
|
|
capabilities_needed.extend(["bdd-scenario-design", "test-execution"])
|
|
priority_subagents.extend(["test-architect", "test-executor"])
|
|
|
|
# Security tasks
|
|
if any(keyword in task_description.lower() for keyword in ["security", "vulnerability", "scan", "audit"]):
|
|
capabilities_needed.extend(["security-scanning", "vulnerability-assessment"])
|
|
priority_subagents.append("security-auditor")
|
|
|
|
# Performance tasks
|
|
if any(keyword in task_description.lower() for keyword in ["performance", "optimize", "profile", "benchmark"]):
|
|
capabilities_needed.extend(["performance-profiling", "optimization"])
|
|
priority_subagents.append("performance-optimizer")
|
|
|
|
# Container/deployment tasks
|
|
if any(keyword in task_description.lower() for keyword in ["docker", "container", "deploy", "kubernetes"]):
|
|
capabilities_needed.extend(["container-design", "deployment"])
|
|
priority_subagents.extend(["container-architect", "kubernetes-specialist"])
|
|
|
|
return {
|
|
"capabilities_needed": capabilities_needed,
|
|
"priority_subagents": priority_subagents,
|
|
"complexity": self._assess_task_complexity(task_description),
|
|
"estimated_duration": self._estimate_task_duration(task_description),
|
|
}
|
|
|
|
def _assess_task_complexity(self, task_description: str) -> str:
|
|
"""Assess task complexity based on description"""
|
|
complexity_indicators = {
|
|
"simple": ["fix", "update", "change"],
|
|
"medium": ["implement", "create", "design", "optimize"],
|
|
"complex": ["architect", "integrate", "migrate", "refactor"],
|
|
"very_complex": ["orchestrate", "coordinate", "comprehensive"],
|
|
}
|
|
|
|
task_lower = task_description.lower()
|
|
for complexity, indicators in complexity_indicators.items():
|
|
if any(indicator in task_lower for indicator in indicators):
|
|
return complexity
|
|
|
|
return "medium" # default
|
|
|
|
def _estimate_task_duration(self, task_description: str) -> str:
|
|
"""Estimate task duration based on complexity and scope"""
|
|
task_lower = task_description.lower()
|
|
|
|
if any(keyword in task_lower for keyword in ["quick", "simple", "minor"]):
|
|
return "short" # < 30 minutes
|
|
elif any(keyword in task_lower for keyword in ["comprehensive", "complete", "full"]):
|
|
return "long" # > 2 hours
|
|
else:
|
|
return "medium" # 30 minutes - 2 hours
|
|
|
|
async def execute_workflow(self, workflow_name: str, context: dict[str, Any]) -> dict[str, Any]:
|
|
"""Execute a predefined workflow pattern"""
|
|
if workflow_name not in self.collaboration_patterns:
|
|
raise ValueError(f"Unknown workflow: {workflow_name}")
|
|
|
|
workflow_id = f"{workflow_name}_{asyncio.get_event_loop().time()}"
|
|
self.active_workflows[workflow_id] = {
|
|
"name": workflow_name,
|
|
"status": "running",
|
|
"context": context,
|
|
"start_time": asyncio.get_event_loop().time(),
|
|
"subagents": self.collaboration_patterns[workflow_name],
|
|
"results": {},
|
|
}
|
|
|
|
try:
|
|
results = await self._execute_workflow_steps(workflow_id)
|
|
self.active_workflows[workflow_id]["status"] = "completed"
|
|
self.active_workflows[workflow_id]["results"] = results
|
|
return results
|
|
|
|
except Exception as e:
|
|
self.active_workflows[workflow_id]["status"] = "failed"
|
|
self.active_workflows[workflow_id]["error"] = str(e)
|
|
logger.error(f"Workflow {workflow_name} failed: {e}")
|
|
raise
|
|
|
|
async def _execute_workflow_steps(self, workflow_id: str) -> dict[str, Any]:
|
|
"""Execute workflow steps with proper coordination"""
|
|
workflow = self.active_workflows[workflow_id]
|
|
subagent_names = workflow["subagents"]
|
|
context = workflow["context"]
|
|
|
|
results = {}
|
|
|
|
# For now, execute sequentially - in real implementation,
|
|
# this would use intelligent parallel execution
|
|
for subagent_name in subagent_names:
|
|
if subagent_name in self.subagents:
|
|
subagent = self.subagents[subagent_name]
|
|
|
|
# Update subagent status
|
|
subagent.status = SubagentStatus.BUSY
|
|
subagent.current_task = workflow_id
|
|
|
|
try:
|
|
# Simulate subagent execution
|
|
result = await self._simulate_subagent_execution(subagent_name, context)
|
|
results[subagent_name] = result
|
|
|
|
# Update context with results for next subagent
|
|
context.update({f"{subagent_name}_result": result})
|
|
|
|
finally:
|
|
# Reset subagent status
|
|
subagent.status = SubagentStatus.AVAILABLE
|
|
subagent.current_task = None
|
|
|
|
return results
|
|
|
|
async def _simulate_subagent_execution(self, subagent_name: str, _context: dict[str, Any]) -> dict[str, Any]:
|
|
"""Simulate subagent execution (placeholder for actual implementation)"""
|
|
subagent = self.subagents[subagent_name]
|
|
|
|
# Simulate processing time based on subagent complexity
|
|
processing_time = len(subagent.config.get("capabilities", [])) * 0.5
|
|
await asyncio.sleep(processing_time)
|
|
|
|
return {
|
|
"status": "completed",
|
|
"subagent": subagent_name,
|
|
"capabilities_used": subagent.config.get("capabilities", []),
|
|
"processing_time": processing_time,
|
|
"recommendations": f"Recommendations from {subagent_name}",
|
|
"artifacts": [f"artifact_{subagent_name}.json"],
|
|
}
|
|
|
|
def get_system_status(self) -> dict[str, Any]:
|
|
"""Get comprehensive system status"""
|
|
status_summary = {
|
|
"total_subagents": len(self.subagents),
|
|
"available_subagents": sum(1 for s in self.subagents.values() if s.status == SubagentStatus.AVAILABLE),
|
|
"busy_subagents": sum(1 for s in self.subagents.values() if s.status == SubagentStatus.BUSY),
|
|
"active_workflows": len([w for w in self.active_workflows.values() if w["status"] == "running"]),
|
|
"subagent_details": {},
|
|
}
|
|
|
|
for name, instance in self.subagents.items():
|
|
status_summary["subagent_details"][name] = {
|
|
"status": instance.status.value,
|
|
"current_task": instance.current_task,
|
|
"capabilities": len(instance.config.get("capabilities", [])),
|
|
"load_factor": instance.load_factor,
|
|
}
|
|
|
|
return status_summary
|
|
|
|
def get_available_workflows(self) -> list[str]:
|
|
"""Get list of available workflow patterns"""
|
|
return list(self.collaboration_patterns.keys())
|
|
|
|
def recommend_subagents(self, task_description: str) -> dict[str, Any]:
|
|
"""Recommend optimal subagents for a given task"""
|
|
analysis = self.analyze_task_requirements(task_description)
|
|
|
|
# Find subagents with required capabilities
|
|
recommended_subagents = []
|
|
for capability in analysis["capabilities_needed"]:
|
|
capable_subagents = self.find_subagents_by_capability(capability)
|
|
recommended_subagents.extend(capable_subagents)
|
|
|
|
# Remove duplicates and prioritize
|
|
recommended_subagents = list(set(recommended_subagents))
|
|
|
|
# Add priority subagents
|
|
for priority_subagent in analysis["priority_subagents"]:
|
|
if priority_subagent not in recommended_subagents:
|
|
recommended_subagents.append(priority_subagent)
|
|
|
|
# Check availability
|
|
available_subagents = [
|
|
name
|
|
for name in recommended_subagents
|
|
if name in self.subagents and self.subagents[name].status == SubagentStatus.AVAILABLE
|
|
]
|
|
|
|
return {
|
|
"task_analysis": analysis,
|
|
"recommended_subagents": recommended_subagents,
|
|
"available_subagents": available_subagents,
|
|
"coordination_recommendation": self._recommend_coordination_pattern(recommended_subagents),
|
|
}
|
|
|
|
def _recommend_coordination_pattern(self, subagents: list[str]) -> dict[str, Any]:
|
|
"""Recommend coordination pattern based on subagent combination"""
|
|
# Find matching workflow patterns
|
|
matching_patterns = []
|
|
for pattern_name, pattern_subagents in self.collaboration_patterns.items():
|
|
overlap = set(subagents) & set(pattern_subagents)
|
|
if overlap:
|
|
matching_patterns.append(
|
|
{
|
|
"pattern": pattern_name,
|
|
"overlap_count": len(overlap),
|
|
"coverage": len(overlap) / len(set(subagents) | set(pattern_subagents)),
|
|
}
|
|
)
|
|
|
|
# Sort by coverage
|
|
matching_patterns.sort(key=lambda x: x["coverage"], reverse=True)
|
|
|
|
return {
|
|
"recommended_pattern": matching_patterns[0]["pattern"] if matching_patterns else None,
|
|
"pattern_options": matching_patterns,
|
|
"custom_coordination_needed": len(matching_patterns) == 0,
|
|
}
|
|
|
|
|
|
# CLI Interface
|
|
def main():
|
|
"""Command-line interface for subagent management"""
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="Claude Code Subagent Manager")
|
|
parser.add_argument("--status", action="store_true", help="Show system status")
|
|
parser.add_argument("--workflows", action="store_true", help="List available workflows")
|
|
parser.add_argument("--recommend", type=str, help="Recommend subagents for task")
|
|
parser.add_argument("--execute", type=str, help="Execute workflow")
|
|
|
|
args = parser.parse_args()
|
|
|
|
manager = SubagentManager()
|
|
|
|
if args.status:
|
|
status = manager.get_system_status()
|
|
print(json.dumps(status, indent=2))
|
|
|
|
elif args.workflows:
|
|
workflows = manager.get_available_workflows()
|
|
print("Available workflows:")
|
|
for workflow in workflows:
|
|
print(f" - {workflow}")
|
|
|
|
elif args.recommend:
|
|
recommendations = manager.recommend_subagents(args.recommend)
|
|
print(json.dumps(recommendations, indent=2))
|
|
|
|
elif args.execute:
|
|
|
|
async def run_workflow():
|
|
result = await manager.execute_workflow(args.execute, {})
|
|
print(json.dumps(result, indent=2))
|
|
|
|
asyncio.run(run_workflow())
|
|
|
|
else:
|
|
parser.print_help()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|