dify/api/core/schemas/resolver.py

import logging
import re
import threading
from collections import deque
from dataclasses import dataclass
from typing import Any, Optional, Union

from core.schemas.registry import SchemaRegistry

logger = logging.getLogger(__name__)

# Type aliases for better clarity
SchemaType = Union[dict[str, Any], list[Any], str, int, float, bool, None]
SchemaDict = dict[str, Any]

# Pre-compiled pattern for better performance
_DIFY_SCHEMA_PATTERN = re.compile(r"^https://dify\.ai/schemas/(v\d+)/(.+)\.json$")


class SchemaResolutionError(Exception):
    """Base exception for schema resolution errors"""

    pass


class CircularReferenceError(SchemaResolutionError):
    """Raised when a circular reference is detected"""

    def __init__(self, ref_uri: str, ref_path: list[str]):
        self.ref_uri = ref_uri
        self.ref_path = ref_path
        super().__init__(f"Circular reference detected: {ref_uri} in path {' -> '.join(ref_path)}")


class MaxDepthExceededError(SchemaResolutionError):
    """Raised when maximum resolution depth is exceeded"""

    def __init__(self, max_depth: int):
        self.max_depth = max_depth
        super().__init__(f"Maximum resolution depth ({max_depth}) exceeded")


class SchemaNotFoundError(SchemaResolutionError):
    """Raised when a referenced schema cannot be found"""

    def __init__(self, ref_uri: str):
        self.ref_uri = ref_uri
        super().__init__(f"Schema not found: {ref_uri}")


@dataclass
class QueueItem:
    """Represents an item in the BFS queue"""

    current: Any
    parent: Optional[Any]
    key: Optional[Union[str, int]]
    depth: int
    ref_path: set[str]


class SchemaResolver:
    """Resolver for Dify schema references with caching and optimizations"""

    _cache: dict[str, SchemaDict] = {}
    _cache_lock = threading.Lock()

    def __init__(self, registry: Optional[SchemaRegistry] = None, max_depth: int = 10):
        """
        Initialize the schema resolver

        Args:
            registry: Schema registry to use (defaults to default registry)
            max_depth: Maximum depth for reference resolution
        """
        self.registry = registry or SchemaRegistry.default_registry()
        self.max_depth = max_depth

    @classmethod
    def clear_cache(cls) -> None:
        """Clear the global schema cache"""
        with cls._cache_lock:
            cls._cache.clear()

    def resolve(self, schema: SchemaType) -> SchemaType:
        """
        Resolve all $ref references in the schema

        Performance optimization: quickly checks for $ref presence before processing.

        Args:
            schema: Schema to resolve

        Returns:
            Resolved schema with all references expanded

        Raises:
            CircularReferenceError: If circular reference detected
            MaxDepthExceededError: If max depth exceeded
            SchemaNotFoundError: If referenced schema not found
        """
        if not isinstance(schema, (dict, list)):
            return schema

        # Fast path: if no Dify refs found, return original schema unchanged
        # This avoids expensive deepcopy and BFS traversal for schemas without refs
        if not _has_dify_refs(schema):
            return schema

        # Slow path: schema contains refs, perform full resolution
        import copy

        result = copy.deepcopy(schema)

        # Initialize BFS queue
        queue = deque([QueueItem(current=result, parent=None, key=None, depth=0, ref_path=set())])

        while queue:
            item = queue.popleft()

            # Process the current item
            self._process_queue_item(queue, item)

        return result

    def _process_queue_item(self, queue: deque, item: QueueItem) -> None:
        """Process a single queue item"""
        if isinstance(item.current, dict):
            self._process_dict(queue, item)
        elif isinstance(item.current, list):
            self._process_list(queue, item)

    def _process_dict(self, queue: deque, item: QueueItem) -> None:
        """Process a dictionary item"""
        ref_uri = item.current.get("$ref")

        if ref_uri and _is_dify_schema_ref(ref_uri):
            # Handle $ref resolution
            self._resolve_ref(queue, item, ref_uri)
        else:
            # Process nested items
            for key, value in item.current.items():
                if isinstance(value, (dict, list)):
                    next_depth = item.depth + 1
                    if next_depth >= self.max_depth:
                        raise MaxDepthExceededError(self.max_depth)
                    queue.append(
                        QueueItem(current=value, parent=item.current, key=key, depth=next_depth, ref_path=item.ref_path)
                    )

    def _process_list(self, queue: deque, item: QueueItem) -> None:
        """Process a list item"""
        for idx, value in enumerate(item.current):
            if isinstance(value, (dict, list)):
                next_depth = item.depth + 1
                if next_depth >= self.max_depth:
                    raise MaxDepthExceededError(self.max_depth)
                queue.append(
                    QueueItem(current=value, parent=item.current, key=idx, depth=next_depth, ref_path=item.ref_path)
                )

    def _resolve_ref(self, queue: deque, item: QueueItem, ref_uri: str) -> None:
        """Resolve a $ref reference"""
        # Check for circular reference
        if ref_uri in item.ref_path:
            # Mark as circular and skip
            item.current["$circular_ref"] = True
            logger.warning("Circular reference detected: %s", ref_uri)
            return

        # Get resolved schema (from cache or registry)
        resolved_schema = self._get_resolved_schema(ref_uri)
        if not resolved_schema:
            logger.warning("Schema not found: %s", ref_uri)
            return

        # Update ref path
        new_ref_path = item.ref_path | {ref_uri}

        # Replace the reference with resolved schema
        next_depth = item.depth + 1
        if next_depth >= self.max_depth:
            raise MaxDepthExceededError(self.max_depth)

        if item.parent is None:
            # Root level replacement
            item.current.clear()
            item.current.update(resolved_schema)
            queue.append(
                QueueItem(current=item.current, parent=None, key=None, depth=next_depth, ref_path=new_ref_path)
            )
        else:
            # Update parent container
            item.parent[item.key] = resolved_schema.copy()
            queue.append(
                QueueItem(
                    current=item.parent[item.key],
                    parent=item.parent,
                    key=item.key,
                    depth=next_depth,
                    ref_path=new_ref_path,
                )
            )

    def _get_resolved_schema(self, ref_uri: str) -> Optional[SchemaDict]:
        """Get resolved schema from cache or registry"""
        # Check cache first
        with self._cache_lock:
            if ref_uri in self._cache:
                return self._cache[ref_uri].copy()

        # Fetch from registry
        schema = self.registry.get_schema(ref_uri)
        if not schema:
            return None

        # Clean and cache
        cleaned = _remove_metadata_fields(schema)
        with self._cache_lock:
            self._cache[ref_uri] = cleaned

        return cleaned.copy()


def resolve_dify_schema_refs(
    schema: SchemaType, registry: Optional[SchemaRegistry] = None, max_depth: int = 30
) -> SchemaType:
    """
    Resolve $ref references in Dify schema to actual schema content

    This is a convenience function that creates a resolver and resolves the schema.
    Performance optimization: quickly checks for $ref presence before processing.

    Args:
        schema: Schema object that may contain $ref references
        registry: Optional schema registry, defaults to default registry
        max_depth: Maximum depth to prevent infinite loops (default: 30)

    Returns:
        Schema with all $ref references resolved to actual content

    Raises:
        CircularReferenceError: If circular reference detected
        MaxDepthExceededError: If maximum depth exceeded
        SchemaNotFoundError: If referenced schema not found
    """
    # Fast path: if no Dify refs found, return original schema unchanged
    # This avoids expensive deepcopy and BFS traversal for schemas without refs
    if not _has_dify_refs(schema):
        return schema

    # Slow path: schema contains refs, perform full resolution
    resolver = SchemaResolver(registry, max_depth)
    return resolver.resolve(schema)


def _remove_metadata_fields(schema: dict) -> dict:
    """
    Remove metadata fields from schema that shouldn't be included in resolved output

    Args:
        schema: Schema dictionary

    Returns:
        Cleaned schema without metadata fields
    """
    # Create a copy and remove metadata fields
    cleaned = schema.copy()
    metadata_fields = ["$id", "$schema", "version"]

    for field in metadata_fields:
        cleaned.pop(field, None)

    return cleaned


def _is_dify_schema_ref(ref_uri: Any) -> bool:
    """
    Check if the reference URI is a Dify schema reference

    Args:
        ref_uri: URI to check

    Returns:
        True if it's a Dify schema reference
    """
    if not isinstance(ref_uri, str):
        return False

    # Use pre-compiled pattern for better performance
    return bool(_DIFY_SCHEMA_PATTERN.match(ref_uri))


def _has_dify_refs_recursive(schema: SchemaType) -> bool:
    """
    Recursively check if a schema contains any Dify $ref references

    This is the fallback method when string-based detection is not possible.

    Args:
        schema: Schema to check for references

    Returns:
        True if any Dify $ref is found, False otherwise
    """
    if isinstance(schema, dict):
        # Check if this dict has a $ref field
        ref_uri = schema.get("$ref")
        if ref_uri and _is_dify_schema_ref(ref_uri):
            return True

        # Check nested values
        for value in schema.values():
            if _has_dify_refs_recursive(value):
                return True

    elif isinstance(schema, list):
        # Check each item in the list
        for item in schema:
            if _has_dify_refs_recursive(item):
                return True

    # Primitive types don't contain refs
    return False


def _has_dify_refs_hybrid(schema: SchemaType) -> bool:
    """
    Hybrid detection: fast string scan followed by precise recursive check

    Performance optimization using two-phase detection:
    1. Fast string scan to quickly eliminate schemas without $ref
    2. Precise recursive validation only for potential candidates

    Args:
        schema: Schema to check for references

    Returns:
        True if any Dify $ref is found, False otherwise
    """
    # Phase 1: Fast string-based pre-filtering
    try:
        import json

        schema_str = json.dumps(schema, separators=(",", ":"))

        # Quick elimination: no $ref at all
        if '"$ref"' not in schema_str:
            return False

        # Quick elimination: no Dify schema URLs
        if "https://dify.ai/schemas/" not in schema_str:
            return False

    except (TypeError, ValueError, OverflowError):
        # JSON serialization failed (e.g., circular references, non-serializable objects)
        # Fall back to recursive detection
        logger.debug("JSON serialization failed for schema, using recursive detection")
        return _has_dify_refs_recursive(schema)

    # Phase 2: Precise recursive validation
    # Only executed for schemas that passed string pre-filtering
    return _has_dify_refs_recursive(schema)


def _has_dify_refs(schema: SchemaType) -> bool:
    """
    Check if a schema contains any Dify $ref references

    Uses hybrid detection for optimal performance:
    - Fast string scan for quick elimination
    - Precise recursive check for validation

    Args:
        schema: Schema to check for references

    Returns:
        True if any Dify $ref is found, False otherwise
    """
    return _has_dify_refs_hybrid(schema)


def parse_dify_schema_uri(uri: str) -> tuple[str, str]:
    """
    Parse a Dify schema URI to extract version and schema name

    Args:
        uri: Schema URI to parse

    Returns:
        Tuple of (version, schema_name) or ("", "") if invalid
    """
    match = _DIFY_SCHEMA_PATTERN.match(uri)
    if not match:
        return "", ""

    return match.group(1), match.group(2)