onetool-mcp 1.0.0b1__py3-none-any.whl

ot/server.py

@@ -0,0 +1,315 @@
+"""FastMCP server implementation with a single 'run' tool.
+
+The agent generates function call syntax with __ot prefix:
+  __ot context7.search(query="next.js")
+  __ot context7.doc(library_key="vercel/next.js", topic="routing")
+  __ot `demo.upper(text="hello")`
+
+Or Python code blocks:
+  __ot
+  ```python
+  metals = ["Gold", "Silver", "Bronze"]
+  results = {}
+  for metal in metals:
+      results[metal] = brave.web_search(query=f"{metal} price", count=3)
+  return results
+  ```
+
+Or direct MCP calls:
+  mcp__onetool__run(command='brave.web_search(query="test")')
+
+Supported prefixes: __ot, __ot__run, __onetool, __onetool__run, mcp__onetool__run
+Note: mcp__ot__run is NOT valid.
+"""
+
+from __future__ import annotations
+
+import time
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+from fastmcp import Context, FastMCP
+from loguru import logger
+
+from ot.config.loader import get_config
+from ot.executor import SimpleExecutor, execute_command
+from ot.executor.runner import prepare_command
+
+# Import logging first to remove Loguru's default console handler
+from ot.logging import LogSpan, configure_logging
+from ot.prompts import get_prompts, get_tool_description, get_tool_examples
+from ot.proxy import get_proxy_manager
+from ot.registry import get_registry
+from ot.stats import (
+    JsonlStatsWriter,
+    get_client_name,
+    set_stats_writer,
+)
+from ot.support import get_startup_message
+
+_config = get_config()
+
+# Initialize logging to serve.log
+configure_logging(log_name="serve")
+
+# Global stats writer (unified JSONL for both run and tool stats)
+_stats_writer: JsonlStatsWriter | None = None
+
+
+def _get_instructions() -> str:
+    """Generate MCP server instructions.
+
+    Note: Tool descriptions are NOT included here - they come through
+    the MCP tool definitions which the client converts to function calling format.
+    """
+    # Load prompts from config (loaded via include: or inline prompts:)
+    prompts = get_prompts(inline_prompts=_config.prompts)
+
+    # Return instructions from prompts.yaml
+    return prompts.instructions.strip()
+
+
+@asynccontextmanager
+async def _lifespan(_server: FastMCP) -> AsyncIterator[None]:
+    """Manage server lifecycle - startup and shutdown."""
+    global _stats_writer
+
+    with LogSpan(span="mcp.server.start") as start_span:
+        # Startup: connect to proxy MCP servers
+        proxy = get_proxy_manager()
+        if _config.servers:
+            with LogSpan(span="server.startup.proxy", serverCount=len(_config.servers)):
+                await proxy.connect(_config.servers)
+            start_span.add("proxyCount", len(_config.servers))
+
+        # Log tool count from registry
+        registry = get_registry()
+        start_span.add("toolCount", len(registry.tools))
+
+        # Startup: initialize unified JSONL stats writer if enabled
+        if _config.stats.enabled:
+            stats_path = _config.get_stats_file_path()
+            flush_interval = _config.stats.flush_interval_seconds
+
+            _stats_writer = JsonlStatsWriter(
+                path=stats_path,
+                flush_interval=flush_interval,
+            )
+            await _stats_writer.start()
+            set_stats_writer(_stats_writer)
+
+            start_span.add("statsEnabled", True)
+            start_span.add("statsPath", str(stats_path))
+
+        # Log support message
+        logger.info(get_startup_message())
+
+    yield
+
+    with LogSpan(span="mcp.server.stop") as stop_span:
+        # Shutdown: stop stats writer
+        if _stats_writer is not None:
+            await _stats_writer.stop()
+            set_stats_writer(None)
+            stop_span.add("statsStopped", True)
+
+        # Shutdown: disconnect from proxy MCP servers
+        if proxy.servers:
+            with LogSpan(span="server.shutdown.proxy", serverCount=len(proxy.servers)):
+                await proxy.shutdown()
+            stop_span.add("proxyCount", len(proxy.servers))
+
+
+mcp = FastMCP(
+    name="ot",
+    instructions=_get_instructions(),
+    lifespan=_lifespan,
+)
+
+
+# =============================================================================
+# MCP Logging - Dynamic log level control
+# =============================================================================
+
+
+@mcp._mcp_server.set_logging_level()  # type: ignore[no-untyped-call,untyped-decorator]
+async def handle_set_logging_level(level: str) -> None:
+    """Handle logging/setLevel requests from MCP clients.
+
+    Allows clients to dynamically change the server's log level.
+    """
+    # Map MCP LoggingLevel to Python logging levels
+    level_map = {
+        "debug": "DEBUG",
+        "info": "INFO",
+        "notice": "INFO",  # MCP notice -> INFO
+        "warning": "WARNING",
+        "error": "ERROR",
+        "critical": "CRITICAL",
+        "alert": "CRITICAL",  # MCP alert -> CRITICAL
+        "emergency": "CRITICAL",  # MCP emergency -> CRITICAL
+    }
+
+    log_level = level_map.get(str(level).lower(), "INFO")
+    logger.info(f"Log level change requested: {level} -> {log_level}")
+
+    # Reconfigure logging with new level
+    configure_logging(log_name="serve", level=log_level)
+    logger.info(f"Logging reconfigured at level {log_level}")
+
+
+# =============================================================================
+# MCP Resources - Tool discoverability
+# =============================================================================
+
+
+@mcp.resource("ot://tools")
+def list_tools_resource() -> list[dict[str, str]]:
+    """List all available tools with signatures and descriptions."""
+    registry = get_registry()
+    prompts = get_prompts(inline_prompts=_config.prompts)
+
+    tools_list = []
+
+    # Add local tools
+    for tool in registry.tools.values():
+        desc = get_tool_description(prompts, tool.name, tool.description)
+        tools_list.append(
+            {
+                "name": tool.name,
+                "signature": tool.signature,
+                "description": desc,
+            }
+        )
+
+    # Add proxied tools
+    proxy = get_proxy_manager()
+    for proxy_tool in proxy.list_tools():
+        tools_list.append(
+            {
+                "name": f"{proxy_tool.server}.{proxy_tool.name}",
+                "signature": f"{proxy_tool.server}.{proxy_tool.name}(...)",
+                "description": f"[proxy] {proxy_tool.description}",
+            }
+        )
+
+    return tools_list
+
+
+@mcp.resource("ot://tool/{name}")
+def get_tool_resource(name: str) -> dict[str, Any]:
+    """Get detailed information about a specific tool."""
+    registry = get_registry()
+    prompts = get_prompts(inline_prompts=_config.prompts)
+
+    tool = registry.tools.get(name)
+    if not tool:
+        return {"error": f"Tool '{name}' not found"}
+
+    desc = get_tool_description(prompts, tool.name, tool.description)
+    examples = get_tool_examples(prompts, tool.name)
+
+    return {
+        "name": tool.name,
+        "module": tool.module,
+        "signature": tool.signature,
+        "description": desc,
+        "args": [
+            {
+                "name": arg.name,
+                "type": arg.type,
+                "default": arg.default,
+                "description": arg.description,
+            }
+            for arg in tool.args
+        ],
+        "returns": tool.returns,
+        "examples": examples or tool.examples,
+        "tags": tool.tags,
+        "enabled": tool.enabled,
+        "deprecated": tool.deprecated,
+        "deprecated_message": tool.deprecated_message,
+    }
+
+
+# Global executor instance
+_executor: SimpleExecutor | None = None
+
+
+def _get_executor() -> SimpleExecutor:
+    """Get or create the executor."""
+    global _executor
+
+    if _executor is None:
+        _executor = SimpleExecutor()
+
+    return _executor
+
+
+def _get_run_description() -> str:
+    """Get run tool description from prompts config.
+
+    Raises:
+        ValueError: If run tool description not found in prompts.yaml
+    """
+    prompts = get_prompts(inline_prompts=_config.prompts)
+    desc = get_tool_description(prompts, "run", "")
+    if not desc:
+        raise ValueError("Missing 'run' tool description in prompts.yaml")
+    return desc
+
+
+@mcp.tool(
+    description=_get_run_description(),
+    annotations={
+        "title": "Execute OneTool Command",
+        "readOnlyHint": False,
+        "destructiveHint": False,
+        "idempotentHint": False,
+        "openWorldHint": True,
+    },
+)
+async def run(command: str, ctx: Context) -> str:  # noqa: ARG001
+    # Get registry (cached, no rescan per request) and executor
+    registry = get_registry()
+    executor = _get_executor()
+
+    # Record start time for stats
+    start_time = time.monotonic()
+
+    # Step 1: Prepare and validate command
+    prepared = prepare_command(command)
+
+    if prepared.error:
+        return f"Error: {prepared.error}"
+
+    # Step 2: Execute through unified runner (skip validation since already done)
+    result = await execute_command(
+        command,
+        registry,
+        executor,
+        prepared_code=prepared.code,
+        skip_validation=True,
+    )
+
+    # Record run-level stats if enabled
+    if _stats_writer is not None:
+        duration_ms = int((time.monotonic() - start_time) * 1000)
+        _stats_writer.record_run(
+            client=get_client_name(),
+            chars_in=len(command),
+            chars_out=len(result.result),
+            duration_ms=duration_ms,
+            success=result.success,
+            error_type=result.error_type,
+        )
+
+    return result.result
+
+
+def main() -> None:
+    """Run the MCP server over stdio transport."""
+    mcp.run(show_banner=False)

ot/shortcuts/__init__.py

@@ -0,0 +1,15 @@
+"""Shortcuts system for OneTool.
+
+Provides aliases and snippets for simplified tool invocation:
+- Aliases: Short names mapping to full function names (e.g., ws -> brave.web_search)
+- Snippets: Jinja2 templates with variable substitution ($wsq q1=AI q2=ML p=Compare)
+"""
+
+from ot.shortcuts.aliases import resolve_alias
+from ot.shortcuts.snippets import expand_snippet, parse_snippet
+
+__all__ = [
+    "expand_snippet",
+    "parse_snippet",
+    "resolve_alias",
+]

ot/shortcuts/aliases.py

@@ -0,0 +1,87 @@
+"""Alias resolution for OneTool shortcuts.
+
+Resolves short alias names to their full namespaced function names.
+E.g., ws(query="test") -> brave.web_search(query="test")
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ot.config import OneToolConfig
+
+
+def resolve_alias(code: str, config: OneToolConfig) -> str:
+    """Resolve aliases in code to their full function names.
+
+    Replaces alias function calls with their target:
+    - ws(query="test") -> brave.web_search(query="test")
+    - c7(query="react") -> context7.search(query="react")
+
+    Args:
+        code: Python code potentially containing alias calls
+        config: Configuration with alias mappings
+
+    Returns:
+        Code with aliases resolved to full names
+    """
+    if not config.alias:
+        return code
+
+    result = code
+
+    # Sort aliases by length (longest first) to avoid partial matches
+    # e.g., "wsb" should be matched before "ws"
+    sorted_aliases = sorted(config.alias.keys(), key=len, reverse=True)
+
+    for alias_name in sorted_aliases:
+        target = config.alias[alias_name]
+
+        # Match alias followed by ( but not preceded by . or alphanumeric
+        # This prevents matching "foo.ws(" or "aws("
+        pattern = rf"(?<![.\w]){re.escape(alias_name)}\("
+
+        if re.search(pattern, result):
+            result = re.sub(pattern, f"{target}(", result)
+
+    return result
+
+
+def validate_aliases(config: OneToolConfig) -> list[str]:
+    """Validate alias configuration for circular references.
+
+    Args:
+        config: Configuration with alias mappings
+
+    Returns:
+        List of validation errors (empty if valid)
+    """
+    errors: list[str] = []
+
+    # Check for circular aliases
+    for alias_name, target in config.alias.items():
+        # Extract just the function name from target (before any dot)
+        target_base = target.split(".")[0] if "." in target else target
+
+        # Check if target points to another alias
+        if target_base in config.alias:
+            # Follow the chain to detect cycles
+            visited = {alias_name}
+            current = target_base
+
+            while current in config.alias:
+                if current in visited:
+                    errors.append(
+                        f"Circular alias detected: '{alias_name}' -> '{target}' "
+                        f"creates a cycle through '{current}'"
+                    )
+                    break
+                visited.add(current)
+                next_target = config.alias[current]
+                current = (
+                    next_target.split(".")[0] if "." in next_target else next_target
+                )
+
+    return errors

ot/shortcuts/snippets.py

@@ -0,0 +1,258 @@
+"""Snippet parsing and expansion for OneTool shortcuts.
+
+Handles snippet syntax parsing and Jinja2 template expansion:
+- Single-line: $wsq q1=AI q2=ML p=Compare
+- Multi-line: $wsq\nq1: AI\nq2: ML\np: Compare
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from loguru import logger
+
+if TYPE_CHECKING:
+    from ot.config import OneToolConfig, SnippetDef
+
+try:
+    from jinja2 import Environment, StrictUndefined, TemplateSyntaxError
+except ImportError as e:
+    raise ImportError(
+        "jinja2 is required for snippets. Install with: pip install jinja2"
+    ) from e
+
+
+@dataclass
+class ParsedSnippet:
+    """Result of parsing a snippet invocation."""
+
+    name: str
+    params: dict[str, str]
+    raw: str
+
+
+def is_snippet(code: str) -> bool:
+    """Check if code is a snippet invocation (starts with $).
+
+    Args:
+        code: Code to check
+
+    Returns:
+        True if code starts with $ (snippet syntax)
+    """
+    stripped = code.strip()
+    # Must start with $ but not be $variable inside other code
+    return stripped.startswith("$") and not stripped.startswith("${")
+
+
+def parse_snippet(code: str) -> ParsedSnippet:
+    """Parse a snippet invocation into name and parameters.
+
+    Supports two syntaxes:
+    - Single-line: $name key=value key2=value2
+    - Multi-line: $name\\nkey: value\\nkey2: value2
+
+    Args:
+        code: Snippet invocation string
+
+    Returns:
+        ParsedSnippet with name and extracted parameters
+
+    Raises:
+        ValueError: If snippet syntax is invalid
+    """
+    stripped = code.strip()
+
+    if not stripped.startswith("$"):
+        raise ValueError(f"Snippet must start with $: {stripped[:50]}")
+
+    # Remove $ prefix
+    content = stripped[1:]
+
+    # Check for multi-line (has newline after snippet name)
+    lines = content.split("\n")
+    first_line = lines[0].strip()
+
+    # Extract snippet name (first word)
+    name_match = re.match(r"^(\w+)", first_line)
+    if not name_match:
+        raise ValueError(f"Invalid snippet name: {first_line[:50]}")
+
+    name = name_match.group(1)
+
+    # Check if multi-line or single-line
+    if len(lines) > 1:
+        return _parse_multiline_snippet(name, lines[1:], stripped)
+    else:
+        return _parse_singleline_snippet(name, first_line[len(name) :], stripped)
+
+
+def _strip_quotes(value: str) -> str:
+    """Strip matching outer quotes from a value.
+
+    Handles both single and double quotes. Only strips if quotes are balanced.
+
+    Args:
+        value: String that may have outer quotes
+
+    Returns:
+        String with outer quotes removed if present and balanced
+    """
+    if len(value) >= 2 and (
+        (value.startswith('"') and value.endswith('"'))
+        or (value.startswith("'") and value.endswith("'"))
+    ):
+        return value[1:-1]
+    return value
+
+
+def _parse_singleline_snippet(name: str, params_str: str, raw: str) -> ParsedSnippet:
+    """Parse single-line snippet parameters: key=value key2="value with spaces".
+
+    Values extend until the next key= or end of string.
+    Outer quotes are stripped from values (key="value" becomes key=value).
+    Escaped equals (\\=) are preserved in values.
+    """
+    params: dict[str, str] = {}
+    params_str = params_str.strip()
+
+    if not params_str:
+        return ParsedSnippet(name=name, params=params, raw=raw)
+
+    # Replace escaped equals with placeholder
+    placeholder = "\x00EQUALS\x00"
+    params_str = params_str.replace("\\=", placeholder)
+
+    # Find all key=value pairs
+    # Pattern: word followed by = and then value until next word= or end
+    pattern = r"(\w+)=((?:[^=]|$)*?)(?=\s+\w+=|$)"
+    matches = re.findall(pattern, params_str)
+
+    for key, value in matches:
+        # Restore escaped equals and strip whitespace
+        value = value.replace(placeholder, "=").strip()
+        # Strip outer quotes from value (e.g., packages="react" -> packages=react)
+        value = _strip_quotes(value)
+        params[key] = value
+
+    return ParsedSnippet(name=name, params=params, raw=raw)
+
+
+def _parse_multiline_snippet(name: str, lines: list[str], raw: str) -> ParsedSnippet:
+    """Parse multi-line snippet parameters: key: value.
+
+    Blank line terminates the snippet parameters.
+    Only the first colon is the separator (colons in values are preserved).
+    Outer quotes are stripped from values for consistency with single-line format.
+    """
+    params: dict[str, str] = {}
+
+    for line in lines:
+        stripped = line.strip()
+
+        # Blank line terminates
+        if not stripped:
+            break
+
+        # Parse key: value (only first colon is separator)
+        colon_idx = stripped.find(":")
+        if colon_idx == -1:
+            logger.warning(f"Invalid snippet line (no colon): {stripped}")
+            continue
+
+        key = stripped[:colon_idx].strip()
+        value = stripped[colon_idx + 1 :].strip()
+
+        if not key:
+            logger.warning(f"Empty key in snippet line: {stripped}")
+            continue
+
+        # Strip outer quotes from value for consistency
+        value = _strip_quotes(value)
+        params[key] = value
+
+    return ParsedSnippet(name=name, params=params, raw=raw)
+
+
+def expand_snippet(
+    parsed: ParsedSnippet,
+    config: OneToolConfig,
+) -> str:
+    """Expand a parsed snippet using Jinja2 templating.
+
+    Args:
+        parsed: Parsed snippet with name and parameters
+        config: Configuration with snippet definitions
+
+    Returns:
+        Expanded Python code from the snippet template
+
+    Raises:
+        ValueError: If snippet not found, missing required params, or Jinja2 error
+    """
+    if parsed.name not in config.snippets:
+        available = ", ".join(sorted(config.snippets.keys())) or "(none)"
+        raise ValueError(f"Unknown snippet '{parsed.name}'. Available: {available}")
+
+    snippet_def: SnippetDef = config.snippets[parsed.name]
+
+    # Build context with defaults and provided values
+    context: dict[str, Any] = {}
+
+    # Apply defaults first
+    for param_name, param_def in snippet_def.params.items():
+        if param_def.default is not None:
+            context[param_name] = param_def.default
+
+    # Apply provided values
+    for key, value in parsed.params.items():
+        if key not in snippet_def.params:
+            logger.warning(
+                f"Unknown parameter '{key}' for snippet '{parsed.name}' (ignored)"
+            )
+        context[key] = value
+
+    # Check required parameters
+    for param_name, param_def in snippet_def.params.items():
+        if param_def.required and param_name not in context:
+            raise ValueError(
+                f"Snippet '{parsed.name}' requires parameter '{param_name}'"
+            )
+
+    # Render template with Jinja2
+    try:
+        env = Environment(undefined=StrictUndefined)
+        template = env.from_string(snippet_def.body)
+        return template.render(**context)
+    except TemplateSyntaxError as e:
+        raise ValueError(f"Jinja2 syntax error in snippet '{parsed.name}': {e}") from e
+    except Exception as e:
+        # StrictUndefined raises UndefinedError for missing variables
+        if "undefined" in str(e).lower():
+            raise ValueError(
+                f"Undefined variable in snippet '{parsed.name}': {e}"
+            ) from e
+        raise ValueError(f"Error expanding snippet '{parsed.name}': {e}") from e
+
+
+def validate_snippets(config: OneToolConfig) -> list[str]:
+    """Validate snippet definitions for Jinja2 syntax errors.
+
+    Args:
+        config: Configuration with snippet definitions
+
+    Returns:
+        List of validation errors (empty if valid)
+    """
+    errors: list[str] = []
+    env = Environment(undefined=StrictUndefined)
+
+    for name, snippet_def in config.snippets.items():
+        try:
+            env.from_string(snippet_def.body)
+        except TemplateSyntaxError as e:
+            errors.append(f"Snippet '{name}' has invalid Jinja2 syntax: {e}")
+
+    return errors