diff --git a/src/schemas.py b/src/schemas.py index 55685e64..de3acdba 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -8,18 +8,84 @@ import tiktoken from pydantic import ( AliasChoices, BaseModel, + BeforeValidator, ConfigDict, Field, PrivateAttr, field_validator, model_validator, ) +from pydantic_core import PydanticCustomError from src.config import ReasoningLevel, settings from src.utils.types import DocumentLevel RESOURCE_NAME_PATTERN = r"^[a-zA-Z0-9_-]+$" +_METADATA_MAX_KEYS = 100 +_METADATA_MAX_DEPTH = 5 + + +def strip_nul_bytes(value: Any) -> Any: + """Strip NUL bytes from string inputs without touching other types.""" + if isinstance(value, str): + return value.replace("\x00", "") + return value + + +def _sanitize_value(v: Any) -> Any: + """Recursively strip NUL bytes from strings in nested data structures.""" + if isinstance(v, str): + return strip_nul_bytes(v) + if isinstance(v, dict): + data = cast(dict[str, Any], v) + return {_sanitize_value(k): _sanitize_value(val) for k, val in data.items()} + if isinstance(v, list): + items = cast(list[Any], v) + return [_sanitize_value(item) for item in items] + return v + + +def _check_metadata_limits(value: Any, *, _current_depth: int = 1) -> None: + """Validate metadata doesn't exceed key count or nesting depth limits.""" + if _current_depth > _METADATA_MAX_DEPTH: + raise ValueError( + f"Metadata nesting exceeds maximum depth of {_METADATA_MAX_DEPTH}" + ) + + if isinstance(value, dict): + data = cast(dict[str, Any], value) + if _current_depth == 1 and len(data) > _METADATA_MAX_KEYS: + raise ValueError( + f"Metadata exceeds maximum of {_METADATA_MAX_KEYS} top-level keys" + ) + for item in data.values(): + if isinstance(item, (dict, list)): + _check_metadata_limits(item, _current_depth=_current_depth + 1) + return + + if isinstance(value, list): + items = cast(list[Any], value) + for item in items: + if isinstance(item, (dict, list)): + _check_metadata_limits(item, _current_depth=_current_depth + 1) + return + + if _current_depth == 1: + raise ValueError("Metadata must be a dict") + + +def _validate_metadata(v: Any) -> Any: + """Validate and sanitize a metadata dict before field parsing.""" + if not isinstance(v, dict): + return v + data = cast(dict[str, Any], v) + _check_metadata_limits(data) + return _sanitize_value(data) + + +_SanitizedMetadata = Annotated[dict[str, Any], BeforeValidator(_validate_metadata)] + class DreamType(str, Enum): """Types of dreams that can be triggered.""" @@ -213,7 +279,7 @@ class WorkspaceCreate(WorkspaceBase): str, Field(alias="id", min_length=1, max_length=100, pattern=RESOURCE_NAME_PATTERN), ] - metadata: dict[str, Any] = {} + metadata: _SanitizedMetadata = {} configuration: WorkspaceConfiguration = Field( default_factory=WorkspaceConfiguration ) @@ -226,7 +292,7 @@ class WorkspaceGet(WorkspaceBase): class WorkspaceUpdate(WorkspaceBase): - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None configuration: WorkspaceConfiguration | None = None @@ -252,7 +318,7 @@ class PeerCreate(PeerBase): str, Field(alias="id", min_length=1, max_length=100, pattern=RESOURCE_NAME_PATTERN), ] - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None configuration: dict[str, Any] | None = None model_config = ConfigDict(populate_by_name=True) # pyright: ignore @@ -263,7 +329,7 @@ class PeerGet(PeerBase): class PeerUpdate(PeerBase): - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None configuration: dict[str, Any] | None = None @@ -338,7 +404,7 @@ class MessageBase(BaseModel): class MessageCreate(MessageBase): content: Annotated[str, Field(min_length=0, max_length=settings.MAX_MESSAGE_SIZE)] peer_name: str = Field(alias="peer_id") - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None configuration: MessageConfiguration | None = None created_at: datetime.datetime | None = None @@ -362,7 +428,7 @@ class MessageGet(MessageBase): class MessageUpdate(MessageBase): - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None class Message(MessageBase): @@ -392,7 +458,7 @@ class MessageUploadCreate(BaseModel): """Schema for message creation from file uploads""" peer_id: str = Field(..., description="ID of the peer creating the message") - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None configuration: MessageConfiguration | None = None created_at: datetime.datetime | None = None @@ -408,7 +474,7 @@ class SessionCreate(SessionBase): str, Field(alias="id", min_length=1, max_length=100, pattern=RESOURCE_NAME_PATTERN), ] - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None peer_names: dict[str, SessionPeerConfig] | None = Field(default=None, alias="peers") configuration: SessionConfiguration | None = None @@ -420,7 +486,7 @@ class SessionGet(SessionBase): class SessionUpdate(SessionBase): - metadata: dict[str, Any] | None = None + metadata: _SanitizedMetadata | None = None configuration: SessionConfiguration | None = None @@ -576,6 +642,17 @@ class ObservationInput(BaseModel): ) = None confidence: Literal["high", "medium", "low"] | None = None + @field_validator("content", mode="after") + @classmethod + def sanitize_content(cls, v: str) -> str: + sanitized = cast(str, strip_nul_bytes(v)) + if not sanitized: + raise PydanticCustomError( + "string_too_short", + "String should have at least 1 character", + ) + return sanitized + @model_validator(mode="after") def validate_level_fields(self) -> Self: """Validate that level-specific fields are present when required.""" @@ -659,6 +736,17 @@ class ConclusionCreate(BaseModel): _token_count: int = PrivateAttr(default=0) + @field_validator("content", mode="after") + @classmethod + def sanitize_content(cls, v: str) -> str: + sanitized = cast(str, strip_nul_bytes(v)) + if not sanitized: + raise PydanticCustomError( + "string_too_short", + "String should have at least 1 character", + ) + return sanitized + @model_validator(mode="after") def validate_token_count(self) -> Self: """Validate that content doesn't exceed embedding token limit.""" @@ -697,6 +785,11 @@ class MessageSearchOptions(BaseModel): description="Number of results to return", ) + @field_validator("query", mode="after") + @classmethod + def sanitize_query(cls, v: str) -> str: + return cast(str, strip_nul_bytes(v)) + class DialecticOptions(BaseModel): session_id: str | None = Field( @@ -715,6 +808,17 @@ class DialecticOptions(BaseModel): description="Level of reasoning to apply: minimal, low, medium, high, or max", ) + @field_validator("query", mode="after") + @classmethod + def sanitize_query(cls, v: str) -> str: + sanitized = cast(str, strip_nul_bytes(v)) + if not sanitized: + raise PydanticCustomError( + "string_too_short", + "String should have at least 1 character", + ) + return sanitized + class DialecticResponse(BaseModel): content: str | None diff --git a/tests/routes/test_validation_api.py b/tests/routes/test_validation_api.py index 38bc115f..fe3d1aa2 100644 --- a/tests/routes/test_validation_api.py +++ b/tests/routes/test_validation_api.py @@ -34,6 +34,16 @@ def test_workspace_validations_api(client: TestClient): assert error["loc"] == ["body", "metadata"] assert error["type"] == "dict_type" + # Test deeply nested list metadata + response = client.post( + "/v3/workspaces", + json={"name": "test-nested", "metadata": {"a": [[[[["too deep"]]]]]}}, + ) + assert response.status_code == 422 + error = response.json()["detail"][0] + assert error["loc"] == ["body", "metadata"] + assert "Metadata nesting exceeds maximum depth" in error["msg"] + def test_peer_validations_api(client: TestClient, sample_data: tuple[Workspace, Peer]): test_workspace, _ = sample_data @@ -189,6 +199,17 @@ def test_agent_query_validations_api( ) assert response.status_code == 200 + # Test NUL-only query is rejected after sanitization + response = client.post( + f"/v3/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat", + params={"session_id": session_id, "target": "test_target"}, + json={"query": "\x00", "stream": False}, + ) + assert response.status_code == 422 + error = response.json()["detail"][0] + assert error["loc"] == ["body", "query"] + assert error["type"] == "string_too_short" + def test_required_field_validations_api( client: TestClient, sample_data: tuple[Workspace, Peer] diff --git a/tests/test_schema_validations.py b/tests/test_schema_validations.py index 281522a9..c667b99c 100644 --- a/tests/test_schema_validations.py +++ b/tests/test_schema_validations.py @@ -4,9 +4,12 @@ import pytest from pydantic import ValidationError from src.schemas import ( + ConclusionCreate, + DialecticOptions, DocumentCreate, DocumentMetadata, MessageCreate, + ObservationInput, PeerCreate, ResolvedConfiguration, SessionCreate, @@ -41,6 +44,14 @@ class TestWorkspaceValidations: error_dict = exc_info.value.errors()[0] assert error_dict["type"] == "dict_type" + def test_app_metadata_rejects_deeply_nested_lists(self): + nested_metadata = {"a": [[[[["too deep"]]]]]} + + with pytest.raises(ValidationError) as exc_info: + WorkspaceCreate(name="test", metadata=nested_metadata) + + assert "Metadata nesting exceeds maximum depth" in str(exc_info.value) + class TestPeerValidations: def test_valid_peer_create(self): @@ -203,3 +214,33 @@ class TestResolvedConfigurationMigration: ResolvedConfiguration.model_validate(payload) assert any(e["loc"] == ("reasoning",) for e in exc_info.value.errors()) + + +class TestSanitizedRequiredFields: + def test_conclusion_content_rejects_nul_only_input(self): + with pytest.raises(ValidationError) as exc_info: + ConclusionCreate( + content="\x00", + observer_id="observer", + observed_id="observed", + ) + + error_dict = exc_info.value.errors()[0] + assert error_dict["loc"] == ("content",) + assert error_dict["type"] == "string_too_short" + + def test_dialectic_query_rejects_nul_only_input(self): + with pytest.raises(ValidationError) as exc_info: + DialecticOptions.model_validate({"query": "\x00"}) + + error_dict = exc_info.value.errors()[0] + assert error_dict["loc"] == ("query",) + assert error_dict["type"] == "string_too_short" + + def test_observation_content_rejects_nul_only_input(self): + with pytest.raises(ValidationError) as exc_info: + ObservationInput(content="\x00") + + error_dict = exc_info.value.errors()[0] + assert error_dict["loc"] == ("content",) + assert error_dict["type"] == "string_too_short" diff --git a/tests/utils/test_prompt_cache_layouts.py b/tests/utils/test_prompt_cache_layouts.py index c76e4cb6..7091423e 100644 --- a/tests/utils/test_prompt_cache_layouts.py +++ b/tests/utils/test_prompt_cache_layouts.py @@ -6,6 +6,26 @@ from src.utils.prompt_cache_layouts import ( merge_system_prompt_with_rolling_context, provider_default_prompt_cache_layout, ) +from src.utils.types import SupportedProviders + + +class DialecticAgentHarness(DialecticAgent): + _provider: SupportedProviders + + @property + def provider(self) -> SupportedProviders: + return self._provider + + @provider.setter + def provider(self, value: SupportedProviders) -> None: + self._provider = value + + @property + def base_system_prompt(self) -> str: + return self._base_system_prompt + + def set_system_messages(self, session_history_section: str | None = None) -> None: + self._set_system_messages(session_history_section) def test_provider_default_prompt_cache_layout_is_google_specific() -> None: @@ -61,7 +81,7 @@ def test_merge_system_prompt_with_rolling_context_strips_noise() -> None: def test_dialectic_agent_rebuilds_google_system_messages() -> None: - agent = DialecticAgent( + agent = DialecticAgentHarness( db=AsyncMock(), workspace_name="workspace", session_name="session", @@ -70,14 +90,14 @@ def test_dialectic_agent_rebuilds_google_system_messages() -> None: reasoning_level="low", ) - agent._provider = "google" - agent._set_system_messages("rolling history") + agent.provider = "google" + agent.set_system_messages("rolling history") assert agent.messages == [ { "role": "system", "content": ( - agent._base_system_prompt.strip() + agent.base_system_prompt.strip() + "\n\n\nrolling history\n" ), } @@ -85,7 +105,7 @@ def test_dialectic_agent_rebuilds_google_system_messages() -> None: def test_dialectic_agent_rebuilds_anthropic_system_messages() -> None: - agent = DialecticAgent( + agent = DialecticAgentHarness( db=AsyncMock(), workspace_name="workspace", session_name="session", @@ -94,10 +114,10 @@ def test_dialectic_agent_rebuilds_anthropic_system_messages() -> None: reasoning_level="medium", ) - agent._provider = "anthropic" - agent._set_system_messages("rolling history") + agent.provider = "anthropic" + agent.set_system_messages("rolling history") assert agent.messages == [ - {"role": "system", "content": agent._base_system_prompt.strip()}, + {"role": "system", "content": agent.base_system_prompt.strip()}, {"role": "system", "content": "rolling history"}, ]