Proto Server

# MCP Proto Server - Architecture ## Overview The MCP Proto Server is a Model Context Protocol (MCP) server implementation that indexes Protocol Buffer (.proto) files and exposes their structure to AI agents through standardized tools. It enables AI assistants to understand and navigate protobuf APIs without manual documentation. ## Architecture Diagram ``` ┌─────────────────────────────────────────────────────────────┐ │ AI Agent / MCP Client │ │ (Cursor, Claude, etc.) │ └───────────────────────────┬─────────────────────────────────┘ │ MCP Protocol (stdio) │ ┌───────────────────────────▼─────────────────────────────────┐ │ MCP Proto Server │ │ ┌───────────────────────────────────────────────────────┐ │ │ │ mcp_proto_server.py (MCP Interface) │ │ │ │ │ │ │ │ Tools: │ │ │ │ - search_proto(query, limit, min_score) │ │ │ │ - get_service_definition(name) │ │ │ │ - get_message_definition(name) │ │ │ └────────────────┬──────────────────────────────────────┘ │ │ │ │ │ ┌────────────────▼──────────────────────────────────────┐ │ │ │ proto_indexer.py (Search & Index) │ │ │ │ │ │ │ │ - In-memory index (services, messages, enums) │ │ │ │ - Fuzzy search with RapidFuzz │ │ │ │ - Name/field/comment search │ │ │ │ - File path tracking │ │ │ └────────────────┬──────────────────────────────────────┘ │ │ │ │ │ ┌────────────────▼──────────────────────────────────────┐ │ │ │ proto_parser.py (Proto File Parser) │ │ │ │ │ │ │ │ - Custom lexer/parser for .proto files │ │ │ │ - Extracts: packages, services, messages, enums │ │ │ │ - Comment extraction │ │ │ │ - Proto2 & Proto3 support │ │ │ └────────────────┬──────────────────────────────────────┘ │ └────────────────────┼──────────────────────────────────────┘ │ │ File I/O │ ┌────────────────────▼──────────────────────────────────────┐ │ Local Filesystem (.proto files) │ │ │ │ examples/api/v1/ │ │ ├── user.proto │ │ ├── auth.proto │ │ └── product.proto │ └─────────────────────────────────────────────────────────────┘ ``` ## Component Details ### 1. MCP Proto Server (`mcp_proto_server.py`) **Purpose:** MCP protocol handler and main entry point **Key Features:** - Implements MCP server using official Python SDK - Exposes three main tools as MCP resources - Handles stdio communication with MCP clients - Validates inputs and formats JSON responses - Manages server lifecycle and initialization **Tools Exposed:** 1. **search_proto**: Fuzzy search across all definitions - Inputs: query (string), limit (int), min_score (int) - Output: Array of search results with scores - Use case: "Find all authentication-related services" 2. **get_service_definition**: Get complete service structure - Input: name (string) - Output: Service with all RPC methods - Use case: "Show me UserService methods" 3. **get_message_definition**: Get message/enum structure - Input: name (string) - Output: Message fields or enum values - Use case: "What fields does User have?" **Technology:** - Python 3.11+ - MCP Python SDK (official) - Asyncio for async handling - Argparse for CLI ### 2. Proto Indexer (`proto_indexer.py`) **Purpose:** In-memory index with search capabilities **Data Structures:** ```python { 'files': Dict[str, ProtoFile], # file_path → ProtoFile 'services': Dict[str, ProtoService], # full_name → Service 'messages': Dict[str, ProtoMessage], # full_name → Message 'enums': Dict[str, ProtoEnum], # full_name → Enum '_search_entries': List[tuple] # (name, type, obj, file) } ``` **Search Algorithm:** 1. **Name matching**: RapidFuzz WRatio scorer on full names 2. **Comment matching**: Partial ratio on comment text 3. **Field matching**: Search within message field names 4. **RPC matching**: Search within service RPC names **Performance:** - O(1) exact lookups by full name - O(n) fuzzy search (optimized with RapidFuzz) - ~1-2 KB memory per definition - Sub-millisecond search times **Key Methods:** - `index_directory(root)`: Scan and index all .proto files - `search(query, limit, min_score)`: Fuzzy search - `get_service(name)`: Retrieve service definition - `get_message(name)`: Retrieve message definition ### 3. Proto Parser (`proto_parser.py`) **Purpose:** Parse .proto files and extract structured data **Parsing Strategy:** - Custom regex-based lexer/parser - No external proto compiler dependencies - Handles both proto2 and proto3 syntax - Preserves comments and metadata **Data Models:** ```python @dataclass class ProtoFile: path: str package: str syntax: str services: List[ProtoService] messages: List[ProtoMessage] enums: List[ProtoEnum] imports: List[str] @dataclass class ProtoService: name: str full_name: str rpcs: List[ProtoRPC] comment: Optional[str] @dataclass class ProtoMessage: name: str full_name: str fields: List[ProtoField] nested_messages: List[ProtoMessage] nested_enums: List[ProtoEnum] comment: Optional[str] @dataclass class ProtoRPC: name: str request_type: str response_type: str request_streaming: bool response_streaming: bool comment: Optional[str] ``` **Parsing Steps:** 1. Preprocess: Extract comments and clean content 2. Extract syntax, package, imports 3. Parse service blocks → RPC methods 4. Parse message blocks → fields 5. Parse enum blocks → values 6. Associate comments with definitions **Supported Features:** - ✅ proto2 and proto3 syntax - ✅ Services with RPC methods - ✅ Messages with fields - ✅ Enums with values - ✅ Comments (inline and block) - ✅ Optional/required/repeated labels - ✅ Streaming RPCs - ❌ Nested messages (partial) - ❌ Extensions - ❌ Options (partial) ## Data Flow ### Indexing Flow ``` 1. User starts server: python mcp_proto_server.py --root /path/to/protos 2. Server calls: index.index_directory(root) 3. Indexer walks filesystem: Path.rglob("*.proto") 4. For each file: a. Parser.parse_file(path) b. Extract services, messages, enums c. Add to index dictionaries d. Create search entries 5. Log statistics 6. Ready for queries ``` ### Search Flow ``` 1. AI agent calls: search_proto(query="auth") 2. MCP server receives request 3. Call indexer.search("auth", limit=20, min_score=60) 4. Indexer: a. Run fuzzy match on names (RapidFuzz) b. Search in comments (partial ratio) c. Search in fields/RPCs d. Combine and sort by score 5. Format as JSON 6. Return to agent via MCP ``` ### Definition Retrieval Flow ``` 1. AI agent calls: get_service_definition(name="UserService") 2. MCP server receives request 3. Call indexer.get_service("UserService") 4. Indexer: a. Try exact match: services["UserService"] b. Try qualified match: services["api.v1.UserService"] c. Try partial match: endswith(".UserService") 5. Format complete service with all RPCs 6. Return to agent via MCP ``` ## Design Decisions ### Why Python? - ✅ Excellent MCP SDK support - ✅ Rich ecosystem (RapidFuzz, asyncio) - ✅ Easy to extend and maintain - ✅ Cross-platform compatibility - ❌ Slower than Go (but fast enough) ### Why Custom Parser? - ✅ No protoc compiler needed - ✅ Simple installation (pip only) - ✅ Full control over comment extraction - ✅ Fast for our use case - ❌ May not support all proto features ### Why In-Memory Index? - ✅ Sub-millisecond search - ✅ Simple implementation - ✅ No database dependencies - ✅ Fast startup - ❌ Limited to available RAM - ❌ No persistence across restarts ### Why RapidFuzz? - ✅ Fastest fuzzy matching library - ✅ Multiple scoring algorithms - ✅ Good "did you mean?" behavior - ✅ Handles typos well ## Performance Characteristics | Operation | Time Complexity | Actual Performance | |-----------|----------------|-------------------| | Index single file | O(n) where n = file size | ~1ms per file | | Index directory | O(m*n) where m = files | ~1000 files/sec | | Exact lookup | O(1) | < 0.1ms | | Fuzzy search | O(n) where n = entries | < 10ms for 1000s entries | | Get definition | O(1) | < 0.1ms | **Memory Usage:** - Base: ~5 MB - Per file: ~10 KB - Per definition: ~1-2 KB - 1000 files: ~15-20 MB total **Scalability:** - Tested with 1000+ proto files - Supports 10,000+ definitions - Memory limit: ~1-2 GB (100K+ files) ## Extension Points ### 1. Semantic Search Add embedding-based search for better semantic matching: ```python # Add to proto_indexer.py import sentence_transformers class SemanticIndex: def __init__(self): self.model = sentence_transformers.SentenceTransformer('all-MiniLM-L6-v2') self.embeddings = [] def embed_definitions(self, definitions): texts = [d.name + " " + (d.comment or "") for d in definitions] self.embeddings = self.model.encode(texts) ``` ### 2. Persistent Cache Add SQLite caching to avoid re-parsing: ```python # Add to proto_indexer.py import sqlite3 import hashlib class CachedIndex: def should_reindex(self, file_path): file_hash = hashlib.md5(open(file_path, 'rb').read()).hexdigest() cached_hash = self.db.get_hash(file_path) return file_hash != cached_hash ``` ### 3. File Watching Already scaffolded, implement with watchdog: ```python from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler class ProtoFileHandler(FileSystemEventHandler): def on_modified(self, event): if event.src_path.endswith('.proto'): self.index.remove_file(event.src_path) self.index.index_file(event.src_path) ``` ### 4. Import Resolution Resolve imports to build dependency graphs: ```python class ImportResolver: def resolve_type(self, type_name, current_package): # Look in current package # Look in imports # Return fully qualified name pass ``` ## Security Considerations 1. **File System Access**: Only reads files, no writes 2. **Path Traversal**: Uses Path().resolve() for safety 3. **Input Validation**: All user inputs validated 4. **Resource Limits**: Configurable limits on search results 5. **No Code Execution**: Pure data processing, no eval/exec ## Testing Strategy 1. **Unit Tests**: Test each component independently 2. **Integration Tests**: Test full indexing → search flow 3. **Example Data**: Real-world proto files included 4. **Manual Testing**: test_server.py for validation ## Future Roadmap - [ ] Semantic search with embeddings - [ ] SQLite caching for faster startups - [ ] File watching for auto-reindexing - [ ] Import resolution and dependency graphs - [ ] REST API wrapper (optional) - [ ] Docker container - [ ] gRPC reflection support - [ ] Web UI for browsing - [ ] VS Code extension