Isolate dataframes and tasks in a chat session the model opens once.

config/runtime.py

@@ -11,6 +11,7 @@
 )
 from config.cache import CachePort, build_cache_from_env
 from config.service_auth import service_caller_token
+from config.session import HttpSessionProvider, InMemorySessionProvider, SessionPort
 from config.storage import (
     DefaultSessionScopeResolver,
     HttpSessionStorageProvider,
@@ -69,6 +70,7 @@ class AppRuntime:
     storage: SessionStoragePort
     scope_resolver: SessionScopeResolverPort
     cache: CachePort
+    sessions: SessionPort
 
     def resolve_user_config(self, ctx: Any) -> dict[str, Any]:
         user_config = dict(self.user_config)
@@ -127,7 +129,24 @@ def _storage_base_url() -> str:
     return os.getenv("PERFECTO_STORAGE_API_BASE_URL", "").strip()
 
 
-def _build_storage_and_cache(transport: Transport) -> tuple[SessionStoragePort, CachePort]:
+def _sessions_for(
+        transport: Transport,
+        storage: SessionStoragePort,
+        storage_base_url: str,
+        caller_token: str,
+) -> SessionPort:
+    """
+    Stdio, and HTTP without a storage URL, keep the session registry in this
+    process and drop a purged session's partition. HTTP with a storage URL
+    uses the storage service ``/sessions`` registry.
+    """
+    if transport == "streamable-http" and storage_base_url:
+        return HttpSessionProvider(base_url=storage_base_url, caller_token=caller_token)
+    on_purge = getattr(storage, "discard_session", None)
+    return InMemorySessionProvider(on_purge=on_purge)
+
+
+def _build_storage_and_cache(transport: Transport) -> tuple[SessionStoragePort, CachePort, str, str]:
     """
     Stdio, and HTTP without a storage URL, keep partitions and the method cache
     in this process. HTTP with ``PERFECTO_STORAGE_API_BASE_URL`` uses the shared
@@ -150,24 +169,27 @@ def _build_storage_and_cache(transport: Transport) -> tuple[SessionStoragePort,
             "streamable-http",
             storage_base_url=storage_base_url,
             caller_token=caller_token,
-        )
+        ), storage_base_url, caller_token
 
     # In-memory cache factory is the stdio path of build_cache_from_env.
-    return InMemorySessionStorageProvider(), build_cache_from_env("stdio")
+    return InMemorySessionStorageProvider(), build_cache_from_env("stdio"), "", ""
 
 
 def build_runtime(
         transport: Transport,
         startup_token: Optional[PerfectoToken] = None,
 ) -> AppRuntime:
     """
-    Compose auth, session storage and the method cache for the selected transport.
+    Compose auth, chat sessions, session storage and the method cache.
 
-    - stdio: process-lifetime ``startup_token`` and in-memory partitions.
-    - streamable-http: request-scoped Bearer auth. Partitions stay in memory
-      unless ``PERFECTO_STORAGE_API_BASE_URL`` points at the storage service.
+    - stdio: process-lifetime ``startup_token``. Sessions and partitions stay
+      in memory. A purged session drops its partition.
+    - streamable-http: request-scoped Bearer auth. Sessions and partitions stay
+      in memory unless ``PERFECTO_STORAGE_API_BASE_URL`` points at the storage
+      service, which then owns ``/sessions`` and ``/session-partitions``.
     """
-    storage, cache = _build_storage_and_cache(transport)
+    storage, cache, storage_base_url, caller_token = _build_storage_and_cache(transport)
+    sessions = _sessions_for(transport, storage, storage_base_url, caller_token)
     if transport == "stdio":
         stdio_user_config = {
             "startup_token": startup_token,
@@ -181,6 +203,7 @@ def build_runtime(
             storage=storage,
             scope_resolver=DefaultSessionScopeResolver(),
             cache=cache,
+            sessions=sessions,
         )
 
     if transport == "streamable-http":
@@ -191,6 +214,7 @@ def build_runtime(
             storage=storage,
             scope_resolver=DefaultSessionScopeResolver(),
             cache=cache,
+            sessions=sessions,
         )
 
     raise ValueError(f"Unknown transport: {transport}")

config/session.py

@@ -0,0 +1,298 @@
+"""Chat sessions that own Perfecto dataframes and tasks."""
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import hmac
+import logging
+import re
+import secrets
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, replace
+from datetime import datetime, timedelta, timezone
+from enum import Enum
+from typing import Any, Callable, Optional
+from urllib.parse import quote
+
+import httpx
+
+from config.env import env_float, env_int
+from config.http_clients import SharedAsyncClient
+from config.ids import SIMPLE_ID_ALPHABET, random_id
+from config.service_auth import service_headers
+
+logger = logging.getLogger(__name__)
+
+# "pfs_" + 32 chars from the simple-id alphabet (160 random bits).
+SESSION_ID_PREFIX = "pfs_"
+SESSION_ID_LENGTH = 32
+SESSION_ID_RE = re.compile(rf"^{SESSION_ID_PREFIX}[{SIMPLE_ID_ALPHABET}]{{{SESSION_ID_LENGTH}}}$")
+SESSION_TOOL_NAME = "perfecto_session"
+
+SESSION_STORAGE_TIMEOUT_SECONDS = env_float("SESSION_STORAGE_TIMEOUT_SECONDS", 15.0, minimum=0.1)
+# stdio lifecycle; the same rules and defaults as a storage-api sweeper.
+SESSION_IDLE_TIMEOUT_SECONDS = env_int("SESSION_IDLE_TIMEOUT_SECONDS", 7 * 24 * 3600, minimum=1)
+SESSION_PURGE_GRACE_SECONDS = env_int("SESSION_PURGE_GRACE_SECONDS", 3600, minimum=0)
+SESSION_SWEEP_INTERVAL_SECONDS = env_int("SESSION_SWEEP_INTERVAL_SECONDS", 60, minimum=1)
+
+
+def generate_session_id() -> str:
+    return f"{SESSION_ID_PREFIX}{random_id(SESSION_ID_LENGTH)}"
+
+
+def is_valid_session_id(value: Any) -> bool:
+    return isinstance(value, str) and bool(SESSION_ID_RE.fullmatch(value))
+
+
+class SessionState(str, Enum):
+    ACTIVE = "ACTIVE"
+    EXPIRED = "EXPIRED"
+
+
+@dataclass(frozen=True)
+class ChatSession:
+    """One conversation's workspace for tasks and dataframes. Expired, then purged."""
+
+    session_id: str
+    owner_id: str
+    state: SessionState
+    created_at: datetime
+    last_seen_at: datetime
+    expired_at: Optional[datetime] = None
+
+
+class SessionErrorCode(str, Enum):
+    REQUIRED = "SESSION_REQUIRED"
+    INVALID = "SESSION_INVALID"
+    EXPIRED = "SESSION_EXPIRED"
+    UNAVAILABLE = "SESSION_UNAVAILABLE"
+    SERVICE_ERROR = "SESSION_SERVICE_ERROR"
+
+
+_RECOVERY_HINT = (
+    f"Call {SESSION_TOOL_NAME} with action 'get' to obtain a new session_id, use it in every "
+    "following tool call, and re-run the tool calls that produced any task or dataframe you "
+    "still need: ids from a previous session are not available in the new one."
+)
+
+SESSION_ERROR_MESSAGES = {
+    SessionErrorCode.REQUIRED: (
+        f"'session_id' is a required argument. Call {SESSION_TOOL_NAME} with action 'get' once "
+        "at the start of this conversation and pass the returned session_id in every tool call."
+    ),
+    SessionErrorCode.INVALID: f"Invalid session_id. {_RECOVERY_HINT}",
+    SessionErrorCode.EXPIRED: f"This session expired. {_RECOVERY_HINT}",
+    SessionErrorCode.UNAVAILABLE: "The session service is unavailable right now. Retry the same call in a moment.",
+    SessionErrorCode.SERVICE_ERROR: (
+        "The session service rejected this server's request. This is a server configuration "
+        "problem and retrying will not help; report it to the server administrator."
+    ),
+}
+
+
+class SessionError(Exception):
+    def __init__(self, code: SessionErrorCode):
+        super().__init__(SESSION_ERROR_MESSAGES[code])
+        self.code = code
+        self.message = SESSION_ERROR_MESSAGES[code]
+
+
+class SessionPort(ABC):
+    """
+    Chat-session registry. Lookups are scoped to the owner and to the credential
+    the session was created with (the Perfecto security token). A session owned
+    by someone else, or used with another token, is reported exactly like an
+    unknown one (INVALID). Failures are SessionError, with UNAVAILABLE when the
+    registry cannot answer.
+    """
+
+    @abstractmethod
+    async def open(self, owner_id: str, credential: str) -> ChatSession:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def touch(self, session_id: str, owner_id: str, credential: str) -> ChatSession:
+        """Validate an ACTIVE session of ``owner_id`` + ``credential`` and record the keep-alive."""
+        raise NotImplementedError
+
+
+def _utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+@dataclass
+class _Entry:
+    session: ChatSession
+    # HMAC of the creating credential (process-local key; the credential is never kept).
+    credential_hash: bytes
+    # First "expired, renegotiate" answer: the purge grace starts here.
+    expired_notified_at: Optional[datetime] = None
+
+    def purgeable(self, now: datetime, grace: timedelta) -> bool:
+        since = self.expired_notified_at or self.session.expired_at
+        return since is not None and now - since >= grace
+
+
+class InMemorySessionProvider(SessionPort):
+    """
+    Process-local registry used by stdio and by HTTP when no storage URL is set.
+
+    ACTIVE → EXPIRED after the idle timeout → purged once the grace has passed
+    since the first "expired" answer (or since expiry if nobody asked).
+    ``on_purge(owner_id, session_id)`` releases what the session held (its
+    partition). Sweeps run on open/touch, at most once per sweep interval.
+    """
+
+    def __init__(
+            self,
+            idle_timeout_seconds: float = SESSION_IDLE_TIMEOUT_SECONDS,
+            purge_grace_seconds: float = SESSION_PURGE_GRACE_SECONDS,
+            sweep_interval_seconds: float = SESSION_SWEEP_INTERVAL_SECONDS,
+            on_purge: Optional[Callable[[str, str], None]] = None,
+            clock: Callable[[], datetime] = _utc_now,
+    ) -> None:
+        self._entries: dict[str, _Entry] = {}
+        self._hash_key = secrets.token_bytes(32)
+        self._lock = asyncio.Lock()
+        self._idle_timeout = timedelta(seconds=idle_timeout_seconds)
+        self._purge_grace = timedelta(seconds=purge_grace_seconds)
+        self._sweep_interval = timedelta(seconds=sweep_interval_seconds)
+        self._on_purge = on_purge
+        self._clock = clock
+        self._last_sweep: Optional[datetime] = None
+
+    def _credential_hash(self, credential: str) -> bytes:
+        return hmac.new(self._hash_key, credential.encode("utf-8"), hashlib.sha256).digest()
+
+    def _expire_if_idle(self, entry: _Entry, now: datetime) -> None:
+        session = entry.session
+        if session.state is SessionState.ACTIVE and now - session.last_seen_at >= self._idle_timeout:
+            entry.session = replace(session, state=SessionState.EXPIRED, expired_at=now)
+
+    def _sweep(self, now: datetime) -> None:
+        if self._last_sweep is not None and now - self._last_sweep < self._sweep_interval:
+            return
+        self._last_sweep = now
+        for session_id, entry in list(self._entries.items()):
+            self._expire_if_idle(entry, now)
+            if entry.session.state is SessionState.EXPIRED and entry.purgeable(now, self._purge_grace):
+                del self._entries[session_id]
+                if self._on_purge is not None:
+                    self._on_purge(entry.session.owner_id, session_id)
+
+    async def open(self, owner_id: str, credential: str) -> ChatSession:
+        async with self._lock:
+            now = self._clock()
+            self._sweep(now)
+            session_id = generate_session_id()
+            while session_id in self._entries:
+                session_id = generate_session_id()
+            session = ChatSession(
+                session_id=session_id,
+                owner_id=owner_id,
+                state=SessionState.ACTIVE,
+                created_at=now,
+                last_seen_at=now,
+            )
+            self._entries[session_id] = _Entry(session, self._credential_hash(credential))
+            return session
+
+    async def touch(self, session_id: str, owner_id: str, credential: str) -> ChatSession:
+        async with self._lock:
+            now = self._clock()
+            self._sweep(now)
+            entry = self._entries.get(session_id)
+            if (
+                entry is None
+                or entry.session.owner_id != owner_id
+                or not hmac.compare_digest(entry.credential_hash, self._credential_hash(credential))
+            ):
+                raise SessionError(SessionErrorCode.INVALID)
+            self._expire_if_idle(entry, now)
+            if entry.session.state is SessionState.EXPIRED:
+                if entry.expired_notified_at is None:
+                    entry.expired_notified_at = now
+                raise SessionError(SessionErrorCode.EXPIRED)
+            entry.session = replace(entry.session, last_seen_at=now)
+            return entry.session
+
+
+_TRANSIENT_STATUSES = frozenset({408, 429})
+
+
+class HttpSessionProvider(SessionPort):
+    """Hosted: the storage service owns the session registry (``/sessions``)."""
+
+    def __init__(
+            self,
+            base_url: str,
+            caller_token: str,
+            timeout_seconds: float = SESSION_STORAGE_TIMEOUT_SECONDS,
+            transport: Optional[httpx.AsyncBaseTransport] = None,
+    ) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._caller_token = caller_token
+        if transport is not None:
+            self._http = SharedAsyncClient(
+                lambda: httpx.AsyncClient(transport=transport, timeout=timeout_seconds)
+            )
+        else:
+            self._http = SharedAsyncClient(
+                lambda: httpx.AsyncClient(http2=True, timeout=timeout_seconds)
+            )
+
+    def _headers(self, credential: str) -> dict[str, str]:
+        # Explicit credential: open/touch run before the session is bound to the context.
+        return service_headers(self._caller_token, credential)
+
+    async def _post(self, path: str, owner_id: str, credential: str) -> httpx.Response:
+        try:
+            return await self._http.get().post(
+                f"{self._base_url}{path}",
+                headers=self._headers(credential),
+                json={"owner_id": owner_id},
+            )
+        except httpx.HTTPError as exc:
+            logger.warning("session storage %s unreachable: %s", path, type(exc).__name__)
+            raise SessionError(SessionErrorCode.UNAVAILABLE) from exc
+
+    @staticmethod
+    def _session_or_unavailable(response: httpx.Response) -> ChatSession:
+        status_code = response.status_code
+        if status_code in _TRANSIENT_STATUSES or status_code >= 500:
+            logger.warning("session storage answered %s", status_code)
+            raise SessionError(SessionErrorCode.UNAVAILABLE)
+        if status_code >= 400:
+            logger.error("session storage rejected the MCP request with %s", status_code)
+            raise SessionError(SessionErrorCode.SERVICE_ERROR)
+        try:
+            return _session_from_payload(response.json())
+        except (ValueError, KeyError, TypeError, AttributeError) as exc:
+            logger.warning("session storage returned an invalid session payload")
+            raise SessionError(SessionErrorCode.UNAVAILABLE) from exc
+
+    async def open(self, owner_id: str, credential: str) -> ChatSession:
+        response = await self._post("/sessions", owner_id, credential)
+        return self._session_or_unavailable(response)
+
+    async def touch(self, session_id: str, owner_id: str, credential: str) -> ChatSession:
+        response = await self._post(
+            f"/sessions/{quote(session_id, safe='')}/touch", owner_id, credential
+        )
+        if response.status_code == 404:
+            raise SessionError(SessionErrorCode.INVALID)
+        if response.status_code == 410:
+            raise SessionError(SessionErrorCode.EXPIRED)
+        return self._session_or_unavailable(response)
+
+
+def _session_from_payload(payload: dict[str, Any]) -> ChatSession:
+    expired_at = payload.get("expired_at")
+    return ChatSession(
+        session_id=str(payload["id"]),
+        owner_id=str(payload["owner_id"]),
+        state=SessionState(payload["state"]),
+        created_at=datetime.fromisoformat(payload["created_at"]),
+        last_seen_at=datetime.fromisoformat(payload["last_seen_at"]),
+        expired_at=datetime.fromisoformat(expired_at) if expired_at else None,
+    )