polyrob 0.5.0__py3-none-any.whl

tools/mcp/validation_tracker.py

@@ -0,0 +1,89 @@
+"""MCP validation-failure tracking — schema-injection policy.
+
+Extracted from ToolCallTracker (PR11). Repeated MCP-tool validation failures
+should escalate to injecting the tool's full schema to help the LLM. That is
+MCP error *policy*, not tool-call ID lifecycle, so it lives here in tools/mcp/
+rather than on the generic tracker. ToolCallTracker composes one of these and
+delegates to it, preserving its public method names for existing callers.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timedelta
+from threading import RLock
+from typing import Dict, Optional
+
+
+class MCPValidationTracker:
+    """Per-(server, tool) validation-failure counter with TTL expiry.
+
+    After ``failure_threshold`` failures within ``failure_ttl_minutes``, callers
+    should inject the full tool schema. A successful call clears the counter.
+    Self-contained and thread-safe (its own lock).
+    """
+
+    def __init__(
+        self,
+        logger: Optional[logging.Logger] = None,
+        failure_threshold: int = 2,
+        failure_ttl_minutes: int = 30,
+    ) -> None:
+        self.logger = logger or logging.getLogger(__name__)
+        self._failures: Dict[str, int] = {}
+        self._timestamps: Dict[str, datetime] = {}
+        self._threshold = failure_threshold
+        self._ttl_minutes = failure_ttl_minutes
+        self._lock = RLock()
+
+    @staticmethod
+    def _key(server_name: str, tool_name: str) -> str:
+        return f"{server_name}:{tool_name}"
+
+    def track_failure(self, server_name: str, tool_name: str) -> int:
+        """Record a validation failure; return the current count for this tool."""
+        with self._lock:
+            key = self._key(server_name, tool_name)
+            self._cleanup_expired()
+            self._failures[key] = self._failures.get(key, 0) + 1
+            self._timestamps[key] = datetime.now()
+            count = self._failures[key]
+            self.logger.warning(f"MCP validation failure #{count} for {key}")
+            return count
+
+    def should_inject_schema(self, server_name: str, tool_name: str) -> bool:
+        """True once failures for this tool reach the threshold.
+
+        P2-12: TTL-clean on the READ path too. Previously only track_failure cleaned,
+        so once a (server, tool) reached the threshold it stayed there for the whole
+        session — the consumer blocks the action, so clear_failures (success-only) never
+        runs and, absent another MCP failure to trigger track_failure, the 30-min TTL
+        never fired. Cleaning here lets an expired failure record lapse.
+        """
+        with self._lock:
+            self._cleanup_expired()
+            return self._failures.get(self._key(server_name, tool_name), 0) >= self._threshold
+
+    def clear_failures(self, server_name: str, tool_name: str) -> None:
+        """Clear the failure count after a successful execution."""
+        with self._lock:
+            key = self._key(server_name, tool_name)
+            if key in self._failures:
+                del self._failures[key]
+                self._timestamps.pop(key, None)
+                self.logger.debug(f"Cleared MCP validation failures for {key}")
+
+    def reset(self) -> None:
+        """Drop all tracked failures."""
+        with self._lock:
+            self._failures.clear()
+            self._timestamps.clear()
+
+    def _cleanup_expired(self) -> None:
+        """Remove failures older than the TTL. Caller holds the lock."""
+        cutoff = datetime.now() - timedelta(minutes=self._ttl_minutes)
+        expired = [k for k, ts in self._timestamps.items() if ts < cutoff]
+        for key in expired:
+            self._failures.pop(key, None)
+            self._timestamps.pop(key, None)
+            self.logger.debug(f"Expired MCP validation failures for {key}")

tools/mcp/views.py

@@ -0,0 +1,346 @@
+"""Pydantic models for MCP service actions."""
+
+from typing import Dict, Any, Optional, List
+from pydantic import BaseModel, Field, field_validator, ConfigDict
+
+
+class MCPExecuteToolAction(BaseModel):
+    """Parameters for executing an MCP tool."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server (e.g., 'anysite', 'ghost')")
+    tool_name: str = Field(..., description="Name of the tool to execute on the server")
+    # FIX (Jan 2026): Make arguments REQUIRED so LLMs cannot omit it
+    # Previously optional with default={}, but LLMs were completely omitting the field,
+    # causing all MCP tools to receive empty parameters.
+    # Making it required forces LLMs to include arguments={...} even if empty.
+    arguments: Dict[str, Any] = Field(
+        ...,  # Required - LLM MUST include this field
+        description=(
+            "REQUIRED: The parameters to pass to the target tool. "
+            "This dict is passed directly to the tool - PUT YOUR SEARCH PARAMETERS HERE! "
+            "Example: {'query': 'AI startups', 'count': 10} for search tools. "
+            "Example: {'keywords': 'CEO founder', 'count': 5} for LinkedIn search. "
+            "Use mcp_list_tools to see each tool's required parameters. "
+            "For tools that need no parameters, pass empty dict {}."
+        )
+    )
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+    @field_validator('tool_name')
+    @classmethod
+    def validate_tool_name(cls, v):
+        """Validate tool name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("tool_name cannot be empty")
+        return v.strip()
+
+    @field_validator('arguments')
+    @classmethod
+    def validate_arguments(cls, v):
+        """Warn about empty arguments - most tools require them."""
+        if v is None:
+            v = {}
+        if not v:
+            # Log warning but allow empty for tools that don't need params
+            import logging
+            logging.getLogger(__name__).warning(
+                f"⚠️ Empty arguments provided for MCP tool. "
+                f"Most MCP tools require parameters - check mcp_list_tools for schema."
+            )
+        return v
+
+
+class MCPReadResourceAction(BaseModel):
+    """Parameters for reading an MCP resource."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server")
+    resource_uri: str = Field(..., description="URI of the resource to read")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+    @field_validator('resource_uri')
+    @classmethod
+    def validate_resource_uri(cls, v):
+        """Validate resource URI is not empty."""
+        if not v or not v.strip():
+            raise ValueError("resource_uri cannot be empty")
+        return v.strip()
+
+
+class MCPListToolsAction(BaseModel):
+    """Parameters for listing MCP tools."""
+    model_config = ConfigDict(extra='ignore')  # Allow extra fields but ignore them
+
+    server_name: Optional[str] = Field(None, description="Specific server name (optional, lists all if not provided)")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty if provided."""
+        if v is not None and (not v or not v.strip()):
+            raise ValueError("server_name cannot be empty if provided")
+        return v.strip() if v else None
+
+
+class MCPListResourcesAction(BaseModel):
+    """Parameters for listing MCP resources."""
+    model_config = ConfigDict(extra='ignore')  # Allow extra fields but ignore them
+
+    server_name: Optional[str] = Field(None, description="Specific server name (optional, lists all if not provided)")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty if provided."""
+        if v is not None and (not v or not v.strip()):
+            raise ValueError("server_name cannot be empty if provided")
+        return v.strip() if v else None
+
+
+class MCPListServersAction(BaseModel):
+    """Parameters for listing MCP servers."""
+    model_config = ConfigDict(extra='ignore')  # Allow extra fields but ignore them
+    
+    include_disabled: bool = Field(False, description="Whether to include disabled servers")
+    include_details: bool = Field(True, description="Whether to include detailed server information")
+
+
+class MCPServerStatusAction(BaseModel):
+    """Parameters for getting MCP server status."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+
+class MCPConnectServerAction(BaseModel):
+    """Parameters for connecting to an MCP server."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server to connect")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+
+class MCPDisconnectServerAction(BaseModel):
+    """Parameters for disconnecting from an MCP server."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server to disconnect")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+
+class MCPReloadServerAction(BaseModel):
+    """Parameters for reloading an MCP server connection."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server to reload")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+
+class MCPGetCapabilitiesAction(BaseModel):
+    """Parameters for getting MCP server capabilities."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+
+class MCPSubscribeResourceAction(BaseModel):
+    """Parameters for subscribing to MCP resource updates."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server")
+    resource_uri: str = Field(..., description="URI of the resource to subscribe to")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+    @field_validator('resource_uri')
+    @classmethod
+    def validate_resource_uri(cls, v):
+        """Validate resource URI is not empty."""
+        if not v or not v.strip():
+            raise ValueError("resource_uri cannot be empty")
+        return v.strip()
+
+
+class MCPUnsubscribeResourceAction(BaseModel):
+    """Parameters for unsubscribing from MCP resource updates."""
+    model_config = ConfigDict(extra='forbid')
+
+    server_name: str = Field(..., description="Name of the MCP server")
+    resource_uri: str = Field(..., description="URI of the resource to unsubscribe from")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("server_name cannot be empty")
+        return v.strip()
+
+    @field_validator('resource_uri')
+    @classmethod
+    def validate_resource_uri(cls, v):
+        """Validate resource URI is not empty."""
+        if not v or not v.strip():
+            raise ValueError("resource_uri cannot be empty")
+        return v.strip()
+
+
+class MCPHealthCheckAction(BaseModel):
+    """Parameters for performing MCP server health check."""
+    model_config = ConfigDict(extra='ignore')
+
+    server_name: Optional[str] = Field(None, description="Specific server name (optional, checks all if not provided)")
+
+    @field_validator('server_name')
+    @classmethod
+    def validate_server_name(cls, v):
+        """Validate server name is not empty if provided."""
+        if v is not None and (not v or not v.strip()):
+            raise ValueError("server_name cannot be empty if provided")
+        return v.strip() if v else None
+
+
+# Response models for better type safety and documentation
+class MCPToolInfo(BaseModel):
+    """Information about an MCP tool."""
+    
+    name: str
+    description: str
+    server_name: str
+    input_schema: Dict[str, Any]
+    
+    def __str__(self) -> str:
+        """Format tool info for LLM visibility with schema details."""
+        # Extract required parameters from schema
+        props = self.input_schema.get('properties', {})
+        required = self.input_schema.get('required', [])
+        
+        # Build parameter summary
+        params = []
+        for param_name, param_schema in props.items():
+            param_type = param_schema.get('type', 'any')
+            is_required = param_name in required
+            req_marker = '(required)' if is_required else '(optional)'
+            param_desc = param_schema.get('description', '')
+            params.append(f"  - {param_name}: {param_type} {req_marker} - {param_desc}")
+        
+        params_str = '\n'.join(params) if params else '  (no parameters)'
+        
+        return f"""
+Tool: {self.name}
+Server: {self.server_name}
+Description: {self.description}
+Parameters:
+{params_str}
+"""
+
+
+class MCPResourceInfo(BaseModel):
+    """Information about an MCP resource."""
+    
+    uri: str
+    name: str
+    server_name: str
+    description: Optional[str] = None
+    mime_type: Optional[str] = None
+
+
+class MCPServerInfo(BaseModel):
+    """Information about an MCP server."""
+    
+    name: str
+    status: str
+    type: str
+    enabled: bool
+    connected_at: Optional[float] = None
+    last_error: Optional[str] = None
+    retry_count: int = 0
+    tools_count: int = 0
+    resources_count: int = 0
+    capabilities: Dict[str, Any] = Field(default_factory=dict)
+
+
+class MCPExecutionResult(BaseModel):
+    """Result of MCP tool execution."""
+    
+    success: bool
+    result: Any = None
+    error: Optional[str] = None
+    execution_time: Optional[float] = None
+
+
+class MCPResourceContent(BaseModel):
+    """Content of an MCP resource."""
+    
+    uri: str
+    content: Any
+    mime_type: Optional[str] = None
+    size: Optional[int] = None
+    last_modified: Optional[float] = None
+
+
+class MCPHealthStatus(BaseModel):
+    """Health status of MCP servers."""
+    
+    healthy_servers: List[str] = Field(default_factory=list)
+    unhealthy_servers: List[str] = Field(default_factory=list)
+    total_servers: int = 0
+    overall_health: str = "unknown"  # "healthy", "degraded", "unhealthy"

tools/oauth/__init__.py

@@ -0,0 +1,17 @@
+"""Unified OAuth manager (Item 4 — minimal WS-G).
+
+Library-only: an ``OAuthProvider`` seam + ``OAuthManager`` registry on the existing
+Fernet token store, plus one ``GenericOAuth2Provider``. No existing tool is migrated
+(deferred). Mint/store(encrypted)/refresh a token per user through the manager.
+"""
+from tools.oauth.provider import OAuthError, OAuthProvider, OAuthToken
+from tools.oauth.manager import OAuthManager
+from tools.oauth.providers.generic_oauth2 import GenericOAuth2Provider
+
+__all__ = [
+    "OAuthError",
+    "OAuthProvider",
+    "OAuthToken",
+    "OAuthManager",
+    "GenericOAuth2Provider",
+]

tools/oauth/manager.py

@@ -0,0 +1,75 @@
+"""OAuth manager (Item 4 — WS-G1, trimmed).
+
+Holds registered ``OAuthProvider`` instances and a per-``(user_id, provider)`` token
+store encrypted with the EXISTING Fernet helper (``tools/mcp/security.py::
+MCPEncryption``). ``get_token`` returns a cached valid token, auto-refreshing on
+expiry. Library-only — no tool migration; mirrors ``MemoryProviderRegistry``'s shape.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, MutableMapping, Optional, Tuple
+
+from tools.oauth.provider import OAuthError, OAuthProvider, OAuthToken
+
+Key = Tuple[str, str]
+
+
+class OAuthManager:
+    """Provider registry + encrypted token store keyed by ``(user_id, provider)``."""
+
+    def __init__(self, encryption: Any = None, store: Optional[MutableMapping[Key, bytes]] = None) -> None:
+        self._providers: Dict[str, OAuthProvider] = {}
+        # Encrypted-at-rest store: (user_id, provider) -> Fernet ciphertext bytes.
+        self._store: MutableMapping[Key, bytes] = store if store is not None else {}
+        self._encryption = encryption or self._default_encryption()
+
+    @staticmethod
+    def _default_encryption():
+        from tools.mcp.security import MCPEncryption
+        return MCPEncryption()
+
+    # -- provider registry ----------------------------------------------------
+
+    def register(self, provider: OAuthProvider) -> OAuthProvider:
+        self._providers[provider.name] = provider
+        return provider
+
+    def get_provider(self, name: str) -> OAuthProvider:
+        prov = self._providers.get(name)
+        if prov is None:
+            raise OAuthError(
+                f"unknown oauth provider '{name}' (known: {sorted(self._providers)})"
+            )
+        return prov
+
+    # -- encrypted token store ------------------------------------------------
+
+    def store_token(self, user_id: str, provider: str, token: OAuthToken) -> None:
+        self._store[(user_id, provider)] = self._encryption.encrypt_dict(token.to_dict())
+
+    def load_token(self, user_id: str, provider: str) -> Optional[OAuthToken]:
+        blob = self._store.get((user_id, provider))
+        if not blob:
+            return None
+        return OAuthToken.from_dict(self._encryption.decrypt_dict(blob))
+
+    def has_token(self, user_id: str, provider: str) -> bool:
+        return (user_id, provider) in self._store
+
+    # -- main entrypoint ------------------------------------------------------
+
+    async def get_token(self, user_id: str, provider: str) -> OAuthToken:
+        """Return a valid token for ``(user_id, provider)``, refreshing if expired."""
+        prov = self.get_provider(provider)
+        token = self.load_token(user_id, provider)
+        if token is None:
+            raise OAuthError(f"no token stored for user='{user_id}' provider='{provider}'")
+        if prov.valid_token(token):
+            return token
+        if not token.refresh_token:
+            raise OAuthError(
+                f"token expired and no refresh_token for user='{user_id}' provider='{provider}'"
+            )
+        refreshed = await prov.refresh(token)
+        self.store_token(user_id, provider, refreshed)
+        return refreshed

tools/oauth/provider.py

@@ -0,0 +1,69 @@
+"""OAuth provider seam (Item 4 — WS-G1, trimmed).
+
+An ``OAuthProvider`` knows how to build an authorize URL, exchange an auth code for
+tokens, and refresh them. ``OAuthToken`` is the normalised token shape stored
+(encrypted) by the ``OAuthManager``.
+"""
+from __future__ import annotations
+
+import time
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, Optional
+
+
+class OAuthError(RuntimeError):
+    """Raised for unknown-provider / missing-token / refresh failures."""
+
+
+@dataclass
+class OAuthToken:
+    """A normalised OAuth2 token set."""
+    access_token: str
+    refresh_token: Optional[str] = None
+    expires_at: Optional[float] = None  # epoch seconds; None = non-expiring
+    token_type: str = "Bearer"
+    scope: Optional[str] = None
+    extra: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "OAuthToken":
+        d = dict(d or {})
+        known = {f for f in cls.__dataclass_fields__}  # type: ignore[attr-defined]
+        extra = {k: d.pop(k) for k in list(d) if k not in known}
+        token = cls(**{k: v for k, v in d.items() if k in known})
+        if extra:
+            token.extra.update(extra)
+        return token
+
+
+class OAuthProvider(ABC):
+    """Auth-code + refresh-token OAuth2 provider."""
+
+    name: str = "oauth"
+
+    #: refresh this many seconds before actual expiry
+    leeway_sec: float = 60.0
+
+    @abstractmethod
+    def authorize_url(self, *, state: Optional[str] = None, redirect_uri: Optional[str] = None) -> str:
+        """Build the user-facing authorization URL."""
+
+    @abstractmethod
+    async def exchange_code(self, code: str, *, redirect_uri: Optional[str] = None) -> OAuthToken:
+        """Exchange an authorization code for a token set."""
+
+    @abstractmethod
+    async def refresh(self, token: OAuthToken) -> OAuthToken:
+        """Exchange a refresh token for a fresh token set."""
+
+    def valid_token(self, token: Optional[OAuthToken]) -> bool:
+        """True if ``token`` is present and not within ``leeway_sec`` of expiry."""
+        if not token or not token.access_token:
+            return False
+        if token.expires_at is None:
+            return True
+        return time.time() < (token.expires_at - self.leeway_sec)

tools/oauth/providers/__init__.py

@@ -0,0 +1 @@
+"""OAuth providers (Item 4). Only the generic auth-code provider ships."""