1218 lines
41 KiB
Python
1218 lines
41 KiB
Python
"""CRUD helpers for sessions and session-related relationship data."""
|
|
|
|
from dataclasses import dataclass
|
|
from logging import getLogger
|
|
from typing import Any
|
|
from typing import cast as typing_cast
|
|
|
|
from cashews import NOT_NONE
|
|
from nanoid import generate as generate_nanoid
|
|
from sqlalchemy import Select, and_, case, cast, delete, func, insert, select, update
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
from sqlalchemy.engine import CursorResult
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.orm import make_transient_to_detached
|
|
from sqlalchemy.types import BigInteger, Boolean
|
|
|
|
from src import models, schemas
|
|
from src.cache.client import (
|
|
cache,
|
|
get_cache_namespace,
|
|
safe_cache_delete,
|
|
safe_cache_set,
|
|
)
|
|
from src.config import settings
|
|
from src.exceptions import (
|
|
ConflictException,
|
|
ObserverException,
|
|
ResourceNotFoundException,
|
|
)
|
|
from src.utils.filter import apply_filter
|
|
from src.utils.types import GetOrCreateResult
|
|
from src.vector_store import get_external_vector_store
|
|
|
|
from .peer import get_or_create_peers, get_peer
|
|
from .workspace import get_or_create_workspace
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class SessionDeletionResult:
|
|
"""Result of a session deletion including cascade counts."""
|
|
|
|
messages_deleted: int
|
|
conclusions_deleted: int
|
|
|
|
|
|
SESSION_CACHE_KEY_TEMPLATE = "v2:workspace:{workspace_name}:session:{session_name}"
|
|
SESSION_LOCK_PREFIX = f"{get_cache_namespace()}:lock:v2"
|
|
|
|
|
|
def session_cache_key(workspace_name: str, session_name: str) -> str:
|
|
"""Generate cache key for session."""
|
|
return (
|
|
get_cache_namespace()
|
|
+ ":"
|
|
+ SESSION_CACHE_KEY_TEMPLATE.format(
|
|
workspace_name=workspace_name,
|
|
session_name=session_name,
|
|
)
|
|
)
|
|
|
|
|
|
@cache(
|
|
key=SESSION_CACHE_KEY_TEMPLATE,
|
|
ttl=f"{settings.CACHE.DEFAULT_TTL_SECONDS}s",
|
|
prefix=get_cache_namespace(),
|
|
condition=NOT_NONE,
|
|
)
|
|
@cache.locked(
|
|
key=SESSION_CACHE_KEY_TEMPLATE,
|
|
ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s",
|
|
prefix=SESSION_LOCK_PREFIX,
|
|
)
|
|
async def _fetch_session(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
) -> dict[str, Any] | None:
|
|
"""Fetch a session from the database and return as a plain dict for safe caching."""
|
|
obj = await db.scalar(
|
|
select(models.Session)
|
|
.where(models.Session.workspace_name == workspace_name)
|
|
.where(models.Session.name == session_name)
|
|
)
|
|
if obj is None:
|
|
return None
|
|
return {
|
|
"id": obj.id,
|
|
"name": obj.name,
|
|
"workspace_name": obj.workspace_name,
|
|
"is_active": obj.is_active,
|
|
"h_metadata": obj.h_metadata,
|
|
"internal_metadata": obj.internal_metadata,
|
|
"configuration": obj.configuration,
|
|
"created_at": obj.created_at,
|
|
}
|
|
|
|
|
|
def count_observers_in_config(
|
|
peer_configs: dict[str, schemas.SessionPeerConfig],
|
|
) -> int:
|
|
"""
|
|
Count the number of peers that will be observing others based on their configurations.
|
|
|
|
Args:
|
|
peer_configs: Dictionary of peer names to their session configurations
|
|
|
|
Returns:
|
|
Number of peers that will be observing others
|
|
"""
|
|
return sum(1 for config in peer_configs.values() if config.observe_others)
|
|
|
|
|
|
async def get_sessions(
|
|
workspace_name: str,
|
|
filters: dict[str, Any] | None = None,
|
|
reverse: bool = False,
|
|
) -> Select[tuple[models.Session]]:
|
|
"""
|
|
Get all active sessions in a workspace.
|
|
|
|
Args:
|
|
workspace_name: Name of the workspace
|
|
filters: Optional filters to apply to the query
|
|
reverse: If True, order by created_at descending; if False, ascending
|
|
|
|
Returns:
|
|
Select statement for Session objects
|
|
"""
|
|
stmt = (
|
|
select(models.Session)
|
|
.where(models.Session.workspace_name == workspace_name)
|
|
.where(models.Session.is_active == True) # noqa: E712
|
|
)
|
|
|
|
stmt = apply_filter(stmt, models.Session, filters)
|
|
|
|
if reverse:
|
|
return stmt.order_by(models.Session.created_at.desc(), models.Session.id.desc())
|
|
return stmt.order_by(models.Session.created_at.asc(), models.Session.id.asc())
|
|
|
|
|
|
async def get_or_create_session(
|
|
db: AsyncSession,
|
|
session: schemas.SessionCreate,
|
|
workspace_name: str,
|
|
*,
|
|
_retry: bool = False,
|
|
) -> GetOrCreateResult[models.Session]:
|
|
"""
|
|
Get an active session in a workspace or create it if it does not exist.
|
|
|
|
If the session already exists, provided metadata replaces the current
|
|
metadata, provided configuration keys are merged into the existing
|
|
configuration, and any provided peers are ensured to be members of the
|
|
session. If the session does not exist, the workspace and peers are created
|
|
as needed before the session is created.
|
|
|
|
Args:
|
|
db: Database session
|
|
session: Session creation payload, including optional metadata,
|
|
configuration, and session-peer configuration
|
|
workspace_name: Name of the workspace
|
|
_retry: Whether to retry after a concurrent create conflict
|
|
|
|
Returns:
|
|
GetOrCreateResult containing the session and whether it was created
|
|
|
|
Raises:
|
|
ValueError: If session.name is empty
|
|
ResourceNotFoundException: If the named session exists but is inactive
|
|
ObserverException: If adding peers would exceed the observer limit
|
|
ConflictException: If concurrent creation prevents fetching or creating
|
|
the session
|
|
"""
|
|
|
|
if not session.name:
|
|
raise ValueError("Session name must be provided")
|
|
|
|
session_data = await _fetch_session(db, workspace_name, session.name)
|
|
|
|
# Reconstruct and merge cached dict into session if it exists
|
|
honcho_session: models.Session | None = None
|
|
if session_data is not None:
|
|
obj = models.Session(**session_data)
|
|
make_transient_to_detached(obj)
|
|
honcho_session = await db.merge(obj, load=False)
|
|
|
|
# Reject operations on inactive sessions (marked for deletion)
|
|
if not honcho_session.is_active:
|
|
raise ResourceNotFoundException(
|
|
f"Session {session.name} not found in workspace {workspace_name}"
|
|
)
|
|
|
|
# Track if we need to update cache and if session was created
|
|
needs_cache_update = False
|
|
created = False
|
|
ws_result = None
|
|
peers_result = None
|
|
|
|
# Check if session already exists
|
|
if honcho_session is None:
|
|
if session.peer_names:
|
|
# Count peers that will be observing others
|
|
observer_count = count_observers_in_config(session.peer_names)
|
|
if observer_count > settings.SESSION_OBSERVERS_LIMIT:
|
|
raise ObserverException(session.name, observer_count)
|
|
|
|
# Get or create workspace to ensure it exists
|
|
ws_result = await get_or_create_workspace(
|
|
db,
|
|
schemas.WorkspaceCreate(name=workspace_name),
|
|
)
|
|
|
|
# Create honcho session
|
|
honcho_session = models.Session(
|
|
workspace_name=workspace_name,
|
|
name=session.name,
|
|
h_metadata=session.metadata or {},
|
|
configuration=session.configuration.model_dump(exclude_none=True)
|
|
if session.configuration
|
|
else {},
|
|
)
|
|
try:
|
|
async with db.begin_nested():
|
|
db.add(honcho_session)
|
|
needs_cache_update = True
|
|
created = True
|
|
|
|
except IntegrityError:
|
|
logger.debug(
|
|
"Race condition detected for session: %s, retrying get", session.name
|
|
)
|
|
if _retry:
|
|
raise ConflictException(
|
|
f"Unable to create or get session: {session.name}"
|
|
) from None
|
|
return await get_or_create_session(db, session, workspace_name, _retry=True)
|
|
else:
|
|
# Update existing session with metadata and feature flags if provided
|
|
if (
|
|
session.metadata is not None
|
|
and honcho_session.h_metadata != session.metadata
|
|
):
|
|
honcho_session.h_metadata = session.metadata
|
|
needs_cache_update = True
|
|
if session.configuration is not None:
|
|
# Merge configuration instead of replacing to preserve existing keys
|
|
existing_config = (honcho_session.configuration or {}).copy()
|
|
incoming_config = session.configuration.model_dump(exclude_none=True)
|
|
merged_config = {**existing_config, **incoming_config}
|
|
if honcho_session.configuration != merged_config:
|
|
honcho_session.configuration = merged_config
|
|
needs_cache_update = True
|
|
|
|
# Add all peers to session
|
|
if session.peer_names:
|
|
peers_result = await get_or_create_peers(
|
|
db,
|
|
workspace_name=workspace_name,
|
|
peers=[
|
|
schemas.PeerCreate(name=peer_name) for peer_name in session.peer_names
|
|
],
|
|
)
|
|
await _get_or_add_peers_to_session(
|
|
db,
|
|
workspace_name=workspace_name,
|
|
session_name=session.name,
|
|
peer_names=session.peer_names,
|
|
fetch_after_upsert=False,
|
|
)
|
|
|
|
await db.commit()
|
|
|
|
# Run deferred cache operations from workspace/peer creation
|
|
if ws_result is not None:
|
|
await ws_result.post_commit()
|
|
if peers_result is not None:
|
|
await peers_result.post_commit()
|
|
|
|
# Only update cache if session data changed or was newly created
|
|
if needs_cache_update:
|
|
cache_key = session_cache_key(workspace_name, session.name)
|
|
await safe_cache_set(
|
|
cache_key,
|
|
{
|
|
"id": honcho_session.id,
|
|
"name": honcho_session.name,
|
|
"workspace_name": honcho_session.workspace_name,
|
|
"is_active": honcho_session.is_active,
|
|
"h_metadata": honcho_session.h_metadata,
|
|
"internal_metadata": honcho_session.internal_metadata,
|
|
"configuration": honcho_session.configuration,
|
|
"created_at": honcho_session.created_at,
|
|
},
|
|
expire=settings.CACHE.DEFAULT_TTL_SECONDS,
|
|
)
|
|
logger.debug(
|
|
"Session %s cache updated in workspace %s", session.name, workspace_name
|
|
)
|
|
|
|
return GetOrCreateResult(honcho_session, created=created)
|
|
|
|
|
|
async def get_session(
|
|
db: AsyncSession,
|
|
session_name: str,
|
|
workspace_name: str,
|
|
*,
|
|
include_inactive: bool = False,
|
|
) -> models.Session:
|
|
"""
|
|
Get a session in a workspace.
|
|
|
|
Args:
|
|
db: Database session
|
|
session_name: Name of the session
|
|
workspace_name: Name of the workspace
|
|
include_inactive: If True, return sessions even if they are marked for deletion.
|
|
This should only be used for internal operations like the deletion task.
|
|
|
|
Returns:
|
|
The session
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If the session does not exist or is inactive
|
|
"""
|
|
data = await _fetch_session(db, workspace_name, session_name)
|
|
|
|
if data is None:
|
|
raise ResourceNotFoundException(
|
|
f"Session {session_name} not found in workspace {workspace_name}"
|
|
)
|
|
|
|
# Check if session is active (unless include_inactive is True)
|
|
# Check on the dict before constructing the ORM object
|
|
if not include_inactive and not data["is_active"]:
|
|
raise ResourceNotFoundException(
|
|
f"Session {session_name} not found in workspace {workspace_name}"
|
|
)
|
|
|
|
# Reconstruct ORM object from cached dict and merge into session
|
|
obj = models.Session(**data)
|
|
make_transient_to_detached(obj)
|
|
session = await db.merge(obj, load=False)
|
|
|
|
return session
|
|
|
|
|
|
async def update_session(
|
|
db: AsyncSession,
|
|
session: schemas.SessionUpdate,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
) -> models.Session:
|
|
"""
|
|
Get or create a session, then apply metadata and configuration updates.
|
|
|
|
Provided metadata replaces the current metadata when present. Provided
|
|
configuration keys are merged into the existing configuration instead of
|
|
replacing it wholesale.
|
|
|
|
Args:
|
|
db: Database session
|
|
session: Session update schema
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
|
|
Returns:
|
|
The updated session
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If the named session exists but is inactive
|
|
ConflictException: If concurrent creation prevents fetching or creating
|
|
the session
|
|
"""
|
|
honcho_session: models.Session = (
|
|
await get_or_create_session(
|
|
db, schemas.SessionCreate(name=session_name), workspace_name=workspace_name
|
|
)
|
|
).resource
|
|
|
|
# Track if anything changed
|
|
needs_update = False
|
|
|
|
if session.metadata is not None and honcho_session.h_metadata != session.metadata:
|
|
honcho_session.h_metadata = session.metadata
|
|
needs_update = True
|
|
|
|
if session.configuration is not None:
|
|
# Merge configuration instead of replacing to preserve existing keys
|
|
base_config = (honcho_session.configuration or {}).copy()
|
|
merged_config = {
|
|
**base_config,
|
|
**session.configuration.model_dump(exclude_none=True),
|
|
}
|
|
if honcho_session.configuration != merged_config:
|
|
honcho_session.configuration = merged_config
|
|
needs_update = True
|
|
|
|
if not needs_update:
|
|
logger.debug(
|
|
"Session %s unchanged in workspace %s, skipping update",
|
|
session_name,
|
|
workspace_name,
|
|
)
|
|
return honcho_session
|
|
|
|
await db.commit()
|
|
|
|
# Only invalidate if we actually updated
|
|
cache_key = session_cache_key(workspace_name, session_name)
|
|
await safe_cache_delete(cache_key)
|
|
|
|
logger.debug("Session %s updated successfully", session_name)
|
|
return honcho_session
|
|
|
|
|
|
async def _batch_delete_matching(
|
|
db: AsyncSession,
|
|
model: Any,
|
|
filter_conditions: list[Any],
|
|
batch_size: int = 5000,
|
|
) -> int:
|
|
"""
|
|
Delete records in batches that match the given filter conditions.
|
|
|
|
Args:
|
|
db: Database session
|
|
model: SQLAlchemy model class
|
|
filter_conditions: List of SQLAlchemy filter conditions
|
|
batch_size: Number of records to delete per batch
|
|
|
|
Returns:
|
|
Total number of records deleted
|
|
"""
|
|
total_deleted = 0
|
|
primary_key_column = model.__table__.primary_key.columns.values()[0]
|
|
|
|
while True:
|
|
subquery = (
|
|
select(primary_key_column).where(and_(*filter_conditions)).limit(batch_size)
|
|
)
|
|
delete_stmt = delete(model).where(primary_key_column.in_(subquery))
|
|
delete_result = typing_cast(CursorResult[Any], await db.execute(delete_stmt))
|
|
batch_deleted = delete_result.rowcount or 0
|
|
total_deleted += batch_deleted
|
|
|
|
if batch_deleted == 0:
|
|
break
|
|
|
|
return total_deleted
|
|
|
|
|
|
async def delete_session(
|
|
db: AsyncSession, workspace_name: str, session_name: str
|
|
) -> SessionDeletionResult:
|
|
"""
|
|
Delete a session and all associated data (hard delete).
|
|
|
|
This performs cascading deletes for all session-related data including:
|
|
- Active queue sessions
|
|
- Queue items
|
|
- Message embeddings (batched)
|
|
- Documents (theory-of-mind data, batched)
|
|
- Messages (batched)
|
|
- Session peer associations
|
|
- The session itself
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
|
|
Returns:
|
|
SessionDeletionResult containing cascade counts
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If the session does not exist
|
|
"""
|
|
honcho_session = await get_session(
|
|
db, session_name, workspace_name, include_inactive=True
|
|
)
|
|
|
|
# Perform cascading deletes in order
|
|
# Order is important to avoid foreign key constraint violations
|
|
try:
|
|
# Delete ActiveQueueSession entries
|
|
# Work unit keys have format: {task_type}:{workspace_name}:{session_name}:{...}
|
|
await db.execute(
|
|
delete(models.ActiveQueueSession).where(
|
|
and_(
|
|
func.split_part(models.ActiveQueueSession.work_unit_key, ":", 2)
|
|
== workspace_name,
|
|
func.split_part(models.ActiveQueueSession.work_unit_key, ":", 3)
|
|
== session_name,
|
|
)
|
|
)
|
|
)
|
|
|
|
# Delete QueueItem entries
|
|
await db.execute(
|
|
delete(models.QueueItem).where(
|
|
models.QueueItem.session_id == honcho_session.id
|
|
)
|
|
)
|
|
|
|
# Delete message vectors from vector store before deleting DB records
|
|
# Fetch all MessageEmbedding records to build vector IDs with {message_id}_{chunk_index}
|
|
embedding_result = await db.execute(
|
|
select(models.MessageEmbedding).where(
|
|
models.MessageEmbedding.session_name == session_name,
|
|
models.MessageEmbedding.workspace_name == workspace_name,
|
|
)
|
|
)
|
|
embeddings = list(embedding_result.scalars().all())
|
|
external_vector_store = get_external_vector_store()
|
|
|
|
# Only delete from external vector store if one exists
|
|
if external_vector_store is not None and embeddings:
|
|
# Compute chunk_index for each embedding based on message_id ordering
|
|
message_chunks: dict[str, list[models.MessageEmbedding]] = {}
|
|
for emb in embeddings:
|
|
message_chunks.setdefault(emb.message_id, []).append(emb)
|
|
|
|
# Sort each message's chunks by id and build vector IDs
|
|
vector_ids: list[str] = []
|
|
for chunks in message_chunks.values():
|
|
chunks.sort(key=lambda e: e.id)
|
|
for chunk_idx, chunk in enumerate(chunks):
|
|
vector_ids.append(f"{chunk.message_id}_{chunk_idx}")
|
|
|
|
# Try to delete from external vector store (best effort)
|
|
try:
|
|
namespace = external_vector_store.get_vector_namespace(
|
|
"message", workspace_name
|
|
)
|
|
await external_vector_store.delete_many(namespace, vector_ids)
|
|
logger.debug(
|
|
f"Deleted {len(vector_ids)} message vectors for session {session_name}"
|
|
)
|
|
except Exception as e:
|
|
# Log warning but continue - workspace deletion will clean up eventually
|
|
logger.warning(
|
|
f"Failed to delete message vectors for session {session_name}: {e}"
|
|
)
|
|
|
|
# Delete MessageEmbedding entries in batches
|
|
await _batch_delete_matching(
|
|
db,
|
|
models.MessageEmbedding,
|
|
[
|
|
models.MessageEmbedding.session_name == session_name,
|
|
models.MessageEmbedding.workspace_name == workspace_name,
|
|
],
|
|
batch_size=5000,
|
|
)
|
|
|
|
# Delete document vectors from vector store before deleting DB records
|
|
# Fetch all Document records to get IDs and namespaces
|
|
doc_result = await db.execute(
|
|
select(
|
|
models.Document.id,
|
|
models.Document.observer,
|
|
models.Document.observed,
|
|
).where(
|
|
models.Document.session_name == session_name,
|
|
models.Document.workspace_name == workspace_name,
|
|
)
|
|
)
|
|
documents = doc_result.all()
|
|
|
|
# Only delete from external vector store if one exists
|
|
if external_vector_store is not None and documents:
|
|
# Group document IDs by namespace (observer/observed)
|
|
docs_by_namespace: dict[str, list[str]] = {}
|
|
for doc in documents:
|
|
namespace = external_vector_store.get_vector_namespace(
|
|
"document",
|
|
workspace_name,
|
|
doc.observer,
|
|
doc.observed,
|
|
)
|
|
docs_by_namespace.setdefault(namespace, []).append(doc.id)
|
|
|
|
# Try to delete from external vector store (best effort, per namespace)
|
|
for namespace, doc_ids in docs_by_namespace.items():
|
|
try:
|
|
await external_vector_store.delete_many(namespace, doc_ids)
|
|
logger.debug(
|
|
f"Deleted {len(doc_ids)} document vectors from {namespace}"
|
|
)
|
|
except Exception as e:
|
|
# Log warning but continue - workspace deletion will clean up eventually
|
|
logger.warning(
|
|
f"Failed to delete document vectors from {namespace}: {e}"
|
|
)
|
|
|
|
# Delete Document entries associated with this session in batches
|
|
conclusions_deleted = await _batch_delete_matching(
|
|
db,
|
|
models.Document,
|
|
[
|
|
models.Document.session_name == session_name,
|
|
models.Document.workspace_name == workspace_name,
|
|
],
|
|
batch_size=5000,
|
|
)
|
|
|
|
# Delete Message entries in batches
|
|
messages_deleted = await _batch_delete_matching(
|
|
db,
|
|
models.Message,
|
|
[
|
|
models.Message.session_name == session_name,
|
|
models.Message.workspace_name == workspace_name,
|
|
],
|
|
batch_size=5000,
|
|
)
|
|
|
|
# Delete SessionPeer associations
|
|
await db.execute(
|
|
delete(models.SessionPeer).where(
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
)
|
|
)
|
|
|
|
# Finally, delete the session itself
|
|
await db.delete(honcho_session)
|
|
await db.commit()
|
|
|
|
# Invalidate session cache
|
|
await safe_cache_delete(session_cache_key(workspace_name, session_name))
|
|
|
|
logger.debug("Session %s and all associated data deleted", session_name)
|
|
except Exception as e:
|
|
logger.error("Failed to delete session %s: %s", session_name, e)
|
|
await db.rollback()
|
|
raise e
|
|
|
|
return SessionDeletionResult(
|
|
messages_deleted=messages_deleted,
|
|
conclusions_deleted=conclusions_deleted,
|
|
)
|
|
|
|
|
|
async def clone_session(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
original_session_name: str,
|
|
cutoff_message_id: str | None = None,
|
|
) -> models.Session:
|
|
"""
|
|
Clone a session and its data. If cutoff_message_id is provided,
|
|
only clone messages up to and including that message.
|
|
|
|
The following data is copied to the new session:
|
|
- Session metadata
|
|
- Session configuration
|
|
- All messages (or up to cutoff_message_id) with their content, metadata, and peer associations
|
|
- Session-peer associations with their configurations (observe_me, observe_others)
|
|
|
|
The new session gets a unique ID (nanoid) and fresh timestamps.
|
|
|
|
Args:
|
|
db: SQLAlchemy session
|
|
workspace_name: Name of the workspace the target session is in
|
|
original_session_name: Name of the session to clone
|
|
cutoff_message_id: Optional ID of the last message to include in the clone
|
|
|
|
Returns:
|
|
The newly created session
|
|
"""
|
|
# Get the original session (must be active)
|
|
stmt = (
|
|
select(models.Session)
|
|
.where(models.Session.workspace_name == workspace_name)
|
|
.where(models.Session.name == original_session_name)
|
|
.where(models.Session.is_active == True) # noqa: E712
|
|
)
|
|
result = await db.execute(stmt)
|
|
original_session = result.scalar_one_or_none()
|
|
if original_session is None:
|
|
raise ResourceNotFoundException("Original session not found")
|
|
|
|
# If cutoff_message_id is provided, verify it belongs to the session
|
|
cutoff_message = None
|
|
if cutoff_message_id is not None:
|
|
stmt = select(models.Message).where(
|
|
models.Message.public_id == cutoff_message_id,
|
|
models.Message.session_name == original_session_name,
|
|
)
|
|
cutoff_message = await db.scalar(stmt)
|
|
if not cutoff_message:
|
|
raise ValueError(
|
|
"Message not found or doesn't belong to the specified session"
|
|
)
|
|
|
|
# Create new session
|
|
new_session = models.Session(
|
|
workspace_name=workspace_name,
|
|
name=generate_nanoid(),
|
|
h_metadata=original_session.h_metadata,
|
|
configuration=original_session.configuration,
|
|
)
|
|
db.add(new_session)
|
|
await db.flush() # Flush to get the new session ID
|
|
|
|
# Build query for messages to clone
|
|
stmt = select(models.Message).where(
|
|
models.Message.session_name == original_session_name
|
|
)
|
|
if cutoff_message_id is not None and cutoff_message is not None:
|
|
stmt = stmt.where(models.Message.id <= cast(cutoff_message.id, BigInteger))
|
|
stmt = stmt.order_by(models.Message.id)
|
|
|
|
# Fetch messages to clone
|
|
messages_to_clone_scalars = await db.scalars(stmt)
|
|
messages_to_clone = messages_to_clone_scalars.all()
|
|
|
|
if not messages_to_clone:
|
|
return new_session
|
|
|
|
# Prepare bulk insert data
|
|
new_messages = [
|
|
{
|
|
"session_name": new_session.name,
|
|
"content": message.content,
|
|
"h_metadata": message.h_metadata,
|
|
"workspace_name": workspace_name,
|
|
"peer_name": message.peer_name,
|
|
"seq_in_session": message.seq_in_session,
|
|
}
|
|
for message in messages_to_clone
|
|
]
|
|
|
|
insert_stmt = insert(models.Message).returning(models.Message)
|
|
result = await db.execute(insert_stmt, new_messages)
|
|
|
|
# Clone peers from original session to new session (including their configurations)
|
|
stmt = select(models.SessionPeer).where(
|
|
models.SessionPeer.session_name == original_session_name
|
|
)
|
|
result = await db.execute(stmt)
|
|
session_peers = result.scalars().all()
|
|
for session_peer in session_peers:
|
|
new_session_peer = models.SessionPeer(
|
|
session_name=new_session.name,
|
|
peer_name=session_peer.peer_name,
|
|
workspace_name=workspace_name,
|
|
configuration=session_peer.configuration,
|
|
)
|
|
db.add(new_session_peer)
|
|
|
|
await db.commit()
|
|
logger.debug("Session %s cloned successfully", original_session_name)
|
|
|
|
# Cache will be populated on next read - read-through pattern
|
|
return new_session
|
|
|
|
|
|
async def remove_peers_from_session(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
peer_names: set[str],
|
|
) -> bool:
|
|
"""
|
|
Remove specified peers from a session.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
peer_names: Set of peer names to remove from the session
|
|
|
|
Returns:
|
|
True if peers were removed successfully
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If the session does not exist
|
|
"""
|
|
# Verify session exists
|
|
await get_session(db, session_name, workspace_name)
|
|
|
|
# Soft delete specified session peers by setting left_at timestamp
|
|
update_stmt = (
|
|
update(models.SessionPeer)
|
|
.where(
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
models.SessionPeer.peer_name.in_(peer_names),
|
|
models.SessionPeer.left_at.is_(None), # Only update active peers
|
|
)
|
|
.values(left_at=func.now())
|
|
)
|
|
await db.execute(update_stmt)
|
|
|
|
await db.commit()
|
|
return True
|
|
|
|
|
|
async def get_peers_from_session(
|
|
workspace_name: str,
|
|
session_name: str,
|
|
) -> Select[tuple[models.Peer]]:
|
|
"""
|
|
Get all peers from a session.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
|
|
Returns:
|
|
Paginated list of Peer objects in the session
|
|
"""
|
|
# Get all active peers in the session (where left_at is NULL)
|
|
return (
|
|
select(models.Peer)
|
|
.join(
|
|
models.SessionPeer,
|
|
and_(
|
|
models.Peer.name == models.SessionPeer.peer_name,
|
|
models.Peer.workspace_name == models.SessionPeer.workspace_name,
|
|
),
|
|
)
|
|
.where(models.SessionPeer.session_name == session_name)
|
|
.where(models.Peer.workspace_name == workspace_name)
|
|
.where(models.SessionPeer.left_at.is_(None)) # Only active peers
|
|
)
|
|
|
|
|
|
async def is_peer_in_session(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
peer_name: str,
|
|
) -> bool:
|
|
"""Return whether a peer is an active member of a session.
|
|
|
|
Active membership means a `SessionPeer` row exists with `left_at IS NULL`.
|
|
Used by the auth layer to grant a peer-scoped key read access to the
|
|
sessions that peer belongs to.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
peer_name: Name of the peer
|
|
|
|
Returns:
|
|
True if the peer is currently a member of the session.
|
|
"""
|
|
result = await db.scalar(
|
|
select(models.SessionPeer.peer_name)
|
|
.where(models.SessionPeer.workspace_name == workspace_name)
|
|
.where(models.SessionPeer.session_name == session_name)
|
|
.where(models.SessionPeer.peer_name == peer_name)
|
|
.where(models.SessionPeer.left_at.is_(None))
|
|
.limit(1)
|
|
)
|
|
return result is not None
|
|
|
|
|
|
async def get_session_peer_configuration(
|
|
workspace_name: str,
|
|
session_name: str,
|
|
) -> Select[tuple[str, dict[str, Any], dict[str, Any], bool]]:
|
|
"""
|
|
Get configuration from both SessionPeer and Peer tables for all peers in a session.
|
|
NOTE: does not filter for active peers. Will return peers that have left the session.
|
|
|
|
Args:
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
|
|
Returns:
|
|
Select statement returning peer_name, peer_configuration, session_peer_configuration,
|
|
and a boolean indicating if the peer is currently in the session
|
|
"""
|
|
stmt: Select[tuple[str, dict[str, Any], dict[str, Any], bool]] = (
|
|
select(
|
|
models.Peer.name.label("peer_name"),
|
|
models.Peer.configuration.label("peer_configuration"),
|
|
models.SessionPeer.configuration.label("session_peer_configuration"),
|
|
(models.SessionPeer.left_at.is_(None)).label("is_active"),
|
|
)
|
|
.join(
|
|
models.SessionPeer,
|
|
and_(
|
|
models.Peer.name == models.SessionPeer.peer_name,
|
|
models.Peer.workspace_name == models.SessionPeer.workspace_name,
|
|
),
|
|
)
|
|
.where(models.SessionPeer.session_name == session_name)
|
|
.where(models.Peer.workspace_name == workspace_name)
|
|
.where(models.SessionPeer.workspace_name == workspace_name)
|
|
)
|
|
|
|
return stmt
|
|
|
|
|
|
async def set_peers_for_session(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
peer_names: dict[str, schemas.SessionPeerConfig],
|
|
) -> list[models.SessionPeer]:
|
|
"""
|
|
Set peers for a session, overwriting any existing peers.
|
|
If peers don't exist, they will be created.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
peer_names: Set of peer names to set for the session
|
|
|
|
Returns:
|
|
List of SessionPeer objects for all peers in the session
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If the session does not exist
|
|
"""
|
|
# Validate observer limit before making any changes
|
|
observer_count = count_observers_in_config(peer_names)
|
|
if observer_count > settings.SESSION_OBSERVERS_LIMIT:
|
|
raise ObserverException(session_name, observer_count)
|
|
|
|
# Verify session exists
|
|
stmt = (
|
|
select(models.Session)
|
|
.where(models.Session.workspace_name == workspace_name)
|
|
.where(models.Session.name == session_name)
|
|
)
|
|
result = await db.execute(stmt)
|
|
session = result.scalar_one_or_none()
|
|
|
|
if session is None:
|
|
raise ResourceNotFoundException(
|
|
f"Session {session_name} not found in workspace {workspace_name}"
|
|
)
|
|
|
|
# Soft delete specified session peers by setting left_at timestamp
|
|
update_stmt = (
|
|
update(models.SessionPeer)
|
|
.where(
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
models.SessionPeer.left_at.is_(None), # Only update active peers
|
|
)
|
|
.values(left_at=func.now())
|
|
)
|
|
result = await db.execute(update_stmt)
|
|
|
|
# Get or create peers
|
|
peers_result = await get_or_create_peers(
|
|
db,
|
|
workspace_name=workspace_name,
|
|
peers=[schemas.PeerCreate(name=peer_name) for peer_name in peer_names],
|
|
)
|
|
|
|
# Add new peers to session
|
|
peers = await _get_or_add_peers_to_session(
|
|
db,
|
|
workspace_name=workspace_name,
|
|
session_name=session_name,
|
|
peer_names=peer_names,
|
|
)
|
|
|
|
await db.commit()
|
|
await peers_result.post_commit()
|
|
return peers
|
|
|
|
|
|
async def _get_or_add_peers_to_session(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
peer_names: dict[str, schemas.SessionPeerConfig],
|
|
*,
|
|
fetch_after_upsert: bool = True,
|
|
) -> list[models.SessionPeer]:
|
|
"""
|
|
Upsert session-peer memberships for a session and optionally fetch the
|
|
active memberships afterward.
|
|
|
|
New peers are inserted, peers that previously left the session are rejoined,
|
|
and already-active peers keep their existing session-level configuration.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
peer_names: Mapping of peer names to session-level configuration
|
|
fetch_after_upsert: If True, query and return the active session peers
|
|
after the upsert. If False, skip that read and return an empty list.
|
|
|
|
Returns:
|
|
Active SessionPeer objects after the upsert, or an empty list when the
|
|
post-upsert fetch is skipped
|
|
|
|
Raises:
|
|
ObserverException: If adding peers would exceed the observer limit
|
|
"""
|
|
# If no peers to add, skip the insert and just return existing active session peers
|
|
if not peer_names:
|
|
if not fetch_after_upsert:
|
|
return []
|
|
select_stmt = select(models.SessionPeer).where(
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
models.SessionPeer.left_at.is_(None), # Only active peers
|
|
)
|
|
result = await db.execute(select_stmt)
|
|
return list(result.scalars().all())
|
|
|
|
# Only validate observer limit if we're adding peers with observe_others=True
|
|
new_observer_count = count_observers_in_config(peer_names)
|
|
|
|
if new_observer_count > 0:
|
|
# Use a single efficient query to count existing observers not being updated
|
|
# This uses PostgreSQL's JSONB operators to check the observe_others field directly
|
|
existing_observers_stmt = select(func.count()).where(
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
models.SessionPeer.left_at.is_(None), # Only active peers
|
|
models.SessionPeer.peer_name.notin_(
|
|
peer_names.keys()
|
|
), # Exclude peers being updated
|
|
models.SessionPeer.configuration["observe_others"].astext.cast(
|
|
Boolean
|
|
), # Only observers
|
|
)
|
|
result = await db.execute(existing_observers_stmt)
|
|
existing_observer_count = result.scalar() or 0
|
|
|
|
total_observers = existing_observer_count + new_observer_count
|
|
|
|
if total_observers > settings.SESSION_OBSERVERS_LIMIT:
|
|
raise ObserverException(session_name, total_observers)
|
|
|
|
# Use upsert to handle both new peers and rejoining peers
|
|
stmt = pg_insert(models.SessionPeer).values(
|
|
[
|
|
{
|
|
"session_name": session_name,
|
|
"peer_name": peer_name,
|
|
"workspace_name": workspace_name,
|
|
"joined_at": func.now(),
|
|
"left_at": None,
|
|
"configuration": configuration.model_dump(),
|
|
}
|
|
for peer_name, configuration in peer_names.items()
|
|
]
|
|
)
|
|
|
|
# On conflict, update joined_at and clear left_at (rejoin scenario)
|
|
# If left_at is not None (peer has left the session): Use the new configuration (stmt.excluded.configuration)
|
|
# If left_at is None (peer is still active): Keep the existing configuration (models.SessionPeer.configuration)
|
|
stmt = stmt.on_conflict_do_update(
|
|
index_elements=["session_name", "peer_name", "workspace_name"],
|
|
set_={
|
|
"joined_at": func.now(),
|
|
"left_at": None,
|
|
"configuration": case(
|
|
(models.SessionPeer.left_at.is_not(None), stmt.excluded.configuration),
|
|
else_=models.SessionPeer.configuration,
|
|
),
|
|
},
|
|
)
|
|
await db.execute(stmt)
|
|
|
|
if not fetch_after_upsert:
|
|
return []
|
|
|
|
# Return all active session peers after the upsert
|
|
select_stmt = select(models.SessionPeer).where(
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
models.SessionPeer.left_at.is_(None), # Only active peers
|
|
)
|
|
result = await db.execute(select_stmt)
|
|
return list(result.scalars().all())
|
|
|
|
|
|
async def get_peer_config(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
peer_id: str,
|
|
) -> schemas.SessionPeerConfig:
|
|
"""
|
|
Get the configuration for a peer in a session.
|
|
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
peer_id: Name of the peer
|
|
|
|
|
|
Returns:
|
|
Configuration for the peer
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If the session or peer does not exist
|
|
"""
|
|
# Get row from session_peer table
|
|
stmt = select(models.SessionPeer).where(
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.peer_name == peer_id,
|
|
)
|
|
result = await db.execute(stmt)
|
|
session_peer = result.scalar_one_or_none()
|
|
|
|
if session_peer is None:
|
|
raise ResourceNotFoundException(
|
|
f"Session peer {peer_id} not found in session {session_name} in workspace {workspace_name}"
|
|
)
|
|
|
|
return schemas.SessionPeerConfig(**session_peer.configuration)
|
|
|
|
|
|
async def set_peer_config(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
session_name: str,
|
|
peer_name: str,
|
|
config: schemas.SessionPeerConfig,
|
|
) -> None:
|
|
"""
|
|
Set the configuration for a specific peer in a session.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
session_name: Name of the session
|
|
peer_name: Name of the peer
|
|
config: The peer configuration to set
|
|
|
|
Raises:
|
|
ObserverException: If the update would exceed the observer limit
|
|
"""
|
|
# First, get the session and peer to ensure they exist
|
|
await get_session(db, session_name, workspace_name)
|
|
await get_peer(db, workspace_name, schemas.PeerCreate(name=peer_name))
|
|
|
|
# Check if a SessionPeer entry already exists
|
|
stmt = (
|
|
select(models.SessionPeer)
|
|
.where(models.SessionPeer.session_name == session_name)
|
|
.where(models.SessionPeer.peer_name == peer_name)
|
|
.where(models.SessionPeer.workspace_name == workspace_name)
|
|
)
|
|
result = await db.execute(stmt)
|
|
session_peer = result.scalar_one_or_none()
|
|
|
|
# Check if this update would exceed observer limits
|
|
if config.observe_others:
|
|
# Check if peer is already an observer
|
|
is_currently_observer = (
|
|
session_peer.configuration.get("observe_others", False)
|
|
if session_peer and session_peer.configuration
|
|
else False
|
|
)
|
|
|
|
# Only need to check limit if peer is becoming a new observer
|
|
if not is_currently_observer:
|
|
# Use a single efficient query to count existing observers
|
|
existing_observers_stmt = select(func.count()).where(
|
|
models.SessionPeer.session_name == session_name,
|
|
models.SessionPeer.workspace_name == workspace_name,
|
|
models.SessionPeer.left_at.is_(None), # Only active peers
|
|
models.SessionPeer.peer_name
|
|
!= peer_name, # Exclude the peer being updated
|
|
models.SessionPeer.configuration["observe_others"].astext.cast(
|
|
Boolean
|
|
), # Only observers
|
|
)
|
|
result = await db.execute(existing_observers_stmt)
|
|
observer_count = result.scalar() or 0
|
|
|
|
# Add one for this peer becoming an observer
|
|
observer_count += 1
|
|
|
|
if observer_count > settings.SESSION_OBSERVERS_LIMIT:
|
|
raise ObserverException(session_name, observer_count)
|
|
|
|
update_data = config.model_dump(exclude_none=True)
|
|
|
|
if session_peer:
|
|
# Update existing configuration
|
|
if session_peer.configuration:
|
|
# Create a new dictionary and update it to ensure SQLAlchemy tracks the change
|
|
new_config = session_peer.configuration.copy()
|
|
new_config.update(update_data)
|
|
session_peer.configuration = new_config
|
|
else:
|
|
session_peer.configuration = update_data
|
|
else:
|
|
# Create a new SessionPeer entry
|
|
session_peer = models.SessionPeer(
|
|
session_name=session_name,
|
|
peer_name=peer_name,
|
|
workspace_name=workspace_name,
|
|
configuration=update_data,
|
|
)
|
|
db.add(session_peer)
|
|
|
|
await db.commit()
|