"""
TTG Scratchpad MCP Server - Pydantic Models
Input validation and sanitization for all tool inputs.
"""
import re
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field, field_validator, model_validator
# ============================================================
# CONSTANTS
# ============================================================
MAX_TASK_LENGTH = 1000
MAX_WORKING_ON_LENGTH = 500
MAX_SUMMARY_LENGTH = 2000
MAX_FILE_PATH_LENGTH = 255
MAX_FILE_CONTENT_SIZE = 1024 * 1024 # 1MB
MAX_ACTIVITY_MESSAGE_LENGTH = 500
MAX_STEPS = 1000
# Characters allowed in file paths (prevent injection)
VALID_PATH_PATTERN = re.compile(r'^[a-zA-Z0-9_\-./]+$')
# Dangerous patterns to block
DANGEROUS_PATTERNS = [
'<script',
'javascript:',
'data:',
'onclick=',
'onerror=',
'../', # Path traversal
'..\\',
]
# ============================================================
# INPUT MODELS
# ============================================================
class TaskInput(BaseModel):
"""Validated input for scratchpad_start."""
task_description: str = Field(
...,
min_length=1,
max_length=MAX_TASK_LENGTH,
description="Description of the task to begin"
)
@field_validator('task_description')
@classmethod
def sanitize_task(cls, v: str) -> str:
"""Sanitize task description to prevent XSS."""
v = v.strip()
# HTML entity encode dangerous characters
v = v.replace('<', '<').replace('>', '>')
v = v.replace('"', '"').replace("'", ''')
return v
class UpdateInput(BaseModel):
"""Validated input for scratchpad_update."""
task_description: str = Field(
...,
min_length=1,
max_length=MAX_TASK_LENGTH
)
current_step: int = Field(
...,
ge=0,
le=MAX_STEPS,
description="Current step number"
)
total_steps: int = Field(
...,
ge=1,
le=MAX_STEPS,
description="Total number of steps"
)
working_on: str = Field(
...,
min_length=1,
max_length=MAX_WORKING_ON_LENGTH,
description="What is currently being processed"
)
files: Optional[str] = Field(
None,
max_length=10000,
description="Optional JSON string of files array"
)
activity_message: Optional[str] = Field(
None,
max_length=MAX_ACTIVITY_MESSAGE_LENGTH,
description="Optional new activity to log"
)
@field_validator('task_description', 'working_on')
@classmethod
def sanitize_text(cls, v: str) -> str:
"""Sanitize text fields."""
v = v.strip()
v = v.replace('<', '<').replace('>', '>')
return v
@model_validator(mode='after')
def validate_steps(self) -> 'UpdateInput':
"""Ensure current_step doesn't exceed total_steps."""
if self.current_step > self.total_steps:
raise ValueError('current_step cannot exceed total_steps')
return self
class CompleteInput(BaseModel):
"""Validated input for scratchpad_complete."""
task_description: str = Field(
...,
min_length=1,
max_length=MAX_TASK_LENGTH
)
summary: str = Field(
...,
min_length=1,
max_length=MAX_SUMMARY_LENGTH,
description="Summary of what was accomplished"
)
files: Optional[str] = Field(
None,
max_length=10000,
description="Optional JSON string of final files array"
)
@field_validator('task_description', 'summary')
@classmethod
def sanitize_text(cls, v: str) -> str:
"""Sanitize text fields."""
v = v.strip()
v = v.replace('<', '<').replace('>', '>')
return v
class FilePathInput(BaseModel):
"""Validated file path input."""
path: str = Field(
...,
min_length=1,
max_length=MAX_FILE_PATH_LENGTH,
description="File path (e.g., 'analysis/AAPL.md')"
)
@field_validator('path')
@classmethod
def validate_path(cls, v: str) -> str:
"""Validate and sanitize file path."""
v = v.strip()
# Check for dangerous patterns
v_lower = v.lower()
for pattern in DANGEROUS_PATTERNS:
if pattern.lower() in v_lower:
raise ValueError(f'Invalid path: contains forbidden pattern')
# Validate path characters
if not VALID_PATH_PATTERN.match(v):
raise ValueError(
'Invalid path: only alphanumeric, underscore, hyphen, dot, '
'and forward slash characters allowed'
)
# Normalize slashes
v = v.replace('\\', '/')
# Remove leading/trailing slashes
v = v.strip('/')
# Prevent empty path after normalization
if not v:
raise ValueError('Path cannot be empty')
return v
class WriteFileInput(FilePathInput):
"""Validated input for scratchpad_write_file."""
content: str = Field(
...,
max_length=MAX_FILE_CONTENT_SIZE,
description="File content to write"
)
@field_validator('content')
@classmethod
def validate_content_size(cls, v: str) -> str:
"""Validate content is not too large."""
if len(v.encode('utf-8')) > MAX_FILE_CONTENT_SIZE:
raise ValueError(f'File content exceeds maximum size of {MAX_FILE_CONTENT_SIZE} bytes')
return v
class ReadFileInput(FilePathInput):
"""Validated input for scratchpad_read_file."""
pass
class DeleteFileInput(FilePathInput):
"""Validated input for scratchpad_delete_file."""
pass
# ============================================================
# RESPONSE MODELS
# ============================================================
class ProgressModel(BaseModel):
"""Progress tracking model."""
current: int = Field(..., ge=0)
total: int = Field(..., ge=1)
class FileModel(BaseModel):
"""File representation for responses."""
name: str
type: str = "file" # "file" or "folder"
path: Optional[str] = None
updated: Optional[bool] = None
children: Optional[int] = None # For folders
class ActivityModel(BaseModel):
"""Activity entry for responses."""
action: str
icon: str = "write" # write, complete, read, delete
timestamp: Optional[str] = None
class ScratchpadResponse(BaseModel):
"""Standardized scratchpad response."""
task: str
status: str # idle, active, complete
files: List[FileModel] = []
activity: List[ActivityModel] = []
progress: Optional[ProgressModel] = None
workingOn: Optional[str] = None
def to_json(self) -> str:
"""Convert to JSON string for tool response."""
return self.model_dump_json(exclude_none=True)
class ErrorResponse(BaseModel):
"""Error response model."""
error: str
status: str = "error"
task: str = "Error"
files: List[FileModel] = []
activity: List[ActivityModel] = []
def to_json(self) -> str:
"""Convert to JSON string for tool response."""
return self.model_dump_json()
# ============================================================
# HELPER FUNCTIONS
# ============================================================
def validate_task_input(task_description: str) -> TaskInput:
"""Validate task input, raising ValueError on failure."""
return TaskInput(task_description=task_description)
def validate_update_input(
task_description: str,
current_step: int,
total_steps: int,
working_on: str,
files: Optional[str] = None,
activity_message: Optional[str] = None
) -> UpdateInput:
"""Validate update input, raising ValueError on failure."""
return UpdateInput(
task_description=task_description,
current_step=current_step,
total_steps=total_steps,
working_on=working_on,
files=files,
activity_message=activity_message
)
def validate_complete_input(
task_description: str,
summary: str,
files: Optional[str] = None
) -> CompleteInput:
"""Validate complete input, raising ValueError on failure."""
return CompleteInput(
task_description=task_description,
summary=summary,
files=files
)
def validate_file_path(path: str) -> str:
"""Validate file path, returning sanitized path."""
validated = FilePathInput(path=path)
return validated.path
def validate_write_file_input(path: str, content: str) -> WriteFileInput:
"""Validate write file input."""
return WriteFileInput(path=path, content=content)
def create_error_response(message: str) -> str:
"""Create a standardized error response JSON."""
return ErrorResponse(
error=message,
activity=[ActivityModel(action=message, icon="delete")]
).to_json()
