server/models.py

from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, ConfigDict, Field


class DefenseActionType(str, Enum):
    ALLOW = "ALLOW"
    BLOCK = "BLOCK"
    APPLY_SPOTLIGHTING = "APPLY_SPOTLIGHTING"
    SEMANTIC_DIFF = "SEMANTIC_DIFF"
    CAPABILITY_MEDIATION = "CAPABILITY_MEDIATION"


class MCPToolContext(BaseModel):
    tool_name: str
    requested_parameters: Dict[str, Any]
    target_endpoint: str
    metadata: Dict[str, Any] = Field(default_factory=dict)
    advertised_capabilities: List[str] = Field(default_factory=list)
    token_audience: Optional[str] = None
    tool_version: Optional[str] = None


class AuthContext(BaseModel):
    token_subject: str
    token_audience: str
    requested_audience: Optional[str] = None
    scopes: List[str] = Field(default_factory=list)


class EcosystemObservation(BaseModel):
    model_config = ConfigDict(extra="forbid")

    task_id: str
    step_id: int
    incoming_user_prompt: str
    mcp_tool_request: Optional[MCPToolContext]
    system_context: str
    anomaly_hints: List[str] = Field(default_factory=list)
    historical_baseline: Dict[str, Any] = Field(default_factory=dict)
    auth_context: Optional[AuthContext] = None


class DefenseAction(BaseModel):
    action_type: DefenseActionType = Field(
        ...,
        description=(
            "One of: ALLOW, BLOCK, APPLY_SPOTLIGHTING, SEMANTIC_DIFF, CAPABILITY_MEDIATION"
        ),
    )
    justification: str = Field(min_length=5, max_length=280)


class StepReward(BaseModel):
    security_score: float = Field(ge=0.0, le=1.0)
    usability_penalty: float = Field(ge=0.0, le=1.0)
    latency_penalty: float = Field(ge=0.0, le=1.0)
    total: float = Field(ge=0.0, le=1.0)
    verdict: str


class StepInfo(BaseModel):
    error: Optional[str] = None
    expected_action: DefenseActionType
    risk_level: float = Field(ge=0.0, le=1.0)
    compromised: bool = False
    reward_breakdown: StepReward


class StepResponse(BaseModel):
    observation: EcosystemObservation
    reward: float
    done: bool
    info: StepInfo


class ResetResponse(BaseModel):
    observation: EcosystemObservation