kolega-code 0.1.0__py3-none-any.whl

kolega_code/agent/tests/test_parallel_tool_calls.py

@@ -0,0 +1,291 @@
+"""Tests for parallel tool-call execution and task-local tool-call IDs."""
+
+import asyncio
+import os
+import uuid
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from dotenv import load_dotenv
+
+from kolega_code.agent.baseagent import BaseAgent
+from kolega_code.config import AgentConfig, ModelConfig, ModelProvider, RateLimitConfig
+from kolega_code.events import AgentConnectionManager
+from kolega_code.llm.models import ToolCall
+
+# Load environment variables
+load_dotenv()
+
+
+@pytest.fixture
+def agent_config():
+    return AgentConfig(
+        anthropic_api_key=os.getenv("ANTHROPIC_API_KEY", "test_key"),
+        openai_api_key="test-key",
+        long_context_config=ModelConfig(
+            provider=ModelProvider.ANTHROPIC,
+            model="claude-haiku-4-5-20251001",
+            rate_limits=RateLimitConfig(),
+        ),
+        fast_config=ModelConfig(
+            provider=ModelProvider.ANTHROPIC,
+            model="claude-haiku-4-5-20251001",
+            rate_limits=RateLimitConfig(),
+        ),
+        thinking_config=ModelConfig(
+            provider=ModelProvider.ANTHROPIC,
+            model="claude-haiku-4-5-20251001",
+            rate_limits=RateLimitConfig(),
+            thinking_tokens=1024,
+        ),
+    )
+
+
+@pytest.fixture
+def mock_connection_manager():
+    return AsyncMock(spec=AgentConnectionManager)
+
+
+@pytest.fixture
+def base_agent(tmp_path, mock_connection_manager, agent_config):
+    agent = BaseAgent(
+        project_path=tmp_path,
+        workspace_id="test_workspace",
+        thread_id=str(uuid.uuid4()),
+        connection_manager=mock_connection_manager,
+        config=agent_config,
+    )
+    agent.send_chat_message = AsyncMock()
+    agent.log_info = AsyncMock()
+    agent.log_error = AsyncMock()
+    return agent
+
+
+def make_tool_call(name: str, index: int) -> ToolCall:
+    return ToolCall(
+        id=f"{name}_{index}",
+        name=name,
+        input={"task": f"task {index}"},
+        execution_id=f"exec_{name}_{index}",
+    )
+
+
+class ConcurrencyTracker:
+    """Tracks how many fake tool executions overlap."""
+
+    def __init__(self):
+        self.active = 0
+        self.max_active = 0
+
+    async def run(self, duration: float = 0.01):
+        self.active += 1
+        self.max_active = max(self.max_active, self.active)
+        await asyncio.sleep(duration)
+        self.active -= 1
+
+
+class FakeToolCollection:
+    """Test stand-in for ToolCollection: builds a registry from get_tool_list + methods."""
+
+    def registry(self):
+        from kolega_code.agent.tools import ToolCollection
+        from kolega_code.llm.models import ToolDefinition
+        from kolega_code.tools import Tool, ToolRegistry
+
+        parallel = set(ToolCollection.read_only_tools) | set(ToolCollection.agent_dispatch_tools)
+        registry = ToolRegistry()
+        for spec in self.get_tool_list():
+            registry.add(
+                Tool(
+                    name=spec.name,
+                    definition=ToolDefinition(name=spec.name, description="", parameters=[]),
+                    handler=getattr(self, spec.name),
+                    parallel_safe=spec.name in parallel,
+                )
+            )
+        return registry
+
+
+class TestParallelToolCalls:
+    @pytest.mark.asyncio
+    async def test_dispatch_tools_run_concurrently(self, base_agent):
+        first_started = asyncio.Event()
+        second_started = asyncio.Event()
+
+        class Tools(FakeToolCollection):
+            def __init__(self):
+                self.calls = 0
+
+            def get_tool_list(self):
+                return [SimpleNamespace(name="dispatch_general_agent")]
+
+            async def dispatch_general_agent(self, task: str):
+                if "task 0" in task:
+                    first_started.set()
+                    # Deadlocks unless the second call runs concurrently
+                    await asyncio.wait_for(second_started.wait(), timeout=2)
+                    return "first done"
+                second_started.set()
+                await asyncio.wait_for(first_started.wait(), timeout=2)
+                return "second done"
+
+        base_agent.tool_collection = Tools()
+        blocks = [make_tool_call("dispatch_general_agent", i) for i in range(2)]
+
+        results = await asyncio.wait_for(base_agent.process_tool_calls(blocks), timeout=5)
+
+        assert [r.content for r in results] == ["first done", "second done"]
+        assert [r.tool_use_id for r in results] == ["dispatch_general_agent_0", "dispatch_general_agent_1"]
+        assert not any(r.is_error for r in results)
+
+    @pytest.mark.asyncio
+    async def test_mixed_read_only_and_dispatch_parallelize(self, base_agent):
+        tracker = ConcurrencyTracker()
+
+        class Tools(FakeToolCollection):
+            def get_tool_list(self):
+                return [
+                    SimpleNamespace(name="dispatch_general_agent"),
+                    SimpleNamespace(name="read_entire_file"),
+                ]
+
+            async def dispatch_general_agent(self, task: str):
+                await tracker.run()
+                return "dispatched"
+
+            async def read_entire_file(self, task: str):
+                await tracker.run()
+                return "file contents"
+
+        base_agent.tool_collection = Tools()
+        blocks = [
+            make_tool_call("read_entire_file", 0),
+            make_tool_call("dispatch_general_agent", 1),
+            make_tool_call("read_entire_file", 2),
+        ]
+
+        results = await base_agent.process_tool_calls(blocks)
+
+        assert tracker.max_active > 1
+        assert len(results) == 3
+
+    @pytest.mark.asyncio
+    async def test_write_tool_forces_sequential(self, base_agent):
+        tracker = ConcurrencyTracker()
+
+        class Tools(FakeToolCollection):
+            def get_tool_list(self):
+                return [
+                    SimpleNamespace(name="dispatch_general_agent"),
+                    SimpleNamespace(name="create_file"),
+                ]
+
+            async def dispatch_general_agent(self, task: str):
+                await tracker.run()
+                return "dispatched"
+
+            async def create_file(self, task: str):
+                await tracker.run()
+                return "created"
+
+        base_agent.tool_collection = Tools()
+        blocks = [
+            make_tool_call("dispatch_general_agent", 0),
+            make_tool_call("create_file", 1),
+        ]
+
+        results = await base_agent.process_tool_calls(blocks)
+
+        assert tracker.max_active == 1
+        assert len(results) == 2
+
+    @pytest.mark.asyncio
+    async def test_semaphore_caps_concurrency(self, base_agent):
+        tracker = ConcurrencyTracker()
+        total = BaseAgent.PARALLEL_TOOL_LIMIT + 4
+
+        class Tools(FakeToolCollection):
+            def get_tool_list(self):
+                return [SimpleNamespace(name="dispatch_general_agent")]
+
+            async def dispatch_general_agent(self, task: str):
+                await tracker.run()
+                return "done"
+
+        base_agent.tool_collection = Tools()
+        blocks = [make_tool_call("dispatch_general_agent", i) for i in range(total)]
+
+        results = await base_agent.process_tool_calls(blocks)
+
+        assert len(results) == total
+        assert tracker.max_active > 1
+        assert tracker.max_active <= BaseAgent.PARALLEL_TOOL_LIMIT
+
+    @pytest.mark.asyncio
+    async def test_contextvar_isolation_under_gather(self, base_agent):
+        captured: dict[str, str] = {}
+
+        class Tools(FakeToolCollection):
+            def get_tool_list(self):
+                return [SimpleNamespace(name="dispatch_general_agent")]
+
+            async def dispatch_general_agent(self, task: str):
+                # Yield so concurrent executions interleave before reading the ID
+                await asyncio.sleep(0.01)
+                captured[task] = base_agent.current_tool_execution_id
+                return "done"
+
+        base_agent.tool_collection = Tools()
+        blocks = [make_tool_call("dispatch_general_agent", i) for i in range(3)]
+
+        await base_agent.process_tool_calls(blocks)
+
+        assert captured == {
+            "task 0": "exec_dispatch_general_agent_0",
+            "task 1": "exec_dispatch_general_agent_1",
+            "task 2": "exec_dispatch_general_agent_2",
+        }
+        assert base_agent.current_tool_call_id is None
+        assert base_agent.current_tool_execution_id is None
+        assert base_agent.current_provider_tool_call_id is None
+
+    @pytest.mark.asyncio
+    async def test_nested_agent_does_not_clobber_parent_ids(
+        self, tmp_path, mock_connection_manager, agent_config, base_agent
+    ):
+        nested_agent = BaseAgent(
+            project_path=tmp_path,
+            workspace_id="test_workspace",
+            thread_id=str(uuid.uuid4()),
+            connection_manager=mock_connection_manager,
+            config=agent_config,
+        )
+        nested_agent.send_chat_message = AsyncMock()
+        nested_agent.log_info = AsyncMock()
+
+        class NestedTools(FakeToolCollection):
+            def get_tool_list(self):
+                return [SimpleNamespace(name="read_entire_file")]
+
+            async def read_entire_file(self, task: str):
+                return "nested contents"
+
+        nested_agent.tool_collection = NestedTools()
+        observed = {}
+
+        class ParentTools(FakeToolCollection):
+            def get_tool_list(self):
+                return [SimpleNamespace(name="dispatch_general_agent")]
+
+            async def dispatch_general_agent(self, task: str):
+                # Simulate a sub-agent running its own tool within the same asyncio task
+                await nested_agent.execute_single_tool(make_tool_call("read_entire_file", 99))
+                observed["parent_id"] = base_agent.current_tool_execution_id
+                return "done"
+
+        base_agent.tool_collection = ParentTools()
+
+        await base_agent.process_tool_calls([make_tool_call("dispatch_general_agent", 0)])
+
+        assert observed["parent_id"] == "exec_dispatch_general_agent_0"

kolega_code/agent/tests/test_planning_agent.py

@@ -0,0 +1,227 @@
+import base64
+import uuid
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+from kolega_code.agent.baseagent import BaseAgent
+from kolega_code.config import AgentConfig, ModelConfig, ModelProvider, RateLimitConfig
+from kolega_code.events import AgentConnectionManager
+from kolega_code.llm.models import Message, TextBlock, ToolCall
+from kolega_code.llm.providers.models import TokenCount
+from kolega_code.agent.planningagent import PlanningAgent
+from kolega_code.agent.prompt_provider import AgentMode
+from kolega_code.agent.tools import ToolExtension
+
+
+def _deepseek_config() -> AgentConfig:
+    model_config = ModelConfig(
+        provider=ModelProvider.DEEPSEEK,
+        model="deepseek-v4-pro",
+        rate_limits=RateLimitConfig(),
+    )
+    return AgentConfig(
+        deepseek_api_key="test-key",
+        long_context_config=model_config,
+        fast_config=model_config,
+        edit_model_config=model_config,
+        thinking_config=model_config,
+    )
+
+
+def _image_attachment() -> dict:
+    return {
+        "type": "image",
+        "media_type": "image/png",
+        "data": base64.b64encode(b"fake-image-data").decode("utf-8"),
+        "filename": "test-image.png",
+    }
+
+
+class _EmptyStream:
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        raise StopAsyncIteration
+
+    async def get_final_message(self):
+        return Message("assistant", [TextBlock("done")], stop_reason="end_turn")
+
+
+@pytest.fixture
+def mock_connection_manager():
+    manager = Mock(spec=AgentConnectionManager)
+    manager.workspace_id = "test_workspace"
+    manager.send_message = AsyncMock()
+    return manager
+
+
+@pytest.fixture
+def agent_config():
+    config = Mock(spec=AgentConfig)
+    config.long_context_config = Mock()
+    config.long_context_config.provider = "anthropic"
+    config.long_context_config.model = "claude-sonnet-4-5-20250929"
+    config.long_context_config.thinking_tokens = None
+    config.openai_api_key = "test_key"
+    config.anthropic_api_key = "test_key"
+    config.browser_use_headless = True
+    return config
+
+
+@pytest.mark.asyncio
+async def test_planning_agent_uses_host_task_list_extension(tmp_path, mock_connection_manager, agent_config):
+    task_list = ""
+
+    async def update_task_list(task_list_markdown: str) -> str:
+        nonlocal task_list
+        task_list = task_list_markdown.strip()
+        return "Task list updated."
+
+    async def get_task_list() -> str:
+        return task_list or "No task list has been set."
+
+    agent = PlanningAgent(
+        project_path=tmp_path,
+        workspace_id="test_workspace",
+        thread_id=str(uuid.uuid4()),
+        connection_manager=mock_connection_manager,
+        config=agent_config,
+        agent_mode=AgentMode.CLI,
+        tool_extensions=[
+            ToolExtension(
+                name="host-task-list",
+                tools={"update_task_list": update_task_list, "get_task_list": get_task_list},
+                tool_groups={"planning_tools": ["update_task_list", "get_task_list"]},
+            )
+        ],
+    )
+
+    result = await agent.tool_collection.update_task_list("- [ ] inspect CLI\n- [x] choose tool shape")
+    tool_names = {tool.name for tool in agent.tool_collection.get_tool_list()}
+
+    assert result == "Task list updated."
+    assert task_list == "- [ ] inspect CLI\n- [x] choose tool shape"
+    assert await agent.tool_collection.get_task_list() == task_list
+    assert {"write_plan", "get_task_list", "update_task_list"}.issubset(tool_names)
+    assert "replace_entire_file" not in tool_names
+    assert "create_file" not in tool_names
+
+
+def test_planning_agent_only_exposes_write_plan_without_host_task_tools(
+    tmp_path, mock_connection_manager, agent_config
+):
+    agent = PlanningAgent(
+        project_path=tmp_path,
+        workspace_id="test_workspace",
+        thread_id=str(uuid.uuid4()),
+        connection_manager=mock_connection_manager,
+        config=agent_config,
+        agent_mode=AgentMode.CLI,
+    )
+
+    tool_names = {tool.name for tool in agent.tool_collection.get_tool_list()}
+
+    assert "write_plan" in tool_names
+    assert "update_task_list" not in tool_names
+    assert "get_task_list" not in tool_names
+
+
+@pytest.mark.asyncio
+async def test_planning_agent_rejects_unavailable_file_edit_tool(
+    tmp_path, mock_connection_manager, agent_config
+):
+    target = tmp_path / "notes.txt"
+    target.write_text("original\n", encoding="utf-8")
+    agent = PlanningAgent(
+        project_path=tmp_path,
+        workspace_id="test_workspace",
+        thread_id=str(uuid.uuid4()),
+        connection_manager=mock_connection_manager,
+        config=agent_config,
+        agent_mode=AgentMode.CLI,
+    )
+
+    result = await agent.execute_single_tool(
+        ToolCall(
+            id="tool-call-1",
+            name="replace_entire_file",
+            input={"relative_path": "notes.txt", "content": "mutated\n"},
+        )
+    )
+
+    assert result.is_error is True
+    assert result.content == "Tool 'replace_entire_file' is not available in this mode."
+    assert target.read_text(encoding="utf-8") == "original\n"
+
+
+@pytest.mark.asyncio
+async def test_planning_agent_write_plan_is_consumable(tmp_path, mock_connection_manager, agent_config):
+    agent = PlanningAgent(
+        project_path=tmp_path,
+        workspace_id="test_workspace",
+        thread_id=str(uuid.uuid4()),
+        connection_manager=mock_connection_manager,
+        config=agent_config,
+        agent_mode=AgentMode.CLI,
+    )
+
+    result = await agent.tool_collection.write_plan("# Plan\n\nImplement planning mode.")
+
+    assert result == "Plan captured."
+    assert agent.consume_completed_plan() == "# Plan\n\nImplement planning mode."
+    assert agent.consume_completed_plan() is None
+
+
+@pytest.mark.asyncio
+async def test_planning_agent_rejects_deepseek_image_without_llm_call(tmp_path, mock_connection_manager):
+    agent = PlanningAgent(
+        project_path=tmp_path,
+        workspace_id="test_workspace",
+        thread_id=str(uuid.uuid4()),
+        connection_manager=mock_connection_manager,
+        config=_deepseek_config(),
+        agent_mode=AgentMode.CLI,
+    )
+    agent.llm = Mock()
+
+    chunks = [
+        chunk
+        async for chunk in agent.process_message_stream("Plan from this screenshot", [_image_attachment()])
+    ]
+
+    assert len(chunks) == 1
+    assert chunks[0]["type"] == "response"
+    assert chunks[0]["content"] == BaseAgent.deepseek_image_unsupported_message
+    assert chunks[0]["complete"] is True
+    assert agent.history == []
+    agent.llm.stream.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_planning_agent_does_not_print_context_token_counts(
+    tmp_path, mock_connection_manager, capsys
+):
+    agent = PlanningAgent(
+        project_path=tmp_path,
+        workspace_id="test_workspace",
+        thread_id=str(uuid.uuid4()),
+        connection_manager=mock_connection_manager,
+        config=_deepseek_config(),
+        agent_mode=AgentMode.CLI,
+    )
+    agent.count_current_context = AsyncMock(return_value=TokenCount(input_tokens=42))
+    agent.llm = Mock()
+    agent.llm.stream = AsyncMock(return_value=_EmptyStream())
+
+    chunks = [chunk async for chunk in agent.process_message_stream("hello")]
+
+    assert chunks[-1]["complete"] is True
+    assert capsys.readouterr().out == ""