agent-os-kernel 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

modules/control-plane/tests/test_adapter.py

@@ -0,0 +1,422 @@
+"""
+Tests for the OpenAI Adapter - Drop-in Middleware
+"""
+
+import json
+import unittest
+from agent_control_plane import (
+    AgentControlPlane,
+    ControlPlaneAdapter,
+    create_governed_client,
+    ActionType,
+    PermissionLevel,
+)
+
+
+class MockOpenAIClient:
+    """
+    Mock OpenAI client for testing the adapter.
+    
+    This class mimics the structure of the OpenAI SDK client,
+    providing the nested classes (MockFunction, MockToolCall, etc.)
+    that represent the OpenAI API response structure.
+    
+    It allows testing the adapter without making actual API calls.
+    """
+    
+    class MockFunction:
+        """Represents a function call in OpenAI's tool call structure"""
+        def __init__(self, name, arguments):
+            self.name = name
+            self.arguments = arguments
+    
+    class MockToolCall:
+        """Represents a tool call in OpenAI's response"""
+        def __init__(self, name, arguments):
+            self.id = f"call_{name}"
+            self.type = "function"
+            self.function = MockOpenAIClient.MockFunction(name, arguments)
+    
+    class MockMessage:
+        """Represents the message object in OpenAI's response"""
+        def __init__(self, tool_calls=None):
+            self.role = "assistant"
+            self.content = None
+            self.tool_calls = tool_calls or []
+    
+    class MockChoice:
+        """Represents a choice in OpenAI's response"""
+        def __init__(self, message):
+            self.index = 0
+            self.message = message
+            self.finish_reason = "tool_calls"
+    
+    class MockResponse:
+        """Represents a complete OpenAI API response"""
+        def __init__(self, choices):
+            self.id = "chatcmpl-mock"
+            self.object = "chat.completion"
+            self.model = "gpt-4"
+            self.choices = choices
+    
+    class MockCompletions:
+        """Represents the chat.completions API endpoint"""
+        def __init__(self, tool_calls):
+            self.tool_calls_to_return = tool_calls
+        
+        def create(self, **kwargs):
+            """Mock the create() method that returns tool calls"""
+            message = MockOpenAIClient.MockMessage(self.tool_calls_to_return)
+            choice = MockOpenAIClient.MockChoice(message)
+            return MockOpenAIClient.MockResponse([choice])
+    
+    class MockChat:
+        """Represents the chat API endpoint"""
+        def __init__(self, tool_calls):
+            self.completions = MockOpenAIClient.MockCompletions(tool_calls)
+    
+    def __init__(self, tool_calls=None):
+        self.chat = MockOpenAIClient.MockChat(tool_calls or [])
+
+
+class TestOpenAIAdapter(unittest.TestCase):
+    """Test suite for OpenAI Adapter"""
+
+    def test_adapter_blocks_unauthorized_action(self):
+        """Test that adapter blocks actions the agent doesn't have permission for"""
+        control_plane = AgentControlPlane()
+        
+        # Agent can only read files, not write
+        permissions = {
+            ActionType.FILE_READ: PermissionLevel.READ_ONLY,
+            ActionType.FILE_WRITE: PermissionLevel.NONE,
+        }
+        agent_context = control_plane.create_agent("test-agent", permissions)
+        
+        # Mock OpenAI wants to write a file
+        tool_calls = [
+            MockOpenAIClient.MockToolCall(
+                "write_file",
+                json.dumps({"path": "/test.txt", "content": "data"})
+            )
+        ]
+        mock_client = MockOpenAIClient(tool_calls)
+        
+        # Wrap with adapter
+        governed = ControlPlaneAdapter(
+            control_plane=control_plane,
+            agent_context=agent_context,
+            original_client=mock_client
+        )
+        
+        # Make API call
+        response = governed.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Write a file"}]
+        )
+        
+        # Verify the tool call was blocked
+        self.assertEqual(len(response.choices[0].message.tool_calls), 1)
+        blocked_call = response.choices[0].message.tool_calls[0]
+        self.assertEqual(blocked_call.function.name, "blocked_action")
+        
+        args = json.loads(blocked_call.function.arguments)
+        self.assertEqual(args["original_tool"], "write_file")
+        self.assertIn("blocked", args["reason"].lower())
+
+    def test_adapter_allows_authorized_action(self):
+        """Test that adapter allows actions the agent has permission for"""
+        control_plane = AgentControlPlane()
+        
+        # Agent can read files
+        permissions = {
+            ActionType.FILE_READ: PermissionLevel.READ_ONLY,
+        }
+        agent_context = control_plane.create_agent("test-agent", permissions)
+        
+        # Mock OpenAI wants to read a file
+        tool_calls = [
+            MockOpenAIClient.MockToolCall(
+                "read_file",
+                json.dumps({"path": "/test.txt"})
+            )
+        ]
+        mock_client = MockOpenAIClient(tool_calls)
+        
+        # Wrap with adapter
+        governed = ControlPlaneAdapter(
+            control_plane=control_plane,
+            agent_context=agent_context,
+            original_client=mock_client
+        )
+        
+        # Make API call
+        response = governed.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Read a file"}]
+        )
+        
+        # Verify the tool call was allowed (not modified to blocked_action)
+        self.assertEqual(len(response.choices[0].message.tool_calls), 1)
+        tool_call = response.choices[0].message.tool_calls[0]
+        self.assertEqual(tool_call.function.name, "read_file")
+
+    def test_adapter_mixed_permissions(self):
+        """Test adapter with multiple tool calls, some allowed and some blocked"""
+        control_plane = AgentControlPlane()
+        
+        # Agent can query database but not write
+        permissions = {
+            ActionType.DATABASE_QUERY: PermissionLevel.READ_ONLY,
+            ActionType.DATABASE_WRITE: PermissionLevel.NONE,
+        }
+        agent_context = control_plane.create_agent("test-agent", permissions)
+        
+        # Mock OpenAI wants to both query and write
+        tool_calls = [
+            MockOpenAIClient.MockToolCall(
+                "database_query",
+                json.dumps({"query": "SELECT * FROM users"})
+            ),
+            MockOpenAIClient.MockToolCall(
+                "database_write",
+                json.dumps({"query": "DELETE FROM users"})
+            ),
+        ]
+        mock_client = MockOpenAIClient(tool_calls)
+        
+        # Wrap with adapter
+        governed = ControlPlaneAdapter(
+            control_plane=control_plane,
+            agent_context=agent_context,
+            original_client=mock_client
+        )
+        
+        # Make API call
+        response = governed.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Query and delete"}]
+        )
+        
+        # Verify mixed results
+        tool_calls_result = response.choices[0].message.tool_calls
+        self.assertEqual(len(tool_calls_result), 2)
+        
+        # First should be allowed
+        self.assertEqual(tool_calls_result[0].function.name, "database_query")
+        
+        # Second should be blocked
+        self.assertEqual(tool_calls_result[1].function.name, "blocked_action")
+
+    def test_custom_tool_mapping(self):
+        """Test that custom tool name mappings work"""
+        control_plane = AgentControlPlane()
+        
+        permissions = {
+            ActionType.DATABASE_QUERY: PermissionLevel.READ_ONLY,
+        }
+        agent_context = control_plane.create_agent("test-agent", permissions)
+        
+        # Custom mapping
+        custom_mapping = {
+            "my_custom_db_tool": ActionType.DATABASE_QUERY,
+        }
+        
+        tool_calls = [
+            MockOpenAIClient.MockToolCall(
+                "my_custom_db_tool",
+                json.dumps({"query": "SELECT 1"})
+            )
+        ]
+        mock_client = MockOpenAIClient(tool_calls)
+        
+        governed = ControlPlaneAdapter(
+            control_plane=control_plane,
+            agent_context=agent_context,
+            original_client=mock_client,
+            tool_mapping=custom_mapping
+        )
+        
+        response = governed.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Query"}]
+        )
+        
+        # Should be recognized and allowed
+        tool_call = response.choices[0].message.tool_calls[0]
+        self.assertEqual(tool_call.function.name, "my_custom_db_tool")
+
+    def test_on_block_callback(self):
+        """Test that on_block callback is called when actions are blocked"""
+        control_plane = AgentControlPlane()
+        
+        permissions = {
+            ActionType.FILE_READ: PermissionLevel.READ_ONLY,
+        }
+        agent_context = control_plane.create_agent("test-agent", permissions)
+        
+        blocked_calls = []
+        
+        def on_block(tool_name, tool_args, result):
+            blocked_calls.append({
+                "tool": tool_name,
+                "args": tool_args,
+                "result": result
+            })
+        
+        tool_calls = [
+            MockOpenAIClient.MockToolCall(
+                "write_file",
+                json.dumps({"path": "/test.txt"})
+            )
+        ]
+        mock_client = MockOpenAIClient(tool_calls)
+        
+        governed = ControlPlaneAdapter(
+            control_plane=control_plane,
+            agent_context=agent_context,
+            original_client=mock_client,
+            on_block=on_block
+        )
+        
+        governed.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Write"}]
+        )
+        
+        # Verify callback was called
+        self.assertEqual(len(blocked_calls), 1)
+        self.assertEqual(blocked_calls[0]["tool"], "write_file")
+        self.assertFalse(blocked_calls[0]["result"]["success"])
+
+    def test_create_governed_client_convenience(self):
+        """Test the convenience function for creating governed clients"""
+        control_plane = AgentControlPlane()
+        
+        tool_calls = [
+            MockOpenAIClient.MockToolCall(
+                "database_query",
+                json.dumps({"query": "SELECT 1"})
+            )
+        ]
+        mock_client = MockOpenAIClient(tool_calls)
+        
+        # Use convenience function
+        governed = create_governed_client(
+            control_plane=control_plane,
+            agent_id="convenience-test",
+            openai_client=mock_client,
+            permissions={
+                ActionType.DATABASE_QUERY: PermissionLevel.READ_ONLY,
+            }
+        )
+        
+        # Should work just like manual setup
+        response = governed.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Query"}]
+        )
+        
+        tool_call = response.choices[0].message.tool_calls[0]
+        self.assertEqual(tool_call.function.name, "database_query")
+
+    def test_adapter_statistics(self):
+        """Test that statistics are properly collected"""
+        control_plane = AgentControlPlane()
+        
+        tool_calls = [
+            MockOpenAIClient.MockToolCall(
+                "read_file",
+                json.dumps({"path": "/test.txt"})
+            )
+        ]
+        mock_client = MockOpenAIClient(tool_calls)
+        
+        governed = create_governed_client(
+            control_plane=control_plane,
+            agent_id="stats-test",
+            openai_client=mock_client,
+            permissions={
+                ActionType.FILE_READ: PermissionLevel.READ_ONLY,
+            }
+        )
+        
+        # Make some calls
+        for _ in range(3):
+            governed.chat.completions.create(
+                model="gpt-4",
+                messages=[{"role": "user", "content": "Read"}]
+            )
+        
+        # Get statistics
+        stats = governed.get_statistics()
+        
+        self.assertIn("agent_id", stats)
+        self.assertEqual(stats["agent_id"], "stats-test")
+        self.assertIn("session_id", stats)
+        self.assertIn("control_plane_audit", stats)
+        self.assertGreater(len(stats["control_plane_audit"]), 0)
+
+    def test_no_tool_calls_response(self):
+        """Test that responses without tool calls pass through unchanged"""
+        control_plane = AgentControlPlane()
+        agent_context = control_plane.create_agent("test-agent", {})
+        
+        # Mock client with no tool calls
+        mock_client = MockOpenAIClient(tool_calls=[])
+        
+        governed = ControlPlaneAdapter(
+            control_plane=control_plane,
+            agent_context=agent_context,
+            original_client=mock_client
+        )
+        
+        response = governed.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}]
+        )
+        
+        # Should pass through without errors
+        self.assertIsNotNone(response)
+        self.assertEqual(len(response.choices[0].message.tool_calls), 0)
+
+    def test_pattern_matching_tool_names(self):
+        """Test that pattern matching recognizes common tool name variations"""
+        control_plane = AgentControlPlane()
+        
+        permissions = {
+            ActionType.FILE_READ: PermissionLevel.READ_ONLY,
+        }
+        agent_context = control_plane.create_agent("test-agent", permissions)
+        
+        # Try various file read naming patterns
+        tool_names = ["read_file", "get_file", "fetch_document", "load_file"]
+        
+        for tool_name in tool_names:
+            tool_calls = [
+                MockOpenAIClient.MockToolCall(
+                    tool_name,
+                    json.dumps({"path": "/test.txt"})
+                )
+            ]
+            mock_client = MockOpenAIClient(tool_calls)
+            
+            governed = ControlPlaneAdapter(
+                control_plane=control_plane,
+                agent_context=agent_context,
+                original_client=mock_client
+            )
+            
+            response = governed.chat.completions.create(
+                model="gpt-4",
+                messages=[{"role": "user", "content": "Read"}]
+            )
+            
+            # All should be recognized and allowed
+            tool_call = response.choices[0].message.tool_calls[0]
+            self.assertEqual(tool_call.function.name, tool_name, 
+                           f"Pattern matching failed for {tool_name}")
+
+
+if __name__ == "__main__":
+    unittest.main()