From 6f6f904340e3d45a1eab81479b0c6107aca43776 Mon Sep 17 00:00:00 2001 From: diego-ferrand Date: Tue, 3 Nov 2026 15:23:44 -0300 Subject: [PATCH 1/2] Store large tool results as session dataframes and park slow calls as tasks. --- config/auth.py | 12 +- config/cache.py | 677 +++++++++++++++ config/cache_codec.py | 117 +++ config/call_scope.py | 53 ++ config/env.py | 77 ++ config/http_clients.py | 72 ++ config/ids.py | 25 + config/runtime.py | 59 +- config/service_auth.py | 22 + config/session_context.py | 75 ++ config/storage.py | 233 +++++ docs/hosted-http.md | 18 +- main.py | 3 + models/result.py | 17 +- pyproject.toml | 1 + server.py | 5 + tests/conftest.py | 14 + tests/test_dataframes_tasks_cache.py | 220 +++++ tools/ai_scriptless_manager.py | 18 +- tools/async_task_manager.py | 666 ++++++++++++++ tools/call_state.py | 55 ++ tools/data_tools.py | 606 +++++++++++++ tools/dataframe_manager.py | 1193 ++++++++++++++++++++++++++ tools/device_manager.py | 20 +- tools/execution_manager.py | 20 +- tools/help_manager.py | 18 +- tools/method_cache.py | 213 +++++ tools/runtime_tools.py | 75 ++ tools/skills_manager.py | 6 +- tools/task_runtime.py | 181 ++++ tools/tools_manager.py | 116 ++- tools/user_manager.py | 12 +- tools/utils.py | 43 + uv.lock | 32 +- 34 files changed, 4946 insertions(+), 28 deletions(-) create mode 100644 config/cache.py create mode 100644 config/cache_codec.py create mode 100644 config/call_scope.py create mode 100644 config/env.py create mode 100644 config/http_clients.py create mode 100644 config/ids.py create mode 100644 config/service_auth.py create mode 100644 config/session_context.py create mode 100644 config/storage.py create mode 100644 tests/test_dataframes_tasks_cache.py create mode 100644 tools/async_task_manager.py create mode 100644 tools/call_state.py create mode 100644 tools/data_tools.py create mode 100644 tools/dataframe_manager.py create mode 100644 tools/method_cache.py create mode 100644 tools/runtime_tools.py create mode 100644 tools/task_runtime.py diff --git a/config/auth.py b/config/auth.py index 5ef261d..21c4c62 100644 --- a/config/auth.py +++ b/config/auth.py @@ -156,6 +156,8 @@ def run_streamable_http(mcp: FastMCP) -> None: register_health_routes(mcp) async def _serve() -> None: + from config.http_clients import aclose_http_clients + app = BearerAuthMiddleware(mcp.streamable_http_app()) config = uvicorn.Config( app, @@ -163,6 +165,14 @@ async def _serve() -> None: port=mcp.settings.port, log_level=mcp.settings.log_level.lower(), ) - await uvicorn.Server(config).serve() + try: + await uvicorn.Server(config).serve() + finally: + from config.cache import get_cache + + try: + await get_cache().close() + finally: + await aclose_http_clients() anyio.run(_serve) diff --git a/config/cache.py b/config/cache.py new file mode 100644 index 0000000..f6e4b62 --- /dev/null +++ b/config/cache.py @@ -0,0 +1,677 @@ +""" +Copyright 2025 Perforce Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" +from __future__ import annotations + +import asyncio +import copy +import heapq +import json +import logging +import time +from abc import ABC, abstractmethod +from collections import OrderedDict +from dataclasses import dataclass +from enum import Enum +from typing import Any, Awaitable, Callable, Optional +from urllib.parse import quote + +import httpx + +from config import cache_codec +from config.env import env_bool, env_float, env_int +from config.http_clients import SharedAsyncClient +from config.service_auth import service_headers + +logger = logging.getLogger(__name__) + +DEFAULT_MAX_ENTRIES = 2048 +DEFAULT_SWEEP_INTERVAL_SECONDS = 30.0 +DEFAULT_SWEEP_BATCH_SIZE = 500 +DEFAULT_HTTP_TIMEOUT_SECONDS = 2.0 +# Match the storage API cache body limit (request body limit). +DEFAULT_HTTP_MAX_VALUE_BYTES = 1048576 + +CacheObserver = Callable[[str, int], None] + + +class CacheScope(str, Enum): + """ + Who may reuse a cache entry. + + USER: only the user identified in the request context. Without a user id the + cache is bypassed, so entries can never be shared across users by accident. + GLOBAL: any caller. Only for content that does not depend on the caller + (help pages, static assets). + """ + + USER = "user" + GLOBAL = "global" + + +def build_cache_key( + scope: CacheScope, + namespace: str, + *parts: str, + user_id: Optional[str] = None, +) -> str: + """Build a namespaced key; USER keys require a non-empty user id.""" + if scope is CacheScope.USER: + if not user_id or not str(user_id).strip(): + raise ValueError("USER-scoped cache keys require a user id.") + prefix = f"{scope.value}:{str(user_id).strip()}" + else: + prefix = scope.value + return ":".join([prefix, namespace, *parts]) + + +@dataclass(frozen=True) +class CacheEntry: + value: Any + # Set when the value is only valid for one session binding (e.g. a task + # snapshot or a session dataframe reference). None means shareable. + bound_to: Optional[str] = None + # Opaque owner marker so a writer can later delete only its own entry. + tag: Optional[str] = None + + def accepts(self, binding: Optional[str]) -> bool: + return self.bound_to is None or self.bound_to == binding + + +@dataclass(frozen=True) +class StorePlan: + """How get_or_load stores a freshly loaded value.""" + + ttl_seconds: Optional[float] = None + bound_to: Optional[str] = None + tag: Optional[str] = None + # False for large read-only values (e.g. static indexes): no copy in or out, + # so every reader must treat the value as immutable. + copy_values: bool = True + # Applied to a copy of the loaded value before storing (e.g. strip per-call debug). + # Waiters sharing the load receive the transformed value as well. + transform: Optional[Callable[[Any], Any]] = None + after_store: Optional[Callable[[], Awaitable[None]]] = None + # False: share the load with concurrent waiters (binding rules apply) but keep nothing, + # for values that can go stale without the cache knowing (e.g. a removable dataframe). + store: bool = True + + +StorePolicy = Callable[[Any], Optional[StorePlan]] + + +def _store_always(_value: Any) -> StorePlan: + return StorePlan() + + +def _notify(observer: Optional[CacheObserver], metric: str, value: int) -> None: + if observer is None: + return + try: + observer(metric, value) + except Exception: + logger.debug("cache observer failed", exc_info=True) + + +class CachePort(ABC): + """ + Key/value cache with per-entry TTL. + + Adapters own storage and expiry. ``get_or_load`` adds process-local + single-flight on top, so an adapter backed by a shared cache service only + has to implement get/set/delete/clear. + """ + + def __init__(self) -> None: + self._inflight: dict[str, asyncio.Future] = {} + + @abstractmethod + async def get(self, key: str) -> Optional[CacheEntry]: + raise NotImplementedError + + @abstractmethod + async def set( + self, + key: str, + value: Any, + ttl_seconds: float, + *, + bound_to: Optional[str] = None, + tag: Optional[str] = None, + copy_values: bool = True, + ) -> None: + raise NotImplementedError + + @abstractmethod + async def delete(self, key: str, *, if_tag: Optional[str] = None) -> bool: + """Delete ``key``; with ``if_tag``, only while the entry still carries that tag.""" + raise NotImplementedError + + @abstractmethod + async def clear(self) -> None: + raise NotImplementedError + + async def start(self) -> None: + """Start background maintenance, if the adapter needs it.""" + + async def close(self) -> None: + """Stop background maintenance and release resources.""" + + async def get_or_load( + self, + key: str, + loader: Callable[[], Awaitable[Any]], + ttl_seconds: float, + *, + store_policy: Optional[StorePolicy] = None, + binding: Optional[str] = None, + observer: Optional[CacheObserver] = None, + ) -> Any: + """ + Return the cached value for ``key`` or load, store and return it. + + Concurrent misses on the same key share one ``loader`` call. The owner + releases the key in every exit path, cancellation included, so a + cancelled load never leaves waiters hanging. ``store_policy`` decides + whether and how the loaded value is stored (None = do not store). + Entries bound to another ``binding`` count as a miss. + """ + policy = store_policy or _store_always + loop = asyncio.get_running_loop() + while True: + entry = await self.get(key) + if entry is not None and entry.accepts(binding): + _notify(observer, "hits", 1) + return entry.value + + future = self._inflight.get(key) + if future is None or future.get_loop() is not loop: + break + + wait_started = time.monotonic() + try: + shared_value, plan = await asyncio.shield(future) + except asyncio.CancelledError: + current = asyncio.current_task() + if future.cancelled() and not (current is not None and current.cancelling()): + # The owner was cancelled, not us: retry and possibly become the owner. + continue + raise + _notify(observer, "shared_wait_ms", int((time.monotonic() - wait_started) * 1000)) + if plan is None or plan.bound_to is None or plan.bound_to == binding: + if plan is not None and not plan.copy_values: + return shared_value + return copy.deepcopy(shared_value) + # The shared value belongs to another binding: go round again, so the waiters + # of this binding share one load (one of them becomes the owner). + continue + + future = loop.create_future() + self._inflight[key] = future + _notify(observer, "misses", 1) + try: + value = await loader() + except BaseException as exc: + if self._inflight.get(key) is future: + del self._inflight[key] + if not future.done(): + if isinstance(exc, asyncio.CancelledError): + future.cancel() + else: + future.set_exception(exc) + future.exception() # mark retrieved; waiters still receive it + raise + + # Caching is best effort: a failure here never fails a load that already succeeded. + plan: Optional[StorePlan] = None + stored = False + shared_value = value + try: + plan = policy(value) + if plan is not None and plan.transform is not None: + shared_value = plan.transform(copy.deepcopy(value)) + if plan is not None and plan.store: + ttl = plan.ttl_seconds if plan.ttl_seconds is not None else ttl_seconds + await self.set( + key, + shared_value, + ttl, + bound_to=plan.bound_to, + tag=plan.tag, + copy_values=plan.copy_values, + ) + stored = True + if plan.after_store is not None: + await plan.after_store() + except Exception: + logger.warning("cache store failed for key %s; serving uncached result", key, exc_info=True) + if plan is None: + # Unknown binding: share only with waiters of the same binding. + plan = StorePlan(bound_to=binding) + if stored: + # The entry may lack its invalidation hook (after_store failed): drop it. + try: + await self.delete(key, if_tag=plan.tag) + except Exception: + logger.warning("cache delete failed for key %s", key, exc_info=True) + finally: + if self._inflight.get(key) is future: + del self._inflight[key] + if not future.done(): + copy_shared = plan is None or plan.copy_values + future.set_result((copy.deepcopy(shared_value) if copy_shared else shared_value, plan)) + return value + + +class NullCache(CachePort): + """Cache disabled: every lookup misses and nothing is shared.""" + + async def get(self, key: str) -> Optional[CacheEntry]: + return None + + async def set(self, key: str, value: Any, ttl_seconds: float, *, bound_to=None, tag=None, copy_values=True) -> None: + return None + + async def delete(self, key: str, *, if_tag: Optional[str] = None) -> bool: + return False + + async def clear(self) -> None: + return None + + async def get_or_load(self, key, loader, ttl_seconds, *, store_policy=None, binding=None, observer=None): + return await loader() + + +@dataclass +class _Slot: + value: Any + expires_at: float + bound_to: Optional[str] + generation: int + tag: Optional[str] = None + copy_values: bool = True + + +class InMemoryTTLCache(CachePort): + """ + Process-local cache (stdio and a single hosted worker). + + - LRU order in an OrderedDict: hits and overflow eviction are O(1). + - Expiry heap (expires_at, generation, key): the sweeper pops only expired + entries; stale heap items are skipped by generation. + - Reads drop an expired entry on the spot, so correctness never waits for + the sweeper. + - The sweeper is an asyncio task on the running loop. All mutations happen + on that loop between awaits, so the store needs no locks; the sweeper + yields every ``sweep_batch_size`` entries to keep the loop responsive. + - Values are deep-copied in and out so callers cannot mutate shared state, + unless stored with ``copy_values=False`` (read-only values). + """ + + def __init__( + self, + *, + max_entries: int = DEFAULT_MAX_ENTRIES, + sweep_interval_seconds: float = DEFAULT_SWEEP_INTERVAL_SECONDS, + sweep_batch_size: int = DEFAULT_SWEEP_BATCH_SIZE, + clock: Callable[[], float] = time.monotonic, + ) -> None: + super().__init__() + self._max_entries = max(1, int(max_entries)) + self._sweep_interval = float(sweep_interval_seconds) + self._sweep_batch_size = max(1, int(sweep_batch_size)) + self._clock = clock + self._data: OrderedDict[str, _Slot] = OrderedDict() + self._heap: list[tuple[float, int, str]] = [] + self._generation = 0 + self._gc_task: Optional[asyncio.Task] = None + self._gc_loop: Optional[asyncio.AbstractEventLoop] = None + + def __len__(self) -> int: + return len(self._data) + + async def get(self, key: str) -> Optional[CacheEntry]: + self._ensure_gc() + slot = self._data.get(key) + if slot is None: + return None + if slot.expires_at <= self._clock(): + del self._data[key] + return None + self._data.move_to_end(key) + value = copy.deepcopy(slot.value) if slot.copy_values else slot.value + return CacheEntry(value=value, bound_to=slot.bound_to, tag=slot.tag) + + async def set( + self, + key: str, + value: Any, + ttl_seconds: float, + *, + bound_to: Optional[str] = None, + tag: Optional[str] = None, + copy_values: bool = True, + ) -> None: + if ttl_seconds is None or ttl_seconds <= 0: + self._data.pop(key, None) + return + self._ensure_gc() + self._generation += 1 + expires_at = self._clock() + float(ttl_seconds) + self._data[key] = _Slot( + value=copy.deepcopy(value) if copy_values else value, + expires_at=expires_at, + bound_to=bound_to, + generation=self._generation, + tag=tag, + copy_values=copy_values, + ) + self._data.move_to_end(key) + heapq.heappush(self._heap, (expires_at, self._generation, key)) + while len(self._data) > self._max_entries: + self._data.popitem(last=False) + self._maybe_compact_heap() + + async def delete(self, key: str, *, if_tag: Optional[str] = None) -> bool: + slot = self._data.get(key) + if slot is None or (if_tag is not None and slot.tag != if_tag): + return False + del self._data[key] + return True + + async def clear(self) -> None: + self._data.clear() + self._heap.clear() + + async def sweep(self) -> int: + """Remove expired entries in batches; returns how many were removed.""" + removed = 0 + processed = 0 + now = self._clock() + while self._heap and self._heap[0][0] <= now: + _, generation, key = heapq.heappop(self._heap) + slot = self._data.get(key) + if slot is not None and slot.generation == generation: + del self._data[key] + removed += 1 + processed += 1 + if processed % self._sweep_batch_size == 0: + await asyncio.sleep(0) + now = self._clock() + return removed + + async def start(self) -> None: + self._ensure_gc() + + async def close(self) -> None: + task = self._gc_task + self._gc_task = None + self._gc_loop = None + if task is None or task.done(): + return + task.cancel() + try: + await task + except (asyncio.CancelledError, RuntimeError): + pass + + def _maybe_compact_heap(self) -> None: + # Overwrites and LRU evictions leave stale heap items; rebuild when they dominate. + if len(self._heap) <= max(64, 2 * len(self._data)): + return + self._heap = [(slot.expires_at, slot.generation, key) for key, slot in self._data.items()] + heapq.heapify(self._heap) + + def _ensure_gc(self) -> None: + if self._sweep_interval <= 0: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + task = self._gc_task + if task is not None and not task.done() and self._gc_loop is loop: + return + self._gc_loop = loop + self._gc_task = loop.create_task(self._gc_run(), name="perfecto-cache-sweeper") + + async def _gc_run(self) -> None: + while True: + await asyncio.sleep(self._sweep_interval) + try: + await self.sweep() + except asyncio.CancelledError: + raise + except Exception: + logger.exception("cache sweeper failed; continuing") + + +class CacheValueTooLarge(ValueError): + """The encoded value exceeds what the remote cache accepts; it is not cached.""" + + +class HttpCache(CachePort): + """ + Hosted: the cache lives behind the storage API (``/cache/entries``). + + What backs the API (in-memory today, a shared cache service later) can change + without touching the MCP. Values go through ``cache_codec`` so a hit returns + the same types as a local cache. Reads never raise (a failure is a miss); + writes raise and ``get_or_load``'s best-effort store serves the value uncached. + Values larger than ``max_value_bytes`` are rejected here, without a request. + """ + + def __init__( + self, + base_url: str, + caller_token: str, + *, + timeout_seconds: float = DEFAULT_HTTP_TIMEOUT_SECONDS, + max_value_bytes: int = DEFAULT_HTTP_MAX_VALUE_BYTES, + transport: Optional[httpx.AsyncBaseTransport] = None, + ) -> None: + super().__init__() + self._base_url = base_url.rstrip("/") + # Same contract as the session, partition and ticket clients: caller identity + # headers and a shared, cookie-less client closed on server shutdown. + self._headers = service_headers(caller_token) + self._max_value_bytes = max_value_bytes + if transport is not None: + self._shared = SharedAsyncClient( + lambda: httpx.AsyncClient(transport=transport, timeout=timeout_seconds) + ) + else: + self._shared = SharedAsyncClient( + lambda: httpx.AsyncClient(http2=True, timeout=timeout_seconds) + ) + self._client_loop: Optional[asyncio.AbstractEventLoop] = None + self._degraded = False + + def _http(self) -> httpx.AsyncClient: + loop = asyncio.get_running_loop() + if self._client_loop is not loop: + # A pool cannot be shared across event loops: forget the previous loop's client. + self._shared.discard() + self._client_loop = loop + return self._shared.get() + + def _url(self, key: str) -> str: + return f"{self._base_url}/cache/entries/{quote(key, safe='')}" + + def _report_failure(self, operation: str, key: Optional[str], exc: BaseException) -> None: + status = exc.response.status_code if isinstance(exc, httpx.HTTPStatusError) else None + if status in (401, 403): + if not self._degraded: + logger.error( + "remote cache rejected the MCP caller (HTTP %s): check the storage caller " + "token. The cache is effectively disabled until this is fixed.", status, + ) + elif not self._degraded: + logger.warning( + "remote cache %s failed (key %s); serving without cache until it recovers", + operation, key, exc_info=True, + ) + else: + logger.debug("remote cache %s failed (key %s): %s", operation, key, exc) + self._degraded = True + + def _report_success(self) -> None: + if self._degraded: + logger.info("remote cache recovered") + self._degraded = False + + async def get(self, key: str) -> Optional[CacheEntry]: + try: + response = await self._http().get(self._url(key), headers=self._headers) + if response.status_code == 404: + self._report_success() + return None + response.raise_for_status() + body = response.json() + entry = CacheEntry( + value=cache_codec.decode(body.get("value")), + bound_to=body.get("bound_to"), + tag=body.get("tag"), + ) + except Exception as exc: + self._report_failure("read", key, exc) + return None + self._report_success() + return entry + + async def set( + self, + key: str, + value: Any, + ttl_seconds: float, + *, + bound_to: Optional[str] = None, + tag: Optional[str] = None, + copy_values: bool = True, + ) -> None: + # copy_values does not apply: a remote value is always a fresh copy. + if ttl_seconds is None or ttl_seconds <= 0: + await self.delete(key) + return + body = json.dumps( + { + "value": cache_codec.encode(value), + "ttl_seconds": float(ttl_seconds), + "bound_to": bound_to, + "tag": tag, + }, + separators=(",", ":"), + ).encode("utf-8") + if len(body) > self._max_value_bytes: + raise CacheValueTooLarge( + f"Encoded cache value is {len(body)} bytes; the remote cache accepts {self._max_value_bytes}." + ) + try: + response = await self._http().put( + self._url(key), + headers={**self._headers, "Content-Type": "application/json"}, + content=body, + ) + response.raise_for_status() + except Exception as exc: + self._report_failure("write", key, exc) + raise + self._report_success() + + async def delete(self, key: str, *, if_tag: Optional[str] = None) -> bool: + params = {"if_tag": if_tag} if if_tag is not None else None + try: + response = await self._http().delete(self._url(key), headers=self._headers, params=params) + response.raise_for_status() + deleted = bool(response.json().get("deleted")) + except Exception as exc: + self._report_failure("delete", key, exc) + return False + self._report_success() + return deleted + + async def clear(self) -> None: + try: + response = await self._http().delete(f"{self._base_url}/cache/entries", headers=self._headers) + response.raise_for_status() + except Exception as exc: + self._report_failure("clear", None, exc) + return + self._report_success() + + async def close(self) -> None: + self._client_loop = None + await self._shared.aclose() + + +def build_cache_from_env( + transport: str = "stdio", + storage_base_url: Optional[str] = None, + caller_token: Optional[str] = None, +) -> CachePort: + """ + Build the process cache from PERFECTO_MCP_CACHE_* env vars (config.env helpers). + + - stdio: in-memory (PERFECTO_MCP_CACHE_MAX_ENTRIES 2048, PERFECTO_MCP_CACHE_SWEEP_INTERVAL_SECONDS 30, + PERFECTO_MCP_CACHE_SWEEP_BATCH_SIZE 500). + - streamable-http: only the storage API cache (PERFECTO_MCP_CACHE_HTTP_TIMEOUT_SECONDS 2, + PERFECTO_MCP_CACHE_HTTP_MAX_VALUE_BYTES 1048576); no in-process state, so what backs + the API can change without touching the MCP. + - PERFECTO_MCP_CACHE_ENABLED=false (default true): no cache in either transport. + """ + if not env_bool("CACHE_ENABLED", True): + return NullCache() + if transport == "streamable-http": + if not storage_base_url or not caller_token: + raise ValueError("The hosted cache needs the storage API base URL and caller token.") + return HttpCache( + storage_base_url, + caller_token, + timeout_seconds=env_float( + "CACHE_HTTP_TIMEOUT_SECONDS", DEFAULT_HTTP_TIMEOUT_SECONDS, minimum=0.1 + ), + max_value_bytes=env_int( + "CACHE_HTTP_MAX_VALUE_BYTES", DEFAULT_HTTP_MAX_VALUE_BYTES, minimum=1 + ), + ) + return _build_in_memory_cache() + + +def _build_in_memory_cache() -> InMemoryTTLCache: + return InMemoryTTLCache( + max_entries=env_int("CACHE_MAX_ENTRIES", DEFAULT_MAX_ENTRIES, minimum=1), + # 0 disables the background sweeper (expired entries still drop on read). + sweep_interval_seconds=env_float( + "CACHE_SWEEP_INTERVAL_SECONDS", DEFAULT_SWEEP_INTERVAL_SECONDS, minimum=0.0 + ), + sweep_batch_size=env_int("CACHE_SWEEP_BATCH_SIZE", DEFAULT_SWEEP_BATCH_SIZE, minimum=1), + ) + + +_cache: Optional[CachePort] = None + + +def configure_cache(cache: Optional[CachePort]) -> None: + """Bind the process cache from AppRuntime (composition root).""" + global _cache + _cache = cache + + +def get_cache() -> CachePort: + """Process cache; lazily built from env when AppRuntime did not configure one.""" + global _cache + if _cache is None: + _cache = build_cache_from_env() + return _cache diff --git a/config/cache_codec.py b/config/cache_codec.py new file mode 100644 index 0000000..9d36d67 --- /dev/null +++ b/config/cache_codec.py @@ -0,0 +1,117 @@ +""" +Copyright 2025 Perforce Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" +""" +JSON codec for values stored in a remote cache. + +Cached results are typed (BaseResult holding Account/Workspace/Test models that +callers read as attributes) and some cached structures use non-string dict keys +(the help index), so plain JSON would change what a cache hit returns. Values +are encoded with small tagged wrappers and decoded back to the same types. + +Only pydantic models from ``ALLOWED_MODEL_MODULES`` are rebuilt: the store never +decides which class gets instantiated. +""" +import importlib +from datetime import date, datetime +from typing import Any + +from pydantic import BaseModel + +ALLOWED_MODEL_MODULES = ("models.",) + +_MODEL = "__model__" +_DICT = "__dict__" +_TUPLE = "__tuple__" +_SET = "__set__" +_DATETIME = "__datetime__" +_DATE = "__date__" +_RESERVED = frozenset({_MODEL, _DICT, _TUPLE, _SET, _DATETIME, _DATE}) + + +class CacheCodecError(TypeError): + """The value cannot be stored remotely (it is then simply not cached).""" + + +def _model_path(cls: type) -> str: + return f"{cls.__module__}:{cls.__qualname__}" + + +def encode(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, BaseModel): + cls = type(value) + if not cls.__module__.startswith(ALLOWED_MODEL_MODULES): + raise CacheCodecError(f"Model {_model_path(cls)} is not allowed in the remote cache.") + return { + _MODEL: _model_path(cls), + "fields": {name: encode(getattr(value, name)) for name in cls.model_fields}, + "set": sorted(value.model_fields_set), + } + if isinstance(value, datetime): + return {_DATETIME: value.isoformat()} + if isinstance(value, date): + return {_DATE: value.isoformat()} + if isinstance(value, dict): + if all(isinstance(key, str) for key in value) and not (_RESERVED & value.keys()): + return {key: encode(item) for key, item in value.items()} + return {_DICT: [[encode(key), encode(item)] for key, item in value.items()]} + if isinstance(value, list): + return [encode(item) for item in value] + if isinstance(value, tuple): + return {_TUPLE: [encode(item) for item in value]} + if isinstance(value, (set, frozenset)): + return {_SET: [encode(item) for item in value]} + raise CacheCodecError(f"Type {type(value).__name__} cannot be stored in the remote cache.") + + +def _load_model(path: str) -> type[BaseModel]: + module_name, _, qualname = path.partition(":") + if not module_name.startswith(ALLOWED_MODEL_MODULES) or not qualname: + raise CacheCodecError(f"Model {path} is not allowed in the remote cache.") + target: Any = importlib.import_module(module_name) + for part in qualname.split("."): + target = getattr(target, part) + if not (isinstance(target, type) and issubclass(target, BaseModel)): + raise CacheCodecError(f"{path} is not a model.") + return target + + +def decode(value: Any) -> Any: + if isinstance(value, list): + return [decode(item) for item in value] + if not isinstance(value, dict): + return value + if _MODEL in value: + cls = _load_model(value[_MODEL]) + fields = {name: decode(item) for name, item in value.get("fields", {}).items()} + # Built from a real instance's fields: restore them as they were. + return cls.model_construct(_fields_set=set(value.get("set", fields)), **fields) + if _DICT in value: + return {_hashable(decode(key)): decode(item) for key, item in value[_DICT]} + if _TUPLE in value: + return tuple(decode(item) for item in value[_TUPLE]) + if _SET in value: + return {_hashable(decode(item)) for item in value[_SET]} + if _DATETIME in value: + return datetime.fromisoformat(value[_DATETIME]) + if _DATE in value: + return date.fromisoformat(value[_DATE]) + return {key: decode(item) for key, item in value.items()} + + +def _hashable(value: Any) -> Any: + return tuple(_hashable(item) for item in value) if isinstance(value, list) else value diff --git a/config/call_scope.py b/config/call_scope.py new file mode 100644 index 0000000..b68efe8 --- /dev/null +++ b/config/call_scope.py @@ -0,0 +1,53 @@ +"""Derive a storage partition from the Perfecto credential and MCP session.""" +from __future__ import annotations + +import hashlib +from typing import Any, Optional + +from config.token import PerfectoToken + + +def user_id_for_token(token: Optional[PerfectoToken]) -> str: + """ + Stable partition owner for a security token and cloud. + + The raw token is never used as a storage key. Anonymous calls share one + owner so local tests without a token still have a scope. + """ + raw = getattr(token, "token", None) if token is not None else None + if not raw: + return "anonymous" + cloud = getattr(token, "cloud_name", None) or "" + digest = hashlib.sha256(f"{cloud}\n{raw}".encode("utf-8")).hexdigest() + return digest[:32] + + +def resolve_mcp_session_id(ctx: Any) -> str: + """ + MCP session that owns dataframes and tasks for this call. + + Streamable HTTP clients send ``mcp-session-id``. FastMCP also exposes it as + ``ctx.session_id``. Stdio and calls without a session share ``local``, which + is one partition per process for that credential. + """ + if ctx is None: + return "local" + + request = None + try: + request = ctx.request_context.request + except Exception: + request = None + headers = getattr(request, "headers", None) + if headers is not None: + raw = headers.get("mcp-session-id") + if raw and str(raw).strip(): + return str(raw).strip() + + try: + session_id = ctx.session_id + if session_id: + return str(session_id) + except Exception: + pass + return "local" diff --git a/config/env.py b/config/env.py new file mode 100644 index 0000000..45a3a7a --- /dev/null +++ b/config/env.py @@ -0,0 +1,77 @@ +"""Environment helpers for Perfecto MCP tunables. + +Names are read as ``PERFECTO_MCP_`` (for example +``PERFECTO_MCP_MAX_PARALLEL_TASKS``). +""" +from __future__ import annotations + +import logging +import os +from typing import Optional + +logger = logging.getLogger(__name__) + +ENV_PREFIX = "PERFECTO_MCP_" + + +def _raw(name: str) -> Optional[str]: + value = os.getenv(f"{ENV_PREFIX}{name}") + if value is None or not value.strip(): + return None + return value.strip() + + +def env_str(name: str, default: str = "") -> str: + value = _raw(name) + return default if value is None else value + + +def env_int(name: str, default: int, *, minimum: Optional[int] = None) -> int: + value = _raw(name) + if value is None: + return default + try: + parsed = int(value) + except ValueError: + logger.warning("Invalid %s%s=%r; using default %s", ENV_PREFIX, name, value, default) + return default + if minimum is not None and parsed < minimum: + logger.warning( + "%s%s=%s is below %s; using default %s", ENV_PREFIX, name, parsed, minimum, default, + ) + return default + return parsed + + +def env_float(name: str, default: float, *, minimum: Optional[float] = None) -> float: + value = _raw(name) + if value is None: + return default + try: + parsed = float(value) + except ValueError: + logger.warning("Invalid %s%s=%r; using default %s", ENV_PREFIX, name, value, default) + return default + if minimum is not None and parsed < minimum: + logger.warning( + "%s%s=%s is below %s; using default %s", ENV_PREFIX, name, parsed, minimum, default, + ) + return default + return parsed + + +_TRUE_VALUES = {"1", "true", "yes", "on"} +_FALSE_VALUES = {"0", "false", "no", "off"} + + +def env_bool(name: str, default: bool) -> bool: + value = _raw(name) + if value is None: + return default + normalized = value.lower() + if normalized in _TRUE_VALUES: + return True + if normalized in _FALSE_VALUES: + return False + logger.warning("Invalid %s%s=%r; using default %s", ENV_PREFIX, name, value, default) + return default diff --git a/config/http_clients.py b/config/http_clients.py new file mode 100644 index 0000000..763da9b --- /dev/null +++ b/config/http_clients.py @@ -0,0 +1,72 @@ +""" +Copyright 2025 Perforce Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" +""" +Reused ``httpx.AsyncClient`` for the clients called on every tool call, so each +request reuses the connection instead of paying a new TCP/TLS (and HTTP/2) +handshake. Built on first use, rebuilt if closed, closed on server shutdown. + +A shared client serves every caller (hosted: every user), so it never keeps +cookies: BlazeMeter answers with a ``bzm_sess`` session cookie, and a cookie +jar would send one user's session with the next user's requests. +""" +import weakref +from http.cookiejar import CookieJar, DefaultCookiePolicy +from typing import Callable, Optional + +import httpx + +_shared: "weakref.WeakSet[SharedAsyncClient]" = weakref.WeakSet() + + +class _NoCookies(DefaultCookiePolicy): + def set_ok(self, cookie, request) -> bool: + return False + + +class SharedAsyncClient: + def __init__(self, factory: Callable[[], httpx.AsyncClient]) -> None: + self._factory = factory + self._client: Optional[httpx.AsyncClient] = None + _shared.add(self) + + def get(self) -> httpx.AsyncClient: + if self._client is None or self._client.is_closed: + client = self._factory() + # A CookieJar is used as is; httpx.Cookies(...) would copy it into a default jar. + client.cookies = CookieJar(policy=_NoCookies()) + self._client = client + return self._client + + async def aclose(self) -> None: + client, self._client = self._client, None + if client is not None: + await client.aclose() + + def discard(self) -> None: + """Forget the client without closing it (its event loop is already gone).""" + self._client = None + + +async def aclose_http_clients() -> None: + """Close every shared client (server shutdown).""" + for shared in list(_shared): + await shared.aclose() + + +def discard_http_clients() -> None: + """Forget every shared client; the next ``get()`` builds a new one.""" + for shared in list(_shared): + shared.discard() diff --git a/config/ids.py b/config/ids.py new file mode 100644 index 0000000..c25de2b --- /dev/null +++ b/config/ids.py @@ -0,0 +1,25 @@ +""" +Copyright 2025 Perforce Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" +"""Random id format shared by tasks, dataframes and chat sessions.""" +import secrets + +# Lowercase base32 without i, l, o, u: unambiguous when an agent reads it back. +SIMPLE_ID_ALPHABET = "0123456789abcdefghjkmnpqrstvwxyz" +SIMPLE_ID_LENGTH = 8 + + +def random_id(length: int = SIMPLE_ID_LENGTH) -> str: + return "".join(secrets.choice(SIMPLE_ID_ALPHABET) for _ in range(length)) diff --git a/config/runtime.py b/config/runtime.py index f786c31..5a69c84 100644 --- a/config/runtime.py +++ b/config/runtime.py @@ -9,6 +9,15 @@ HttpAuthProvider, StdioAuthProvider, ) +from config.cache import CachePort, build_cache_from_env +from config.service_auth import service_caller_token +from config.storage import ( + DefaultSessionScopeResolver, + HttpSessionStorageProvider, + InMemorySessionStorageProvider, + SessionScopeResolverPort, + SessionStoragePort, +) from config.token import PerfectoToken Transport = Literal["stdio", "streamable-http"] @@ -57,6 +66,9 @@ class AppRuntime: transport: Transport auth: AuthPort user_config: dict[str, Any] + storage: SessionStoragePort + scope_resolver: SessionScopeResolverPort + cache: CachePort def resolve_user_config(self, ctx: Any) -> dict[str, Any]: user_config = dict(self.user_config) @@ -111,16 +123,51 @@ def _hydrate_ctx_user_config(ctx: Any, user_config: dict[str, Any]) -> None: setattr(target, PERFECTO_USER_CONFIG_STATE_ATTR, dict(config_copy)) +def _storage_base_url() -> str: + return os.getenv("PERFECTO_STORAGE_API_BASE_URL", "").strip() + + +def _build_storage_and_cache(transport: Transport) -> tuple[SessionStoragePort, CachePort]: + """ + 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 + storage service so dataframes, tasks and cache entries survive across workers. + """ + storage_base_url = _storage_base_url() + if transport == "streamable-http" and storage_base_url: + caller_token = service_caller_token() + if not caller_token: + raise ValueError( + "PERFECTO_MCP_STORAGE_CALLER_TOKEN is required when " + "PERFECTO_STORAGE_API_BASE_URL is set." + ) + storage: SessionStoragePort = HttpSessionStorageProvider( + base_url=storage_base_url, + caller_token=caller_token, + ) + storage.ensure_available() + return storage, build_cache_from_env( + "streamable-http", + storage_base_url=storage_base_url, + caller_token=caller_token, + ) + + # In-memory cache factory is the stdio path of build_cache_from_env. + return InMemorySessionStorageProvider(), build_cache_from_env("stdio") + + def build_runtime( transport: Transport, startup_token: Optional[PerfectoToken] = None, ) -> AppRuntime: """ - Compose auth for the selected transport. + Compose auth, session storage and the method cache for the selected transport. - - stdio: process-lifetime ``startup_token``. - - streamable-http: request-scoped Bearer auth. + - 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. """ + storage, cache = _build_storage_and_cache(transport) if transport == "stdio": stdio_user_config = { "startup_token": startup_token, @@ -131,6 +178,9 @@ def build_runtime( transport=transport, auth=StdioAuthProvider(startup_token), user_config=stdio_user_config, + storage=storage, + scope_resolver=DefaultSessionScopeResolver(), + cache=cache, ) if transport == "streamable-http": @@ -138,6 +188,9 @@ def build_runtime( transport=transport, auth=HttpAuthProvider(), user_config={}, + storage=storage, + scope_resolver=DefaultSessionScopeResolver(), + cache=cache, ) raise ValueError(f"Unknown transport: {transport}") diff --git a/config/service_auth.py b/config/service_auth.py new file mode 100644 index 0000000..ae19530 --- /dev/null +++ b/config/service_auth.py @@ -0,0 +1,22 @@ +"""How Perfecto MCP authenticates to the hosted storage service.""" +from typing import Optional + +from config.env import env_str + +# End-user credential of the current call. The storage service treats it as opaque. +CREDENTIAL_HEADER = "X-Perfecto-Credential" + + +def service_caller_token() -> str: + """MCP caller identity for the hosted storage service (``PERFECTO_MCP_STORAGE_CALLER_TOKEN``).""" + return env_str("STORAGE_CALLER_TOKEN") + + +def service_headers(caller_token: Optional[str], credential: Optional[str] = None) -> dict[str, str]: + """Headers for a hosted-service request: caller identity, plus the end-user credential if any.""" + headers: dict[str, str] = {} + if caller_token: + headers["Authorization"] = f"Bearer {caller_token}" + if credential: + headers[CREDENTIAL_HEADER] = credential + return headers diff --git a/config/session_context.py b/config/session_context.py new file mode 100644 index 0000000..ae4de70 --- /dev/null +++ b/config/session_context.py @@ -0,0 +1,75 @@ +"""Request-scoped identity and 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. +""" +from __future__ import annotations + +import contextvars +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Iterator, Optional + + +@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 +) +_current_session: contextvars.ContextVar[Optional[ChatSession]] = contextvars.ContextVar( + "perfecto_current_session", default=None +) +_current_credential: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( + "perfecto_current_credential", default=None +) + + +class SessionContextMissing(RuntimeError): + """Session-scoped state was used outside a tool call that bound a scope.""" + + +def current_identity() -> Optional[Identity]: + return _current_identity.get() + + +def current_session() -> Optional[ChatSession]: + return _current_session.get() + + +def current_credential() -> Optional[str]: + return _current_credential.get() + + +@contextmanager +def bind_session_context( + identity: Identity, + session: Optional[ChatSession] = None, + credential: Optional[str] = None, +) -> Iterator[None]: + """Bind identity, session and credential for the duration of one tool call.""" + identity_token = _current_identity.set(identity) + session_token = _current_session.set(session) + credential_token = _current_credential.set(credential) + try: + yield + finally: + _current_credential.reset(credential_token) + _current_session.reset(session_token) + _current_identity.reset(identity_token) + + +@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): + yield diff --git a/config/storage.py b/config/storage.py new file mode 100644 index 0000000..769ec9a --- /dev/null +++ b/config/storage.py @@ -0,0 +1,233 @@ +""" +Copyright 2025 Perforce Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Optional +from urllib.parse import quote + +import httpx + + + +class StorageNotConfiguredError(RuntimeError): + """Raised when session storage is used before AppRuntime wiring.""" + + +@dataclass(frozen=True) +class SessionScope: + user_id: str + mcp_session_id: str + + +@dataclass(frozen=True) +class SessionPartitionPayload: + metadata: dict[str, Any] | None = None + dataframes: dict[str, Any] | None = None + tasks: dict[str, Any] | None = None + uploaded_files: list[dict[str, Any]] | None = None + + def to_dict(self) -> dict[str, Any]: + body: dict[str, Any] = {} + if self.metadata is not None: + body["metadata"] = self.metadata + if self.dataframes is not None: + body["dataframes"] = self.dataframes + if self.tasks is not None: + body["tasks"] = self.tasks + if self.uploaded_files is not None: + body["uploaded_files"] = self.uploaded_files + return body + + +@dataclass(frozen=True) +class SessionPartition: + user_id: str + mcp_session_id: str + metadata: dict[str, Any] + dataframes: dict[str, Any] + tasks: dict[str, Any] + uploaded_files: list[dict[str, Any]] + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "SessionPartition": + return cls( + user_id=str(data.get("user_id", "")), + mcp_session_id=str(data.get("mcp_session_id", "")), + metadata=data.get("metadata", {}) or {}, + dataframes=data.get("dataframes", {}) or {}, + tasks=data.get("tasks", {}) or {}, + uploaded_files=data.get("uploaded_files", []) or [], + ) + + +class SessionStoragePort(ABC): + @abstractmethod + async def put_partition(self, scope: SessionScope, payload: SessionPartitionPayload) -> None: + raise NotImplementedError + + @abstractmethod + async def get_partition(self, scope: SessionScope) -> SessionPartition | None: + raise NotImplementedError + + @abstractmethod + async def delete_partition(self, scope: SessionScope) -> bool: + raise NotImplementedError + + +class SessionScopeResolverPort(ABC): + @abstractmethod + def resolve(self) -> SessionScope: + raise NotImplementedError + + +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. + """ + + def resolve(self) -> SessionScope: + from config.session_context import SessionContextMissing, current_identity, current_session + + identity = current_identity() + session = current_session() + if identity is None or session is None: + raise SessionContextMissing( + "No call scope is bound; session-scoped data is only available inside a tool call." + ) + return SessionScope(user_id=identity.user_id, mcp_session_id=session.session_id) + + +def resolve_session_scope( + scope_resolver: Optional[SessionScopeResolverPort] = None, +) -> SessionScope: + """Resolve partition keys for the current validated chat session.""" + resolver = scope_resolver or DefaultSessionScopeResolver() + return resolver.resolve() + + +class InMemorySessionStorageProvider(SessionStoragePort): + def __init__(self) -> None: + self._partitions: dict[tuple[str, str], SessionPartition] = {} + + async def put_partition(self, scope: SessionScope, payload: SessionPartitionPayload) -> None: + existing = self._partitions.get((scope.user_id, scope.mcp_session_id)) + metadata = existing.metadata if existing else {} + dataframes = existing.dataframes if existing else {} + tasks = existing.tasks if existing else {} + uploaded_files = existing.uploaded_files if existing else [] + + if payload.metadata is not None: + metadata = payload.metadata + if payload.dataframes is not None: + dataframes = payload.dataframes + if payload.tasks is not None: + tasks = payload.tasks + if payload.uploaded_files is not None: + uploaded_files = payload.uploaded_files + + self._partitions[(scope.user_id, scope.mcp_session_id)] = SessionPartition( + user_id=scope.user_id, + mcp_session_id=scope.mcp_session_id, + metadata=metadata, + dataframes=dataframes, + tasks=tasks, + uploaded_files=uploaded_files, + ) + + async def get_partition(self, scope: SessionScope) -> SessionPartition | None: + return self._partitions.get((scope.user_id, scope.mcp_session_id)) + + async def delete_partition(self, scope: SessionScope) -> bool: + return self._partitions.pop((scope.user_id, scope.mcp_session_id), None) is not None + + def discard_session(self, user_id: str, mcp_session_id: str) -> None: + """Drop a purged session's partition (stdio session provider ``on_purge``).""" + self._partitions.pop((user_id, mcp_session_id), None) + + +class HttpSessionStorageProvider(SessionStoragePort): + def __init__( + self, + base_url: str, + timeout_seconds: Optional[float] = None, + caller_token: Optional[str] = None, + ) -> None: + from config.env import env_float + from config.http_clients import SharedAsyncClient + + self._base_url = base_url.rstrip("/") + default_timeout = env_float("SESSION_STORAGE_TIMEOUT_SECONDS", 15.0, minimum=0.1) + self._timeout = timeout_seconds if timeout_seconds is not None else default_timeout + # Shared: partitions are read and written on most tool calls. + self._http = SharedAsyncClient(lambda: httpx.AsyncClient(http2=True, timeout=self._timeout)) + # The storage-api only serves partitions to the MCP caller identity. + self._caller_token = caller_token + + def _request_headers(self) -> dict[str, str]: + """Caller identity plus the end-user credential of the current validated call.""" + from config.session_context import current_credential + from config.service_auth import service_headers + + return service_headers(self._caller_token, current_credential()) + + def _url_for_scope(self, scope: SessionScope) -> str: + user_id = quote(scope.user_id, safe="") + mcp_session_id = quote(scope.mcp_session_id, safe="") + return f"{self._base_url}/session-partitions/{user_id}/{mcp_session_id}" + + def _health_url(self) -> str: + return f"{self._base_url}/health" + + def ensure_available(self) -> None: + """Fail fast if the storage API is unreachable.""" + with httpx.Client(http2=True, timeout=min(self._timeout, 5.0)) as client: + response = client.get(self._health_url()) + response.raise_for_status() + + async def put_partition(self, scope: SessionScope, payload: SessionPartitionPayload) -> None: + response = await self._http.get().put( + self._url_for_scope(scope), + headers=self._request_headers(), + json=payload.to_dict(), + ) + response.raise_for_status() + + async def get_partition(self, scope: SessionScope) -> SessionPartition | None: + response = await self._http.get().get( + self._url_for_scope(scope), + headers=self._request_headers(), + ) + if response.status_code == 404: + return None + response.raise_for_status() + return SessionPartition.from_dict(response.json()) + + async def delete_partition(self, scope: SessionScope) -> bool: + response = await self._http.get().delete( + self._url_for_scope(scope), + headers=self._request_headers(), + ) + response.raise_for_status() + payload = response.json() + return bool(payload.get("deleted")) + + diff --git a/docs/hosted-http.md b/docs/hosted-http.md index 5fa4bd7..db567d5 100644 --- a/docs/hosted-http.md +++ b/docs/hosted-http.md @@ -68,8 +68,24 @@ docker run --rm -p 8000:8000 \ | `FASTMCP_PORT` | Listen port (HTTP only). Also accepts `PORT` (e.g. Cloud Run) | `8000` | | `FASTMCP_STREAMABLE_HTTP_PATH` | HTTP path for the MCP endpoint | `/mcp` | | `PERFECTO_CLOUD_NAME` | Fallback cloud when the `Perfecto-Cloud-Name` header is absent | *(unset)* | +| `PERFECTO_STORAGE_API_BASE_URL` | Storage service base URL. When set on HTTP, session dataframes, async tasks and the method cache are stored there (`/session-partitions/...`, `/cache/entries`) so they survive across workers. Unset keeps them in process memory. | *(unset)* | +| `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` | + +## 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. + +- 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`). +- 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 -- No session storage service: session state is not shared across server instances. +- Without the storage service, session dataframes and tasks are not shared across server instances. - File upload / local file access is not supported by Perfecto MCP on any transport. +- Cancelling a task only stops the worker that is running it. Another instance may still finish and write the result. diff --git a/main.py b/main.py index 8f7f135..73e5cc7 100644 --- a/main.py +++ b/main.py @@ -157,6 +157,9 @@ def build_mcp_server( ## Important Guidelines - **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. +- **Result format**: Optional `result_format` on tool args: `auto` (default, store only large results), `dataframe` (always store), `raw` (always return inline). - **Don't assume / Don't invent**: If something is unclear, consult Skills/Help before responding. - **Provide resources**: Include markdown-formatted links to Perfecto help documentation when relevant. """ diff --git a/models/result.py b/models/result.py index db35e66..d8c3085 100644 --- a/models/result.py +++ b/models/result.py @@ -1,13 +1,28 @@ from typing import Any, Optional, List -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, PrivateAttr class BaseResult(BaseModel): result: Optional[Any] = Field(description="Result", default=None) + 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) 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) + tool_call_finished_at: Optional[str] = Field(description="ISO timestamp when tool action finished", default=None) + tool_call_duration_ms: Optional[int] = Field(description="Tool action duration in milliseconds", default=None) + debug: Optional[dict[str, Any]] = Field(description="Optional debug metrics for tool calls", default=None) + # In-process only (never serialized): (task_id, scope) while a task is still running. + _pending_task: Optional[tuple[str, Any]] = PrivateAttr(default=None) + + def mark_pending_task(self, task_id: str, scope: Any) -> None: + self._pending_task = (task_id, scope) + + @property + def pending_task(self) -> Optional[tuple[str, Any]]: + return self._pending_task def append_warnings(self, messages: List[str]): if not self.warning: diff --git a/pyproject.toml b/pyproject.toml index 7b3aefc..3a8c100 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "pydantic-core>=2.33.2", "pydantic-settings>=2.10.1", "lxml>=5.3.0", + "polars>=1.40.1", ] [project.scripts] diff --git a/server.py b/server.py index c21af55..7f25ebc 100644 --- a/server.py +++ b/server.py @@ -1,5 +1,7 @@ +from config.cache import configure_cache from config.runtime import AppRuntime from tools.ai_scriptless_manager import register as register_ai_scriptless_manager +from tools.async_task_manager import configure_task_storage 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 @@ -16,6 +18,9 @@ def register_tools(mcp, runtime: AppRuntime): mcp: The MCP server instance runtime: App runtime (transport + auth port and shared collaborators) """ + configure_task_storage(runtime.storage) + if runtime.cache is not None: + configure_cache(runtime.cache) register_user_manager(mcp, runtime) register_device_manager(mcp, runtime) register_execution_manager(mcp, runtime) diff --git a/tests/conftest.py b/tests/conftest.py index bbee662..baa6d1a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,6 +24,20 @@ from tools.ai_scriptless import definitions +@pytest.fixture(autouse=True) +def isolated_session_cache(): + """Keep dataframe tasks and the method cache from leaking across tests.""" + from config.cache import InMemoryTTLCache, configure_cache + from config.storage import InMemorySessionStorageProvider + from tools.async_task_manager import configure_task_storage + + configure_task_storage(InMemorySessionStorageProvider()) + configure_cache(InMemoryTTLCache(max_entries=64, sweep_interval_seconds=0)) + yield + configure_task_storage(InMemorySessionStorageProvider()) + configure_cache(InMemoryTTLCache(max_entries=64, sweep_interval_seconds=0)) + + @pytest.fixture def perfecto_token() -> PerfectoToken: return PerfectoToken("test-token", "demo") diff --git a/tests/test_dataframes_tasks_cache.py b/tests/test_dataframes_tasks_cache.py new file mode 100644 index 0000000..0163f8e --- /dev/null +++ b/tests/test_dataframes_tasks_cache.py @@ -0,0 +1,220 @@ +"""Session dataframes, async tasks and the method cache ported from BlazeMeter MCP.""" +import asyncio +from types import SimpleNamespace + +from config.call_scope import resolve_mcp_session_id, user_id_for_token +from config.runtime import build_runtime +from config.session_context import bind_call_scope +from config.storage import InMemorySessionStorageProvider, SessionScope +from config.token import PerfectoToken +from models.manager import Manager +from models.result import BaseResult, PaginationResult +from tools.async_task_manager import configure_task_storage, get_task_record +from tools.data_tools import DataTools +from tools.dataframe_manager import finalize_tool_result, query_dataframes +from tools.method_cache import ttl_cache_method +from tools.task_runtime import run_as_task + + +def _scope(user_id: str, session_id: str = "local") -> SessionScope: + return SessionScope(user_id=user_id, mcp_session_id=session_id) + + +def test_user_and_session_scope_are_stable_and_separated(): + demo = PerfectoToken("token-a", "demo") + other_cloud = PerfectoToken("token-a", "other") + other_token = PerfectoToken("token-b", "demo") + + assert user_id_for_token(demo) == user_id_for_token(PerfectoToken("token-a", "demo")) + assert user_id_for_token(demo) != user_id_for_token(other_cloud) + assert user_id_for_token(demo) != user_id_for_token(other_token) + assert user_id_for_token(None) == "anonymous" + assert resolve_mcp_session_id(None) == "local" + assert resolve_mcp_session_id(SimpleNamespace(session_id="sess-1")) == "sess-1" + request = SimpleNamespace(headers={"mcp-session-id": "from-header"}) + ctx = SimpleNamespace(session_id="from-ctx", request_context=SimpleNamespace(request=request)) + assert resolve_mcp_session_id(ctx) == "from-header" + + +def test_large_results_become_dataframes_and_stay_session_scoped(): + async def scenario(): + storage = InMemorySessionStorageProvider() + left = _scope("user-a", "s1") + right = _scope("user-b", "s1") + small = await finalize_tool_result( + BaseResult(result=[{"id": 1}]), + action="list_real_devices", + args={"result_format": "auto"}, + origin_manager="perfecto_devices", + session_storage=storage, + scope=left, + ) + assert small.result == [{"id": 1}] + + rows = [{"id": index, "name": f"device-{index}", "pad": "x" * 40} for index in range(200)] + stored = await finalize_tool_result( + BaseResult(result=rows), + action="list_real_devices", + args={}, + origin_manager="perfecto_devices", + session_storage=storage, + scope=left, + ) + assert stored.result[0]["stored_as_dataframe"] is True + table = stored.result[0]["table_name"] + + paged = await finalize_tool_result( + BaseResult(result=PaginationResult(items=rows, count=len(rows), total=len(rows), has_more=False)), + action="list_report_executions", + args={"result_format": "dataframe"}, + origin_manager="perfecto_execution", + session_storage=storage, + scope=left, + ) + assert paged.result[0]["stored_as_dataframe"] is True + + raw = await finalize_tool_result( + BaseResult(result=rows), + action="list_real_devices", + args={"result_format": "raw"}, + origin_manager="perfecto_devices", + session_storage=storage, + scope=left, + ) + assert raw.result == rows + + queried = await query_dataframes( + f"SELECT id, name FROM {table} ORDER BY id LIMIT 2 OFFSET 1", + session_storage=storage, + scope=left, + output_format="records", + ) + assert queried.get("error") is None + assert queried["result"] == [ + {"id": 1, "name": "device-1"}, + {"id": 2, "name": "device-2"}, + ] + + other = await query_dataframes( + f"SELECT id FROM {table} ORDER BY id LIMIT 1 OFFSET 0", + session_storage=storage, + scope=right, + output_format="records", + ) + assert other.get("error") + + blocked = await query_dataframes( + f"DELETE FROM {table}", + session_storage=storage, + scope=left, + ) + assert blocked.get("error") + + asyncio.run(scenario()) + + +def test_fast_task_returns_result_and_slow_task_returns_a_snapshot(): + async def scenario(): + storage = InMemorySessionStorageProvider() + configure_task_storage(storage) + scope = _scope("user-a", "local") + + class Work(Manager): + @run_as_task(fast_response_threshold_seconds=1) + async def quick(self): + return BaseResult(result=[{"done": True}]) + + @run_as_task(fast_response_threshold_seconds=0.05) + async def slow(self): + await asyncio.sleep(0.3) + return BaseResult(result=[{"done": True}]) + + ctx = SimpleNamespace() + worker = Work(ctx) + with bind_call_scope(scope.user_id, scope.mcp_session_id): + quick = await worker.quick() + assert quick.result == [{"done": True}] + assert await get_task_record("missing", scope=scope) is None + + pending = await worker.slow() + assert pending.pending_task is not None + task_id, task_scope = pending.pending_task + assert task_scope == scope + assert pending.result[0]["task_id"] == task_id + assert pending.result[0]["status"] in {"parking", "working"} + + record = await get_task_record(task_id, scope=scope) + assert record is not None + await asyncio.wait_for(record.asyncio_task, timeout=2) + finished = await get_task_record(task_id, scope=scope) + assert finished.status == "completed" + assert finished.result.result == [{"done": True}] + + tools = DataTools(ctx, storage, scope_resolver=_FixedScope(scope)) + status = await tools.tasks_status(task_id) + assert status.error is None + assert status.result[0]["status"] == "completed" + fetched = await tools.tasks_get(task_id) + assert fetched.result[0]["task_result"]["result"] == [{"done": True}] + missing = await tools.tasks_get(task_id) + assert missing.error + + asyncio.run(scenario()) + + +def test_method_cache_reuses_user_scoped_results(): + async def scenario(): + calls = {"n": 0} + + class Reader: + def __init__(self, token): + self.token = token + + @ttl_cache_method(ttl_seconds=30) + async def read(self): + calls["n"] += 1 + return BaseResult(result=[{"n": calls["n"]}]) + + reader = Reader(PerfectoToken("token-a", "demo")) + with bind_call_scope(user_id_for_token(reader.token), "local", reader.token.token): + first = await reader.read() + second = await reader.read() + assert first.result == [{"n": 1}] + assert second.result == [{"n": 1}] + assert calls["n"] == 1 + + anonymous = Reader(None) + with bind_call_scope("anonymous", "local"): + await anonymous.read() + await anonymous.read() + assert calls["n"] == 3 + + asyncio.run(scenario()) + + +def test_runtime_uses_memory_until_a_storage_url_is_configured(monkeypatch): + monkeypatch.delenv("PERFECTO_STORAGE_API_BASE_URL", raising=False) + 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.cache.__class__.__name__ == "InMemoryTTLCache" + + http_runtime = build_runtime("streamable-http") + assert http_runtime.storage.__class__.__name__ == "InMemorySessionStorageProvider" + + monkeypatch.setenv("PERFECTO_STORAGE_API_BASE_URL", "http://storage.internal") + try: + build_runtime("streamable-http") + raised = False + except ValueError as exc: + raised = True + assert "PERFECTO_MCP_STORAGE_CALLER_TOKEN" in str(exc) + assert raised + + +class _FixedScope: + def __init__(self, scope: SessionScope): + self._scope = scope + + def resolve(self) -> SessionScope: + return self._scope diff --git a/tools/ai_scriptless_manager.py b/tools/ai_scriptless_manager.py index 4074885..6045b53 100644 --- a/tools/ai_scriptless_manager.py +++ b/tools/ai_scriptless_manager.py @@ -17,7 +17,7 @@ format_step_detail, format_test_structure, format_test_variables from models.manager import Manager from models.result import BaseResult, PaginationResult -from telemetry import run_tool +from tools.runtime_tools import run_tool_with_runtime from tools.ai_scriptless_script import ( DUT_NAME, add_script_variable, @@ -63,7 +63,7 @@ validate_variable_bindings, ) from tools.ai_scriptless.definitions import SUPPORTED_ERROR_POLICIES -from tools.utils import api_request, format_sanitized_traceback, normalize_action_args +from tools.utils import api_request, format_sanitized_traceback, normalize_action_args, run_as_task STEP_PATH_REFRESH_NOTES = [ "step_path values are dot-separated positional paths (e.g. 0, 2.0, 5.b0.1); Perfecto does not persist them.", @@ -132,6 +132,7 @@ def __init__(self, ctx: Context): super().__init__(ctx) @token_verify + @run_as_task() async def list_tests(self, args: dict[str, Any]) -> BaseResult: page_size = 50 page_index = args.get("page_index", 1) @@ -163,6 +164,7 @@ async def list_tests(self, args: dict[str, Any]) -> BaseResult: return _append_ui_access_info(result, self.token.cloud_name) @token_verify + @run_as_task() async def list_filter_values(self, filter_names: list[str]) -> BaseResult: tree_url = perfecto.get_ai_scriptless_api_url(self.token.cloud_name) tree_url = tree_url + "/scripts/tree" @@ -381,6 +383,7 @@ async def _device_preflight( ) @token_verify + @run_as_task() async def execute_test( self, test_id: str, @@ -426,6 +429,7 @@ async def execute_test( return result @token_verify + @run_as_task() async def list_commands(self, checkpoint: bool = False) -> BaseResult: commands_url = perfecto.get_ai_scriptless_command_repository_url(self.token.cloud_name) commands_url = commands_url + "/commands" @@ -438,6 +442,7 @@ async def list_commands(self, checkpoint: bool = False) -> BaseResult: return result @token_verify + @run_as_task() async def get_command_definitions(self, command_ids: list[str]) -> BaseResult: if not command_ids: return BaseResult(error="command_ids is required and must not be empty") @@ -448,6 +453,7 @@ async def get_command_definitions(self, command_ids: list[str]) -> BaseResult: result_formatter=format_command_definitions) @token_verify + @run_as_task() async def view_test_structure(self, test_id: str) -> BaseResult: if not test_id: return BaseResult(error="test_id is required (itemKey from list_tests)") @@ -459,6 +465,7 @@ async def view_test_structure(self, test_id: str) -> BaseResult: return _append_ui_access_info(result, self.token.cloud_name, test_id) @token_verify + @run_as_task() async def view_test_step(self, test_id: str, step_path: str) -> BaseResult: if not test_id: return BaseResult(error="test_id is required (itemKey from list_tests)") @@ -849,6 +856,7 @@ async def move_test( return result @token_verify + @run_as_task() async def list_snapshots(self, test_id: str) -> BaseResult: if not test_id: return BaseResult(error="test_id is required (itemKey from list_tests)") @@ -869,6 +877,7 @@ async def list_snapshots(self, test_id: str) -> BaseResult: ) @token_verify + @run_as_task() async def view_snapshot(self, snapshot_id: str) -> BaseResult: if not snapshot_id: return BaseResult(error="snapshot_id is required (key from list_snapshots)") @@ -887,6 +896,7 @@ async def view_snapshot(self, snapshot_id: str) -> BaseResult: ) @token_verify + @run_as_task() async def list_test_variables(self, test_id: str) -> BaseResult: if not test_id: return BaseResult(error="test_id is required (itemKey from list_tests)") @@ -1430,7 +1440,9 @@ async def _dispatch(): return BaseResult(error=_unknown_action_error(action, args)) try: - return await run_tool(f"{TOOLS_PREFIX}_ai_scriptless", action, ctx, _dispatch) + return await run_tool_with_runtime( + runtime, f"{TOOLS_PREFIX}_ai_scriptless", action, ctx, _dispatch, tool_args=args, + ) except httpx.HTTPStatusError: return BaseResult( error=f"Error: {format_sanitized_traceback()}" diff --git a/tools/async_task_manager.py b/tools/async_task_manager.py new file mode 100644 index 0000000..dec6c97 --- /dev/null +++ b/tools/async_task_manager.py @@ -0,0 +1,666 @@ +""" +Copyright 2025 Perforce Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import asyncio +import copy +import json +import logging +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Awaitable, Callable, Dict, List, Optional, Set + +from config.env import env_int +from config.storage import ( + SessionPartitionPayload, + SessionScope, + SessionStoragePort, + StorageNotConfiguredError, + resolve_session_scope, +) +from models.result import BaseResult +from tools.dataframe_manager import finalize_tool_result +from tools.utils import ( + SIMPLE_ID_ALPHABET, + SIMPLE_ID_LENGTH, + TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME, + generate_simple_id, + normalize_simple_id, +) + +# Crockford-like base32 alphabet used by generate_simple_id / task ids (tests assert against this). +TASK_ID_ALPHABET = SIMPLE_ID_ALPHABET + +STATUS_WORKING = "working" +STATUS_PARKING = "parking" +STATUS_INPUT_REQUIRED = "input_required" +STATUS_COMPLETED = "completed" +STATUS_FAILED = "failed" +STATUS_CANCELLED = "cancelled" + +TERMINAL_STATES = {STATUS_COMPLETED, STATUS_FAILED, STATUS_CANCELLED} +ACTIVE_STATES = {STATUS_PARKING, STATUS_WORKING, STATUS_INPUT_REQUIRED} +MAX_PARALLEL_TASKS = env_int("MAX_PARALLEL_TASKS", 10, minimum=1) + +TASK_ID_MAX_ATTEMPTS = 10 +_MAX_SESSION_CACHES = env_int("MAX_SESSION_TASK_CACHES", 256, minimum=1) + +logger = logging.getLogger(__name__) + +STATUS_INFO = { + STATUS_WORKING: ( + "The request is currently being processed." + ), + STATUS_PARKING: ( + "The request is queued and waiting for an execution slot." + ), + STATUS_INPUT_REQUIRED: ( + "The receiver needs input from the requestor. " + "Use tasks_status for lightweight tracking and tasks_get to receive input requests." + ), + STATUS_COMPLETED: ( + "The request completed successfully and results are available." + ), + STATUS_FAILED: ( + "The associated request did not complete successfully." + ), + STATUS_CANCELLED: ( + "The request was cancelled before completion." + ), +} + +_semaphore = asyncio.Semaphore(MAX_PARALLEL_TASKS) + +_storage: Optional[SessionStoragePort] = None +_registry_lock = asyncio.Lock() +_session_caches: OrderedDict[tuple[str, str], "_SessionTaskCache"] = OrderedDict() +# Process-local hooks run once a task reaches a terminal state on this worker +# (e.g. drop a cached task snapshot). Not persisted: a task owned by another +# worker never fires them, and the caller's TTL applies instead. +_terminal_callbacks: Dict[tuple[str, str, str], List[Callable[[], Awaitable[None]]]] = {} + + +@dataclass +class TaskRecord: + task_id: str + action: Dict[str, Any] + created_at: float + last_updated_at: float + time_to_live_ms: Optional[int] + status: str + status_message: str + status_info: str + result: Optional[BaseResult] = None + asyncio_task: Optional[asyncio.Task] = None + started_running_at: Optional[float] = None + finished_at: Optional[float] = None + # The chat session that created the task (partition owner + session id). + user_id: str = "" + mcp_session_id: str = "" + # Process-local. Not stored. Re-raised on the fast path so tool wrappers + # keep their existing traceback and support-message formatting. + failure: Optional[BaseException] = None + + def set_status(self, status: str, status_message: str): + self.status = status + self.status_message = status_message + self.status_info = STATUS_INFO.get(status, "") + self.last_updated_at = time.time() + if status == STATUS_WORKING and self.started_running_at is None: + self.started_running_at = self.last_updated_at + if status in TERMINAL_STATES: + self.finished_at = self.last_updated_at + + def scope(self) -> SessionScope: + return SessionScope(user_id=self.user_id, mcp_session_id=self.mcp_session_id) + + +@dataclass +class _SessionTaskCache: + """In-process task handles + hydrated records for one session partition.""" + + tasks: Dict[str, TaskRecord] = field(default_factory=dict) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + hydrated: bool = False + + +def configure_task_storage(storage: SessionStoragePort) -> None: + """Bind session storage from AppRuntime (composition root).""" + global _storage, _session_caches + _storage = storage + _session_caches = OrderedDict() + _terminal_callbacks.clear() + + +def _get_storage() -> SessionStoragePort: + if _storage is None: + raise StorageNotConfiguredError( + "Session storage is not configured. " + "Wire AppRuntime via configure_task_storage() before using tasks." + ) + return _storage + + +def _session_key(scope: SessionScope) -> tuple[str, str]: + return (str(scope.user_id), str(scope.mcp_session_id)) + + +def _task_key(task_id: str) -> str: + return normalize_simple_id(task_id) + + +def _cache_is_idle(cache: _SessionTaskCache) -> bool: + if cache.lock.locked(): + return False + for record in cache.tasks.values(): + handle = record.asyncio_task + if handle is not None and not handle.done(): + return False + return True + + +def _evict_idle_session_caches() -> None: + """Drop least-recent idle caches so the map cannot grow without bound.""" + while len(_session_caches) >= _MAX_SESSION_CACHES: + evicted = False + for key, cache in list(_session_caches.items()): + if _cache_is_idle(cache): + del _session_caches[key] + evicted = True + break + if not evicted: + return + + +async def _get_or_create_cache(scope: SessionScope) -> _SessionTaskCache: + key = _session_key(scope) + async with _registry_lock: + cache = _session_caches.get(key) + if cache is not None: + _session_caches.move_to_end(key) + return cache + _evict_idle_session_caches() + cache = _SessionTaskCache() + _session_caches[key] = cache + return cache + + +def _to_iso(timestamp: float) -> str: + return datetime.fromtimestamp(timestamp).isoformat() + + +def _normalize_result(result: Any) -> BaseResult: + if isinstance(result, BaseResult): + return result + return BaseResult(result=[result]) + + +def _generate_task_id() -> str: + return generate_simple_id() + + +def _json_default(value: Any) -> Any: + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if hasattr(value, "isoformat"): + return value.isoformat() + return str(value) + + +def _result_to_storage_dict(result: Optional[BaseResult]) -> Optional[Dict[str, Any]]: + """Serialize a BaseResult to a JSON-compatible dict for Storage.""" + if result is None: + return None + try: + dumped = result.model_dump(mode="json") + except (TypeError, ValueError): + dumped = result.model_dump(mode="python") + return json.loads(json.dumps(dumped, default=_json_default)) + + +def _serialize_record(record: TaskRecord) -> Dict[str, Any]: + payload = { + "task_id": record.task_id, + "action": copy.deepcopy(record.action), + "created_at": record.created_at, + "last_updated_at": record.last_updated_at, + "time_to_live_ms": record.time_to_live_ms, + "status": record.status, + "status_message": record.status_message, + "status_info": record.status_info, + "started_running_at": record.started_running_at, + "finished_at": record.finished_at, + "user_id": record.user_id, + "mcp_session_id": record.mcp_session_id, + } + if record.result is not None: + payload["result"] = _result_to_storage_dict(record.result) + return payload + + +def _deserialize_record(payload: Dict[str, Any]) -> TaskRecord: + raw_result = payload.get("result") + result: Optional[BaseResult] = None + if isinstance(raw_result, dict): + result = BaseResult.model_validate(raw_result) + return TaskRecord( + task_id=str(payload["task_id"]), + action=dict(payload.get("action") or {}), + created_at=float(payload.get("created_at") or 0.0), + last_updated_at=float(payload.get("last_updated_at") or 0.0), + time_to_live_ms=payload.get("time_to_live_ms"), + status=str(payload.get("status") or STATUS_PARKING), + status_message=str(payload.get("status_message") or ""), + status_info=str(payload.get("status_info") or ""), + result=result, + asyncio_task=None, + started_running_at=( + float(payload["started_running_at"]) + if payload.get("started_running_at") is not None + else None + ), + finished_at=( + float(payload["finished_at"]) + if payload.get("finished_at") is not None + else None + ), + user_id=str(payload.get("user_id") or ""), + mcp_session_id=str(payload.get("mcp_session_id") or ""), + ) + + +def _merge_remote_into_cache(cache: _SessionTaskCache, remote_tasks: Dict[str, Any]) -> None: + """ + Merge Storage tasks into the in-process cache. + + Prefer live local records that still own an asyncio.Task handle (running or + finished on this worker) so cancel/wait/result identity stay intact. Purely + hydrated records (no local handle) refresh from Storage for cross-request + polling. + """ + merged: Dict[str, TaskRecord] = {} + for raw in remote_tasks.values(): + if not isinstance(raw, dict) or "task_id" not in raw: + continue + remote = _deserialize_record(raw) + key = _task_key(remote.task_id) + local = cache.tasks.get(key) + if local is not None and local.asyncio_task is not None: + merged[key] = local + else: + merged[key] = remote + for task_id, local in cache.tasks.items(): + key = _task_key(task_id) + if key not in merged: + merged[key] = local + cache.tasks = merged + + +async def _hydrate_cache(cache: _SessionTaskCache, scope: SessionScope) -> None: + """Load partition tasks into the session cache (caller holds cache.lock).""" + partition = await _get_storage().get_partition(scope) + remote = partition.tasks if partition else {} + _merge_remote_into_cache(cache, remote) + cache.hydrated = True + + +async def _commit_cache( + cache: _SessionTaskCache, + scope: SessionScope, + snapshot_ids: Optional[Set[str]] = None, +) -> None: + """ + Persist the tasks map for this SessionScope. + + Atomicity: one put_partition of the full tasks map. + Consistency: partial payload (tasks only) so dataframes/metadata/files stay. + Isolation: caller holds the in-process session lock. Before PUT, re-read + and union keys added by other workers; drop ids this operation removed. + Live local asyncio handles win on overlapping keys via cache.tasks overlay. + Durability: delegated to SessionStoragePort. + + Same-key concurrent writes and the GET/PUT race can still last-write-win; + closing that window requires Storage CAS/etag (not implemented here). + """ + current_ids = set(cache.tasks.keys()) + removed_ids = ( + {_task_key(task_id) for task_id in snapshot_ids} - current_ids + if snapshot_ids is not None + else set() + ) + partition = await _get_storage().get_partition(scope) + remote = partition.tasks if partition else {} + + merged: Dict[str, TaskRecord] = {} + for raw in remote.values(): + if not isinstance(raw, dict) or "task_id" not in raw: + continue + record = _deserialize_record(raw) + merged[_task_key(record.task_id)] = record + + for removed_id in removed_ids: + merged.pop(removed_id, None) + + merged.update(cache.tasks) + for removed_id in removed_ids: + merged.pop(removed_id, None) + + cache.tasks = merged + tasks = { + task_id: _serialize_record(record) + for task_id, record in merged.items() + } + await _get_storage().put_partition( + scope, + SessionPartitionPayload(tasks=tasks), + ) + + +async def _allocate_task_id(cache: _SessionTaskCache) -> str: + for _ in range(TASK_ID_MAX_ATTEMPTS): + candidate = _generate_task_id() + if candidate not in cache.tasks: + return candidate + + logger.error( + "Unable to allocate task id. attempts=%s id_length=%s alphabet=crockford32 active_pool_size=%s", + TASK_ID_MAX_ATTEMPTS, + SIMPLE_ID_LENGTH, + len(cache.tasks), + ) + raise RuntimeError( + f"Unable to allocate unique {SIMPLE_ID_LENGTH}-char task id after {TASK_ID_MAX_ATTEMPTS} attempts." + ) + + +async def _set_status_and_persist( + record: TaskRecord, + status: str, + status_message: str, +) -> None: + record.set_status(status, status_message) + scope = record.scope() + cache = await _get_or_create_cache(scope) + async with cache.lock: + cache.tasks[_task_key(record.task_id)] = record + await _commit_cache(cache, scope) + + +async def _task_runner(task_record: TaskRecord, coro_factory: Callable[[], Awaitable[Any]]): + try: + # Inside the try so a failure or cancellation here still drains terminal callbacks. + await _set_status_and_persist( + task_record, + STATUS_PARKING, + "Task is waiting for an available execution slot.", + ) + async with _semaphore: + await _set_status_and_persist( + task_record, + STATUS_WORKING, + "Task is currently running.", + ) + if task_record.time_to_live_ms is None: + action_result = await coro_factory() + else: + action_result = await asyncio.wait_for( + coro_factory(), + timeout=task_record.time_to_live_ms / 1000, + ) + normalized = _normalize_result(action_result) + if not task_record.action.get("disable_dataframe_materialization"): + result_format = str(task_record.action.get("result_format", "auto")).strip().lower() + origin_action = str(task_record.action.get("method", "unknown")).strip() or "unknown" + # Task runner owns materialization for parked work (pollers read Storage). + # run_tool_with_runtime may finalize again on the fast path; that path is + # idempotent for already-stored dataframe payloads. + normalized = await finalize_tool_result( + normalized, + action=origin_action, + args={"result_format": result_format}, + origin_manager=str(task_record.action.get("manager", "unknown")), + session_storage=_get_storage(), + scope=task_record.scope(), + excluded_actions=TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME, + ) + task_record.result = normalized + if normalized.error: + await _set_status_and_persist( + task_record, + STATUS_FAILED, + f"Task finished with error: {normalized.error}", + ) + else: + await _set_status_and_persist( + task_record, + STATUS_COMPLETED, + "Task finished successfully.", + ) + except asyncio.TimeoutError: + timeout_message = ( + f"Task timed out after {task_record.time_to_live_ms} ms." + if task_record.time_to_live_ms is not None + else "Task timed out." + ) + task_record.result = BaseResult(error=timeout_message) + await _set_status_and_persist(task_record, STATUS_CANCELLED, timeout_message) + except asyncio.CancelledError: + cancel_message = "Task was cancelled." + task_record.result = BaseResult(error=cancel_message) + await _set_status_and_persist(task_record, STATUS_CANCELLED, cancel_message) + except Exception as exc: + error_message = f"Task failed with exception: {str(exc)}" + task_record.failure = exc + task_record.result = BaseResult(error=error_message) + await _set_status_and_persist(task_record, STATUS_FAILED, error_message) + finally: + await _run_terminal_callbacks(task_record) + + +def _terminal_callback_key(task_id: str, scope: SessionScope) -> tuple[str, str, str]: + user_id, mcp_session_id = _session_key(scope) + return user_id, mcp_session_id, _task_key(task_id) + + +async def _run_terminal_callbacks(task_record: TaskRecord) -> None: + callbacks = _terminal_callbacks.pop( + _terminal_callback_key(task_record.task_id, task_record.scope()), [] + ) + for callback in callbacks: + try: + await callback() + except Exception: + logger.exception("Task terminal callback failed for task %s", task_record.task_id) + + +async def add_task_terminal_callback( + task_id: str, + scope: SessionScope, + callback: Callable[[], Awaitable[None]], +) -> None: + """ + Run ``callback`` once the task reaches a terminal state on this worker. + + Runs it right away when the task is already terminal or unknown here. + """ + record = await get_task_record(task_id, scope=scope) + # No await between the status check and the registration: the runner sets the + # terminal status before it drains callbacks, so a callback cannot be missed. + if record is None or is_terminal_status(record.status): + await callback() + return + _terminal_callbacks.setdefault(_terminal_callback_key(task_id, scope), []).append(callback) + + +async def submit_task( + action: Dict[str, Any], + coro_factory: Callable[[], Awaitable[Any]], + time_to_live_ms: Optional[int] = None, + *, + scope: SessionScope, +) -> str: + cache = await _get_or_create_cache(scope) + async with cache.lock: + await _hydrate_cache(cache, scope) + now = time.time() + task_id = await _allocate_task_id(cache) + task_record = TaskRecord( + task_id=task_id, + action=action, + created_at=now, + last_updated_at=now, + time_to_live_ms=time_to_live_ms, + status=STATUS_PARKING, + status_message="Task accepted and pending scheduling.", + status_info=STATUS_INFO[STATUS_PARKING], + user_id=scope.user_id, + mcp_session_id=scope.mcp_session_id, + ) + cache.tasks[task_id] = task_record + await _commit_cache(cache, scope) + # Assign the handle before releasing the lock so concurrent get/cancel + # cannot hydrate a deserialized copy that drops asyncio.Task identity. + async_task = asyncio.create_task(_task_runner(task_record, coro_factory)) + task_record.asyncio_task = async_task + return task_id + + +async def get_task_record( + task_id: str, + *, + scope: SessionScope, +) -> Optional[TaskRecord]: + normalized = _task_key(task_id) + cache = await _get_or_create_cache(scope) + async with cache.lock: + await _hydrate_cache(cache, scope) + return cache.tasks.get(normalized) + + +async def remove_task( + task_id: str, + *, + scope: SessionScope, +) -> bool: + normalized = _task_key(task_id) + cache = await _get_or_create_cache(scope) + async with cache.lock: + await _hydrate_cache(cache, scope) + snapshot_ids = set(cache.tasks.keys()) + removed = cache.tasks.pop(normalized, None) is not None + if removed: + await _commit_cache(cache, scope, snapshot_ids=snapshot_ids) + return removed + + +def task_snapshot(task_record: TaskRecord, include_result: bool = False) -> Dict[str, Any]: + snapshot = { + "task_id": task_record.task_id, + "action": task_record.action, + "created_at": task_record.created_at, + "created_at_iso": _to_iso(task_record.created_at), + "last_updated_at": task_record.last_updated_at, + "last_updated_at_iso": _to_iso(task_record.last_updated_at), + "time_to_live_ms": task_record.time_to_live_ms, + "status": task_record.status, + "status_message": task_record.status_message, + "status_info": task_record.status_info, + "started_running_at": task_record.started_running_at, + "started_running_at_iso": _to_iso(task_record.started_running_at) if task_record.started_running_at else None, + "finished_at": task_record.finished_at, + "finished_at_iso": _to_iso(task_record.finished_at) if task_record.finished_at else None, + } + if include_result and task_record.result is not None: + snapshot["task_result"] = task_record.result.model_dump() + return snapshot + + +async def list_tasks( + status_list: Optional[List[str]] = None, + *, + scope: SessionScope, +) -> List[TaskRecord]: + cache = await _get_or_create_cache(scope) + async with cache.lock: + await _hydrate_cache(cache, scope) + records = list(cache.tasks.values()) + if not status_list: + return records + expected = {status.lower() for status in status_list} + return [task for task in records if task.status.lower() in expected] + + +def is_terminal_status(status: str) -> bool: + return status in TERMINAL_STATES + + +def is_active_status(status: str) -> bool: + return status in ACTIVE_STATES + + +async def cancel_task( + task_id: str, + *, + scope: SessionScope, +) -> Optional[TaskRecord]: + """ + Request cancellation for a task in this session partition. + + Terminal tasks (completed/failed/cancelled) are left unchanged. + Live local asyncio handles are cancelled on this worker only. + Active tasks without a local handle (typical hosted multi-worker case) are + marked cancelled in Storage for status visibility, but the owning worker's + coroutine may still finish and overwrite status — see hosted runbook. + """ + normalized_task_id = _task_key(task_id) + cache = await _get_or_create_cache(scope) + async with cache.lock: + await _hydrate_cache(cache, scope) + task_record = cache.tasks.get(normalized_task_id) + if not task_record: + return None + if is_terminal_status(task_record.status): + return task_record + if task_record.asyncio_task and not task_record.asyncio_task.done(): + task_record.asyncio_task.cancel() + # CancelledError path in _task_runner persists terminal state. + return task_record + # Active in Storage but no cancelable handle on this process. + task_record.set_status( + STATUS_CANCELLED, + ( + "Cancel recorded in session Storage, but this worker has no local " + "asyncio handle. If another worker owns the coroutine, it may still " + "run to completion and overwrite status." + ), + ) + if task_record.result is None: + task_record.result = BaseResult( + error=( + "Task cancel was recorded without a local execution handle. " + "Execution affinity is process-local in hosted MCP." + ) + ) + await _commit_cache(cache, scope) + return task_record + + +def session_scope_from_manager(manager: Any) -> SessionScope: + """Storage partition keys of the validated session (using the manager's resolver, if any).""" + return resolve_session_scope(scope_resolver=getattr(manager, "scope_resolver", None)) diff --git a/tools/call_state.py b/tools/call_state.py new file mode 100644 index 0000000..2a70558 --- /dev/null +++ b/tools/call_state.py @@ -0,0 +1,55 @@ +"""Process-wide context for dataframe materialization and async tasks.""" +from __future__ import annotations + +import contextvars + +from config.ids import SIMPLE_ID_ALPHABET, SIMPLE_ID_LENGTH, random_id + +_task_management_enabled = contextvars.ContextVar("task_management_enabled", default=False) +_result_format_context = contextvars.ContextVar("result_format_context", default="auto") +_disable_dataframe_materialization = contextvars.ContextVar( + "disable_dataframe_materialization", default=False +) +_cache_bypass = contextvars.ContextVar("cache_bypass", default=False) + +# Tool actions that must stay inline under result_format=auto. +TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME = frozenset({ + "tasks_get", + "tasks_list", + "tasks_status", + "tasks_cancel", + "tasks_remove", + "dataframes_list", + "dataframes_get", + "dataframes_schema_groups", + "dataframes_query", + "dataframes_remove", + "dataframes_clear", + "dataframes_sql_help", + "version", + "check_updates", +}) + + +def generate_simple_id() -> str: + return random_id(SIMPLE_ID_LENGTH) + + +def normalize_simple_id(simple_id: str) -> str: + return str(simple_id).strip().lower() + + +def set_disable_dataframe_materialization(disabled: bool) -> contextvars.Token: + return _disable_dataframe_materialization.set(bool(disabled)) + + +def reset_disable_dataframe_materialization(token: contextvars.Token) -> None: + _disable_dataframe_materialization.reset(token) + + +def set_result_format(result_format: str) -> contextvars.Token: + return _result_format_context.set(result_format or "auto") + + +def reset_result_format(token: contextvars.Token) -> None: + _result_format_context.reset(token) diff --git a/tools/data_tools.py b/tools/data_tools.py new file mode 100644 index 0000000..379438c --- /dev/null +++ b/tools/data_tools.py @@ -0,0 +1,606 @@ +"""Session dataframe and async task actions exposed on perfecto_tools.""" +from __future__ import annotations + +import asyncio +import re +import time +from typing import Any, Dict, Optional + +from mcp.server.fastmcp import Context + +from config.storage import SessionScope, SessionScopeResolverPort, SessionStoragePort +from models.manager import Manager +from models.result import BaseResult +from tools.async_task_manager import ( + cancel_task, + get_task_record, + is_active_status, + is_terminal_status, + list_tasks, + remove_task, + task_snapshot, +) +from tools.dataframe_manager import ( + INVALID_RESULT_FORMAT_ERROR, + clear_dataframes, + get_dataframe_metadata, + get_sql_capabilities, + group_dataframe_schemas, + list_dataframes_metadata, + normalize_result_format, + query_dataframes, + register_dataframe, + remove_dataframes, + serialize_result_to_compact_json, + stored_as_dataframe_payload, +) +from tools.task_runtime import run_as_task + +class DataTools(Manager): + """Session-scoped dataframe + async task tools backed by SessionStoragePort.""" + + def __init__( + self, + ctx: Context, + session_storage: SessionStoragePort, + scope_resolver: SessionScopeResolverPort, + ): + super().__init__(ctx) + self.session_storage = session_storage + self.scope_resolver = scope_resolver + + def _scope(self) -> SessionScope: + return self.scope_resolver.resolve() + + @staticmethod + def _poll_args_error(wait_for_terminal_ms: int, poll_interval_ms: int) -> Optional[BaseResult]: + if wait_for_terminal_ms < 0: + return BaseResult(error="wait_for_terminal_ms must be greater than or equal to 0.") + if poll_interval_ms <= 0: + return BaseResult(error="poll_interval_ms must be greater than 0.") + return None + + @staticmethod + def _should_continue_polling(status: str) -> bool: + return status in {"parking", "working", "input_required"} + + @staticmethod + def _to_snake_case(value: str) -> str: + s1 = re.sub(r"(.)([A-Z][a-z]+)", r"\1_\2", value) + return re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", s1).lower() + + @classmethod + def _operation_name(cls, action_payload: Dict[str, Any]) -> str: + manager = str(action_payload.get("manager", "tool")) + method = str(action_payload.get("method", "action")) + tool_name = manager[:-7] if manager.endswith("Manager") else manager + tool_name = cls._to_snake_case(tool_name) + return f"{tool_name}.{method}" + + + async def _batch_summary_line(self) -> str: + records = await list_tasks(scope=self._scope()) + counts = { + "completed": 0, + "working": 0, + "parking": 0, + "failed": 0, + "cancelled": 0, + "input_required": 0, + } + for record in records: + status = str(record.status).lower() + if status in counts: + counts[status] += 1 + summary = ( + f"batch summary: total={len(records)} completed={counts['completed']} " + f"working={counts['working']} parking={counts['parking']} failed={counts['failed']}" + ) + if counts["cancelled"] > 0: + summary += f" cancelled={counts['cancelled']}" + if counts["input_required"] > 0: + summary += f" input_required={counts['input_required']}" + return summary + + @classmethod + def _task_status_line( + cls, + task_record, + poll_count: Optional[int], + elapsed_seconds: int, + next_poll_seconds: Optional[float], + window_seconds: Optional[float] = None, + include_polling_prefix: bool = True, + ) -> str: + operation = cls._operation_name(task_record.action) + prefix = "Polling " if include_polling_prefix else "" + line = f"{prefix}{task_record.task_id}[{operation}] ({task_record.status})" + if poll_count is not None: + line += f" attempt={poll_count}" + line += f" elapsed={elapsed_seconds}s" + if window_seconds is not None: + line += f"/{int(window_seconds)}s" + if next_poll_seconds is not None and cls._should_continue_polling(task_record.status): + line += f" next={int(next_poll_seconds)}s" + if task_record.status == "parking" and task_record.status_message: + line += f" note={repr(task_record.status_message)}" + return line + + async def _polling_message( + self, + task_record, + poll_count: int, + elapsed_seconds: int, + next_poll_seconds: float, + window_seconds: float, + ) -> str: + line = self._task_status_line( + task_record=task_record, + poll_count=poll_count, + elapsed_seconds=elapsed_seconds, + next_poll_seconds=next_poll_seconds, + window_seconds=window_seconds, + include_polling_prefix=True, + ) + return f"{line} | {await self._batch_summary_line()}" + + async def _polling_finished_message(self, task_record, elapsed_seconds: int) -> str: + line = self._task_status_line( + task_record=task_record, + poll_count=None, + elapsed_seconds=elapsed_seconds, + next_poll_seconds=None, + window_seconds=None, + include_polling_prefix=True, + ) + return f"{line} | {await self._batch_summary_line()}" + + async def _wait_for_terminal( + self, + task_id: str, + task_record, + wait_for_terminal_ms: int, + poll_interval_ms: int, + ): + """Poll Storage until terminal, timeout, or the record disappears. + + Returns (record, polling_exhausted, error_result). error_result is set + when the task vanishes mid-poll. + """ + if wait_for_terminal_ms <= 0 or is_terminal_status(task_record.status): + return task_record, False, None + + scope = self._scope() + start_time = time.monotonic() + wait_seconds = wait_for_terminal_ms / 1000.0 + poll_seconds = poll_interval_ms / 1000.0 + attempt = 0 + polling_exhausted = False + + while True: + task_record = await get_task_record(task_id, scope=scope) + if not task_record: + return None, False, BaseResult(error=f"Task ID {task_id} was not found.") + if is_terminal_status(task_record.status): + break + + elapsed = time.monotonic() - start_time + if elapsed >= wait_seconds: + polling_exhausted = True + break + + attempt += 1 + progress = min(100.0, (elapsed / wait_seconds) * 100.0) if wait_seconds > 0 else 100.0 + try: + await self.ctx.report_progress( + progress=progress, + total=100.0, + message=await self._polling_message( + task_record=task_record, + poll_count=attempt, + elapsed_seconds=int(elapsed), + next_poll_seconds=poll_seconds, + window_seconds=wait_seconds, + ), + ) + except Exception: + pass + + remaining = wait_seconds - elapsed + await asyncio.sleep(min(poll_seconds, remaining)) + + try: + final_elapsed = min(wait_seconds, time.monotonic() - start_time) + await self.ctx.report_progress( + progress=100.0, + total=100.0, + message=await self._polling_finished_message( + task_record=task_record, + elapsed_seconds=int(final_elapsed), + ), + ) + except Exception: + pass + + return task_record, polling_exhausted, None + + async def tasks_get( + self, + task_id: str, + remove_on_terminal: bool = True, + wait_for_terminal_ms: int = 0, + poll_interval_ms: int = 1000, + ) -> BaseResult: + if error := self._poll_args_error(wait_for_terminal_ms, poll_interval_ms): + return error + + scope = self._scope() + task_record = await get_task_record(task_id, scope=scope) + if not task_record: + return BaseResult(error=f"Task ID {task_id} was not found.") + + task_record, polling_exhausted, missing = await self._wait_for_terminal( + task_id, task_record, wait_for_terminal_ms, poll_interval_ms, + ) + if missing: + return missing + + terminal = is_terminal_status(task_record.status) + snapshot = task_snapshot(task_record, include_result=terminal) + snapshot["should_continue_polling"] = self._should_continue_polling(task_record.status) + snapshot["next_poll_after_ms"] = poll_interval_ms if snapshot["should_continue_polling"] else 0 + + if terminal and remove_on_terminal: + await remove_task(task_id, scope=scope) + return BaseResult( + result=[snapshot], + info=[ + "Task result retrieved successfully and removed automatically from the task registry. " + "It will not be available in subsequent queries." + ], + ) + + if terminal: + return BaseResult( + result=[snapshot], + info=[ + "Task result retrieved successfully and kept in the task registry. " + "Use tasks_remove to delete it when no longer needed." + ], + ) + + return BaseResult( + result=[snapshot], + info=[ + ( + "Task is still in progress after the polling window. Query tasks_status again in a few moments." + if polling_exhausted + else "Task is still in progress. Query tasks_status again in a few moments to check updated state." + ) + ], + ) + + async def tasks_status( + self, + task_id: str, + wait_for_terminal_ms: int = 0, + poll_interval_ms: int = 1000, + ) -> BaseResult: + if error := self._poll_args_error(wait_for_terminal_ms, poll_interval_ms): + return error + + task_record = await get_task_record(task_id, scope=self._scope()) + if not task_record: + return BaseResult(error=f"Task ID {task_id} was not found.") + + task_record, polling_exhausted, missing = await self._wait_for_terminal( + task_id, task_record, wait_for_terminal_ms, poll_interval_ms, + ) + if missing: + return missing + + snapshot = task_snapshot(task_record, include_result=False) + snapshot["should_continue_polling"] = self._should_continue_polling(task_record.status) + snapshot["next_poll_after_ms"] = poll_interval_ms if snapshot["should_continue_polling"] else 0 + info_message = ( + "Task is terminal. Use tasks_get to retrieve task_result when needed." + if is_terminal_status(task_record.status) + else ( + "Task is still in progress after the polling window. Query tasks_status again in a few moments." + if polling_exhausted + else "Task is still in progress. Query tasks_status again in a few moments to check updated state." + ) + ) + return BaseResult(result=[snapshot], info=[info_message]) + + async def tasks_list( + self, + status: Optional[str] = None, + status_list: Optional[list[str]] = None, + ) -> BaseResult: + filters = status_list if status_list else ([status] if status else None) + records = await list_tasks(filters, scope=self._scope()) + snapshots = [] + for record in records: + base_snapshot = task_snapshot(record, include_result=False) + snapshots.append( + { + "task_id": record.task_id, + "operation": self._operation_name(record.action), + "status": record.status, + "status_message": record.status_message, + "created_at": record.created_at, + "created_at_iso": base_snapshot["created_at_iso"], + "last_updated_at": record.last_updated_at, + "last_updated_at_iso": base_snapshot["last_updated_at_iso"], + "started_running_at": record.started_running_at, + "started_running_at_iso": base_snapshot["started_running_at_iso"], + "finished_at": record.finished_at, + "finished_at_iso": base_snapshot["finished_at_iso"], + "time_to_live_ms": record.time_to_live_ms, + } + ) + return BaseResult(result=snapshots, total=len(snapshots), has_more=False) + + async def tasks_cancel(self, task_id: str) -> BaseResult: + scope = self._scope() + prior = await get_task_record(task_id, scope=scope) + if not prior: + return BaseResult(error=f"Task ID {task_id} was not found.") + was_terminal = is_terminal_status(prior.status) + had_local_handle = bool(prior.asyncio_task and not prior.asyncio_task.done()) + + task_record = await cancel_task(task_id, scope=scope) + if not task_record: + return BaseResult(error=f"Task ID {task_id} was not found.") + + if was_terminal: + info = ( + f"Task was already terminal ({prior.status}); status was left unchanged. " + "Cancel does not rewrite completed/failed tasks." + ) + elif had_local_handle: + info = ( + "Task cancellation was requested on this worker's local asyncio handle." + ) + else: + info = ( + "Cancel was recorded in session Storage, but this worker has no local " + "asyncio handle. Hosted execution affinity is process-local: the owning " + "worker may still run the coroutine to completion and overwrite status." + ) + return BaseResult( + result=[task_snapshot(task_record, include_result=True)], + info=[info], + ) + + async def tasks_remove(self, task_id: str) -> BaseResult: + scope = self._scope() + task_record = await get_task_record(task_id, scope=scope) + if not task_record: + return BaseResult(error=f"Task ID {task_id} was not found.") + + cancel_requested = False + if is_active_status(task_record.status): + cancel_requested = True + await cancel_task(task_id, scope=scope) + if task_record.asyncio_task: + try: + await asyncio.wait_for(asyncio.shield(task_record.asyncio_task), timeout=0.2) + except asyncio.TimeoutError: + pass + except asyncio.CancelledError: + pass + + snapshot = task_snapshot(task_record, include_result=is_terminal_status(task_record.status)) + removed = await remove_task(task_id, scope=scope) + if not removed: + return BaseResult(error=f"Task ID {task_id} could not be removed.") + + info_message = ( + "Task was active. Cancellation was requested before removal." + if cancel_requested + else "Task was removed from task registry." + ) + return BaseResult( + result=[{ + "task_id": task_id, + "removed": True, + "cancel_requested": cancel_requested, + "task_snapshot": snapshot, + }], + info=[info_message], + ) + + @run_as_task() + async def dataframes_list(self) -> BaseResult: + scope = self._scope() + metadata = await list_dataframes_metadata( + session_storage=self.session_storage, + scope=scope, + include_schema=False, + ) + return BaseResult( + result=metadata, + total=len(metadata), + has_more=False, + info=[ + "Schema is omitted in dataframes_list to reduce payload size.", + "Use dataframes_schema_groups to compare shared/different schemas across dataframes.", + "Use dataframes_get for full metadata and schema of a specific dataframe.", + ], + ) + + @run_as_task() + async def dataframes_get(self, dataframe_id: str) -> BaseResult: + scope = self._scope() + metadata = await get_dataframe_metadata( + dataframe_id, + session_storage=self.session_storage, + scope=scope, + ) + if not metadata: + return BaseResult( + error=( + f"Dataframe ID {dataframe_id} was not found. " + "Use dataframes_list to discover available dataframes." + ) + ) + return BaseResult(result=[metadata]) + + @run_as_task() + async def dataframes_schema_groups( + self, + dataframe_id_list: Optional[list[str]] = None, + ) -> BaseResult: + scope = self._scope() + grouped = await group_dataframe_schemas( + session_storage=self.session_storage, + scope=scope, + dataframe_id_list=dataframe_id_list, + ) + mandatory_review_groups = [ + grp for grp in grouped.get("groups", []) + if isinstance(grp, dict) and str(grp.get("varying_columns", "")).strip() + ] + info_messages = [ + "Grouped dataframe schemas by top-level contract and per-column schema variations.", + "Dataframe ID lists are deduplicated in 'df_sets'; groups and variations reference them via 'df_ref'.", + "If dataframe_id_list is omitted, all current dataframes are included.", + ] + if mandatory_review_groups: + info_messages.append( + "CRITICAL: Column variations were detected. Perform mandatory detailed schema " + "review for varying columns before the final query." + ) + info_messages.append( + "IMPORTANT: Before planning and executing the final dataframe query, " + "call dataframes_sql_help synchronously in a separate call." + ) + return BaseResult(result=[grouped], info=info_messages) + + @run_as_task() + async def dataframes_query( + self, + sql: str, + output_format: str = "matrix", + result_format: str = "auto", + ) -> BaseResult: + scope = self._scope() + normalized_result_format = normalize_result_format(result_format) + if normalized_result_format == "invalid": + return BaseResult(error=INVALID_RESULT_FORMAT_ERROR) + # Store path always queries as records so register_dataframe can rebuild rows. + effective_output_format = ( + "records" if normalized_result_format == "dataframe" else output_format + ) + info_messages = [ + "Query executed successfully against the session SQL context.", + "ORDER BY + LIMIT + OFFSET are mandatory in every dataframe query.", + "Use a prudent default page size of up to 100 rows (for example, LIMIT 100 OFFSET 0), " + "then continue paging as needed.", + ] + query_response = await query_dataframes( + sql, + session_storage=self.session_storage, + scope=scope, + output_format=effective_output_format, + ) + if query_response.get("error"): + return BaseResult(error=query_response["error"]) + + if normalized_result_format == "dataframe": + rows = query_response["result"] or [] + try: + json_size_chars = len(serialize_result_to_compact_json(rows)) + except Exception as exc: + return BaseResult(error=f"Could not serialize query result: {exc}") + metadata = await register_dataframe( + result=rows, + origin_manager="perfecto_tools", + origin_action="dataframes_query", + json_size_chars=json_size_chars, + session_storage=self.session_storage, + scope=scope, + ) + return BaseResult( + result=[stored_as_dataframe_payload(metadata)], + info=info_messages + [ + "result_format=dataframe stored the query output as a new session dataframe." + ], + ) + + return BaseResult( + result=query_response["result"], + total=query_response["rows"], + has_more=False, + info=info_messages, + ) + + @run_as_task() + async def dataframes_remove(self, dataframe_id_list: list[str]) -> BaseResult: + empty_list_error = ( + "Missing required args for action 'dataframes_remove': " + "dataframe_id_list must be a non-empty list of dataframe IDs." + ) + if not dataframe_id_list or not isinstance(dataframe_id_list, list): + return BaseResult(error=empty_list_error) + ids = [str(df_id).strip() for df_id in dataframe_id_list if str(df_id).strip()] + if not ids: + return BaseResult(error=empty_list_error) + + unique_ids = list(dict.fromkeys(ids)) + outcome = await remove_dataframes(unique_ids, self.session_storage, self._scope()) + removed_ids = set(outcome["removed"]) + missing_ids = outcome["missing"] + removed_count = len(outcome["removed"]) + removed_results = [ + {"dataframe_id": df_id, "removed": df_id in removed_ids} + for df_id in unique_ids + ] + + if len(unique_ids) == 1 and removed_count == 0: + only_id = unique_ids[0] + return BaseResult( + error=( + f"Dataframe ID {only_id} was not found. " + "Use dataframes_list to discover available dataframes." + ) + ) + + info_messages = [ + f"Requested removal for {len(unique_ids)} dataframe(s). " + f"Removed: {removed_count}. Missing: {len(missing_ids)}." + ] + if removed_count > 0: + info_messages.append("Removed dataframes were unregistered from SQL context.") + if missing_ids: + info_messages.append( + "Some dataframe IDs were not found: " + ", ".join(missing_ids) + "." + ) + return BaseResult( + result=removed_results, + total=len(removed_results), + has_more=False, + info=info_messages, + ) + + @run_as_task() + async def dataframes_clear(self) -> BaseResult: + removed_count = await clear_dataframes( + session_storage=self.session_storage, + scope=self._scope(), + ) + return BaseResult( + result=[{"removed_count": removed_count, "remaining": 0}], + info=[ + "All session dataframes were removed and unregistered from SQL context." + ], + ) + + async def dataframes_sql_help(self) -> BaseResult: + return BaseResult( + result=[get_sql_capabilities()], + info=["Only read-only SQL is allowed in dataframe queries."], + ) + diff --git a/tools/dataframe_manager.py b/tools/dataframe_manager.py new file mode 100644 index 0000000..f658d94 --- /dev/null +++ b/tools/dataframe_manager.py @@ -0,0 +1,1193 @@ +""" +Copyright 2025 Perforce Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" +import asyncio +import hashlib +import json +import logging +import re +from collections import OrderedDict +from contextlib import asynccontextmanager +from dataclasses import dataclass, asdict, field +from datetime import datetime, UTC +from typing import Any, AsyncIterator, Dict, List, Optional, Set + +import polars as pl + +from config.storage import ( + SessionPartitionPayload, + SessionScope, + SessionScopeResolverPort, + SessionStoragePort, + resolve_session_scope, +) +from config.session_context import SessionContextMissing +from models.result import BaseResult +from tools.utils import generate_simple_id, SIMPLE_ID_LENGTH + +logger = logging.getLogger(__name__) + +DATAFRAME_JSON_SIZE_THRESHOLD = 8000 + +DATAFRAME_ID_MAX_ATTEMPTS = 10 + +ALLOWED_RESULT_FORMATS = {"auto", "dataframe", "raw"} +INVALID_RESULT_FORMAT_ERROR = ( + "Invalid result_format value. Allowed values: auto, dataframe, raw." +) +MISSING_STORAGE_ERROR = "Session storage is required to persist dataframes." + +# In-process mutexes only. Hosted MCP workers do not share this map; +# SessionStoragePort merge-on-commit is the cross-instance safeguard (not CAS). +_MAX_SESSION_LOCKS = 256 +_locks_guard = asyncio.Lock() +_overflow_lock = asyncio.Lock() +_session_locks: OrderedDict[tuple[str, str], asyncio.Lock] = OrderedDict() +_overflow_lock: Optional[asyncio.Lock] = None + + +@dataclass +class _SessionWorkingSet: + """Ephemeral Polars/SQL state loaded from Storage for one operation.""" + + dataframes: Dict[str, "DataFrameRecord"] = field(default_factory=dict) + sql_context: Any = field(default_factory=pl.SQLContext) + + +def _scope_key(scope: SessionScope) -> tuple[str, str]: + return (str(scope.user_id), str(scope.mcp_session_id)) + + +def _evict_unlocked_session_locks() -> None: + """Drop least-recent unlocked locks so the map cannot grow without bound.""" + while len(_session_locks) >= _MAX_SESSION_LOCKS: + evicted = False + for key, lock in list(_session_locks.items()): + if not lock.locked(): + del _session_locks[key] + evicted = True + break + if not evicted: + return + + +def _shared_overflow_lock() -> asyncio.Lock: + global _overflow_lock + if _overflow_lock is None: + _overflow_lock = asyncio.Lock() + return _overflow_lock + + +async def _lock_for(scope: SessionScope) -> asyncio.Lock: + key = _scope_key(scope) + async with _locks_guard: + lock = _session_locks.get(key) + if lock is not None: + _session_locks.move_to_end(key) + return lock + _evict_unlocked_session_locks() + if len(_session_locks) >= _MAX_SESSION_LOCKS: + return _shared_overflow_lock() + lock = asyncio.Lock() + _session_locks[key] = lock + return lock + + +def _serialize_record(record: "DataFrameRecord") -> Dict[str, Any]: + payload = record.to_metadata(include_schema=True) + payload["data"] = record.dataframe.to_dicts() + return payload + + +def _deserialize_record(payload: Dict[str, Any]) -> "DataFrameRecord": + rows = payload.get("data") or [] + dataframe = pl.DataFrame(rows) if rows else pl.DataFrame() + schema_rows = payload.get("schema") or _to_schema_rows(dataframe) + return DataFrameRecord( + dataframe_id=str(payload["dataframe_id"]), + table_name=str(payload.get("table_name") or f"df_{payload['dataframe_id']}"), + created_at=str(payload.get("created_at") or datetime.now(UTC).isoformat()), + origin_manager=str(payload.get("origin_manager") or ""), + origin_action=str(payload.get("origin_action") or ""), + rows=int(payload.get("rows") if payload.get("rows") is not None else dataframe.height), + columns=int(payload.get("columns") if payload.get("columns") is not None else dataframe.width), + schema=list(schema_rows), + schema_hash=str(payload.get("schema_hash") or _schema_hash(schema_rows)), + json_size_chars=int(payload.get("json_size_chars") or 0), + dataframe=dataframe, + session_id=str(payload.get("session_id") or ""), + owner_id=str(payload.get("owner_id") or ""), + ) + + +async def _load_working_set( + session_storage: SessionStoragePort, + scope: SessionScope, +) -> _SessionWorkingSet: + """Always reload dataframes from SessionStoragePort (no process-lifetime cache).""" + working_set = _SessionWorkingSet() + partition = await session_storage.get_partition(scope) + if not partition: + return working_set + for raw in partition.dataframes.values(): + if not isinstance(raw, dict) or "dataframe_id" not in raw: + continue + record = _deserialize_record(raw) + working_set.sql_context.register(record.table_name, record.dataframe) + working_set.dataframes[record.dataframe_id] = record + return working_set + + +async def _commit_working_set( + session_storage: SessionStoragePort, + scope: SessionScope, + working_set: _SessionWorkingSet, + snapshot_ids: Optional[Set[str]] = None, +) -> None: + """ + Commit the dataframes map for this SessionScope. + + Atomicity: one put_partition of the full dataframes map. + Consistency: partial payload (dataframes only) so tasks/metadata/files stay. + Isolation: caller holds the in-process session lock. Before PUT, re-read + and union keys added by other workers; drop ids this operation removed. + Durability: delegated to SessionStoragePort. + + Same-key concurrent writes and the GET/PUT race can still last-write-win; + closing that window requires Storage CAS/etag (not implemented here). + """ + current_ids = set(working_set.dataframes.keys()) + removed_ids = (snapshot_ids - current_ids) if snapshot_ids is not None else set() + latest = await _load_working_set(session_storage, scope) + merged: Dict[str, "DataFrameRecord"] = dict(latest.dataframes) + for removed_id in removed_ids: + merged.pop(removed_id, None) + merged.update(working_set.dataframes) + working_set.dataframes = merged + dataframes = { + dataframe_id: _serialize_record(record) + for dataframe_id, record in merged.items() + } + await session_storage.put_partition( + scope, SessionPartitionPayload(dataframes=dataframes), + ) + + +@asynccontextmanager +async def _locked_working_set( + session_storage: SessionStoragePort, + scope: SessionScope, +) -> AsyncIterator[_SessionWorkingSet]: + lock = await _lock_for(scope) + async with lock: + yield await _load_working_set(session_storage, scope) + + +def _unregister_table(working_set: _SessionWorkingSet, table_name: str) -> None: + try: + working_set.sql_context.unregister(table_name) + except Exception: + logger.warning("Failed to unregister SQL table %s", table_name, exc_info=True) + + +def stored_as_dataframe_payload(metadata: Dict[str, Any]) -> Dict[str, Any]: + return { + "stored_as_dataframe": True, + "dataframe_id": metadata["dataframe_id"], + "table_name": metadata["table_name"], + "rows": metadata["rows"], + "columns": metadata["columns"], + "schema_hash": metadata["schema_hash"], + "json_size_chars": metadata["json_size_chars"], + } + + +def is_dataframe_reference(base_result: Any) -> bool: + """True when the result is a stored_as_dataframe_payload (a session dataframe reference).""" + return ( + isinstance(base_result, BaseResult) + and isinstance(base_result.result, list) + and len(base_result.result) == 1 + and isinstance(base_result.result[0], dict) + and base_result.result[0].get("stored_as_dataframe") is True + and bool(base_result.result[0].get("dataframe_id")) + ) + +_DISALLOWED_SQL_PATTERN = re.compile( + r"\b(insert|update|delete|create|drop|alter|truncate|replace|merge|call|copy|grant|revoke)\b", + re.IGNORECASE, +) +_LEADING_SQL_COMMENTS_PATTERN = re.compile( + r"^(?:\s*(?:--[^\n]*\n|/\*.*?\*/))*\s*", + re.DOTALL, +) +_SQL_LINE_COMMENT_PATTERN = re.compile(r"--[^\n]*") +_SQL_BLOCK_COMMENT_PATTERN = re.compile(r"/\*.*?\*/", re.DOTALL) +_ORDER_BY_PATTERN = re.compile(r"\border\s+by\b", re.IGNORECASE) +_LIMIT_PATTERN = re.compile(r"\blimit\b", re.IGNORECASE) +_OFFSET_PATTERN = re.compile(r"\boffset\b", re.IGNORECASE) + + +@dataclass +class DataFrameRecord: + dataframe_id: str + table_name: str + created_at: str + origin_manager: str + origin_action: str + rows: int + columns: int + schema: List[Dict[str, str]] + schema_hash: str + json_size_chars: int + dataframe: pl.DataFrame + # Chat session that created it (the partition it lives in); kept in its metadata. + session_id: str = "" + owner_id: str = "" + + def to_metadata(self, include_schema: bool = True) -> Dict[str, Any]: + metadata = asdict(self) + metadata.pop("dataframe", None) + if not include_schema: + metadata.pop("schema", None) + return metadata + + +def _json_default_serializer(value: Any) -> Any: + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if hasattr(value, "isoformat"): + return value.isoformat() + return str(value) + + +def serialize_result_to_compact_json(result: List[Any]) -> str: + return json.dumps(result, separators=(",", ":"), ensure_ascii=False, default=_json_default_serializer) + + +def build_dataframe_from_result(result: List[Any]) -> pl.DataFrame: + normalized = json.loads(serialize_result_to_compact_json(result)) + + # matrix envelope: [{"columns":[...], "rows":[...]}] + if ( + isinstance(normalized, list) + and len(normalized) == 1 + and isinstance(normalized[0], dict) + and set(normalized[0].keys()) == {"columns", "rows"} + and isinstance(normalized[0]["columns"], list) + and isinstance(normalized[0]["rows"], list) + ): + matrix = normalized[0] + return pl.DataFrame(matrix["rows"], schema=[str(c) for c in matrix["columns"]], orient="row") + + # columnar envelope: [{"colA":[...], "colB":[...]}] + if ( + isinstance(normalized, list) + and len(normalized) == 1 + and isinstance(normalized[0], dict) + and normalized[0] + and all(isinstance(v, list) for v in normalized[0].values()) + ): + col_lengths = {len(v) for v in normalized[0].values()} + if len(col_lengths) == 1: + return pl.DataFrame(normalized[0]) + + if isinstance(normalized, list): + if not normalized: + return pl.DataFrame() + if all(isinstance(item, dict) for item in normalized): + return pl.DataFrame(normalized) + return pl.DataFrame({"value": normalized}) + if isinstance(normalized, dict): + return pl.DataFrame([normalized]) + return pl.DataFrame({"value": [normalized]}) + + +def auto_flatten_wide( + df: pl.DataFrame, + max_passes: int = 30, + sep: str = "__", +) -> pl.DataFrame: + """ + Flatten nested structures in a DataFrame for SQL queryability. + + - Nested structs: expanded into flat columns with path-style names + (e.g. configuration__threads, config__inner__b). Only flattens down to leaf scalars. + - List columns: flattened to scalar (first element). List of structs becomes + the struct fields of the first element with path prefix; list of scalars + becomes the first scalar. + - Preserves original row count. + - Safe for schemas with configuration, override_executions, and similar nested structures. + """ + for _ in range(max_passes): + struct_cols = [c for c, dt in df.schema.items() if isinstance(dt, pl.Struct)] + list_cols = [c for c, dt in df.schema.items() if isinstance(dt, pl.List)] + + if not struct_cols and not list_cols: + break + + # Flatten list columns: take first element, then unnest if struct + for col in list_cols: + inner = getattr(df.schema[col], "inner", None) + is_struct_inner = inner is not None and isinstance(inner, pl.Struct) + + temp = f"{col}{sep}temp" + expr = pl.col(col).fill_null([]).list.first() + df = df.with_columns(expr.alias(temp)) + + if is_struct_inner: + df = df.unnest(temp).drop(col) + # Rename to path format: col__field_name + fields = getattr(inner, "fields", []) + rename_map = {f.name: f"{col}{sep}{f.name}" for f in fields} + df = df.rename(rename_map) + else: + df = df.drop(col).rename({temp: col}) + + # Unnest struct columns one at a time, renaming to path format + struct_cols = [c for c, dt in df.schema.items() if isinstance(dt, pl.Struct)] + for col in struct_cols: + struct_dtype = df.schema[col] + fields = getattr(struct_dtype, "fields", []) + rename_map = {f.name: f"{col}{sep}{f.name}" for f in fields} + df = df.unnest(col) + df = df.rename(rename_map) + + return df + + +def _to_schema_rows(dataframe: pl.DataFrame) -> List[Dict[str, str]]: + schema = dataframe.schema + return [{"name": name, "dtype": str(dtype)} for name, dtype in schema.items()] + + +def _stable_hash(payload: str) -> str: + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _schema_hash(schema_rows: List[Dict[str, str]]) -> str: + payload = json.dumps(schema_rows, separators=(",", ":"), ensure_ascii=False) + return _stable_hash(payload) + + +def _normalize_root_dtype(dtype: str) -> str: + dtype_str = str(dtype or "").strip() + if dtype_str.startswith("Struct("): + return "Struct" + if dtype_str.startswith("List("): + inner = dtype_str[5:-1].strip() if dtype_str.endswith(")") else "" + if inner.startswith("Struct("): + return "List(Struct)" + return "List" + if dtype_str.startswith("Array("): + inner = dtype_str[6:-1].strip() if dtype_str.endswith(")") else "" + if inner.startswith("Struct("): + return "Array(Struct)" + return "Array" + return dtype_str + + +def _canonicalize_top_schema(schema_rows: List[Dict[str, str]]) -> List[Dict[str, str]]: + top_level = [ + {"name": str(row.get("name", "")), "dtype": _normalize_root_dtype(str(row.get("dtype", "")))} + for row in schema_rows + ] + return sorted(top_level, key=lambda col: col["name"]) + + +async def register_dataframe( + result: List[Any], + origin_manager: str, + origin_action: str, + json_size_chars: int, + session_storage: SessionStoragePort, + scope: SessionScope, + flatten: bool = True, +) -> Dict[str, Any]: + dataframe = build_dataframe_from_result(result) + return await _register_dataframe_instance( + dataframe, + origin_manager, + origin_action, + json_size_chars, + session_storage=session_storage, + scope=scope, + flatten=flatten, + ) + + +def _allocate_dataframe_id(working_set: _SessionWorkingSet) -> str: + for _ in range(DATAFRAME_ID_MAX_ATTEMPTS): + candidate = generate_simple_id() + if candidate not in working_set.dataframes: + return candidate + + logger.error( + "Unable to allocate dataframe id. attempts=%s id_length=%s active_pool_size=%s", + DATAFRAME_ID_MAX_ATTEMPTS, + SIMPLE_ID_LENGTH, + len(working_set.dataframes), + ) + raise RuntimeError(f"Unable to allocate dataframe id after {DATAFRAME_ID_MAX_ATTEMPTS} attempts.") + + +async def _register_dataframe_instance( + dataframe: pl.DataFrame, + origin_manager: str, + origin_action: str, + json_size_chars: int, + session_storage: SessionStoragePort, + scope: SessionScope, + flatten: bool = True, +) -> Dict[str, Any]: + if flatten: + try: + dataframe = auto_flatten_wide(dataframe) + except Exception: + logger.warning( + "Dataframe flatten failed; storing original shape. origin=%s action=%s", + origin_manager, + origin_action, + exc_info=True, + ) + async with _locked_working_set(session_storage, scope) as working_set: + snapshot_ids = set(working_set.dataframes.keys()) + dataframe_id = _allocate_dataframe_id(working_set) + table_name = f"df_{dataframe_id}" + record = DataFrameRecord( + dataframe_id=dataframe_id, + table_name=table_name, + created_at=datetime.now(UTC).isoformat(), + origin_manager=origin_manager, + origin_action=origin_action, + rows=dataframe.height, + columns=dataframe.width, + schema=(schema_rows := _to_schema_rows(dataframe)), + schema_hash=_schema_hash(schema_rows), + json_size_chars=json_size_chars, + dataframe=dataframe, + session_id=scope.mcp_session_id, + owner_id=scope.user_id, + ) + working_set.sql_context.register(table_name, dataframe) + working_set.dataframes[dataframe_id] = record + await _commit_working_set( + session_storage, scope, working_set, snapshot_ids=snapshot_ids, + ) + return record.to_metadata() + + +async def materialize_large_result_if_needed( + base_result: BaseResult, + origin_manager: str, + origin_action: str, + session_storage: SessionStoragePort, + scope: SessionScope, + force: bool = False, +) -> BaseResult: + if not isinstance(base_result, BaseResult) or base_result.error or base_result.result is None: + return base_result + if is_dataframe_reference(base_result): + # Avoid rematerializing a payload that is already a dataframe reference. + return base_result + rows = base_result.result + page_note = None + page_items = getattr(rows, "items", None) + if isinstance(page_items, list) and hasattr(rows, "model_dump"): + # Pagination envelopes (report pages) store the row list, not one nested object. + try: + page_note = rows.model_dump(exclude={"items"}, exclude_none=True) + except Exception: + page_note = None + rows = page_items + + try: + compact_json = serialize_result_to_compact_json(rows) + json_size_chars = len(compact_json) + except Exception as exc: + base_result.append_warnings( + [f"Result size check failed, skipping dataframe materialization: {exc}"] + ) + return base_result + + if not force and json_size_chars <= DATAFRAME_JSON_SIZE_THRESHOLD: + return base_result + + try: + dataframe_preview = build_dataframe_from_result(rows) + except Exception as exc: + base_result.append_warnings( + [f"Result dataframe preview failed, skipping dataframe materialization: {exc}"] + ) + return base_result + + if dataframe_preview.height == 0: + base_result.append_info([ + "Result contains no rows; dataframe was not created." + ]) + return base_result + + metadata = await _register_dataframe_instance( + dataframe=dataframe_preview, + origin_manager=origin_manager, + origin_action=origin_action, + json_size_chars=json_size_chars, + session_storage=session_storage, + scope=scope, + ) + base_result.result = [stored_as_dataframe_payload(metadata)] + notes = [ + "Large result was stored as a session dataframe. Use perfecto_tools with action " + "'dataframes_list'/'dataframes_get' to inspect metadata and 'dataframes_query' to read data with SQL.", + "ORDER BY + LIMIT + OFFSET are mandatory in every dataframe query.", + "Use a prudent default page size of up to 100 rows (for example, LIMIT 100 OFFSET 0), then continue paging as needed.", + "When the dataframe is no longer needed, free resources with 'dataframes_remove' or 'dataframes_clear'.", + ] + if page_note: + notes.append(f"Page metadata from the original result: {page_note}.") + base_result.append_info(notes) + return base_result + + +def normalize_result_format(value: Any) -> str: + result_format = str(value or "auto").strip().lower() + if result_format not in ALLOWED_RESULT_FORMATS: + return "invalid" + return result_format + + +def _extract_result_format(action: Any, args_dict: Any) -> str: + if not isinstance(args_dict, dict): + return "auto" + return normalize_result_format(args_dict.get("result_format", "auto")) + + +async def finalize_tool_result( + result: Any, + *, + action: Any, + args: Any, + origin_manager: str, + session_storage: Optional[SessionStoragePort] = None, + scope_resolver: Optional[SessionScopeResolverPort] = None, + scope: Optional[SessionScope] = None, + excluded_actions: Optional[set[str]] = None, +) -> Any: + """ + Post-process a tool BaseResult: honor result_format and materialize large payloads. + + excluded_actions: skip auto materialization (still honors result_format=dataframe). + SessionStoragePort is required once this path would persist; missing session storage fails closed. + Pass ``scope`` to skip resolving the current session (used by the async task runner). + """ + if not isinstance(result, BaseResult) or result.error or result.result is None: + return result + + excluded = excluded_actions or set() + result_format = _extract_result_format(action, args) + if result_format == "invalid": + return BaseResult(error=INVALID_RESULT_FORMAT_ERROR) + if isinstance(action, str) and action == "batch": + result_format = "raw" + if result_format == "raw": + return result + if result_format == "auto" and isinstance(action, str) and action in excluded: + return result + if session_storage is None: + return BaseResult(error=MISSING_STORAGE_ERROR) + + if scope is None: + try: + scope = resolve_session_scope(scope_resolver=scope_resolver) + except SessionContextMissing: + # Public tools running without a session keep results inline. + return result + try: + return await materialize_large_result_if_needed( + base_result=result, + origin_manager=origin_manager, + origin_action=str(action) if action is not None else "unknown", + session_storage=session_storage, + scope=scope, + force=(result_format == "dataframe"), + ) + except Exception as exc: + return BaseResult( + error=( + f"Large result materialization failed: {exc}. " + "Try reducing the scope or filters and retry." + ) + ) + + +async def list_dataframes_metadata( + session_storage: SessionStoragePort, + scope: SessionScope, + include_schema: bool = False, +) -> List[Dict[str, Any]]: + async with _locked_working_set(session_storage, scope) as working_set: + return [ + record.to_metadata(include_schema=include_schema) + for record in working_set.dataframes.values() + ] + + +async def get_dataframe_metadata( + dataframe_id: str, + session_storage: SessionStoragePort, + scope: SessionScope, + include_schema: bool = True, +) -> Optional[Dict[str, Any]]: + async with _locked_working_set(session_storage, scope) as working_set: + record = working_set.dataframes.get(dataframe_id) + if not record: + return None + return record.to_metadata(include_schema=include_schema) + + +async def group_dataframe_schemas( + session_storage: SessionStoragePort, + scope: SessionScope, + dataframe_id_list: Optional[List[str]] = None, +) -> Dict[str, Any]: + async with _locked_working_set(session_storage, scope) as working_set: + if dataframe_id_list: + requested = [str(df_id) for df_id in dataframe_id_list] + selected = [ + record for df_id in requested + if (record := working_set.dataframes.get(df_id)) + ] + missing = [df_id for df_id in requested if df_id not in working_set.dataframes] + else: + selected = list(working_set.dataframes.values()) + missing = [] + top_groups: Dict[str, Dict[str, Any]] = {} + for record in selected: + top_schema = _canonicalize_top_schema(record.schema) + top_signature = json.dumps(top_schema, separators=(",", ":"), ensure_ascii=False) + top_hash = _stable_hash(top_signature) + group = top_groups.setdefault( + top_hash, + { + "dataframes": [], + "_columns": {}, + }, + ) + group["dataframes"].append( + record.dataframe_id + ) + + schema_by_name = {str(col.get("name", "")): str(col.get("dtype", "")) for col in record.schema} + for top_col in top_schema: + column_name = top_col["name"] + full_dtype = schema_by_name.get(column_name, "__MISSING__") + schema_preview = full_dtype + version_signature = json.dumps({"dtype": full_dtype}, separators=(",", ":"), ensure_ascii=False) + column_hash = _stable_hash(version_signature) + + column_group = group["_columns"].setdefault( + column_name, + { + "name": column_name, + "_versions": {}, + }, + ) + version_group = column_group["_versions"].setdefault( + column_hash, + { + "hash": column_hash, + "column_schema": schema_preview, + "dataframes": [], + }, + ) + version_group["dataframes"].append( + record.dataframe_id + ) + + top_level_groups = [] + df_sets: Dict[str, str] = {} + dataframe_set_index: Dict[tuple[str, ...], str] = {} + next_df_set_id = 1 + + def _register_dataframe_set(ids: List[str]) -> str: + nonlocal next_df_set_id + normalized = tuple(sorted(set(ids))) + if normalized in dataframe_set_index: + return dataframe_set_index[normalized] + set_id = str(next_df_set_id) + next_df_set_id += 1 + dataframe_set_index[normalized] = set_id + df_sets[set_id] = ",".join(normalized) + return set_id + + for top_hash in sorted(top_groups.keys()): + group = top_groups[top_hash] + columns = [] + varying_columns: List[str] = [] + for column_name in sorted(group["_columns"].keys()): + column_group = group["_columns"][column_name] + variations = [] + for column_hash in sorted(column_group["_versions"].keys()): + version = dict(column_group["_versions"][column_hash]) + version["df_ref"] = _register_dataframe_set(version.pop("dataframes")) + version.pop("hash", None) + variations.append(version) + if len(variations) > 1: + varying_columns.append(column_group["name"]) + columns.append( + { + "name": column_group["name"], + "variations": variations, + } + ) + top_level_groups.append( + { + "df_ref": _register_dataframe_set(group["dataframes"]), + "varying_columns": ",".join(varying_columns), + "columns": columns, + } + ) + + return { + "groups": top_level_groups, + "df_sets": df_sets, + "missing_df_ids": ",".join(missing), + } + + +def _sanitize_sql_for_keyword_scan(sql: str) -> str: + without_comments = _SQL_BLOCK_COMMENT_PATTERN.sub(" ", _SQL_LINE_COMMENT_PATTERN.sub(" ", sql)) + # Remove string and quoted identifier contents to avoid false positives, e.g. SELECT 'delete' + sanitized = re.sub(r"'(?:''|[^'])*'", "''", without_comments) + sanitized = re.sub(r'"(?:""|[^"])*"', '""', sanitized) + sanitized = re.sub(r"`(?:``|[^`])*`", "``", sanitized) + return sanitized + + +def _validate_sql_read_only(sql: str) -> Optional[str]: + without_comments = _LEADING_SQL_COMMENTS_PATTERN.sub("", sql or "") + lowered = without_comments.strip().lower() + if not lowered: + return "SQL query is empty. Provide a SELECT query." + if not (lowered.startswith("select") or lowered.startswith("with")): + return ( + "Only read-only SQL is allowed. Start queries with SELECT or WITH (CTE + SELECT). " + "ORDER BY + LIMIT + OFFSET are mandatory for all dataframe queries." + ) + sanitized = _sanitize_sql_for_keyword_scan(lowered) + if not _ORDER_BY_PATTERN.search(sanitized): + return ( + "Deterministic pagination required. ORDER BY is mandatory for dataframe queries. " + "Use ORDER BY + LIMIT + OFFSET in every query, for example: " + "SELECT * FROM df_x ORDER BY created_at DESC LIMIT 100 OFFSET 0." + ) + if not _LIMIT_PATTERN.search(sanitized): + return ( + "Deterministic pagination required. LIMIT is mandatory for dataframe queries. " + "Use ORDER BY + LIMIT + OFFSET in every query, for example: " + "SELECT * FROM df_x ORDER BY created_at DESC LIMIT 100 OFFSET 0." + ) + if not _OFFSET_PATTERN.search(sanitized): + return ( + "Deterministic pagination required. OFFSET is mandatory for dataframe queries. " + "Use ORDER BY + LIMIT + OFFSET in every query, for example: " + "SELECT * FROM df_x ORDER BY created_at DESC LIMIT 100 OFFSET 0." + ) + disallowed = _DISALLOWED_SQL_PATTERN.search(sanitized) + if disallowed: + return ( + f"SQL statement '{disallowed.group(1).upper()}' is not allowed in dataframe queries. " + "Allowed entry points are SELECT and WITH." + ) + return None + + +async def query_dataframes( + sql: str, + session_storage: SessionStoragePort, + scope: SessionScope, + output_format: str = "matrix", +) -> Dict[str, Any]: + sql_error = _validate_sql_read_only(sql) + if sql_error: + return {"error": sql_error} + normalized_output_format = str(output_format or "matrix").strip().lower() + if normalized_output_format not in {"matrix", "columnar", "records"}: + return {"error": "Invalid output_format. Allowed values: matrix, columnar, records."} + async with _locked_working_set(session_storage, scope) as working_set: + return _execute_sql_on_working_set(working_set, sql, normalized_output_format) + + +def _execute_sql_on_working_set( + working_set: _SessionWorkingSet, + sql: str, + normalized_output_format: str, +) -> Dict[str, Any]: + try: + query_result = working_set.sql_context.execute(sql) + if hasattr(query_result, "collect"): + query_result = query_result.collect() + if not isinstance(query_result, pl.DataFrame): + query_result = pl.DataFrame(query_result) + if normalized_output_format == "columnar": + result_payload = [query_result.to_dict(as_series=False)] + elif normalized_output_format == "records": + result_payload = query_result.to_dicts() + else: + result_payload = [{ + "columns": query_result.columns, + "rows": [list(row) for row in query_result.rows()], + }] + return { + "result": result_payload, + "rows": query_result.height, + "columns": query_result.width, + "schema": _to_schema_rows(query_result), + "output_format": normalized_output_format, + } + except Exception as exc: + error_text = str(exc) + lowered = error_text.lower() + if "not found" in lowered and ("table" in lowered or "relation" in lowered): + guidance = ( + "Table not found in SQL context. Use dataframes_list to discover available table_name values, " + "then retry your query." + ) + elif "column" in lowered and "not found" in lowered: + guidance = ( + "Column not found. Use dataframes_get to inspect schema and exact column names before querying." + ) + elif "syntax" in lowered or "parse" in lowered: + guidance = ( + "SQL syntax error. Use dataframes_sql_help for allowed SQL operations and examples." + ) + else: + guidance = ( + "Use dataframes_sql_help first for supported SQL semantics. " + "For multi-dataframe queries, run dataframes_schema_groups before broad schema inspection. " + "Use dataframes_get selectively for outliers or ambiguous fields." + ) + return { + "error": ( + f"SQL query failed: {exc}. {guidance}" + ) + } + + +def _remove_ids_from_working_set( + working_set: _SessionWorkingSet, + dataframe_ids: List[str], +) -> tuple[List[str], List[str]]: + unique_ids = list(dict.fromkeys(dataframe_ids)) + removed: List[str] = [] + missing: List[str] = [] + for dataframe_id in unique_ids: + record = working_set.dataframes.pop(dataframe_id, None) + if not record: + missing.append(dataframe_id) + continue + _unregister_table(working_set, record.table_name) + removed.append(dataframe_id) + return removed, missing + + +async def remove_dataframes( + dataframe_ids: List[str], + session_storage: SessionStoragePort, + scope: SessionScope, +) -> Dict[str, List[str]]: + unique_ids = list(dict.fromkeys(dataframe_ids)) + async with _locked_working_set(session_storage, scope) as working_set: + snapshot_ids = set(working_set.dataframes.keys()) + removed, missing = _remove_ids_from_working_set(working_set, unique_ids) + if removed: + await _commit_working_set( + session_storage, scope, working_set, snapshot_ids=snapshot_ids, + ) + return {"removed": removed, "missing": missing} + + +async def remove_dataframe( + dataframe_id: str, + session_storage: SessionStoragePort, + scope: SessionScope, +) -> bool: + result = await remove_dataframes([dataframe_id], session_storage, scope) + return bool(result["removed"]) + + +async def clear_dataframes( + session_storage: SessionStoragePort, + scope: SessionScope, +) -> int: + async with _locked_working_set(session_storage, scope) as working_set: + snapshot_ids = set(working_set.dataframes.keys()) + removed, _ = _remove_ids_from_working_set( + working_set, list(working_set.dataframes.keys()), + ) + if removed: + await _commit_working_set( + session_storage, scope, working_set, snapshot_ids=snapshot_ids, + ) + return len(removed) + + +def get_sql_capabilities() -> Dict[str, Any]: + return { + "engine_scope": { + "query_entrypoints": ["SELECT", "WITH"], + "mode": "read-only", + "description": "SQL support is available for SELECT/WITH queries with BlazeMeter MCP query constraints.", + }, + "allowed_entrypoints": ["SELECT", "WITH"], + "disallowed_statements": [ + "INSERT", + "UPDATE", + "DELETE", + "CREATE", + "DROP", + "ALTER", + "TRUNCATE", + "REPLACE", + "MERGE", + "CALL", + "COPY", + "GRANT", + "REVOKE", + ], + "allowed_features": [ + "JOIN", + "UNION", + "UNION ALL", + "CTE (WITH)", + "GROUP BY", + "HAVING", + "ORDER BY", + "LIMIT", + "OFFSET", + "aggregations", + "UNNEST", + ], + "supported_functions": [ + "ABS", "ACOS", "ACOSD", "ARRAY_CONTAINS", "ARRAY_GET", "ARRAY_LENGTH", "ARRAY_LOWER", "ARRAY_MEAN", + "ARRAY_REVERSE", "ARRAY_SUM", "ARRAY_TO_STRING", "ARRAY_UNIQUE", "ARRAY_UPPER", "ASIN", "ASIND", + "ATAN", "ATAN2", "ATAN2D", "ATAND", "AVG", "BIT_LENGTH", "CBRT", "CEIL", "COALESCE", "CONCAT", + "CONCAT_WS", "COS", "COSD", "COT", "COTD", "COUNT", "DATE", "DATE_PART", "DEGREES", "ENDS_WITH", + "EXP", "EXTRACT", "FIRST", "FLOOR", "GREATEST", "IF", "IFNULL", "INITCAP", "LAST", "LEAST", "LEFT", + "LENGTH", "LN", "LOG", "LOG1P", "LOG10", "LOG2", "LOWER", "LTRIM", "MAX", "MEDIAN", "MIN", "MOD", + "NULLIF", "OCTET_LENGTH", "PI", "POW", "RADIANS", "REGEXP_LIKE", "REPLACE", "REVERSE", "RIGHT", + "ROUND", "RTRIM", "SIGN", "SIN", "SIND", "SQRT", "STARTS_WITH", "STDDEV", "STRPOS", "SUBSTRING", + "SUM", "TAN", "TAND", "UNNEST", "UPPER", "VARIANCE" + ], + "unsupported_functions": [ + "GENERATE_SERIES", + "STRING_AGG", + "GROUP_CONCAT", + "LISTAGG", + "PERCENTILE_CONT", + "PERCENTILE_DISC", + "NTILE", + "CUME_DIST", + "PERCENT_RANK", + "WIDTH_BUCKET", + "JSON_EXTRACT", + "JSON_EXTRACT_PATH", + "JSON_EXTRACT_STRING", + "TO_JSON", + "TYPEOF", + "STRUCT_EXTRACT", + "MONTHS_BETWEEN", + "MODE", + "ROLLUP", + "CUBE", + "GROUPING SETS", + ], + "limited_or_unstable_functions": [ + {"name": "LAG", + "reason": "Unstable, especially with complex CTEs or multiple windows. Avoid when possible."}, + {"name": "LEAD", + "reason": "Unstable, especially with complex CTEs or multiple windows. Avoid when possible."}, + {"name": "DATE_TRUNC", + "reason": "Partial and inconsistent support. Better to use DATE_PART combined with CAST or manual date arithmetic."}, + {"name": "RANK", + "reason": "Partial support through window functions. Results may differ from PostgreSQL/BigQuery."}, + {"name": "DENSE_RANK", + "reason": "Partial support through window functions. Results may differ from PostgreSQL/BigQuery."}, + {"name": "ROW_NUMBER", + "reason": "Works in simple cases but can be unstable with complex queries or multiple CTEs."}, + {"name": "FIRST_VALUE", "reason": "Limited support as window function."}, + {"name": "LAST_VALUE", "reason": "Limited support as window function."}, + {"name": "NTH_VALUE", "reason": "Very limited and unstable support."} + ], + "unsupported_or_unstable_patterns": [ + "Complex chained nested access with mixed subscript and dot notation in a single expression", + "Casting LIST/STRUCT directly to STRING for inspection", + "Nested extraction without staged CTE when list expansion is required", + "Unqualified join keys that create ambiguous column references", + ], + "ai_common_mistakes": [ + "Assuming generic warehouse helper functions are available", + "Building one very large query instead of staged CTEs", + "Skipping aliases in JOIN/CTE steps", + "Trying direct list aggregations (for example list_max on nested overrides) instead of UNNEST + staged CTE", + "Trying unsupported helper functions before checking supported_functions", + "Assuming nested/scalar fields are homogeneous across dataframes without checking schemas first", + "Inspecting every dataframe with dataframes_get before checking grouped schema differences", + "Using direct nested extraction in the first multi-dataframe query after schema groups reports column variations", + "Assuming single dataframe justifies bypassing the robust UNNEST/CTE pattern for nested/list fields", + "Trying the 'fast' direct nested access first when the query touches nested/list fields", + "Try-fast: attempting the simplest path first and retrying on failure instead of reasoning through the design before executing", + "Not considering all values in a nested list when searching for max/min, which can miss important extreme values", + "Using only the first element of a nested list instead of aggregating over all its values", + "Using ANSI date literals (DATE '2026-03-30') inside VALUES clause", + ], + "query_rules": [ + "CRITICAL: Before writing queries that combine 2 or more dataframes, run dataframes_schema_groups first to validate schema compatibility across all involved dataframes.", + "CRITICAL: Use dataframes_get only for targeted drill-down on dataframes flagged by schema groups as different or ambiguous for required fields.", + "CRITICAL: Hard gate: if schema groups reports column variations, direct nested extraction is forbidden in the first query.", + "CRITICAL: If the query touches nested/list fields, direct nested access is forbidden. Always use the robust pattern: UNNEST -> aggregate -> join-back in CTEs. No exception for single dataframe.", + "IMPORTANT: Validate schema compatibility before using nested fields.", + "ORDER BY + LIMIT + OFFSET are mandatory in every dataframe query.", + "Use deterministic pagination: ORDER BY + LIMIT + OFFSET.", + "Recommended default page size: LIMIT 100 OFFSET 0, then continue paging.", + "If loading data with result_format=dataframe, prefer one initial fetch with the maximum allowed tool limit, then paginate/filter in dataframes_query.", + "CRITICAL: When a query includes UNNEST + CTE + JOIN, always enforce explicit join-key renaming and qualification. Rename the base key in the UNNEST CTE (e.g. test_id AS base_test_id) and use only that renamed key downstream.", + "For CTE-heavy joins, rename join keys in intermediate CTEs (for example test_id AS t_id or base_test_id).", + "Single dataframe query flow (scalar-only): dataframes_sql_help -> dataframes_get -> dataframes_query.", + "Single dataframe query flow (nested/list fields): dataframes_sql_help -> dataframes_get -> staged CTE (UNNEST -> aggregate -> join-back) -> dataframes_query. Same robust pattern as multi-dataframe.", + "Multi-dataframe nested flow: dataframes_sql_help -> dataframes_schema_groups -> targeted dataframes_get -> staged CTE (UNNEST -> aggregate -> join-back) -> final query.", + "If schema groups returns a CRITICAL variation warning, call dataframes_sql_help again immediately before writing the final query.", + "Direct nested access is allowed only when each required nested column has exactly one variation across all relevant dataframes in schema groups.", + "For date literals in VALUES → always use: CAST('YYYY-MM-DD' AS DATE)", + "DATE('YYYY-MM-DD') is also supported and often cleaner", + "Never use: DATE '2026-03-30' inside VALUES", + ], + "nested_unnest_intro": ( + "To query and aggregate data from a list of structs (e.g., override_executions), use UNNEST in a CTE to flatten the list, " + "then aggregate and compare with scalar fields using GREATEST/LEAST. See query_examples.good for the compact pattern." + ), + "nested_list_pre_sql_checklist": [ + "Step 1: Identify if the query touches nested/list (List, Struct, Array in schema). Step 2: If yes, confirm robust pattern. Step 3: Design the CTE structure. Step 4: Execute. Do not skip to Step 4.", + "Before launching SQL that touches nested/list fields, explicitly confirm: 'There are nested/list fields; I use the robust UNNEST -> aggregate -> join-back pattern.'", + "Do not attempt the 'fast' direct nested extraction first. Start with the robust CTE pattern.", + "Single dataframe is NOT an exception: use the same robust pattern when querying nested/list columns.", + "Anti-ambiguity checklist (UNNEST+CTE+JOIN): No unqualified key columns in SELECT, JOIN, GROUP BY, or ORDER BY; UNNEST CTE key is renamed (base_* or src_*); join-back uses different left/right key names; final projection is scalar-only; query ends with ORDER BY ... LIMIT ... OFFSET.", + ], + "pre_execution_reasoning": [ + "Before dataframes_query: reason step-by-step. (1) Schema check: what columns and types? (2) Nested/list? If List, Struct, Array → robust pattern. (3) Pattern selection: scalar-only vs UNNEST/CTE. (4) Design the query structure. (5) Confirm, then execute.", + "Do not try-fast. Design before do.", + ], + "recommended_patterns": [ + "Prefer one final aggregation query over multiple partial queries when feasible", + "Build queries incrementally: base SELECT -> UNNEST CTE -> aggregate -> join -> final sort/page", + "Use a dedicated CTE for UNNEST operations on nested arrays/lists", + "If a nested field fails, use the robust pattern: UNNEST -> aggregate -> join-back.", + "First nested-field query must use the robust UNNEST/CTE pattern. Never try direct nested access first. No exception for single dataframe.", + "Alias every table and CTE explicitly", + "Rename join keys in CTEs (for example: t_id, base_test_id, src_test_id) before joins to avoid ambiguous references", + "Join-key hygiene for UNNEST+CTE+JOIN: (1) In UNNEST CTE rename base key immediately (test_id AS base_test_id). (2) In downstream CTEs use only the renamed key (GROUP BY base_test_id). (3) In join-back use fully-qualified names (ON b.test_id = a.base_test_id). (4) In final SELECT prefix columns with table alias. (5) Never reuse generic key names across CTE boundaries.", + "Use COALESCE/CASE for fallback values after joins", + "UNION ALL only scalar projections; avoid UNION over nested struct/list columns", + "Validate each CTE with a small LIMIT before composing final query", + "For multi-dataframe analysis, use schema groups first, then perform targeted per-dataframe inspection only when needed.", + "To get the maximum value between a scalar field and all values in a nested list per record, use UNNEST on the list, then GROUP BY and GREATEST(MAX(list.field), MAX(scalar)).", + "Before UNION ALL, normalize each branch to the same concrete type your next step expects (e.g. INTEGER year, not “string then parse after union”).", + ], + "known_engine_pitfalls": [ + "CTE + JOIN resolution may treat same-name keys as ambiguous even when aliases are present; rename join keys in the UNNEST stage (base_*/src_*) to guarantee deterministic resolution", + "Alias/join-key resolution may fail in some CTE + JOIN combinations", + "Ambiguous join keys are common if columns are not fully qualified", + "Nested schema drift across tables can break field resolution", + "UNION over nested struct/list columns is fragile; normalize to scalar output first", + "Direct list aggregation over nested overrides is brittle; UNNEST + MAX + join-back is more reliable", + "VALUES clause is very strict: does not accept DATE '2026-03-30' literal. Must use CAST('2026-03-30' AS DATE) or DATE('2026-03-30')", + "Temporal literals inside VALUES frequently cause 'expects literals' errors", + "CAST to DATE/DATETIME is more reliable than the typed literal syntax in Polars SQL" + ], + "nested_query_recipe": [ + "Base table CTE", + "UNNEST CTE", + "Aggregate CTE (for example MAX over nested field)", + "Join aggregate back to base", + "Apply null-safe metric expression (for example GREATEST(COALESCE(default1,0), COALESCE(default2,0), COALESCE(override_max,0)))", + "Emit scalar projection only", + "UNION ALL scalar projections only", + "Final ORDER BY + LIMIT + OFFSET", + ], + "debug_ladder": [ + "A) Run schema groups for all candidate dataframes and identify only the outliers to inspect with dataframes_get.", + "B) base SELECT LIMIT 10", + "C) UNNEST stage LIMIT 10", + "D) aggregate stage LIMIT 10", + "E) join result LIMIT 10", + "F) add ranking and pagination", + "G) add next table to UNION ALL and repeat", + ], + "query_examples": { + "good": [ + "SELECT * FROM df_tests ORDER BY test_id LIMIT 100 OFFSET 0", + "WITH expanded AS (SELECT t.test_id, UNNEST(t.override_executions) AS ov FROM df_tests t), " + "agg AS (SELECT e.test_id, MAX(e.ov.concurrency) AS max_concurrency FROM expanded e GROUP BY e.test_id) " + "SELECT t.test_id, t.test_name, " + "GREATEST(COALESCE(t.configuration.threads, 0), COALESCE(a.max_concurrency, 0)) AS max_concurrency_used " + "FROM df_tests t LEFT JOIN agg a ON t.test_id = a.test_id " + "ORDER BY max_concurrency_used DESC, t.test_id ASC LIMIT 10 OFFSET 0", + "WITH a AS (SELECT test_id, test_name FROM df_a), b AS (SELECT test_id, test_name FROM df_b) " + "SELECT * FROM a UNION ALL SELECT * FROM b ORDER BY test_id LIMIT 100 OFFSET 0", + "WITH s1_exp AS (SELECT t.test_id AS t_id, UNNEST(t.override_executions) AS ov FROM df_a t), " + "s1_agg AS (SELECT e.t_id, MAX(e.ov.concurrency) AS ov_max FROM s1_exp e GROUP BY e.t_id), " + "s1 AS (SELECT t.test_id, t.test_name, GREATEST(COALESCE(t.configuration.threads,0), COALESCE(a.ov_max,0)) " + "AS max_concurrency_used FROM df_a t LEFT JOIN s1_agg a ON t.test_id = a.t_id), " + "s2_exp AS (SELECT t.test_id AS t_id, UNNEST(t.override_executions) AS ov FROM df_b t), " + "s2_agg AS (SELECT e.t_id, MAX(e.ov.concurrency) AS ov_max FROM s2_exp e GROUP BY e.t_id), " + "s2 AS (SELECT t.test_id, t.test_name, GREATEST(COALESCE(t.configuration.threads,0), COALESCE(a.ov_max,0)) " + "AS max_concurrency_used FROM df_b t LEFT JOIN s2_agg a ON t.test_id = a.t_id), " + "all_rows AS (SELECT test_id, test_name, max_concurrency_used FROM s1 UNION ALL " + "SELECT test_id, test_name, max_concurrency_used FROM s2) " + "SELECT test_name, test_id, max_concurrency_used FROM all_rows " + "ORDER BY max_concurrency_used DESC, test_id ASC LIMIT 10 OFFSET 0", + "WITH expanded AS (SELECT t.test_id, t.test_name, t.configuration.threads AS threads, UNNEST(t.override_executions) AS ov FROM df_tests t), " + "agg AS (SELECT test_id, test_name, GREATEST(COALESCE(MAX(ov.concurrency), 0), COALESCE(MAX(threads), 0)) AS max_concurrency FROM expanded GROUP BY test_id, test_name) " + "SELECT test_id, test_name, max_concurrency FROM agg ORDER BY max_concurrency DESC, test_id ASC LIMIT 10 OFFSET 0", + ], + "bad": [ + "SELECT * FROM df_tests", + "SELECT * FROM df_tests WHERE status = 'ERROR'", + "SELECT test_id, MAX(UNNEST(override_executions).concurrency) FROM df_tests GROUP BY test_id", + "SELECT * FROM df_a JOIN df_b ON test_id = test_id ORDER BY test_id LIMIT 100 OFFSET 0", + "SELECT * FROM df_a UNION ALL SELECT * FROM df_b ORDER BY test_id LIMIT 100 OFFSET 0", + "SELECT TO_JSON(configuration) FROM df_tests ORDER BY test_id LIMIT 10 OFFSET 0", + ], + }, + "troubleshooting_hints": [ + "If you get ambiguous column errors, alias every table/CTE and qualify join keys.", + "If nested field access fails, move expansion into a dedicated CTE using UNNEST.", + "If a query is too complex, split it into 2-4 CTE stages and validate each stage independently.", + "If an inferred function fails, verify against supported_functions and unsupported_functions.", + "If aggregation results over nested lists do not reflect expected values, check that you are using UNNEST and aggregation (MAX, MIN, etc.) correctly, and that you compare against the scalar field with GREATEST/LEAST.", + ], + "notes": [ + "All session dataframe tables (loaded from Storage for the current partition) are queryable in the same SQL context.", + "Function name typo seen in some sources: STRPOST; use STRPOS.", + "This help defines practical usage constraints for BlazeMeter MCP SQL queries.", + ], + "references": [ + "https://docs.pola.rs/py-polars/html/reference/sql/index.html", + "https://docs.pola.rs/py-polars/html/reference/sql/functions/index.html", + "https://docs.pola.rs/py-polars/html/reference/sql/clauses.html", + "https://docs.pola.rs/py-polars/html/reference/sql/table_operations.html", + "https://docs.pola.rs/py-polars/html/reference/sql/set_operations.html" + ], + } diff --git a/tools/device_manager.py b/tools/device_manager.py index b499f5b..d49838f 100644 --- a/tools/device_manager.py +++ b/tools/device_manager.py @@ -12,8 +12,10 @@ from formatters.grid import format_grid_info from models.manager import Manager from models.result import BaseResult -from telemetry import run_tool -from tools.utils import api_request, format_sanitized_traceback, normalize_action_args +from config.cache import CacheScope +from tools.method_cache import ttl_cache_method +from tools.runtime_tools import run_tool_with_runtime +from tools.utils import api_request, format_sanitized_traceback, normalize_action_args, run_as_task class DeviceManager(Manager): @@ -21,6 +23,8 @@ def __init__(self, ctx: Context): super().__init__(ctx) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def read_selenium_grid_info(self) -> BaseResult: tenant_url = perfecto.get_tenant_management_api_url(self.token.cloud_name) tenant_response = await api_request(self.token, "GET", endpoint=tenant_url, result_formatter=format_grid_info) @@ -34,6 +38,8 @@ async def read_selenium_grid_info(self) -> BaseResult: return tenant_response @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def list_real_devices(self) -> BaseResult: devices_url = perfecto.get_real_device_management_api_url(self.token.cloud_name) # List all devices @@ -45,17 +51,23 @@ async def list_real_devices(self) -> BaseResult: result_formatter=format_real_device) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def read_real_device_info(self, device_id: str) -> BaseResult: devices_url = perfecto.get_real_device_management_api_url(self.token.cloud_name) devices_url = f"{devices_url}/{device_id}" return await api_request(self.token, "GET", endpoint=devices_url) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def list_virtual_devices(self) -> BaseResult: virtual_device_url = perfecto.get_virtual_device_management_api_url(self.token.cloud_name) return await api_request(self.token, "GET", endpoint=virtual_device_url, result_formatter=format_virtual_device) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def list_desktop_devices(self) -> BaseResult: virtual_web_url = perfecto.get_web_desktop_management_api_url(self.token.cloud_name) return await api_request(self.token, "GET", endpoint=virtual_web_url) @@ -102,7 +114,9 @@ async def _dispatch(): ) try: - return await run_tool(f"{TOOLS_PREFIX}_devices", action, ctx, _dispatch) + return await run_tool_with_runtime( + runtime, f"{TOOLS_PREFIX}_devices", action, ctx, _dispatch, tool_args=args, + ) except httpx.HTTPStatusError: return BaseResult( error=f"Error: {format_sanitized_traceback()}" diff --git a/tools/execution_manager.py b/tools/execution_manager.py index f393d3f..e5e6b6c 100644 --- a/tools/execution_manager.py +++ b/tools/execution_manager.py @@ -12,8 +12,10 @@ from formatters.execution import format_executions from models.manager import Manager from models.result import BaseResult, PaginationResult -from telemetry import run_tool -from tools.utils import api_request, format_sanitized_traceback, normalize_action_args +from config.cache import CacheScope +from tools.method_cache import ttl_cache_method +from tools.runtime_tools import run_tool_with_runtime +from tools.utils import api_request, format_sanitized_traceback, normalize_action_args, run_as_task class ExecutionManager(Manager): @@ -53,6 +55,8 @@ def __init__(self, ctx: Context): } @token_verify + @ttl_cache_method(ttl_seconds=15, scope=CacheScope.USER) + @run_as_task() async def list_live_executions(self) -> BaseResult: execution_management_url = perfecto.get_execution_management_api_url(self.token.cloud_name) execution_management_url = execution_management_url + "/search" @@ -77,12 +81,16 @@ async def stop_live_executions(self, execution_id_list: list[str]) -> BaseResult ) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def list_report_names(self) -> BaseResult: report_management_url = perfecto.get_test_execution_name_api_url(self.token.cloud_name) body = {} return await api_request(self.token, "POST", endpoint=report_management_url, json=body) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def list_filter_values(self, filter_names: list[str]) -> BaseResult: metadata_management_url = perfecto.get_test_execution_metadata_api_url(self.token.cloud_name) @@ -112,6 +120,8 @@ async def list_filter_values(self, filter_names: list[str]) -> BaseResult: ) @token_verify + @ttl_cache_method(ttl_seconds=30, scope=CacheScope.USER) + @run_as_task() async def list_report_executions(self, args: dict[str, Any]) -> BaseResult: page_size = 50 page_index = args.get("page_index", 1) @@ -193,6 +203,8 @@ async def list_report_executions(self, args: dict[str, Any]) -> BaseResult: ) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def red_report_execution(self, execution_id: str) -> BaseResult: report_commands_url = perfecto.get_test_execution_commands_api_url(self.token.cloud_name) + "/" @@ -283,7 +295,9 @@ async def _dispatch(): ) try: - return await run_tool(f"{TOOLS_PREFIX}_execution", action, ctx, _dispatch) + return await run_tool_with_runtime( + runtime, f"{TOOLS_PREFIX}_execution", action, ctx, _dispatch, tool_args=args, + ) except httpx.HTTPStatusError: return BaseResult( error=f"Error: {format_sanitized_traceback()}" diff --git a/tools/help_manager.py b/tools/help_manager.py index 68cc429..adfc871 100644 --- a/tools/help_manager.py +++ b/tools/help_manager.py @@ -14,9 +14,11 @@ format_read_real_devices_extended_command_info, format_help_info from models.manager import Manager from models.result import BaseResult -from telemetry import run_tool +from config.cache import CacheScope from tools.help_utils import convert_js_to_py_dict -from tools.utils import http_request, format_sanitized_traceback, normalize_action_args +from tools.method_cache import ttl_cache_method +from tools.runtime_tools import run_tool_with_runtime +from tools.utils import http_request, format_sanitized_traceback, normalize_action_args, run_as_task class HelpManager(Manager): @@ -118,6 +120,8 @@ async def fetch_chunk(chunk_url: str): } HelpManager.help_tree = help_tree + @ttl_cache_method(ttl_seconds=300, scope=CacheScope.GLOBAL) + @run_as_task() async def list_help_categories(self) -> BaseResult: if HelpManager.help_tree is None: await self._load_help_tree() @@ -133,6 +137,8 @@ async def list_help_categories(self) -> BaseResult: info=["A list of subcategories is provided for each category"] ) + @ttl_cache_method(ttl_seconds=300, scope=CacheScope.GLOBAL) + @run_as_task() async def list_help_category_content(self, category_id: str, subcategory_id_list: List[str]) -> BaseResult: if not isinstance(subcategory_id_list, list) or not subcategory_id_list: return BaseResult( @@ -153,6 +159,8 @@ async def list_help_category_content(self, category_id: str, subcategory_id_list result=results, ) + @ttl_cache_method(ttl_seconds=300, scope=CacheScope.GLOBAL) + @run_as_task() async def read_help_info(self, category_id: str, subcategory_id: str, help_id_list: List[str]) -> BaseResult: if not isinstance(help_id_list, list) or not help_id_list: return BaseResult( @@ -212,6 +220,8 @@ async def read_help_info(self, category_id: str, subcategory_id: str, help_id_li }, ) + @ttl_cache_method(ttl_seconds=300, scope=CacheScope.GLOBAL) + @run_as_task() async def list_real_devices_extended_commands(self) -> BaseResult: real_devices_extended_commands_help_url = get_real_devices_extended_commands_help_url() commands_result = await http_request("GET", endpoint=real_devices_extended_commands_help_url, @@ -298,7 +308,9 @@ async def _dispatch(): ) try: - return await run_tool(f"{TOOLS_PREFIX}_help", action, ctx, _dispatch) + return await run_tool_with_runtime( + runtime, f"{TOOLS_PREFIX}_help", action, ctx, _dispatch, tool_args=args, + ) except httpx.HTTPStatusError: return BaseResult( error=f"Error: {format_sanitized_traceback()}" diff --git a/tools/method_cache.py b/tools/method_cache.py new file mode 100644 index 0000000..f7b0a30 --- /dev/null +++ b/tools/method_cache.py @@ -0,0 +1,213 @@ +"""TTL cache for manager methods. Hits skip another Perfecto API call.""" +from __future__ import annotations + +import contextlib +import functools +from typing import Any, Awaitable, Callable, Optional + +from config.cache import CacheScope, StorePlan, build_cache_key, get_cache +from config.session_context import current_identity +from models.result import BaseResult +from tools.call_state import ( + _cache_bypass, + _disable_dataframe_materialization, + _result_format_context, + _task_management_enabled, +) + + +def _accumulate_cache_debug(metric: str, value: int = 1) -> None: + """Observer hook. Debug counters are optional and never fail a call.""" + return None + + +def resolve_cache_user_id(manager: Any) -> Optional[str]: + """ + Verified user id that scopes USER cache entries. + """ + if getattr(manager, "token", None) is None: + return None + identity = current_identity() + user_id = getattr(identity, "user_id", None) + if user_id is None or not str(user_id).strip(): + return None + return str(user_id).strip() + + +def _cache_compact_value(value: Any) -> str: + if isinstance(value, (str, int, float, bool)) or value is None: + return repr(value) + if isinstance(value, dict): + keys = sorted(value.keys(), key=lambda x: str(x)) + return "{" + ",".join(f"{k}:{_cache_compact_value(value[k])}" for k in keys) + "}" + if isinstance(value, (list, tuple, set)): + items = sorted(value, key=repr) if isinstance(value, set) else value + return "[" + ",".join(_cache_compact_value(v) for v in items) + "]" + return repr(value) + + +def _strip_per_call_fields(value: Any) -> Any: + """Drop timing/debug that describe one call, so cache hits do not replay them.""" + if isinstance(value, BaseResult): + value.debug = None + value.tool_call_started_at = None + value.tool_call_finished_at = None + value.tool_call_duration_ms = None + return value + + +_RESULT_FORMATS = ("auto", "dataframe", "raw") + + +def _method_cache_key( + scope: CacheScope, + qualified_name: str, + args: tuple, + kwargs: dict, + user_id: Optional[str], + *, + result_format: str, + materialize: int, + inline: int, +) -> str: + return build_cache_key( + scope, + qualified_name, + f"args={_cache_compact_value(args)}", + f"kwargs={_cache_compact_value(kwargs)}", + f"format={result_format}", + f"materialize={materialize}", + f"inline={inline}", + user_id=user_id, + ) + + +@contextlib.contextmanager +def bypass_method_cache(): + """Cached methods called inside load fresh and store nothing (read-modify-write paths).""" + token = _cache_bypass.set(True) + try: + yield + finally: + _cache_bypass.reset(token) + + +async def invalidate_cached_method(manager: Any, method: Any, *args, **kwargs) -> None: + """ + Drop every cached variant of ``method(*args, **kwargs)`` for the caller, after a write. + + Pass the arguments exactly as callers do (same positional/keyword form): they are part of the key. + """ + scope = method.cache_scope + if scope is CacheScope.USER: + user_id = resolve_cache_user_id(manager) + if user_id is None: + return # nothing was cached for this caller + else: + user_id = None + cache = get_cache() + keys = [ + _method_cache_key( + scope, method.cache_qualified_name, args, kwargs, user_id, + result_format=result_format, materialize=materialize, inline=inline, + ) + for result_format in _RESULT_FORMATS + for materialize in (0, 1) + for inline in (0, 1) + ] + # Best effort, like every cache write: a failed delete never fails the write it follows. + await asyncio.gather(*(cache.delete(key) for key in keys), return_exceptions=True) + + +def ttl_cache_method(ttl_seconds: float = 30, scope: CacheScope = CacheScope.USER): + """ + Async TTL cache for manager instance methods (bursts and internal re-reads). + + - USER scope keys on the user id from the request context; without one the + cache is bypassed, so entries are never shared across users. + - GLOBAL scope is for results that do not depend on the caller. + - Errors are never cached. Concurrent misses share one call. + - Place it above @run_as_task: a hit skips the task entirely. When the call + returns a snapshot of a task that is still running, the snapshot is cached + for this session only and dropped as soon as the task finishes; session + dataframe references are never stored (the dataframe can be removed), only + shared with concurrent calls of the same session. + - Writes read their base state under ``bypass_method_cache()`` and call + ``invalidate_cached_method`` once they succeed. + """ + + def decorator(func: Callable[..., Awaitable[Any]]): + qualified_name = f"{func.__module__}.{func.__qualname__}" + + @functools.wraps(func) + async def wrapper(self, *args, **kwargs): + if _cache_bypass.get(): + return await func(self, *args, **kwargs) + if scope is CacheScope.USER: + user_id = resolve_cache_user_id(self) + if user_id is None: + return await func(self, *args, **kwargs) + else: + user_id = None + + cache = get_cache() + # Inline calls (nested inside a running tool/task, e.g. bridge validations) always + # get the domain result; top-level calls may get a task snapshot or a dataframe + # reference. Separate keys keep one from ever being served to the other. + key = _method_cache_key( + scope, qualified_name, args, kwargs, user_id, + result_format=_result_format_context.get(), + materialize=int(not _disable_dataframe_materialization.get()), + inline=int(_task_management_enabled.get()), + ) + + from config.session_context import SessionContextMissing + from tools.async_task_manager import add_task_terminal_callback, session_scope_from_manager + from tools.dataframe_manager import is_dataframe_reference + + try: + session_scope = session_scope_from_manager(self) + except SessionContextMissing: + # Unit tests and calls outside a bound tool scope run uncached. + return await func(self, *args, **kwargs) + binding = f"{session_scope.user_id}/{session_scope.mcp_session_id}" + + def _store_policy(value: Any) -> Optional[StorePlan]: + if not isinstance(value, BaseResult): + return StorePlan() + if value.error: + return None + if value.pending_task is not None: + task_id, task_scope = value.pending_task + + async def _invalidate() -> None: + # Only our snapshot: another session may have stored a newer value since. + await cache.delete(key, if_tag=task_id) + + async def _drop_when_task_finishes() -> None: + await add_task_terminal_callback(task_id, task_scope, _invalidate) + + return StorePlan( + bound_to=binding, + tag=task_id, + transform=_strip_per_call_fields, + after_store=_drop_when_task_finishes, + ) + if is_dataframe_reference(value): + return StorePlan(bound_to=binding, store=False, transform=_strip_per_call_fields) + return StorePlan(transform=_strip_per_call_fields) + + return await cache.get_or_load( + key, + lambda: func(self, *args, **kwargs), + ttl_seconds, + store_policy=_store_policy, + binding=binding, + observer=_accumulate_cache_debug, + ) + + wrapper.cache_scope = scope + wrapper.cache_qualified_name = qualified_name + return wrapper + + return decorator diff --git a/tools/runtime_tools.py b/tools/runtime_tools.py new file mode 100644 index 0000000..eedcf6f --- /dev/null +++ b/tools/runtime_tools.py @@ -0,0 +1,75 @@ +"""Run a tool action with telemetry, session scope, 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 models.result import BaseResult +from telemetry import run_tool +from tools.call_state import ( + reset_disable_dataframe_materialization, + reset_result_format, + set_disable_dataframe_materialization, + set_result_format, +) +from tools.dataframe_manager import INVALID_RESULT_FORMAT_ERROR, normalize_result_format + + +def _result_format_from_args(tool_args: Any) -> str: + if not isinstance(tool_args, dict): + return "auto" + return normalize_result_format(tool_args.get("result_format", "auto")) + + +async def run_tool_with_runtime( + runtime: AppRuntime, + tool_name: str, + action: str, + ctx: Any, + dispatch: Callable[[], Awaitable[Any]], + *, + tool_args: Any = None, + dataframe_excluded_actions: Optional[set[str]] = None, + disable_dataframe_materialization: bool = False, +) -> Any: + """ + Bind this call's storage scope, 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. + """ + 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) diff --git a/tools/skills_manager.py b/tools/skills_manager.py index 49763d8..d1ab069 100644 --- a/tools/skills_manager.py +++ b/tools/skills_manager.py @@ -26,7 +26,7 @@ from config.runtime import AppRuntime from models.manager import Manager from models.result import BaseResult -from telemetry import run_tool +from tools.runtime_tools import run_tool_with_runtime from tools.skills_utils import list_skills, read_skill_definition, read_skill_file, parse_skill_uri, \ is_skill_uri, list_skill_resources_uri from tools.utils import format_sanitized_traceback, normalize_action_args @@ -254,7 +254,9 @@ async def process_call(call: Dict[str, Any]) -> BaseResult | List[BaseResult]: ) try: - return await run_tool(f"{TOOLS_PREFIX}_skills", action, ctx, _dispatch) + return await run_tool_with_runtime( + runtime, f"{TOOLS_PREFIX}_skills", action, ctx, _dispatch, tool_args=args, + ) except httpx.HTTPStatusError: return BaseResult( error=f"Error: {format_sanitized_traceback()}" diff --git a/tools/task_runtime.py b/tools/task_runtime.py new file mode 100644 index 0000000..5179125 --- /dev/null +++ b/tools/task_runtime.py @@ -0,0 +1,181 @@ +"""Schedule manager methods as session tasks when they run longer than a few seconds.""" +from __future__ import annotations + +import asyncio +import functools +import inspect +import time +from typing import Any, Awaitable, Callable, Dict, Optional + +from models.result import BaseResult +from tools.call_state import ( + _disable_dataframe_materialization, + _result_format_context, + _task_management_enabled, +) + +_result_debug_enabled = False + +def _attach_task_debug(result: BaseResult, task_record: Any): + if not _result_debug_enabled: + return + if not isinstance(result, BaseResult) or task_record is None: + return + if not hasattr(result, "debug"): + return + debug = result.debug if isinstance(result.debug, dict) else {} + task_debug: Dict[str, int] = {} + if task_record.started_running_at is not None: + task_debug["queue_wait_ms"] = int((task_record.started_running_at - task_record.created_at) * 1000) + end_ts = task_record.finished_at if task_record.finished_at is not None else task_record.last_updated_at + task_debug["run_ms"] = int((end_ts - task_record.started_running_at) * 1000) + task_debug["lifecycle_ms"] = int((task_record.last_updated_at - task_record.created_at) * 1000) + debug["task"] = task_debug + result.debug = debug + + +def _serialize_action_value(value: Any) -> Any: + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, dict): + return {str(k): _serialize_action_value(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [_serialize_action_value(v) for v in value] + return repr(value) + + +async def execute_with_task_management( + action_payload: Dict[str, Any], + coro_factory: Callable[[], Awaitable[Any]], + time_to_live_ms: Optional[int] = None, + fast_response_threshold_seconds: float = 5.0, + scope: Optional[Any] = None, +) -> BaseResult: + # Deferred import avoids circular dependency: utils → async_task_manager → dataframe_manager → utils. + from config.storage import SessionScope + from tools.async_task_manager import ( + submit_task, + get_task_record, + remove_task, + task_snapshot, + ) + + if not isinstance(scope, SessionScope): + # Fail closed: a task without its chat session would land in a shared partition. + return BaseResult(error="Task could not be scheduled: no session scope for this call.") + resolved_scope = scope + wait_started = time.monotonic() + try: + task_id = await submit_task( + action_payload, + coro_factory, + time_to_live_ms=time_to_live_ms, + scope=resolved_scope, + ) + except RuntimeError as exc: + return BaseResult(error=str(exc)) + task_record = await get_task_record(task_id, scope=resolved_scope) + if not task_record or not task_record.asyncio_task: + return BaseResult(error="Task could not be scheduled.") + + try: + await asyncio.wait_for( + asyncio.shield(task_record.asyncio_task), + timeout=fast_response_threshold_seconds, + ) + latest_record = await get_task_record(task_id, scope=resolved_scope) + if not latest_record or latest_record.result is None: + await remove_task(task_id, scope=resolved_scope) + return BaseResult(error="Task finished without result.") + final_result = latest_record.result + failure = getattr(latest_record, "failure", None) + _attach_task_debug(final_result, latest_record) + debug = getattr(final_result, "debug", None) + if isinstance(debug, dict): + debug.setdefault("task", {}) + debug["task"]["sync_wait_ms"] = int((time.monotonic() - wait_started) * 1000) + await remove_task(task_id, scope=resolved_scope) + if failure is not None: + raise failure + return final_result + except asyncio.TimeoutError: + latest_record = await get_task_record(task_id, scope=resolved_scope) + if not latest_record: + return BaseResult(error="Task was not found after scheduling.") + snapshot = task_snapshot(latest_record, include_result=False) + timeout_result = BaseResult( + result=[snapshot], + info=[ + "Long-running operation accepted. Use perfecto_tools with action 'tasks_status' to monitor status." + ], + ) + _attach_task_debug(timeout_result, latest_record) + debug = getattr(timeout_result, "debug", None) + if isinstance(debug, dict): + debug.setdefault("task", {}) + debug["task"]["sync_wait_ms"] = int((time.monotonic() - wait_started) * 1000) + timeout_result.mark_pending_task(task_id, resolved_scope) + return timeout_result + + +def run_as_task( + time_to_live_ms: Optional[int] = None, + fast_response_threshold_seconds: float = 5.0, +): + def decorator(func: Callable[..., Awaitable[Any]]): + @functools.wraps(func) + async def wrapper(self, *args, **kwargs): + if _task_management_enabled.get(): + return await func(self, *args, **kwargs) + + try: + signature = inspect.signature(func) + bound = signature.bind(self, *args, **kwargs) + bound.apply_defaults() + named_params = { + key: _serialize_action_value(value) + for key, value in bound.arguments.items() + if key != "self" + } + except Exception: + named_params = {} + + action_payload = { + "manager": self.__class__.__name__, + "method": func.__name__, + "args": _serialize_action_value(args), + "kwargs": _serialize_action_value(kwargs), + "params": named_params, + "result_format": _result_format_context.get(), + "disable_dataframe_materialization": bool( + _disable_dataframe_materialization.get() + ), + } + + from config.session_context import SessionContextMissing + from tools.async_task_manager import session_scope_from_manager + + try: + scope = session_scope_from_manager(self) + except SessionContextMissing: + # Only public tools (help/skills) run without a session: no task and + # no partition, just run inline. Gated tools always have a session. + return await func(self, *args, **kwargs) + token = _task_management_enabled.set(True) + try: + coro_factory = lambda: func(self, *args, **kwargs) + return await execute_with_task_management( + action_payload=action_payload, + coro_factory=coro_factory, + time_to_live_ms=time_to_live_ms, + fast_response_threshold_seconds=fast_response_threshold_seconds, + scope=scope, + ) + finally: + _task_management_enabled.reset(token) + + return wrapper + + return decorator + + diff --git a/tools/tools_manager.py b/tools/tools_manager.py index b87c600..cc1e972 100644 --- a/tools/tools_manager.py +++ b/tools/tools_manager.py @@ -20,8 +20,16 @@ from config.version import __bundle__, __executable__, __uvx__, __version__ from models.manager import Manager from models.result import BaseResult -from telemetry import run_tool -from tools.utils import normalize_action_args, timeout, user_agent +from tools.call_state import TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME +from tools.data_tools import DataTools +from tools.runtime_tools import run_tool_with_runtime +from tools.utils import ( + normalize_action_args, + timeout, + user_agent, + validate_non_empty_str_arg, + validate_required_args, +) def _normalize_system(system: str) -> str: @@ -318,12 +326,37 @@ def register(mcp, runtime: AppRuntime): @mcp.tool( name=f"{TOOLS_PREFIX}_tools", description=""" -Operations on Perfecto MCP tooling metadata (versioning and updates). +Operations on Perfecto MCP tooling: version checks, session tasks, and session dataframes. Actions: - version: Return the current MCP version and runtime/platform information used by the user-agent and `--version` / console display. - check_updates: Query the GitHub repository for the latest release, compare it with the current version, and return release notes plus download links when an update is available. +- tasks_get: Get task metadata by task ID and return task_result when terminal. + args(dict): task_id (str, required); remove_on_terminal (bool, optional); + wait_for_terminal_ms (int, optional); poll_interval_ms (int, optional) +- tasks_status: Lightweight task status by task ID (no task_result payload). + args(dict): task_id (str, required); wait_for_terminal_ms (int, optional); + poll_interval_ms (int, optional) +- tasks_list: List tasks for the current MCP session. + args(dict): status (str, optional); status_list (list[str], optional) +- tasks_cancel: Cancel a running or queued task when this process owns it. + args(dict): task_id (str, required) +- tasks_remove: Remove a task from the session registry. + args(dict): task_id (str, required) +- dataframes_list: List dataframes and metadata for the current MCP session. +- dataframes_get: Get dataframe metadata and schema by dataframe ID. + args(dict): dataframe_id (str, required) +- dataframes_schema_groups: Group dataframe schemas for multi-dataframe queries. + args(dict): dataframe_id_list (list[str], optional) +- dataframes_query: Execute read-only SQL against session dataframe tables. + args(dict): sql (str, required); output_format (matrix|columnar|records); + result_format (auto|dataframe|raw, optional) + requirement: ORDER BY + LIMIT + OFFSET are mandatory in every query. +- dataframes_remove: Remove one or more dataframes from the session store. + args(dict): dataframe_id_list (list[str], required) +- dataframes_clear: Remove all dataframes for the current session. +- dataframes_sql_help: Describe supported SQL usage and blocked operations. Hints: - Prefer `version` first when the user asks what build they are running. - Prefer `check_updates` when the user asks whether a newer MCP release exists. @@ -332,6 +365,9 @@ def register(mcp, runtime: AppRuntime): check is unavailable, still share the current version, and point to the releases page for a manual check — do not treat it as a hard failure. - Render release and download URLs as markdown links. +- Before writing any dataframe SQL query, call `dataframes_sql_help`. +- ORDER BY + LIMIT + OFFSET are mandatory in every dataframe query. +- After a long-running tool returns a task snapshot, poll with tasks_status then tasks_get. """ ) async def tools( @@ -341,6 +377,7 @@ async def tools( action, args = normalize_action_args(arguments) runtime.configure_context(ctx) tools_manager = ToolsManager(ctx) + data_tools = DataTools(ctx, runtime.storage, runtime.scope_resolver) async def _dispatch(): match action: @@ -348,13 +385,84 @@ async def _dispatch(): return await tools_manager.version() case "check_updates": return await tools_manager.check_updates() + case "tasks_get": + if validation_error := validate_required_args(action, args, ["task_id"]): + return validation_error + if err := validate_non_empty_str_arg(action, args, "task_id"): + return err + return await data_tools.tasks_get( + task_id=str(args.get("task_id", "")).strip(), + remove_on_terminal=bool(args.get("remove_on_terminal", True)), + wait_for_terminal_ms=int(args.get("wait_for_terminal_ms", 0) or 0), + poll_interval_ms=int(args.get("poll_interval_ms", 1000) or 1000), + ) + case "tasks_list": + return await data_tools.tasks_list(args.get("status"), args.get("status_list")) + case "tasks_status": + if validation_error := validate_required_args(action, args, ["task_id"]): + return validation_error + if err := validate_non_empty_str_arg(action, args, "task_id"): + return err + return await data_tools.tasks_status( + task_id=str(args.get("task_id", "")).strip(), + wait_for_terminal_ms=int(args.get("wait_for_terminal_ms", 0) or 0), + poll_interval_ms=int(args.get("poll_interval_ms", 1000) or 1000), + ) + case "tasks_cancel": + if validation_error := validate_required_args(action, args, ["task_id"]): + return validation_error + if err := validate_non_empty_str_arg(action, args, "task_id"): + return err + return await data_tools.tasks_cancel(str(args.get("task_id", "")).strip()) + case "tasks_remove": + if validation_error := validate_required_args(action, args, ["task_id"]): + return validation_error + if err := validate_non_empty_str_arg(action, args, "task_id"): + return err + return await data_tools.tasks_remove(str(args.get("task_id", "")).strip()) + case "dataframes_list": + return await data_tools.dataframes_list() + case "dataframes_get": + if validation_error := validate_required_args(action, args, ["dataframe_id"]): + return validation_error + if err := validate_non_empty_str_arg(action, args, "dataframe_id"): + return err + return await data_tools.dataframes_get(str(args.get("dataframe_id", "")).strip()) + case "dataframes_schema_groups": + return await data_tools.dataframes_schema_groups(args.get("dataframe_id_list")) + case "dataframes_query": + if validation_error := validate_required_args(action, args, ["sql"]): + return validation_error + if err := validate_non_empty_str_arg(action, args, "sql"): + return err + return await data_tools.dataframes_query( + sql=str(args.get("sql", "")).strip(), + output_format=str(args.get("output_format", "matrix")), + result_format=str(args.get("result_format", "auto")), + ) + case "dataframes_remove": + if validation_error := validate_required_args(action, args, ["dataframe_id_list"]): + return validation_error + return await data_tools.dataframes_remove(args.get("dataframe_id_list")) + case "dataframes_clear": + return await data_tools.dataframes_clear() + case "dataframes_sql_help": + return await data_tools.dataframes_sql_help() case _: return BaseResult( error=f"Action {action} not found in tools manager tool" ) try: - return await run_tool(f"{TOOLS_PREFIX}_tools", action, ctx, _dispatch) + return await run_tool_with_runtime( + runtime, + f"{TOOLS_PREFIX}_tools", + action, + ctx, + _dispatch, + tool_args=args, + dataframe_excluded_actions=set(TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME), + ) except httpx.HTTPStatusError: return BaseResult( error=f"Error: {traceback.format_exc()}" diff --git a/tools/user_manager.py b/tools/user_manager.py index faac69b..2c3df4a 100644 --- a/tools/user_manager.py +++ b/tools/user_manager.py @@ -11,14 +11,18 @@ from formatters.user import format_users from models.manager import Manager from models.result import BaseResult -from telemetry import run_tool -from tools.utils import api_request, format_sanitized_traceback, normalize_action_args +from config.cache import CacheScope +from tools.method_cache import ttl_cache_method +from tools.runtime_tools import run_tool_with_runtime +from tools.utils import api_request, format_sanitized_traceback, normalize_action_args, run_as_task class UserManager(Manager): def __init__(self, ctx: Context): super().__init__(ctx) @token_verify + @ttl_cache_method(ttl_seconds=60, scope=CacheScope.USER) + @run_as_task() async def read_user(self) -> BaseResult: user_url = perfecto.get_user_management_api_url(self.token.cloud_name) user_url = user_url + "/current" @@ -64,7 +68,9 @@ async def _dispatch(): ) try: - return await run_tool(f"{TOOLS_PREFIX}_user", action, ctx, _dispatch) + return await run_tool_with_runtime( + runtime, f"{TOOLS_PREFIX}_user", action, ctx, _dispatch, tool_args=args, + ) except httpx.HTTPStatusError: return BaseResult( error=f"Error: {format_sanitized_traceback()}" diff --git a/tools/utils.py b/tools/utils.py index ec11bba..efeebb9 100644 --- a/tools/utils.py +++ b/tools/utils.py @@ -16,10 +16,21 @@ import httpx +from config.ids import SIMPLE_ID_ALPHABET, SIMPLE_ID_LENGTH from config.security import validate_http_request_endpoint from config.token import PerfectoToken from config.version import __version__ from models.result import BaseResult +from tools.call_state import ( + TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME, + generate_simple_id, + normalize_simple_id, + reset_disable_dataframe_materialization, + reset_result_format, + set_disable_dataframe_materialization, + set_result_format, +) +from tools.task_runtime import run_as_task so = platform.system() # "Windows", "Linux", "Darwin" version = platform.version() # kernel / build version @@ -213,6 +224,38 @@ def get_mcp_icon_uri(): icon_data = base64.standard_b64encode(icon_path.read_bytes()).decode() return f"data:image/png;base64,{icon_data}" + +def validate_required_args(action: str, args: Optional[Dict[str, Any]], required: list[str]) -> Optional[BaseResult]: + args = args or {} + missing = [key for key in required if key not in args or args[key] is None] + if not missing: + return None + missing_str = ", ".join(missing) + required_str = ", ".join(required) + return BaseResult( + error=( + f"Missing required args for action '{action}': {missing_str} not found within 'args'. " + f"Required args: {required_str}. Ensure parameters are passed inside the 'args' argument." + ) + ) + + +def validate_non_empty_str_arg( + action: str, args: Optional[Dict[str, Any]], key: str +) -> Optional[BaseResult]: + """Return BaseResult error if args[key] is missing, not a str, or only whitespace.""" + args = args or {} + value = args.get(key) + if not isinstance(value, str) or not value.strip(): + return BaseResult( + error=( + f"Missing required args for action '{action}': {key} must be a non-empty string " + f"within 'args'. Required args: {key}." + ) + ) + return None + + def normalize_action_args(arguments: Optional[Dict[str, Any]] = None) -> tuple[str, Dict[str, Any]]: """ Normalize tool arguments to (action, args) format. diff --git a/uv.lock b/uv.lock index 3171e83..88d7821 100644 --- a/uv.lock +++ b/uv.lock @@ -614,7 +614,7 @@ name = "macholib" version = "1.16.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "altgraph" }, + { name = "altgraph", marker = "sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/10/2f/97589876ea967487978071c9042518d28b958d87b17dceb7cdc1d881f963/macholib-1.16.4.tar.gz", hash = "sha256:f408c93ab2e995cd2c46e34fe328b130404be143469e41bc366c807448979362", size = 59427, upload-time = "2025-11-22T08:28:38.373Z" } wheels = [ @@ -814,6 +814,7 @@ dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-sdk" }, + { name = "polars" }, { name = "pydantic" }, { name = "pydantic-core" }, { name = "pydantic-settings" }, @@ -835,6 +836,7 @@ requires-dist = [ { name = "opentelemetry-api", specifier = ">=1.20.0" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.20.0" }, { name = "opentelemetry-sdk", specifier = ">=1.20.0" }, + { name = "polars", specifier = ">=1.40.1" }, { name = "pydantic", specifier = ">=2.11.7" }, { name = "pydantic-core", specifier = ">=2.33.2" }, { name = "pydantic-settings", specifier = ">=2.10.1" }, @@ -855,6 +857,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "polars" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "polars-runtime-32" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b9/da/c27daf205a824f89f71172a75ad0098c5837a5a31ddeaac7f8737ea8574b/polars-2.1.0.tar.gz", hash = "sha256:9eee488d6e10773c028753dfc8dac647fd60a708108dcb818a733e85d33df253", size = 811508, upload-time = "2026-10-25T18:00:33.43Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/d4/e696a46c74ec8cac6434764deaf72f7edcc2e91c2145a9b3f61900a7aa7a/polars-2.1.0-py3-none-any.whl", hash = "sha256:63c50ce3533f5a03d17c3a2a24f1d0e0c7eda98fbde4268a3cba900d34060e67", size = 911474, upload-time = "2026-10-25T17:58:26.863Z" }, +] + +[[package]] +name = "polars-runtime-32" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/a6/46e09af693326735b1fba3a8bf5a725b57136fb1342e707e75d3113c826b/polars_runtime_32-2.1.0.tar.gz", hash = "sha256:c575f5f3cdb8a5f9be65477b9b4202b062593687fe62a0082add87fbd4571067", size = 3878004, upload-time = "2026-10-25T18:00:34.932Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/40/b424f3946c702ff48eef14e44d35df81a25a784f20f77355f39c746528fc/polars_runtime_32-2.1.0-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:de243b6e5a0b9f3abeaefe542136fcdc6dd8df5be3c988a40838d075314e246c", size = 52169672, upload-time = "2026-10-25T17:58:32.437Z" }, + { url = "https://files.pythonhosted.org/packages/fe/38/07066430ac655a8cab2f9e96ea1fec6ca1f20ec47844afe117662d43ace0/polars_runtime_32-2.1.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:317314d65949f1d58c4939adf83aa0d2cfaf8da50fd4a725a2de27bde2c10018", size = 47626883, upload-time = "2026-10-25T17:58:38.688Z" }, + { url = "https://files.pythonhosted.org/packages/49/a8/a773e20b451bbaab4d8267867b7c43a0c3aa67ebb6f781423a42a5a1af37/polars_runtime_32-2.1.0-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:777228cc8bf9cda36dee85ba74d13a25974ad4caffc5da93e10b1da9d810ce2d", size = 50033514, upload-time = "2026-10-25T17:58:44.261Z" }, + { url = "https://files.pythonhosted.org/packages/d2/e6/649c524252d32b24d0ba92278033823fa8e0040ee3748ce5afbeda107403/polars_runtime_32-2.1.0-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b838056cf365d7497971d9d9825c547766923338771c48c545b86ead2ebc572", size = 53863954, upload-time = "2026-10-25T17:58:48.737Z" }, + { url = "https://files.pythonhosted.org/packages/90/f7/2eb405bf5ead6373a33fa5dd4637fe996565a1f07ccce8911ed167673740/polars_runtime_32-2.1.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:04cfc3db1571b74ea4145944de7feceed4bfc9b26e5b2c265e27ad3b2d3866a9", size = 50204290, upload-time = "2026-10-25T17:58:53.33Z" }, + { url = "https://files.pythonhosted.org/packages/33/4d/d76cd42042605e079d6df91cc4fea09ed5f4cb836c72c952a17f585860a9/polars_runtime_32-2.1.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:af58b5f6e5cc44b578f2a93d752f87c135c3a53205b4360301bc3a9b89dfbfd1", size = 53794253, upload-time = "2026-10-25T17:58:57.99Z" }, + { url = "https://files.pythonhosted.org/packages/4d/4e/2a419d0def2908255dd1b7318c48430319859dc92ca156e2f81ff7bcd78b/polars_runtime_32-2.1.0-cp311-abi3-win_amd64.whl", hash = "sha256:f2fe32f48af2395cd6ec52a433a03970a69aab1dbb66a613729fd549ce92c904", size = 53926265, upload-time = "2026-10-25T17:59:03.13Z" }, + { url = "https://files.pythonhosted.org/packages/c8/1b/15ffba3264baf21d2f04c51f54b01159cb6f4fad4e84de1404be552d601f/polars_runtime_32-2.1.0-cp311-abi3-win_arm64.whl", hash = "sha256:0af30dd53a7725f703d0d21dd421a39e0165378bfaa41431008f286608376d71", size = 48636436, upload-time = "2026-10-25T17:59:08.476Z" }, +] + [[package]] name = "protobuf" version = "7.35.1" From 4593e236a5a997bb951772a40d5923648dd16dd6 Mon Sep 17 00:00:00 2001 From: diego-ferrand Date: Thu, 5 Nov 2026 11:52:33 -0300 Subject: [PATCH 2/2] Update comments --- config/http_clients.py | 5 ++--- config/session_context.py | 2 +- tests/test_dataframes_tasks_cache.py | 2 +- tools/dataframe_manager.py | 4 ++-- 4 files changed, 6 insertions(+), 7 deletions(-) diff --git a/config/http_clients.py b/config/http_clients.py index 763da9b..f425bbb 100644 --- a/config/http_clients.py +++ b/config/http_clients.py @@ -18,9 +18,8 @@ request reuses the connection instead of paying a new TCP/TLS (and HTTP/2) handshake. Built on first use, rebuilt if closed, closed on server shutdown. -A shared client serves every caller (hosted: every user), so it never keeps -cookies: BlazeMeter answers with a ``bzm_sess`` session cookie, and a cookie -jar would send one user's session with the next user's requests. +A shared client serves every caller, so it never keeps cookies. A session +cookie from one caller must not be sent with the next caller's requests. """ import weakref from http.cookiejar import CookieJar, DefaultCookiePolicy diff --git a/config/session_context.py b/config/session_context.py index ae4de70..ba79816 100644 --- a/config/session_context.py +++ b/config/session_context.py @@ -3,7 +3,7 @@ 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. +context vars. """ from __future__ import annotations diff --git a/tests/test_dataframes_tasks_cache.py b/tests/test_dataframes_tasks_cache.py index 0163f8e..f30e048 100644 --- a/tests/test_dataframes_tasks_cache.py +++ b/tests/test_dataframes_tasks_cache.py @@ -1,4 +1,4 @@ -"""Session dataframes, async tasks and the method cache ported from BlazeMeter MCP.""" +"""Session dataframes, async tasks and the method cache.""" import asyncio from types import SimpleNamespace diff --git a/tools/dataframe_manager.py b/tools/dataframe_manager.py index f658d94..98cce73 100644 --- a/tools/dataframe_manager.py +++ b/tools/dataframe_manager.py @@ -949,7 +949,7 @@ def get_sql_capabilities() -> Dict[str, Any]: "engine_scope": { "query_entrypoints": ["SELECT", "WITH"], "mode": "read-only", - "description": "SQL support is available for SELECT/WITH queries with BlazeMeter MCP query constraints.", + "description": "SQL support is available for SELECT/WITH queries with Perfecto MCP query constraints.", }, "allowed_entrypoints": ["SELECT", "WITH"], "disallowed_statements": [ @@ -1181,7 +1181,7 @@ def get_sql_capabilities() -> Dict[str, Any]: "notes": [ "All session dataframe tables (loaded from Storage for the current partition) are queryable in the same SQL context.", "Function name typo seen in some sources: STRPOST; use STRPOS.", - "This help defines practical usage constraints for BlazeMeter MCP SQL queries.", + "This help defines practical usage constraints for Perfecto MCP SQL queries.", ], "references": [ "https://docs.pola.rs/py-polars/html/reference/sql/index.html",