"""
Strict validation layer for PageContext and tool safety.
This module provides validation and safety checks for the context-aware MCP system.
All context and tool policy validation must go through these checks.
"""
from __future__ import annotations
from typing import Optional
from pydantic import BaseModel, Field, field_validator
from src.core.tool_context import PageType, EntityType, PageContext
from src.observability import get_logger
logger = get_logger(__name__)
class ValidationError(Exception):
"""Raised when context validation fails."""
pass
# Page type groupings for validation
TABLE_PAGES = {
PageType.INVOICE_TABLE,
PageType.CONTRACT_TABLE,
PageType.VENDOR_TABLE
}
DETAIL_PAGES = {
PageType.INVOICE_DETAIL,
PageType.CONTRACT_DETAIL,
PageType.VENDOR_DETAIL
}
CREATE_PAGES = {
PageType.INVOICE_CREATE,
PageType.CONTRACT_CREATE,
PageType.VENDOR_CREATE
}
# Valid entity types per page family
PAGE_ENTITY_MAPPING = {
"invoice": {
"table": PageType.INVOICE_TABLE,
"detail": PageType.INVOICE_DETAIL,
"create": PageType.INVOICE_CREATE,
"entity": EntityType.INVOICE
},
"contract": {
"table": PageType.CONTRACT_TABLE,
"detail": PageType.CONTRACT_DETAIL,
"create": PageType.CONTRACT_CREATE,
"entity": EntityType.CONTRACT
},
"vendor": {
"table": PageType.VENDOR_TABLE,
"detail": PageType.VENDOR_DETAIL,
"create": PageType.VENDOR_CREATE,
"entity": EntityType.VENDOR
}
}
class ValidatedPageContext(PageContext):
"""
Strictly validated page context.
This wraps the basic PageContext with additional validation rules
to ensure context consistency.
"""
page_type: PageType
entity_type: Optional[EntityType] = None
entity_id: Optional[str] = None
@field_validator("entity_type")
@classmethod
def validate_entity_type(cls, v: Optional[EntityType], info) -> Optional[EntityType]:
"""Validate entity type matches page type."""
page_type = info.data.get("page_type")
if not page_type:
raise ValidationError("page_type is required")
# Table pages should not have entities
if page_type in TABLE_PAGES and v is not None:
raise ValidationError("Table pages cannot have entities")
# Detail pages must have matching entity type
if page_type in DETAIL_PAGES:
if v is None:
raise ValidationError("Detail pages require an entity type")
# Check entity type matches page family
for family in PAGE_ENTITY_MAPPING.values():
if page_type == family["detail"] and v != family["entity"]:
raise ValidationError(
f"Invalid entity type {v} for page {page_type}"
)
# Create pages should not have entities
if page_type in CREATE_PAGES and v is not None:
raise ValidationError("Create pages cannot have entities")
return v
@field_validator("entity_id")
@classmethod
def validate_entity_id(cls, v: Optional[str], info) -> Optional[str]:
"""Validate entity ID presence matches requirements."""
entity_type = info.data.get("entity_type")
if entity_type and not v:
raise ValidationError("entity_id required when entity_type is set")
if not entity_type and v:
raise ValidationError("entity_id not allowed without entity_type")
return v
@classmethod
def from_page_context(cls, context: PageContext) -> ValidatedPageContext:
"""Create validated context from basic context."""
try:
return cls(
page_type=context.page_type,
entity_type=context.entity_type,
entity_id=context.entity_id
)
except ValidationError as e:
logger.error(
"page_context_validation_failed",
error=str(e),
context=context.dict()
)
raise
def to_page_context(self) -> PageContext:
"""Convert back to basic PageContext."""
return PageContext(
page_type=self.page_type,
entity_type=self.entity_type,
entity_id=self.entity_id
)
