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

control_plane_api/worker/tests/unit/test_tool_enforcement.py

@@ -0,0 +1,324 @@
+"""Unit tests for tool enforcement service."""
+
+import pytest
+import asyncio
+from unittest.mock import AsyncMock, Mock, patch
+from datetime import datetime, timezone
+
+from control_plane_api.worker.services.tool_enforcement import ToolEnforcementService
+
+
+@pytest.fixture
+def mock_enforcer_client():
+    """Mock policy enforcer client."""
+    client = AsyncMock()
+    client.evaluation = AsyncMock()
+    return client
+
+
+@pytest.fixture
+def enforcement_service(mock_enforcer_client):
+    """Create enforcement service with mock client."""
+    return ToolEnforcementService(mock_enforcer_client)
+
+
+@pytest.fixture
+def enforcement_context():
+    """Sample enforcement context."""
+    return {
+        "user_email": "test@example.com",
+        "user_id": "user-123",
+        "user_roles": ["developer"],
+        "organization_id": "org-456",
+        "team_id": "team-789",
+        "agent_id": "agent-xyz",
+        "execution_id": "exec-abc",
+        "environment": "production"
+    }
+
+
+class TestToolEnforcementService:
+    """Test suite for ToolEnforcementService."""
+
+    @pytest.mark.asyncio
+    async def test_enforce_tool_allowed(self, enforcement_service, enforcement_context, mock_enforcer_client):
+        """Test enforcement check when tool is allowed."""
+        # Mock enforcer response
+        mock_enforcer_client.evaluation.enforce.return_value = {
+            "allow": True,
+            "id": "enforcement-123",
+            "policies": ["role_based_access"]
+        }
+
+        allow, violation, metadata = await enforcement_service.enforce_tool_execution(
+            tool_name="Read",
+            tool_args={"file_path": "/tmp/test.txt"},
+            enforcement_context=enforcement_context
+        )
+
+        assert allow is True
+        assert violation is None
+        assert metadata["enforcer"] == "allowed"
+        assert "role_based_access" in metadata["policies"]
+        assert "enforcement_id" in metadata
+
+    @pytest.mark.asyncio
+    async def test_enforce_tool_blocked(self, enforcement_service, enforcement_context, mock_enforcer_client):
+        """Test enforcement check when tool is blocked."""
+        mock_enforcer_client.evaluation.enforce.return_value = {
+            "allow": False,
+            "id": "enforcement-456",
+            "policies": ["production_safeguards"]
+        }
+
+        allow, violation, metadata = await enforcement_service.enforce_tool_execution(
+            tool_name="Bash",
+            tool_args={"command": "rm -rf /tmp/*"},
+            enforcement_context=enforcement_context
+        )
+
+        assert allow is False
+        assert violation is not None
+        assert "blocked by policy enforcement" in violation.lower()
+        assert "Bash" in violation
+        assert metadata["enforcer"] == "blocked"
+        assert "production_safeguards" in metadata["policies"]
+
+    @pytest.mark.asyncio
+    async def test_enforce_timeout_fails_open(self, enforcement_service, enforcement_context, mock_enforcer_client):
+        """Test that enforcement timeout fails open (allows execution)."""
+        # Mock timeout
+        async def slow_enforce(*args, **kwargs):
+            await asyncio.sleep(5)  # Longer than timeout
+            return {"allow": False}
+
+        mock_enforcer_client.evaluation.enforce = slow_enforce
+
+        allow, violation, metadata = await enforcement_service.enforce_tool_execution(
+            tool_name="Bash",
+            tool_args={"command": "ls"},
+            enforcement_context=enforcement_context,
+            timeout=0.1  # Very short timeout
+        )
+
+        assert allow is True  # Fails open
+        assert violation is None
+        assert metadata["enforcer"] == "timeout"
+
+    @pytest.mark.asyncio
+    async def test_enforce_error_fails_open(self, enforcement_service, enforcement_context, mock_enforcer_client):
+        """Test that enforcement errors fail open (allows execution)."""
+        # Mock error
+        mock_enforcer_client.evaluation.enforce.side_effect = Exception("Enforcer unavailable")
+
+        allow, violation, metadata = await enforcement_service.enforce_tool_execution(
+            tool_name="Bash",
+            tool_args={"command": "ls"},
+            enforcement_context=enforcement_context
+        )
+
+        assert allow is True  # Fails open
+        assert violation is None
+        assert metadata["enforcer"] == "error"
+        assert "error" in metadata
+
+    @pytest.mark.asyncio
+    async def test_disabled_enforcer(self):
+        """Test that disabled enforcer allows all tools."""
+        service = ToolEnforcementService(None)
+
+        allow, violation, metadata = await service.enforce_tool_execution(
+            tool_name="Bash",
+            tool_args={"command": "rm -rf /"},
+            enforcement_context={}
+        )
+
+        assert allow is True
+        assert violation is None
+        assert metadata["enforcer"] == "disabled"
+
+    def test_build_enforcement_payload(self, enforcement_service, enforcement_context):
+        """Test enforcement payload construction."""
+        payload = enforcement_service._build_enforcement_payload(
+            tool_name="Bash",
+            tool_args={"command": "kubectl get pods"},
+            context=enforcement_context
+        )
+
+        assert payload["action"] == "tool_execution"
+        assert payload["tool"]["name"] == "Bash"
+        assert payload["tool"]["source"] == "builtin"
+        assert payload["tool"]["category"] == "command_execution"
+        assert payload["tool"]["risk_level"] == "high"
+        assert payload["user"]["email"] == "test@example.com"
+        assert payload["organization"]["id"] == "org-456"
+        assert payload["execution"]["environment"] == "production"
+        assert "timestamp" in payload["execution"]
+
+    def test_determine_tool_source(self, enforcement_service):
+        """Test tool source detection."""
+        assert enforcement_service._determine_tool_source("mcp__github__list_repos") == "mcp"
+        assert enforcement_service._determine_tool_source("Bash") == "builtin"
+        assert enforcement_service._determine_tool_source("Read") == "builtin"
+        assert enforcement_service._determine_tool_source("custom_tool") == "skill"
+
+    def test_determine_tool_category(self, enforcement_service):
+        """Test tool category detection."""
+        assert enforcement_service._determine_tool_category("Bash") == "command_execution"
+        assert enforcement_service._determine_tool_category("Read") == "file_operation"
+        assert enforcement_service._determine_tool_category("Write") == "file_operation"
+        assert enforcement_service._determine_tool_category("Grep") == "file_search"
+        assert enforcement_service._determine_tool_category("WebFetch") == "network"
+        assert enforcement_service._determine_tool_category("custom_tool") == "general"
+
+    def test_determine_risk_level_critical(self, enforcement_service):
+        """Test risk level assessment for critical commands."""
+        # Critical risk - destructive commands
+        assert enforcement_service._determine_risk_level(
+            "Bash",
+            {"command": "rm -rf /"}
+        ) == "critical"
+
+        assert enforcement_service._determine_risk_level(
+            "Bash",
+            {"command": "dd if=/dev/zero of=/dev/sda"}
+        ) == "critical"
+
+    def test_determine_risk_level_high(self, enforcement_service):
+        """Test risk level assessment for high-risk operations."""
+        # High risk - command execution
+        assert enforcement_service._determine_risk_level(
+            "Bash",
+            {"command": "kubectl delete deployment"}
+        ) == "high"
+
+        # High risk - sensitive file access
+        assert enforcement_service._determine_risk_level(
+            "Read",
+            {"file_path": "/etc/passwd"}
+        ) == "high"
+
+        assert enforcement_service._determine_risk_level(
+            "Read",
+            {"file_path": "~/.ssh/id_rsa"}
+        ) == "high"
+
+    def test_determine_risk_level_medium(self, enforcement_service):
+        """Test risk level assessment for medium-risk operations."""
+        assert enforcement_service._determine_risk_level(
+            "Write",
+            {"file_path": "/tmp/test.txt", "content": "test"}
+        ) == "medium"
+
+        assert enforcement_service._determine_risk_level(
+            "Edit",
+            {"file_path": "/tmp/config.yaml"}
+        ) == "medium"
+
+    def test_determine_risk_level_low(self, enforcement_service):
+        """Test risk level assessment for low-risk operations."""
+        assert enforcement_service._determine_risk_level(
+            "Read",
+            {"file_path": "/tmp/test.txt"}
+        ) == "low"
+
+        assert enforcement_service._determine_risk_level(
+            "Grep",
+            {"pattern": "error", "path": "/var/log"}
+        ) == "low"
+
+    def test_format_violation_message(self, enforcement_service):
+        """Test violation message formatting."""
+        enforcement_result = {
+            "id": "enf-123",
+            "allow": False,
+            "policies": ["policy1", "policy2"]
+        }
+
+        message = enforcement_service._format_violation_message(
+            tool_name="Bash",
+            policies=["policy1", "policy2"],
+            enforcement_result=enforcement_result
+        )
+
+        assert "Tool execution blocked" in message
+        assert "Bash" in message
+        assert "policy1, policy2" in message
+        assert "enf-123" in message
+        assert "administrator" in message.lower()
+
+
+class TestToolEnforcementPayloadValidation:
+    """Test payload structure validation."""
+
+    def test_payload_has_all_required_fields(self, enforcement_service, enforcement_context):
+        """Verify payload contains all required fields."""
+        payload = enforcement_service._build_enforcement_payload(
+            tool_name="Bash",
+            tool_args={"command": "ls"},
+            context=enforcement_context
+        )
+
+        # Top-level fields
+        assert "action" in payload
+        assert "tool" in payload
+        assert "user" in payload
+        assert "organization" in payload
+        assert "team" in payload
+        assert "execution" in payload
+
+        # Tool fields
+        assert "name" in payload["tool"]
+        assert "arguments" in payload["tool"]
+        assert "source" in payload["tool"]
+        assert "category" in payload["tool"]
+        assert "risk_level" in payload["tool"]
+
+        # User fields
+        assert "email" in payload["user"]
+        assert "id" in payload["user"]
+        assert "roles" in payload["user"]
+
+        # Organization fields
+        assert "id" in payload["organization"]
+
+        # Execution fields
+        assert "execution_id" in payload["execution"]
+        assert "agent_id" in payload["execution"]
+        assert "environment" in payload["execution"]
+        assert "timestamp" in payload["execution"]
+
+    def test_payload_timestamp_format(self, enforcement_service, enforcement_context):
+        """Verify timestamp is in ISO format."""
+        payload = enforcement_service._build_enforcement_payload(
+            tool_name="Test",
+            tool_args={},
+            context=enforcement_context
+        )
+
+        timestamp = payload["execution"]["timestamp"]
+        # Should be able to parse back
+        parsed = datetime.fromisoformat(timestamp.replace('Z', '+00:00'))
+        assert isinstance(parsed, datetime)
+
+    def test_payload_handles_missing_context_fields(self, enforcement_service):
+        """Test payload construction with minimal context."""
+        minimal_context = {
+            "organization_id": "org-123",
+            "agent_id": "agent-456"
+        }
+
+        payload = enforcement_service._build_enforcement_payload(
+            tool_name="Test",
+            tool_args={},
+            context=minimal_context
+        )
+
+        # Should not crash, just have None values
+        assert payload["user"]["email"] is None
+        assert payload["user"]["roles"] == []
+        assert payload["organization"]["id"] == "org-123"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

control_plane_api/worker/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility functions and helpers"""

control_plane_api/worker/utils/chunk_batcher.py

@@ -0,0 +1,330 @@
+"""
+Smart chunk batching for streaming to reduce HTTP requests.
+
+Instead of sending one POST per chunk (50-70 requests), batch chunks
+with configurable time/size windows (5-10 requests).
+
+Batching Strategy:
+- Time window: Flush after X ms (default: 100ms)
+- Size window: Flush when batch reaches Y bytes (default: 100 bytes)
+- Immediate flush: On tool events, errors, or completion
+
+This provides:
+- 90%+ reduction in HTTP requests
+- Still feels real-time (100ms is imperceptible)
+- Lower latency (fewer round trips)
+- Better serverless performance (fewer cold starts)
+- Lower costs (fewer invocations)
+"""
+
+import asyncio
+import time
+from typing import Dict, Any, Optional, Callable
+from dataclasses import dataclass, field
+import structlog
+
+logger = structlog.get_logger()
+
+
+@dataclass
+class BatchConfig:
+    """Configuration for chunk batching behavior."""
+
+    # Time-based batching: flush after this many milliseconds
+    time_window_ms: int = 100
+
+    # Size-based batching: flush when accumulated content reaches this size
+    size_window_bytes: int = 100
+
+    # Maximum batch size before forced flush (safety limit)
+    max_batch_size_bytes: int = 1000
+
+    # Enable/disable batching (for testing/debugging)
+    enabled: bool = True
+
+    @classmethod
+    def from_env(cls) -> "BatchConfig":
+        """
+        Create configuration from environment variables.
+
+        Environment variables:
+            CHUNK_BATCHING_ENABLED: Enable/disable batching (default: true)
+            CHUNK_BATCHING_TIME_WINDOW_MS: Time window in ms (default: 100)
+            CHUNK_BATCHING_SIZE_WINDOW_BYTES: Size window in bytes (default: 100)
+            CHUNK_BATCHING_MAX_SIZE_BYTES: Max batch size in bytes (default: 1000)
+
+        Returns:
+            BatchConfig instance with values from environment
+        """
+        import os
+
+        return cls(
+            enabled=os.getenv("CHUNK_BATCHING_ENABLED", "true").lower() == "true",
+            time_window_ms=int(os.getenv("CHUNK_BATCHING_TIME_WINDOW_MS", "100")),
+            size_window_bytes=int(os.getenv("CHUNK_BATCHING_SIZE_WINDOW_BYTES", "100")),
+            max_batch_size_bytes=int(os.getenv("CHUNK_BATCHING_MAX_SIZE_BYTES", "1000")),
+        )
+
+
+@dataclass
+class ContentBatch:
+    """Accumulated content chunks waiting to be flushed."""
+
+    chunks: list[str] = field(default_factory=list)
+    total_size: int = 0
+    first_chunk_time: Optional[float] = None
+
+    def add(self, content: str) -> None:
+        """Add content to the batch."""
+        self.chunks.append(content)
+        self.total_size += len(content.encode('utf-8'))
+
+        if self.first_chunk_time is None:
+            self.first_chunk_time = time.time()
+
+    def get_combined_content(self) -> str:
+        """Get all chunks combined into single string."""
+        return ''.join(self.chunks)
+
+    def clear(self) -> None:
+        """Clear the batch after flushing."""
+        self.chunks.clear()
+        self.total_size = 0
+        self.first_chunk_time = None
+
+    def is_empty(self) -> bool:
+        """Check if batch is empty."""
+        return len(self.chunks) == 0
+
+    def age_ms(self) -> float:
+        """Get age of batch in milliseconds."""
+        if self.first_chunk_time is None:
+            return 0
+        return (time.time() - self.first_chunk_time) * 1000
+
+
+class ChunkBatcher:
+    """
+    Smart batching for streaming chunks to reduce HTTP requests.
+
+    Usage:
+        batcher = ChunkBatcher(
+            publish_func=control_plane.publish_event,
+            execution_id=execution_id,
+            message_id=message_id,
+            config=BatchConfig(time_window_ms=100, size_window_bytes=100)
+        )
+
+        # Add chunks as they arrive
+        await batcher.add_chunk("Hello")
+        await batcher.add_chunk(" world")
+
+        # Flush remaining chunks when done
+        await batcher.flush()
+    """
+
+    def __init__(
+        self,
+        publish_func: Callable,
+        execution_id: str,
+        message_id: str,
+        config: Optional[BatchConfig] = None
+    ):
+        self.publish_func = publish_func
+        self.execution_id = execution_id
+        self.message_id = message_id
+        self.config = config or BatchConfig()
+
+        self.batch = ContentBatch()
+        self._flush_task: Optional[asyncio.Task] = None
+        self._stats = {
+            "chunks_received": 0,
+            "batches_sent": 0,
+            "bytes_sent": 0,
+            "flushes_by_time": 0,
+            "flushes_by_size": 0,
+            "flushes_manual": 0,
+        }
+
+    async def add_chunk(self, content: str) -> None:
+        """
+        Add a chunk to the batch.
+
+        Automatically flushes if:
+        - Batch size exceeds size_window_bytes
+        - Batch age exceeds time_window_ms
+        - Max batch size is reached (safety)
+        """
+        if not self.config.enabled:
+            # Batching disabled - send immediately
+            await self._publish_batch([content])
+            return
+
+        self._stats["chunks_received"] += 1
+        self.batch.add(content)
+
+        # Check if we should flush immediately due to size
+        should_flush_size = self.batch.total_size >= self.config.size_window_bytes
+        should_flush_max = self.batch.total_size >= self.config.max_batch_size_bytes
+
+        if should_flush_max:
+            # Safety: flush immediately if max size reached
+            logger.debug(
+                "Flushing batch (max size reached)",
+                execution_id=self.execution_id[:8],
+                batch_size=self.batch.total_size,
+                chunk_count=len(self.batch.chunks),
+            )
+            await self.flush(reason="max_size")
+        elif should_flush_size:
+            # Size threshold reached - flush now
+            await self.flush(reason="size")
+        else:
+            # Start/reset timer for time-based flush
+            await self._schedule_time_flush()
+
+    async def _schedule_time_flush(self) -> None:
+        """Schedule a time-based flush if not already scheduled."""
+        if self._flush_task is not None and not self._flush_task.done():
+            # Timer already running
+            return
+
+        self._flush_task = asyncio.create_task(self._time_based_flush())
+
+    async def _time_based_flush(self) -> None:
+        """Wait for time window, then flush."""
+        await asyncio.sleep(self.config.time_window_ms / 1000.0)
+
+        if not self.batch.is_empty():
+            await self.flush(reason="time")
+
+    async def flush(self, reason: str = "manual") -> None:
+        """
+        Flush current batch immediately.
+
+        Args:
+            reason: Why flush was triggered (for stats/debugging)
+        """
+        if self.batch.is_empty():
+            return
+
+        # Cancel pending timer if any
+        if self._flush_task is not None and not self._flush_task.done():
+            self._flush_task.cancel()
+            try:
+                await self._flush_task
+            except asyncio.CancelledError:
+                pass
+
+        # Publish the batch
+        chunks = self.batch.chunks.copy()
+        await self._publish_batch(chunks)
+
+        # Update stats
+        if reason == "time":
+            self._stats["flushes_by_time"] += 1
+        elif reason == "size" or reason == "max_size":
+            self._stats["flushes_by_size"] += 1
+        else:
+            self._stats["flushes_manual"] += 1
+
+        # Clear batch
+        self.batch.clear()
+
+    async def _publish_batch(self, chunks: list[str]) -> None:
+        """Publish a batch of chunks as single event with retry logic."""
+        combined_content = ''.join(chunks)
+        max_retries = 3
+        base_delay = 0.1  # 100ms
+
+        for attempt in range(max_retries):
+            try:
+                # CRITICAL: Always await async functions immediately
+                # publish_func is an async function, so call it with await
+                await self.publish_func(
+                    execution_id=self.execution_id,
+                    event_type="message_chunk",
+                    data={
+                        "role": "assistant",
+                        "content": combined_content,
+                        "is_chunk": True,
+                        "message_id": self.message_id,
+                        # Metadata for debugging
+                        "batch_info": {
+                            "chunk_count": len(chunks),
+                            "batch_size": len(combined_content.encode('utf-8')),
+                        } if len(chunks) > 1 else None,
+                    }
+                )
+
+                self._stats["batches_sent"] += 1
+                self._stats["bytes_sent"] += len(combined_content.encode('utf-8'))
+
+                # Success - exit retry loop
+                if attempt > 0:
+                    logger.debug(
+                        "Batch published after retry",
+                        execution_id=self.execution_id[:8],
+                        attempt=attempt + 1,
+                        chunk_count=len(chunks),
+                    )
+                return
+
+            except Exception as e:
+                is_last_attempt = attempt == max_retries - 1
+
+                if is_last_attempt:
+                    logger.error(
+                        "Failed to publish batch after all retries",
+                        execution_id=self.execution_id[:8],
+                        error=str(e),
+                        chunk_count=len(chunks),
+                        attempts=max_retries,
+                    )
+                else:
+                    # Exponential backoff
+                    delay = base_delay * (2 ** attempt)
+                    logger.debug(
+                        "Retrying batch publish",
+                        execution_id=self.execution_id[:8],
+                        error=str(e),
+                        attempt=attempt + 1,
+                        next_delay_ms=int(delay * 1000),
+                    )
+                    await asyncio.sleep(delay)
+
+    def get_stats(self) -> Dict[str, Any]:
+        """
+        Get batching statistics.
+
+        Returns:
+            Dict with stats about batching performance
+        """
+        chunks_received = self._stats["chunks_received"]
+        batches_sent = self._stats["batches_sent"]
+
+        return {
+            **self._stats,
+            "reduction_percent": round(
+                (1 - batches_sent / max(chunks_received, 1)) * 100, 1
+            ) if chunks_received > 0 else 0,
+            "avg_batch_size": round(
+                chunks_received / max(batches_sent, 1), 1
+            ) if batches_sent > 0 else 0,
+        }
+
+    async def close(self) -> None:
+        """
+        Close the batcher and flush remaining chunks.
+
+        Call this when streaming is complete.
+        """
+        await self.flush(reason="close")
+
+        # Log stats
+        stats = self.get_stats()
+        logger.info(
+            "Chunk batching stats",
+            execution_id=self.execution_id[:8],
+            **stats
+        )