RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Harden schema sanitization and prompt cache tests

This commit is contained in:
adavyas 2026-03-27 15:08:59 -07:00
parent 2c4d88c316
commit aa326c8870
4 changed files with 203 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

122
src/schemas.py
View File

@ -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

21
tests/routes/test_validation_api.py
View File

@ -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]

41
tests/test_schema_validations.py
View File

@ -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"

36
tests/utils/test_prompt_cache_layouts.py
View File

@ -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<rolling_history>\nrolling history\n</rolling_history>"
),
}
@ -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"},
]