massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/tests/test_claude_code_context_sharing.py

@@ -0,0 +1,233 @@
+# -*- coding: utf-8 -*-
+"""
+Test script for Claude Code Context Sharing functionality.
+
+This script tests the context sharing capabilities between multiple Claude Code agents,
+ensuring that:
+1. Workspace snapshots are saved correctly
+2. Snapshots are restored with anonymization
+3. Agents can access each other's work through shared context
+"""
+
+import shutil
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+from massgen.chat_agent import ChatAgent
+from massgen.orchestrator import Orchestrator
+
+
+class MockClaudeCodeBackend:
+    """Mock Claude Code backend for testing."""
+
+    def __init__(self, cwd: str = None):
+        self._cwd = cwd or "test_workspace"
+
+    def get_provider_name(self) -> str:
+        return "claude_code"
+
+
+class MockClaudeCodeAgent(ChatAgent):
+    """Mock Claude Code agent for testing."""
+
+    def __init__(self, agent_id: str, cwd: str = None):
+        super().__init__(session_id=f"session_{agent_id}")
+        self.agent_id = agent_id
+        self.backend = MockClaudeCodeBackend(cwd)
+
+    async def chat(self, messages, tools=None, reset_chat=False, clear_history=False):
+        """Mock chat implementation."""
+        for _ in range(3):
+            yield {
+                "type": "content",
+                "content": f"Working on task from {self.agent_id}",
+            }
+        yield {"type": "result", "data": ("answer", f"Solution from {self.agent_id}")}
+
+
+@pytest.fixture
+def test_workspace(tmp_path):
+    """Create temporary test workspace."""
+    workspace = tmp_path / "test_context_sharing"
+    workspace.mkdir(exist_ok=True)
+
+    # Create test directories
+    snapshot_storage = workspace / "snapshots"
+    temp_workspace = workspace / "temp_workspaces"
+
+    snapshot_storage.mkdir(exist_ok=True)
+    temp_workspace.mkdir(exist_ok=True)
+
+    yield {
+        "workspace": workspace,
+        "snapshot_storage": str(snapshot_storage),
+        "temp_workspace": str(temp_workspace),
+    }
+
+    # Cleanup
+    if workspace.exists():
+        shutil.rmtree(workspace)
+
+
+@pytest.fixture
+def mock_agents():
+    """Create mock Claude Code agents."""
+    agents = {}
+    for i in range(1, 4):
+        agent_id = f"claude_code_{i}"
+        cwd = f"test_workspace/agent_{i}"
+        agents[agent_id] = MockClaudeCodeAgent(agent_id, cwd)
+    return agents
+
+
+def test_orchestrator_initialization_with_context_sharing(test_workspace, mock_agents):
+    """Test orchestrator initializes with context sharing parameters."""
+
+    orchestrator = Orchestrator(
+        agents=mock_agents,
+        snapshot_storage=test_workspace["snapshot_storage"],
+        agent_temporary_workspace=test_workspace["temp_workspace"],
+    )
+
+    # Check that parameters are set
+    assert orchestrator._snapshot_storage == test_workspace["snapshot_storage"]
+    assert orchestrator._agent_temporary_workspace == test_workspace["temp_workspace"]
+
+    # Check that agent ID mappings are created
+    assert len(orchestrator._agent_id_mapping) == 3
+    assert "claude_code_1" in orchestrator._agent_id_mapping
+    assert "claude_code_2" in orchestrator._agent_id_mapping
+    assert "claude_code_3" in orchestrator._agent_id_mapping
+
+    # Check anonymized IDs
+    assert orchestrator._agent_id_mapping["claude_code_1"] == "agent_1"
+    assert orchestrator._agent_id_mapping["claude_code_2"] == "agent_2"
+    assert orchestrator._agent_id_mapping["claude_code_3"] == "agent_3"
+
+    # Check directories are created
+    assert Path(test_workspace["snapshot_storage"]).exists()
+    assert Path(test_workspace["temp_workspace"]).exists()
+
+    for agent_id in mock_agents.keys():
+        snapshot_dir = Path(test_workspace["snapshot_storage"]) / agent_id
+        temp_dir = Path(test_workspace["temp_workspace"]) / agent_id
+        assert snapshot_dir.exists()
+        assert temp_dir.exists()
+
+
+@pytest.mark.asyncio
+async def test_snapshot_saving(test_workspace, mock_agents):
+    """Test that snapshots are saved correctly when agents complete tasks."""
+
+    orchestrator = Orchestrator(
+        agents=mock_agents,
+        snapshot_storage=test_workspace["snapshot_storage"],
+        agent_temporary_workspace=test_workspace["temp_workspace"],
+    )
+
+    # Create some test files in agent workspaces
+    for agent_id, agent in mock_agents.items():
+        workspace = Path(agent.backend._cwd)
+        workspace.mkdir(parents=True, exist_ok=True)
+
+        # Create test files
+        (workspace / f"code_{agent_id}.py").write_text(f"# Code from {agent_id}")
+        (workspace / f"test_{agent_id}.txt").write_text(f"Test data from {agent_id}")
+
+    # Save snapshot for one agent
+    await orchestrator._save_claude_code_snapshot("claude_code_1")
+
+    # Check snapshot was saved
+    snapshot_dir = Path(test_workspace["snapshot_storage"]) / "claude_code_1"
+    assert (snapshot_dir / "code_claude_code_1.py").exists()
+    assert (snapshot_dir / "test_claude_code_1.txt").exists()
+
+
+@pytest.mark.asyncio
+async def test_workspace_restoration_with_anonymization(test_workspace, mock_agents):
+    """Test that workspaces are restored with anonymized agent IDs."""
+
+    orchestrator = Orchestrator(
+        agents=mock_agents,
+        snapshot_storage=test_workspace["snapshot_storage"],
+        agent_temporary_workspace=test_workspace["temp_workspace"],
+    )
+
+    # Create snapshots for all agents
+    for agent_id, agent in mock_agents.items():
+        workspace = Path(agent.backend._cwd)
+        workspace.mkdir(parents=True, exist_ok=True)
+        (workspace / f"work_{agent_id}.txt").write_text(f"Work from {agent_id}")
+
+        # Save to snapshot
+        await orchestrator._save_claude_code_snapshot(agent_id)
+
+    # Restore workspace for agent 2
+    workspace_path = await orchestrator._restore_snapshots_to_workspace("claude_code_2")
+
+    assert workspace_path is not None
+    workspace_dir = Path(workspace_path)
+
+    # Check that all snapshots are restored with anonymized names
+    assert (workspace_dir / "agent_1" / "work_claude_code_1.txt").exists()
+    assert (workspace_dir / "agent_2" / "work_claude_code_2.txt").exists()
+    assert (workspace_dir / "agent_3" / "work_claude_code_3.txt").exists()
+
+    # Verify content
+    content1 = (workspace_dir / "agent_1" / "work_claude_code_1.txt").read_text()
+    assert content1 == "Work from claude_code_1"
+
+
+@pytest.mark.asyncio
+async def test_save_all_snapshots(test_workspace, mock_agents):
+    """Test saving all Claude Code agent snapshots at once."""
+
+    orchestrator = Orchestrator(
+        agents=mock_agents,
+        snapshot_storage=test_workspace["snapshot_storage"],
+        agent_temporary_workspace=test_workspace["temp_workspace"],
+    )
+
+    # Create test files in all agent workspaces
+    for agent_id, agent in mock_agents.items():
+        workspace = Path(agent.backend._cwd)
+        workspace.mkdir(parents=True, exist_ok=True)
+        (workspace / "shared.py").write_text(f"# Shared code from {agent_id}")
+
+    # Save all snapshots
+    await orchestrator._save_all_claude_code_snapshots()
+
+    # Verify all snapshots were saved
+    for agent_id in mock_agents.keys():
+        snapshot_dir = Path(test_workspace["snapshot_storage"]) / agent_id
+        assert (snapshot_dir / "shared.py").exists()
+        content = (snapshot_dir / "shared.py").read_text()
+        assert agent_id in content
+
+
+def test_non_claude_code_agents_ignored(test_workspace):
+    """Test that non-Claude Code agents are ignored for context sharing."""
+
+    # Create mixed agents (some Claude Code, some not)
+    agents = {
+        "claude_code_1": MockClaudeCodeAgent("claude_code_1"),
+        "regular_agent": MagicMock(backend=MagicMock(get_provider_name=lambda: "openai")),
+    }
+
+    orchestrator = Orchestrator(
+        agents=agents,
+        snapshot_storage=test_workspace["snapshot_storage"],
+        agent_temporary_workspace=test_workspace["temp_workspace"],
+    )
+
+    # Only Claude Code agents should have mappings
+    assert "claude_code_1" in orchestrator._agent_id_mapping
+    assert "regular_agent" not in orchestrator._agent_id_mapping
+    assert len(orchestrator._agent_id_mapping) == 1
+
+
+if __name__ == "__main__":
+    # Run tests
+    pytest.main([__file__, "-v"])

massgen/tests/test_claude_code_orchestrator.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Test ClaudeCodeBackend with MassGen Orchestrator.
+This test demonstrates a workflow with Claude Code backend.
+"""
+
+import asyncio
+import os
+import sys
+
+from massgen.agent_config import AgentConfig
+from massgen.backend.claude_code import ClaudeCodeBackend
+from massgen.chat_agent import ConfigurableAgent
+from massgen.frontend.coordination_ui import CoordinationUI
+from massgen.orchestrator import Orchestrator
+
+sys.path.insert(0, "/workspaces/MassGen")
+
+
+async def test_claude_code_with_orchestrator():
+    """Test Claude Code backend with MassGen Orchestrator."""
+
+    print("🚀 Testing Claude Code Backend with MassGen Orchestrator")
+    print("=" * 60)
+
+    # Create Claude Code backend
+    backend = ClaudeCodeBackend()
+
+    print(f"✅ Backend initialized: {backend.get_provider_name()}")
+    print(f"📊 Stateful backend: {backend.is_stateful()}")
+    print(f"🛠️  Supported tools: {len(backend.get_supported_builtin_tools())} tools")
+
+    # Create agent configuration
+    agent_config = AgentConfig(
+        agent_id="claude_code_agent",
+        # custom_system_instruction="You are a helpful AI assistant. Provide clear, accurate answers.",
+        backend_params={
+            "model": "claude-sonnet-4-20250514",
+        },
+    )
+
+    # Create configurable agent
+    agent = ConfigurableAgent(config=agent_config, backend=backend)
+
+    print(f"✅ Agent created: {agent.agent_id}")
+    print(f"📋 Agent status: {agent.get_status()}")
+
+    # Create orchestrator
+    orchestrator = Orchestrator(agents={agent.agent_id: agent})
+    print(f"✅ Orchestrator created with {len(orchestrator.agents)} agent(s)")
+
+    # Create coordination UI
+    ui = CoordinationUI(display_type="rich_terminal", logging_enabled=True)
+    print(f"✅ Coordination UI created: {ui.display_type}")
+
+    # Test question
+    question = "2+2=?"
+
+    print(f"\n📝 Test Question: {question}")
+    print("\n🔄 Starting orchestrator coordination...")
+    print("=" * 60)
+
+    try:
+        # Run orchestrator coordination
+        final_response = await ui.coordinate(orchestrator, question)
+
+        print("\n" + "=" * 60)
+        print("✅ Orchestrator coordination completed!")
+        print(f"📊 Final response length: {len(final_response)} characters")
+
+        # Display backend statistics
+        token_usage = backend.get_token_usage()
+        print("\n📈 Backend Statistics:")
+        print(f"   Input tokens: {token_usage.input_tokens}")
+        print(f"   Output tokens: {token_usage.output_tokens}")
+        print(f"   Estimated cost: ${token_usage.estimated_cost:.4f}")
+        print(f"   Session ID: {backend.get_current_session_id()}")
+
+        # Display final response
+        if final_response:
+            print("\n📄 Final Response:")
+            print("-" * 40)
+            print(final_response)
+            print("-" * 40)
+
+    except Exception as e:
+        print(f"❌ Error during orchestrator test: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return
+
+    print("\n✅ Orchestrator test completed successfully!")
+
+
+async def test_stateful_behavior():
+    """Test the stateful behavior of Claude Code backend."""
+
+    if not os.getenv("ANTHROPIC_API_KEY"):
+        print("❌ ANTHROPIC_API_KEY not found - skipping stateful test")
+        return
+
+    print("\n" + "=" * 60)
+    print("🚀 Testing Stateful Behavior")
+    print("=" * 60)
+
+    # Create backend and agent
+    backend = ClaudeCodeBackend(model="claude-sonnet-4-20250514")
+    config = AgentConfig(
+        agent_id="stateful_agent",
+        backend_params={
+            "model": "claude-sonnet-4-20250514",
+            "append_system_prompt": "You are a helpful assistant. Remember our conversation context.",
+        },
+    )
+    agent = ConfigurableAgent(config=config, backend=backend)
+
+    print(f"✅ Stateful agent created: {backend.is_stateful()}")
+    print(f"🔗 Initial session ID: {backend.get_current_session_id()}")
+
+    # Test conversation continuity
+    messages1 = [{"role": "user", "content": "My favorite color is blue. Please remember this."}]
+    messages2 = [{"role": "user", "content": "What did I just tell you about my favorite color?"}]
+
+    try:
+        print("\n📝 Turn 1: Setting context...")
+        response1 = ""
+        async for chunk in agent.chat(messages1):
+            if chunk.type == "content" and chunk.content:
+                response1 += chunk.content
+
+        print(f"✅ Turn 1 completed. Session: {backend.get_current_session_id()}")
+        print(f"📄 Response 1 preview: {response1[:100]}...")
+
+        print("\n📝 Turn 2: Testing memory...")
+        response2 = ""
+        async for chunk in agent.chat(messages2):
+            if chunk.type == "content" and chunk.content:
+                response2 += chunk.content
+
+        print(f"✅ Turn 2 completed. Session: {backend.get_current_session_id()}")
+        print(f"📄 Response 2 preview: {response2[:100]}...")
+
+        # Check if context was maintained
+        if "blue" in response2.lower():
+            print("✅ Stateful behavior confirmed - context maintained!")
+        else:
+            print("⚠️  Context may not have been maintained")
+
+        print(f"\n📊 Token usage after 2 turns: {backend.get_token_usage()}")
+
+    except Exception as e:
+        print(f"❌ Error during stateful test: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return
+
+    print("\n✅ Stateful behavior test completed!")
+
+
+async def main():
+    """Run all Claude Code tests."""
+    print("🧪 Claude Code Backend + Orchestrator Tests")
+    print("=" * 60)
+
+    await test_claude_code_with_orchestrator()
+    # await test_stateful_behavior()
+
+    print("\n🎉 All tests completed!")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

massgen/tests/test_cli_backends.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Test script for CLI backends - Claude Code CLI and Gemini CLI integration.
+
+This script tests the basic functionality of CLI backends without requiring
+the actual CLI tools to be installed (mocked for testing).
+
+TODO: This file is outdated - ClaudeCodeCLIBackend was removed, only SDK-based ClaudeCodeBackend remains.
+Update tests to reflect current backend architecture.
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+# Add project root to path
+project_root = Path(__file__).parent
+sys.path.insert(0, str(project_root))
+
+try:
+    # from massgen.backend.claude_code_cli import ClaudeCodeCLIBackend  # File removed
+    from massgen.backend.cli_base import CLIBackend
+except ImportError as e:
+    print(f"❌ Import error: {e}")
+    print("Make sure you're running from the project root directory")
+    sys.exit(1)
+
+
+class MockCLIBackend(CLIBackend):
+    """Mock CLI backend for testing purposes."""
+
+    def __init__(self, cli_command: str, mock_output: str = "Mock response", **kwargs):
+        self.mock_output = mock_output
+        # Skip the actual CLI tool check
+        self.cli_command = cli_command
+        self.working_dir = kwargs.get("working_dir", Path.cwd())
+        self.timeout = kwargs.get("timeout", 300)
+        self.config = kwargs
+        from massgen.backend.base import TokenUsage
+
+        self.token_usage = TokenUsage()
+
+    def _build_command(self, messages, tools, **kwargs):
+        return ["echo", "mock command"]
+
+    def _parse_output(self, output):
+        return {"content": self.mock_output, "tool_calls": [], "raw_response": output}
+
+    async def _execute_cli_command(self, command):
+        """Mock command execution."""
+        await asyncio.sleep(0.1)  # Simulate some delay
+        return self.mock_output
+
+    def get_cost_per_token(self):
+        """Mock cost per token."""
+        return {"input": 0.001, "output": 0.002}
+
+
+async def test_cli_base_functionality():
+    """Test the CLI base class functionality."""
+    print("🧪 Testing CLI base functionality...")
+
+    backend = MockCLIBackend("mock-cli", "Hello from mock CLI!")
+
+    messages = [{"role": "user", "content": "Test message"}]
+    tools = []
+
+    chunks = []
+    async for chunk in backend.stream_with_tools(messages, tools):
+        chunks.append(chunk)
+
+    assert len(chunks) > 0, "Should produce at least one chunk"
+    assert any(chunk.type == "content" for chunk in chunks), "Should have content chunk"
+    assert any(chunk.type == "done" for chunk in chunks), "Should have done chunk"
+
+    print("✅ CLI base functionality test passed")
+
+
+def test_claude_code_cli_command_building():
+    """Test Claude Code CLI command building (without executing) - SKIPPED: File removed."""
+    print("🧪 Testing Claude Code CLI command building... SKIPPED (file removed)")
+    print("✅ Claude Code CLI command building test skipped")
+
+    # NOTE: ClaudeCodeCLIBackend was removed, only ClaudeCodeBackend (SDK-based) remains
+
+
+def test_configuration_files():
+    """Test that configuration files are valid."""
+    print("🧪 Testing configuration files...")
+
+    import yaml
+
+    config_files = [
+        "massgen/configs/claude_code_cli.yaml",
+        "massgen/configs/cli_backends_mixed.yaml",
+    ]
+
+    for config_file in config_files:
+        if Path(config_file).exists():
+            try:
+                with open(config_file, "r") as f:
+                    config = yaml.safe_load(f)
+                assert config is not None, f"Config {config_file} should not be empty"
+                print(f"✅ {config_file} is valid")
+            except Exception as e:
+                print(f"❌ {config_file} is invalid: {e}")
+                raise
+        else:
+            print(f"⚠️  {config_file} not found, skipping")
+
+
+async def test_end_to_end_mock():
+    """Test end-to-end functionality with mocked CLI execution."""
+    print("🧪 Testing end-to-end with mock execution...")
+
+    # Test Claude Code CLI mock
+    claude_backend = MockCLIBackend("claude", '{"response": "4", "reasoning": "2+2 equals 4"}')
+
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    tools = []
+
+    chunks = []
+    async for chunk in claude_backend.stream_with_tools(messages, tools):
+        chunks.append(chunk)
+        print(f"  📝 Chunk: {chunk.type} - {chunk.content}")
+
+    assert len(chunks) >= 3, "Should have content, complete_message, and done chunks"
+
+    print("✅ End-to-end mock test passed")
+
+
+async def main():
+    """Run all tests."""
+    print("🚀 Starting CLI backend tests...\n")
+
+    try:
+        # Test basic functionality
+        await test_cli_base_functionality()
+        print()
+
+        # Test command building
+        test_claude_code_cli_command_building()
+        print()
+
+        # Test configuration files
+        test_configuration_files()
+        print()
+
+        # Test end-to-end mock
+        await test_end_to_end_mock()
+        print()
+
+        print("🎉 All CLI backend tests passed!")
+
+        # Show usage information
+        print("\n📋 Usage Information:")
+        print("CLI backends are now available in MassGen!")
+        print()
+        print("Prerequisites:")
+        print("  • Claude Code CLI: npm install -g @anthropic-ai/claude-code")
+        print("  • Gemini CLI: npm install -g @google/gemini-cli")
+        print()
+        print("Usage examples:")
+        print("  # Claude Code (SDK-based)")
+        print("  uv run python -m massgen.cli --backend claude_code --model claude-sonnet-4-20250514 'What is 2+2?'")
+        print()
+        print("  # Mixed CLI backends")
+        print("  uv run python -m massgen.cli --config massgen/configs/cli_backends_mixed.yaml 'Complex question'")
+
+    except Exception as e:
+        print(f"❌ Test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())