diff --git a/finir_agent/mcp_servers/ponytail_mcp_server.py b/finir_agent/mcp_servers/ponytail_mcp_server.py new file mode 100644 index 0000000..75420b5 --- /dev/null +++ b/finir_agent/mcp_servers/ponytail_mcp_server.py @@ -0,0 +1,525 @@ +"""Ponytail MCP Server — help agents simplify code.""" +import ast +import os +import re +import sys +from pathlib import Path +from typing import Any + +from mcp.server import Server, stdio_server +from mcp.types import TextContent, Tool + +# --- helpers --- + +PYTHON_KEYWORDS = frozenset( + "False None True and as assert async await break class continue " + "def del elif else except finally for from global if import in " + "is lambda nonlocal not or pass raise return try while with yield".split() +) + + +def walk_python_files(root: str) -> list[Path]: + """Recursively find all .py files under root.""" + found = [] + for p in Path(root).rglob("*.py"): + try: + rel = p.relative_to(root) + except ValueError: + rel = p + found.append(rel) + return sorted(found) + + +def read_lines(path: str) -> list[str]: + try: + with open(path) as f: + return f.readlines() + except OSError: + return [] + + +# --- tool 1: find reinvented stdlib --- + +# patterns: function names / class names that reimplement stdlib +REINVENTION_PATTERNS = { + "flatten": re.compile(r"def\s+(flatten|flatten_list)\s*\("), + "average": re.compile(r"def\s+(average|avg|mean)\s*\("), + "url_parse": re.compile( + r"(def\s+(parse_url|url_parse|parse_uri|url_parser)\s*\(|" + r"class\s+(URLParser|URIParser|UrlParser))" + ), + "deep_merge": re.compile(r"def\s+(deep_merge|merge_dicts)\s*\("), + "chunk": re.compile(r"def\s+(chunk|chunked|chunks)\s*\("), + "debounce": re.compile(r"def\s+(debounce|throttle)\s*\("), + "slugify": re.compile(r"def\s+(slugify|to_slug)\s*\("), + "unique": re.compile(r"def\s+(unique|deduplicate|distinct)\s*\("), + "flatten_json": re.compile(r"def\s+(flatten_json|flatten_dict)\s*\("), + "retry": re.compile(r"def\s+(retry|retry_with|with_retry)\s*\("), +} + + +def tool_find_reinvented_stdlib(_params: dict[str, Any]) -> list[dict[str, Any]]: + """Find reimplemented Python standard library functions.""" + cwd = os.getcwd() + results: list[dict[str, Any]] = [] + for pyfile in walk_python_files(cwd): + lines = read_lines(os.path.join(cwd, str(pyfile))) + for i, line in enumerate(lines, 1): + stripped = line.strip() + if stripped.startswith("#") or stripped.startswith("'''") or stripped.startswith('"""'): + continue + for name, pattern in REINVENTION_PATTERNS.items(): + if pattern.search(line): + results.append( + {"file": str(pyfile), "line": i, "code": stripped} + ) + break + return results + + +# --- tool 2: find dead code --- + +def _collect_imports(root: str, pyfiles: list[Path]) -> dict[str, str]: + """Map module name -> relative path.""" + imports: dict[str, str] = {} + for pyfile in pyfiles: + src = "\n".join(read_lines(os.path.join(root, str(pyfile)))) + try: + tree = ast.parse(src, filename=str(pyfile)) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + mod = alias.asname or alias.name + imports[mod] = str(pyfile) + elif isinstance(node, ast.ImportFrom): + mod = node.module or "" + if node.level > 0: + parts = str(pyfile).split("/")[: -node.level] + mod = "/".join(parts) + ("/" + mod if mod else "") + imports[mod] = str(pyfile) + return imports + + +def tool_find_dead_code(_params: dict[str, Any]) -> list[dict[str, Any]]: + """Find files never imported and symbols never referenced.""" + cwd = os.getcwd() + pyfiles = walk_python_files(cwd) + if not pyfiles: + return [] + + # Build set of all imported module names + imports = _collect_imports(cwd, pyfiles) + + results: list[dict[str, Any]] = [] + + # Dead files: .py files not imported anywhere (excluding __init__) + for pyfile in pyfiles: + name = pyfile.stem + if name == "__init__": + continue + stem = pyfile.stem + if stem not in imports: + # Check if any file imports this module + found_import = False + for pyfile2 in pyfiles: + if pyfile2 == pyfile: + continue + src = "\n".join(read_lines(os.path.join(cwd, str(pyfile2)))) + try: + tree = ast.parse(src, filename=str(pyfile2)) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == stem or alias.name.startswith(stem + "."): + found_import = True + break + elif isinstance(node, ast.ImportFrom): + if node.module and (node.module == stem or node.module.startswith(stem + ".")): + found_import = True + break + if not found_import: + results.append( + { + "file": str(pyfile), + "line": 1, + "symbol": stem, + "reason": "module never imported", + } + ) + + # Dead symbols: top-level functions/classes not referenced elsewhere + for pyfile in pyfiles: + src = "\n".join(read_lines(os.path.join(cwd, str(pyfile)))) + try: + tree = ast.parse(src, filename=str(pyfile)) + except SyntaxError: + continue + + # Collect defined top-level names + defined: list[tuple[str, int]] = [] + for node in ast.iter_child_nodes(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + defined.append((node.name, node.lineno)) + elif isinstance(node, ast.ClassDef): + defined.append((node.name, node.lineno)) + + # Collect all Name references across the whole project + all_refs: set[str] = set() + for other in pyfiles: + other_src = "\n".join(read_lines(os.path.join(cwd, str(other)))) + try: + other_tree = ast.parse(other_src, filename=str(other)) + except SyntaxError: + continue + for n in ast.walk(other_tree): + if isinstance(n, ast.Name): + all_refs.add(n.id) + elif isinstance(n, ast.Attribute): + if isinstance(n.value, ast.Name): + all_refs.add(n.value.id) + + # Also check self-references within same file (imports count as refs) + for imp_name in imports: + base = imp_name.split(".")[-1] + all_refs.add(base) + + for sym_name, lineno in defined: + if sym_name.startswith("_"): + continue + if sym_name in PYTHON_KEYWORDS: + continue + # Check if referenced outside its own file + refs_outside = False + for other in pyfiles: + if other == pyfile: + continue + other_src = "\n".join(read_lines(os.path.join(cwd, str(other)))) + try: + other_tree = ast.parse(other_src, filename=str(other)) + except SyntaxError: + continue + for n in ast.walk(other_tree): + if isinstance(n, ast.Name) and n.id == sym_name: + refs_outside = True + break + if isinstance(n, ast.Attribute) and isinstance(n.value, ast.Name) and n.value.id == sym_name: + refs_outside = True + break + if refs_outside: + break + + if not refs_outside: + results.append( + { + "file": str(pyfile), + "line": lineno, + "symbol": sym_name, + "reason": "symbol never referenced", + } + ) + + return results + + +# --- tool 3: find single-impl interfaces --- + +def tool_find_single_impl_interface(_params: dict[str, Any]) -> list[dict[str, Any]]: + """Find abstract classes with exactly one concrete implementation.""" + cwd = os.getcwd() + pyfiles = walk_python_files(cwd) + if not pyfiles: + return [] + + # Collect all abstract base classes + abstract_classes: dict[str, tuple[str, int]] = {} + concrete_implementations: dict[str, list[tuple[str, int]]] = {} + + for pyfile in pyfiles: + src = "\n".join(read_lines(os.path.join(cwd, str(pyfile)))) + try: + tree = ast.parse(src, filename=str(pyfile)) + except SyntaxError: + continue + + for node in ast.iter_child_nodes(tree): + if not isinstance(node, ast.ClassDef): + continue + + # Check if it's an abstract class + is_abstract = False + for base in node.bases: + base_name = "" + if isinstance(base, ast.Name): + base_name = base.id + elif isinstance(base, ast.Attribute): + base_name = base.attr + + if base_name in ("ABC", "ABCMeta"): + is_abstract = True + if isinstance(base, ast.Attribute): + full = f"{ast.dump(base)}" + if "ABCMeta" in full or "ABC" in full: + is_abstract = True + + # Check for abstract methods (decorated with @abstractmethod) + has_abstract_methods = False + for child in ast.iter_child_nodes(node): + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): + for dec in child.decorator_list: + dec_name = "" + if isinstance(dec, ast.Name): + dec_name = dec.id + elif isinstance(dec, ast.Attribute): + dec_name = dec.attr + if "abstractmethod" in dec_name.lower(): + has_abstract_methods = True + break + + class_name = node.name + if is_abstract or has_abstract_methods: + abstract_classes[class_name] = (str(pyfile), node.lineno) + concrete_implementations[class_name] = [] + + # Check if this class inherits from a known abstract class + for base in node.bases: + base_name = "" + if isinstance(base, ast.Name): + base_name = base.id + elif isinstance(base, ast.Attribute): + base_name = base.attr + if base_name in abstract_classes: + concrete_implementations[base_name].append( + (str(pyfile), node.lineno) + ) + + # Also check for ABCMeta metaclass usage + for pyfile in pyfiles: + src = "\n".join(read_lines(os.path.join(cwd, str(pyfile)))) + try: + tree = ast.parse(src, filename=str(pyfile)) + except SyntaxError: + continue + for node in ast.iter_child_nodes(tree): + if isinstance(node, ast.ClassDef): + for kw in node.keywords: + if kw.arg == "metaclass": + if isinstance(kw.value, ast.Name) and kw.value.id == "ABCMeta": + cn = node.name + if cn not in abstract_classes: + abstract_classes[cn] = (str(pyfile), node.lineno) + concrete_implementations[cn] = [] + + results: list[dict[str, Any]] = [] + for cls_name, (file_path, line_no) in abstract_classes.items(): + impls = concrete_implementations.get(cls_name, []) + if len(impls) == 1: + impl_file, impl_line = impls[0] + results.append( + { + "abstract": cls_name, + "implementation": impls[0][0].rsplit("/", 1)[-1].replace(".py", ""), + "file": impl_file, + "line": impl_line, + } + ) + return results + + +# --- tool 4: find ponytail comments --- + +def tool_find_ponytail_comments(_params: dict[str, Any]) -> list[dict[str, Any]]: + """Extract all //ponytail: or #ponytail: comments.""" + cwd = os.getcwd() + results: list[dict[str, Any]] = [] + ponytail_re = re.compile(r"(?://|#)ponytail:\s*(.*)") + + for pyfile in walk_python_files(cwd): + lines = read_lines(os.path.join(cwd, str(pyfile))) + for i, line in enumerate(lines, 1): + # Skip string literals (simple check) + stripped = line.strip() + if stripped.startswith("#!") or stripped.startswith('"""') or stripped.startswith("'''"): + continue + match = ponytail_re.search(line) + if match: + results.append( + { + "file": str(pyfile), + "line": i, + "comment": match.group(1).strip(), + } + ) + return results + + +# --- tool 5: analyze bloat --- + +def _nesting_depth(node: ast.AST) -> int: + """Compute max nesting depth of control structures.""" + max_depth = 0 + for child in ast.iter_child_nodes(node): + if isinstance(child, (ast.For, ast.AsyncFor, ast.While, ast.If, ast.With, + ast.AsyncWith, ast.Try, ast.ExceptHandler)): + child_depth = _nesting_depth(child) + 1 + max_depth = max(max_depth, child_depth) + else: + child_depth = _nesting_depth(child) + max_depth = max(max_depth, child_depth) + return max_depth + + +def _func_lengths(tree: ast.Module) -> list[int]: + """Get lengths of all top-level and nested functions.""" + lengths: list[int] = [] + + def walk(node: ast.AST): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + end = max( + (getattr(n, "end_lineno", n.lineno) for n in ast.walk(node)), + default=node.lineno, + ) + lengths.append(end - node.lineno + 1) + for child in ast.iter_child_nodes(node): + walk(child) + + walk(tree) + return lengths + + +def tool_analyze_bloat(params: dict[str, Any]) -> dict[str, Any]: + """Analyze a file for bloat: lines, function lengths, nesting, imports.""" + filepath = params.get("path", "") + if not filepath: + return {"error": "path required"} + + if not os.path.exists(filepath): + return {"error": f"file not found: {filepath}"} + + lines = read_lines(filepath) + total_lines = len(lines) + if total_lines == 0: + return { + "file": filepath, + "total_lines": 0, + "avg_function_length": 0, + "max_nesting_depth": 0, + "import_count": 0, + } + + src = "\n".join(lines) + try: + tree = ast.parse(src, filename=filepath) + except SyntaxError: + return { + "file": filepath, + "total_lines": total_lines, + "avg_function_length": 0, + "max_nesting_depth": 0, + "import_count": 0, + "parse_error": True, + } + + func_lengths = _func_lengths(tree) + avg_func_len = ( + sum(func_lengths) / len(func_lengths) if func_lengths else 0 + ) + + max_nesting = _nesting_depth(tree) + + import_count = 0 + for node in ast.walk(tree): + if isinstance(node, (ast.Import, ast.ImportFrom)): + import_count += 1 + + return { + "file": filepath, + "total_lines": total_lines, + "avg_function_length": round(avg_func_len, 1), + "max_nesting_depth": max_nesting, + "import_count": import_count, + } + + +# --- server setup --- + +async def run_server() -> None: + """Run the ponytail MCP server.""" + server = Server("ponytail") + + @server.list_tools() + async def list_tools() -> list[Tool]: + return [ + Tool( + name="ponytail_find_reinvented_stdlib", + description="Find reimplemented Python standard library functions (flatten, average, URL parser, etc.)", + inputSchema={ + "type": "object", + "properties": {}, + }, + ), + Tool( + name="ponytail_find_dead_code", + description="Find files never imported and symbols never referenced", + inputSchema={ + "type": "object", + "properties": {}, + }, + ), + Tool( + name="ponytail_find_single_impl_interface", + description="Find abstract classes with exactly one concrete implementation", + inputSchema={ + "type": "object", + "properties": {}, + }, + ), + Tool( + name="ponytail_find_ponytail_comments", + description="Extract all //ponytail: or #ponytail: comments from the workspace", + inputSchema={ + "type": "object", + "properties": {}, + }, + ), + Tool( + name="ponytail_analyze_bloat", + description="Analyze a file for bloat metrics: total lines, avg function length, max nesting depth, import count", + inputSchema={ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the Python file to analyze", + } + }, + "required": ["path"], + }, + ), + ] + + @server.call_tool() + async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: + handlers = { + "ponytail_find_reinvented_stdlib": tool_find_reinvented_stdlib, + "ponytail_find_dead_code": tool_find_dead_code, + "ponytail_find_single_impl_interface": tool_find_single_impl_interface, + "ponytail_find_ponytail_comments": tool_find_ponytail_comments, + "ponytail_analyze_bloat": tool_analyze_bloat, + } + handler = handlers.get(name) + if not handler: + return [TextContent(type="text", text=f"Unknown tool: {name}")] + result = handler(arguments or {}) + return [TextContent(type="text", text=__import__("json").dumps(result, separators=(",", ":")))] + + async with stdio_server() as (read_stream, write_stream): + await server.run(read_stream, write_stream) + + +if __name__ == "__main__": + import asyncio + asyncio.run(run_server())