From 10c7f60e810236c794a60fe69121620683445931 Mon Sep 17 00:00:00 2001 From: diego-ferrand Date: Thu, 5 Nov 2026 11:46:35 -0300 Subject: [PATCH] Isolate dataframes and tasks in a chat session the model opens once. --- config/runtime.py | 40 ++- config/session.py | 298 ++++++++++++++++++++++ config/session_context.py | 27 +- config/storage.py | 4 +- docs/hosted-http.md | 11 +- main.py | 1 + models/result.py | 8 + server.py | 4 + tests/test_ai_scriptless_manager.py | 22 +- tests/test_dataframes_tasks_cache.py | 2 + tests/test_session_manager.py | 366 +++++++++++++++++++++++++++ tools/runtime_tools.py | 142 ++++++++--- tools/session_gate.py | 109 ++++++++ tools/session_manager.py | 92 +++++++ 14 files changed, 1056 insertions(+), 70 deletions(-) create mode 100644 config/session.py create mode 100644 tests/test_session_manager.py create mode 100644 tools/session_gate.py create mode 100644 tools/session_manager.py diff --git a/config/runtime.py b/config/runtime.py index 5a69c84..6f16956 100644 --- a/config/runtime.py +++ b/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,10 +169,10 @@ 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( @@ -161,13 +180,16 @@ def build_runtime( 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}") diff --git a/config/session.py b/config/session.py new file mode 100644 index 0000000..5bd5569 --- /dev/null +++ b/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, + ) diff --git a/config/session_context.py b/config/session_context.py index ae4de70..c205788 100644 --- a/config/session_context.py +++ b/config/session_context.py @@ -1,28 +1,25 @@ -"""Request-scoped identity and session used to partition dataframes and tasks. +"""Request-scoped identity and chat session used to partition dataframes and tasks. -Perfecto does not ask the model for a separate chat session id. The tool -entrypoint binds a scope derived from the security token and the MCP session -before any manager code runs. Storage then reads that scope from these -context vars, the same way the BlazeMeter MCP reads a validated chat session. +The tool entrypoint binds these after the security token is present and, for +every tool except ``perfecto_session``, after ``session_id`` is an ACTIVE +session of that caller. Storage reads the scope from these context vars. """ from __future__ import annotations import contextvars from contextlib import contextmanager from dataclasses import dataclass +from datetime import datetime, timezone from typing import Iterator, Optional +from config.session import ChatSession, SessionState + @dataclass(frozen=True) class Identity: user_id: str -@dataclass(frozen=True) -class ChatSession: - session_id: str - - _current_identity: contextvars.ContextVar[Optional[Identity]] = contextvars.ContextVar( "perfecto_current_identity", default=None ) @@ -71,5 +68,13 @@ def bind_session_context( @contextmanager def bind_call_scope(user_id: str, session_id: str, credential: Optional[str] = None) -> Iterator[None]: """Bind the partition key for one Perfecto tool call.""" - with bind_session_context(Identity(user_id), ChatSession(session_id), credential): + now = datetime.now(timezone.utc) + session = ChatSession( + session_id=session_id, + owner_id=user_id, + state=SessionState.ACTIVE, + created_at=now, + last_seen_at=now, + ) + with bind_session_context(Identity(user_id), session, credential): yield diff --git a/config/storage.py b/config/storage.py index 769ec9a..9a260f3 100644 --- a/config/storage.py +++ b/config/storage.py @@ -100,8 +100,8 @@ class DefaultSessionScopeResolver(SessionScopeResolverPort): Resolve the partition from the scope bound for the current tool call. The tool entrypoint binds a user id (hash of cloud plus security token) and - the MCP session id before manager code runs. Outside that call this fails - closed instead of writing into a shared partition. + the chat session id from ``perfecto_session`` before manager code runs. + Outside that call this fails closed instead of writing into a shared partition. """ def resolve(self) -> SessionScope: diff --git a/docs/hosted-http.md b/docs/hosted-http.md index db567d5..62b25d4 100644 --- a/docs/hosted-http.md +++ b/docs/hosted-http.md @@ -72,16 +72,19 @@ docker run --rm -p 8000:8000 \ | `PERFECTO_MCP_STORAGE_CALLER_TOKEN` | Bearer token the MCP sends to the storage service. Required when `PERFECTO_STORAGE_API_BASE_URL` is set. | *(unset)* | | `PERFECTO_MCP_CACHE_ENABLED` | Method cache on or off | `true` | | `PERFECTO_MCP_MAX_PARALLEL_TASKS` | Async tasks running at once in this process | `10` | -| `PERFECTO_MCP_SESSION_STORAGE_TIMEOUT_SECONDS` | Timeout for storage partition calls | `15` | +| `PERFECTO_MCP_SESSION_STORAGE_TIMEOUT_SECONDS` | Timeout for storage partition and hosted session calls | `15` | +| `PERFECTO_MCP_SESSION_IDLE_TIMEOUT_SECONDS` | In-process session idle time before it expires | `604800` (7 days) | +| `PERFECTO_MCP_SESSION_PURGE_GRACE_SECONDS` | How long an expired in-process session is kept after the first expired answer | `3600` | +| `PERFECTO_MCP_SESSION_SWEEP_INTERVAL_SECONDS` | Minimum time between in-process session sweeps | `60` | ## Session dataframes, tasks and cache Large tool results are stored as a session dataframe instead of being pasted into the model context. Slow calls return a task id; the model polls `perfecto_tools` (`tasks_status`, then `tasks_get`) and reads rows with `dataframes_query`. -Partitions are keyed by a hash of the security token plus cloud, and by the MCP session id (`mcp-session-id`, otherwise `local` for stdio). Two credentials never share dataframes. +Call `perfecto_session` action `get` once per conversation and pass the returned `session_id` on every later tool call. Partitions are keyed by a hash of the security token plus cloud, and by that chat session id. Two credentials never share dataframes, and two conversations of the same credential do not either. An idle session expires, then its partition is dropped. `SESSION_REQUIRED`, `SESSION_INVALID` and `SESSION_EXPIRED` mean the model must call `perfecto_session` `get` again. -- Stdio and HTTP without a storage URL keep partitions and the method cache in memory. Restarting the process drops them. One HTTP process can serve that data; another process cannot see it. -- HTTP with `PERFECTO_STORAGE_API_BASE_URL` writes partitions and cache entries to the storage service. Set `PERFECTO_MCP_STORAGE_CALLER_TOKEN` to the caller token that service expects. The service must be reachable at startup (`GET /health`). +- Stdio and HTTP without a storage URL keep sessions, partitions and the method cache in memory. Restarting the process drops them. One HTTP process can serve that data; another process cannot see it. +- HTTP with `PERFECTO_STORAGE_API_BASE_URL` writes sessions (`/sessions`), partitions and cache entries to the storage service. Set `PERFECTO_MCP_STORAGE_CALLER_TOKEN` to the caller token that service expects. The service must be reachable at startup (`GET /health`). - The method cache (`@ttl_cache_method`) stores help pages and other repeated reads. A cache failure is a miss: the tool still runs. User-scoped entries are keyed by the same credential hash as partitions. ## Limitations diff --git a/main.py b/main.py index 73e5cc7..7f0dfa7 100644 --- a/main.py +++ b/main.py @@ -156,6 +156,7 @@ def build_mcp_server( - **Golden rule**: If you're not 100% certain about something related to Perfecto, consult Skills or Help first, and if you need to search online, prioritize help.perfecto.io. ## Important Guidelines +- **Chat session**: Call `perfecto_session` action `get` once at the start of the conversation. Pass the returned `session_id` on every later Perfecto tool call (next to `action`, or inside `args`). Tasks and dataframes belong to that session. If a call returns `SESSION_REQUIRED`, `SESSION_INVALID` or `SESSION_EXPIRED`, get a new session and re-run the calls whose results you still need. - **Batch Operations**: When making multiple calls to the same tool, check if that tool supports a `batch` action and use it instead of separate calls. - **Large results**: Results that are too large for the conversation are stored as a session dataframe. Use `perfecto_tools` actions `dataframes_list`, `dataframes_get`, `dataframes_schema_groups` and `dataframes_query` to inspect them. Every SQL query must include ORDER BY, LIMIT and OFFSET. Call `dataframes_sql_help` before writing SQL. Remove dataframes with `dataframes_remove` or `dataframes_clear` when finished. - **Long-running calls**: If a tool returns a task snapshot instead of the final payload, poll `perfecto_tools` action `tasks_status`, then `tasks_get` for the result. Do not start the same operation again while its task is still working. diff --git a/models/result.py b/models/result.py index d8c3085..c21f9d8 100644 --- a/models/result.py +++ b/models/result.py @@ -8,6 +8,14 @@ class BaseResult(BaseModel): total: Optional[int] = Field(description="Total available records", default=None) has_more: Optional[bool] = Field(description="More records are available", default=None) error: Optional[str] = Field(description="Error message", default=None) + error_code: Optional[str] = Field( + description="Machine-readable error code (e.g. SESSION_REQUIRED, SESSION_INVALID, SESSION_EXPIRED)", + default=None, + ) + session_id: Optional[str] = Field( + description="Chat session this call ran in; keep passing it as session_id in every tool call", + default=None, + ) info: Optional[List[str]] = Field(description="Info messages", default=None) warning: Optional[List[str]] = Field(description="Warning messages", default=None) tool_call_started_at: Optional[str] = Field(description="ISO timestamp when tool action started", default=None) diff --git a/server.py b/server.py index 7f25ebc..ba05920 100644 --- a/server.py +++ b/server.py @@ -5,6 +5,8 @@ from tools.device_manager import register as register_device_manager from tools.execution_manager import register as register_execution_manager from tools.help_manager import register as register_help_manager +from tools.session_gate import annotate_session_requirements +from tools.session_manager import register as register_session_manager from tools.skills_manager import register as register_skills_manager from tools.tools_manager import register as register_tools_manager from tools.user_manager import register as register_user_manager @@ -21,6 +23,7 @@ def register_tools(mcp, runtime: AppRuntime): configure_task_storage(runtime.storage) if runtime.cache is not None: configure_cache(runtime.cache) + register_session_manager(mcp, runtime) register_user_manager(mcp, runtime) register_device_manager(mcp, runtime) register_execution_manager(mcp, runtime) @@ -28,3 +31,4 @@ def register_tools(mcp, runtime: AppRuntime): register_skills_manager(mcp, runtime) register_ai_scriptless_manager(mcp, runtime) register_tools_manager(mcp, runtime) + annotate_session_requirements(mcp) diff --git a/tests/test_ai_scriptless_manager.py b/tests/test_ai_scriptless_manager.py index 2453399..6cc60b4 100644 --- a/tests/test_ai_scriptless_manager.py +++ b/tests/test_ai_scriptless_manager.py @@ -23,6 +23,7 @@ import pytest from config.perfecto import SUPPORT_MESSAGE +from config.call_scope import user_id_for_token from models.result import BaseResult from tools import ai_scriptless_manager from tools.ai_scriptless import definitions @@ -1846,10 +1847,17 @@ def test_dispatcher_routes_registered_action(self, perfecto_token, monkeypatch, assert "not found in AI Scriptless manager tool" not in (result.error or "") -def _call_tool(tool, action, args, ctx=None): +async def _call_tool(tool, action, args, ctx=None): if ctx is None: - ctx = make_ctx() - return tool(arguments={"action": action, "args": args}, ctx=ctx) + ctx = make_ctx(tool._perfecto_token) + session = await tool._perfecto_runtime.sessions.open( + user_id_for_token(tool._perfecto_token), + tool._perfecto_token.token, + ) + payload = {"action": action, "session_id": session.session_id} + if args is not None: + payload["args"] = args + return await tool(arguments=payload, ctx=ctx) def _register_tool(token): @@ -1864,8 +1872,12 @@ def decorator(fn): return decorator mcp = _McpStub() - ai_scriptless_manager.register(mcp, build_runtime("stdio", startup_token=token)) - return mcp.tools["perfecto_ai_scriptless"] + runtime = build_runtime("stdio", startup_token=token) + ai_scriptless_manager.register(mcp, runtime) + tool = mcp.tools["perfecto_ai_scriptless"] + tool._perfecto_runtime = runtime + tool._perfecto_token = token + return tool class TestCmdArgumentsValidation: diff --git a/tests/test_dataframes_tasks_cache.py b/tests/test_dataframes_tasks_cache.py index 0163f8e..2f6c9a0 100644 --- a/tests/test_dataframes_tasks_cache.py +++ b/tests/test_dataframes_tasks_cache.py @@ -197,10 +197,12 @@ def test_runtime_uses_memory_until_a_storage_url_is_configured(monkeypatch): monkeypatch.delenv("PERFECTO_MCP_STORAGE_CALLER_TOKEN", raising=False) runtime = build_runtime("stdio", startup_token=PerfectoToken("token-a", "demo")) assert runtime.storage.__class__.__name__ == "InMemorySessionStorageProvider" + assert runtime.sessions.__class__.__name__ == "InMemorySessionProvider" assert runtime.cache.__class__.__name__ == "InMemoryTTLCache" http_runtime = build_runtime("streamable-http") assert http_runtime.storage.__class__.__name__ == "InMemorySessionStorageProvider" + assert http_runtime.sessions.__class__.__name__ == "InMemorySessionProvider" monkeypatch.setenv("PERFECTO_STORAGE_API_BASE_URL", "http://storage.internal") try: diff --git a/tests/test_session_manager.py b/tests/test_session_manager.py new file mode 100644 index 0000000..8562281 --- /dev/null +++ b/tests/test_session_manager.py @@ -0,0 +1,366 @@ +"""Chat session registry and the tool gate that requires session_id.""" +import asyncio +import json +from dataclasses import replace +from datetime import datetime, timedelta, timezone + +import httpx +import pytest + +from config.call_scope import user_id_for_token +from config.ids import SIMPLE_ID_ALPHABET +from config.runtime import build_runtime +from config.session import ( + SESSION_ID_PREFIX, + SESSION_TOOL_NAME, + HttpSessionProvider, + InMemorySessionProvider, + SessionError, + SessionErrorCode, + SessionState, + generate_session_id, + is_valid_session_id, +) +from config.session_context import ( + Identity, + SessionContextMissing, + bind_session_context, + current_identity, + current_session, +) +from config.storage import InMemorySessionStorageProvider, SessionPartitionPayload, SessionScope +from config.token import PerfectoToken +from models.result import BaseResult +from tests.conftest import make_ctx +from tools.dataframe_manager import list_dataframes_metadata, register_dataframe +from tools.runtime_tools import run_tool_with_runtime +from tools.session_gate import SESSION_HINT +from tools.session_manager import SessionManager, register as register_session_tool +from tools.utils import normalize_action_args + +ALICE = PerfectoToken("alice-token", "demo") +BOB = PerfectoToken("bob-token", "demo") +ALICE_OTHER_CLOUD = PerfectoToken("alice-token", "other") + + +class _Clock: + def __init__(self): + self.now = datetime(2026, 1, 1, tzinfo=timezone.utc) + + def __call__(self): + return self.now + + def advance(self, seconds): + self.now += timedelta(seconds=seconds) + + +def _provider(clock, purged=None, idle=100, grace=10): + return InMemorySessionProvider( + idle_timeout_seconds=idle, + purge_grace_seconds=grace, + sweep_interval_seconds=1, + on_purge=(lambda owner, sid: purged.append((owner, sid))) if purged is not None else None, + clock=clock, + ) + + +def test_keep_alive_postpones_idle_expiry(): + clock = _Clock() + provider = _provider(clock) + + async def scenario(): + session = await provider.open("alice", "secret") + clock.advance(90) + await provider.touch(session.session_id, "alice", "secret") + clock.advance(90) + touched = await provider.touch(session.session_id, "alice", "secret") + assert touched.state is SessionState.ACTIVE + + asyncio.run(scenario()) + + +def test_idle_session_expires_then_purges_after_the_grace_with_its_partition(): + clock = _Clock() + purged = [] + provider = _provider(clock, purged) + + async def scenario(): + session = await provider.open("alice", "secret") + clock.advance(100) + await provider.open("bob", "other") + clock.advance(8) + with pytest.raises(SessionError) as expired: + await provider.touch(session.session_id, "alice", "secret") + assert expired.value.code is SessionErrorCode.EXPIRED + assert purged == [] + + clock.advance(10) + with pytest.raises(SessionError) as gone: + await provider.touch(session.session_id, "alice", "secret") + assert gone.value.code is SessionErrorCode.INVALID + assert purged == [("alice", session.session_id)] + + asyncio.run(scenario()) + + +def test_stdio_runtime_releases_the_partition_of_a_purged_session(): + runtime = build_runtime("stdio", startup_token=ALICE) + scope = SessionScope(user_id="alice", mcp_session_id="pfs_" + "a" * 32) + + async def scenario(): + await runtime.storage.put_partition(scope, SessionPartitionPayload(metadata={"x": 1})) + runtime.sessions._on_purge("alice", scope.mcp_session_id) + assert await runtime.storage.get_partition(scope) is None + + asyncio.run(scenario()) + + +def test_session_ids_use_the_shared_alphabet(): + session_id = generate_session_id() + assert session_id.startswith(SESSION_ID_PREFIX) + assert set(session_id[len(SESSION_ID_PREFIX):]) <= set(SIMPLE_ID_ALPHABET) + assert is_valid_session_id(session_id) + for ambiguous in "ilou": + assert not is_valid_session_id(SESSION_ID_PREFIX + ambiguous * 32) + + +def _http_provider(handler): + return HttpSessionProvider( + "https://storage.internal", "caller", transport=httpx.MockTransport(handler) + ) + + +@pytest.mark.parametrize("status,code", [(404, SessionErrorCode.INVALID), (410, SessionErrorCode.EXPIRED)]) +def test_http_session_touch_keeps_invalid_and_expired(status, code): + provider = _http_provider(lambda request: httpx.Response(status)) + with pytest.raises(SessionError) as failed: + asyncio.run(provider.touch("pfs_" + "a" * 32, "alice", "secret")) + assert failed.value.code is code + + +def test_http_provider_maps_storage_api_contract(): + requests = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + body = json.loads(request.content) + record = { + "id": "pfs_" + "b" * 32, + "owner_id": body["owner_id"], + "state": "ACTIVE", + "created_at": "2026-09-01T00:00:00+00:00", + "last_seen_at": "2026-09-01T00:00:00+00:00", + "expired_at": None, + } + if request.url.path == "/sessions": + return httpx.Response(201, json=record) + if "unknown" in body["owner_id"]: + return httpx.Response(404, json={"detail": "Session not found."}) + if "expired" in body["owner_id"]: + return httpx.Response(410, json={"detail": "Session expired."}) + return httpx.Response(200, json=record) + + provider = HttpSessionProvider("http://storage", "mcp-caller", transport=httpx.MockTransport(handler)) + + async def scenario(): + opened = await provider.open("owner-alice", "secret") + assert opened.session_id == "pfs_" + "b" * 32 + assert (await provider.touch(opened.session_id, "owner-alice", "secret")).owner_id == "owner-alice" + for owner, code in (("unknown-user", SessionErrorCode.INVALID), ("expired-user", SessionErrorCode.EXPIRED)): + with pytest.raises(SessionError) as failed: + await provider.touch(opened.session_id, owner, "secret") + assert failed.value.code is code + + asyncio.run(scenario()) + assert all(request.headers["authorization"] == "Bearer mcp-caller" for request in requests) + assert all(request.headers["x-perfecto-credential"] == "secret" for request in requests) + assert requests[1].url.path == f"/sessions/{'pfs_' + 'b' * 32}/touch" + + +class _TokenAuth: + def __init__(self): + self.token = None + + def get_token(self, ctx): + return self.token + + +class _FakeMcp: + def __init__(self): + self.tools = {} + self.descriptions = {} + + def tool(self, *, name, description): + def decorator(fn): + self.tools[name] = fn + self.descriptions[name] = description + return fn + return decorator + + +@pytest.fixture +def harness(): + auth = _TokenAuth() + storage = InMemorySessionStorageProvider() + sessions = InMemorySessionProvider(on_purge=storage.discard_session) + runtime = replace( + build_runtime("stdio", startup_token=ALICE), + auth=auth, + storage=storage, + sessions=sessions, + ) + mcp = _FakeMcp() + register_session_tool(mcp, runtime) + seen = [] + + @mcp.tool(name="probe", description="Probe tool.") + async def probe(arguments=None, ctx=None): + action, args = normalize_action_args(arguments) + runtime.configure_context(ctx) + + async def _dispatch(): + match action: + case "whoami": + seen.append((current_identity(), current_session())) + return BaseResult(result=[{"ok": True}]) + case "store": + from config.storage import resolve_session_scope + scope = resolve_session_scope() + await register_dataframe( + result=[{"row": 1}], + origin_manager="probe", + origin_action="store", + json_size_chars=10, + session_storage=storage, + scope=scope, + ) + return BaseResult(result=[{"stored": True}]) + case "count": + from config.storage import resolve_session_scope + listed = await list_dataframes_metadata(storage, resolve_session_scope()) + return BaseResult(result=[{"dataframes": len(listed)}]) + case "batch": + return await mcp.tools["probe"]( + {"action": "whoami", "args": args.get("sub_args", {})}, + ctx, + ) + return BaseResult(error="unknown") + + return await run_tool_with_runtime( + runtime, "probe", action, ctx, _dispatch, tool_args=args, + ) + + def call(tool, action, token, **args): + auth.token = token + result = asyncio.run(mcp.tools[tool]( + {"action": action, "args": args}, + make_ctx(token), + )) + return result.model_dump() + + def open_session(token=ALICE): + body = call(SESSION_TOOL_NAME, "get", token) + return body["session_id"] + + return { + "call": call, + "open_session": open_session, + "sessions": sessions, + "seen": seen, + "mcp": mcp, + } + + +def test_session_tool_returns_a_new_session_for_the_caller(harness): + body = harness["call"](SESSION_TOOL_NAME, "get", ALICE) + assert body.get("error") is None + assert is_valid_session_id(body["session_id"]) + assert body["result"][0]["session_id"] == body["session_id"] + assert harness["open_session"]() != body["session_id"] + + +def test_session_tool_rejects_a_missing_token(harness): + body = harness["call"](SESSION_TOOL_NAME, "get", None) + assert "Perfecto Security Token not set" in body["error"] + assert "session_id" not in body + + +def test_missing_session_id_is_required_and_points_to_the_session_tool(harness): + body = harness["call"]("probe", "whoami", ALICE) + assert body["error_code"] == SessionErrorCode.REQUIRED.value + assert SESSION_TOOL_NAME in body["error"] + assert harness["seen"] == [] + + +@pytest.mark.parametrize("bad", ["abc", "pfs_short", "PFS_" + "a" * 32, 123]) +def test_malformed_session_id_is_invalid(harness, bad): + body = harness["call"]("probe", "whoami", ALICE, session_id=bad) + assert body["error_code"] == SessionErrorCode.INVALID.value + + +def test_unknown_and_foreign_sessions_are_indistinguishable(harness): + alice_session = harness["open_session"](ALICE) + foreign = harness["call"]("probe", "whoami", BOB, session_id=alice_session) + unknown = harness["call"]("probe", "whoami", BOB, session_id="pfs_" + "0" * 32) + other_cloud = harness["call"]("probe", "whoami", ALICE_OTHER_CLOUD, session_id=alice_session) + assert foreign["error_code"] == unknown["error_code"] == other_cloud["error_code"] == SessionErrorCode.INVALID.value + assert foreign["error"] == unknown["error"] == other_cloud["error"] + assert harness["seen"] == [] + + +def test_valid_session_binds_context_and_echoes_session_id(harness): + session_id = harness["open_session"]() + body = harness["call"]("probe", "whoami", ALICE, session_id=session_id) + assert body.get("error") is None + assert body["session_id"] == session_id + identity, session = harness["seen"][0] + assert identity.user_id == user_id_for_token(ALICE) + assert session.session_id == session_id + assert current_identity() is None and current_session() is None + + +def test_sessions_of_the_same_credential_do_not_share_dataframes(harness): + first = harness["open_session"]() + second = harness["open_session"]() + harness["call"]("probe", "store", ALICE, session_id=first) + assert harness["call"]("probe", "count", ALICE, session_id=first)["result"][0]["dataframes"] == 1 + assert harness["call"]("probe", "count", ALICE, session_id=second)["result"][0]["dataframes"] == 0 + + +def test_nested_calls_inherit_the_session_and_reject_another_one(harness): + session_id = harness["open_session"]() + inherited = harness["call"]("probe", "batch", ALICE, session_id=session_id) + assert inherited.get("error") is None + assert harness["seen"][-1][1].session_id == session_id + + other = harness["open_session"]() + body = harness["call"]("probe", "batch", ALICE, session_id=session_id, sub_args={"session_id": other}) + assert body["error_code"] == SessionErrorCode.INVALID.value + + +def test_session_tool_description_says_to_call_it_once(harness): + assert "ONCE" in harness["mcp"].descriptions[SESSION_TOOL_NAME] + assert SESSION_HINT not in harness["mcp"].descriptions[SESSION_TOOL_NAME] + + +def test_session_tool_reports_an_outage_without_a_traceback(): + class _Broken: + async def open(self, owner_id, credential): + raise RuntimeError("boom") + + async def touch(self, session_id, owner_id, credential): + raise RuntimeError("boom") + + async def scenario(): + with bind_session_context(Identity(user_id="alice"), None, "secret"): + return await SessionManager(_Broken()).get() + + result = asyncio.run(scenario()) + assert result.error_code == SessionErrorCode.UNAVAILABLE.value + assert "Traceback" not in result.error + assert "boom" not in result.error + + +def test_missing_identity_on_get_fails_closed(): + with pytest.raises(SessionContextMissing): + asyncio.run(SessionManager(InMemorySessionProvider()).get()) diff --git a/tools/runtime_tools.py b/tools/runtime_tools.py index eedcf6f..b40cd91 100644 --- a/tools/runtime_tools.py +++ b/tools/runtime_tools.py @@ -1,9 +1,8 @@ -"""Run a tool action with telemetry, session scope, and dataframe materialization.""" +"""Run a tool action with telemetry, a chat session, and dataframe materialization.""" from typing import Any, Awaitable, Callable, Optional -from config.call_scope import resolve_mcp_session_id, user_id_for_token from config.runtime import AppRuntime -from config.session_context import bind_call_scope +from config.session_context import bind_session_context, current_identity, current_session from models.result import BaseResult from telemetry import run_tool from tools.call_state import ( @@ -13,6 +12,12 @@ set_result_format, ) from tools.dataframe_manager import INVALID_RESULT_FORMAT_ERROR, normalize_result_format +from tools.session_gate import ( + SESSION_ID_ARG, + credential_for_token, + identity_for_token, + validate_session, +) def _result_format_from_args(tool_args: Any) -> str: @@ -31,45 +36,102 @@ async def run_tool_with_runtime( tool_args: Any = None, dataframe_excluded_actions: Optional[set[str]] = None, disable_dataframe_materialization: bool = False, + requires_session: bool = True, ) -> Any: """ - Bind this call's storage scope, run the tool, then store oversized results. + Validate the caller and chat session, run the tool, then store oversized results. - Fast calls return their result directly. Methods wrapped with ``run_as_task`` - park work that exceeds a few seconds and return a task snapshot instead. + ``perfecto_session`` passes ``requires_session=False``. Every other tool + needs the ``session_id`` from that tool. Nested calls inherit the session + already bound for this call. Fast calls return their result directly. + Methods wrapped with ``run_as_task`` park work that exceeds a few seconds. """ + call_args = dict(tool_args) if isinstance(tool_args, dict) else {} + requested = call_args.pop(SESSION_ID_ARG, None) + if isinstance(requested, str): + requested = requested.strip() or None + result_format = _result_format_from_args(call_args) + + inherited = current_identity() + if inherited is not None: + inherited_session = current_session() + if requires_session: + if inherited_session is None: + from tools.session_gate import session_error + from config.session import SessionErrorCode + return session_error(SessionErrorCode.REQUIRED) + if requested is not None and requested != inherited_session.session_id: + from tools.session_gate import session_error + from config.session import SessionErrorCode + return session_error(SessionErrorCode.INVALID) + return await _execute( + runtime, tool_name, action, ctx, dispatch, call_args, result_format, + dataframe_excluded_actions, disable_dataframe_materialization, + inherited_session.session_id if inherited_session is not None else None, + ) + token = runtime.auth.get_token(ctx) if ctx is not None else None - user_id = user_id_for_token(token) - session_id = resolve_mcp_session_id(ctx) - credential = token.token if token is not None else None - result_format = _result_format_from_args(tool_args) - - with bind_call_scope(user_id, session_id, credential): - format_token = set_result_format("auto" if result_format == "invalid" else result_format) - try: - if result_format == "invalid": - return BaseResult(error=INVALID_RESULT_FORMAT_ERROR) - - async def _dispatch_and_finalize() -> Any: - policy_token = set_disable_dataframe_materialization(disable_dataframe_materialization) - try: - result = await dispatch() - if disable_dataframe_materialization or result is None: - return result - from tools.dataframe_manager import finalize_tool_result - - return await finalize_tool_result( - result, - action=action, - args=tool_args, - origin_manager=tool_name, - session_storage=runtime.storage, - scope_resolver=runtime.scope_resolver, - excluded_actions=dataframe_excluded_actions, - ) - finally: - reset_disable_dataframe_materialization(policy_token) - - return await run_tool(tool_name, action, ctx, _dispatch_and_finalize) - finally: - reset_result_format(format_token) + identity = identity_for_token(token) + if isinstance(identity, BaseResult): + return identity + credential = credential_for_token(token) + + session = None + if requires_session: + validated = await validate_session(runtime, requested, identity, credential) + if isinstance(validated, BaseResult): + return validated + session = validated + + with bind_session_context(identity, session, credential): + return await _execute( + runtime, tool_name, action, ctx, dispatch, call_args, result_format, + dataframe_excluded_actions, disable_dataframe_materialization, + session.session_id if session is not None else None, + ) + + +async def _execute( + runtime: AppRuntime, + tool_name: str, + action: str, + ctx: Any, + dispatch: Callable[[], Awaitable[Any]], + call_args: dict, + result_format: str, + dataframe_excluded_actions: Optional[set[str]], + disable_dataframe_materialization: bool, + session_id: Optional[str], +) -> Any: + format_token = set_result_format("auto" if result_format == "invalid" else result_format) + try: + if result_format == "invalid": + return BaseResult(error=INVALID_RESULT_FORMAT_ERROR) + + async def _dispatch_and_finalize() -> Any: + policy_token = set_disable_dataframe_materialization(disable_dataframe_materialization) + try: + result = await dispatch() + if disable_dataframe_materialization or result is None: + return result + from tools.dataframe_manager import finalize_tool_result + + return await finalize_tool_result( + result, + action=action, + args=call_args, + origin_manager=tool_name, + session_storage=runtime.storage, + scope_resolver=runtime.scope_resolver, + excluded_actions=dataframe_excluded_actions, + ) + finally: + reset_disable_dataframe_materialization(policy_token) + + result = await run_tool(tool_name, action, ctx, _dispatch_and_finalize) + finally: + reset_result_format(format_token) + + if isinstance(result, BaseResult) and not result.error and session_id and not result.session_id: + result.session_id = session_id + return result diff --git a/tools/session_gate.py b/tools/session_gate.py new file mode 100644 index 0000000..69be2e1 --- /dev/null +++ b/tools/session_gate.py @@ -0,0 +1,109 @@ +"""Require a chat session on every Perfecto tool except the session tool itself.""" +from __future__ import annotations + +import logging +from typing import Any, Optional, Union + +from config.perfecto import ( + PERFECTO_CLOUD_NAME_NOT_SET_MESSAGE, + SECURITY_TOKEN_NOT_SET_MESSAGE, +) +from config.runtime import AppRuntime +from config.session import ( + SESSION_TOOL_NAME, + ChatSession, + SessionError, + SessionErrorCode, + is_valid_session_id, +) +from config.session_context import ( + Identity, + current_credential, + current_identity, + current_session, +) +from config.call_scope import user_id_for_token +from config.token import PerfectoToken +from models.result import BaseResult + +logger = logging.getLogger(__name__) + +SESSION_ID_ARG = "session_id" +SESSION_HINT = ( + f"- **CRITICAL (session)**: Every call requires `{SESSION_ID_ARG}` next to `action` " + f"(or inside args). Get it once per conversation with `{SESSION_TOOL_NAME}` action 'get' " + "and reuse it in every call; never reuse a session_id from another conversation. If a call " + "returns error_code SESSION_REQUIRED, SESSION_INVALID or SESSION_EXPIRED, get a new " + "session_id, use it from then on, and re-run the calls whose tasks or dataframes you still need." +) + + +def session_error(code: SessionErrorCode) -> BaseResult: + error = SessionError(code) + return BaseResult(error=error.message, error_code=error.code.value) + + +def credential_for_token(token: Optional[PerfectoToken]) -> Optional[str]: + """Opaque security token the session is bound to. The cloud is part of the owner id.""" + if token is None or not token.token: + return None + return token.token + + +def identity_for_token(token: Optional[PerfectoToken]) -> Union[Identity, BaseResult]: + """ + Owner of sessions and partitions for this call. + + Perfecto already authenticated the security token at the transport. The owner + id is a hash of that token plus the cloud, so two credentials never share a + session. There is no separate identity API. + """ + if token is None or not getattr(token, "token", None): + return BaseResult(error=SECURITY_TOKEN_NOT_SET_MESSAGE) + if not getattr(token, "cloud_name", None): + return BaseResult(error=PERFECTO_CLOUD_NAME_NOT_SET_MESSAGE) + return Identity(user_id=user_id_for_token(token)) + + +async def validate_session( + runtime: AppRuntime, + requested: Any, + identity: Identity, + credential: Optional[str], +) -> Union[BaseResult, ChatSession]: + if requested is None or (isinstance(requested, str) and not requested.strip()): + return session_error(SessionErrorCode.REQUIRED) + if not is_valid_session_id(requested) or not credential: + return session_error(SessionErrorCode.INVALID) + try: + return await runtime.sessions.touch(requested, identity.user_id, credential) + except SessionError as exc: + return session_error(exc.code) + except Exception: + logger.exception("Session validation failed") + return session_error(SessionErrorCode.UNAVAILABLE) + + +def annotate_session_requirements(mcp: Any) -> None: + """Tell agents, from the tool description, that session_id is required.""" + try: + from mcp.server.fastmcp import FastMCP + except Exception: + return + if not isinstance(mcp, FastMCP): + return + for tool in mcp._tool_manager.list_tools(): + if tool.name == SESSION_TOOL_NAME: + continue + description = tool.description or "" + if SESSION_HINT not in description: + tool.description = f"{description.rstrip()}\n{SESSION_HINT}\n" + parameters = tool.parameters if isinstance(tool.parameters, dict) else None + if parameters is None: + continue + arguments = parameters.setdefault("properties", {}).get("arguments") + if isinstance(arguments, dict): + note = f" Include `{SESSION_ID_ARG}` from `{SESSION_TOOL_NAME}` action 'get'." + current = arguments.get("description") or "Dictionary with arguments." + if SESSION_ID_ARG not in current: + arguments["description"] = f"{current.rstrip()}.{note}" if not current.endswith(".") else f"{current}{note}" diff --git a/tools/session_manager.py b/tools/session_manager.py new file mode 100644 index 0000000..3ba1952 --- /dev/null +++ b/tools/session_manager.py @@ -0,0 +1,92 @@ +"""Open a chat session for the Perfecto caller of the current tool call.""" +import logging + +from config.perfecto import TOOLS_PREFIX +from config.runtime import AppRuntime +from config.session import SESSION_TOOL_NAME, SessionError, SessionErrorCode +from config.session_context import SessionContextMissing, current_credential, current_identity +from models.result import BaseResult +from tools.runtime_tools import run_tool_with_runtime +from tools.utils import normalize_action_args + +logger = logging.getLogger(__name__) + +_HEADER = ( + "Chat session for this conversation. Call it ONCE at the start of the conversation, " + "before any other Perfecto tool, and keep the returned session_id in context." +) + + +class SessionManager: + def __init__(self, sessions): + self.sessions = sessions + + async def get(self) -> BaseResult: + identity = current_identity() + credential = current_credential() + if identity is None or not credential: + raise SessionContextMissing("The caller identity was not bound for this call.") + try: + session = await self.sessions.open(identity.user_id, credential) + except SessionError as exc: + return BaseResult(error=exc.message, error_code=exc.code.value) + except Exception: + logger.exception("Opening a chat session failed") + unavailable = SessionError(SessionErrorCode.UNAVAILABLE) + return BaseResult(error=unavailable.message, error_code=unavailable.code.value) + return BaseResult( + result=[{ + "session_id": session.session_id, + "state": session.state.value, + "created_at": session.created_at.isoformat(), + }], + session_id=session.session_id, + info=[ + "Keep this session_id in context and pass it as 'session_id' in every Perfecto " + "tool call of this conversation. Do not call this tool again unless a tool returns " + "SESSION_REQUIRED, SESSION_INVALID or SESSION_EXPIRED.", + ], + ) + + +def register(mcp, runtime: AppRuntime): + manager = SessionManager(runtime.sessions) + + @mcp.tool( + name=SESSION_TOOL_NAME, + description=f""" +{_HEADER} +Actions: +- get: Start a new chat session and return its session_id. + args(dict): No arguments. +Hints: +- **CRITICAL**: Pass the returned `session_id` in every call to every other {TOOLS_PREFIX} tool + of this conversation (next to action, or inside args). Do not call 'get' again while it keeps working. +- **CRITICAL**: Never reuse a session_id from another conversation: tasks and dataframes are + isolated per session. +- Only call 'get' again when a tool returns error_code SESSION_REQUIRED, SESSION_INVALID or + SESSION_EXPIRED; then use the new session_id and re-run the calls whose tasks or dataframes + you still need. +""", + ) + async def session_tool(arguments: dict = None, ctx=None) -> BaseResult: + action, args = normalize_action_args(arguments) + runtime.configure_context(ctx) + + async def _dispatch(): + match action: + case "get": + return await manager.get() + case _: + return BaseResult(error=f"Action {action} not found in session tool") + + return await run_tool_with_runtime( + runtime, + SESSION_TOOL_NAME, + action, + ctx, + _dispatch, + tool_args=args, + requires_session=False, + disable_dataframe_materialization=True, + )