"""Session-lifecycle operations — thin wrappers around SessionRecord. All public functions are async and idempotent where possible: * ``ensure_and_touch_session`` upserts the row and bumps ``last_activity_at`` in a single DB round trip. * ``accumulate_usage`` atomically adds tokens / cost to the session row and, when a model is named, updates the ``session_model_usage`` table so ``cost-by-model`` attributes mixed-model sessions correctly. * ``mark_ended`` transitions to a terminal status. Writes to ``running`` sessions only — terminal sessions (completed / failed) stay frozen so late tool-calls don't resurrect or distort them. The ``abandoned`` transition is never written — it's computed at read time via ``get_effective_status_sql`` against ``SESSION_ABANDON_AFTER_SECONDS`` (defaults to 30 min). """ import os from dataclasses import dataclass from datetime import datetime, timedelta, timezone from enum import Enum from typing import Optional, Sequence from uuid import UUID as UUIDType from sqlalchemy import and_, case, func, or_, select, update from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.dialects.sqlite import insert as sqlite_insert from cognee.infrastructure.databases.relational import get_relational_engine from cognee.shared.logging_utils import get_logger from .models import SessionModelUsage, SessionRecord logger = get_logger("session_lifecycle") class SessionStatus(str, Enum): RUNNING = "running" COMPLETED = "completed" FAILED = "failed" ABANDONED = "abandoned" # computed at read time, not stored # 30 minutes by default; overridable via env var for tests. def _abandon_after_seconds() -> int: raw = os.environ.get("SESSION_ABANDON_AFTER_SECONDS", "") try: return int(raw) if raw else 1800 except ValueError: return 1800 def _dialect_name(bind) -> str: try: return bind.dialect.name except Exception: return "" async def ensure_and_touch_session( *, session_id: str, user_id: UUIDType, dataset_id: Optional[UUIDType] = None, ) -> None: """Upsert the session row in one round trip. Creates the row if absent (status=running). If present AND still running, bumps ``last_activity_at``. Terminal sessions are left untouched so a late straggler can't accidentally resurrect them. Also fills in ``dataset_id`` when currently null. """ now = datetime.now(timezone.utc) engine = get_relational_engine() async with engine.get_async_session() as session: bind = await session.connection() dialect = _dialect_name(bind) values = { "session_id": session_id, "user_id": user_id, "dataset_id": dataset_id, "status": SessionStatus.RUNNING.value, "started_at": now, "last_activity_at": now, "tokens_in": 0, "tokens_out": 0, "cost_usd": 0.0, "error_count": 0, } if dialect in ("sqlite", "postgresql"): insert = sqlite_insert if dialect == "sqlite" else pg_insert stmt = insert(SessionRecord).values(**values) set_ = {"last_activity_at": now} # Back-fill a previously-unset dataset_id. set_["dataset_id"] = case( (SessionRecord.dataset_id.is_(None), dataset_id), else_=SessionRecord.dataset_id, ) stmt = stmt.on_conflict_do_update( index_elements=["session_id", "user_id"], set_=set_, where=SessionRecord.status == SessionStatus.RUNNING.value, ) await session.execute(stmt) await session.commit() return # Portable fallback: SELECT-then-INSERT/UPDATE. Two round trips. existing = ( await session.execute( select(SessionRecord).where( and_( SessionRecord.session_id == session_id, SessionRecord.user_id == user_id, ) ) ) ).scalar_one_or_none() if existing is None: session.add(SessionRecord(**values)) elif existing.status == SessionStatus.RUNNING.value: existing.last_activity_at = now if existing.dataset_id is None and dataset_id is not None: existing.dataset_id = dataset_id await session.commit() async def accumulate_usage( *, session_id: str, user_id: UUIDType, tokens_in: int = 0, tokens_out: int = 0, cost_usd: float = 0.0, model: Optional[str] = None, errored: bool = False, ) -> None: """Atomically add usage counters to the session row + per-model row. Only mutates sessions in ``running`` state — terminal sessions are frozen. Per-model accumulation runs against ``session_model_usage`` via an upsert so mixed-model sessions attribute correctly in ``cost-by-model``. """ if tokens_in == 0 and tokens_out == 0 and cost_usd == 0.0 and not errored and model is None: return engine = get_relational_engine() async with engine.get_async_session() as session: # 1) Session-level aggregate. Gated on running status so a # terminal session doesn't accrue straggler charges. values = {} if tokens_in: values["tokens_in"] = SessionRecord.tokens_in + tokens_in if tokens_out: values["tokens_out"] = SessionRecord.tokens_out + tokens_out if cost_usd: values["cost_usd"] = SessionRecord.cost_usd + cost_usd if errored: values["error_count"] = SessionRecord.error_count + 1 if model: values["last_model"] = model if values: await session.execute( update(SessionRecord) .where( and_( SessionRecord.session_id == session_id, SessionRecord.user_id == user_id, SessionRecord.status == SessionStatus.RUNNING.value, ) ) .values(**values) ) # 2) Per-model row — only when there's usage to credit. if model and (tokens_in or tokens_out or cost_usd): now = datetime.now(timezone.utc) bind = await session.connection() dialect = _dialect_name(bind) if dialect in ("sqlite", "postgresql"): insert = sqlite_insert if dialect == "sqlite" else pg_insert stmt = insert(SessionModelUsage).values( session_id=session_id, user_id=user_id, model=model, tokens_in=tokens_in, tokens_out=tokens_out, cost_usd=cost_usd, updated_at=now, ) stmt = stmt.on_conflict_do_update( index_elements=["session_id", "user_id", "model"], set_={ "tokens_in": SessionModelUsage.tokens_in + tokens_in, "tokens_out": SessionModelUsage.tokens_out + tokens_out, "cost_usd": SessionModelUsage.cost_usd + cost_usd, "updated_at": now, }, ) await session.execute(stmt) else: # Portable fallback: SELECT then INSERT/UPDATE. existing = ( await session.execute( select(SessionModelUsage).where( and_( SessionModelUsage.session_id == session_id, SessionModelUsage.user_id == user_id, SessionModelUsage.model == model, ) ) ) ).scalar_one_or_none() if existing is None: session.add( SessionModelUsage( session_id=session_id, user_id=user_id, model=model, tokens_in=tokens_in, tokens_out=tokens_out, cost_usd=cost_usd, updated_at=now, ) ) else: existing.tokens_in = existing.tokens_in + tokens_in existing.tokens_out = existing.tokens_out + tokens_out existing.cost_usd = existing.cost_usd + cost_usd existing.updated_at = now await session.commit() async def mark_ended( *, session_id: str, user_id: UUIDType, status: SessionStatus, ) -> None: """Transition to a terminal status (completed / failed).""" if status == SessionStatus.RUNNING or status == SessionStatus.ABANDONED: raise ValueError(f"mark_ended requires a terminal status (completed/failed), got {status}") now = datetime.now(timezone.utc) engine = get_relational_engine() async with engine.get_async_session() as session: stmt = ( update(SessionRecord) .where( and_( SessionRecord.session_id == session_id, SessionRecord.user_id == user_id, ) ) .values(status=status.value, ended_at=now) ) await session.execute(stmt) await session.commit() def get_effective_status_sql(): """Return a SQL expression that evaluates to the effective status. Rule: if stored status is ``running`` AND last_activity_at is older than ``SESSION_ABANDON_AFTER_SECONDS`` → ``abandoned``. Otherwise use the stored status verbatim. The threshold is computed in Python and passed as a bound parameter, so this works uniformly on SQLite, Postgres, etc. """ threshold_seconds = _abandon_after_seconds() threshold_ts = datetime.now(timezone.utc) - timedelta(seconds=threshold_seconds) return case( ( and_( SessionRecord.status == SessionStatus.RUNNING.value, SessionRecord.last_activity_at < threshold_ts, ), SessionStatus.ABANDONED.value, ), else_=SessionRecord.status, ) async def get_session_row( *, session_id: str, user_id: UUIDType, user_ids: Optional[list[UUIDType]] = None, permitted_dataset_ids: Optional[list[UUIDType]] = None, prefer_other_owner: bool = False, ) -> Optional[SessionRecord]: """Fetch a session row visible to the caller. Returns the row if the caller (or their child agents, via ``user_ids``) owns the session OR if the session's dataset is in ``permitted_dataset_ids``. Returns None otherwise. The same ``session_id`` can exist under multiple owners (it's only unique per user in the composite PK). When the query matches multiple rows and ``prefer_other_owner`` is True, returns one whose owner is NOT the caller (useful for cache-reads via dataset grants). Otherwise returns the first match. """ engine = get_relational_engine() async with engine.get_async_session() as session: if user_ids is not None and len(user_ids) > 0: visibility_terms = [SessionRecord.user_id.in_(user_ids)] else: visibility_terms = [SessionRecord.user_id == user_id] if permitted_dataset_ids: visibility_terms.append(SessionRecord.dataset_id.in_(permitted_dataset_ids)) result = await session.execute( select(SessionRecord).where( and_( SessionRecord.session_id == session_id, or_(*visibility_terms) if len(visibility_terms) > 1 else visibility_terms[0], ) ) ) rows = list(result.scalars().all()) if not rows: return None if prefer_other_owner: non_owner = [r for r in rows if r.user_id != user_id] if non_owner: return non_owner[0] return rows[0] @dataclass(slots=True) class SessionRowWithStatus: """List envelope for ``list_session_rows`` — the SessionRecord plus the computed effective status, without attaching dynamic attributes to the ORM instance.""" record: SessionRecord effective_status: str def to_dict(self) -> dict: d = self.record.to_dict() d["effective_status"] = self.effective_status return d @dataclass(slots=True) class SessionListPage: """Paginated list envelope.""" sessions: list[SessionRowWithStatus] total: int limit: int offset: int @property def has_more(self) -> bool: return self.offset + len(self.sessions) < self.total async def list_session_rows( *, user_id: Optional[UUIDType] = None, user_ids: Optional[list[UUIDType]] = None, permitted_dataset_ids: Optional[list[UUIDType]] = None, since: Optional[datetime] = None, status_filter: Optional[str] = None, limit: int = 50, offset: int = 0, order_by: str = "last_activity_at", descending: bool = True, ) -> SessionListPage: """List sessions with pagination metadata. Visibility: returns sessions the caller owns (or their child agents own, via ``user_ids``) OR sessions whose ``dataset_id`` is in ``permitted_dataset_ids`` (read permission granted at the dataset level). status_filter accepts the effective status (including ``abandoned``) — the SQL predicate handles the abandoned-by-time inference. """ engine = get_relational_engine() async with engine.get_async_session() as session: eff = get_effective_status_sql() # Ownership / permission predicate. visibility_terms = [] if user_ids is not None and len(user_ids) > 0: visibility_terms.append(SessionRecord.user_id.in_(user_ids)) elif user_id is not None: visibility_terms.append(SessionRecord.user_id == user_id) if permitted_dataset_ids: visibility_terms.append(SessionRecord.dataset_id.in_(permitted_dataset_ids)) filters = [] if visibility_terms: filters.append( or_(*visibility_terms) if len(visibility_terms) > 1 else visibility_terms[0] ) if since is not None: filters.append(SessionRecord.last_activity_at >= since) if status_filter: filters.append(eff == status_filter) # Count before pagination so the caller can render # "showing N of M". count_stmt = select(func.count()).select_from(SessionRecord) if filters: count_stmt = count_stmt.where(and_(*filters)) total = (await session.execute(count_stmt)).scalar_one() sortable = { "last_activity_at": SessionRecord.last_activity_at, "started_at": SessionRecord.started_at, "ended_at": SessionRecord.ended_at, "cost_usd": SessionRecord.cost_usd, "tokens_in": SessionRecord.tokens_in, "tokens_out": SessionRecord.tokens_out, } sort_col = sortable.get(order_by, SessionRecord.last_activity_at) rows_stmt = select(SessionRecord, eff.label("effective_status")) if filters: rows_stmt = rows_stmt.where(and_(*filters)) rows_stmt = rows_stmt.order_by(sort_col.desc() if descending else sort_col.asc()) rows_stmt = rows_stmt.limit(limit).offset(offset) rows = (await session.execute(rows_stmt)).all() return SessionListPage( sessions=[SessionRowWithStatus(record=r[0], effective_status=r[1]) for r in rows], total=int(total), limit=limit, offset=offset, ) # Backward-compatibility shims ------------------------------------------------ async def ensure_session(*, session_id, user_id, dataset_id=None): """Deprecated: prefer ``ensure_and_touch_session``. Kept for callers that only want the "row must exist" half.""" await ensure_and_touch_session(session_id=session_id, user_id=user_id, dataset_id=dataset_id) async def touch_session(*, session_id, user_id, dataset_id=None): """Deprecated: prefer ``ensure_and_touch_session``.""" await ensure_and_touch_session(session_id=session_id, user_id=user_id, dataset_id=dataset_id) _session_record_write_failed = False async def record_session_activity(user_id: str, session_id: str, *, errored: bool = False) -> None: """Write a lifecycle heartbeat for a session: upsert + touch the SessionRecord row. Accepts a string ``user_id`` (coerced to UUID). Swallows failures — the session_records table is optional for SessionManager correctness — but logs once at WARNING per process so silent breakage stays visible in ops. """ global _session_record_write_failed try: try: user_uuid = UUIDType(str(user_id)) except (ValueError, TypeError): return await ensure_and_touch_session(session_id=session_id, user_id=user_uuid) if errored: await accumulate_usage(session_id=session_id, user_id=user_uuid, errored=True) except Exception as exc: if not _session_record_write_failed: _session_record_write_failed = True logger.warning( "session_records write failed (%s); subsequent failures will log at debug. " "Check alembic migrations for the session_records table.", exc, ) else: logger.debug("session_records write failed (%s)", exc)