Sanitization and Memory Bug Fixes (#419)
* fix: use WeakValueDictionary for _observation_locks to prevent memory leak * fix: harden input sanitization across API surface (DEV-1400) - Parameterize SQL in set_config calls to prevent injection via request context - Strip NUL bytes from string inputs (message content, queries, peer cards) - Add JSONB metadata validation (100 key limit, 5 depth limit) - Add filter recursion depth limit (max 5) to prevent stack overflow - Update changelogs with unreleased entries Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: Refactor Schemas into separate files * fix: Code Rabbit Comments --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Vineeth Voruganti
GitHub
parent
122ebfe998
commit
09a980c2fb
13
CHANGELOG.md
13
CHANGELOG.md
|
|
@ -5,6 +5,19 @@ All notable changes to this project will be documented in this file.
|
|||
The format is based on [Keep a Changelog](http://keepachangelog.com/)
|
||||
and this project adheres to [Semantic Versioning](http://semver.org/).
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Fixed
|
||||
|
||||
- Memory leak in deriver: `_observation_locks` dict grew unboundedly with every unique (workspace, observer, observed) combination; replaced with `WeakValueDictionary` so locks are automatically evicted when no longer in use (DEV-1412)
|
||||
- SQL injection vector in `dependencies.py`: parameterized `SET application_name` queries using `set_config()` instead of f-string interpolation (DEV-1400)
|
||||
- NUL byte (`\x00`) crashes: all user-facing text inputs (message content, metadata, peer cards, queries) are now sanitized at the Pydantic schema level before reaching PostgreSQL (DEV-1400)
|
||||
|
||||
### Added
|
||||
|
||||
- JSONB metadata validation: max 100 top-level keys and max nesting depth of 5 on all metadata input fields (DEV-1400)
|
||||
- Filter recursion depth limit: `_build_filter_conditions()` now enforces a max depth of 5 to prevent stack overflow from deeply nested filter dicts (DEV-1400)
|
||||
|
||||
## [3.0.3] - 2026-02-25
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -27,6 +27,19 @@ Welcome to the Honcho changelog! This section documents all notable changes to t
|
|||
### Honcho API and SDK Changelogs
|
||||
<Tabs>
|
||||
<Tab title="Honcho API">
|
||||
<Update label="v3.0.4 (Unreleased)">
|
||||
### Fixed
|
||||
|
||||
- Memory leak in deriver: `_observation_locks` dict grew unboundedly with every unique (workspace, observer, observed) combination; replaced with `WeakValueDictionary` so locks are automatically evicted when no longer in use (DEV-1412)
|
||||
- SQL injection vector in `dependencies.py`: parameterized `SET application_name` queries using `set_config()` instead of f-string interpolation (DEV-1400)
|
||||
- NUL byte (`\x00`) crashes: all user-facing text inputs (message content, metadata, peer cards, queries) are now sanitized at the Pydantic schema level before reaching PostgreSQL (DEV-1400)
|
||||
|
||||
### Added
|
||||
|
||||
- JSONB metadata validation: max 100 top-level keys and max nesting depth of 5 on all metadata input fields (DEV-1400)
|
||||
- Filter recursion depth limit: `_build_filter_conditions()` now enforces a max depth of 5 to prevent stack overflow from deeply nested filter dicts (DEV-1400)
|
||||
</Update>
|
||||
|
||||
<Update label="v3.0.3 (Current)">
|
||||
### Added
|
||||
|
||||
|
|
|
|||
|
|
@ -17,7 +17,10 @@ async def get_db():
|
|||
db: AsyncSession = SessionLocal()
|
||||
try:
|
||||
if settings.DB.TRACING:
|
||||
await db.execute(text(f"SET application_name = '{context}'"))
|
||||
await db.execute(
|
||||
text("SELECT set_config('application_name', :name, false)"),
|
||||
{"name": context},
|
||||
)
|
||||
yield db
|
||||
except Exception:
|
||||
await db.rollback()
|
||||
|
|
@ -45,7 +48,8 @@ async def tracked_db(operation_name: str | None = None):
|
|||
try:
|
||||
if settings.DB.TRACING:
|
||||
await db.execute(
|
||||
text(f"SET application_name = '{context or f'task:{operation_name}'}'")
|
||||
text("SELECT set_config('application_name', :name, false)"),
|
||||
{"name": context or f"task:{operation_name}"},
|
||||
)
|
||||
|
||||
yield db
|
||||
|
|
|
|||
|
|
@ -0,0 +1,162 @@
|
|||
"""Pydantic schemas for Honcho.
|
||||
|
||||
Re-exports all public names from submodules so that existing
|
||||
``from src.schemas import X`` imports continue to work unchanged.
|
||||
"""
|
||||
|
||||
from src.schemas.api import (
|
||||
RESOURCE_NAME_PATTERN,
|
||||
Conclusion,
|
||||
ConclusionBatchCreate,
|
||||
ConclusionCreate,
|
||||
ConclusionGet,
|
||||
ConclusionQuery,
|
||||
DialecticOptions,
|
||||
DialecticResponse,
|
||||
DialecticStreamChunk,
|
||||
DialecticStreamDelta,
|
||||
Message,
|
||||
MessageBase,
|
||||
MessageBatchCreate,
|
||||
MessageCreate,
|
||||
MessageGet,
|
||||
MessageSearchOptions,
|
||||
MessageUpdate,
|
||||
MessageUploadCreate,
|
||||
Peer,
|
||||
PeerBase,
|
||||
PeerCardResponse,
|
||||
PeerCardSet,
|
||||
PeerContext,
|
||||
PeerCreate,
|
||||
PeerGet,
|
||||
PeerRepresentationGet,
|
||||
PeerUpdate,
|
||||
QueueStatus,
|
||||
RepresentationResponse,
|
||||
ScheduleDreamRequest,
|
||||
Session,
|
||||
SessionBase,
|
||||
SessionContext,
|
||||
SessionCreate,
|
||||
SessionGet,
|
||||
SessionQueueStatus,
|
||||
SessionSummaries,
|
||||
SessionUpdate,
|
||||
Summary,
|
||||
WebhookEndpoint,
|
||||
WebhookEndpointBase,
|
||||
WebhookEndpointCreate,
|
||||
Workspace,
|
||||
WorkspaceBase,
|
||||
WorkspaceCreate,
|
||||
WorkspaceGet,
|
||||
WorkspaceUpdate,
|
||||
)
|
||||
from src.schemas.configuration import (
|
||||
DreamConfiguration,
|
||||
DreamType,
|
||||
MessageConfiguration,
|
||||
PeerCardConfiguration,
|
||||
PeerConfig,
|
||||
ReasoningConfiguration,
|
||||
ResolvedConfiguration,
|
||||
ResolvedDreamConfiguration,
|
||||
ResolvedPeerCardConfiguration,
|
||||
ResolvedReasoningConfiguration,
|
||||
ResolvedSummaryConfiguration,
|
||||
SessionConfiguration,
|
||||
SessionPeerConfig,
|
||||
SummaryConfiguration,
|
||||
WorkspaceConfiguration,
|
||||
)
|
||||
from src.schemas.internal import (
|
||||
DocumentBase,
|
||||
DocumentCreate,
|
||||
DocumentMetadata,
|
||||
MessageBulkData,
|
||||
ObservationInput,
|
||||
QueueCounts,
|
||||
QueueStatusRow,
|
||||
ReconcilerType,
|
||||
SessionCounts,
|
||||
SessionPeerData,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# configuration
|
||||
"DreamConfiguration",
|
||||
"DreamType",
|
||||
"MessageConfiguration",
|
||||
"PeerCardConfiguration",
|
||||
"PeerConfig",
|
||||
"ReasoningConfiguration",
|
||||
"ResolvedConfiguration",
|
||||
"ResolvedDreamConfiguration",
|
||||
"ResolvedPeerCardConfiguration",
|
||||
"ResolvedReasoningConfiguration",
|
||||
"ResolvedSummaryConfiguration",
|
||||
"SessionConfiguration",
|
||||
"SessionPeerConfig",
|
||||
"SummaryConfiguration",
|
||||
"WorkspaceConfiguration",
|
||||
# api
|
||||
"Conclusion",
|
||||
"ConclusionBatchCreate",
|
||||
"ConclusionCreate",
|
||||
"ConclusionGet",
|
||||
"ConclusionQuery",
|
||||
"DialecticOptions",
|
||||
"DialecticResponse",
|
||||
"DialecticStreamChunk",
|
||||
"DialecticStreamDelta",
|
||||
"Message",
|
||||
"MessageBase",
|
||||
"MessageBatchCreate",
|
||||
"MessageCreate",
|
||||
"MessageGet",
|
||||
"MessageSearchOptions",
|
||||
"MessageUpdate",
|
||||
"MessageUploadCreate",
|
||||
"Peer",
|
||||
"PeerBase",
|
||||
"PeerCardResponse",
|
||||
"PeerCardSet",
|
||||
"PeerContext",
|
||||
"PeerCreate",
|
||||
"PeerGet",
|
||||
"PeerRepresentationGet",
|
||||
"PeerUpdate",
|
||||
"QueueStatus",
|
||||
"RESOURCE_NAME_PATTERN",
|
||||
"RepresentationResponse",
|
||||
"ScheduleDreamRequest",
|
||||
"Session",
|
||||
"SessionBase",
|
||||
"SessionContext",
|
||||
"SessionCreate",
|
||||
"SessionGet",
|
||||
"SessionQueueStatus",
|
||||
"SessionSummaries",
|
||||
"SessionUpdate",
|
||||
"Summary",
|
||||
"WebhookEndpoint",
|
||||
"WebhookEndpointBase",
|
||||
"WebhookEndpointCreate",
|
||||
"Workspace",
|
||||
"WorkspaceBase",
|
||||
"WorkspaceCreate",
|
||||
"WorkspaceGet",
|
||||
"WorkspaceUpdate",
|
||||
# internal
|
||||
"DocumentBase",
|
||||
"DocumentCreate",
|
||||
"DocumentMetadata",
|
||||
"MessageBulkData",
|
||||
"ObservationInput",
|
||||
"QueueCounts",
|
||||
"QueueStatusRow",
|
||||
"ReconcilerType",
|
||||
"SessionCounts",
|
||||
"SessionPeerData",
|
||||
]
|
||||
|
|
@ -1,13 +1,19 @@
|
|||
"""Pydantic schemas for API request/response validation.
|
||||
|
||||
These schemas are consumed by the FastAPI routers and define the public
|
||||
API contract.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import ipaddress
|
||||
from enum import Enum
|
||||
from typing import Annotated, Any, Literal, Self, cast
|
||||
from typing import Annotated, Any, Self, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import tiktoken
|
||||
from pydantic import (
|
||||
AliasChoices,
|
||||
BaseModel,
|
||||
BeforeValidator,
|
||||
ConfigDict,
|
||||
Field,
|
||||
PrivateAttr,
|
||||
|
|
@ -16,192 +22,72 @@ from pydantic import (
|
|||
)
|
||||
|
||||
from src.config import ReasoningLevel, settings
|
||||
from src.utils.types import DocumentLevel
|
||||
from src.schemas.configuration import (
|
||||
DreamType,
|
||||
MessageConfiguration,
|
||||
SessionConfiguration,
|
||||
SessionPeerConfig,
|
||||
WorkspaceConfiguration,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Metadata validation helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
RESOURCE_NAME_PATTERN = r"^[a-zA-Z0-9_-]+$"
|
||||
|
||||
|
||||
class DreamType(str, Enum):
|
||||
"""Types of dreams that can be triggered."""
|
||||
|
||||
OMNI = "omni"
|
||||
_METADATA_MAX_KEYS = 100
|
||||
_METADATA_MAX_DEPTH = 5
|
||||
|
||||
|
||||
class ReconcilerType(str, Enum):
|
||||
"""Types of reconciler tasks that can be performed."""
|
||||
|
||||
SYNC_VECTORS = "sync_vectors"
|
||||
CLEANUP_QUEUE = "cleanup_queue"
|
||||
def _sanitize_value(v: Any) -> Any:
|
||||
"""Recursively strip NUL bytes from strings in nested data structures."""
|
||||
if isinstance(v, str):
|
||||
return v.replace("\x00", "")
|
||||
if isinstance(v, dict):
|
||||
d = cast(dict[str, Any], v)
|
||||
return {_sanitize_value(k): _sanitize_value(val) for k, val in d.items()}
|
||||
if isinstance(v, list):
|
||||
lst = cast(list[Any], v)
|
||||
return [_sanitize_value(item) for item in lst]
|
||||
return v
|
||||
|
||||
|
||||
class ReasoningConfiguration(BaseModel):
|
||||
enabled: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to enable reasoning functionality.",
|
||||
)
|
||||
custom_instructions: str | None = Field(
|
||||
default=None,
|
||||
description="TODO: currently unused. Custom instructions to use for the reasoning system on this workspace/session/message.",
|
||||
)
|
||||
|
||||
|
||||
class PeerCardConfiguration(BaseModel):
|
||||
use: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to use peer card related to this peer during reasoning process.",
|
||||
)
|
||||
create: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to generate peer card based on content.",
|
||||
)
|
||||
|
||||
|
||||
class SummaryConfiguration(BaseModel):
|
||||
enabled: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to enable summary functionality.",
|
||||
)
|
||||
messages_per_short_summary: int | None = Field(
|
||||
default=None,
|
||||
ge=10,
|
||||
description="Number of messages per short summary. Must be positive, greater than or equal to 10, and less than messages_per_long_summary.",
|
||||
)
|
||||
messages_per_long_summary: int | None = Field(
|
||||
default=None,
|
||||
ge=20,
|
||||
description="Number of messages per long summary. Must be positive, greater than or equal to 20, and greater than messages_per_short_summary.",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_summary_thresholds(self) -> Self:
|
||||
"""Validate that short summary threshold <= long summary threshold."""
|
||||
short = self.messages_per_short_summary
|
||||
long = self.messages_per_long_summary
|
||||
|
||||
if short is not None and long is not None and short >= long:
|
||||
raise ValueError(
|
||||
"messages_per_short_summary must be less than messages_per_long_summary"
|
||||
def _check_metadata_limits(
|
||||
data: dict[str, Any],
|
||||
*,
|
||||
_current_depth: int = 1,
|
||||
) -> None:
|
||||
"""Validate metadata dict 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 _current_depth == 1 and len(data) > _METADATA_MAX_KEYS:
|
||||
raise ValueError(
|
||||
f"Metadata exceeds maximum of {_METADATA_MAX_KEYS} top-level keys"
|
||||
)
|
||||
for v in data.values():
|
||||
if isinstance(v, dict):
|
||||
_check_metadata_limits(
|
||||
cast(dict[str, Any], v), _current_depth=_current_depth + 1
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
def _validate_metadata(v: Any) -> Any:
|
||||
"""Validate and sanitize a metadata dict: enforce limits and strip NUL bytes."""
|
||||
if not isinstance(v, dict):
|
||||
return v
|
||||
data = cast(dict[str, Any], v)
|
||||
_check_metadata_limits(data)
|
||||
return _sanitize_value(data)
|
||||
|
||||
|
||||
class DreamConfiguration(BaseModel):
|
||||
enabled: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to enable dream functionality. If reasoning is disabled, dreams will also be disabled and this setting will be ignored.",
|
||||
)
|
||||
_SanitizedMetadata = Annotated[dict[str, Any], BeforeValidator(_validate_metadata)]
|
||||
|
||||
|
||||
class WorkspaceConfiguration(BaseModel):
|
||||
"""
|
||||
The set of options that can be in a workspace DB-level configuration dictionary.
|
||||
|
||||
All fields are optional. Session-level configuration overrides workspace-level configuration, which overrides global configuration.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="allow") # pyright: ignore
|
||||
|
||||
reasoning: ReasoningConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for reasoning functionality.",
|
||||
)
|
||||
peer_card: PeerCardConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for peer card functionality. If reasoning is disabled, peer cards will also be disabled and these settings will be ignored.",
|
||||
)
|
||||
summary: SummaryConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for summary functionality.",
|
||||
)
|
||||
dream: DreamConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for dream functionality. If reasoning is disabled, dreams will also be disabled and these settings will be ignored.",
|
||||
)
|
||||
|
||||
|
||||
class SessionConfiguration(WorkspaceConfiguration):
|
||||
"""
|
||||
The set of options that can be in a session DB-level configuration dictionary.
|
||||
|
||||
All fields are optional. Session-level configuration overrides workspace-level configuration, which overrides global configuration.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class MessageConfiguration(BaseModel):
|
||||
"""
|
||||
The set of options that can be in a message DB-level configuration dictionary.
|
||||
|
||||
All fields are optional. Message-level configuration overrides all other configurations.
|
||||
"""
|
||||
|
||||
reasoning: ReasoningConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for reasoning functionality.",
|
||||
)
|
||||
|
||||
|
||||
class ResolvedReasoningConfiguration(BaseModel):
|
||||
enabled: bool
|
||||
|
||||
|
||||
class ResolvedPeerCardConfiguration(BaseModel):
|
||||
use: bool
|
||||
create: bool
|
||||
|
||||
|
||||
class ResolvedSummaryConfiguration(BaseModel):
|
||||
enabled: bool
|
||||
messages_per_short_summary: int
|
||||
messages_per_long_summary: int
|
||||
|
||||
|
||||
class ResolvedDreamConfiguration(BaseModel):
|
||||
enabled: bool
|
||||
|
||||
|
||||
class ResolvedConfiguration(BaseModel):
|
||||
"""
|
||||
The final resolved configuration for a given message.
|
||||
Hierarchy: message > session > workspace > global configuration
|
||||
"""
|
||||
|
||||
reasoning: ResolvedReasoningConfiguration
|
||||
peer_card: ResolvedPeerCardConfiguration
|
||||
summary: ResolvedSummaryConfiguration
|
||||
dream: ResolvedDreamConfiguration
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def migrate_deriver_to_reasoning(cls, data: Any) -> Any:
|
||||
"""Handle v3.0.0 migration: 'deriver' was renamed to 'reasoning'."""
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
|
||||
config = cast(dict[str, Any], data)
|
||||
|
||||
if "deriver" in config and "reasoning" not in config:
|
||||
config["reasoning"] = config.pop("deriver")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
class PeerConfig(BaseModel):
|
||||
# TODO: Update description - should say "Whether honcho forms a representation of the peer itself"
|
||||
observe_me: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether Honcho will use reasoning to form a representation of this peer",
|
||||
)
|
||||
|
||||
|
||||
class SessionPeerConfig(PeerConfig):
|
||||
# TODO: Update description - should say "Whether this peer forms representations of other peers in the session"
|
||||
observe_others: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether this peer should form a session-level theory-of-mind representation of other peers in the session",
|
||||
)
|
||||
# ---------------------------------------------------------------------------
|
||||
# Workspace schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WorkspaceBase(BaseModel):
|
||||
|
|
@ -213,7 +99,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 +112,7 @@ class WorkspaceGet(WorkspaceBase):
|
|||
|
||||
|
||||
class WorkspaceUpdate(WorkspaceBase):
|
||||
metadata: dict[str, Any] | None = None
|
||||
metadata: _SanitizedMetadata | None = None
|
||||
configuration: WorkspaceConfiguration | None = None
|
||||
|
||||
|
||||
|
|
@ -243,6 +129,11 @@ class Workspace(WorkspaceBase):
|
|||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Peer schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PeerBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
|
@ -252,7 +143,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 +154,7 @@ class PeerGet(PeerBase):
|
|||
|
||||
|
||||
class PeerUpdate(PeerBase):
|
||||
metadata: dict[str, Any] | None = None
|
||||
metadata: _SanitizedMetadata | None = None
|
||||
configuration: dict[str, Any] | None = None
|
||||
|
||||
|
||||
|
|
@ -330,6 +221,21 @@ class PeerCardResponse(BaseModel):
|
|||
class PeerCardSet(BaseModel):
|
||||
peer_card: list[str] = Field(..., description="The peer card content to set")
|
||||
|
||||
@field_validator("peer_card", mode="before")
|
||||
@classmethod
|
||||
def sanitize_peer_card(cls, v: Any) -> Any:
|
||||
if isinstance(v, list):
|
||||
return [
|
||||
item.replace("\x00", "") if isinstance(item, str) else item
|
||||
for item in cast(list[Any], v)
|
||||
]
|
||||
return v
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Message schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MessageBase(BaseModel):
|
||||
pass
|
||||
|
|
@ -338,12 +244,17 @@ 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
|
||||
|
||||
_encoded_message: list[int] = PrivateAttr(default=[])
|
||||
|
||||
@field_validator("content", mode="after")
|
||||
@classmethod
|
||||
def sanitize_content(cls, v: str) -> str:
|
||||
return v.replace("\x00", "")
|
||||
|
||||
@property
|
||||
def encoded_message(self) -> list[int]:
|
||||
return self._encoded_message
|
||||
|
|
@ -362,7 +273,7 @@ class MessageGet(MessageBase):
|
|||
|
||||
|
||||
class MessageUpdate(MessageBase):
|
||||
metadata: dict[str, Any] | None = None
|
||||
metadata: _SanitizedMetadata | None = None
|
||||
|
||||
|
||||
class Message(MessageBase):
|
||||
|
|
@ -392,13 +303,18 @@ 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
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True) # pyright: ignore
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SessionBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
|
@ -408,7 +324,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 +336,7 @@ class SessionGet(SessionBase):
|
|||
|
||||
|
||||
class SessionUpdate(SessionBase):
|
||||
metadata: dict[str, Any] | None = None
|
||||
metadata: _SanitizedMetadata | None = None
|
||||
configuration: SessionConfiguration | None = None
|
||||
|
||||
|
||||
|
|
@ -505,95 +421,9 @@ class SessionSummaries(SessionBase):
|
|||
)
|
||||
|
||||
|
||||
class DocumentBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class DocumentMetadata(BaseModel):
|
||||
message_ids: list[int] = Field(
|
||||
description="The ID range(s) of the messages that this document was derived from. Acts as a link to the primary source of the document. Note that as a document gets deduplicated, additional ranges will be added, because the same document could be derived from completely separate message ranges."
|
||||
)
|
||||
message_created_at: str = Field(
|
||||
description="The timestamp of the message that this document was derived from. Note that this is not the same as the created_at timestamp of the document. This timestamp is usually only saved with second-level precision."
|
||||
)
|
||||
source_ids: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Document IDs of source documents for tree traversal -- required for deductive and inductive documents",
|
||||
)
|
||||
premises: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Human-readable premise text for display -- only applicable for deductive documents",
|
||||
)
|
||||
sources: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Human-readable source text for display -- only applicable for inductive documents",
|
||||
)
|
||||
pattern_type: str | None = Field(
|
||||
default=None,
|
||||
description="Type of pattern identified (preference, behavior, personality, tendency, correlation) -- only applicable for inductive documents",
|
||||
)
|
||||
confidence: str | None = Field(
|
||||
default=None,
|
||||
description="Confidence level (high, medium, low) -- only applicable for inductive documents",
|
||||
)
|
||||
|
||||
|
||||
class DocumentCreate(DocumentBase):
|
||||
content: Annotated[str, Field(min_length=1, max_length=100000)]
|
||||
session_name: str | None = Field(
|
||||
default=None,
|
||||
description="The session from which the document was derived (NULL for global observations)",
|
||||
)
|
||||
level: DocumentLevel = Field(
|
||||
default="explicit",
|
||||
description="The level of the document (explicit, deductive, inductive, or contradiction)",
|
||||
)
|
||||
times_derived: int = Field(
|
||||
default=1,
|
||||
ge=1,
|
||||
description="The number of times that a semantic duplicate document to this one has been derived",
|
||||
)
|
||||
metadata: DocumentMetadata = Field()
|
||||
embedding: list[float] = Field()
|
||||
# Tree linkage field
|
||||
source_ids: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Document IDs of source/premise documents -- for deductive and inductive documents",
|
||||
)
|
||||
|
||||
|
||||
class ObservationInput(BaseModel):
|
||||
"""Validated observation input from LLM tool calls."""
|
||||
|
||||
content: str = Field(min_length=1)
|
||||
level: DocumentLevel = "explicit"
|
||||
source_ids: list[str] | None = None
|
||||
premises: list[str] | None = None
|
||||
sources: list[str] | None = None
|
||||
pattern_type: (
|
||||
Literal["preference", "behavior", "personality", "tendency", "correlation"]
|
||||
| None
|
||||
) = None
|
||||
confidence: Literal["high", "medium", "low"] | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_level_fields(self) -> Self:
|
||||
"""Validate that level-specific fields are present when required."""
|
||||
if self.level == "deductive" and not self.source_ids:
|
||||
raise ValueError(
|
||||
"deductive observations require 'source_ids' field with document IDs of premises"
|
||||
)
|
||||
if self.level == "inductive" and not self.source_ids:
|
||||
raise ValueError(
|
||||
"inductive observations require 'source_ids' field with document IDs of sources"
|
||||
)
|
||||
if self.level == "contradiction" and (
|
||||
not self.source_ids or len(self.source_ids) < 2
|
||||
):
|
||||
raise ValueError(
|
||||
"contradiction observations require 'source_ids' field with at least 2 IDs of contradicting observations"
|
||||
)
|
||||
return self
|
||||
# ---------------------------------------------------------------------------
|
||||
# Conclusion schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ConclusionGet(BaseModel):
|
||||
|
|
@ -659,6 +489,11 @@ class ConclusionCreate(BaseModel):
|
|||
|
||||
_token_count: int = PrivateAttr(default=0)
|
||||
|
||||
@field_validator("content", mode="after")
|
||||
@classmethod
|
||||
def sanitize_content(cls, v: str) -> str:
|
||||
return v.replace("\x00", "")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_token_count(self) -> Self:
|
||||
"""Validate that content doesn't exceed embedding token limit."""
|
||||
|
|
@ -685,8 +520,13 @@ class ConclusionBatchCreate(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Search schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MessageSearchOptions(BaseModel):
|
||||
query: str = Field(..., description="Search query")
|
||||
query: Annotated[str, Field(..., description="Search query")]
|
||||
filters: dict[str, Any] | None = Field(
|
||||
default=None, description="Filters to scope the search"
|
||||
)
|
||||
|
|
@ -697,6 +537,16 @@ class MessageSearchOptions(BaseModel):
|
|||
description="Number of results to return",
|
||||
)
|
||||
|
||||
@field_validator("query", mode="after")
|
||||
@classmethod
|
||||
def sanitize_query(cls, v: str) -> str:
|
||||
return v.replace("\x00", "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dialectic schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DialecticOptions(BaseModel):
|
||||
session_id: str | None = Field(
|
||||
|
|
@ -715,6 +565,11 @@ 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:
|
||||
return v.replace("\x00", "")
|
||||
|
||||
|
||||
class DialecticResponse(BaseModel):
|
||||
content: str | None
|
||||
|
|
@ -737,50 +592,9 @@ class DialecticStreamChunk(BaseModel):
|
|||
done: bool = False
|
||||
|
||||
|
||||
class SessionCounts(BaseModel):
|
||||
"""Counts for a specific session in queue processing."""
|
||||
|
||||
completed: int
|
||||
in_progress: int
|
||||
pending: int
|
||||
|
||||
|
||||
class QueueCounts(BaseModel):
|
||||
"""Aggregated counts for queue processing status."""
|
||||
|
||||
total: int
|
||||
completed: int
|
||||
in_progress: int
|
||||
pending: int
|
||||
sessions: dict[str, SessionCounts]
|
||||
|
||||
|
||||
class QueueStatusRow(BaseModel):
|
||||
"""Represents a row from the queue status SQL query result."""
|
||||
|
||||
session_id: str | None
|
||||
total: int
|
||||
completed: int
|
||||
in_progress: int
|
||||
pending: int
|
||||
session_total: int
|
||||
session_completed: int
|
||||
session_in_progress: int
|
||||
session_pending: int
|
||||
|
||||
|
||||
class SessionPeerData(BaseModel):
|
||||
"""Data for managing session peer relationships."""
|
||||
|
||||
peer_names: dict[str, SessionPeerConfig]
|
||||
|
||||
|
||||
class MessageBulkData(BaseModel):
|
||||
"""Data for bulk message operations."""
|
||||
|
||||
messages: list[MessageCreate]
|
||||
session_name: str
|
||||
workspace_name: str
|
||||
# ---------------------------------------------------------------------------
|
||||
# Queue status schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SessionQueueStatus(BaseModel):
|
||||
|
|
@ -822,6 +636,11 @@ class QueueStatus(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dream scheduling schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ScheduleDreamRequest(BaseModel):
|
||||
observer: str = Field(..., description="Observer peer name")
|
||||
observed: str | None = Field(
|
||||
|
|
@ -833,7 +652,11 @@ class ScheduleDreamRequest(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
# Webhook endpoint schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
# Webhook schemas
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WebhookEndpointBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
|
@ -0,0 +1,186 @@
|
|||
"""Configuration schemas for hierarchical settings resolution.
|
||||
|
||||
Covers workspace, session, and message-level configuration as well as
|
||||
the fully-resolved variants used at runtime.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, Self, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
|
||||
class DreamType(str, Enum):
|
||||
"""Types of dreams that can be triggered."""
|
||||
|
||||
OMNI = "omni"
|
||||
|
||||
|
||||
class ReasoningConfiguration(BaseModel):
|
||||
enabled: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to enable reasoning functionality.",
|
||||
)
|
||||
custom_instructions: str | None = Field(
|
||||
default=None,
|
||||
description="TODO: currently unused. Custom instructions to use for the reasoning system on this workspace/session/message.",
|
||||
)
|
||||
|
||||
|
||||
class PeerCardConfiguration(BaseModel):
|
||||
use: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to use peer card related to this peer during reasoning process.",
|
||||
)
|
||||
create: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to generate peer card based on content.",
|
||||
)
|
||||
|
||||
|
||||
class SummaryConfiguration(BaseModel):
|
||||
enabled: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to enable summary functionality.",
|
||||
)
|
||||
messages_per_short_summary: int | None = Field(
|
||||
default=None,
|
||||
ge=10,
|
||||
description="Number of messages per short summary. Must be positive, greater than or equal to 10, and less than messages_per_long_summary.",
|
||||
)
|
||||
messages_per_long_summary: int | None = Field(
|
||||
default=None,
|
||||
ge=20,
|
||||
description="Number of messages per long summary. Must be positive, greater than or equal to 20, and greater than messages_per_short_summary.",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_summary_thresholds(self) -> Self:
|
||||
"""Validate that short summary threshold <= long summary threshold."""
|
||||
short = self.messages_per_short_summary
|
||||
long = self.messages_per_long_summary
|
||||
|
||||
if short is not None and long is not None and short >= long:
|
||||
raise ValueError(
|
||||
"messages_per_short_summary must be less than messages_per_long_summary"
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class DreamConfiguration(BaseModel):
|
||||
enabled: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether to enable dream functionality. If reasoning is disabled, dreams will also be disabled and this setting will be ignored.",
|
||||
)
|
||||
|
||||
|
||||
class WorkspaceConfiguration(BaseModel):
|
||||
"""
|
||||
The set of options that can be in a workspace DB-level configuration dictionary.
|
||||
|
||||
All fields are optional. Session-level configuration overrides workspace-level configuration, which overrides global configuration.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="allow") # pyright: ignore
|
||||
|
||||
reasoning: ReasoningConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for reasoning functionality.",
|
||||
)
|
||||
peer_card: PeerCardConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for peer card functionality. If reasoning is disabled, peer cards will also be disabled and these settings will be ignored.",
|
||||
)
|
||||
summary: SummaryConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for summary functionality.",
|
||||
)
|
||||
dream: DreamConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for dream functionality. If reasoning is disabled, dreams will also be disabled and these settings will be ignored.",
|
||||
)
|
||||
|
||||
|
||||
class SessionConfiguration(WorkspaceConfiguration):
|
||||
"""
|
||||
The set of options that can be in a session DB-level configuration dictionary.
|
||||
|
||||
All fields are optional. Session-level configuration overrides workspace-level configuration, which overrides global configuration.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class MessageConfiguration(BaseModel):
|
||||
"""
|
||||
The set of options that can be in a message DB-level configuration dictionary.
|
||||
|
||||
All fields are optional. Message-level configuration overrides all other configurations.
|
||||
"""
|
||||
|
||||
reasoning: ReasoningConfiguration | None = Field(
|
||||
default=None,
|
||||
description="Configuration for reasoning functionality.",
|
||||
)
|
||||
|
||||
|
||||
class ResolvedReasoningConfiguration(BaseModel):
|
||||
enabled: bool
|
||||
|
||||
|
||||
class ResolvedPeerCardConfiguration(BaseModel):
|
||||
use: bool
|
||||
create: bool
|
||||
|
||||
|
||||
class ResolvedSummaryConfiguration(BaseModel):
|
||||
enabled: bool
|
||||
messages_per_short_summary: int
|
||||
messages_per_long_summary: int
|
||||
|
||||
|
||||
class ResolvedDreamConfiguration(BaseModel):
|
||||
enabled: bool
|
||||
|
||||
|
||||
class ResolvedConfiguration(BaseModel):
|
||||
"""
|
||||
The final resolved configuration for a given message.
|
||||
Hierarchy: message > session > workspace > global configuration
|
||||
"""
|
||||
|
||||
reasoning: ResolvedReasoningConfiguration
|
||||
peer_card: ResolvedPeerCardConfiguration
|
||||
summary: ResolvedSummaryConfiguration
|
||||
dream: ResolvedDreamConfiguration
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def migrate_deriver_to_reasoning(cls, data: Any) -> Any:
|
||||
"""Handle v3.0.0 migration: 'deriver' was renamed to 'reasoning'."""
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
|
||||
config = cast(dict[str, Any], data)
|
||||
|
||||
if "deriver" in config and "reasoning" not in config:
|
||||
config["reasoning"] = config.pop("deriver")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
class PeerConfig(BaseModel):
|
||||
# TODO: Update description - should say "Whether honcho forms a representation of the peer itself"
|
||||
observe_me: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether Honcho will use reasoning to form a representation of this peer",
|
||||
)
|
||||
|
||||
|
||||
class SessionPeerConfig(PeerConfig):
|
||||
# TODO: Update description - should say "Whether this peer forms representations of other peers in the session"
|
||||
observe_others: bool | None = Field(
|
||||
default=None,
|
||||
description="Whether this peer should form a session-level theory-of-mind representation of other peers in the session",
|
||||
)
|
||||
|
|
@ -0,0 +1,177 @@
|
|||
"""Internal schemas used by the deriver, dreamer, and other background systems.
|
||||
|
||||
These are not part of the public API contract and may change without notice.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Annotated, Literal, Self
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
from src.schemas.api import MessageCreate
|
||||
from src.schemas.configuration import SessionPeerConfig
|
||||
from src.utils.types import DocumentLevel
|
||||
|
||||
|
||||
class ReconcilerType(str, Enum):
|
||||
"""Types of reconciler tasks that can be performed."""
|
||||
|
||||
SYNC_VECTORS = "sync_vectors"
|
||||
CLEANUP_QUEUE = "cleanup_queue"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Document / observation schemas (vector storage internals)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DocumentBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class DocumentMetadata(BaseModel):
|
||||
message_ids: list[int] = Field(
|
||||
description="The ID range(s) of the messages that this document was derived from. Acts as a link to the primary source of the document. Note that as a document gets deduplicated, additional ranges will be added, because the same document could be derived from completely separate message ranges."
|
||||
)
|
||||
message_created_at: str = Field(
|
||||
description="The timestamp of the message that this document was derived from. Note that this is not the same as the created_at timestamp of the document. This timestamp is usually only saved with second-level precision."
|
||||
)
|
||||
source_ids: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Document IDs of source documents for tree traversal -- required for deductive and inductive documents",
|
||||
)
|
||||
premises: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Human-readable premise text for display -- only applicable for deductive documents",
|
||||
)
|
||||
sources: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Human-readable source text for display -- only applicable for inductive documents",
|
||||
)
|
||||
pattern_type: str | None = Field(
|
||||
default=None,
|
||||
description="Type of pattern identified (preference, behavior, personality, tendency, correlation) -- only applicable for inductive documents",
|
||||
)
|
||||
confidence: str | None = Field(
|
||||
default=None,
|
||||
description="Confidence level (high, medium, low) -- only applicable for inductive documents",
|
||||
)
|
||||
|
||||
|
||||
class DocumentCreate(DocumentBase):
|
||||
content: Annotated[str, Field(min_length=1, max_length=100000)]
|
||||
session_name: str | None = Field(
|
||||
default=None,
|
||||
description="The session from which the document was derived (NULL for global observations)",
|
||||
)
|
||||
level: DocumentLevel = Field(
|
||||
default="explicit",
|
||||
description="The level of the document (explicit, deductive, inductive, or contradiction)",
|
||||
)
|
||||
times_derived: int = Field(
|
||||
default=1,
|
||||
ge=1,
|
||||
description="The number of times that a semantic duplicate document to this one has been derived",
|
||||
)
|
||||
metadata: DocumentMetadata = Field()
|
||||
embedding: list[float] = Field()
|
||||
# Tree linkage field
|
||||
source_ids: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Document IDs of source/premise documents -- for deductive and inductive documents",
|
||||
)
|
||||
|
||||
|
||||
class ObservationInput(BaseModel):
|
||||
"""Validated observation input from LLM tool calls."""
|
||||
|
||||
content: Annotated[str, Field(min_length=1)]
|
||||
level: DocumentLevel = "explicit"
|
||||
source_ids: list[str] | None = None
|
||||
premises: list[str] | None = None
|
||||
sources: list[str] | None = None
|
||||
pattern_type: (
|
||||
Literal["preference", "behavior", "personality", "tendency", "correlation"]
|
||||
| None
|
||||
) = None
|
||||
confidence: Literal["high", "medium", "low"] | None = None
|
||||
|
||||
@field_validator("content", mode="after")
|
||||
@classmethod
|
||||
def sanitize_content(cls, v: str) -> str:
|
||||
return v.replace("\x00", "")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_level_fields(self) -> Self:
|
||||
"""Validate that level-specific fields are present when required."""
|
||||
if self.level == "deductive" and not self.source_ids:
|
||||
raise ValueError(
|
||||
"deductive observations require 'source_ids' field with document IDs of premises"
|
||||
)
|
||||
if self.level == "inductive" and not self.source_ids:
|
||||
raise ValueError(
|
||||
"inductive observations require 'source_ids' field with document IDs of sources"
|
||||
)
|
||||
if self.level == "contradiction" and (
|
||||
not self.source_ids or len(self.source_ids) < 2
|
||||
):
|
||||
raise ValueError(
|
||||
"contradiction observations require 'source_ids' field with at least 2 IDs of contradicting observations"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Queue internals
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SessionCounts(BaseModel):
|
||||
"""Counts for a specific session in queue processing."""
|
||||
|
||||
completed: int
|
||||
in_progress: int
|
||||
pending: int
|
||||
|
||||
|
||||
class QueueCounts(BaseModel):
|
||||
"""Aggregated counts for queue processing status."""
|
||||
|
||||
total: int
|
||||
completed: int
|
||||
in_progress: int
|
||||
pending: int
|
||||
sessions: dict[str, SessionCounts]
|
||||
|
||||
|
||||
class QueueStatusRow(BaseModel):
|
||||
"""Represents a row from the queue status SQL query result."""
|
||||
|
||||
session_id: str | None
|
||||
total: int
|
||||
completed: int
|
||||
in_progress: int
|
||||
pending: int
|
||||
session_total: int
|
||||
session_completed: int
|
||||
session_in_progress: int
|
||||
session_pending: int
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal data containers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SessionPeerData(BaseModel):
|
||||
"""Data for managing session peer relationships."""
|
||||
|
||||
peer_names: dict[str, SessionPeerConfig]
|
||||
|
||||
|
||||
class MessageBulkData(BaseModel):
|
||||
"""Data for bulk message operations."""
|
||||
|
||||
messages: list[MessageCreate]
|
||||
session_name: str
|
||||
workspace_name: str
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import logging
|
||||
import weakref
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
|
@ -46,7 +47,14 @@ def _safe_int(value: Any, default: int) -> int:
|
|||
# Module-level lock registry for thread-safe observation creation.
|
||||
# Keyed by (workspace_name, observer, observed) to ensure all tool executors
|
||||
# operating on the same data share the same lock.
|
||||
_observation_locks: dict[tuple[str, str, str], asyncio.Lock] = {}
|
||||
#
|
||||
# Uses WeakValueDictionary so entries are automatically removed when no
|
||||
# ToolContext holds a reference to the lock (i.e., all executors for that
|
||||
# key have finished and been garbage collected). This prevents unbounded
|
||||
# growth over the lifetime of a long-running deriver process.
|
||||
_observation_locks: weakref.WeakValueDictionary[tuple[str, str, str], asyncio.Lock] = (
|
||||
weakref.WeakValueDictionary()
|
||||
)
|
||||
_registry_lock = asyncio.Lock()
|
||||
|
||||
|
||||
|
|
@ -59,6 +67,11 @@ async def get_observation_lock(
|
|||
This ensures that concurrent tool executors operating on the same observation
|
||||
space share a lock, preventing race conditions during document creation.
|
||||
|
||||
The lock is stored as a weak reference — it stays alive as long as at least
|
||||
one ToolContext (via create_tool_executor) holds a strong reference. Once all
|
||||
executors for a key finish and are garbage collected, the entry is
|
||||
automatically removed from the registry.
|
||||
|
||||
Args:
|
||||
workspace_name: Workspace identifier
|
||||
observer: The observing peer
|
||||
|
|
@ -69,9 +82,11 @@ async def get_observation_lock(
|
|||
"""
|
||||
key = (workspace_name, observer, observed)
|
||||
async with _registry_lock:
|
||||
if key not in _observation_locks:
|
||||
_observation_locks[key] = asyncio.Lock()
|
||||
return _observation_locks[key]
|
||||
lock = _observation_locks.get(key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
_observation_locks[key] = lock
|
||||
return lock
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
|
|||
|
|
@ -107,7 +107,7 @@ def apply_filter(
|
|||
|
||||
|
||||
def _build_filter_conditions(
|
||||
filter_dict: dict[str, Any], model_class: type[Any]
|
||||
filter_dict: dict[str, Any], model_class: type[Any], *, _depth: int = 0
|
||||
) -> ColumnElement[bool] | None:
|
||||
"""
|
||||
Recursively build filter conditions from a filter dictionary.
|
||||
|
|
@ -119,6 +119,9 @@ def _build_filter_conditions(
|
|||
Returns:
|
||||
SQLAlchemy condition object or None
|
||||
"""
|
||||
if _depth > 5:
|
||||
raise FilterError("Filter nesting exceeds maximum depth of 5")
|
||||
|
||||
conditions: list[ColumnElement[bool]] = []
|
||||
|
||||
# Handle logical operators
|
||||
|
|
@ -129,7 +132,11 @@ def _build_filter_conditions(
|
|||
)
|
||||
and_conditions: list[ColumnElement[bool]] = []
|
||||
for sub_filter in filter_dict["AND"]: # pyright: ignore
|
||||
sub_condition = _build_filter_conditions(sub_filter, model_class) # pyright: ignore
|
||||
sub_condition = _build_filter_conditions(
|
||||
sub_filter, # pyright: ignore[reportUnknownArgumentType]
|
||||
model_class,
|
||||
_depth=_depth + 1,
|
||||
)
|
||||
if sub_condition is not None:
|
||||
and_conditions.append(sub_condition)
|
||||
if and_conditions:
|
||||
|
|
@ -142,7 +149,11 @@ def _build_filter_conditions(
|
|||
)
|
||||
or_conditions: list[ColumnElement[bool]] = []
|
||||
for sub_filter in filter_dict["OR"]: # pyright: ignore
|
||||
sub_condition = _build_filter_conditions(sub_filter, model_class) # pyright: ignore
|
||||
sub_condition = _build_filter_conditions(
|
||||
sub_filter, # pyright: ignore[reportUnknownArgumentType]
|
||||
model_class,
|
||||
_depth=_depth + 1,
|
||||
)
|
||||
if sub_condition is not None:
|
||||
or_conditions.append(sub_condition)
|
||||
if or_conditions:
|
||||
|
|
@ -157,7 +168,11 @@ def _build_filter_conditions(
|
|||
)
|
||||
not_conditions: list[ColumnElement[bool]] = []
|
||||
for sub_filter in filter_dict["NOT"]: # pyright: ignore
|
||||
sub_condition = _build_filter_conditions(sub_filter, model_class) # pyright: ignore
|
||||
sub_condition = _build_filter_conditions(
|
||||
sub_filter, # pyright: ignore[reportUnknownArgumentType]
|
||||
model_class,
|
||||
_depth=_depth + 1,
|
||||
)
|
||||
if sub_condition is not None:
|
||||
not_conditions.append(
|
||||
not_(sub_condition)
|
||||
|
|
|
|||
|
|
@ -13,12 +13,12 @@ from src.dependencies import tracked_db as real_tracked_db
|
|||
class FakeSession:
|
||||
def __init__(self, *, in_transaction: bool = False):
|
||||
self._in_transaction: bool = in_transaction
|
||||
self.execute_calls: list[Any] = []
|
||||
self.execute_calls: list[tuple[Any, ...]] = []
|
||||
self.rollback_calls: int = 0
|
||||
self.close_calls: int = 0
|
||||
|
||||
async def execute(self, statement: Any) -> None:
|
||||
self.execute_calls.append(statement)
|
||||
async def execute(self, statement: Any, params: Any = None) -> None:
|
||||
self.execute_calls.append((statement, params))
|
||||
|
||||
async def rollback(self) -> None:
|
||||
self.rollback_calls += 1
|
||||
|
|
@ -45,9 +45,9 @@ async def test_get_db_sets_application_name_when_tracing_enabled(
|
|||
db = await anext(dep_gen)
|
||||
assert db is fake_db
|
||||
assert len(fake_db.execute_calls) == 1
|
||||
assert "SET application_name = 'request:test-ctx'" in str(
|
||||
fake_db.execute_calls[0]
|
||||
)
|
||||
stmt, params = fake_db.execute_calls[0]
|
||||
assert "set_config" in str(stmt)
|
||||
assert params == {"name": "request:test-ctx"}
|
||||
finally:
|
||||
await dep_gen.aclose()
|
||||
request_context.reset(context_token)
|
||||
|
|
@ -96,9 +96,9 @@ async def test_tracked_db_creates_and_resets_task_context(
|
|||
|
||||
assert request_context.get() is None
|
||||
assert len(fake_db.execute_calls) == 1
|
||||
assert "SET application_name = 'task:cleanup_job:12345678'" in str(
|
||||
fake_db.execute_calls[0]
|
||||
)
|
||||
stmt, params = fake_db.execute_calls[0]
|
||||
assert "set_config" in str(stmt)
|
||||
assert params == {"name": "task:cleanup_job:12345678"}
|
||||
assert fake_db.rollback_calls == 0
|
||||
assert fake_db.close_calls == 1
|
||||
|
||||
|
|
@ -119,7 +119,9 @@ async def test_tracked_db_preserves_existing_request_context(
|
|||
request_context.reset(context_token)
|
||||
|
||||
assert len(fake_db.execute_calls) == 1
|
||||
assert "SET application_name = 'request:existing'" in str(fake_db.execute_calls[0])
|
||||
stmt, params = fake_db.execute_calls[0]
|
||||
assert "set_config" in str(stmt)
|
||||
assert params == {"name": "request:existing"}
|
||||
assert fake_db.rollback_calls == 0
|
||||
assert fake_db.close_calls == 1
|
||||
|
||||
|
|
|
|||
|
|
@ -1160,3 +1160,147 @@ class TestToolExecutor:
|
|||
if "Found" in result and "observations" in result:
|
||||
# IDs should be included in the output
|
||||
assert "[id:" in result or "observations" in result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Observation Lock Registry Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestObservationLockRegistry:
|
||||
"""Tests for the WeakValueDictionary-based observation lock registry."""
|
||||
|
||||
async def test_same_key_returns_same_lock(self):
|
||||
"""Concurrent callers with the same key get the same Lock instance."""
|
||||
from src.utils.agent_tools import get_observation_lock
|
||||
|
||||
lock_a = await get_observation_lock("ws1", "obs1", "peer1")
|
||||
lock_b = await get_observation_lock("ws1", "obs1", "peer1")
|
||||
|
||||
assert lock_a is lock_b
|
||||
|
||||
async def test_different_keys_return_different_locks(self):
|
||||
"""Different keys produce independent Lock instances."""
|
||||
from src.utils.agent_tools import get_observation_lock
|
||||
|
||||
lock_a = await get_observation_lock("ws_diff_a", "obs", "peer")
|
||||
lock_b = await get_observation_lock("ws_diff_b", "obs", "peer")
|
||||
|
||||
assert lock_a is not lock_b
|
||||
|
||||
async def test_lock_evicted_after_all_references_dropped(self):
|
||||
"""Lock is removed from registry once no strong references remain."""
|
||||
import gc
|
||||
|
||||
from src.utils.agent_tools import (
|
||||
_observation_locks, # pyright: ignore[reportPrivateUsage]
|
||||
get_observation_lock,
|
||||
)
|
||||
|
||||
key = ("ws_evict", "obs_evict", "peer_evict")
|
||||
lock = await get_observation_lock(*key)
|
||||
assert key in _observation_locks
|
||||
|
||||
# Drop the only strong reference and force GC
|
||||
del lock
|
||||
gc.collect()
|
||||
|
||||
assert key not in _observation_locks
|
||||
|
||||
async def test_lock_recreated_after_eviction(self):
|
||||
"""A new lock is created for a key whose previous lock was evicted."""
|
||||
import gc
|
||||
import weakref
|
||||
|
||||
from src.utils.agent_tools import get_observation_lock
|
||||
|
||||
key = ("ws_recreate", "obs_recreate", "peer_recreate")
|
||||
first_lock = await get_observation_lock(*key)
|
||||
first_ref = weakref.ref(first_lock)
|
||||
|
||||
# Evict
|
||||
del first_lock
|
||||
gc.collect()
|
||||
|
||||
# Confirm the old lock was garbage-collected
|
||||
assert first_ref() is None
|
||||
|
||||
# Recreate
|
||||
second_lock = await get_observation_lock(*key)
|
||||
assert isinstance(second_lock, asyncio.Lock)
|
||||
|
||||
async def test_lock_survives_while_any_reference_held(self):
|
||||
"""Lock stays alive as long as at least one strong reference exists."""
|
||||
import gc
|
||||
|
||||
from src.utils.agent_tools import (
|
||||
_observation_locks, # pyright: ignore[reportPrivateUsage]
|
||||
get_observation_lock,
|
||||
)
|
||||
|
||||
key = ("ws_survive", "obs_survive", "peer_survive")
|
||||
ref_a = await get_observation_lock(*key)
|
||||
ref_b = await get_observation_lock(*key)
|
||||
assert ref_a is ref_b
|
||||
|
||||
# Drop one reference — lock should survive via the other
|
||||
del ref_a
|
||||
gc.collect()
|
||||
assert key in _observation_locks
|
||||
|
||||
# Drop the last reference — now it should be evicted
|
||||
del ref_b
|
||||
gc.collect()
|
||||
assert key not in _observation_locks
|
||||
|
||||
async def test_concurrent_executors_share_lock_for_mutual_exclusion(self):
|
||||
"""Two coroutines using the same key are serialized by the shared lock."""
|
||||
from src.utils.agent_tools import get_observation_lock
|
||||
|
||||
key = ("ws_mutex", "obs_mutex", "peer_mutex")
|
||||
shared_lock = await get_observation_lock(*key)
|
||||
|
||||
order: list[str] = []
|
||||
|
||||
async def task(name: str, delay: float):
|
||||
async with shared_lock:
|
||||
order.append(f"{name}_start")
|
||||
await asyncio.sleep(delay)
|
||||
order.append(f"{name}_end")
|
||||
|
||||
# task_a grabs the lock first, task_b must wait
|
||||
task_a = asyncio.create_task(task("a", 0.05))
|
||||
await asyncio.sleep(0.01) # let task_a acquire the lock
|
||||
task_b = asyncio.create_task(task("b", 0.01))
|
||||
|
||||
await asyncio.gather(task_a, task_b)
|
||||
|
||||
# task_a must fully complete before task_b starts
|
||||
assert order == ["a_start", "a_end", "b_start", "b_end"]
|
||||
|
||||
async def test_no_registry_growth_across_many_keys(self):
|
||||
"""Registry does not retain locks after references are dropped."""
|
||||
import gc
|
||||
|
||||
from src.utils.agent_tools import (
|
||||
_observation_locks, # pyright: ignore[reportPrivateUsage]
|
||||
get_observation_lock,
|
||||
)
|
||||
|
||||
locks: list[asyncio.Lock] = []
|
||||
for i in range(100):
|
||||
locks.append(await get_observation_lock(f"ws_growth_{i}", "obs", "peer"))
|
||||
|
||||
count_before = sum(
|
||||
1 for k in _observation_locks if k[0].startswith("ws_growth_")
|
||||
)
|
||||
assert count_before == 100
|
||||
|
||||
# Drop all strong references and force GC
|
||||
locks.clear()
|
||||
gc.collect()
|
||||
|
||||
# All 100 entries should be cleaned up
|
||||
remaining = sum(1 for k in _observation_locks if k[0].startswith("ws_growth_"))
|
||||
assert remaining == 0
|
||||
|
|
|
|||
Loading…
Reference in New Issue