agent_os_kernel 3.1.0__py3-none-any.whl

iatp/tests/test_telemetry.py

@@ -0,0 +1,175 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Unit tests for IATP telemetry module.
+"""
+import shutil
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from iatp.models import (
+    AgentCapabilities,
+    CapabilityManifest,
+    PrivacyContract,
+    RetentionPolicy,
+)
+from iatp.telemetry import FlightRecorder, TraceIDGenerator
+
+
+@pytest.fixture
+def temp_log_dir():
+    """Create a temporary directory for logs."""
+    temp_dir = Path(tempfile.mkdtemp())
+    yield temp_dir
+    shutil.rmtree(temp_dir)
+
+
+def test_trace_id_generation():
+    """Test trace ID generation."""
+    trace_id1 = TraceIDGenerator.generate()
+    trace_id2 = TraceIDGenerator.generate()
+
+    assert trace_id1 != trace_id2
+    assert len(trace_id1) > 0
+    assert len(trace_id2) > 0
+
+
+def test_create_tracing_context():
+    """Test creating a tracing context."""
+    trace_id = TraceIDGenerator.generate()
+    context = TraceIDGenerator.create_context(trace_id, "test-agent")
+
+    assert context.trace_id == trace_id
+    assert context.agent_id == "test-agent"
+    assert context.parent_trace_id is None
+    assert context.timestamp is not None
+
+
+def test_flight_recorder_log_request(temp_log_dir):
+    """Test logging a request."""
+    recorder = FlightRecorder(log_dir=temp_log_dir)
+    trace_id = "test-trace-123"
+
+    recorder.log_request(
+        trace_id=trace_id,
+        agent_id="test-agent",
+        payload={"task": "test"},
+        quarantined=False
+    )
+
+    logs = recorder.get_trace_logs(trace_id)
+    assert len(logs) == 1
+    assert logs[0]["type"] == "request"
+    assert logs[0]["trace_id"] == trace_id
+    assert logs[0]["agent_id"] == "test-agent"
+
+
+def test_flight_recorder_log_response(temp_log_dir):
+    """Test logging a response."""
+    recorder = FlightRecorder(log_dir=temp_log_dir)
+    trace_id = "test-trace-456"
+
+    recorder.log_response(
+        trace_id=trace_id,
+        agent_id="test-agent",
+        response={"status": "success"},
+        status_code=200,
+        latency_ms=123.45
+    )
+
+    logs = recorder.get_trace_logs(trace_id)
+    assert len(logs) == 1
+    assert logs[0]["type"] == "response"
+    assert logs[0]["status_code"] == 200
+    assert logs[0]["latency_ms"] == 123.45
+
+
+def test_flight_recorder_log_error(temp_log_dir):
+    """Test logging an error."""
+    recorder = FlightRecorder(log_dir=temp_log_dir)
+    trace_id = "test-trace-789"
+
+    recorder.log_error(
+        trace_id=trace_id,
+        agent_id="test-agent",
+        error="Connection timeout",
+        details={"timeout": 30}
+    )
+
+    logs = recorder.get_trace_logs(trace_id)
+    assert len(logs) == 1
+    assert logs[0]["type"] == "error"
+    assert logs[0]["error"] == "Connection timeout"
+    assert logs[0]["details"]["timeout"] == 30
+
+
+def test_flight_recorder_log_blocked(temp_log_dir):
+    """Test logging a blocked request."""
+    recorder = FlightRecorder(log_dir=temp_log_dir)
+    trace_id = "test-trace-blocked"
+
+    manifest = CapabilityManifest(
+        agent_id="test-agent",
+        capabilities=AgentCapabilities(),
+        privacy_contract=PrivacyContract(
+            retention=RetentionPolicy.FOREVER
+        )
+    )
+
+    recorder.log_blocked_request(
+        trace_id=trace_id,
+        agent_id="test-agent",
+        payload={"sensitive": "data"},
+        reason="Privacy violation",
+        manifest=manifest
+    )
+
+    logs = recorder.get_trace_logs(trace_id)
+    assert len(logs) == 1
+    assert logs[0]["type"] == "blocked"
+    assert logs[0]["reason"] == "Privacy violation"
+
+
+def test_flight_recorder_multiple_logs(temp_log_dir):
+    """Test multiple log entries for the same trace."""
+    recorder = FlightRecorder(log_dir=temp_log_dir)
+    trace_id = "test-trace-multi"
+
+    recorder.log_request(trace_id, "agent1", {"req": 1}, quarantined=False)
+    recorder.log_response(trace_id, "agent1", {"resp": 1}, 200, 100.0)
+    recorder.log_request(trace_id, "agent2", {"req": 2}, quarantined=False)
+    recorder.log_response(trace_id, "agent2", {"resp": 2}, 200, 150.0)
+
+    logs = recorder.get_trace_logs(trace_id)
+    assert len(logs) == 4
+    assert logs[0]["type"] == "request"
+    assert logs[1]["type"] == "response"
+    assert logs[2]["type"] == "request"
+    assert logs[3]["type"] == "response"
+
+
+def test_flight_recorder_sensitive_data_scrubbed(temp_log_dir):
+    """Test that sensitive data is scrubbed in logs."""
+    recorder = FlightRecorder(log_dir=temp_log_dir)
+    trace_id = "test-trace-sensitive"
+
+    recorder.log_request(
+        trace_id=trace_id,
+        agent_id="test-agent",
+        payload={"card": "4532-1234-5678-9010"},
+        quarantined=False
+    )
+
+    logs = recorder.get_trace_logs(trace_id)
+    payload_str = str(logs[0]["payload"])
+    assert "4532-1234-5678-9010" not in payload_str
+    assert "[CREDIT_CARD_REDACTED]" in payload_str
+
+
+def test_flight_recorder_nonexistent_trace(temp_log_dir):
+    """Test retrieving logs for a nonexistent trace."""
+    recorder = FlightRecorder(log_dir=temp_log_dir)
+    logs = recorder.get_trace_logs("nonexistent-trace")
+    assert logs == []

mcp_kernel_server/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+MCP Kernel Server - Expose Agent OS primitives through Model Context Protocol.
+
+This package provides:
+- CMVK verification as MCP tool
+- IATP trust signing as MCP tool  
+- VFS filesystem as MCP resource
+- Governed execution as MCP tool
+"""
+
+from mcp_kernel_server.server import KernelMCPServer
+from mcp_kernel_server.tools import (
+    CMVKVerifyTool,
+    KernelExecuteTool,
+    IATPSignTool,
+)
+from mcp_kernel_server.resources import VFSResource
+
+__version__ = "3.1.0"
+__all__ = [
+    "KernelMCPServer",
+    "CMVKVerifyTool",
+    "KernelExecuteTool", 
+    "IATPSignTool",
+    "VFSResource",
+]

mcp_kernel_server/cli.py

@@ -0,0 +1,274 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+CLI for MCP Kernel Server.
+
+Usage:
+    # For Claude Desktop (stdio transport)
+    mcp-kernel-server --stdio
+    
+    # For development (HTTP transport)
+    mcp-kernel-server --http --port 8080
+
+Claude Desktop Integration:
+    Add to ~/Library/Application Support/Claude/claude_desktop_config.json:
+    
+    {
+      "mcpServers": {
+        "agent-os": {
+          "command": "mcp-kernel-server",
+          "args": ["--stdio"]
+        }
+      }
+    }
+"""
+
+import argparse
+import asyncio
+import logging
+import sys
+from pathlib import Path
+
+from mcp_kernel_server.server import KernelMCPServer, ServerConfig
+
+# Configure logging to stderr (stdout is for MCP protocol)
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    stream=sys.stderr
+)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="MCP Kernel Server - Agent OS primitives via Model Context Protocol",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # For Claude Desktop integration
+    mcp-kernel-server --stdio
+    
+    # For development/testing
+    mcp-kernel-server --http --port 8080
+    
+    # With custom policy mode
+    mcp-kernel-server --stdio --policy-mode permissive
+
+Claude Desktop Setup:
+    Add to claude_desktop_config.json:
+    {
+      "mcpServers": {
+        "agent-os": {
+          "command": "mcp-kernel-server",
+          "args": ["--stdio"]
+        }
+      }
+    }
+"""
+    )
+    parser.add_argument(
+        "--stdio",
+        action="store_true",
+        help="Use stdio transport (for Claude Desktop)"
+    )
+    parser.add_argument(
+        "--http",
+        action="store_true",
+        help="Use HTTP transport (for development)"
+    )
+    parser.add_argument(
+        "--port", "-p",
+        type=int,
+        default=8080,
+        help="Port for HTTP transport (default: 8080)"
+    )
+    parser.add_argument(
+        "--host", "-H",
+        type=str,
+        default="127.0.0.1",
+        help="Host for HTTP transport (default: 127.0.0.1)"
+    )
+    parser.add_argument(
+        "--policy-mode",
+        choices=["strict", "permissive", "audit"],
+        default="strict",
+        help="Policy enforcement mode (default: strict)"
+    )
+    parser.add_argument(
+        "--cmvk-threshold",
+        type=float,
+        default=0.85,
+        help="CMVK verification threshold (default: 0.85)"
+    )
+    parser.add_argument(
+        "--version", "-v",
+        action="store_true",
+        help="Show version and exit"
+    )
+    parser.add_argument(
+        "--list-tools",
+        action="store_true",
+        help="List available tools and exit"
+    )
+    parser.add_argument(
+        "--list-prompts",
+        action="store_true",
+        help="List available prompts and exit"
+    )
+    parser.add_argument(
+        "--json",
+        action="store_true",
+        help="Output in JSON format (for listing tools/prompts)"
+    )
+    
+    return parser.parse_args()
+
+
+def print_tools():
+    """Print available tools."""
+    print("Available MCP Tools:")
+    print()
+    print("  cmvk_verify")
+    print("    Verify claims across multiple AI models to detect hallucinations")
+    print("    Arguments: claim (required), context, models, threshold")
+    print()
+    print("  kernel_execute")
+    print("    Execute actions through the Agent OS kernel with policy enforcement")
+    print("    Arguments: action (required), params, agent_id (required), policies, context")
+    print()
+    print("  iatp_sign")
+    print("    Create a trust attestation for inter-agent communication")
+    print("    Arguments: attester_id (required), subject_id (required), trust_level, claims")
+    print()
+    print("  iatp_verify")
+    print("    Verify trust relationship before agent-to-agent communication")
+    print("    Arguments: source_agent (required), target_agent (required), action, data_classification")
+    print()
+    print("  iatp_reputation")
+    print("    Query or modify agent reputation across the trust network")
+    print("    Arguments: agent_id (required), action (required), reason")
+
+
+def print_prompts():
+    """Print available prompts."""
+    print("Available MCP Prompts:")
+    print()
+    print("  governed_agent")
+    print("    Instructions for operating as a governed agent under Agent OS")
+    print("    Arguments: agent_id (required), policies")
+    print()
+    print("  verify_claim")
+    print("    Instructions for verifying a claim using CMVK")
+    print("    Arguments: claim (required)")
+    print()
+    print("  safe_execution")
+    print("    Template for executing actions safely through the kernel")
+    print("    Arguments: action (required), params (required)")
+
+
+def main():
+    """Main entry point."""
+    args = parse_args()
+    
+    if args.version:
+        from mcp_kernel_server import __version__
+        if args.json:
+            import json
+            print(json.dumps({
+                "name": "mcp-kernel-server",
+                "version": __version__,
+                "description": "Agent OS MCP Server for kernel-level AI agent governance"
+            }, indent=2))
+        else:
+            print(f"mcp-kernel-server {__version__}")
+            print(f"Agent OS MCP Server for kernel-level AI agent governance")
+        return 0
+    
+    if args.list_tools:
+        if args.json:
+            import json
+            tools = [
+                {"name": "cmvk_verify", "description": "Verify claims across multiple AI models to detect hallucinations"},
+                {"name": "kernel_execute", "description": "Execute actions through the Agent OS kernel with policy enforcement"},
+                {"name": "iatp_sign", "description": "Create a trust attestation for inter-agent communication"},
+                {"name": "iatp_verify", "description": "Verify trust relationship before agent-to-agent communication"},
+                {"name": "iatp_reputation", "description": "Query or modify agent reputation across the trust network"}
+            ]
+            print(json.dumps(tools, indent=2))
+        else:
+            print_tools()
+        return 0
+    
+    if args.list_prompts:
+        if args.json:
+            import json
+            prompts = [
+                {"name": "governed_agent", "description": "Instructions for operating as a governed agent under Agent OS"},
+                {"name": "verify_claim", "description": "Instructions for verifying a claim using CMVK"},
+                {"name": "safe_execution", "description": "Template for executing actions safely through the kernel"}
+            ]
+            print(json.dumps(prompts, indent=2))
+        else:
+            print_prompts()
+        return 0
+    
+    # Default to stdio if neither specified
+    if not args.stdio and not args.http:
+        args.stdio = True
+    
+    # Build config
+    config = ServerConfig(
+        host=args.host,
+        port=args.port,
+        policy_mode=args.policy_mode,
+        cmvk_threshold=args.cmvk_threshold
+    )
+    
+    # Create server
+    server = KernelMCPServer(config)
+    
+    # Run
+    try:
+        if args.stdio:
+            logger.info("Starting MCP Kernel Server with stdio transport")
+            logger.info(f"Policy mode: {args.policy_mode}")
+            logger.info("Tools: cmvk_verify, kernel_execute, iatp_sign, iatp_verify, iatp_reputation")
+            logger.info("Prompts: governed_agent, verify_claim, safe_execution")
+            
+            asyncio.run(server.run_stdio())
+        else:
+            # HTTP transport
+            logger.info(f"Starting MCP Kernel Server on http://{args.host}:{args.port}")
+            logger.info(f"Policy mode: {args.policy_mode}")
+            logger.info("Tools: cmvk_verify, kernel_execute, iatp_sign, iatp_verify, iatp_reputation")
+            logger.info("Prompts: governed_agent, verify_claim, safe_execution")
+            logger.info("Press Ctrl+C to stop")
+            
+            asyncio.run(server.start())
+            asyncio.get_event_loop().run_forever()
+            
+    except KeyboardInterrupt:
+        logger.info("Shutting down...")
+    except Exception as e:
+        # Sanitize startup errors for JSON mode to prevent info leakage
+        is_known = isinstance(e, (ValueError, PermissionError, OSError))
+        msg = "A validation or system error occurred." if is_known else "An internal error occurred during server startup"
+        
+        if getattr(args, "json", False):
+            import json
+            print(json.dumps({
+                "status": "error",
+                "message": msg,
+                "type": "StartupError" if is_known else "InternalError"
+            }, indent=2))
+        else:
+            logger.exception(f"Server error: {msg}")
+        return 1
+    
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)

mcp_kernel_server/resources.py

@@ -0,0 +1,217 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+MCP Resources for Agent OS VFS.
+
+Exposes the Agent Virtual File System as MCP resources,
+allowing agents to read/write structured memory.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+from datetime import datetime, timezone
+import json
+
+
+@dataclass
+class ResourceContent:
+    """Content returned from VFS resource."""
+    uri: str
+    mime_type: str
+    content: Any
+    metadata: dict = field(default_factory=dict)
+
+
+class VFSResource:
+    """
+    Agent Virtual File System as MCP Resource.
+    
+    Provides structured memory access through standard paths:
+    - /vfs/{agent_id}/mem/working/* - Ephemeral working memory
+    - /vfs/{agent_id}/mem/episodic/* - Experience logs
+    - /vfs/{agent_id}/policy/* - Read-only policies
+    
+    Stateless Design:
+    - All state stored in external backend (memory/redis/s3)
+    - No session state maintained in server
+    - Horizontally scalable
+    """
+    
+    uri_template = "vfs://{agent_id}/{path}"
+    
+    # Backend storage (in production: Redis, S3, DynamoDB)
+    _storage: Dict[str, Dict[str, Any]] = {}
+    
+    def __init__(self, config: Optional[dict] = None):
+        self.config = config or {}
+        self.backend = self.config.get("backend", "memory")
+    
+    def list_resources(self, agent_id: str) -> list:
+        """List available resources for an agent."""
+        return [
+            {
+                "uri": f"vfs://{agent_id}/mem/working",
+                "name": "Working Memory",
+                "description": "Ephemeral working memory for current task",
+                "mimeType": "application/json"
+            },
+            {
+                "uri": f"vfs://{agent_id}/mem/episodic",
+                "name": "Episodic Memory", 
+                "description": "Agent experience and history logs",
+                "mimeType": "application/json"
+            },
+            {
+                "uri": f"vfs://{agent_id}/policy",
+                "name": "Policies",
+                "description": "Agent policies and constraints (read-only)",
+                "mimeType": "application/json"
+            }
+        ]
+    
+    async def read(self, uri: str) -> ResourceContent:
+        """Read from VFS path."""
+        agent_id, path = self._parse_uri(uri)
+        
+        # Initialize agent storage if needed
+        if agent_id not in self._storage:
+            self._storage[agent_id] = self._init_agent_storage(agent_id)
+        
+        # Navigate to path
+        content = self._get_path(self._storage[agent_id], path)
+        
+        return ResourceContent(
+            uri=uri,
+            mime_type="application/json",
+            content=content,
+            metadata={
+                "agent_id": agent_id,
+                "path": path,
+                "timestamp": datetime.now(timezone.utc).isoformat()
+            }
+        )
+    
+    async def write(self, uri: str, content: Any) -> ResourceContent:
+        """Write to VFS path (if allowed)."""
+        agent_id, path = self._parse_uri(uri)
+        
+        # Check write permissions
+        if path.startswith("policy"):
+            raise PermissionError(f"Cannot write to policy path: {path}")
+        
+        # Initialize agent storage if needed
+        if agent_id not in self._storage:
+            self._storage[agent_id] = self._init_agent_storage(agent_id)
+        
+        # Write to path
+        self._set_path(self._storage[agent_id], path, content)
+        
+        return ResourceContent(
+            uri=uri,
+            mime_type="application/json",
+            content={"status": "written", "path": path},
+            metadata={
+                "agent_id": agent_id,
+                "path": path,
+                "timestamp": datetime.now(timezone.utc).isoformat()
+            }
+        )
+    
+    def _parse_uri(self, uri: str) -> tuple:
+        """Parse VFS URI into agent_id and path."""
+        # Handle vfs://agent_id/path format
+        if uri.startswith("vfs://"):
+            uri = uri[6:]
+        
+        parts = uri.split("/", 1)
+        agent_id = parts[0]
+        path = parts[1] if len(parts) > 1 else ""
+        
+        return agent_id, path
+    
+    def _init_agent_storage(self, agent_id: str) -> dict:
+        """Initialize storage structure for an agent."""
+        return {
+            "mem": {
+                "working": {
+                    "_meta": {"type": "working_memory", "ephemeral": True}
+                },
+                "episodic": {
+                    "_meta": {"type": "episodic_memory"},
+                    "sessions": []
+                }
+            },
+            "policy": {
+                "_meta": {"type": "policies", "read_only": True},
+                "default": {
+                    "name": "default",
+                    "rules": [
+                        {"action": "*", "effect": "allow"}
+                    ]
+                }
+            }
+        }
+    
+    def _get_path(self, storage: dict, path: str) -> Any:
+        """Navigate storage to get value at path."""
+        if not path:
+            return storage
+        
+        current = storage
+        for part in path.split("/"):
+            if not part:
+                continue
+            if isinstance(current, dict) and part in current:
+                current = current[part]
+            else:
+                return None
+        return current
+    
+    def _set_path(self, storage: dict, path: str, value: Any):
+        """Navigate storage to set value at path."""
+        parts = [p for p in path.split("/") if p]
+        
+        current = storage
+        for part in parts[:-1]:
+            if part not in current:
+                current[part] = {}
+            current = current[part]
+        
+        if parts:
+            current[parts[-1]] = value
+
+
+class VFSResourceTemplate:
+    """
+    MCP Resource Template for dynamic VFS paths.
+    
+    Allows clients to discover available resources dynamically.
+    """
+    
+    uri_template = "vfs://{agent_id}/{path}"
+    name = "Agent VFS"
+    description = "Virtual File System for agent memory and policies"
+    
+    @staticmethod
+    def get_templates() -> list:
+        """Return MCP resource templates."""
+        return [
+            {
+                "uriTemplate": "vfs://{agent_id}/mem/working/{key}",
+                "name": "Working Memory Item",
+                "description": "Read/write ephemeral working memory",
+                "mimeType": "application/json"
+            },
+            {
+                "uriTemplate": "vfs://{agent_id}/mem/episodic/{session_id}",
+                "name": "Episodic Session",
+                "description": "Read/write session history",
+                "mimeType": "application/json"
+            },
+            {
+                "uriTemplate": "vfs://{agent_id}/policy/{policy_name}",
+                "name": "Policy",
+                "description": "Read agent policy (read-only)",
+                "mimeType": "application/json"
+            }
+        ]