From acc6d99fa8bb75565e34fff818a72a98fbe879ec Mon Sep 17 00:00:00 2001 From: Diego Ferrand Date: Tue, 4 Aug 2026 15:20:13 -0300 Subject: [PATCH 1/6] Http auth (#93) * Add otel documentation * Add per-request Bearer auth for hosted HTTP MCP * Match original arquitecture plan naming and refactors. * Integrate streamable HTTP transport with Bearer auth wiring --- README.md | 28 +++++ config/auth.py | 132 ++++++++++++++++++++++++ config/runtime.py | 47 +++++++++ config/token.py | 29 ++++++ main.py | 36 +++++-- server.py | 28 +++-- tests/test_batch_controls.py | 5 +- tests/test_failure_criteria.py | 3 +- tests/test_http_auth.py | 165 ++++++++++++++++++++++++++++++ tests/test_main_transport.py | 109 ++++++++++++++++++-- tests/test_required_args_tools.py | 27 ++--- tools/account_manager.py | 5 +- tools/billing_manager.py | 5 +- tools/execution_manager.py | 7 +- tools/help_manager.py | 5 +- tools/project_manager.py | 5 +- tools/skills_manager.py | 5 +- tools/test_manager.py | 5 +- tools/user_manager.py | 5 +- tools/workspace_manager.py | 5 +- 20 files changed, 589 insertions(+), 67 deletions(-) create mode 100644 config/auth.py create mode 100644 config/runtime.py create mode 100644 tests/test_http_auth.py diff --git a/README.md b/README.md index 899d84f..581db9a 100644 --- a/README.md +++ b/README.md @@ -110,6 +110,34 @@ After installing, set `BLAZEMETER_API_KEY` to your `api-key.json` path in your c --- +**Hosted HTTP (streamable-http) Client Configuration** + +Run a shared server that authenticates each client via `Authorization: Bearer`: + +```bash +uv run python main.py --mcp --transport streamable-http //PLACEHOLDER +``` + +Configure the MCP client with the server URL and your BlazeMeter API key as Bearer credentials (`id:secret` or base64 of `id:secret`): + +```json +{ + "mcpServers": { + "BlazeMeter MCP": { + "url": "http://localhost:8000/mcp", + "headers": { + "Authorization": "Bearer :" + } + } + } +} +``` + +> [!NOTE] +> Over HTTP, credentials are resolved per request from the `Authorization` header. Invalid or missing Bearer credentials return `401` before any tool runs. Well-formed but wrong API keys fail later inside BlazeMeter API calls (same as stdio). Stdio transport uses `api-key.json` / env / Docker secrets. + +--- + **Docker MCP Client Configuration** 1. **Prerequisites:** [Docker]([https://www.docker.com/products/docker-desktop/](https://www.docker.com/products/docker-desktop/)) diff --git a/config/auth.py b/config/auth.py new file mode 100644 index 0000000..b028011 --- /dev/null +++ b/config/auth.py @@ -0,0 +1,132 @@ +""" +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 typing import Optional, Protocol, runtime_checkable + +from mcp.server.fastmcp import Context, FastMCP +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send + +from config.token import BzmToken, BzmTokenError + +BZM_TOKEN_STATE_ATTR = "token" + + +class AuthError(Exception): + """Raised when Authorization cannot be parsed into credentials.""" + + +@runtime_checkable +class AuthPort(Protocol): + """Resolves the BlazeMeter API token for the current tool invocation.""" + + def get_token(self, ctx: Context) -> Optional[BzmToken]: + ... + + +class StdioAuthProvider: + """Process-lifetime token from env / api-key.json / Docker secrets.""" + + def __init__(self, token: Optional[BzmToken]): + self._token = token + + def get_token(self, ctx: Context) -> Optional[BzmToken]: + return self._token + + +class HttpAuthProvider: + """Per-request token attached by Bearer auth middleware to request.state.""" + + def get_token(self, ctx: Context) -> Optional[BzmToken]: + request = ctx.request_context.request + if request is None: + return None + return getattr(request.state, BZM_TOKEN_STATE_ATTR, None) + + +def parse_authorization_header(value: Optional[str]) -> BzmToken: + """ + Parse ``Authorization: Bearer `` into a BzmToken. + + Credentials may be ``id:secret`` or base64(``id:secret``). Does not call + the BlazeMeter API — parse only. + """ + if not value or not value.strip(): + raise AuthError("Missing Authorization header") + + scheme, _, credentials = value.strip().partition(" ") + if scheme.lower() != "bearer" or not credentials.strip(): + raise AuthError("Authorization header must use Bearer scheme") + + try: + return BzmToken.from_bearer_credentials(credentials.strip()) + except BzmTokenError as exc: + raise AuthError("Unparseable Bearer credentials") from exc + + +class BearerAuthMiddleware: + """ + HTTP gate: require a parseable Bearer token on every request. + + Attaches BzmToken to ``request.state``; does not validate against BlazeMeter. + """ + + def __init__(self, app: ASGIApp): + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + if scope.get("method") == "OPTIONS": + await self.app(scope, receive, send) + return + + request = Request(scope, receive) + try: + token = parse_authorization_header(request.headers.get("authorization")) + except AuthError: + response = JSONResponse( + {"error": "Unauthorized"}, + status_code=401, + headers={"WWW-Authenticate": "Bearer"}, + ) + await response(scope, receive, send) + return + + setattr(request.state, BZM_TOKEN_STATE_ATTR, token) + await self.app(scope, receive, send) + + +def run_streamable_http(mcp: FastMCP) -> None: + """Serve FastMCP over streamable HTTP with Bearer auth middleware.""" + import anyio + import uvicorn + + async def _serve() -> None: + app = BearerAuthMiddleware(mcp.streamable_http_app()) + config = uvicorn.Config( + app, + host=mcp.settings.host, + port=mcp.settings.port, + log_level=mcp.settings.log_level.lower(), + ) + await uvicorn.Server(config).serve() + + anyio.run(_serve) diff --git a/config/runtime.py b/config/runtime.py new file mode 100644 index 0000000..956d181 --- /dev/null +++ b/config/runtime.py @@ -0,0 +1,47 @@ +""" +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 dataclasses import dataclass +from typing import Literal, Optional + +from config.auth import AuthPort, HttpAuthProvider, StdioAuthProvider +from config.token import BzmToken + +Transport = Literal["stdio", "streamable-http"] + + +@dataclass(frozen=True) +class AppRuntime: + """Process-level collaborators shared by tool registrations.""" + + transport: Transport + auth: AuthPort + + +def build_runtime( + transport: Transport, + startup_token: Optional[BzmToken] = None, +) -> AppRuntime: + """ + Compose auth for the selected transport. + + - stdio: use process-lifetime ``startup_token`` (from env / api-key.json / Docker). + - streamable-http: resolve credentials per request via Bearer middleware + HttpAuthProvider. + """ + if transport == "stdio": + return AppRuntime(transport=transport, auth=StdioAuthProvider(startup_token)) + if transport == "streamable-http": + return AppRuntime(transport=transport, auth=HttpAuthProvider()) + raise ValueError(f"Unknown transport: {transport}") diff --git a/config/token.py b/config/token.py index b7bdd44..465e0c3 100644 --- a/config/token.py +++ b/config/token.py @@ -58,6 +58,35 @@ def from_file(cls, path: Union[str, Path]) -> "BzmToken": return cls(token_id=id_val, token_secret=secret_val) + @classmethod + def from_bearer_credentials(cls, credentials: str) -> "BzmToken": + """ + Parse Bearer credential material into a BzmToken. + + Accepts plaintext ``id:secret`` or base64-encoded ``id:secret``. + Does not call the BlazeMeter API. + """ + raw = (credentials or "").strip() + if not raw: + raise BzmTokenError("Empty bearer credentials") + + # Plaintext id:secret (base64 alphabet has no ':') + if ":" in raw: + token_id, _, token_secret = raw.partition(":") + return cls(token_id=token_id, token_secret=token_secret) + + try: + pad = "=" * (-len(raw) % 4) + decoded = base64.b64decode(raw + pad, validate=False).decode("utf-8") + except Exception as e: + raise BzmTokenError("Invalid bearer credentials encoding") from e + + if ":" not in decoded: + raise BzmTokenError("Invalid bearer credentials format") + + token_id, _, token_secret = decoded.partition(":") + return cls(token_id=token_id, token_secret=token_secret) + def as_basic_auth(self) -> str: """ Returns the HTTP Basic Authentication header: diff --git a/main.py b/main.py index 1ee9b52..ec0b7c3 100644 --- a/main.py +++ b/main.py @@ -29,6 +29,8 @@ from mcp.server.fastmcp import FastMCP +from config.auth import run_streamable_http +from config.runtime import build_runtime from config.token import BzmToken, BzmTokenError from config.version import __version__, __executable__, __bundle__ from server import register_tools @@ -404,13 +406,23 @@ def resolve_mcp_transport(raw_cli_transport: str) -> str: return normalized -def build_runtime( +def to_wire_transport(logical_transport: str) -> Literal["stdio", "streamable-http"]: + """Map CLI/logical transport (stdio|http|docker) to FastMCP wire transport.""" + return "streamable-http" if logical_transport == "http" else "stdio" + + +def build_mcp_server( log_level: str = "CRITICAL", confirm_mode: ConfirmMode = ConfirmMode.DELETE, transport: str = "stdio", ) -> tuple[FastMCP, str]: + """ + Build FastMCP + auth wiring for a logical CLI transport (stdio|http|docker). + + Returns ``(mcp, wire_transport)`` where ``wire_transport`` is the FastMCP + transport name (``stdio`` or ``streamable-http``). + """ init_telemetry("bzm-mcp", __version__) - token = get_token() host = "127.0.0.1" port = 8000 streamable_http_path = "/mcp" @@ -418,6 +430,13 @@ def build_runtime( host = os.getenv("FASTMCP_HOST", "127.0.0.1").strip() or "127.0.0.1" port = int(os.getenv("FASTMCP_PORT", "8000").strip() or "8000") streamable_http_path = os.getenv("FASTMCP_STREAMABLE_HTTP_PATH", "/mcp").strip() or "/mcp" + + # docker and stdio share process-lifetime credentials; http uses Bearer per request. + wire_transport = to_wire_transport(transport) + app_runtime = build_runtime( + wire_transport, + startup_token=get_token() if wire_transport == "stdio" else None, + ) instructions = """ # BlazeMeter MCP Server A comprehensive integration tool that provides AI assistants with full programmatic access to BlazeMeter's cloud-based performance testing platform. Enables automated management of complete load testing workflows from creation to execution and reporting. Transforms enterprise-grade testing capabilities into an AI-accessible service for intelligent automation of complex performance testing scenarios. @@ -483,18 +502,21 @@ def build_runtime( stateless_http=False, ) register_confirm_mode(confirm_mode) - register_tools(mcp, token) - runtime_transport = "streamable-http" if transport == "http" else "stdio" - return mcp, runtime_transport + register_tools(mcp, app_runtime) + return mcp, wire_transport def run(log_level: str = "CRITICAL", confirm_mode: ConfirmMode = ConfirmMode.DELETE, transport: str = "stdio"): - mcp, runtime_transport = build_runtime( + mcp, runtime_transport = build_mcp_server( log_level=log_level, confirm_mode=confirm_mode, transport=transport, ) - mcp.run(transport=runtime_transport) + if runtime_transport == "stdio": + mcp.run(transport=runtime_transport) + else: + # Hosted HTTP requires Bearer auth middleware around the ASGI app. + run_streamable_http(mcp) def main(): diff --git a/server.py b/server.py index 104272b..c5ed472 100644 --- a/server.py +++ b/server.py @@ -13,9 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. """ -from typing import Optional - -from config.token import BzmToken +from config.runtime import AppRuntime from tools.account_manager import register as register_account_manager from tools.billing_manager import register as register_billing_manager from tools.execution_manager import register as register_execution_manager @@ -27,20 +25,20 @@ from tools.workspace_manager import register as register_workspace_manager -def register_tools(mcp, token: Optional[BzmToken]): +def register_tools(mcp, runtime: AppRuntime): """ Register all available tools with the MCP server. - + Args: mcp: The MCP server instance - token: Optional BlazeMeter token (can be None if not configured) + runtime: App runtime (transport + auth port and shared collaborators) """ - register_user_manager(mcp, token) - register_project_manager(mcp, token) - register_workspace_manager(mcp, token) - register_test_manager(mcp, token) - register_execution_manager(mcp, token) - register_account_manager(mcp, token) - register_billing_manager(mcp, token) - register_help_manager(mcp, token) - register_skills_manager(mcp, token) + register_user_manager(mcp, runtime) + register_project_manager(mcp, runtime) + register_workspace_manager(mcp, runtime) + register_test_manager(mcp, runtime) + register_execution_manager(mcp, runtime) + register_account_manager(mcp, runtime) + register_billing_manager(mcp, runtime) + register_help_manager(mcp, runtime) + register_skills_manager(mcp, runtime) diff --git a/tests/test_batch_controls.py b/tests/test_batch_controls.py index ff5f1cf..a402107 100644 --- a/tests/test_batch_controls.py +++ b/tests/test_batch_controls.py @@ -14,6 +14,7 @@ limitations under the License. """ +from config.runtime import AppRuntime, build_runtime import asyncio from config.blazemeter import TOOLS_PREFIX @@ -41,7 +42,7 @@ def decorator(func): class TestBatchControls: def test_help_batch_respects_concurrency_limit(self, monkeypatch): mcp = FakeMcp() - register_help_tool(mcp, token=None) + register_help_tool(mcp, build_runtime("stdio")) help_tool = mcp.tools[f"{TOOLS_PREFIX}_help"] HelpManager.help_tree = {} monkeypatch.setattr(HelpManager, "MAX_BATCH_CONCURRENCY", 2) @@ -67,7 +68,7 @@ async def slow_list_help_categories(self): def test_skills_batch_respects_concurrency_limit(self, monkeypatch): mcp = FakeMcp() - register_skills_tool(mcp, token=None) + register_skills_tool(mcp, build_runtime("stdio")) skills_tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] monkeypatch.setattr(SkillsManager, "MAX_BATCH_CONCURRENCY", 2) diff --git a/tests/test_failure_criteria.py b/tests/test_failure_criteria.py index 10e206d..a8ad36a 100644 --- a/tests/test_failure_criteria.py +++ b/tests/test_failure_criteria.py @@ -14,6 +14,7 @@ limitations under the License. """ +from config.runtime import AppRuntime, build_runtime import asyncio import pytest @@ -277,7 +278,7 @@ def test_contains_catalog(self): class TestFailureCriteriaMetaAction: def test_tool_returns_catalog_without_api(self): mcp = _FakeMcpForTests() - register_tests_tool(mcp, token=None) + register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] result = asyncio.run(tool("failure_criteria_meta", {}, ctx=None)) assert result.error is None diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py new file mode 100644 index 0000000..9fe95c8 --- /dev/null +++ b/tests/test_http_auth.py @@ -0,0 +1,165 @@ +""" +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 base64 +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.routing import Route +from starlette.testclient import TestClient + +from config.auth import ( + AuthError, + BZM_TOKEN_STATE_ATTR, + BearerAuthMiddleware, + HttpAuthProvider, + StdioAuthProvider, + parse_authorization_header, +) +from config.runtime import build_runtime +from config.token import BzmToken, BzmTokenError + + +class TestBearerCredentialParsing: + def test_plaintext_id_secret(self): + token = BzmToken.from_bearer_credentials("key-id:key-secret") + assert token.id == "key-id" + assert token.secret == "key-secret" + + def test_base64_id_secret(self): + raw = base64.b64encode(b"key-id:key-secret").decode() + token = BzmToken.from_bearer_credentials(raw) + assert token.id == "key-id" + assert token.secret == "key-secret" + + def test_empty_raises(self): + with pytest.raises(BzmTokenError): + BzmToken.from_bearer_credentials(" ") + + def test_unparseable_raises(self): + with pytest.raises(BzmTokenError): + BzmToken.from_bearer_credentials(base64.b64encode(b"no-colon").decode()) + + +class TestAuthorizationHeaderParsing: + def test_bearer_plaintext(self): + token = parse_authorization_header("Bearer key-id:key-secret") + assert token.id == "key-id" + assert token.secret == "key-secret" + + def test_bearer_base64(self): + raw = base64.b64encode(b"key-id:key-secret").decode() + token = parse_authorization_header(f"Bearer {raw}") + assert token.id == "key-id" + assert token.secret == "key-secret" + + def test_missing_header(self): + with pytest.raises(AuthError): + parse_authorization_header(None) + + def test_wrong_scheme(self): + with pytest.raises(AuthError): + parse_authorization_header("Basic abc") + + +class TestAuthProviders: + def test_stdio_returns_startup_token(self): + token = BzmToken("id", "secret") + provider = StdioAuthProvider(token) + assert provider.get_token(ctx=None) is token + + def test_stdio_allows_none(self): + assert StdioAuthProvider(None).get_token(ctx=None) is None + + def test_http_reads_request_state(self): + token_a = BzmToken("a-id", "a-secret") + token_b = BzmToken("b-id", "b-secret") + provider = HttpAuthProvider() + + def ctx_with(token: BzmToken): + request = SimpleNamespace(state=SimpleNamespace(**{BZM_TOKEN_STATE_ATTR: token})) + request_context = SimpleNamespace(request=request) + return SimpleNamespace(request_context=request_context) + + assert provider.get_token(ctx_with(token_a)).id == "a-id" + assert provider.get_token(ctx_with(token_b)).id == "b-id" + + def test_http_concurrent_tokens_are_isolated(self): + """Two contexts with different Bearer-derived tokens resolve independently.""" + provider = HttpAuthProvider() + token_a = BzmToken("account-a", "secret-a") + token_b = BzmToken("account-b", "secret-b") + + def make_ctx(token: BzmToken): + request = MagicMock() + setattr(request.state, BZM_TOKEN_STATE_ATTR, token) + ctx = MagicMock() + ctx.request_context.request = request + return ctx + + assert provider.get_token(make_ctx(token_a)).id == "account-a" + assert provider.get_token(make_ctx(token_b)).id == "account-b" + + +class TestBearerAuthMiddleware: + def _app(self): + async def ok(request: Request): + token = getattr(request.state, BZM_TOKEN_STATE_ATTR, None) + return JSONResponse({"id": token.id if token else None}) + + return BearerAuthMiddleware(Starlette(routes=[Route("/mcp", endpoint=ok, methods=["POST"])])) + + def test_missing_authorization_returns_401(self): + client = TestClient(self._app()) + response = client.post("/mcp") + assert response.status_code == 401 + assert response.json()["error"] == "Unauthorized" + + def test_invalid_bearer_returns_401(self): + client = TestClient(self._app()) + response = client.post("/mcp", headers={"Authorization": "Bearer not-valid"}) + assert response.status_code == 401 + + def test_valid_bearer_attaches_token(self): + client = TestClient(self._app()) + response = client.post( + "/mcp", + headers={"Authorization": "Bearer key-id:key-secret"}, + ) + assert response.status_code == 200 + assert response.json()["id"] == "key-id" + + def test_options_bypasses_auth(self): + async def ok(_request: Request): + return JSONResponse({"ok": True}) + + app = BearerAuthMiddleware(Starlette(routes=[Route("/mcp", endpoint=ok, methods=["OPTIONS"])])) + client = TestClient(app) + assert client.options("/mcp").status_code == 200 + + +class TestBuildRuntime: + def test_build_runtime_stdio_and_http(self): + stdio = build_runtime("stdio") + assert stdio.transport == "stdio" + assert isinstance(stdio.auth, StdioAuthProvider) + + http = build_runtime("streamable-http") + assert http.transport == "streamable-http" + assert isinstance(http.auth, HttpAuthProvider) diff --git a/tests/test_main_transport.py b/tests/test_main_transport.py index dd3c6a5..ed51904 100644 --- a/tests/test_main_transport.py +++ b/tests/test_main_transport.py @@ -1,6 +1,8 @@ import pytest import main +from config.auth import HttpAuthProvider, StdioAuthProvider +from config.runtime import AppRuntime class _DummyFastMCP: @@ -13,7 +15,7 @@ def run(self, transport="stdio", mount_path=None): self.run_calls.append({"transport": transport, "mount_path": mount_path}) -def _patch_runtime_dependencies(monkeypatch): +def _patch_mcp_server_dependencies(monkeypatch): monkeypatch.setattr(main, "init_telemetry", lambda *a, **k: None) monkeypatch.setattr(main, "get_token", lambda: object()) monkeypatch.setattr(main, "register_confirm_mode", lambda *a, **k: None) @@ -39,14 +41,21 @@ def test_invalid_transport_raises_clear_error(self, monkeypatch): main.resolve_mcp_transport("banana") -class TestBuildRuntimeHttp: - def test_http_runtime_uses_env_settings_and_stateful_http(self, monkeypatch): - _patch_runtime_dependencies(monkeypatch) +class TestToWireTransport: + def test_maps_logical_transports(self): + assert main.to_wire_transport("http") == "streamable-http" + assert main.to_wire_transport("stdio") == "stdio" + assert main.to_wire_transport("docker") == "stdio" + + +class TestBuildMcpServerHttp: + def test_http_uses_env_settings_and_stateful_http(self, monkeypatch): + _patch_mcp_server_dependencies(monkeypatch) monkeypatch.setenv("FASTMCP_HOST", "0.0.0.0") monkeypatch.setenv("FASTMCP_PORT", "8012") monkeypatch.setenv("FASTMCP_STREAMABLE_HTTP_PATH", "/custom-mcp") - mcp, runtime_transport = main.build_runtime(transport="http") + mcp, runtime_transport = main.build_mcp_server(transport="http") assert runtime_transport == "streamable-http" assert isinstance(mcp, _DummyFastMCP) @@ -56,14 +65,94 @@ def test_http_runtime_uses_env_settings_and_stateful_http(self, monkeypatch): assert mcp.kwargs["stateless_http"] is False -class TestBuildRuntimeTransportMapping: +class TestBuildMcpServerTransportMapping: def test_transport_mapping_keeps_docker_stdio_and_http_streamable(self, monkeypatch): - _patch_runtime_dependencies(monkeypatch) + _patch_mcp_server_dependencies(monkeypatch) - _mcp_http, runtime_transport_http = main.build_runtime(transport="http") - _mcp_docker, runtime_transport_docker = main.build_runtime(transport="docker") - _mcp_stdio, runtime_transport_stdio = main.build_runtime(transport="stdio") + _mcp_http, runtime_transport_http = main.build_mcp_server(transport="http") + _mcp_docker, runtime_transport_docker = main.build_mcp_server(transport="docker") + _mcp_stdio, runtime_transport_stdio = main.build_mcp_server(transport="stdio") assert runtime_transport_http == "streamable-http" assert runtime_transport_docker == "stdio" assert runtime_transport_stdio == "stdio" + + +class TestBuildMcpServerAuthWiring: + def test_http_registers_http_auth_provider(self, monkeypatch): + captured = {} + + def capture_register(mcp, runtime): + captured["runtime"] = runtime + + _patch_mcp_server_dependencies(monkeypatch) + monkeypatch.setattr(main, "register_tools", capture_register) + + main.build_mcp_server(transport="http") + + runtime = captured["runtime"] + assert isinstance(runtime, AppRuntime) + assert runtime.transport == "streamable-http" + assert isinstance(runtime.auth, HttpAuthProvider) + + def test_stdio_and_docker_register_stdio_auth_provider(self, monkeypatch): + captured = {} + + def capture_register(mcp, runtime): + captured.setdefault("runtimes", []).append(runtime) + + token = object() + _patch_mcp_server_dependencies(monkeypatch) + monkeypatch.setattr(main, "get_token", lambda: token) + monkeypatch.setattr(main, "register_tools", capture_register) + + main.build_mcp_server(transport="stdio") + main.build_mcp_server(transport="docker") + + assert len(captured["runtimes"]) == 2 + for runtime in captured["runtimes"]: + assert isinstance(runtime.auth, StdioAuthProvider) + assert runtime.auth.get_token(ctx=None) is token + + +class TestRunTransportDispatch: + def test_http_uses_bearer_middleware_server(self, monkeypatch): + calls = {"stdio": 0, "http": 0} + + class _Mcp: + def run(self, transport="stdio", mount_path=None): + calls["stdio"] += 1 + + _patch_mcp_server_dependencies(monkeypatch) + monkeypatch.setattr(main, "build_mcp_server", lambda **k: (_Mcp(), "streamable-http")) + monkeypatch.setattr( + main, + "run_streamable_http", + lambda mcp: calls.__setitem__("http", calls["http"] + 1), + ) + + main.run(transport="http") + + assert calls["http"] == 1 + assert calls["stdio"] == 0 + + def test_stdio_uses_mcp_run(self, monkeypatch): + calls = {"stdio": 0, "http": 0} + + class _Mcp: + def run(self, transport="stdio", mount_path=None): + calls["stdio"] += 1 + assert transport == "stdio" + + _patch_mcp_server_dependencies(monkeypatch) + monkeypatch.setattr(main, "build_mcp_server", lambda **k: (_Mcp(), "stdio")) + monkeypatch.setattr( + main, + "run_streamable_http", + lambda mcp: calls.__setitem__("http", calls["http"] + 1), + ) + + main.run(transport="stdio") + + assert calls["stdio"] == 1 + assert calls["http"] == 0 diff --git a/tests/test_required_args_tools.py b/tests/test_required_args_tools.py index 265abcc..05ffde6 100644 --- a/tests/test_required_args_tools.py +++ b/tests/test_required_args_tools.py @@ -14,6 +14,7 @@ limitations under the License. """ +from config.runtime import AppRuntime, build_runtime import asyncio from config.blazemeter import TOOLS_PREFIX @@ -45,7 +46,7 @@ def decorator(func): class TestRequiredArgumentsForTools: def test_account_read_requires_account_id(self): mcp = FakeMcp() - register_account_tool(mcp, token=None) + register_account_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_account"] result = asyncio.run(tool("read", {}, ctx=None)) @@ -54,7 +55,7 @@ def test_account_read_requires_account_id(self): def test_workspace_list_requires_account_id(self): mcp = FakeMcp() - register_workspaces_tool(mcp, token=None) + register_workspaces_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_workspaces"] result = asyncio.run(tool("list", {}, ctx=None)) @@ -63,7 +64,7 @@ def test_workspace_list_requires_account_id(self): def test_project_read_requires_project_id(self): mcp = FakeMcp() - register_project_tool(mcp, token=None) + register_project_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_project"] result = asyncio.run(tool("read", {}, ctx=None)) @@ -72,7 +73,7 @@ def test_project_read_requires_project_id(self): def test_tests_create_requires_test_name(self): mcp = FakeMcp() - register_tests_tool(mcp, token=None) + register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] result = asyncio.run(tool("create", {"project_id": 123}, ctx=None)) @@ -81,7 +82,7 @@ def test_tests_create_requires_test_name(self): def test_tests_upload_assets_requires_file_paths(self): mcp = FakeMcp() - register_tests_tool(mcp, token=None) + register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] result = asyncio.run(tool("upload_assets", {"test_id": 123}, ctx=None)) @@ -90,7 +91,7 @@ def test_tests_upload_assets_requires_file_paths(self): def test_tests_configure_failure_criteria_requires_enabled_and_rules(self): mcp = FakeMcp() - register_tests_tool(mcp, token=None) + register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] result = asyncio.run(tool("configure_failure_criteria", {"test_id": 123}, ctx=None)) @@ -105,7 +106,7 @@ def test_tests_configure_failure_criteria_requires_enabled_and_rules(self): def test_execution_read_requires_execution_id(self): mcp = FakeMcp() - register_execution_tool(mcp, token=None) + register_execution_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_execution"] result = asyncio.run(tool("read", {}, ctx=None)) @@ -114,7 +115,7 @@ def test_execution_read_requires_execution_id(self): def test_execution_read_summary_requires_execution_id(self): mcp = FakeMcp() - register_execution_tool(mcp, token=None) + register_execution_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_execution"] result = asyncio.run(tool("read_summary", {}, ctx=None)) @@ -123,7 +124,7 @@ def test_execution_read_summary_requires_execution_id(self): def test_skills_read_skill_requires_skill_name(self): mcp = FakeMcp() - register_skills_tool(mcp, token=None) + register_skills_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] result = asyncio.run(tool("read_skill", {}, ctx=None)) @@ -132,7 +133,7 @@ def test_skills_read_skill_requires_skill_name(self): def test_skills_read_skill_resource_uri_requires_uri(self): mcp = FakeMcp() - register_skills_tool(mcp, token=None) + register_skills_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] result = asyncio.run(tool("read_skill_resource_uri", {}, ctx=None)) @@ -141,7 +142,7 @@ def test_skills_read_skill_resource_uri_requires_uri(self): def test_skills_read_skill_resource_uri_list_requires_non_empty_list(self): mcp = FakeMcp() - register_skills_tool(mcp, token=None) + register_skills_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] result = asyncio.run(tool("read_skill_resource_uri_list", {}, ctx=None)) @@ -150,7 +151,7 @@ def test_skills_read_skill_resource_uri_list_requires_non_empty_list(self): def test_help_read_help_info_requires_help_id_list(self): mcp = FakeMcp() - register_help_tool(mcp, token=None) + register_help_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_help"] result = asyncio.run(tool("read_help_info", {}, ctx=None)) @@ -159,7 +160,7 @@ def test_help_read_help_info_requires_help_id_list(self): def test_help_list_help_category_content_requires_subcategory_list(self): mcp = FakeMcp() - register_help_tool(mcp, token=None) + register_help_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_help"] result = asyncio.run(tool("list_help_category_content", {}, ctx=None)) diff --git a/tools/account_manager.py b/tools/account_manager.py index 63d3986..4befc05 100644 --- a/tools/account_manager.py +++ b/tools/account_manager.py @@ -19,6 +19,7 @@ from config.blazemeter import ACCOUNTS_ENDPOINT, TOOLS_PREFIX, SUPPORT_MESSAGE from config.token import BzmToken +from config.runtime import AppRuntime from formatters.account import format_accounts from models.manager import Manager from models.result import BaseResult @@ -76,7 +77,7 @@ async def list(self, limit: int = 50, offset: int = 0) -> BaseResult: params=parameters ) -def register(mcp, token: Optional[BzmToken]) -> None: +def register(mcp, runtime: AppRuntime) -> None: @mcp.tool( name=f"{TOOLS_PREFIX}_account", description=""" @@ -97,7 +98,7 @@ def register(mcp, token: Optional[BzmToken]) -> None: """ ) async def account(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - account_manager = AccountManager(token, ctx) + account_manager = AccountManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/billing_manager.py b/tools/billing_manager.py index 9575637..afde982 100644 --- a/tools/billing_manager.py +++ b/tools/billing_manager.py @@ -20,6 +20,7 @@ from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE from config.token import BzmToken +from config.runtime import AppRuntime from models.manager import Manager from models.result import BaseResult from tools.billing_utils import calculate_test_cost @@ -46,7 +47,7 @@ async def calculate_cost_from_config(self, args: Dict) -> BaseResult: ]) -def register(mcp, token: Optional[BzmToken]) -> None: +def register(mcp, runtime: AppRuntime) -> None: @mcp.tool( name=f"{TOOLS_PREFIX}_billing", description=""" @@ -87,7 +88,7 @@ def register(mcp, token: Optional[BzmToken]) -> None: """ ) async def billing(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - billing_manager = BillingManager(token, ctx) + billing_manager = BillingManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/execution_manager.py b/tools/execution_manager.py index 1392f43..b84a15a 100644 --- a/tools/execution_manager.py +++ b/tools/execution_manager.py @@ -20,6 +20,7 @@ from config.blazemeter import TOOLS_PREFIX, EXECUTIONS_ENDPOINT, SUPPORT_MESSAGE from config.token import BzmToken +from config.runtime import AppRuntime from formatters.execution import format_executions, format_executions_detailed, format_executions_status from models.manager import Manager from models.result import BaseResult @@ -430,7 +431,7 @@ def _get_analysis_context_message(is_ready: bool, status_message: str) -> str: ) -def register(mcp, token: Optional[BzmToken]): +def register(mcp, runtime: AppRuntime): @mcp.tool( name=f"{TOOLS_PREFIX}_execution", description=""" @@ -496,8 +497,8 @@ def register(mcp, token: Optional[BzmToken]): """ ) async def execution(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - execution_manager = ExecutionManager(token, ctx) - report_manager = ReportManager(token, ctx) + execution_manager = ExecutionManager(runtime.auth.get_token(ctx), ctx) + report_manager = ReportManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/help_manager.py b/tools/help_manager.py index 1a911d3..53d8823 100644 --- a/tools/help_manager.py +++ b/tools/help_manager.py @@ -25,6 +25,7 @@ from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE, \ HELP_INDEX_URL, HELP_TOC_URL, HELP_BASE_CONTENT_URL from config.token import BzmToken +from config.runtime import AppRuntime from formatters.help import format_help_info from models.manager import Manager from models.result import BaseResult @@ -255,7 +256,7 @@ async def read_help_info(self, category_id: str, subcategory_id: str, help_id_li ) -def register(mcp, token: Optional[BzmToken]): +def register(mcp, runtime: AppRuntime): @mcp.tool( name=f"{TOOLS_PREFIX}_help", description=""" @@ -290,7 +291,7 @@ async def help_main( if args is None: args = {} - help_manager = HelpManager(token, ctx) + help_manager = HelpManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/project_manager.py b/tools/project_manager.py index 89cb4f4..3b59e0e 100644 --- a/tools/project_manager.py +++ b/tools/project_manager.py @@ -20,6 +20,7 @@ from config.blazemeter import TOOLS_PREFIX, PROJECTS_ENDPOINT from config.token import BzmToken +from config.runtime import AppRuntime from formatters.project import format_projects from models.manager import Manager from models.result import BaseResult @@ -83,7 +84,7 @@ async def list(self, workspace_id: Optional[int], limit: int = 50, offset: int = params=parameters ) -def register(mcp, token: Optional[BzmToken]): +def register(mcp, runtime: AppRuntime): @mcp.tool( name=f"{TOOLS_PREFIX}_project", description=""" @@ -105,7 +106,7 @@ def register(mcp, token: Optional[BzmToken]): """ ) async def project(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - project_manager = ProjectManager(token, ctx) + project_manager = ProjectManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/skills_manager.py b/tools/skills_manager.py index d2a60c2..8fbdc14 100644 --- a/tools/skills_manager.py +++ b/tools/skills_manager.py @@ -23,6 +23,7 @@ from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE from config.token import BzmToken +from config.runtime import AppRuntime from models.manager import Manager from models.result import BaseResult from telemetry import run_tool @@ -157,7 +158,7 @@ async def read_skill_resource_uri_list(skill_uri_list: Optional[List[str]]) -> B ) -def register(mcp, token: Optional[BzmToken]): +def register(mcp, runtime: AppRuntime): @mcp.resource("blazemeter-skill-{skill_name}://{path}") def universal_skills_handler(skill_name: str, path: str) -> str: path = unquote(path) @@ -203,7 +204,7 @@ async def skills( if args is None: args = {} - skills_manager = SkillsManager(token, ctx) + skills_manager = SkillsManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/test_manager.py b/tools/test_manager.py index 2c3db48..196277d 100644 --- a/tools/test_manager.py +++ b/tools/test_manager.py @@ -27,6 +27,7 @@ from config.path_mapper import PathMapperFactory from config.security import detect_sensitive_upload_path_reason from config.token import BzmToken +from config.runtime import AppRuntime from formatters.failure_criteria_labels import failure_criteria_meta_payload from formatters.test import format_tests from models.failure_criteria import ( @@ -549,7 +550,7 @@ async def failure_criteria_meta(self, args: Dict[str, Any]) -> BaseResult: return BaseResult(result=[failure_criteria_meta_payload()]) -def register(mcp, token: Optional[BzmToken]): +def register(mcp, runtime: AppRuntime): @mcp.tool( name=f"{TOOLS_PREFIX}_tests", description=""" @@ -653,7 +654,7 @@ def register(mcp, token: Optional[BzmToken]): """, ) async def tests(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - test_manager = TestManager(token, ctx) + test_manager = TestManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/user_manager.py b/tools/user_manager.py index 4a28a06..e314b12 100644 --- a/tools/user_manager.py +++ b/tools/user_manager.py @@ -21,6 +21,7 @@ from config.blazemeter import TOOLS_PREFIX, USER_ENDPOINT from config.token import BzmToken +from config.runtime import AppRuntime from formatters.user import format_users from models.manager import Manager from models.result import BaseResult @@ -42,7 +43,7 @@ async def read(self) -> BaseResult: ) -def register(mcp, token: Optional[BzmToken]): +def register(mcp, runtime: AppRuntime): @mcp.tool( name=f"{TOOLS_PREFIX}_user", description=""" @@ -60,7 +61,7 @@ async def user( ctx: Context = Field(description="Context object providing access to MCP capabilities") ) -> BaseResult: - user_manager = UserManager(token, ctx) + user_manager = UserManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: diff --git a/tools/workspace_manager.py b/tools/workspace_manager.py index 686c049..8a8a497 100644 --- a/tools/workspace_manager.py +++ b/tools/workspace_manager.py @@ -21,6 +21,7 @@ from config.blazemeter import WORKSPACES_ENDPOINT, TOOLS_PREFIX from config.token import BzmToken +from config.runtime import AppRuntime from formatters.workspace import format_workspaces, format_workspaces_detailed, format_workspaces_locations from models.manager import Manager from models.result import BaseResult @@ -109,7 +110,7 @@ async def read_locations(self, workspace_id: Optional[int], purpose: str = "load else: return locations_result -def register(mcp, token: Optional[BzmToken]): +def register(mcp, runtime: AppRuntime): @mcp.tool( name=f"{TOOLS_PREFIX}_workspaces", description=""" @@ -138,7 +139,7 @@ async def workspace( ctx: Context = Field(description="Context object providing access to MCP capabilities") ) -> BaseResult: - workspace_manager = WorkspaceManager(token, ctx) + workspace_manager = WorkspaceManager(runtime.auth.get_token(ctx), ctx) async def _dispatch(): match action: From 8c606b5c02d31747d43854adbb5dcc886daf8239 Mon Sep 17 00:00:00 2001 From: Diego Ferrand Date: Wed, 12 Aug 2026 10:31:53 -0300 Subject: [PATCH 2/6] Hosted mcp (#97) * Add otel documentation * Add per-request Bearer auth for hosted HTTP MCP * Match original arquitecture plan naming and refactors. * Integrate streamable HTTP transport with Bearer auth wiring * Add hosted HTTP storage port, health probes, and Cloud Run image * Moved to new repo * Remove hosted Docker image; document HTTP env vars on :latest * Reverted dockerfile modifications to previous version * Refactor documentation to split core from hosted docs. --- README.md | 14 +-- config/auth.py | 22 +++++ config/runtime.py | 23 ++++- config/storage.py | 141 +++++++++++++++++++++++++++ docs/hosted-http.md | 70 +++++++++++++ main.py | 8 +- tests/test_http_auth.py | 25 ++++- tests/test_main_transport.py | 15 +++ tests/test_storage.py | 118 ++++++++++++++++++++++ tests/test_upload_assets_security.py | 6 +- tools/test_manager.py | 51 ++++++---- 11 files changed, 454 insertions(+), 39 deletions(-) create mode 100644 config/storage.py create mode 100644 docs/hosted-http.md create mode 100644 tests/test_storage.py diff --git a/README.md b/README.md index 581db9a..8969be3 100644 --- a/README.md +++ b/README.md @@ -110,21 +110,15 @@ After installing, set `BLAZEMETER_API_KEY` to your `api-key.json` path in your c --- -**Hosted HTTP (streamable-http) Client Configuration** +**Hosted HTTP (streamable-http)** -Run a shared server that authenticates each client via `Authorization: Bearer`: - -```bash -uv run python main.py --mcp --transport streamable-http //PLACEHOLDER -``` - -Configure the MCP client with the server URL and your BlazeMeter API key as Bearer credentials (`id:secret` or base64 of `id:secret`): +Connect to the shared hosted server with Bearer auth. Production URL: `https://mcp.blazemeter.com/mcp` ```json { "mcpServers": { "BlazeMeter MCP": { - "url": "http://localhost:8000/mcp", + "url": "https://mcp.blazemeter.com/mcp", "headers": { "Authorization": "Bearer :" } @@ -134,7 +128,7 @@ Configure the MCP client with the server URL and your BlazeMeter API key as Bear ``` > [!NOTE] -> Over HTTP, credentials are resolved per request from the `Authorization` header. Invalid or missing Bearer credentials return `401` before any tool runs. Well-formed but wrong API keys fail later inside BlazeMeter API calls (same as stdio). Stdio transport uses `api-key.json` / env / Docker secrets. +> For local/operator HTTP setup, env vars, auth behavior, and MVP limits (including `upload_assets`), see [docs/hosted-http.md](docs/hosted-http.md). --- diff --git a/config/auth.py b/config/auth.py index b028011..373abcb 100644 --- a/config/auth.py +++ b/config/auth.py @@ -26,6 +26,9 @@ BZM_TOKEN_STATE_ATTR = "token" +# Unauthenticated probe paths for orchestrators / load balancers. +HEALTH_PATHS = frozenset({"/health", "/healthz"}) + class AuthError(Exception): """Raised when Authorization cannot be parsed into credentials.""" @@ -98,6 +101,11 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: await self.app(scope, receive, send) return + path = scope.get("path", "") or "" + if path in HEALTH_PATHS: + await self.app(scope, receive, send) + return + request = Request(scope, receive) try: token = parse_authorization_header(request.headers.get("authorization")) @@ -114,11 +122,25 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: await self.app(scope, receive, send) +def register_health_routes(mcp: FastMCP) -> None: + """Register unauthenticated health probes on the FastMCP ASGI app.""" + + @mcp.custom_route("/health", methods=["GET"]) + async def health(_request: Request) -> JSONResponse: + return JSONResponse({"status": "ok"}) + + @mcp.custom_route("/healthz", methods=["GET"]) + async def healthz(_request: Request) -> JSONResponse: + return JSONResponse({"status": "ok"}) + + def run_streamable_http(mcp: FastMCP) -> None: """Serve FastMCP over streamable HTTP with Bearer auth middleware.""" import anyio import uvicorn + register_health_routes(mcp) + async def _serve() -> None: app = BearerAuthMiddleware(mcp.streamable_http_app()) config = uvicorn.Config( diff --git a/config/runtime.py b/config/runtime.py index 956d181..90f8e11 100644 --- a/config/runtime.py +++ b/config/runtime.py @@ -17,6 +17,7 @@ from typing import Literal, Optional from config.auth import AuthPort, HttpAuthProvider, StdioAuthProvider +from config.storage import StoragePort, build_storage from config.token import BzmToken Transport = Literal["stdio", "streamable-http"] @@ -28,20 +29,32 @@ class AppRuntime: transport: Transport auth: AuthPort + storage: StoragePort def build_runtime( transport: Transport, startup_token: Optional[BzmToken] = None, + storage_backend: Optional[str] = None, ) -> AppRuntime: """ - Compose auth for the selected transport. + Compose auth and storage for the selected transport. - - stdio: use process-lifetime ``startup_token`` (from env / api-key.json / Docker). - - streamable-http: resolve credentials per request via Bearer middleware + HttpAuthProvider. + - stdio: process-lifetime ``startup_token``; local/memory file storage by default. + - streamable-http: Bearer middleware + HttpAuthProvider; HttpStorageClient + (local paths rejected) regardless of BZM_STORAGE_BACKEND for MVP. """ + storage = build_storage(transport, backend=storage_backend) if transport == "stdio": - return AppRuntime(transport=transport, auth=StdioAuthProvider(startup_token)) + return AppRuntime( + transport=transport, + auth=StdioAuthProvider(startup_token), + storage=storage, + ) if transport == "streamable-http": - return AppRuntime(transport=transport, auth=HttpAuthProvider()) + return AppRuntime( + transport=transport, + auth=HttpAuthProvider(), + storage=storage, + ) raise ValueError(f"Unknown transport: {transport}") diff --git a/config/storage.py b/config/storage.py new file mode 100644 index 0000000..a3887de --- /dev/null +++ b/config/storage.py @@ -0,0 +1,141 @@ +""" +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 os +from pathlib import Path +from typing import List, Literal, Optional, Protocol, runtime_checkable + +from config.path_mapper import PathMapperFactory, PathMappingStrategy + +StorageBackend = Literal["memory", "http"] + +HOSTED_FILE_ACCESS_MESSAGE = ( + "Local file lookup and upload are not supported on the hosted MCP server. " + "Use a local stdio/Docker MCP installation for upload_assets, or wait for " + "Phase 2 remote Storage." +) + + +class StorageNotSupportedError(NotImplementedError): + """Raised when a storage backend cannot fulfill a file operation.""" + + +@runtime_checkable +class StoragePort(Protocol): + """ + Contract for resolving and reading files used by MCP tools (e.g. upload_assets). + + MVP backends: + - memory/local: process-local disk via path mapping (stdio / local Docker) + - http: fail-closed stub for hosted streamable-http (no local disk) + """ + + def map_paths(self, file_paths: List[str]) -> List[str]: + ... + + def exists(self, path: str) -> bool: + ... + + def is_file(self, path: str) -> bool: + ... + + def read_bytes(self, path: str) -> bytes: + ... + + def basename(self, path: str) -> str: + ... + + +class LocalStorageClient: + """ + Process-local file access (BZM_STORAGE_BACKEND=memory for MVP). + + Uses the existing path mapper so Docker volume mounts keep working. + No external Storage Service — suitable for single-instance / stdio MVP. + """ + + def __init__(self, path_mapper: Optional[PathMappingStrategy] = None): + self._path_mapper = path_mapper or PathMapperFactory.create_strategy() + + def map_paths(self, file_paths: List[str]) -> List[str]: + return self._path_mapper.map_paths(file_paths) + + def exists(self, path: str) -> bool: + return os.path.exists(path) + + def is_file(self, path: str) -> bool: + return os.path.isfile(path) + + def read_bytes(self, path: str) -> bytes: + with open(path, "rb") as handle: + return handle.read() + + def basename(self, path: str) -> str: + return Path(path).name + + +class HttpStorageClient: + """ + Fail-closed file storage for hosted HTTP (and future remote Storage Service). + + Every file lookup/upload path raises so local client paths cannot be used + against a shared hosted instance. Phase 2 can replace these stubs with + real remote Storage API calls. + """ + + def map_paths(self, file_paths: List[str]) -> List[str]: + raise StorageNotSupportedError(HOSTED_FILE_ACCESS_MESSAGE) + + def exists(self, path: str) -> bool: + raise StorageNotSupportedError(HOSTED_FILE_ACCESS_MESSAGE) + + def is_file(self, path: str) -> bool: + raise StorageNotSupportedError(HOSTED_FILE_ACCESS_MESSAGE) + + def read_bytes(self, path: str) -> bytes: + raise StorageNotSupportedError(HOSTED_FILE_ACCESS_MESSAGE) + + def basename(self, path: str) -> str: + raise StorageNotSupportedError(HOSTED_FILE_ACCESS_MESSAGE) + + +def resolve_storage_backend(raw: Optional[str] = None) -> StorageBackend: + """Resolve BZM_STORAGE_BACKEND (default: memory).""" + candidate = (raw if raw is not None else os.getenv("BZM_STORAGE_BACKEND", "memory")).strip().lower() + if not candidate: + return "memory" + if candidate not in ("memory", "http"): + raise ValueError( + f"Invalid BZM_STORAGE_BACKEND '{candidate}'. Valid values: memory, http." + ) + return candidate # type: ignore[return-value] + + +def build_storage( + transport: Literal["stdio", "streamable-http"], + backend: Optional[str] = None, +) -> StoragePort: + """ + Select storage for the process. + + Hosted streamable-http always uses HttpStorageClient so local paths are + rejected. Stdio uses LocalStorageClient when backend is memory (MVP default). + """ + resolved = resolve_storage_backend(backend) + if transport == "streamable-http" or resolved == "http": + return HttpStorageClient() + return LocalStorageClient() diff --git a/docs/hosted-http.md b/docs/hosted-http.md new file mode 100644 index 0000000..83ffc95 --- /dev/null +++ b/docs/hosted-http.md @@ -0,0 +1,70 @@ +# Hosted HTTP (streamable-http) + +Operator and advanced client guide for running BlazeMeter MCP over HTTP. For the standard local install (binary, uvx, Docker stdio), see the [README](../README.md). + +## Hosted endpoint (clients) + +Production URL: + +`https://mcp.blazemeter.com/mcp` + +Configure the MCP client with that URL and your BlazeMeter API key as Bearer credentials (`id:secret` or base64 of `id:secret`): + +```json +{ + "mcpServers": { + "BlazeMeter MCP": { + "url": "https://mcp.blazemeter.com/mcp", + "headers": { + "Authorization": "Bearer :" + } + } + } +} +``` + +For a locally run server, use `"url": "http://localhost:8000/mcp"` instead. + +### Auth behavior + +- Over HTTP, credentials are resolved **per request** from the `Authorization` header. +- Invalid or missing Bearer credentials return `401` before any tool runs. +- Well-formed but wrong API keys fail later inside BlazeMeter API calls (same as stdio). +- Stdio / local Docker transport uses `api-key.json` / env / Docker secrets instead of Bearer auth. + +## Local / operator run + +Transport resolution precedence: **CLI `--mcp` > `BZM_MCP_TRANSPORT` > stdio**. + +```bash +# From source +uv run python main.py --mcp http +# or +BZM_MCP_TRANSPORT=http FASTMCP_HOST=0.0.0.0 FASTMCP_PORT=8000 uv run python main.py --mcp + +# Container image (:latest is stdio by default; pass hosted HTTP env vars) +docker run --rm -p 8000:8000 \ + -e BZM_MCP_TRANSPORT=http \ + -e FASTMCP_HOST=0.0.0.0 \ + -e FASTMCP_PORT=8000 \ + -e FASTMCP_STREAMABLE_HTTP_PATH=/mcp \ + -e BZM_STORAGE_BACKEND=memory \ + ghcr.io/blazemeter/bzm-mcp:latest +``` + +### Environment variables + +| Variable | Description | Default | +|----------|-------------|---------| +| `BZM_MCP_TRANSPORT` | Logical transport: `stdio`, `http`, or `docker` | `stdio` | +| `FASTMCP_HOST` | Bind address (HTTP only) | `127.0.0.1` | +| `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` | +| `BZM_STORAGE_BACKEND` | `memory` or `http` | `memory` | + +On streamable-http, local file paths are always rejected regardless of `BZM_STORAGE_BACKEND` (hosted fail-closed storage). + +## Hosted MVP limitations + +- In-memory / fail-closed storage: no local disk access on the shared hosted server. +- `upload_assets` and other local file lookup/upload paths are rejected. Use a local stdio or Docker MCP installation for those workflows, or wait for remote Storage (Phase 2). diff --git a/main.py b/main.py index ec0b7c3..dc79815 100644 --- a/main.py +++ b/main.py @@ -428,7 +428,13 @@ def build_mcp_server( streamable_http_path = "/mcp" if transport == "http": host = os.getenv("FASTMCP_HOST", "127.0.0.1").strip() or "127.0.0.1" - port = int(os.getenv("FASTMCP_PORT", "8000").strip() or "8000") + # Cloud Run injects PORT; prefer FASTMCP_PORT when set, else PORT, else 8000. + port_raw = ( + os.getenv("FASTMCP_PORT") + or os.getenv("PORT") + or "8000" + ).strip() or "8000" + port = int(port_raw) streamable_http_path = os.getenv("FASTMCP_STREAMABLE_HTTP_PATH", "/mcp").strip() or "/mcp" # docker and stdio share process-lifetime credentials; http uses Bearer per request. diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py index 9fe95c8..6ffc501 100644 --- a/tests/test_http_auth.py +++ b/tests/test_http_auth.py @@ -33,6 +33,7 @@ parse_authorization_header, ) from config.runtime import build_runtime +from config.storage import HttpStorageClient, LocalStorageClient from config.token import BzmToken, BzmTokenError @@ -153,13 +154,35 @@ async def ok(_request: Request): client = TestClient(app) assert client.options("/mcp").status_code == 200 + def test_health_bypasses_auth(self): + async def health(_request: Request): + return JSONResponse({"status": "ok"}) + + app = BearerAuthMiddleware( + Starlette( + routes=[ + Route("/health", endpoint=health, methods=["GET"]), + Route("/healthz", endpoint=health, methods=["GET"]), + ] + ) + ) + client = TestClient(app) + assert client.get("/health").status_code == 200 + assert client.get("/health").json()["status"] == "ok" + assert client.get("/healthz").status_code == 200 + class TestBuildRuntime: - def test_build_runtime_stdio_and_http(self): + def test_build_runtime_stdio_and_http(self, monkeypatch): + monkeypatch.delenv("MCP_DOCKER", raising=False) + monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + stdio = build_runtime("stdio") assert stdio.transport == "stdio" assert isinstance(stdio.auth, StdioAuthProvider) + assert isinstance(stdio.storage, LocalStorageClient) http = build_runtime("streamable-http") assert http.transport == "streamable-http" assert isinstance(http.auth, HttpAuthProvider) + assert isinstance(http.storage, HttpStorageClient) diff --git a/tests/test_main_transport.py b/tests/test_main_transport.py index ed51904..b7a3cb2 100644 --- a/tests/test_main_transport.py +++ b/tests/test_main_transport.py @@ -3,6 +3,7 @@ import main from config.auth import HttpAuthProvider, StdioAuthProvider from config.runtime import AppRuntime +from config.storage import HttpStorageClient, LocalStorageClient class _DummyFastMCP: @@ -64,6 +65,16 @@ def test_http_uses_env_settings_and_stateful_http(self, monkeypatch): assert mcp.kwargs["streamable_http_path"] == "/custom-mcp" assert mcp.kwargs["stateless_http"] is False + def test_http_falls_back_to_cloud_run_port(self, monkeypatch): + _patch_mcp_server_dependencies(monkeypatch) + monkeypatch.delenv("FASTMCP_PORT", raising=False) + monkeypatch.setenv("PORT", "8080") + monkeypatch.setenv("FASTMCP_HOST", "0.0.0.0") + + mcp, _ = main.build_mcp_server(transport="http") + + assert mcp.kwargs["port"] == 8080 + class TestBuildMcpServerTransportMapping: def test_transport_mapping_keeps_docker_stdio_and_http_streamable(self, monkeypatch): @@ -94,6 +105,7 @@ def capture_register(mcp, runtime): assert isinstance(runtime, AppRuntime) assert runtime.transport == "streamable-http" assert isinstance(runtime.auth, HttpAuthProvider) + assert isinstance(runtime.storage, HttpStorageClient) def test_stdio_and_docker_register_stdio_auth_provider(self, monkeypatch): captured = {} @@ -105,6 +117,8 @@ def capture_register(mcp, runtime): _patch_mcp_server_dependencies(monkeypatch) monkeypatch.setattr(main, "get_token", lambda: token) monkeypatch.setattr(main, "register_tools", capture_register) + monkeypatch.delenv("MCP_DOCKER", raising=False) + monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) main.build_mcp_server(transport="stdio") main.build_mcp_server(transport="docker") @@ -113,6 +127,7 @@ def capture_register(mcp, runtime): for runtime in captured["runtimes"]: assert isinstance(runtime.auth, StdioAuthProvider) assert runtime.auth.get_token(ctx=None) is token + assert isinstance(runtime.storage, LocalStorageClient) class TestRunTransportDispatch: diff --git a/tests/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000..6d0dd79 --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,118 @@ +""" +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 pytest + +from config.runtime import build_runtime +from config.storage import ( + HOSTED_FILE_ACCESS_MESSAGE, + HttpStorageClient, + LocalStorageClient, + StorageNotSupportedError, + build_storage, + resolve_storage_backend, +) +from tools.test_manager import TestManager + + +class TestResolveStorageBackend: + def test_default_memory(self, monkeypatch): + monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + assert resolve_storage_backend() == "memory" + + def test_env_http(self, monkeypatch): + monkeypatch.setenv("BZM_STORAGE_BACKEND", "http") + assert resolve_storage_backend() == "http" + + def test_invalid_raises(self): + with pytest.raises(ValueError, match="BZM_STORAGE_BACKEND"): + resolve_storage_backend("s3") + + +class TestBuildStorage: + def test_stdio_memory_uses_local(self, monkeypatch): + monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + monkeypatch.delenv("MCP_DOCKER", raising=False) + storage = build_storage("stdio") + assert isinstance(storage, LocalStorageClient) + + def test_streamable_http_uses_http_client(self, monkeypatch): + monkeypatch.setenv("BZM_STORAGE_BACKEND", "memory") + storage = build_storage("streamable-http") + assert isinstance(storage, HttpStorageClient) + + def test_explicit_http_backend_on_stdio(self): + storage = build_storage("stdio", backend="http") + assert isinstance(storage, HttpStorageClient) + + +class TestLocalStorageClient: + def test_read_roundtrip(self, tmp_path, monkeypatch): + monkeypatch.delenv("MCP_DOCKER", raising=False) + path = tmp_path / "asset.jmx" + path.write_bytes(b"") + client = LocalStorageClient() + mapped = client.map_paths([str(path)]) + assert mapped == [str(path)] + assert client.exists(str(path)) + assert client.is_file(str(path)) + assert client.read_bytes(str(path)) == b"" + assert client.basename(str(path)) == "asset.jmx" + + +class TestHttpStorageClient: + def test_all_file_methods_raise(self): + client = HttpStorageClient() + with pytest.raises(StorageNotSupportedError, match="hosted MCP") as exc_info: + client.map_paths(["/tmp/a.jmx"]) + assert HOSTED_FILE_ACCESS_MESSAGE in str(exc_info.value) + with pytest.raises(StorageNotSupportedError): + client.exists("/tmp/a.jmx") + with pytest.raises(StorageNotSupportedError): + client.is_file("/tmp/a.jmx") + with pytest.raises(StorageNotSupportedError): + client.read_bytes("/tmp/a.jmx") + with pytest.raises(StorageNotSupportedError): + client.basename("/tmp/a.jmx") + + +class TestRuntimeStorageWiring: + def test_http_runtime_gets_http_storage(self): + runtime = build_runtime("streamable-http") + assert isinstance(runtime.storage, HttpStorageClient) + + def test_stdio_runtime_gets_local_storage(self, monkeypatch): + monkeypatch.delenv("MCP_DOCKER", raising=False) + monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + runtime = build_runtime("stdio") + assert isinstance(runtime.storage, LocalStorageClient) + + +class TestUploadAssetsHostedRejection: + def test_upload_assets_returns_clear_error_on_http_storage(self): + manager = TestManager(token=None, ctx=None, storage=HttpStorageClient()) + + async def _fake_read(_test_id): + from models.result import BaseResult + return BaseResult(result=[{"id": 1}]) + + manager.read = _fake_read # type: ignore[method-assign] + result = asyncio.run( + manager.upload_assets(1, ["/tmp/demo.jmx"], main_script=None) + ) + assert "error" in result + assert "not supported on the hosted MCP" in result["error"] diff --git a/tests/test_upload_assets_security.py b/tests/test_upload_assets_security.py index b728e2c..042befc 100644 --- a/tests/test_upload_assets_security.py +++ b/tests/test_upload_assets_security.py @@ -14,8 +14,7 @@ limitations under the License. """ -from pathlib import Path - +from config.storage import LocalStorageClient from tools.test_manager import TestManager as UploadAssetsManager @@ -56,7 +55,8 @@ def test_validate_files_classifies_valid_invalid_and_blocked(self, tmp_path): invalid_files = [] blocked_files = [] - UploadAssetsManager._validate_files( + manager = UploadAssetsManager(token=None, ctx=None, storage=LocalStorageClient()) + manager._validate_files( [str(safe_file), str(env_file), str(missing_file)], valid_files, invalid_files, diff --git a/tools/test_manager.py b/tools/test_manager.py index 196277d..46a55de 100644 --- a/tools/test_manager.py +++ b/tools/test_manager.py @@ -15,7 +15,6 @@ """ import asyncio import logging -import os from pathlib import Path from typing import Any, Dict from typing import Optional, List @@ -24,8 +23,8 @@ from mcp.server.fastmcp import Context from config.blazemeter import TESTS_ENDPOINT, TOOLS_PREFIX -from config.path_mapper import PathMapperFactory from config.security import detect_sensitive_upload_path_reason +from config.storage import LocalStorageClient, StorageNotSupportedError, StoragePort from config.token import BzmToken from config.runtime import AppRuntime from formatters.failure_criteria_labels import failure_criteria_meta_payload @@ -52,9 +51,14 @@ class TestManager(Manager): __test__ = False - def __init__(self, token: Optional[BzmToken], ctx: Context): + def __init__( + self, + token: Optional[BzmToken], + ctx: Context, + storage: Optional[StoragePort] = None, + ): super().__init__(token, ctx) - self.path_mapper = PathMapperFactory.create_strategy() + self.storage = storage or LocalStorageClient() async def read(self, test_id: Optional[int]) -> BaseResult: if not isinstance(test_id, int) or test_id < 1: @@ -142,9 +146,8 @@ async def delete(self, test_id: Optional[int]) -> BaseResult: def _detect_sensitive_path_reason(cls, file_path: str) -> Optional[str]: return detect_sensitive_upload_path_reason(file_path) - @classmethod def _validate_files( - cls, + self, file_paths: List[str], valid_files: List[str], invalid_files: List[str], @@ -159,7 +162,7 @@ def _validate_files( # locations is an administrative responsibility of the UNC share owners/administrators. for file_path in file_paths: logger.debug(f"Checking file: {file_path}") - sensitive_reason = cls._detect_sensitive_path_reason(file_path) + sensitive_reason = self._detect_sensitive_path_reason(file_path) if sensitive_reason: logger.warning( f"Blocked sensitive file path: {file_path} ({sensitive_reason})" @@ -171,7 +174,7 @@ def _validate_files( } ) continue - if os.path.exists(file_path) and os.path.isfile(file_path): + if self.storage.exists(file_path) and self.storage.is_file(file_path): logger.debug(f"File exists: {file_path}") valid_files.append(file_path) else: @@ -218,12 +221,19 @@ async def upload_assets( logger.debug(f"Original file paths: {file_paths}") logger.debug(f"Main script: {main_script}") - mapped_file_paths = self.path_mapper.map_paths(file_paths) + try: + mapped_file_paths = self.storage.map_paths(file_paths) + except StorageNotSupportedError as exc: + return {"error": str(exc)} + logger.debug(f"Mapped file paths: {mapped_file_paths}") mapped_main_script = None if main_script: - mapped_main_script_list = self.path_mapper.map_paths([main_script]) + try: + mapped_main_script_list = self.storage.map_paths([main_script]) + except StorageNotSupportedError as exc: + return {"error": str(exc)} mapped_main_script = ( mapped_main_script_list[0] if mapped_main_script_list else None ) @@ -233,9 +243,12 @@ async def upload_assets( invalid_files = [] blocked_files = [] - self._validate_files( - mapped_file_paths, valid_files, invalid_files, blocked_files - ) + try: + self._validate_files( + mapped_file_paths, valid_files, invalid_files, blocked_files + ) + except StorageNotSupportedError as exc: + return {"error": str(exc)} logger.debug(f"Valid files: {valid_files}") logger.debug(f"Invalid files: {invalid_files}") @@ -285,13 +298,11 @@ async def upload_assets( async def _upload_single_file(self, test_id: int, file_path: str) -> BaseResult: logger.debug(f"Uploading single file: {file_path} to test: {test_id}") try: - file_path_obj = Path(file_path) - file_name = file_path_obj.name + file_name = self.storage.basename(file_path) logger.debug(f"File name: {file_name}") - with open(file_path, "rb") as file: - file_content = file.read() + file_content = self.storage.read_bytes(file_path) logger.debug(f"File size: {len(file_content)} bytes") @@ -306,6 +317,8 @@ async def _upload_single_file(self, test_id: int, file_path: str) -> BaseResult: return result + except StorageNotSupportedError: + raise except Exception as e: logger.error(f"Exception in _upload_single_file: {e}") logger.error(f"Traceback: {format_sanitized_traceback(e)}") @@ -315,7 +328,7 @@ async def _update_test_configuration( self, test_id: int, main_script_path: str ) -> BaseResult: try: - file_name = Path(main_script_path).name + file_name = self.storage.basename(main_script_path) config_update = { "configuration": { "filename": file_name, @@ -654,7 +667,7 @@ def register(mcp, runtime: AppRuntime): """, ) async def tests(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - test_manager = TestManager(runtime.auth.get_token(ctx), ctx) + test_manager = TestManager(runtime.auth.get_token(ctx), ctx, runtime.storage) async def _dispatch(): match action: From 0211dea5055b2a12a2717243a99d6791ed30f88f Mon Sep 17 00:00:00 2001 From: alejandroaires Date: Fri, 14 Aug 2026 17:11:40 -0300 Subject: [PATCH 3/6] Mob 50444 storage service foundation integration (#100) * added storage api sync with streamable-http runtime * update api url path * add fileaccess port * Tests update * updated test_help_live_href_rendering.py * remove storage_backend unused variable * remove legacy compatibility aliases * rename BZM_STORAGE_BACKEND to BZM_STORAGE_STRATEGY across codebase --- config/file_access.py | 138 +++++++++++++++++ config/runtime.py | 40 +++-- config/storage.py | 214 ++++++++++++++++++++++++++- docs/hosted-http.md | 6 +- tests/test_file_access.py | 94 ++++++++++++ tests/test_hierarchy_and_consent.py | 9 +- tests/test_http_auth.py | 27 +++- tests/test_main_transport.py | 12 +- tests/test_storage.py | 34 +++-- tests/test_upload_assets_security.py | 11 +- tools/test_manager.py | 65 ++++---- 11 files changed, 579 insertions(+), 71 deletions(-) create mode 100644 config/file_access.py create mode 100644 tests/test_file_access.py diff --git a/config/file_access.py b/config/file_access.py new file mode 100644 index 0000000..f925ba9 --- /dev/null +++ b/config/file_access.py @@ -0,0 +1,138 @@ +""" +File access abstractions for upload-oriented tools. +""" +from __future__ import annotations + +import os +from abc import ABC, abstractmethod +from pathlib import Path + +from config.storage import SessionScope + + +class FileAccessPort(ABC): + """Abstraction for file path mapping and file content reads.""" + + @abstractmethod + def map_paths(self, file_paths: list[str], scope: SessionScope | None = None) -> list[str]: + raise NotImplementedError + + @abstractmethod + def exists(self, file_path: str, scope: SessionScope | None = None) -> bool: + raise NotImplementedError + + @abstractmethod + def is_file(self, file_path: str, scope: SessionScope | None = None) -> bool: + raise NotImplementedError + + @abstractmethod + def read_bytes(self, file_path: str, scope: SessionScope | None = None) -> bytes: + raise NotImplementedError + + +class LocalPathFileSource(FileAccessPort): + """Use filesystem paths as provided by the caller.""" + + def map_paths(self, file_paths: list[str], scope: SessionScope | None = None) -> list[str]: + return file_paths + + def exists(self, file_path: str, scope: SessionScope | None = None) -> bool: + return os.path.exists(file_path) + + def is_file(self, file_path: str, scope: SessionScope | None = None) -> bool: + return os.path.isfile(file_path) + + def read_bytes(self, file_path: str, scope: SessionScope | None = None) -> bytes: + return Path(file_path).read_bytes() + + +class DockerMappedFileSource(LocalPathFileSource): + """ + Map host paths into the mounted container path and then read locally. + + This mirrors the old path mapper behavior for Docker stdio mode. + """ + + def __init__( + self, + source_working_directory: str, + container_working_directory: str = "/home/bzm-mcp/working_directory", + ) -> None: + self._source_working_directory = Path(source_working_directory).resolve() + self._container_working_directory = container_working_directory.rstrip("/\\") + + def map_paths(self, file_paths: list[str], scope: SessionScope | None = None) -> list[str]: + mapped_paths: list[str] = [] + for file_path in file_paths: + abs_file_path = Path(file_path).resolve() + try: + relative_path = abs_file_path.relative_to(self._source_working_directory) + mapped_path = ( + f"{self._container_working_directory}/{relative_path.as_posix()}" + ) + mapped_paths.append(mapped_path) + except ValueError: + mapped_paths.append(file_path) + return mapped_paths + + +class StorageFileSource(FileAccessPort): + """ + Mock placeholder for streamable-http file access. + + Real implementation will be provided in a future task where file-upload UI + mediates uploads and storage API integration. + """ + + def __init__(self, base_url: str) -> None: + self._base_url = base_url.rstrip("/") + + def ensure_available(self) -> None: + # Mocked source: no network checks for now. + return None + + def map_paths(self, file_paths: list[str], scope: SessionScope | None = None) -> list[str]: + # Keep paths untouched until backend file-source semantics are defined. + return file_paths + + def exists(self, file_path: str, scope: SessionScope | None = None) -> bool: + # Mocked behavior: storage-backed files are not available yet. + return False + + def is_file(self, file_path: str, scope: SessionScope | None = None) -> bool: + return False + + def read_bytes(self, file_path: str, scope: SessionScope | None = None) -> bytes: + raise NotImplementedError( + "StorageFileSource.read_bytes is not implemented yet. " + "Use file-upload UI flow until storage-backed file access is implemented." + ) + + +def build_file_access(transport: str) -> FileAccessPort: + """ + Build file-access implementation for the runtime transport. + + - Docker stdio uses path mapping (host -> mounted container path). + - Streamable HTTP uses storage API-backed file source. + - Other modes use local path access. + """ + if transport == "streamable-http": + base_url = os.getenv("BZM_STORAGE_API_BASE_URL", "").strip() + if not base_url: + raise ValueError( + "BZM_STORAGE_API_BASE_URL is required for streamable-http transport." + ) + return StorageFileSource(base_url=base_url) + + is_docker = os.getenv("MCP_DOCKER", "false").lower() == "true" + if transport == "stdio" and is_docker: + source_dir = os.getenv("SOURCE_WORKING_DIRECTORY") + if not source_dir: + raise ValueError( + "Working directory must be set in the Docker catalog configuration." + "Without volume mount, actions like upload assets will not work." + "Lack of volume mount results in missing SOURCE_WORKING_DIRECTORY environment variable" + ) + return DockerMappedFileSource(source_working_directory=source_dir) + return LocalPathFileSource() diff --git a/config/runtime.py b/config/runtime.py index 90f8e11..37f674a 100644 --- a/config/runtime.py +++ b/config/runtime.py @@ -15,9 +15,17 @@ """ from dataclasses import dataclass from typing import Literal, Optional +import os from config.auth import AuthPort, HttpAuthProvider, StdioAuthProvider -from config.storage import StoragePort, build_storage +from config.file_access import FileAccessPort, build_file_access +from config.storage import ( + DefaultSessionScopeResolver, + HttpSessionStorageProvider, + InMemorySessionStorageProvider, + SessionScopeResolverPort, + SessionStoragePort, +) from config.token import BzmToken Transport = Literal["stdio", "streamable-http"] @@ -29,32 +37,46 @@ class AppRuntime: transport: Transport auth: AuthPort - storage: StoragePort + storage: SessionStoragePort + file_access: FileAccessPort + scope_resolver: SessionScopeResolverPort def build_runtime( transport: Transport, startup_token: Optional[BzmToken] = None, - storage_backend: Optional[str] = None, ) -> AppRuntime: """ - Compose auth and storage for the selected transport. + Compose auth, file access, and session storage for the selected transport. - - stdio: process-lifetime ``startup_token``; local/memory file storage by default. - - streamable-http: Bearer middleware + HttpAuthProvider; HttpStorageClient - (local paths rejected) regardless of BZM_STORAGE_BACKEND for MVP. + - stdio: process-lifetime ``startup_token`` and in-memory session storage. + - streamable-http: request-scoped auth and storage API-backed partitions. """ - storage = build_storage(transport, backend=storage_backend) if transport == "stdio": return AppRuntime( transport=transport, auth=StdioAuthProvider(startup_token), - storage=storage, + storage=InMemorySessionStorageProvider(), + file_access=build_file_access(transport), + scope_resolver=DefaultSessionScopeResolver(), ) + if transport == "streamable-http": + storage_base_url = os.getenv("BZM_STORAGE_API_BASE_URL", "").strip() + if not storage_base_url: + raise ValueError( + "BZM_STORAGE_API_BASE_URL is required for streamable-http transport." + ) + storage: SessionStoragePort = HttpSessionStorageProvider( + base_url=storage_base_url, + ) + storage.ensure_available() return AppRuntime( transport=transport, auth=HttpAuthProvider(), storage=storage, + file_access=build_file_access(transport), + scope_resolver=DefaultSessionScopeResolver(), ) + raise ValueError(f"Unknown transport: {transport}") diff --git a/config/storage.py b/config/storage.py index a3887de..fd3e711 100644 --- a/config/storage.py +++ b/config/storage.py @@ -15,11 +15,18 @@ """ from __future__ import annotations +from abc import ABC, abstractmethod +from dataclasses import dataclass import os from pathlib import Path -from typing import List, Literal, Optional, Protocol, runtime_checkable +from typing import Any, List, Literal, Optional, Protocol, runtime_checkable +from urllib.parse import quote + +import httpx +from mcp.server.fastmcp import Context from config.path_mapper import PathMapperFactory, PathMappingStrategy +from config.token import BzmToken StorageBackend = Literal["memory", "http"] @@ -35,7 +42,7 @@ class StorageNotSupportedError(NotImplementedError): @runtime_checkable -class StoragePort(Protocol): +class FileStoragePort(Protocol): """ Contract for resolving and reading files used by MCP tools (e.g. upload_assets). @@ -62,7 +69,7 @@ def basename(self, path: str) -> str: class LocalStorageClient: """ - Process-local file access (BZM_STORAGE_BACKEND=memory for MVP). + Process-local file access (BZM_STORAGE_STRATEGY=memory for MVP). Uses the existing path mapper so Docker volume mounts keep working. No external Storage Service — suitable for single-instance / stdio MVP. @@ -114,13 +121,15 @@ def basename(self, path: str) -> str: def resolve_storage_backend(raw: Optional[str] = None) -> StorageBackend: - """Resolve BZM_STORAGE_BACKEND (default: memory).""" - candidate = (raw if raw is not None else os.getenv("BZM_STORAGE_BACKEND", "memory")).strip().lower() + """Resolve BZM_STORAGE_STRATEGY (default: memory).""" + candidate = ( + raw if raw is not None else os.getenv("BZM_STORAGE_STRATEGY", "memory") + ).strip().lower() if not candidate: return "memory" if candidate not in ("memory", "http"): raise ValueError( - f"Invalid BZM_STORAGE_BACKEND '{candidate}'. Valid values: memory, http." + f"Invalid BZM_STORAGE_STRATEGY '{candidate}'. Valid values: memory, http." ) return candidate # type: ignore[return-value] @@ -128,9 +137,9 @@ def resolve_storage_backend(raw: Optional[str] = None) -> StorageBackend: def build_storage( transport: Literal["stdio", "streamable-http"], backend: Optional[str] = None, -) -> StoragePort: +) -> FileStoragePort: """ - Select storage for the process. + Select file storage for the process. Hosted streamable-http always uses HttpStorageClient so local paths are rejected. Stdio uses LocalStorageClient when backend is memory (MVP default). @@ -139,3 +148,192 @@ def build_storage( if transport == "streamable-http" or resolved == "http": return HttpStorageClient() return LocalStorageClient() + + +@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, ctx: Context, token: Optional[BzmToken]) -> SessionScope: + raise NotImplementedError + + +class DefaultSessionScopeResolver(SessionScopeResolverPort): + """ + Resolve scope from request/ctx metadata. + + Hosted HTTP receives `Mcp-Session-Id` via header. + Local stdio/docker falls back to FastMCP context session_id when available. + """ + + @staticmethod + def _resolve_session_id(ctx: Context) -> str: + request = getattr(getattr(ctx, "request_context", None), "request", None) + if request is not None: + session_id = request.headers.get("mcp-session-id") + if session_id and session_id.strip(): + return session_id.strip() + session_id = getattr(ctx, "session_id", None) + if session_id is not None and str(session_id).strip(): + return str(session_id).strip() + return "default" + + @staticmethod + def _resolve_user_id(token: Optional[BzmToken]) -> str: + if token is not None and token.id.strip(): + return token.id.strip() + return "anonymous" + + def resolve(self, ctx: Context, token: Optional[BzmToken]) -> SessionScope: + return SessionScope( + user_id=self._resolve_user_id(token), + mcp_session_id=self._resolve_session_id(ctx), + ) + + +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 + + +class HttpSessionStorageProvider(SessionStoragePort): + def __init__( + self, + base_url: str, + timeout_seconds: float = 15.0, + ) -> None: + self._base_url = base_url.rstrip("/") + self._timeout = timeout_seconds + + 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(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: + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.put( + self._url_for_scope(scope), + json=payload.to_dict(), + ) + response.raise_for_status() + + async def get_partition(self, scope: SessionScope) -> SessionPartition | None: + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.get( + self._url_for_scope(scope), + ) + 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: + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.delete( + self._url_for_scope(scope), + ) + 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 83ffc95..79acc33 100644 --- a/docs/hosted-http.md +++ b/docs/hosted-http.md @@ -48,7 +48,7 @@ docker run --rm -p 8000:8000 \ -e FASTMCP_HOST=0.0.0.0 \ -e FASTMCP_PORT=8000 \ -e FASTMCP_STREAMABLE_HTTP_PATH=/mcp \ - -e BZM_STORAGE_BACKEND=memory \ + -e BZM_STORAGE_STRATEGY=memory \ ghcr.io/blazemeter/bzm-mcp:latest ``` @@ -60,9 +60,9 @@ docker run --rm -p 8000:8000 \ | `FASTMCP_HOST` | Bind address (HTTP only) | `127.0.0.1` | | `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` | -| `BZM_STORAGE_BACKEND` | `memory` or `http` | `memory` | +| `BZM_STORAGE_STRATEGY` | `memory` or `http` | `memory` | -On streamable-http, local file paths are always rejected regardless of `BZM_STORAGE_BACKEND` (hosted fail-closed storage). +On streamable-http, local file paths are always rejected regardless of `BZM_STORAGE_STRATEGY` (hosted fail-closed storage). ## Hosted MVP limitations diff --git a/tests/test_file_access.py b/tests/test_file_access.py new file mode 100644 index 0000000..3d7d06f --- /dev/null +++ b/tests/test_file_access.py @@ -0,0 +1,94 @@ +""" +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 pytest + +from config.file_access import ( + DockerMappedFileSource, + LocalPathFileSource, + StorageFileSource, + build_file_access, +) + + +class TestBuildFileAccess: + def test_stdio_defaults_to_local_path_source(self, monkeypatch): + monkeypatch.setenv("MCP_DOCKER", "false") + source = build_file_access("stdio") + assert isinstance(source, LocalPathFileSource) + + def test_stdio_docker_requires_source_working_directory(self, monkeypatch): + monkeypatch.setenv("MCP_DOCKER", "true") + monkeypatch.delenv("SOURCE_WORKING_DIRECTORY", raising=False) + with pytest.raises(ValueError, match="Working directory must be set"): + build_file_access("stdio") + + def test_stdio_docker_uses_mapped_source(self, monkeypatch): + monkeypatch.setenv("MCP_DOCKER", "true") + monkeypatch.setenv("SOURCE_WORKING_DIRECTORY", "/Users/me/work") + source = build_file_access("stdio") + assert isinstance(source, DockerMappedFileSource) + + def test_streamable_http_requires_storage_api_base_url(self, monkeypatch): + monkeypatch.delenv("BZM_STORAGE_API_BASE_URL", raising=False) + with pytest.raises(ValueError, match="BZM_STORAGE_API_BASE_URL is required"): + build_file_access("streamable-http") + + def test_streamable_http_uses_storage_file_source(self, monkeypatch): + monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") + monkeypatch.setattr(StorageFileSource, "ensure_available", lambda self: None) + source = build_file_access("streamable-http") + assert isinstance(source, StorageFileSource) + + +class TestLocalPathFileSource: + def test_exists_is_file_and_read_bytes(self, tmp_path): + test_file = tmp_path / "demo.txt" + test_file.write_text("hello", encoding="utf-8") + + source = LocalPathFileSource() + assert source.map_paths([str(test_file)]) == [str(test_file)] + assert source.exists(str(test_file)) is True + assert source.is_file(str(test_file)) is True + assert source.read_bytes(str(test_file)) == b"hello" + + +class TestDockerMappedFileSource: + def test_map_paths_inside_source_directory(self, tmp_path): + source_root = tmp_path / "workspace" + source_root.mkdir() + test_file = source_root / "suite.jmx" + test_file.write_text("xml", encoding="utf-8") + + source = DockerMappedFileSource( + source_working_directory=str(source_root), + container_working_directory="/home/bzm-mcp/working_directory", + ) + mapped = source.map_paths([str(test_file)]) + assert mapped == ["/home/bzm-mcp/working_directory/suite.jmx"] + + def test_map_paths_outside_source_directory_remains_unchanged(self, tmp_path): + source_root = tmp_path / "workspace" + source_root.mkdir() + external_file = tmp_path / "external.csv" + external_file.write_text("1,2", encoding="utf-8") + + source = DockerMappedFileSource( + source_working_directory=str(source_root), + container_working_directory="/home/bzm-mcp/working_directory", + ) + mapped = source.map_paths([str(external_file)]) + assert mapped == [str(external_file)] diff --git a/tests/test_hierarchy_and_consent.py b/tests/test_hierarchy_and_consent.py index d777667..2a0fb12 100644 --- a/tests/test_hierarchy_and_consent.py +++ b/tests/test_hierarchy_and_consent.py @@ -17,6 +17,8 @@ import asyncio from types import SimpleNamespace +from config.file_access import LocalPathFileSource +from config.storage import DefaultSessionScopeResolver from models.result import BaseResult from tools import account_manager, project_manager, workspace_manager, test_manager, execution_manager from tools.account_manager import AccountManager @@ -84,7 +86,12 @@ async def fake_read_project(*args, **kwargs): monkeypatch.setattr(test_manager, "api_request", fake_api_request) monkeypatch.setattr(test_manager.bridge, "read_project", fake_read_project) - manager = TestManager(token=None, ctx=None) + manager = TestManager( + token=None, + ctx=None, + file_access=LocalPathFileSource(), + scope_resolver=DefaultSessionScopeResolver(), + ) result = asyncio.run(manager.read(50)) diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py index 6ffc501..f21f641 100644 --- a/tests/test_http_auth.py +++ b/tests/test_http_auth.py @@ -32,8 +32,12 @@ StdioAuthProvider, parse_authorization_header, ) +from config.file_access import LocalPathFileSource, StorageFileSource from config.runtime import build_runtime -from config.storage import HttpStorageClient, LocalStorageClient +from config.storage import ( + HttpSessionStorageProvider, + InMemorySessionStorageProvider, +) from config.token import BzmToken, BzmTokenError @@ -175,14 +179,29 @@ async def health(_request: Request): class TestBuildRuntime: def test_build_runtime_stdio_and_http(self, monkeypatch): monkeypatch.delenv("MCP_DOCKER", raising=False) - monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + monkeypatch.delenv("BZM_STORAGE_STRATEGY", raising=False) + monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") + monkeypatch.setattr(HttpSessionStorageProvider, "ensure_available", lambda self: None) stdio = build_runtime("stdio") assert stdio.transport == "stdio" assert isinstance(stdio.auth, StdioAuthProvider) - assert isinstance(stdio.storage, LocalStorageClient) + assert isinstance(stdio.storage, InMemorySessionStorageProvider) + assert isinstance(stdio.file_access, LocalPathFileSource) http = build_runtime("streamable-http") assert http.transport == "streamable-http" assert isinstance(http.auth, HttpAuthProvider) - assert isinstance(http.storage, HttpStorageClient) + assert isinstance(http.storage, HttpSessionStorageProvider) + assert isinstance(http.file_access, StorageFileSource) + + def test_build_runtime_http_uses_storage_api_when_configured(self, monkeypatch): + monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") + monkeypatch.setattr(HttpSessionStorageProvider, "ensure_available", lambda self: None) + monkeypatch.setattr(StorageFileSource, "ensure_available", lambda self: None) + + runtime = build_runtime("streamable-http") + assert runtime.transport == "streamable-http" + assert isinstance(runtime.auth, HttpAuthProvider) + assert isinstance(runtime.storage, HttpSessionStorageProvider) + assert isinstance(runtime.file_access, StorageFileSource) diff --git a/tests/test_main_transport.py b/tests/test_main_transport.py index b7a3cb2..a1ee8fe 100644 --- a/tests/test_main_transport.py +++ b/tests/test_main_transport.py @@ -2,8 +2,9 @@ import main from config.auth import HttpAuthProvider, StdioAuthProvider +from config.file_access import StorageFileSource from config.runtime import AppRuntime -from config.storage import HttpStorageClient, LocalStorageClient +from config.storage import HttpSessionStorageProvider, InMemorySessionStorageProvider class _DummyFastMCP: @@ -22,6 +23,9 @@ def _patch_mcp_server_dependencies(monkeypatch): monkeypatch.setattr(main, "register_confirm_mode", lambda *a, **k: None) monkeypatch.setattr(main, "register_tools", lambda *a, **k: None) monkeypatch.setattr(main, "FastMCP", _DummyFastMCP) + monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") + monkeypatch.setattr(HttpSessionStorageProvider, "ensure_available", lambda self: None) + monkeypatch.setattr(StorageFileSource, "ensure_available", lambda self: None) class TestResolveMcpTransport: @@ -105,7 +109,7 @@ def capture_register(mcp, runtime): assert isinstance(runtime, AppRuntime) assert runtime.transport == "streamable-http" assert isinstance(runtime.auth, HttpAuthProvider) - assert isinstance(runtime.storage, HttpStorageClient) + assert isinstance(runtime.storage, HttpSessionStorageProvider) def test_stdio_and_docker_register_stdio_auth_provider(self, monkeypatch): captured = {} @@ -118,7 +122,7 @@ def capture_register(mcp, runtime): monkeypatch.setattr(main, "get_token", lambda: token) monkeypatch.setattr(main, "register_tools", capture_register) monkeypatch.delenv("MCP_DOCKER", raising=False) - monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + monkeypatch.delenv("BZM_STORAGE_STRATEGY", raising=False) main.build_mcp_server(transport="stdio") main.build_mcp_server(transport="docker") @@ -127,7 +131,7 @@ def capture_register(mcp, runtime): for runtime in captured["runtimes"]: assert isinstance(runtime.auth, StdioAuthProvider) assert runtime.auth.get_token(ctx=None) is token - assert isinstance(runtime.storage, LocalStorageClient) + assert isinstance(runtime.storage, InMemorySessionStorageProvider) class TestRunTransportDispatch: diff --git a/tests/test_storage.py b/tests/test_storage.py index 6d0dd79..649a9cb 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -17,10 +17,14 @@ import pytest +from config.file_access import LocalPathFileSource, StorageFileSource from config.runtime import build_runtime from config.storage import ( + DefaultSessionScopeResolver, HOSTED_FILE_ACCESS_MESSAGE, + HttpSessionStorageProvider, HttpStorageClient, + InMemorySessionStorageProvider, LocalStorageClient, StorageNotSupportedError, build_storage, @@ -31,27 +35,27 @@ class TestResolveStorageBackend: def test_default_memory(self, monkeypatch): - monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + monkeypatch.delenv("BZM_STORAGE_STRATEGY", raising=False) assert resolve_storage_backend() == "memory" def test_env_http(self, monkeypatch): - monkeypatch.setenv("BZM_STORAGE_BACKEND", "http") + monkeypatch.setenv("BZM_STORAGE_STRATEGY", "http") assert resolve_storage_backend() == "http" def test_invalid_raises(self): - with pytest.raises(ValueError, match="BZM_STORAGE_BACKEND"): + with pytest.raises(ValueError, match="BZM_STORAGE_STRATEGY"): resolve_storage_backend("s3") class TestBuildStorage: def test_stdio_memory_uses_local(self, monkeypatch): - monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) + monkeypatch.delenv("BZM_STORAGE_STRATEGY", raising=False) monkeypatch.delenv("MCP_DOCKER", raising=False) storage = build_storage("stdio") assert isinstance(storage, LocalStorageClient) def test_streamable_http_uses_http_client(self, monkeypatch): - monkeypatch.setenv("BZM_STORAGE_BACKEND", "memory") + monkeypatch.setenv("BZM_STORAGE_STRATEGY", "memory") storage = build_storage("streamable-http") assert isinstance(storage, HttpStorageClient) @@ -91,20 +95,28 @@ def test_all_file_methods_raise(self): class TestRuntimeStorageWiring: - def test_http_runtime_gets_http_storage(self): + def test_http_runtime_gets_http_storage(self, monkeypatch): + monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") + monkeypatch.setattr(HttpSessionStorageProvider, "ensure_available", lambda self: None) runtime = build_runtime("streamable-http") - assert isinstance(runtime.storage, HttpStorageClient) + assert isinstance(runtime.storage, HttpSessionStorageProvider) + assert isinstance(runtime.file_access, StorageFileSource) def test_stdio_runtime_gets_local_storage(self, monkeypatch): monkeypatch.delenv("MCP_DOCKER", raising=False) - monkeypatch.delenv("BZM_STORAGE_BACKEND", raising=False) runtime = build_runtime("stdio") - assert isinstance(runtime.storage, LocalStorageClient) + assert isinstance(runtime.storage, InMemorySessionStorageProvider) + assert isinstance(runtime.file_access, LocalPathFileSource) class TestUploadAssetsHostedRejection: def test_upload_assets_returns_clear_error_on_http_storage(self): - manager = TestManager(token=None, ctx=None, storage=HttpStorageClient()) + manager = TestManager( + token=None, + ctx=None, + file_access=StorageFileSource("https://mcp-storage.internal"), + scope_resolver=DefaultSessionScopeResolver(), + ) async def _fake_read(_test_id): from models.result import BaseResult @@ -115,4 +127,4 @@ async def _fake_read(_test_id): manager.upload_assets(1, ["/tmp/demo.jmx"], main_script=None) ) assert "error" in result - assert "not supported on the hosted MCP" in result["error"] + assert "No valid files found to upload" in result["error"] diff --git a/tests/test_upload_assets_security.py b/tests/test_upload_assets_security.py index 042befc..c3a32e6 100644 --- a/tests/test_upload_assets_security.py +++ b/tests/test_upload_assets_security.py @@ -14,7 +14,8 @@ limitations under the License. """ -from config.storage import LocalStorageClient +from config.file_access import LocalPathFileSource +from config.storage import DefaultSessionScopeResolver from tools.test_manager import TestManager as UploadAssetsManager @@ -55,12 +56,18 @@ def test_validate_files_classifies_valid_invalid_and_blocked(self, tmp_path): invalid_files = [] blocked_files = [] - manager = UploadAssetsManager(token=None, ctx=None, storage=LocalStorageClient()) + manager = UploadAssetsManager( + token=None, + ctx=None, + file_access=LocalPathFileSource(), + scope_resolver=DefaultSessionScopeResolver(), + ) manager._validate_files( [str(safe_file), str(env_file), str(missing_file)], valid_files, invalid_files, blocked_files, + file_access=manager.file_access, ) assert valid_files == [str(safe_file)] diff --git a/tools/test_manager.py b/tools/test_manager.py index 46a55de..0b54509 100644 --- a/tools/test_manager.py +++ b/tools/test_manager.py @@ -23,8 +23,9 @@ from mcp.server.fastmcp import Context from config.blazemeter import TESTS_ENDPOINT, TOOLS_PREFIX +from config.file_access import FileAccessPort from config.security import detect_sensitive_upload_path_reason -from config.storage import LocalStorageClient, StorageNotSupportedError, StoragePort +from config.storage import SessionScopeResolverPort from config.token import BzmToken from config.runtime import AppRuntime from formatters.failure_criteria_labels import failure_criteria_meta_payload @@ -55,10 +56,15 @@ def __init__( self, token: Optional[BzmToken], ctx: Context, - storage: Optional[StoragePort] = None, + file_access: FileAccessPort, + scope_resolver: SessionScopeResolverPort, ): super().__init__(token, ctx) - self.storage = storage or LocalStorageClient() + self.file_access = file_access + self.scope_resolver = scope_resolver + + def _current_scope(self): + return self.scope_resolver.resolve(self.ctx, self.token) async def read(self, test_id: Optional[int]) -> BaseResult: if not isinstance(test_id, int) or test_id < 1: @@ -152,6 +158,8 @@ def _validate_files( valid_files: List[str], invalid_files: List[str], blocked_files: List[Dict[str, str]], + file_access: Optional[FileAccessPort] = None, + scope=None, ): # Security design note: # Uploads are intentionally allowed from any user working location (not restricted to one workspace root), @@ -174,7 +182,9 @@ def _validate_files( } ) continue - if self.storage.exists(file_path) and self.storage.is_file(file_path): + exists = file_access.exists(file_path, scope=scope) if file_access else False + is_file = file_access.is_file(file_path, scope=scope) if file_access else False + if exists and is_file: logger.debug(f"File exists: {file_path}") valid_files.append(file_path) else: @@ -220,20 +230,14 @@ async def upload_assets( logger.debug(f"Starting upload_assets for test_id: {test_id}") logger.debug(f"Original file paths: {file_paths}") logger.debug(f"Main script: {main_script}") + scope = self._current_scope() - try: - mapped_file_paths = self.storage.map_paths(file_paths) - except StorageNotSupportedError as exc: - return {"error": str(exc)} - + mapped_file_paths = self.file_access.map_paths(file_paths, scope=scope) logger.debug(f"Mapped file paths: {mapped_file_paths}") mapped_main_script = None if main_script: - try: - mapped_main_script_list = self.storage.map_paths([main_script]) - except StorageNotSupportedError as exc: - return {"error": str(exc)} + mapped_main_script_list = self.file_access.map_paths([main_script], scope=scope) mapped_main_script = ( mapped_main_script_list[0] if mapped_main_script_list else None ) @@ -243,12 +247,14 @@ async def upload_assets( invalid_files = [] blocked_files = [] - try: - self._validate_files( - mapped_file_paths, valid_files, invalid_files, blocked_files - ) - except StorageNotSupportedError as exc: - return {"error": str(exc)} + self._validate_files( + mapped_file_paths, + valid_files, + invalid_files, + blocked_files, + file_access=self.file_access, + scope=scope, + ) logger.debug(f"Valid files: {valid_files}") logger.debug(f"Invalid files: {invalid_files}") @@ -263,9 +269,7 @@ async def upload_assets( } logger.debug("Starting concurrent uploads") - upload_tasks = [ - self._upload_single_file(test_id, file_path) for file_path in valid_files - ] + upload_tasks = [self._upload_single_file(test_id, file_path, scope) for file_path in valid_files] upload_results = await asyncio.gather(*upload_tasks, return_exceptions=True) logger.debug(f"Upload results: {upload_results}") @@ -295,14 +299,14 @@ async def upload_assets( "config_update": config_update_result, } - async def _upload_single_file(self, test_id: int, file_path: str) -> BaseResult: + async def _upload_single_file(self, test_id: int, file_path: str, scope) -> BaseResult: logger.debug(f"Uploading single file: {file_path} to test: {test_id}") try: - file_name = self.storage.basename(file_path) + file_name = Path(file_path).name logger.debug(f"File name: {file_name}") - file_content = self.storage.read_bytes(file_path) + file_content = self.file_access.read_bytes(file_path, scope=scope) logger.debug(f"File size: {len(file_content)} bytes") @@ -317,8 +321,6 @@ async def _upload_single_file(self, test_id: int, file_path: str) -> BaseResult: return result - except StorageNotSupportedError: - raise except Exception as e: logger.error(f"Exception in _upload_single_file: {e}") logger.error(f"Traceback: {format_sanitized_traceback(e)}") @@ -328,7 +330,7 @@ async def _update_test_configuration( self, test_id: int, main_script_path: str ) -> BaseResult: try: - file_name = self.storage.basename(main_script_path) + file_name = Path(main_script_path).name config_update = { "configuration": { "filename": file_name, @@ -667,7 +669,12 @@ def register(mcp, runtime: AppRuntime): """, ) async def tests(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - test_manager = TestManager(runtime.auth.get_token(ctx), ctx, runtime.storage) + test_manager = TestManager( + runtime.auth.get_token(ctx), + ctx, + runtime.file_access, + runtime.scope_resolver, + ) async def _dispatch(): match action: From 9d46ff9c25cee20081da9b5543989c16e1d5f255 Mon Sep 17 00:00:00 2001 From: alejandroaires Date: Tue, 18 Aug 2026 10:43:23 -0300 Subject: [PATCH 4/6] Mob 50447 per session elicitation mode (#101) * refactor confirmation mode to per-session context and runtime user_config Removed global confirmation state and moved elicitation mode resolution to per-session context. HTTP now reads Confirmation-Mode from request headers, while stdio injects startup confirmation mode through runtime user_config, and managers pass that config so require_confirmation resolves mode consistently per invocation. # Conflicts: # config/runtime.py # tests/test_http_auth.py # tools/test_manager.py * persist confirmation mode per http session in auth middleware * update readme * remove http confirmation-mode session persistence * refactor managers to consume token from user_config * update readme * removed "NONE" case on normalized confirme mode * update hosted-http documentation * user config resolution from ctx * standardize context user_config resolution * centralize ctx token/user_config resolution in config module * move context user config setup into AppRuntime.configure_context * Update docs/hosted-http.md Co-authored-by: Joaquin Araujo <48018971+Baraujo25@users.noreply.github.com> * Make TestManager file ports optional for hosted HTTP --------- Co-authored-by: Joaquin Araujo <48018971+Baraujo25@users.noreply.github.com> --- README.md | 22 ----- config/auth.py | 20 ++++ config/context_resolution.py | 50 ++++++++++ config/runtime.py | 67 +++++++++++++- docs/hosted-http.md | 9 +- main.py | 4 +- models/manager.py | 12 +-- tests/test_confirmation_behavior.py | 20 ++-- tests/test_hierarchy_and_consent.py | 48 +++++++--- tests/test_http_auth.py | 134 ++++++++++++++++++++++++++- tests/test_main_transport.py | 1 - tests/test_storage.py | 8 +- tests/test_upload_assets_security.py | 1 - tools/account_manager.py | 11 ++- tools/billing_manager.py | 13 ++- tools/bridge.py | 51 ++++++++-- tools/execution_manager.py | 27 ++++-- tools/help_manager.py | 13 ++- tools/project_manager.py | 11 ++- tools/report_manager.py | 10 +- tools/skills_manager.py | 11 ++- tools/test_manager.py | 27 +++--- tools/user_manager.py | 13 ++- tools/utils.py | 54 ++++++++--- tools/workspace_manager.py | 11 ++- 25 files changed, 511 insertions(+), 137 deletions(-) create mode 100644 config/context_resolution.py diff --git a/README.md b/README.md index 8969be3..899d84f 100644 --- a/README.md +++ b/README.md @@ -110,28 +110,6 @@ After installing, set `BLAZEMETER_API_KEY` to your `api-key.json` path in your c --- -**Hosted HTTP (streamable-http)** - -Connect to the shared hosted server with Bearer auth. Production URL: `https://mcp.blazemeter.com/mcp` - -```json -{ - "mcpServers": { - "BlazeMeter MCP": { - "url": "https://mcp.blazemeter.com/mcp", - "headers": { - "Authorization": "Bearer :" - } - } - } -} -``` - -> [!NOTE] -> For local/operator HTTP setup, env vars, auth behavior, and MVP limits (including `upload_assets`), see [docs/hosted-http.md](docs/hosted-http.md). - ---- - **Docker MCP Client Configuration** 1. **Prerequisites:** [Docker]([https://www.docker.com/products/docker-desktop/](https://www.docker.com/products/docker-desktop/)) diff --git a/config/auth.py b/config/auth.py index 373abcb..eab918d 100644 --- a/config/auth.py +++ b/config/auth.py @@ -25,6 +25,16 @@ from config.token import BzmToken, BzmTokenError BZM_TOKEN_STATE_ATTR = "token" +BZM_USER_CONFIG_STATE_ATTR = "user_config" +BZM_CONFIRMATION_MODE_HEADER = "confirmation-mode" + + +def _normalize_confirmation_mode(raw_value: Optional[str]) -> str: + value = (raw_value or "").strip().upper() + if value in ("DELETE", "CUD", "DISABLE"): + return value + return "DELETE" + # Unauthenticated probe paths for orchestrators / load balancers. HEALTH_PATHS = frozenset({"/health", "/healthz"}) @@ -119,6 +129,16 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: return setattr(request.state, BZM_TOKEN_STATE_ATTR, token) + raw_confirmation_mode = request.headers.get(BZM_CONFIRMATION_MODE_HEADER) + if raw_confirmation_mode is not None and raw_confirmation_mode.strip(): + confirmation_mode = _normalize_confirmation_mode(raw_confirmation_mode) + else: + confirmation_mode = "DELETE" + setattr( + request.state, + BZM_USER_CONFIG_STATE_ATTR, + {"confirmation_mode": confirmation_mode}, + ) await self.app(scope, receive, send) diff --git a/config/context_resolution.py b/config/context_resolution.py new file mode 100644 index 0000000..b86c845 --- /dev/null +++ b/config/context_resolution.py @@ -0,0 +1,50 @@ +""" +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 typing import Any + +from config.auth import BZM_TOKEN_STATE_ATTR, BZM_USER_CONFIG_STATE_ATTR + + +def get_request_context(ctx: Any) -> Any: + return getattr(ctx, "request_context", None) + + +def get_request_state(ctx: Any) -> Any: + request_context = get_request_context(ctx) + request = getattr(request_context, "request", None) + return getattr(request, "state", None) + + +def resolve_ctx_user_config(ctx: Any) -> dict[str, Any]: + request_context = get_request_context(ctx) + request_state = get_request_state(ctx) + + request_context_config = getattr(request_context, BZM_USER_CONFIG_STATE_ATTR, None) + if isinstance(request_context_config, dict): + return request_context_config + + request_state_config = getattr(request_state, BZM_USER_CONFIG_STATE_ATTR, None) + if isinstance(request_state_config, dict): + return request_state_config + + return {} + + +def resolve_ctx_token(ctx: Any) -> Any: + user_config = resolve_ctx_user_config(ctx) + request_state = get_request_state(ctx) + request_state_token = getattr(request_state, BZM_TOKEN_STATE_ATTR, None) + return user_config.get("token") or request_state_token diff --git a/config/runtime.py b/config/runtime.py index 37f674a..31f8650 100644 --- a/config/runtime.py +++ b/config/runtime.py @@ -14,10 +14,15 @@ limitations under the License. """ from dataclasses import dataclass -from typing import Literal, Optional import os +from typing import Any, Literal, Optional -from config.auth import AuthPort, HttpAuthProvider, StdioAuthProvider +from config.auth import ( + AuthPort, + BZM_USER_CONFIG_STATE_ATTR, + HttpAuthProvider, + StdioAuthProvider, +) from config.file_access import FileAccessPort, build_file_access from config.storage import ( DefaultSessionScopeResolver, @@ -27,6 +32,7 @@ SessionStoragePort, ) from config.token import BzmToken +from tools.utils import ConfirmMode Transport = Literal["stdio", "streamable-http"] @@ -40,11 +46,61 @@ class AppRuntime: storage: SessionStoragePort file_access: FileAccessPort scope_resolver: SessionScopeResolverPort + user_config: dict[str, Any] + + def resolve_user_config(self, ctx: Any) -> dict[str, Any]: + user_config = dict(self.user_config) + user_config.update(_read_ctx_user_config(ctx)) + token = self.auth.get_token(ctx) + if token is not None: + user_config["token"] = token + return user_config + + def configure_context(self, ctx: Any) -> dict[str, Any]: + user_config = self.resolve_user_config(ctx) + _hydrate_ctx_user_config(ctx, user_config) + return user_config + + +def _read_ctx_user_config(ctx: Any) -> dict[str, Any]: + if ctx is None: + return {} + + user_config: dict[str, Any] = {} + request_context = getattr(ctx, "request_context", None) + request = getattr(request_context, "request", None) + request_state = getattr(request, "state", None) + + for target, attr_name in ( + (ctx, "user_config"), + (request_context, BZM_USER_CONFIG_STATE_ATTR), + (request_state, BZM_USER_CONFIG_STATE_ATTR), + ): + request_config = getattr(target, attr_name, None) + if isinstance(request_config, dict): + user_config.update(request_config) + + return user_config + + +def _hydrate_ctx_user_config(ctx: Any, user_config: dict[str, Any]) -> None: + if ctx is None: + return + + config_copy = dict(user_config) + request_context = getattr(ctx, "request_context", None) + request = getattr(request_context, "request", None) + request_state = getattr(request, "state", None) + + for target in (request_context, request_state): + if target is not None: + setattr(target, BZM_USER_CONFIG_STATE_ATTR, dict(config_copy)) def build_runtime( transport: Transport, startup_token: Optional[BzmToken] = None, + startup_confirmation_mode: ConfirmMode = ConfirmMode.DELETE, ) -> AppRuntime: """ Compose auth, file access, and session storage for the selected transport. @@ -53,12 +109,18 @@ def build_runtime( - streamable-http: request-scoped auth and storage API-backed partitions. """ if transport == "stdio": + stdio_user_config = { + "startup_token": startup_token, + "token": startup_token, + "confirmation_mode": startup_confirmation_mode.name, + } return AppRuntime( transport=transport, auth=StdioAuthProvider(startup_token), storage=InMemorySessionStorageProvider(), file_access=build_file_access(transport), scope_resolver=DefaultSessionScopeResolver(), + user_config=stdio_user_config, ) if transport == "streamable-http": @@ -77,6 +139,7 @@ def build_runtime( storage=storage, file_access=build_file_access(transport), scope_resolver=DefaultSessionScopeResolver(), + user_config={}, ) raise ValueError(f"Unknown transport: {transport}") diff --git a/docs/hosted-http.md b/docs/hosted-http.md index 79acc33..ea7849f 100644 --- a/docs/hosted-http.md +++ b/docs/hosted-http.md @@ -16,7 +16,8 @@ Configure the MCP client with that URL and your BlazeMeter API key as Bearer cre "BlazeMeter MCP": { "url": "https://mcp.blazemeter.com/mcp", "headers": { - "Authorization": "Bearer :" + "Authorization": "Bearer :", + "confirmation-mode": "DELETE" } } } @@ -32,6 +33,12 @@ For a locally run server, use `"url": "http://localhost:8000/mcp"` instead. - Well-formed but wrong API keys fail later inside BlazeMeter API calls (same as stdio). - Stdio / local Docker transport uses `api-key.json` / env / Docker secrets instead of Bearer auth. +### Confirmation mode header + +- Optional header: `confirmation-mode` +- Allowed values: `DELETE`, `CUD`, `DISABLE` +- If omitted, empty, or invalid, the session falls back to `DELETE`. + ## Local / operator run Transport resolution precedence: **CLI `--mcp` > `BZM_MCP_TRANSPORT` > stdio**. diff --git a/main.py b/main.py index dc79815..472d480 100644 --- a/main.py +++ b/main.py @@ -35,7 +35,7 @@ from config.version import __version__, __executable__, __bundle__ from server import register_tools from telemetry import init_telemetry -from tools.utils import ConfirmMode, register_confirm_mode +from tools.utils import ConfirmMode BLAZEMETER_API_KEY_FILE_PATH = os.getenv('BLAZEMETER_API_KEY') @@ -442,6 +442,7 @@ def build_mcp_server( app_runtime = build_runtime( wire_transport, startup_token=get_token() if wire_transport == "stdio" else None, + startup_confirmation_mode=confirm_mode, ) instructions = """ # BlazeMeter MCP Server @@ -507,7 +508,6 @@ def build_mcp_server( streamable_http_path=streamable_http_path, stateless_http=False, ) - register_confirm_mode(confirm_mode) register_tools(mcp, app_runtime) return mcp, wire_transport diff --git a/models/manager.py b/models/manager.py index ca90575..0b4804d 100644 --- a/models/manager.py +++ b/models/manager.py @@ -13,14 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. """ -from typing import Optional - from mcp.server.fastmcp import Context -from config.token import BzmToken - +from config.context_resolution import resolve_ctx_token class Manager: - def __init__(self, token: Optional[BzmToken], ctx: Context): - self.token = token + def __init__( + self, + ctx: Context, + ): self.ctx = ctx + self.token = resolve_ctx_token(ctx) diff --git a/tests/test_confirmation_behavior.py b/tests/test_confirmation_behavior.py index 2bd9d68..08b1921 100644 --- a/tests/test_confirmation_behavior.py +++ b/tests/test_confirmation_behavior.py @@ -17,12 +17,16 @@ import asyncio from types import SimpleNamespace +from config.auth import BZM_USER_CONFIG_STATE_ATTR from models.result import BaseResult -from tools.utils import ConfirmMode, Operations, register_confirm_mode, require_confirmation +from tools.utils import Operations, require_confirmation class _ManagerWithConfirmation: - def __init__(self, ctx): + def __init__(self, ctx, confirmation_mode: str = "DELETE"): + request_state = SimpleNamespace(**{BZM_USER_CONFIG_STATE_ATTR: {"confirmation_mode": confirmation_mode}}) + request = SimpleNamespace(state=request_state) + setattr(ctx, "request_context", SimpleNamespace(request=request)) self.ctx = ctx @require_confirmation(operation=Operations.CREATE) @@ -47,8 +51,7 @@ async def elicit(self, message, schema): class TestConfirmationBehavior: def test_blocks_when_confirmation_required_and_elicit_unsupported(self): - register_confirm_mode(ConfirmMode.CUD) - manager = _ManagerWithConfirmation(_NoElicitContext()) + manager = _ManagerWithConfirmation(_NoElicitContext(), confirmation_mode="CUD") result = asyncio.run(manager.do_create()) @@ -57,8 +60,7 @@ def test_blocks_when_confirmation_required_and_elicit_unsupported(self): assert result.result is None def test_allows_when_confirmation_required_and_user_accepts(self): - register_confirm_mode(ConfirmMode.CUD) - manager = _ManagerWithConfirmation(_AcceptContext()) + manager = _ManagerWithConfirmation(_AcceptContext(), confirmation_mode="CUD") result = asyncio.run(manager.do_create()) @@ -66,8 +68,7 @@ def test_allows_when_confirmation_required_and_user_accepts(self): assert result.result == ["created"] def test_returns_cancelled_when_confirmation_required_and_user_declines(self): - register_confirm_mode(ConfirmMode.CUD) - manager = _ManagerWithConfirmation(_RejectContext()) + manager = _ManagerWithConfirmation(_RejectContext(), confirmation_mode="CUD") result = asyncio.run(manager.do_create()) @@ -75,8 +76,7 @@ def test_returns_cancelled_when_confirmation_required_and_user_declines(self): assert result.result == ["Action manually cancelled by the user."] def test_allows_when_confirmation_not_required_even_without_elicit(self): - register_confirm_mode(ConfirmMode.DISABLE) - manager = _ManagerWithConfirmation(_NoElicitContext()) + manager = _ManagerWithConfirmation(_NoElicitContext(), confirmation_mode="DISABLE") result = asyncio.run(manager.do_create()) diff --git a/tests/test_hierarchy_and_consent.py b/tests/test_hierarchy_and_consent.py index 2a0fb12..437682e 100644 --- a/tests/test_hierarchy_and_consent.py +++ b/tests/test_hierarchy_and_consent.py @@ -17,15 +17,13 @@ import asyncio from types import SimpleNamespace -from config.file_access import LocalPathFileSource -from config.storage import DefaultSessionScopeResolver from models.result import BaseResult from tools import account_manager, project_manager, workspace_manager, test_manager, execution_manager +from tools import bridge from tools.account_manager import AccountManager from tools.execution_manager import ExecutionManager from tools.project_manager import ProjectManager from tools.test_manager import TestManager -from tools.utils import register_confirm_mode, ConfirmMode from tools.workspace_manager import WorkspaceManager @@ -35,7 +33,7 @@ async def fake_api_request(*args, **kwargs): return BaseResult(result=[SimpleNamespace(ai_consent=False)]) monkeypatch.setattr(account_manager, "api_request", fake_api_request) - manager = AccountManager(token=None, ctx=None) + manager = AccountManager(ctx=None) result = asyncio.run(manager.read(123)) @@ -51,7 +49,7 @@ async def fake_read_account(*args, **kwargs): monkeypatch.setattr(workspace_manager, "api_request", fake_api_request) monkeypatch.setattr(workspace_manager.bridge, "read_account", fake_read_account) - manager = WorkspaceManager(token=None, ctx=None) + manager = WorkspaceManager(ctx=None) result = asyncio.run(manager.read(55)) @@ -70,7 +68,7 @@ async def fake_count_project_tests(*args, **kwargs): monkeypatch.setattr(project_manager, "api_request", fake_api_request) monkeypatch.setattr(project_manager.bridge, "read_workspace", fake_read_workspace) monkeypatch.setattr(project_manager.bridge, "count_project_tests", fake_count_project_tests) - manager = ProjectManager(token=None, ctx=None) + manager = ProjectManager(ctx=None) result = asyncio.run(manager.read(200)) @@ -86,19 +84,39 @@ async def fake_read_project(*args, **kwargs): monkeypatch.setattr(test_manager, "api_request", fake_api_request) monkeypatch.setattr(test_manager.bridge, "read_project", fake_read_project) - manager = TestManager( - token=None, - ctx=None, - file_access=LocalPathFileSource(), - scope_resolver=DefaultSessionScopeResolver(), - ) + manager = TestManager(ctx=None) result = asyncio.run(manager.read(50)) assert result.error == "Project validation failed" + def test_count_project_tests_does_not_require_file_access(self, monkeypatch): + async def fake_api_request(*args, **kwargs): + return BaseResult(result=[], total=4) + + monkeypatch.setattr(test_manager, "api_request", fake_api_request) + + total = asyncio.run(bridge.count_project_tests(token=None, ctx=None, project_id=2554395)) + + assert total == 4 + + def test_create_validates_project_without_file_access(self, monkeypatch): + async def fake_api_request(*args, **kwargs): + return BaseResult(result=[SimpleNamespace(test_id=99, test_name="dummy_test")]) + + async def fake_read_project(*args, **kwargs): + return BaseResult(result=["ok"]) + + monkeypatch.setattr(test_manager, "api_request", fake_api_request) + monkeypatch.setattr(test_manager.bridge, "read_project", fake_read_project) + manager = TestManager(ctx=None) + + result = asyncio.run(manager.create("dummy_test", 2554395)) + + assert result.error is None + assert result.result[0].test_name == "dummy_test" + def test_execution_start_stops_when_test_validation_fails(self, monkeypatch): - register_confirm_mode(ConfirmMode.DISABLE) called = {"api_request": False} async def fake_read_test(*args, **kwargs): @@ -110,7 +128,7 @@ async def fake_api_request(*args, **kwargs): monkeypatch.setattr(execution_manager.bridge, "read_test", fake_read_test) monkeypatch.setattr(execution_manager, "api_request", fake_api_request) - manager = ExecutionManager(token=None, ctx=None) + manager = ExecutionManager(ctx=SimpleNamespace(user_config={"confirmation_mode": "DISABLE"})) result = asyncio.run(manager.start(42)) @@ -130,7 +148,7 @@ async def fake_fetch_execution_status(*args, **kwargs): # pragma: no cover monkeypatch.setattr(execution_manager, "api_request", fake_api_request) monkeypatch.setattr(execution_manager.bridge, "read_project", fake_read_project) monkeypatch.setattr(ExecutionManager, "_fetch_execution_status", fake_fetch_execution_status) - manager = ExecutionManager(token=None, ctx=None) + manager = ExecutionManager(ctx=None) result = asyncio.run(manager.read(909)) diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py index f21f641..3dd6604 100644 --- a/tests/test_http_auth.py +++ b/tests/test_http_auth.py @@ -26,6 +26,7 @@ from config.auth import ( AuthError, + BZM_USER_CONFIG_STATE_ATTR, BZM_TOKEN_STATE_ATTR, BearerAuthMiddleware, HttpAuthProvider, @@ -39,6 +40,7 @@ InMemorySessionStorageProvider, ) from config.token import BzmToken, BzmTokenError +from models.manager import Manager class TestBearerCredentialParsing: @@ -122,11 +124,40 @@ def make_ctx(token: BzmToken): assert provider.get_token(make_ctx(token_b)).id == "account-b" +class TestManagerTokenResolution: + def test_manager_falls_back_to_request_state_token(self): + token = BzmToken("account-a", "secret-a") + request = SimpleNamespace(state=SimpleNamespace(**{BZM_TOKEN_STATE_ATTR: token})) + ctx = SimpleNamespace(request_context=SimpleNamespace(request=request)) + + manager = Manager(ctx) + + assert manager.token is token + + +class _StrictCtx: + """Mimics FastMCP Context where arbitrary attrs are disallowed.""" + + def __init__(self, request_context): + object.__setattr__(self, "request_context", request_context) + + def __setattr__(self, name, value): + if name == "user_config": + raise ValueError('"Context" object has no field "user_config"') + object.__setattr__(self, name, value) + + class TestBearerAuthMiddleware: def _app(self): async def ok(request: Request): token = getattr(request.state, BZM_TOKEN_STATE_ATTR, None) - return JSONResponse({"id": token.id if token else None}) + user_config = getattr(request.state, BZM_USER_CONFIG_STATE_ATTR, {}) + return JSONResponse( + { + "id": token.id if token else None, + "confirmation_mode": user_config.get("confirmation_mode"), + } + ) return BearerAuthMiddleware(Starlette(routes=[Route("/mcp", endpoint=ok, methods=["POST"])])) @@ -149,6 +180,50 @@ def test_valid_bearer_attaches_token(self): ) assert response.status_code == 200 assert response.json()["id"] == "key-id" + assert response.json()["confirmation_mode"] == "DELETE" + + def test_valid_bearer_reads_confirmation_mode_header(self): + client = TestClient(self._app()) + response = client.post( + "/mcp", + headers={ + "Authorization": "Bearer key-id:key-secret", + "Confirmation-Mode": "CUD", + }, + ) + assert response.status_code == 200 + assert response.json()["confirmation_mode"] == "CUD" + + def test_confirmation_mode_none_falls_back_to_delete(self): + client = TestClient(self._app()) + response = client.post( + "/mcp", + headers={ + "Authorization": "Bearer key-id:key-secret", + "Confirmation-Mode": "NONE", + }, + ) + assert response.status_code == 200 + assert response.json()["confirmation_mode"] == "DELETE" + + def test_confirmation_mode_not_persisted_between_requests_without_header(self): + client = TestClient(self._app()) + first = client.post( + "/mcp", + headers={ + "Authorization": "Bearer key-id:key-secret", + "Confirmation-Mode": "CUD", + }, + ) + assert first.status_code == 200 + assert first.json()["confirmation_mode"] == "CUD" + + second = client.post( + "/mcp", + headers={"Authorization": "Bearer key-id:key-secret"}, + ) + assert second.status_code == 200 + assert second.json()["confirmation_mode"] == "DELETE" def test_options_bypasses_auth(self): async def ok(_request: Request): @@ -188,12 +263,14 @@ def test_build_runtime_stdio_and_http(self, monkeypatch): assert isinstance(stdio.auth, StdioAuthProvider) assert isinstance(stdio.storage, InMemorySessionStorageProvider) assert isinstance(stdio.file_access, LocalPathFileSource) + assert stdio.user_config["confirmation_mode"] == "DELETE" http = build_runtime("streamable-http") assert http.transport == "streamable-http" assert isinstance(http.auth, HttpAuthProvider) assert isinstance(http.storage, HttpSessionStorageProvider) assert isinstance(http.file_access, StorageFileSource) + assert http.user_config == {} def test_build_runtime_http_uses_storage_api_when_configured(self, monkeypatch): monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") @@ -205,3 +282,58 @@ def test_build_runtime_http_uses_storage_api_when_configured(self, monkeypatch): assert isinstance(runtime.auth, HttpAuthProvider) assert isinstance(runtime.storage, HttpSessionStorageProvider) assert isinstance(runtime.file_access, StorageFileSource) + + def test_configure_context_injects_request_context_for_stdio(self, monkeypatch): + monkeypatch.delenv("MCP_DOCKER", raising=False) + runtime = build_runtime("stdio") + ctx = SimpleNamespace(request_context=SimpleNamespace(request=None)) + + user_config = runtime.configure_context(ctx) + + assert "token" in user_config + assert user_config["confirmation_mode"] == "DELETE" + assert getattr(ctx.request_context, BZM_USER_CONFIG_STATE_ATTR, None) == user_config + + def test_configure_context_merges_http_request_state(self, monkeypatch): + monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") + monkeypatch.setattr(HttpSessionStorageProvider, "ensure_available", lambda self: None) + runtime = build_runtime("streamable-http") + token = BzmToken("key-id", "key-secret") + request = SimpleNamespace( + state=SimpleNamespace( + **{ + BZM_TOKEN_STATE_ATTR: token, + BZM_USER_CONFIG_STATE_ATTR: {"confirmation_mode": "CUD"}, + } + ) + ) + ctx = SimpleNamespace(request_context=SimpleNamespace(request=request)) + + user_config = runtime.configure_context(ctx) + + assert user_config["confirmation_mode"] == "CUD" + assert user_config["token"] is token + assert getattr(ctx.request_context, BZM_USER_CONFIG_STATE_ATTR, None) == user_config + assert getattr(request.state, BZM_USER_CONFIG_STATE_ATTR, None) == user_config + + def test_configure_context_hydrates_request_context_when_ctx_is_strict(self, monkeypatch): + monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") + monkeypatch.setattr(HttpSessionStorageProvider, "ensure_available", lambda self: None) + runtime = build_runtime("streamable-http") + token = BzmToken("key-id", "key-secret") + request = SimpleNamespace( + state=SimpleNamespace( + **{ + BZM_TOKEN_STATE_ATTR: token, + BZM_USER_CONFIG_STATE_ATTR: {"confirmation_mode": "CUD"}, + } + ) + ) + ctx = _StrictCtx(request_context=SimpleNamespace(request=request)) + + user_config = runtime.configure_context(ctx) + manager = Manager(ctx) + + assert user_config["token"] is token + assert getattr(ctx.request_context, BZM_USER_CONFIG_STATE_ATTR, None) == user_config + assert manager.token is token diff --git a/tests/test_main_transport.py b/tests/test_main_transport.py index a1ee8fe..bd197e1 100644 --- a/tests/test_main_transport.py +++ b/tests/test_main_transport.py @@ -20,7 +20,6 @@ def run(self, transport="stdio", mount_path=None): def _patch_mcp_server_dependencies(monkeypatch): monkeypatch.setattr(main, "init_telemetry", lambda *a, **k: None) monkeypatch.setattr(main, "get_token", lambda: object()) - monkeypatch.setattr(main, "register_confirm_mode", lambda *a, **k: None) monkeypatch.setattr(main, "register_tools", lambda *a, **k: None) monkeypatch.setattr(main, "FastMCP", _DummyFastMCP) monkeypatch.setenv("BZM_STORAGE_API_BASE_URL", "https://mcp-storage.internal") diff --git a/tests/test_storage.py b/tests/test_storage.py index 649a9cb..6815120 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -112,7 +112,6 @@ def test_stdio_runtime_gets_local_storage(self, monkeypatch): class TestUploadAssetsHostedRejection: def test_upload_assets_returns_clear_error_on_http_storage(self): manager = TestManager( - token=None, ctx=None, file_access=StorageFileSource("https://mcp-storage.internal"), scope_resolver=DefaultSessionScopeResolver(), @@ -128,3 +127,10 @@ async def _fake_read(_test_id): ) assert "error" in result assert "No valid files found to upload" in result["error"] + + def test_upload_assets_without_file_ports_returns_hosted_message(self): + manager = TestManager(ctx=None) + result = asyncio.run( + manager.upload_assets(1, ["/tmp/demo.jmx"], main_script=None) + ) + assert result["error"] == HOSTED_FILE_ACCESS_MESSAGE diff --git a/tests/test_upload_assets_security.py b/tests/test_upload_assets_security.py index c3a32e6..c1f5110 100644 --- a/tests/test_upload_assets_security.py +++ b/tests/test_upload_assets_security.py @@ -57,7 +57,6 @@ def test_validate_files_classifies_valid_invalid_and_blocked(self, tmp_path): blocked_files = [] manager = UploadAssetsManager( - token=None, ctx=None, file_access=LocalPathFileSource(), scope_resolver=DefaultSessionScopeResolver(), diff --git a/tools/account_manager.py b/tools/account_manager.py index 4befc05..42662e9 100644 --- a/tools/account_manager.py +++ b/tools/account_manager.py @@ -18,7 +18,6 @@ from mcp.server.fastmcp import Context from config.blazemeter import ACCOUNTS_ENDPOINT, TOOLS_PREFIX, SUPPORT_MESSAGE -from config.token import BzmToken from config.runtime import AppRuntime from formatters.account import format_accounts from models.manager import Manager @@ -33,8 +32,11 @@ class AccountManager(Manager): # the format_accounts only expose minimum information to user # The read operation verify permissions and don't allow to share if don't have permissions. - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) async def read(self, account_id: Optional[int]) -> BaseResult: if not isinstance(account_id, int) or account_id < 1: @@ -98,7 +100,8 @@ def register(mcp, runtime: AppRuntime) -> None: """ ) async def account(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - account_manager = AccountManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + account_manager = AccountManager(ctx) async def _dispatch(): match action: diff --git a/tools/billing_manager.py b/tools/billing_manager.py index afde982..6dbf269 100644 --- a/tools/billing_manager.py +++ b/tools/billing_manager.py @@ -13,13 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. """ -from typing import Optional, Dict, Any +from typing import Dict, Any import httpx from mcp.server.fastmcp import Context from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE -from config.token import BzmToken from config.runtime import AppRuntime from models.manager import Manager from models.result import BaseResult @@ -30,8 +29,11 @@ class BillingManager(Manager): - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) async def calculate_cost_from_config(self, args: Dict) -> BaseResult: result = calculate_test_cost(args) @@ -88,7 +90,8 @@ def register(mcp, runtime: AppRuntime) -> None: """ ) async def billing(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - billing_manager = BillingManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + billing_manager = BillingManager(ctx) async def _dispatch(): match action: diff --git a/tools/bridge.py b/tools/bridge.py index 39a30ee..f8a29b0 100644 --- a/tools/bridge.py +++ b/tools/bridge.py @@ -15,40 +15,79 @@ """ from mcp.server.fastmcp import Context +from config.auth import BZM_TOKEN_STATE_ATTR, BZM_USER_CONFIG_STATE_ATTR from config.token import BzmToken from models.result import BaseResult +from types import SimpleNamespace # NOTE: Imports are performed locally in each method to avoid cyclical import problems. # This file currently acts as a bridge between different managers to access specific methods, # primarily for validation of reference elements. +def _with_token_context(ctx: Context, token: BzmToken) -> Context: + context = ctx or SimpleNamespace() + request_context = getattr(context, "request_context", None) + request = getattr(request_context, "request", None) + request_state = getattr(request, "state", None) + + current_config = ( + getattr(request_context, BZM_USER_CONFIG_STATE_ATTR, None) + or getattr(request_state, BZM_USER_CONFIG_STATE_ATTR, None) + or getattr(context, "user_config", None) + or {} + ) + user_config = dict(current_config) if isinstance(current_config, dict) else {} + user_config["token"] = token + + if request_context is not None: + setattr(request_context, BZM_USER_CONFIG_STATE_ATTR, dict(user_config)) + + if request_state is not None: + setattr(request_state, BZM_USER_CONFIG_STATE_ATTR, dict(user_config)) + setattr(request_state, BZM_TOKEN_STATE_ATTR, token) + + # Keep compatibility for non-FastMCP contexts used in tests. + try: + setattr(context, "user_config", user_config) + except Exception: + pass + + return context + + async def read_account(token: BzmToken, ctx: Context, account_id: int) -> BaseResult: from tools.account_manager import AccountManager - return await AccountManager(token, ctx).read(account_id) + return await AccountManager(_with_token_context(ctx, token)).read(account_id) async def read_project(token: BzmToken, ctx: Context, project_id: int) -> BaseResult: from tools.project_manager import ProjectManager - return await ProjectManager(token, ctx).read(project_id) + return await ProjectManager(_with_token_context(ctx, token)).read(project_id) async def read_workspace(token: BzmToken, ctx: Context, workspace_id: int) -> BaseResult: from tools.workspace_manager import WorkspaceManager - return await WorkspaceManager(token, ctx).read(workspace_id) + return await WorkspaceManager(_with_token_context(ctx, token)).read(workspace_id) async def read_test(token: BzmToken, ctx: Context, test_id: int) -> BaseResult: from tools.test_manager import TestManager - return await TestManager(token, ctx).read(test_id) + return await TestManager(_with_token_context(ctx, token)).read(test_id) async def count_project_tests(token: BzmToken, ctx: Context, project_id: int) -> int: from tools.test_manager import TestManager return ( - await TestManager(token, ctx).list(project_id=project_id, limit=1, offset=0, control_ai_consent=False)).total + await TestManager(_with_token_context(ctx, token)).list( + project_id=project_id, + limit=1, + offset=0, + control_ai_consent=False, + ) + ).total async def read_execution(token: BzmToken, ctx: Context, execution_id: int) -> BaseResult: from tools.execution_manager import ExecutionManager - return await ExecutionManager(token, ctx).read(execution_id) + return await ExecutionManager(_with_token_context(ctx, token)).read(execution_id) diff --git a/tools/execution_manager.py b/tools/execution_manager.py index b84a15a..4f6e427 100644 --- a/tools/execution_manager.py +++ b/tools/execution_manager.py @@ -19,7 +19,6 @@ from mcp.server.fastmcp import Context from config.blazemeter import TOOLS_PREFIX, EXECUTIONS_ENDPOINT, SUPPORT_MESSAGE -from config.token import BzmToken from config.runtime import AppRuntime from formatters.execution import format_executions, format_executions_detailed, format_executions_status from models.manager import Manager @@ -32,19 +31,23 @@ class ExecutionManager(Manager): - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) async def _request_log_analyzer_api(self, method: str, execution_id: int, json_body: Optional[Dict[str, Any]] = None) -> BaseResult: - if not self.token: + token = self.token + if not token: return BaseResult( error="No API token. Set BLAZEMETER_API_KEY env var with file path or API_KEY_ID and API_KEY_SECRET secrets." ) url = f"https://log-analyzer.blazemeter.com/analyzer/{execution_id}" headers = { - "Authorization": self.token.as_basic_auth(), + "Authorization": token.as_basic_auth(), "User-Agent": user_agent, "Accept": "application/json", "Content-Type": "application/json" @@ -229,7 +232,12 @@ async def search_filter_values(self, account_id: int, filter_names: List[str]) - if account_data.error: return account_data - return await search_utils.test_execution_search_filter_values("master", account_id, self.token, filter_names) + return await search_utils.test_execution_search_filter_values( + "master", + account_id, + self.token, + filter_names, + ) async def ai_analysis(self, execution_id: Optional[int]) -> BaseResult: if not isinstance(execution_id, int) or execution_id < 1: @@ -396,7 +404,7 @@ async def read_all_reports(self, execution_id: Optional[int]) -> BaseResult: if not isinstance(execution_id, int) or execution_id < 1: return BaseResult(error="Missing or invalid required argument 'execution_id'. Expected integer.") - report_manager = ReportManager(self.token, self.ctx) + report_manager = ReportManager(self.ctx) summary_result = await report_manager.read_summary(execution_id) error_result = await report_manager.read_error(execution_id) stats_result = await report_manager.read_request_stats(execution_id) @@ -497,8 +505,9 @@ def register(mcp, runtime: AppRuntime): """ ) async def execution(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - execution_manager = ExecutionManager(runtime.auth.get_token(ctx), ctx) - report_manager = ReportManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + execution_manager = ExecutionManager(ctx) + report_manager = ReportManager(ctx) async def _dispatch(): match action: diff --git a/tools/help_manager.py b/tools/help_manager.py index 53d8823..81b9969 100644 --- a/tools/help_manager.py +++ b/tools/help_manager.py @@ -16,7 +16,7 @@ import asyncio from copy import deepcopy from itertools import chain -from typing import Optional, Any, Dict, List +from typing import Any, Dict, List import httpx from mcp.server.fastmcp import Context @@ -24,7 +24,6 @@ from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE, \ HELP_INDEX_URL, HELP_TOC_URL, HELP_BASE_CONTENT_URL -from config.token import BzmToken from config.runtime import AppRuntime from formatters.help import format_help_info from models.manager import Manager @@ -45,8 +44,11 @@ class HelpManager(Manager): "Help content is sourced from curated BlazeMeter documentation domains and is trusted by design." ) - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) async def _load_help_tree(self): help_index_url = HELP_INDEX_URL @@ -291,7 +293,8 @@ async def help_main( if args is None: args = {} - help_manager = HelpManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + help_manager = HelpManager(ctx) async def _dispatch(): match action: diff --git a/tools/project_manager.py b/tools/project_manager.py index 3b59e0e..a8c9d8c 100644 --- a/tools/project_manager.py +++ b/tools/project_manager.py @@ -19,7 +19,6 @@ from mcp.server.fastmcp import Context from config.blazemeter import TOOLS_PREFIX, PROJECTS_ENDPOINT -from config.token import BzmToken from config.runtime import AppRuntime from formatters.project import format_projects from models.manager import Manager @@ -31,8 +30,11 @@ class ProjectManager(Manager): - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) async def read(self, project_id: Optional[int]) -> BaseResult: if not isinstance(project_id, int) or project_id < 1: @@ -106,7 +108,8 @@ def register(mcp, runtime: AppRuntime): """ ) async def project(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - project_manager = ProjectManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + project_manager = ProjectManager(ctx) async def _dispatch(): match action: diff --git a/tools/report_manager.py b/tools/report_manager.py index 8ea5161..8b7bec7 100644 --- a/tools/report_manager.py +++ b/tools/report_manager.py @@ -13,12 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. """ -from typing import Optional +from typing import Any, Optional from mcp.server.fastmcp import Context from config.blazemeter import EXECUTIONS_ENDPOINT -from config.token import BzmToken from formatters.execution import ( format_summary_report, format_request_stats, @@ -36,8 +35,11 @@ class ReportManager(Manager): - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) def _extract_execution_name(self, execution_result: BaseResult) -> Optional[str]: """Extract execution name from execution result if available.""" diff --git a/tools/skills_manager.py b/tools/skills_manager.py index 8fbdc14..993935f 100644 --- a/tools/skills_manager.py +++ b/tools/skills_manager.py @@ -22,7 +22,6 @@ from pydantic import Field from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE -from config.token import BzmToken from config.runtime import AppRuntime from models.manager import Manager from models.result import BaseResult @@ -43,8 +42,11 @@ class SkillsManager(Manager): "Skills content is sourced from curated repository resources and is trusted by design." ) - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) @staticmethod async def list_skills() -> BaseResult: @@ -204,7 +206,8 @@ async def skills( if args is None: args = {} - skills_manager = SkillsManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + skills_manager = SkillsManager(ctx) async def _dispatch(): match action: diff --git a/tools/test_manager.py b/tools/test_manager.py index 0b54509..ccc2190 100644 --- a/tools/test_manager.py +++ b/tools/test_manager.py @@ -25,8 +25,7 @@ from config.blazemeter import TESTS_ENDPOINT, TOOLS_PREFIX from config.file_access import FileAccessPort from config.security import detect_sensitive_upload_path_reason -from config.storage import SessionScopeResolverPort -from config.token import BzmToken +from config.storage import HOSTED_FILE_ACCESS_MESSAGE, SessionScopeResolverPort from config.runtime import AppRuntime from formatters.failure_criteria_labels import failure_criteria_meta_payload from formatters.test import format_tests @@ -54,12 +53,13 @@ class TestManager(Manager): def __init__( self, - token: Optional[BzmToken], ctx: Context, - file_access: FileAccessPort, - scope_resolver: SessionScopeResolverPort, + file_access: Optional[FileAccessPort] = None, + scope_resolver: Optional[SessionScopeResolverPort] = None, ): - super().__init__(token, ctx) + super().__init__(ctx) + # Upload ports are stdio-only today. HTTP create/list/read must work + # without them; hosted file upload will be a separate tool later. self.file_access = file_access self.scope_resolver = scope_resolver @@ -221,6 +221,8 @@ async def upload_assets( return { "error": "Missing or invalid required argument 'file_paths'. Expected non-empty list." } + if self.file_access is None or self.scope_resolver is None: + return {"error": HOSTED_FILE_ACCESS_MESSAGE} # Check if it's valid or allowed test_data = await self.read(test_id) @@ -669,12 +671,13 @@ def register(mcp, runtime: AppRuntime): """, ) async def tests(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - test_manager = TestManager( - runtime.auth.get_token(ctx), - ctx, - runtime.file_access, - runtime.scope_resolver, - ) + runtime.configure_context(ctx) + if runtime.transport == "stdio": + test_manager = TestManager( + ctx, runtime.file_access, runtime.scope_resolver + ) + else: + test_manager = TestManager(ctx) async def _dispatch(): match action: diff --git a/tools/user_manager.py b/tools/user_manager.py index e314b12..d971559 100644 --- a/tools/user_manager.py +++ b/tools/user_manager.py @@ -13,14 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. """ -from typing import Any, Dict, Optional +from typing import Any, Dict import httpx from mcp.server.fastmcp import Context from pydantic import Field from config.blazemeter import TOOLS_PREFIX, USER_ENDPOINT -from config.token import BzmToken from config.runtime import AppRuntime from formatters.user import format_users from models.manager import Manager @@ -31,8 +30,11 @@ class UserManager(Manager): - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) async def read(self) -> BaseResult: return await api_request( @@ -61,7 +63,8 @@ async def user( ctx: Context = Field(description="Context object providing access to MCP capabilities") ) -> BaseResult: - user_manager = UserManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + user_manager = UserManager(ctx) async def _dispatch(): match action: diff --git a/tools/utils.py b/tools/utils.py index 37d8599..049b8d7 100644 --- a/tools/utils.py +++ b/tools/utils.py @@ -24,7 +24,7 @@ import traceback from datetime import datetime, timezone from enum import Enum -from typing import Optional, Callable, Awaitable +from typing import Any, Optional, Callable, Awaitable from importlib import resources from pathlib import Path @@ -32,6 +32,7 @@ from pydantic import BaseModel from config.blazemeter import BZM_API_BASE_URL +from config.context_resolution import resolve_ctx_user_config from config.security import validate_http_request_endpoint from config.token import BzmToken from config.version import __version__ @@ -129,9 +130,6 @@ class ConfirmMode(Enum): DISABLE = "NONE" # No confirmation -_confirm_mode = ConfirmMode.DELETE - - class Operations(Enum): CREATE = "C" # Create READ = "R" # Read @@ -270,18 +268,44 @@ class Confirmation(BaseModel): pass # Empty model with no fields for simple accept/cancel without UI elements -def register_confirm_mode(confirm_mode_value: ConfirmMode): - global _confirm_mode - _confirm_mode = confirm_mode_value +def _to_confirm_mode(value: Any) -> ConfirmMode: + if isinstance(value, ConfirmMode): + return value + if isinstance(value, str): + normalized = value.strip().upper() + if normalized in ConfirmMode.__members__: + return ConfirmMode[normalized] + for mode in ConfirmMode: + if normalized == mode.value: + return mode + return ConfirmMode.DELETE + + +def _get_ctx_user_config(ctx: Any) -> dict[str, Any] | None: + user_config = resolve_ctx_user_config(ctx) + if user_config: + return user_config + return None -def get_confirm_mode() -> ConfirmMode: - global _confirm_mode - return _confirm_mode +def resolve_confirmation_mode(ctx: Any, manager_user_config: Any = None) -> ConfirmMode: + """ + Resolve confirmation mode from runtime/user session context. + + Precedence: + 1) per-request/per-session context user config + 2) manager-level user config (stdio startup config) + 3) DELETE default + """ + ctx_user_config = _get_ctx_user_config(ctx) + if isinstance(ctx_user_config, dict): + return _to_confirm_mode(ctx_user_config.get("confirmation_mode")) + if isinstance(manager_user_config, dict): + return _to_confirm_mode(manager_user_config.get("confirmation_mode")) + return ConfirmMode.DELETE -def operation_need_confirmation(operation: Operations) -> bool: - confirm_mode = get_confirm_mode() +def operation_need_confirmation(operation: Operations, confirm_mode: ConfirmMode) -> bool: if confirm_mode == ConfirmMode.DELETE and operation in [Operations.DELETE]: return True elif confirm_mode == ConfirmMode.CUD and operation in [Operations.CREATE, Operations.UPDATE, Operations.DELETE]: @@ -301,7 +325,11 @@ def require_confirmation(operation: Operations = Operations.READ, def decorator(func: Callable[..., Awaitable]): @functools.wraps(func) async def wrapper(self, *args, **kwargs): - need_confirmation = operation_need_confirmation(operation) + confirm_mode = resolve_confirmation_mode( + getattr(self, "ctx", None), + getattr(self, "user_config", None), + ) + need_confirmation = operation_need_confirmation(operation, confirm_mode) confirmed = True # Run operation by default if need_confirmation: try: diff --git a/tools/workspace_manager.py b/tools/workspace_manager.py index 8a8a497..84da19a 100644 --- a/tools/workspace_manager.py +++ b/tools/workspace_manager.py @@ -20,7 +20,6 @@ from pydantic import Field from config.blazemeter import WORKSPACES_ENDPOINT, TOOLS_PREFIX -from config.token import BzmToken from config.runtime import AppRuntime from formatters.workspace import format_workspaces, format_workspaces_detailed, format_workspaces_locations from models.manager import Manager @@ -36,8 +35,11 @@ class WorkspaceManager(Manager): # the format_workspaces only expose minimum information to user # The read operation verify permissions and don't allow to share details. - def __init__(self, token: Optional[BzmToken], ctx: Context): - super().__init__(token, ctx) + def __init__( + self, + ctx: Context, + ): + super().__init__(ctx) async def read(self, workspace_id: Optional[int]) -> BaseResult: if not isinstance(workspace_id, int) or workspace_id < 1: @@ -139,7 +141,8 @@ async def workspace( ctx: Context = Field(description="Context object providing access to MCP capabilities") ) -> BaseResult: - workspace_manager = WorkspaceManager(runtime.auth.get_token(ctx), ctx) + runtime.configure_context(ctx) + workspace_manager = WorkspaceManager(ctx) async def _dispatch(): match action: From 66586d1f416830e4b85c7cfb80abe4f7663b634e Mon Sep 17 00:00:00 2001 From: alejandroaires Date: Thu, 27 Aug 2026 10:21:15 -0300 Subject: [PATCH 5/6] Use http2 for ensure_availability for storage-api connection --- config/storage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/config/storage.py b/config/storage.py index fd3e711..32c4d10 100644 --- a/config/storage.py +++ b/config/storage.py @@ -305,7 +305,7 @@ def _health_url(self) -> str: def ensure_available(self) -> None: """Fail fast if the storage API is unreachable.""" - with httpx.Client(timeout=min(self._timeout, 5.0)) as client: + with httpx.Client(http2=True, timeout=min(self._timeout, 5.0)) as client: response = client.get(self._health_url()) response.raise_for_status() From 75ab7742f530e53ab020d49e8eb65b8032fe7877 Mon Sep 17 00:00:00 2001 From: Diego Ferrand Date: Mon, 7 Sep 2026 11:10:46 -0300 Subject: [PATCH 6/6] Dataframe storage port (#99) * Add session StoragePort and Storage-backed dataframes for hosted MCP. * Align Storage client path and uploaded_files with storage-api contract. * Isolate session dataframe state and wire result materialization. * Split storage ports and harden session dataframe persistence. * Merge STREAMABLE_HTTP and persist dataframes via SessionStoragePort. * Harden session dataframe persistence against multi-worker Storage and keep tracing free of dataframe policy. * Pass tool_args through managers so result_format can store or skip dataframes. * Drop the JMeter load-test section from the hosted runbook until that plan ships. * Use SessionStoragePort names in dataframe tools so they are not confused with file-access HttpStorageClient. * Tasks storage port (#102) * Rebase async task storage onto SessionScope partitions and shared MCP entrypoints. * Harden hosted task persistence with merge-on-commit and SessionScope. * Fix PR comments * Fixed PR comments * Add overflow lock for session management --- .gitignore | 9 +- config/storage.py | 26 +- docs/hosted-http.md | 12 +- docs/hosted-mvp-runbook.md | 88 ++ main.py | 19 + models/result.py | 77 ++ pyproject.toml | 1 + server.py | 4 + tests/conftest.py | 75 ++ tests/storage_fakes.py | 79 ++ tests/test_async_task_storage.py | 505 +++++++++ tests/test_batch_controls.py | 4 +- tests/test_dataframe_session_fixes.py | 301 ++++++ tests/test_dataframe_storage.py | 367 +++++++ tests/test_failure_criteria.py | 2 +- tests/test_required_args_tools.py | 28 +- tests/test_skills_manager_security.py | 70 +- tests/test_storage.py | 17 +- tests/test_support_message.py | 182 ++++ tests/test_tool_result_wrapper.py | 35 + tests/test_tools_manager_dataframes.py | 84 ++ tests/test_tools_manager_tasks.py | 264 +++++ tests/test_utils_normalize_action_args.py | 72 ++ tests/test_utils_required_args.py | 32 + tools/account_manager.py | 57 +- tools/async_task_manager.py | 621 +++++++++++ tools/billing_manager.py | 49 +- tools/dataframe_manager.py | 1164 +++++++++++++++++++++ tools/execution_manager.py | 105 +- tools/help_manager.py | 155 ++- tools/mcp_entrypoint.py | 98 ++ tools/project_manager.py | 57 +- tools/report_manager.py | 6 +- tools/runtime_tools.py | 69 ++ tools/skills_manager.py | 172 ++- tools/test_manager.py | 157 +-- tools/tools_manager.py | 746 +++++++++++++ tools/user_manager.py | 57 +- tools/utils.py | 433 +++++++- tools/workspace_manager.py | 72 +- uv.lock | 30 + 41 files changed, 5892 insertions(+), 509 deletions(-) create mode 100644 docs/hosted-mvp-runbook.md create mode 100644 tests/conftest.py create mode 100644 tests/storage_fakes.py create mode 100644 tests/test_async_task_storage.py create mode 100644 tests/test_dataframe_session_fixes.py create mode 100644 tests/test_dataframe_storage.py create mode 100644 tests/test_support_message.py create mode 100644 tests/test_tool_result_wrapper.py create mode 100644 tests/test_tools_manager_dataframes.py create mode 100644 tests/test_tools_manager_tasks.py create mode 100644 tests/test_utils_normalize_action_args.py create mode 100644 tests/test_utils_required_args.py create mode 100644 tools/async_task_manager.py create mode 100644 tools/dataframe_manager.py create mode 100644 tools/mcp_entrypoint.py create mode 100644 tools/runtime_tools.py create mode 100644 tools/tools_manager.py diff --git a/.gitignore b/.gitignore index 31ed30c..9232607 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,11 @@ whitesource/ .whitesource.lock mend mend-cli -*.mend.json \ No newline at end of file +*.mend.json +# Local credentials (never commit) +dist/api-key.json +api-key.json + +# Local/PyInstaller build artifacts +bzm-mcp-linux-arm64.spec +*.spec \ No newline at end of file diff --git a/config/storage.py b/config/storage.py index 32c4d10..b3e387d 100644 --- a/config/storage.py +++ b/config/storage.py @@ -41,6 +41,10 @@ class StorageNotSupportedError(NotImplementedError): """Raised when a storage backend cannot fulfill a file operation.""" +class StorageNotConfiguredError(RuntimeError): + """Raised when session storage is used before AppRuntime wiring.""" + + @runtime_checkable class FileStoragePort(Protocol): """ @@ -213,7 +217,7 @@ async def delete_partition(self, scope: SessionScope) -> bool: class SessionScopeResolverPort(ABC): @abstractmethod - def resolve(self, ctx: Context, token: Optional[BzmToken]) -> SessionScope: + def resolve(self, ctx: Optional[Context], token: Optional[BzmToken]) -> SessionScope: raise NotImplementedError @@ -226,7 +230,9 @@ class DefaultSessionScopeResolver(SessionScopeResolverPort): """ @staticmethod - def _resolve_session_id(ctx: Context) -> str: + def _resolve_session_id(ctx: Optional[Context]) -> str: + if ctx is None: + return "default" request = getattr(getattr(ctx, "request_context", None), "request", None) if request is not None: session_id = request.headers.get("mcp-session-id") @@ -243,13 +249,27 @@ def _resolve_user_id(token: Optional[BzmToken]) -> str: return token.id.strip() return "anonymous" - def resolve(self, ctx: Context, token: Optional[BzmToken]) -> SessionScope: + def resolve(self, ctx: Optional[Context], token: Optional[BzmToken]) -> SessionScope: return SessionScope( user_id=self._resolve_user_id(token), mcp_session_id=self._resolve_session_id(ctx), ) +def resolve_session_scope( + ctx: Any, + token: Optional[BzmToken] = None, + scope_resolver: Optional[SessionScopeResolverPort] = None, +) -> SessionScope: + """Resolve partition keys from auth token + MCP session context.""" + resolver = scope_resolver or DefaultSessionScopeResolver() + resolved_token = token + if resolved_token is None: + from config.context_resolution import resolve_ctx_token + resolved_token = resolve_ctx_token(ctx) + return resolver.resolve(ctx, resolved_token) + + class InMemorySessionStorageProvider(SessionStoragePort): def __init__(self) -> None: self._partitions: dict[tuple[str, str], SessionPartition] = {} diff --git a/docs/hosted-http.md b/docs/hosted-http.md index ea7849f..1cc4f17 100644 --- a/docs/hosted-http.md +++ b/docs/hosted-http.md @@ -55,6 +55,7 @@ docker run --rm -p 8000:8000 \ -e FASTMCP_HOST=0.0.0.0 \ -e FASTMCP_PORT=8000 \ -e FASTMCP_STREAMABLE_HTTP_PATH=/mcp \ + -e BZM_STORAGE_API_BASE_URL=https://mcp-storage.internal \ -e BZM_STORAGE_STRATEGY=memory \ ghcr.io/blazemeter/bzm-mcp:latest ``` @@ -65,13 +66,14 @@ docker run --rm -p 8000:8000 \ |----------|-------------|---------| | `BZM_MCP_TRANSPORT` | Logical transport: `stdio`, `http`, or `docker` | `stdio` | | `FASTMCP_HOST` | Bind address (HTTP only) | `127.0.0.1` | -| `FASTMCP_PORT` | Listen port (HTTP only). Also accepts `PORT` (e.g. Cloud Run) | `8000` | +| `FASTMCP_PORT` | Listen port (HTTP only). Also accepts `PORT` | `8000` | | `FASTMCP_STREAMABLE_HTTP_PATH` | HTTP path for the MCP endpoint | `/mcp` | -| `BZM_STORAGE_STRATEGY` | `memory` or `http` | `memory` | +| `BZM_STORAGE_API_BASE_URL` | Storage Service base URL (required for streamable-http) | — | +| `BZM_STORAGE_STRATEGY` | `memory` or `http` (file-access helper; session store follows transport) | `memory` | -On streamable-http, local file paths are always rejected regardless of `BZM_STORAGE_STRATEGY` (hosted fail-closed storage). +On streamable-http, session partitions are stored via `HttpSessionStorageProvider`. Local file paths are always rejected (`StorageFileSource`) regardless of `BZM_STORAGE_STRATEGY` (hosted fail-closed storage). ## Hosted MVP limitations -- In-memory / fail-closed storage: no local disk access on the shared hosted server. -- `upload_assets` and other local file lookup/upload paths are rejected. Use a local stdio or Docker MCP installation for those workflows, or wait for remote Storage (Phase 2). +- Session dataframes/tasks live in the Storage Service keyed by `{user_id}/{mcp_session_id}`. +- `upload_assets` and other local file lookup/upload paths are rejected. Use a local stdio or Docker MCP installation for those workflows, or wait for remote file access. diff --git a/docs/hosted-mvp-runbook.md b/docs/hosted-mvp-runbook.md new file mode 100644 index 0000000..d317cfb --- /dev/null +++ b/docs/hosted-mvp-runbook.md @@ -0,0 +1,88 @@ +# Hosted MCP — session storage + dataframes + +Hosted MCP ops deploy lives in +[`hosted-bzm-mcp`](https://github.com/Blazemeter/hosted-bzm-mcp). +This file documents how **bzm-mcp** uses the Storage Service introduced on +`STREAMABLE_HTTP`. + +## Naming (do not rename STREAMABLE_HTTP types) + +Dataframe tools depend on `AppRuntime.storage: SessionStoragePort`. Story names +map to the types already landed on `STREAMABLE_HTTP`: + +| Story / plan name | STREAMABLE_HTTP type | Used for dataframes | +|-------------------|----------------------|---------------------| +| StoragePort | `SessionStoragePort` | Yes | +| MemoryStorageProvider | `InMemorySessionStorageProvider` | Yes (stdio) | +| HttpStorageProvider | `HttpSessionStorageProvider` | Yes (hosted) | +| HTTPStorageClient (story) | `HttpSessionStorageProvider` | Yes — not `HttpStorageClient` | +| HttpStorageClient (codebase) | `FileStoragePort` stub | No (file access, Phase 3) | + +## Runtime wiring + +| Mode | Transport | Session store | File access | +|------|-----------|---------------|-------------| +| Stdio / local Docker | `stdio` | `InMemorySessionStorageProvider` | `LocalPathFileSource` / Docker mapped paths | +| Hosted HTTP | `streamable-http` | `HttpSessionStorageProvider` | `StorageFileSource` | + +Composition root: `build_runtime` → `AppRuntime.storage` / `scope_resolver`. +`server.register_tools` also calls `configure_task_storage(runtime.storage)` so +async tasks share the same session partitions as dataframes. +Tool registrations call `run_tool_with_runtime(runtime, ...)` so tracing stays +unaware of dataframe types. There is no process-global dataframe store. + +Partition key: `{user_id}/{mcp_session_id}` via `DefaultSessionScopeResolver` +(`Mcp-Session-Id` header, then FastMCP `ctx.session_id`). + +## Session Storage Service + +Env: `BZM_STORAGE_API_BASE_URL` (required for streamable-http). + +| Method | Path | +|--------|------| +| `GET` / `PUT` / `DELETE` | `/session-partitions/{user_id}/{mcp_session_id}` | + +`put_partition` accepts a partial `SessionPartitionPayload`. Dataframe tools +send only `dataframes`; the async task runner sends only `tasks`. Other +sections are preserved by the storage merge. + +MCP workers keep Polars/SQL and asyncio task handles in-process; only the +partition document is remote. + +### Task map concurrency + +`put_partition(tasks=...)` replaces the **entire** tasks map for that +partition. In-process locks serialize mutations on one worker. Across hosted MCP +workers, commit re-reads `SessionStoragePort` and unions keys added by other +writers (and drops ids this operation removed). Live local `asyncio.Task` +handles win on overlapping keys. Same-key concurrent writes and the GET/PUT +race can still last-write-win. Closing that window needs Storage CAS/etag. + +## Task execution affinity (hosted) + +- **Status / list / get** are Storage-backed and work across workers for a session. +- **In-flight execution** (`asyncio.Task`, semaphore, cancel of a live handle) is + **process-local**. Only the worker that started the coroutine can cancel it via + the local handle. +- Calling `tasks_cancel` on a different worker marks cancel in Storage when there + is no local handle, but the owning worker may still finish and overwrite status + to `completed` / `failed`. Do not assume multi-worker cancel stops execution. +- Prefer sticky routing / single-writer affinity for a session while tasks are + active if cancel must be reliable. + +### Dataframe map concurrency + +`put_partition(dataframes=...)` replaces the **entire** dataframes map for that +partition. In-process locks serialize mutations on one worker. Across hosted MCP +workers, commit re-reads `SessionStoragePort` and unions keys added by other +writers (and drops ids this operation removed). Same-key concurrent writes and +the GET/PUT race can still last-write-win. Closing that window needs Storage +CAS/etag, which this service does not expose yet. + +## Dataframe tools + +`dataframes_list`, `dataframes_query`, `dataframes_remove`, and +`dataframes_clear` hydrate/commit through `SessionStoragePort`. Stdio uses +`InMemorySessionStorageProvider`; hosted uses `HttpSessionStorageProvider`. +Missing session storage on a path that would persist fails closed (error, not raw payload). + diff --git a/main.py b/main.py index 472d480..d6772ca 100644 --- a/main.py +++ b/main.py @@ -27,6 +27,25 @@ from pathlib import Path from typing import Literal, cast +# Patch MCP ArgModelBase so tools with an "arguments" param receive the full payload +# when the client sends {"action": "x", "key": "value"} instead of {"arguments": {...}} +from mcp.server.fastmcp.utilities import func_metadata +from pydantic import model_validator + +_OriginalArgModelBase = func_metadata.ArgModelBase + + +class _PatchedArgModelBase(_OriginalArgModelBase): + @model_validator(mode="before") + @classmethod + def _wrap_root_as_arguments(cls, data: object) -> object: + if isinstance(data, dict) and "arguments" not in data: + return {"arguments": data} + return data + + +func_metadata.ArgModelBase = _PatchedArgModelBase + from mcp.server.fastmcp import FastMCP from config.auth import run_streamable_http diff --git a/models/result.py b/models/result.py index b8072b4..9d8e716 100644 --- a/models/result.py +++ b/models/result.py @@ -15,6 +15,7 @@ """ from typing import Any, Optional, List +from mcp.types import CallToolResult, TextContent from pydantic import BaseModel, Field class BaseResult(BaseModel): @@ -24,6 +25,10 @@ class BaseResult(BaseModel): 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) def append_warnings(self, messages: List[str]): if not self.warning: @@ -44,3 +49,75 @@ def model_dump_json(self, **kwargs): class HttpBaseResult(BaseResult): result: Optional[Any] = Field(description="Result", default=None) + + +class ToolResult(CallToolResult): + @classmethod + def from_base_result(cls, base_result: BaseResult) -> "ToolResult": + compact_text = base_result.model_dump_json(indent=2) + structured = base_result.model_dump(mode="json") + return cls( + content=[TextContent(type="text", text=compact_text)], + structuredContent=structured, + isError=bool(base_result.error), + ) + + @property + def result(self) -> Optional[List[Any]]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("result") + + @property + def total(self) -> Optional[int]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("total") + + @property + def has_more(self) -> Optional[bool]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("has_more") + + @property + def error(self) -> Optional[str]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("error") + + @property + def info(self) -> Optional[List[str]]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("info") + + @property + def warning(self) -> Optional[List[str]]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("warning") + + @property + def tool_call_started_at(self) -> Optional[str]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("tool_call_started_at") + + @property + def tool_call_finished_at(self) -> Optional[str]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("tool_call_finished_at") + + @property + def tool_call_duration_ms(self) -> Optional[int]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("tool_call_duration_ms") + + @property + def debug(self) -> Optional[dict[str, Any]]: + if not isinstance(self.structuredContent, dict): + return None + return self.structuredContent.get("debug") diff --git a/pyproject.toml b/pyproject.toml index 53a3970..914ffcb 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 c5ed472..924bd00 100644 --- a/server.py +++ b/server.py @@ -15,12 +15,14 @@ """ from config.runtime import AppRuntime from tools.account_manager import register as register_account_manager +from tools.async_task_manager import configure_task_storage from tools.billing_manager import register as register_billing_manager from tools.execution_manager import register as register_execution_manager from tools.help_manager import register as register_help_manager from tools.project_manager import register as register_project_manager from tools.skills_manager import register as register_skills_manager from tools.test_manager import register as register_test_manager +from tools.tools_manager import register as register_tools_manager from tools.user_manager import register as register_user_manager from tools.workspace_manager import register as register_workspace_manager @@ -33,6 +35,7 @@ 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) register_user_manager(mcp, runtime) register_project_manager(mcp, runtime) register_workspace_manager(mcp, runtime) @@ -42,3 +45,4 @@ def register_tools(mcp, runtime: AppRuntime): register_billing_manager(mcp, runtime) register_help_manager(mcp, runtime) register_skills_manager(mcp, runtime) + register_tools_manager(mcp, runtime) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..a4f3d2e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,75 @@ +""" +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 +from types import SimpleNamespace + +import pytest + +from config.auth import BZM_TOKEN_STATE_ATTR, BZM_USER_CONFIG_STATE_ATTR +from config.storage import InMemorySessionStorageProvider, SessionScope +from config.token import BzmToken + + +def run_async(coro): + return asyncio.run(coro) + + +def make_ctx(token: BzmToken, session_id: str): + request_state = SimpleNamespace( + **{ + BZM_TOKEN_STATE_ATTR: token, + BZM_USER_CONFIG_STATE_ATTR: {"token": token}, + } + ) + request = SimpleNamespace( + state=request_state, + headers={"mcp-session-id": session_id}, + ) + return SimpleNamespace( + session_id=session_id, + request_context=SimpleNamespace(request=request), + ) + + +@pytest.fixture(autouse=True) +def reset_dataframe_session_locks(): + from tools import dataframe_manager as dataframe_manager_module + + dataframe_manager_module._session_locks.clear() + dataframe_manager_module._overflow_lock = None + yield + dataframe_manager_module._session_locks.clear() + dataframe_manager_module._overflow_lock = None + + +@pytest.fixture(autouse=True) +def _configure_session_task_storage(in_memory_session_storage): + """Ensure @run_as_task can persist when manager methods are called in unit tests.""" + from tools.async_task_manager import configure_task_storage + + configure_task_storage(in_memory_session_storage) + yield in_memory_session_storage + + +@pytest.fixture +def in_memory_session_storage(): + """Stdio-equivalent SessionStoragePort (InMemorySessionStorageProvider).""" + return InMemorySessionStorageProvider() + + +@pytest.fixture +def session_scope(): + return SessionScope(user_id="user-1", mcp_session_id="sess-a") diff --git a/tests/storage_fakes.py b/tests/storage_fakes.py new file mode 100644 index 0000000..30cff54 --- /dev/null +++ b/tests/storage_fakes.py @@ -0,0 +1,79 @@ +""" +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 json + +import httpx + +from config.storage import ( + HttpSessionStorageProvider, + InMemorySessionStorageProvider, + SessionPartitionPayload, + SessionScope, +) + + +class RecordingSessionStorageProvider(InMemorySessionStorageProvider): + def __init__(self) -> None: + super().__init__() + self.put_payloads: list[SessionPartitionPayload] = [] + + async def put_partition(self, scope: SessionScope, payload: SessionPartitionPayload) -> None: + self.put_payloads.append(payload) + await super().put_partition(scope, payload) + + +class MergingSessionTransport(httpx.AsyncBaseTransport): + """Simulates Storage API merge-on-PUT for /session-partitions/{user}/{session}.""" + + def __init__(self) -> None: + self.partitions: dict[str, dict] = {} + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + path = request.url.path + prefix = "/session-partitions/" + if not path.startswith(prefix): + return httpx.Response(404, json={"error": "not found"}) + key = path[len(prefix):] + if request.method == "GET": + if key not in self.partitions: + return httpx.Response(404, json={"error": "not found"}) + return httpx.Response(200, json=self.partitions[key]) + if request.method == "PUT": + incoming = json.loads(request.content.decode("utf-8")) + existing = dict(self.partitions.get(key) or {}) + for section in ("metadata", "dataframes", "tasks", "uploaded_files"): + if section in incoming: + existing[section] = incoming[section] + user_id, _, session_id = key.partition("/") + existing["user_id"] = user_id + existing["mcp_session_id"] = session_id + self.partitions[key] = existing + return httpx.Response(200, json=existing) + if request.method == "DELETE": + deleted = key in self.partitions + self.partitions.pop(key, None) + return httpx.Response(200, json={"deleted": deleted}) + return httpx.Response(405) + + +def http_session_storage_provider(monkeypatch, transport: MergingSessionTransport) -> HttpSessionStorageProvider: + class _Client(httpx.AsyncClient): + def __init__(self, *args, **kwargs): + kwargs["transport"] = transport + super().__init__(*args, **kwargs) + + monkeypatch.setattr("config.storage.httpx.AsyncClient", _Client) + return HttpSessionStorageProvider(base_url="http://storage.test") diff --git a/tests/test_async_task_storage.py b/tests/test_async_task_storage.py new file mode 100644 index 0000000..5e9f91f --- /dev/null +++ b/tests/test_async_task_storage.py @@ -0,0 +1,505 @@ +""" +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 pytest + +import tools.async_task_manager as task_manager +from config.runtime import build_runtime +from config.storage import ( + InMemorySessionStorageProvider, + SessionPartitionPayload, + SessionScope, + StorageNotConfiguredError, +) +from models.result import BaseResult +from tests.storage_fakes import ( + MergingSessionTransport, + RecordingSessionStorageProvider, + http_session_storage_provider, +) +from tools.async_task_manager import ( + STATUS_CANCELLED, + STATUS_COMPLETED, + STATUS_WORKING, + configure_task_storage, + cancel_task, + get_task_record, + list_tasks, + remove_task, + submit_task, + task_snapshot, +) +from tools.dataframe_manager import register_dataframe + + +SCOPE_A = SessionScope(user_id="user-1", mcp_session_id="sess-a") +SCOPE_B = SessionScope(user_id="user-1", mcp_session_id="sess-b") +SCOPE_TASKS = SessionScope(user_id="user-9", mcp_session_id="sess-tasks") + + +def _run(coro): + return asyncio.run(coro) + + +def _stub_task_payload(task_id: str, scope: SessionScope, status: str = STATUS_WORKING) -> dict: + return { + "task_id": task_id, + "action": {"manager": "OtherWorker", "method": "list"}, + "created_at": 1.0, + "last_updated_at": 1.0, + "time_to_live_ms": None, + "status": status, + "status_message": "seeded", + "status_info": "seeded", + "started_running_at": 1.0, + "finished_at": None, + "user_id": scope.user_id, + "mcp_session_id": scope.mcp_session_id, + } + + +async def _wait_terminal(task_id: str, scope: SessionScope, timeout: float = 2.0): + deadline = asyncio.get_event_loop().time() + timeout + while asyncio.get_event_loop().time() < deadline: + record = await get_task_record(task_id, scope=scope) + if record and record.status in {"completed", "failed", "cancelled"}: + return record + await asyncio.sleep(0.01) + raise AssertionError(f"Task {task_id} did not reach a terminal state") + + +@pytest.fixture +def memory_store(): + store = InMemorySessionStorageProvider() + configure_task_storage(store) + return store + + +class TestTaskStorageWiring: + def test_requires_configure_before_use(self): + task_manager._storage = None + with pytest.raises(StorageNotConfiguredError, match="not configured"): + _run(list_tasks(scope=SessionScope(user_id="u", mcp_session_id="s"))) + + def test_register_tools_binds_runtime_storage(self): + from server import register_tools + + class _DummyMcp: + def tool(self, *args, **kwargs): + def deco(fn): + return fn + return deco + + runtime = build_runtime("stdio") + register_tools(_DummyMcp(), runtime) + assert task_manager._storage is runtime.storage + + +class TestAsyncTaskManagerMemoryStorage: + def test_submit_lifecycle_and_persist(self, memory_store): + async def scenario(): + async def action(): + await asyncio.sleep(0.05) + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + action={"manager": "TestManager", "method": "read"}, + coro_factory=action, + scope=SCOPE_A, + ) + assert len(task_id) == 8 + assert all(ch in task_manager.TASK_ID_ALPHABET for ch in task_id) + + record = await _wait_terminal(task_id, SCOPE_A) + assert record.status == STATUS_COMPLETED + assert record.result is not None + assert record.result.result == [{"ok": True}] + + partition = await memory_store.get_partition(SCOPE_A) + assert partition is not None + assert task_id in partition.tasks + assert partition.tasks[task_id]["status"] == STATUS_COMPLETED + + assert await remove_task(task_id, scope=SCOPE_A) is True + partition = await memory_store.get_partition(SCOPE_A) + assert partition is not None + assert task_id not in partition.tasks + + _run(scenario()) + + def test_submit_skips_materialization_when_action_disables_it(self, memory_store): + payload = [{"id": index, "note": "x" * 40} for index in range(300)] + + async def scenario(): + async def action(): + return BaseResult(result=payload) + + task_id = await submit_task( + { + "manager": "SkillsManager", + "method": "read_skill", + "result_format": "auto", + "disable_dataframe_materialization": True, + }, + action, + scope=SCOPE_A, + ) + record = await _wait_terminal(task_id, SCOPE_A) + assert record.status == STATUS_COMPLETED + assert record.result is not None + assert record.result.result == payload + assert record.result.result[0].get("stored_as_dataframe") is not True + partition = await memory_store.get_partition(SCOPE_A) + assert partition is not None + assert partition.dataframes == {} + + _run(scenario()) + + def test_session_isolation(self, memory_store): + async def scenario(): + async def action(): + return BaseResult(result=[{"v": 1}]) + + task_a = await submit_task( + {"manager": "A", "method": "list"}, + action, + scope=SCOPE_A, + ) + task_b = await submit_task( + {"manager": "B", "method": "list"}, + action, + scope=SCOPE_B, + ) + await _wait_terminal(task_a, SCOPE_A) + await _wait_terminal(task_b, SCOPE_B) + + listed_a = await list_tasks(scope=SCOPE_A) + listed_b = await list_tasks(scope=SCOPE_B) + assert [t.task_id for t in listed_a] == [task_a] + assert [t.task_id for t in listed_b] == [task_b] + + _run(scenario()) + + def test_collision_policy_fails_after_ten_attempts(self, memory_store, monkeypatch): + async def scenario(): + cache = await task_manager._get_or_create_cache(SCOPE_A) + async with cache.lock: + cache.hydrated = True + cache.tasks["deadbeef"] = task_manager.TaskRecord( + task_id="deadbeef", + action={"manager": "TestManager", "method": "read"}, + created_at=0.0, + last_updated_at=0.0, + time_to_live_ms=None, + status=task_manager.STATUS_PARKING, + status_message="seed", + status_info="seed", + user_id=SCOPE_A.user_id, + mcp_session_id=SCOPE_A.mcp_session_id, + ) + monkeypatch.setattr(task_manager, "_generate_task_id", lambda: "deadbeef") + with pytest.raises( + RuntimeError, + match="Unable to allocate unique 8-char task id after 10 attempts.", + ): + await task_manager._allocate_task_id(cache) + + _run(scenario()) + + def test_task_lookup_is_case_insensitive(self, memory_store): + async def scenario(): + now = 0.0 + cache = await task_manager._get_or_create_cache(SCOPE_A) + async with cache.lock: + cache.hydrated = True + cache.tasks["7k2p9m4q"] = task_manager.TaskRecord( + task_id="7k2p9m4q", + action={"manager": "ExecutionManager", "method": "list"}, + created_at=now, + last_updated_at=now, + time_to_live_ms=None, + status=STATUS_WORKING, + status_message="running", + status_info="running", + user_id=SCOPE_A.user_id, + mcp_session_id=SCOPE_A.mcp_session_id, + ) + await task_manager._commit_cache(cache, SCOPE_A) + + assert await get_task_record("7K2P9M4Q", scope=SCOPE_A) is not None + assert await remove_task("7K2P9M4Q", scope=SCOPE_A) is True + assert await get_task_record("7k2p9m4q", scope=SCOPE_A) is None + + _run(scenario()) + + def test_status_visible_after_cache_drop(self, memory_store): + """Simulate another worker by dropping the in-process cache and hydrating from Storage.""" + async def scenario(): + async def action(): + await asyncio.sleep(0.05) + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=SCOPE_TASKS, + ) + await _wait_terminal(task_id, SCOPE_TASKS) + + task_manager._session_caches.clear() + record = await get_task_record(task_id, scope=SCOPE_TASKS) + assert record is not None + assert record.status == STATUS_COMPLETED + snap = task_snapshot(record, include_result=True) + assert snap["task_result"]["result"] == [{"ok": True}] + + _run(scenario()) + + def test_submit_preserves_existing_dataframes(self, memory_store): + async def scenario(): + await memory_store.put_partition( + SCOPE_TASKS, + SessionPartitionPayload(dataframes={"df1": {"dataframe_id": "df1", "data": []}}), + ) + + async def action(): + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=SCOPE_TASKS, + ) + await _wait_terminal(task_id, SCOPE_TASKS) + partition = await memory_store.get_partition(SCOPE_TASKS) + assert partition is not None + assert "df1" in partition.dataframes + assert task_id in partition.tasks + + _run(scenario()) + + def test_submit_handle_is_visible_under_lock(self, memory_store): + async def scenario(): + started = asyncio.Event() + + async def action(): + started.set() + await asyncio.sleep(0.05) + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=SCOPE_A, + ) + record = await get_task_record(task_id, scope=SCOPE_A) + assert record is not None + assert record.asyncio_task is not None + await started.wait() + await _wait_terminal(task_id, SCOPE_A) + + _run(scenario()) + + +class TestTaskMapMerge: + def test_status_persist_keeps_keys_added_by_other_worker(self, memory_store): + async def scenario(): + started = asyncio.Event() + + async def action(): + started.set() + await asyncio.sleep(0.08) + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=SCOPE_A, + ) + await asyncio.wait_for(started.wait(), timeout=2.0) + + partition = await memory_store.get_partition(SCOPE_A) + tasks = dict(partition.tasks) if partition else {} + tasks["otherwrk"] = _stub_task_payload("otherwrk", SCOPE_A) + await memory_store.put_partition(SCOPE_A, SessionPartitionPayload(tasks=tasks)) + + await _wait_terminal(task_id, SCOPE_A) + latest = await memory_store.get_partition(SCOPE_A) + assert latest is not None + assert task_id in latest.tasks + assert "otherwrk" in latest.tasks + + _run(scenario()) + + def test_put_payload_is_tasks_only(self): + session_storage = RecordingSessionStorageProvider() + configure_task_storage(session_storage) + + async def scenario(): + async def action(): + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=SCOPE_A, + ) + await _wait_terminal(task_id, SCOPE_A) + assert session_storage.put_payloads + for payload in session_storage.put_payloads: + assert payload.tasks is not None + assert payload.dataframes is None + assert payload.metadata is None + assert payload.uploaded_files is None + + _run(scenario()) + + def test_idle_caches_are_evicted(self, memory_store, monkeypatch): + monkeypatch.setattr(task_manager, "_MAX_SESSION_CACHES", 4) + + async def scenario(): + async def action(): + return BaseResult(result=[{"ok": True}]) + + for index in range(10): + scope = SessionScope(user_id="user-1", mcp_session_id=f"sess-{index}") + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=scope, + ) + await _wait_terminal(task_id, scope) + return len(task_manager._session_caches) + + assert _run(scenario()) <= 4 + + +class TestHttpSessionStorageProviderTasks: + def test_http_roundtrip_preserves_dataframes(self, monkeypatch): + transport = MergingSessionTransport() + client = http_session_storage_provider(monkeypatch, transport) + configure_task_storage(client) + scope = SessionScope("user-9", "sess-http") + + async def scenario(): + await register_dataframe( + result=[{"id": 10, "label": "x"}], + origin_manager="tests", + origin_action="seed", + json_size_chars=9001, + session_storage=client, + scope=scope, + ) + + async def action(): + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=scope, + ) + await _wait_terminal(task_id, scope) + partition = await client.get_partition(scope) + assert partition is not None + assert task_id in partition.tasks + assert partition.tasks[task_id]["status"] == STATUS_COMPLETED + assert len(partition.dataframes) == 1 + + _run(scenario()) + + +class TestCancelTaskSemantics: + def test_cancel_leaves_completed_task_unchanged(self, memory_store): + async def scenario(): + async def action(): + return BaseResult(result=[{"ok": True}]) + + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=SessionScope(user_id="user-cancel", mcp_session_id="sess-cancel"), + ) + scope = SessionScope(user_id="user-cancel", mcp_session_id="sess-cancel") + completed = await _wait_terminal(task_id, scope) + assert completed.status == STATUS_COMPLETED + + after = await cancel_task(task_id, scope=scope) + assert after is not None + assert after.status == STATUS_COMPLETED + assert after.result is not None + assert after.result.result == [{"ok": True}] + + partition = await memory_store.get_partition(scope) + assert partition.tasks[task_id]["status"] == STATUS_COMPLETED + + _run(scenario()) + + def test_cancel_without_local_handle_marks_storage_cancelled(self, memory_store): + async def scenario(): + async def action(): + await asyncio.sleep(60) + return BaseResult(result=[{"ok": True}]) + + scope = SessionScope(user_id="user-affinity", mcp_session_id="sess-affinity") + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=scope, + ) + for _ in range(100): + record = await get_task_record(task_id, scope=scope) + if record and record.status == STATUS_WORKING: + break + await asyncio.sleep(0.01) + else: + raise AssertionError("task did not start working") + + record.asyncio_task = None + cancelled = await cancel_task(task_id, scope=scope) + assert cancelled is not None + assert cancelled.status == STATUS_CANCELLED + assert "no local" in (cancelled.status_message or "").lower() + + _run(scenario()) + + def test_cancel_local_running_task(self, memory_store): + async def scenario(): + started = asyncio.Event() + + async def action(): + started.set() + await asyncio.sleep(60) + return BaseResult(result=[{"ok": True}]) + + scope = SessionScope(user_id="user-local-cancel", mcp_session_id="sess-local-cancel") + task_id = await submit_task( + {"manager": "TestManager", "method": "read"}, + action, + scope=scope, + ) + await asyncio.wait_for(started.wait(), timeout=2.0) + record = await get_task_record(task_id, scope=scope) + assert record is not None + assert record.asyncio_task is not None + assert not record.asyncio_task.done() + + await cancel_task(task_id, scope=scope) + terminal = await _wait_terminal(task_id, scope, timeout=2.0) + assert terminal.status == STATUS_CANCELLED + + _run(scenario()) diff --git a/tests/test_batch_controls.py b/tests/test_batch_controls.py index a402107..e572da8 100644 --- a/tests/test_batch_controls.py +++ b/tests/test_batch_controls.py @@ -61,7 +61,7 @@ async def slow_list_help_categories(self): monkeypatch.setattr(HelpManager, "list_help_categories", slow_list_help_categories) batch_calls = [{"action": "list_help_categories", "args": {}} for _ in range(6)] - result = asyncio.run(help_tool("batch", {"batch_calls": batch_calls}, ctx=None)) + result = asyncio.run(help_tool({"action": "batch", "args": {"batch_calls": batch_calls}}, ctx=None)) assert result.error is None assert active_calls["max"] <= 2 @@ -86,7 +86,7 @@ async def slow_list_skills(): monkeypatch.setattr(SkillsManager, "list_skills", staticmethod(slow_list_skills)) batch_calls = [{"action": "list_skills", "args": {}} for _ in range(6)] - result = asyncio.run(skills_tool("batch", {"batch_calls": batch_calls}, ctx=None)) + result = asyncio.run(skills_tool({"action": "batch", "args": {"batch_calls": batch_calls}}, ctx=None)) assert result.error is None assert active_calls["max"] <= 2 diff --git a/tests/test_dataframe_session_fixes.py b/tests/test_dataframe_session_fixes.py new file mode 100644 index 0000000..e38fc96 --- /dev/null +++ b/tests/test_dataframe_session_fixes.py @@ -0,0 +1,301 @@ +""" +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 +from unittest.mock import MagicMock + +from config.runtime import AppRuntime +from config.storage import DefaultSessionScopeResolver, SessionScope +from config.token import BzmToken +from models.result import BaseResult +from tests.conftest import make_ctx, run_async +from tools.dataframe_manager import ( + MISSING_STORAGE_ERROR, + finalize_tool_result, + list_dataframes_metadata, + materialize_large_result_if_needed, + register_dataframe, +) +from tools.runtime_tools import run_tool_with_runtime +from tools.tools_manager import ToolsManager + + +class TestConcurrentSessionIsolation: + def test_concurrent_registers_do_not_cross_write(self, in_memory_session_storage): + async def _register(user_id: str, session_id: str, value: int): + return await register_dataframe( + result=[{"v": value}], + origin_manager="tests", + origin_action="seed", + json_size_chars=9001, + session_storage=in_memory_session_storage, + scope=SessionScope(user_id, session_id), + ) + + async def _exercise(): + await asyncio.gather( + _register("u1", "s1", 1), + _register("u1", "s2", 2), + _register("u2", "s1", 3), + _register("u1", "s1", 11), + ) + listed = await asyncio.gather( + list_dataframes_metadata(in_memory_session_storage, SessionScope("u1", "s1")), + list_dataframes_metadata(in_memory_session_storage, SessionScope("u1", "s2")), + list_dataframes_metadata(in_memory_session_storage, SessionScope("u2", "s1")), + ) + return listed + + s1, s2, s3 = run_async(_exercise()) + assert len(s1) == 2 + assert len(s2) == 1 + assert len(s3) == 1 + assert {row["rows"] for row in s1} == {1} + assert s2[0]["origin_action"] == "seed" + for partition in (s1, s2, s3): + assert all(item["origin_manager"] == "tests" for item in partition) + + +class TestMaterializeWiring: + def test_finalize_materializes_large_auto_result(self, in_memory_session_storage): + token = BzmToken("user-mat", "secret") + ctx = make_ctx(token, "sess-mat") + payload = [{"id": i, "name": f"row-{i}", "note": "x" * 40} for i in range(300)] + base = BaseResult(result=payload) + + finalized = run_async( + finalize_tool_result( + base, + action="list", + args={}, + origin_manager="blazemeter_tests", + session_storage=in_memory_session_storage, + scope_resolver=DefaultSessionScopeResolver(), + token=token, + ctx=ctx, + ) + ) + assert finalized.error is None + assert finalized.result[0]["stored_as_dataframe"] is True + listed = run_async( + list_dataframes_metadata(in_memory_session_storage, SessionScope("user-mat", "sess-mat")) + ) + assert len(listed) == 1 + + def test_force_dataframe_even_when_small(self, in_memory_session_storage): + base = BaseResult(result=[{"id": 1}]) + finalized = run_async( + materialize_large_result_if_needed( + base, + origin_manager="tests", + origin_action="list", + session_storage=in_memory_session_storage, + scope=SessionScope("user-force", "sess-force"), + force=True, + ) + ) + assert finalized.result[0]["stored_as_dataframe"] is True + + def test_finalize_without_storage_fails_closed(self): + payload = [{"id": i, "note": "x" * 40} for i in range(300)] + finalized = run_async( + finalize_tool_result( + BaseResult(result=payload), + action="list", + args={}, + origin_manager="blazemeter_tests", + ) + ) + assert finalized.error == MISSING_STORAGE_ERROR + + def test_excluded_action_skips_without_storage(self): + payload = [{"id": i, "note": "x" * 40} for i in range(300)] + finalized = run_async( + finalize_tool_result( + BaseResult(result=payload), + action="dataframes_list", + args={}, + origin_manager="blazemeter_tools", + excluded_actions={"dataframes_list"}, + ) + ) + assert len(finalized.result) == 300 + + def test_excluded_action_skips_auto_materialize(self, in_memory_session_storage): + token = BzmToken("user-ex", "secret") + ctx = make_ctx(token, "sess-ex") + payload = [{"id": i, "note": "x" * 40} for i in range(300)] + finalized = run_async( + finalize_tool_result( + BaseResult(result=payload), + action="dataframes_list", + args={}, + origin_manager="blazemeter_tools", + session_storage=in_memory_session_storage, + token=token, + ctx=ctx, + excluded_actions={"dataframes_list"}, + ) + ) + assert len(finalized.result) == 300 + + +class TestRunToolWithRuntime: + def test_materializes_large_result_via_runtime_storage(self, in_memory_session_storage): + token = BzmToken("user-rt", "secret") + ctx = make_ctx(token, "sess-rt") + runtime = AppRuntime( + transport="stdio", + auth=MagicMock(get_token=MagicMock(return_value=token)), + storage=in_memory_session_storage, + file_access=MagicMock(), + scope_resolver=DefaultSessionScopeResolver(), + user_config={}, + ) + payload = [{"id": i, "name": f"row-{i}", "note": "x" * 40} for i in range(300)] + + async def _dispatch(): + return BaseResult(result=payload) + + finalized = run_async( + run_tool_with_runtime( + runtime, "blazemeter_tests", "list", ctx, _dispatch, + ) + ) + assert finalized.error is None + assert finalized.result[0]["stored_as_dataframe"] is True + listed = run_async( + list_dataframes_metadata(in_memory_session_storage, SessionScope("user-rt", "sess-rt")) + ) + assert len(listed) == 1 + + def test_force_store_when_result_format_passed_as_tool_args(self, in_memory_session_storage): + token = BzmToken("user-rf", "secret") + ctx = make_ctx(token, "sess-rf") + runtime = AppRuntime( + transport="streamable-http", + auth=MagicMock(get_token=MagicMock(return_value=token)), + storage=in_memory_session_storage, + file_access=MagicMock(), + scope_resolver=DefaultSessionScopeResolver(), + user_config={}, + ) + + async def _dispatch(): + return BaseResult(result=[{"id": 1}]) + + finalized = run_async( + run_tool_with_runtime( + runtime, "blazemeter_user", "read", ctx, _dispatch, + tool_args={"result_format": "dataframe"}, + ) + ) + assert finalized.error is None + assert finalized.result[0]["stored_as_dataframe"] is True + listed = run_async( + list_dataframes_metadata(in_memory_session_storage, SessionScope("user-rf", "sess-rf")) + ) + assert len(listed) == 1 + + def test_raw_skips_materialize_when_result_format_passed_as_tool_args( + self, in_memory_session_storage): + token = BzmToken("user-raw", "secret") + ctx = make_ctx(token, "sess-raw") + runtime = AppRuntime( + transport="streamable-http", + auth=MagicMock(get_token=MagicMock(return_value=token)), + storage=in_memory_session_storage, + file_access=MagicMock(), + scope_resolver=DefaultSessionScopeResolver(), + user_config={}, + ) + payload = [{"id": i, "note": "x" * 40} for i in range(300)] + + async def _dispatch(): + return BaseResult(result=payload) + + finalized = run_async( + run_tool_with_runtime( + runtime, "blazemeter_user", "read", ctx, _dispatch, + tool_args={"result_format": "raw"}, + ) + ) + assert finalized.error is None + assert len(finalized.result) == 300 + listed = run_async( + list_dataframes_metadata(in_memory_session_storage, SessionScope("user-raw", "sess-raw")) + ) + assert listed == [] + + def test_disable_dataframe_materialization_skips_persist(self, in_memory_session_storage): + token = BzmToken("user-off", "secret") + ctx = make_ctx(token, "sess-off") + runtime = AppRuntime( + transport="stdio", + auth=MagicMock(get_token=MagicMock(return_value=token)), + storage=in_memory_session_storage, + file_access=MagicMock(), + scope_resolver=DefaultSessionScopeResolver(), + user_config={}, + ) + payload = [{"id": i, "note": "x" * 40} for i in range(300)] + + async def _dispatch(): + return BaseResult(result=payload) + + finalized = run_async( + run_tool_with_runtime( + runtime, "blazemeter_skills", "read_skill", ctx, _dispatch, + disable_dataframe_materialization=True, + ) + ) + assert finalized.error is None + assert len(finalized.result) == 300 + listed = run_async( + list_dataframes_metadata(in_memory_session_storage, SessionScope("user-off", "sess-off")) + ) + assert listed == [] + + +class TestDataframesQueryResultFormatStore: + def test_result_format_dataframe_registers_new_dataframe(self, in_memory_session_storage): + token = BzmToken("user-q", "secret") + ctx = make_ctx(token, "sess-q") + manager = ToolsManager(ctx, in_memory_session_storage, DefaultSessionScopeResolver()) + scope = SessionScope("user-q", "sess-q") + + meta = run_async( + register_dataframe( + result=[{"id": 1, "name": "a"}, {"id": 2, "name": "b"}], + origin_manager="tests", + origin_action="seed", + json_size_chars=9001, + session_storage=in_memory_session_storage, + scope=scope, + ) + ) + queried = run_async( + manager.dataframes_query( + sql=( + f"SELECT id, name FROM {meta['table_name']} " + f"ORDER BY id LIMIT 100 OFFSET 0" + ), + result_format="dataframe", + ) + ) + assert queried.error is None + assert queried.result[0]["stored_as_dataframe"] is True + listed = run_async(list_dataframes_metadata(in_memory_session_storage, scope)) + assert len(listed) == 2 diff --git a/tests/test_dataframe_storage.py b/tests/test_dataframe_storage.py new file mode 100644 index 0000000..f1faeb2 --- /dev/null +++ b/tests/test_dataframe_storage.py @@ -0,0 +1,367 @@ +""" +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 hashlib +import json + +import pytest + +from config.storage import ( + FileStoragePort, + HttpSessionStorageProvider, + HttpStorageClient, + InMemorySessionStorageProvider, + SessionPartitionPayload, + SessionScope, + SessionStoragePort, +) +from tests.conftest import run_async +from tests.storage_fakes import ( + MergingSessionTransport, + RecordingSessionStorageProvider, + http_session_storage_provider, +) +from tools import dataframe_manager as dataframe_manager_module +from tools.dataframe_manager import ( + _schema_hash, + _stable_hash, + clear_dataframes, + list_dataframes_metadata, + query_dataframes, + register_dataframe, + remove_dataframes, + remove_dataframe, +) + + +def _seed(session_storage, scope, rows, action="seed"): + return run_async( + register_dataframe( + result=rows, + origin_manager="tests", + origin_action=action, + json_size_chars=9001, + session_storage=session_storage, + scope=scope, + ) + ) + + +class TestSessionStoragePortContract: + """Dataframes use STREAMABLE_HTTP SessionStoragePort names, not FileStoragePort.""" + + def test_in_memory_provider_is_session_storage_port(self, in_memory_session_storage): + assert isinstance(in_memory_session_storage, InMemorySessionStorageProvider) + assert isinstance(in_memory_session_storage, SessionStoragePort) + + def test_http_session_provider_is_session_storage_port(self, monkeypatch): + session_storage = http_session_storage_provider( + monkeypatch, MergingSessionTransport(), + ) + assert isinstance(session_storage, HttpSessionStorageProvider) + assert isinstance(session_storage, SessionStoragePort) + + def test_file_http_storage_client_is_not_session_storage_port(self): + file_client = HttpStorageClient() + assert isinstance(file_client, FileStoragePort) + assert not isinstance(file_client, SessionStoragePort) + assert not hasattr(file_client, "put_partition") + assert not hasattr(file_client, "get_partition") + + def test_register_rejects_file_http_storage_client(self, session_scope): + with pytest.raises(AttributeError): + _seed(HttpStorageClient(), session_scope, [{"id": 1}]) + + +class TestDataframeManagerInMemorySessionStorageProvider: + def test_register_list_query_remove_clear(self, in_memory_session_storage, session_scope): + meta = _seed( + in_memory_session_storage, + session_scope, + [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}], + ) + assert meta["rows"] == 2 + + listed = run_async(list_dataframes_metadata(in_memory_session_storage, session_scope)) + assert len(listed) == 1 + assert listed[0]["dataframe_id"] == meta["dataframe_id"] + + sql = ( + f"SELECT id, name FROM {meta['table_name']} " + f"ORDER BY id LIMIT 100 OFFSET 0" + ) + queried = run_async(query_dataframes(sql, in_memory_session_storage, session_scope)) + assert "error" not in queried + assert queried["rows"] == 2 + + assert run_async(remove_dataframe(meta["dataframe_id"], in_memory_session_storage, session_scope)) + assert run_async(list_dataframes_metadata(in_memory_session_storage, session_scope)) == [] + assert run_async(clear_dataframes(in_memory_session_storage, session_scope)) == 0 + + def test_sessions_are_isolated(self, in_memory_session_storage): + scope_a = SessionScope("user-1", "sess-a") + scope_b = SessionScope("user-1", "sess-b") + _seed(in_memory_session_storage, scope_a, [{"id": 1}], action="a") + _seed(in_memory_session_storage, scope_b, [{"id": 2}], action="b") + a = run_async(list_dataframes_metadata(in_memory_session_storage, scope_a)) + b = run_async(list_dataframes_metadata(in_memory_session_storage, scope_b)) + assert len(a) == 1 + assert len(b) == 1 + assert a[0]["dataframe_id"] != b[0]["dataframe_id"] + + def test_second_call_sees_session_storage_without_cache_reset(self, in_memory_session_storage): + scope = SessionScope("user-42", "mcp-session-shared") + meta = _seed( + in_memory_session_storage, + scope, + [{"n": 1}, {"n": 2}], + action="req1", + ) + listed = run_async(list_dataframes_metadata(in_memory_session_storage, scope)) + assert len(listed) == 1 + assert listed[0]["dataframe_id"] == meta["dataframe_id"] + assert listed[0]["rows"] == 2 + + def test_external_put_is_visible_on_next_read(self, monkeypatch): + transport = MergingSessionTransport() + client_a = http_session_storage_provider(monkeypatch, transport) + client_b = HttpSessionStorageProvider(base_url="http://storage.test") + scope = SessionScope("user-1", "sess-a") + first = _seed(client_a, scope, [{"id": 1}]) + injected = _seed(client_b, scope, [{"id": 99}], action="external") + listed_again = run_async(list_dataframes_metadata(client_a, scope)) + ids = {row["dataframe_id"] for row in listed_again} + assert first["dataframe_id"] in ids + assert injected["dataframe_id"] in ids + + +class TestDataframeManagerSessionStoragePort: + def test_register_writes_through_put_partition(self): + session_storage = RecordingSessionStorageProvider() + scope = SessionScope("user-9", "sess-http") + meta = _seed(session_storage, scope, [{"id": 10, "label": "x"}], action="http") + payload = session_storage.put_payloads[-1] + assert meta["dataframe_id"] in payload.dataframes + assert payload.tasks is None + assert payload.metadata is None + assert payload.uploaded_files is None + + listed = run_async(list_dataframes_metadata(session_storage, scope)) + assert len(listed) == 1 + sql = ( + f"SELECT id, label FROM {meta['table_name']} " + f"ORDER BY id LIMIT 10 OFFSET 0" + ) + queried = run_async(query_dataframes(sql, session_storage, scope)) + assert "error" not in queried + assert queried["rows"] == 1 + + partition = run_async(session_storage.get_partition(scope)) + assert partition is not None + assert meta["dataframe_id"] in partition.dataframes + + def test_register_preserves_existing_tasks(self, in_memory_session_storage): + scope = SessionScope("user-9", "sess-tasks") + run_async( + in_memory_session_storage.put_partition( + scope, + SessionPartitionPayload(tasks={"t1": {"status": "running"}}), + ) + ) + meta = _seed(in_memory_session_storage, scope, [{"id": 1}], action="http") + partition = run_async(in_memory_session_storage.get_partition(scope)) + assert partition is not None + assert partition.tasks == {"t1": {"status": "running"}} + assert meta["dataframe_id"] in partition.dataframes + + def test_remove_batch_commits_once(self): + session_storage = RecordingSessionStorageProvider() + scope = SessionScope("user-9", "sess-batch") + a = _seed(session_storage, scope, [{"id": 1}], action="a") + b = _seed(session_storage, scope, [{"id": 2}], action="b") + puts_before = len(session_storage.put_payloads) + outcome = run_async( + remove_dataframes( + [a["dataframe_id"], b["dataframe_id"]], + session_storage, + scope, + ) + ) + assert outcome["removed"] == [a["dataframe_id"], b["dataframe_id"]] + assert len(session_storage.put_payloads) == puts_before + 1 + + +class InterveningGetSessionStorageProvider(InMemorySessionStorageProvider): + """Injects extra dataframes on a chosen GET to simulate another worker.""" + + def __init__(self, extra_dataframes: dict) -> None: + super().__init__() + self.get_count = 0 + self.inject_on_get: int | None = None + self.extra_dataframes = extra_dataframes + + async def get_partition(self, scope: SessionScope): + self.get_count += 1 + if self.get_count == self.inject_on_get: + existing = await super().get_partition(scope) + dataframes = dict(existing.dataframes) if existing else {} + dataframes.update(self.extra_dataframes) + await super().put_partition( + scope, SessionPartitionPayload(dataframes=dataframes) + ) + return await super().get_partition(scope) + + +class TestDataframeMapMerge: + def test_persist_keeps_keys_added_since_hydrate(self, session_scope): + side = InMemorySessionStorageProvider() + extra_meta = _seed(side, session_scope, [{"id": 99}], action="external") + extra_partition = run_async(side.get_partition(session_scope)) + extra_dataframes = { + extra_meta["dataframe_id"]: extra_partition.dataframes[extra_meta["dataframe_id"]] + } + session_storage = InterveningGetSessionStorageProvider(extra_dataframes) + first = _seed(session_storage, session_scope, [{"id": 1}]) + session_storage.get_count = 0 + session_storage.inject_on_get = 2 + second = _seed(session_storage, session_scope, [{"id": 2}], action="local") + listed = run_async(list_dataframes_metadata(session_storage, session_scope)) + ids = {row["dataframe_id"] for row in listed} + assert first["dataframe_id"] in ids + assert second["dataframe_id"] in ids + assert extra_meta["dataframe_id"] in ids + + +class TestHttpSessionStorageProviderDataframes: + def testhttp_session_storage_provider_roundtrip_and_task_merge(self, monkeypatch): + transport = MergingSessionTransport() + client = http_session_storage_provider(monkeypatch, transport) + scope = SessionScope("user-9", "sess-http") + + run_async( + client.put_partition( + scope, + SessionPartitionPayload(tasks={"t1": {"status": "running"}}), + ) + ) + meta = _seed(client, scope, [{"id": 10, "label": "x"}], action="http") + listed = run_async(list_dataframes_metadata(client, scope)) + assert len(listed) == 1 + sql = ( + f"SELECT id, label FROM {meta['table_name']} " + f"ORDER BY id LIMIT 10 OFFSET 0" + ) + queried = run_async(query_dataframes(sql, client, scope)) + assert "error" not in queried + partition = run_async(client.get_partition(scope)) + assert partition is not None + assert partition.tasks == {"t1": {"status": "running"}} + assert meta["dataframe_id"] in partition.dataframes + assert isinstance(client, SessionStoragePort) + + +class TestSqlReadOnlyGate: + def test_requires_select_order_limit_offset(self, in_memory_session_storage, session_scope): + missing_clauses = run_async( + query_dataframes("SELECT * FROM df_x", in_memory_session_storage, session_scope) + ) + assert "error" in missing_clauses + assert "ORDER BY" in missing_clauses["error"] + + delete_stmt = run_async( + query_dataframes("DELETE FROM df_x", in_memory_session_storage, session_scope) + ) + assert "error" in delete_stmt + assert "read-only" in delete_stmt["error"].lower() + + meta = _seed(in_memory_session_storage, session_scope, [{"id": 1, "name": "a"}]) + valid = run_async( + query_dataframes( + f"SELECT * FROM {meta['table_name']} ORDER BY id LIMIT 10 OFFSET 0", + in_memory_session_storage, + session_scope, + ) + ) + assert "error" not in valid + assert valid["rows"] == 1 + + +class TestSessionLockBound: + def test_unlocked_locks_are_evicted(self, in_memory_session_storage, monkeypatch): + monkeypatch.setattr(dataframe_manager_module, "_MAX_SESSION_LOCKS", 4) + + async def _exercise(): + for index in range(10): + await list_dataframes_metadata( + in_memory_session_storage, SessionScope("user-1", f"sess-{index}") + ) + return len(dataframe_manager_module._session_locks) + + assert run_async(_exercise()) <= 4 + + def test_overflow_lock_keeps_map_bounded_when_all_locks_held(self, monkeypatch): + monkeypatch.setattr(dataframe_manager_module, "_MAX_SESSION_LOCKS", 2) + + async def _exercise(): + lock_a = await dataframe_manager_module._lock_for(SessionScope("user-1", "sess-a")) + lock_b = await dataframe_manager_module._lock_for(SessionScope("user-1", "sess-b")) + await lock_a.acquire() + await lock_b.acquire() + try: + lock_c = await dataframe_manager_module._lock_for(SessionScope("user-1", "sess-c")) + lock_d = await dataframe_manager_module._lock_for(SessionScope("user-1", "sess-d")) + assert len(dataframe_manager_module._session_locks) == 2 + overflow = dataframe_manager_module._overflow_lock + assert overflow is not None + assert lock_c is overflow + assert lock_d is overflow + assert overflow is not lock_a + assert overflow is not lock_b + finally: + lock_a.release() + lock_b.release() + + run_async(_exercise()) + + def test_overflow_session_can_still_read_metadata(self, in_memory_session_storage, monkeypatch): + monkeypatch.setattr(dataframe_manager_module, "_MAX_SESSION_LOCKS", 2) + + async def _exercise(): + lock_a = await dataframe_manager_module._lock_for(SessionScope("user-1", "held-a")) + lock_b = await dataframe_manager_module._lock_for(SessionScope("user-1", "held-b")) + await lock_a.acquire() + await lock_b.acquire() + try: + listed = await list_dataframes_metadata( + in_memory_session_storage, SessionScope("user-1", "overflow") + ) + assert listed == [] + assert len(dataframe_manager_module._session_locks) == 2 + finally: + lock_a.release() + lock_b.release() + + run_async(_exercise()) + + +class TestContentHashes: + def test_stable_hash_uses_sha256(self): + payload = '{"name":"id","dtype":"Int64"}' + assert _stable_hash(payload) == hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def test_schema_hash_reuses_stable_hash(self): + schema_rows = [{"name": "id", "dtype": "Int64"}, {"name": "label", "dtype": "String"}] + canonical = json.dumps(schema_rows, separators=(",", ":"), ensure_ascii=False) + assert _schema_hash(schema_rows) == _stable_hash(canonical) diff --git a/tests/test_failure_criteria.py b/tests/test_failure_criteria.py index a8ad36a..40a17ce 100644 --- a/tests/test_failure_criteria.py +++ b/tests/test_failure_criteria.py @@ -280,7 +280,7 @@ def test_tool_returns_catalog_without_api(self): mcp = _FakeMcpForTests() register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] - result = asyncio.run(tool("failure_criteria_meta", {}, ctx=None)) + result = asyncio.run(tool({"action": "failure_criteria_meta", "args": {}}, ctx=None)) assert result.error is None payload = result.result[0] assert "top_level_tool_args" in payload diff --git a/tests/test_required_args_tools.py b/tests/test_required_args_tools.py index 05ffde6..8b995aa 100644 --- a/tests/test_required_args_tools.py +++ b/tests/test_required_args_tools.py @@ -49,7 +49,7 @@ def test_account_read_requires_account_id(self): register_account_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_account"] - result = asyncio.run(tool("read", {}, ctx=None)) + result = asyncio.run(tool({"action": "read", "args": {}}, ctx=None)) assert result.error is not None assert "account_id" in result.error @@ -58,7 +58,7 @@ def test_workspace_list_requires_account_id(self): register_workspaces_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_workspaces"] - result = asyncio.run(tool("list", {}, ctx=None)) + result = asyncio.run(tool({"action": "list", "args": {}}, ctx=None)) assert result.error is not None assert "account_id" in result.error @@ -67,7 +67,7 @@ def test_project_read_requires_project_id(self): register_project_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_project"] - result = asyncio.run(tool("read", {}, ctx=None)) + result = asyncio.run(tool({"action": "read", "args": {}}, ctx=None)) assert result.error is not None assert "project_id" in result.error @@ -76,7 +76,7 @@ def test_tests_create_requires_test_name(self): register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] - result = asyncio.run(tool("create", {"project_id": 123}, ctx=None)) + result = asyncio.run(tool({"action": "create", "args": {"project_id": 123}}, ctx=None)) assert result.error is not None assert "test_name" in result.error @@ -85,7 +85,7 @@ def test_tests_upload_assets_requires_file_paths(self): register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] - result = asyncio.run(tool("upload_assets", {"test_id": 123}, ctx=None)) + result = asyncio.run(tool({"action": "upload_assets", "args": {"test_id": 123}}, ctx=None)) assert result.error is not None assert "file_paths" in result.error @@ -94,12 +94,12 @@ def test_tests_configure_failure_criteria_requires_enabled_and_rules(self): register_tests_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_tests"] - result = asyncio.run(tool("configure_failure_criteria", {"test_id": 123}, ctx=None)) + result = asyncio.run(tool({"action": "configure_failure_criteria", "args": {"test_id": 123}}, ctx=None)) assert result.error is not None assert "enabled" in result.error result = asyncio.run( - tool("configure_failure_criteria", {"test_id": 123, "enabled": True}, ctx=None) + tool({"action": "configure_failure_criteria", "args": {"test_id": 123, "enabled": True}}, ctx=None) ) assert result.error is not None assert "rules" in result.error @@ -109,7 +109,7 @@ def test_execution_read_requires_execution_id(self): register_execution_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_execution"] - result = asyncio.run(tool("read", {}, ctx=None)) + result = asyncio.run(tool({"action": "read", "args": {}}, ctx=None)) assert result.error is not None assert "execution_id" in result.error @@ -118,7 +118,7 @@ def test_execution_read_summary_requires_execution_id(self): register_execution_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_execution"] - result = asyncio.run(tool("read_summary", {}, ctx=None)) + result = asyncio.run(tool({"action": "read_summary", "args": {}}, ctx=None)) assert result.error is not None assert "execution_id" in result.error @@ -127,7 +127,7 @@ def test_skills_read_skill_requires_skill_name(self): register_skills_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] - result = asyncio.run(tool("read_skill", {}, ctx=None)) + result = asyncio.run(tool({"action": "read_skill", "args": {}}, ctx=None)) assert result.error is not None assert "skill_name" in result.error @@ -136,7 +136,7 @@ def test_skills_read_skill_resource_uri_requires_uri(self): register_skills_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] - result = asyncio.run(tool("read_skill_resource_uri", {}, ctx=None)) + result = asyncio.run(tool({"action": "read_skill_resource_uri", "args": {}}, ctx=None)) assert result.error is not None assert "skill_resource_uri" in result.error @@ -145,7 +145,7 @@ def test_skills_read_skill_resource_uri_list_requires_non_empty_list(self): register_skills_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] - result = asyncio.run(tool("read_skill_resource_uri_list", {}, ctx=None)) + result = asyncio.run(tool({"action": "read_skill_resource_uri_list", "args": {}}, ctx=None)) assert result.error is not None assert "skill_resource_uri_list" in result.error @@ -154,7 +154,7 @@ def test_help_read_help_info_requires_help_id_list(self): register_help_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_help"] - result = asyncio.run(tool("read_help_info", {}, ctx=None)) + result = asyncio.run(tool({"action": "read_help_info", "args": {}}, ctx=None)) assert result.error is not None assert "help_id_list" in result.error @@ -163,6 +163,6 @@ def test_help_list_help_category_content_requires_subcategory_list(self): register_help_tool(mcp, build_runtime("stdio")) tool = mcp.tools[f"{TOOLS_PREFIX}_help"] - result = asyncio.run(tool("list_help_category_content", {}, ctx=None)) + result = asyncio.run(tool({"action": "list_help_category_content", "args": {}}, ctx=None)) assert result.error is not None assert "subcategory_id_list" in result.error diff --git a/tests/test_skills_manager_security.py b/tests/test_skills_manager_security.py index 013c12e..e04abd8 100644 --- a/tests/test_skills_manager_security.py +++ b/tests/test_skills_manager_security.py @@ -18,8 +18,16 @@ import pytest +from config.auth import StdioAuthProvider +from config.blazemeter import TOOLS_PREFIX +from config.file_access import build_file_access +from config.runtime import AppRuntime +from config.storage import DefaultSessionScopeResolver, SessionScope +from config.token import BzmToken +from tests.conftest import make_ctx, run_async from tools import skills_utils -from tools.skills_manager import SkillsManager +from tools.dataframe_manager import list_dataframes_metadata +from tools.skills_manager import SkillsManager, register as register_skills_tool @pytest.fixture @@ -44,15 +52,71 @@ def isolated_skills_resources(tmp_path, monkeypatch): class TestSkillsManagerListResourcesErrors: def test_list_skill_resources_returns_controlled_error_for_invalid_skill_name(self, isolated_skills_resources): - result = asyncio.run(SkillsManager.list_skill_resources("../safe-skill")) + manager = SkillsManager(ctx=None) + result = asyncio.run(manager.list_skill_resources("../safe-skill")) assert result.error is not None assert "Invalid skill name" in result.error assert result.result is None def test_list_skill_resources_returns_controlled_error_for_missing_skill(self, isolated_skills_resources): - result = asyncio.run(SkillsManager.list_skill_resources("unknown-skill")) + manager = SkillsManager(ctx=None) + result = asyncio.run(manager.list_skill_resources("unknown-skill")) assert result.error is not None assert "Skill folder not found" in result.error assert result.result is None + + +class FakeMcp: + def __init__(self): + self.tools = {} + + def tool(self, name, description): + def decorator(func): + self.tools[name] = func + return func + + return decorator + + +class TestSkillsSkipDataframeMaterialization: + def test_read_skill_keeps_large_document_inline(self, isolated_skills_resources, in_memory_session_storage): + large_body = "x" * 12000 + skill_dir = isolated_skills_resources / "skills" / "safe-skill" + (skill_dir / "SKILL.md").write_text( + "---\n" + "name: safe-skill\n" + "description: Security test skill\n" + "---\n" + f"{large_body}\n", + encoding="utf-8", + ) + SkillsManager.skills = None + token = BzmToken("user-skills", "secret") + runtime = AppRuntime( + transport="stdio", + auth=StdioAuthProvider(token), + storage=in_memory_session_storage, + file_access=build_file_access("stdio"), + scope_resolver=DefaultSessionScopeResolver(), + user_config={"token": token}, + ) + mcp = FakeMcp() + register_skills_tool(mcp, runtime) + tool = mcp.tools[f"{TOOLS_PREFIX}_skills"] + ctx = make_ctx(token, "sess-skills") + + result = asyncio.run( + tool({"action": "read_skill", "args": {"skill_name": "safe-skill"}}, ctx=ctx) + ) + + assert result.error is None + assert result.result[0].get("stored_as_dataframe") is not True + assert large_body in result.result[0]["content"] + listed = run_async( + list_dataframes_metadata( + in_memory_session_storage, SessionScope("user-skills", "sess-skills") + ) + ) + assert listed == [] diff --git a/tests/test_storage.py b/tests/test_storage.py index 6815120..f9000e5 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -125,12 +125,23 @@ async def _fake_read(_test_id): result = asyncio.run( manager.upload_assets(1, ["/tmp/demo.jmx"], main_script=None) ) - assert "error" in result - assert "No valid files found to upload" in result["error"] + if hasattr(result, "error"): + # @run_as_task normalizes dict errors into BaseResult(result=[{error:...}]). + inner = result.result[0] if result.result else {} + error_text = result.error or (inner.get("error") if isinstance(inner, dict) else None) + else: + error_text = result.get("error") if isinstance(result, dict) else None + assert error_text is not None + assert "No valid files found to upload" in error_text def test_upload_assets_without_file_ports_returns_hosted_message(self): manager = TestManager(ctx=None) result = asyncio.run( manager.upload_assets(1, ["/tmp/demo.jmx"], main_script=None) ) - assert result["error"] == HOSTED_FILE_ACCESS_MESSAGE + if hasattr(result, "error"): + inner = result.result[0] if result.result else {} + error_text = result.error or (inner.get("error") if isinstance(inner, dict) else None) + else: + error_text = result.get("error") if isinstance(result, dict) else None + assert error_text == HOSTED_FILE_ACCESS_MESSAGE diff --git a/tests/test_support_message.py b/tests/test_support_message.py new file mode 100644 index 0000000..ac5a610 --- /dev/null +++ b/tests/test_support_message.py @@ -0,0 +1,182 @@ +""" +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 + + 10|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 ast +import asyncio +import importlib +from pathlib import Path + +import pytest + +from config.blazemeter import SUPPORT_MESSAGE +from config.runtime import build_runtime +from tools.mcp_entrypoint import register_managed_tool + +TOOLS_DIR = Path(__file__).resolve().parents[1] / "tools" +HARDCODED_SUPPORT_SNIPPET = "If you think this is a bug, please contact BlazeMeter support" + +MANAGER_MODULES = ( + "tools.account_manager", + "tools.billing_manager", + "tools.execution_manager", + "tools.help_manager", + "tools.project_manager", + "tools.skills_manager", + "tools.test_manager", + "tools.tools_manager", + "tools.user_manager", + "tools.workspace_manager", +) + +# tools_manager wraps other tools; it opts out so callers do not see the message twice. +OPT_OUT_MODULES = frozenset({"tools.tools_manager"}) + +# Skills return document text, not tabular rows; storing them as dataframes +# would add an extra MCP round-trip to retrieve the same string. +DATAFRAME_MATERIALIZATION_OPT_OUT = frozenset({"tools.skills_manager"}) + + +class FakeMcp: + def __init__(self): + self.tools = {} + + def tool(self, name, description): + def decorator(func): + self.tools[name] = func + return func + + return decorator + + +def _literal_strings(tree: ast.AST) -> list[str]: + found: list[str] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Constant) and isinstance(node.value, str): + found.append(node.value) + elif isinstance(node, ast.JoinedStr): + parts: list[str] = [] + for value in node.values: + if isinstance(value, ast.Constant) and isinstance(value.value, str): + parts.append(value.value) + if parts: + found.append("".join(parts)) + return found + + +def test_support_message_is_defined_once_in_config(): + assert "https://github.com/Blazemeter/bzm-mcp/issues" in SUPPORT_MESSAGE + assert HARDCODED_SUPPORT_SNIPPET in SUPPORT_MESSAGE + + +def test_tool_modules_do_not_hardcode_the_support_message(): + offenders: list[str] = [] + for path in sorted(TOOLS_DIR.glob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + if any(HARDCODED_SUPPORT_SNIPPET in value for value in _literal_strings(tree)): + offenders.append(path.name) + assert offenders == [], ( + f"Hardcoded support text in {offenders}; import SUPPORT_MESSAGE from config.blazemeter" + ) + + +@pytest.mark.parametrize("module_name", MANAGER_MODULES) +def test_managers_pass_central_support_message(module_name, monkeypatch): + captured: dict = {} + + def fake_register(*args, **kwargs): + captured.update(kwargs) + + async def _tool(*_args, **_kwargs): + return None + + return _tool + + module = importlib.import_module(module_name) + monkeypatch.setattr(module, "register_managed_tool", fake_register) + module.register(FakeMcp(), build_runtime("stdio")) + + if module_name in OPT_OUT_MODULES: + assert captured.get("support_message") is None + return + + assert captured.get("support_message", SUPPORT_MESSAGE) is SUPPORT_MESSAGE + + +@pytest.mark.parametrize("module_name", MANAGER_MODULES) +def test_managers_dataframe_materialization_policy(module_name, monkeypatch): + captured: dict = {} + + def fake_register(*args, **kwargs): + captured.update(kwargs) + + async def _tool(*_args, **_kwargs): + return None + + return _tool + + module = importlib.import_module(module_name) + monkeypatch.setattr(module, "register_managed_tool", fake_register) + module.register(FakeMcp(), build_runtime("stdio")) + + if module_name in DATAFRAME_MATERIALIZATION_OPT_OUT: + assert captured.get("disable_materialization") is True + return + + assert not captured.get("disable_materialization") + + +def test_unexpected_error_appends_support_message_by_default(monkeypatch): + async def boom(*_args, **_kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr("tools.mcp_entrypoint.run_tool_with_runtime", boom) + mcp = FakeMcp() + register_managed_tool( + mcp, + build_runtime("stdio"), + name="demo", + description="demo", + dispatch=_unused_dispatch, + ) + + result = asyncio.run(mcp.tools["demo"]({"action": "read"}, ctx=None)) + assert result.error is not None + assert SUPPORT_MESSAGE in result.error + + +def test_opt_out_omits_support_message(monkeypatch): + async def boom(*_args, **_kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr("tools.mcp_entrypoint.run_tool_with_runtime", boom) + mcp = FakeMcp() + register_managed_tool( + mcp, + build_runtime("stdio"), + name="demo", + description="demo", + dispatch=_unused_dispatch, + support_message=None, + ) + + result = asyncio.run(mcp.tools["demo"]({"action": "read"}, ctx=None)) + assert result.error is not None + assert SUPPORT_MESSAGE not in result.error + + +async def _unused_dispatch(action, args, token, ctx): + raise AssertionError("dispatch should not run when run_tool_with_runtime is stubbed") diff --git a/tests/test_tool_result_wrapper.py b/tests/test_tool_result_wrapper.py new file mode 100644 index 0000000..50d2ce1 --- /dev/null +++ b/tests/test_tool_result_wrapper.py @@ -0,0 +1,35 @@ +""" +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 models.result import BaseResult, ToolResult + + +def test_tool_result_from_base_result_builds_pretty_text_and_structured_content(): + base = BaseResult(result=[{"ok": True}], info=["done"]) + wrapped = ToolResult.from_base_result(base) + + assert wrapped.isError is False + assert wrapped.structuredContent == base.model_dump(mode="json") + assert wrapped.content[0].type == "text" + assert wrapped.content[0].text == base.model_dump_json(indent=2) + assert "\n" in wrapped.content[0].text + + +def test_tool_result_from_base_result_marks_error(): + base = BaseResult(error="boom") + wrapped = ToolResult.from_base_result(base) + + assert wrapped.isError is True + assert wrapped.error == "boom" diff --git a/tests/test_tools_manager_dataframes.py b/tests/test_tools_manager_dataframes.py new file mode 100644 index 0000000..38ecdd7 --- /dev/null +++ b/tests/test_tools_manager_dataframes.py @@ -0,0 +1,84 @@ +""" +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 types import SimpleNamespace + +from config.storage import DefaultSessionScopeResolver, SessionScope +from config.token import BzmToken +from tests.conftest import make_ctx, run_async +from tools.dataframe_manager import register_dataframe, resolve_session_scope +from tools.tools_manager import ToolsManager + + +class TestResolveSessionScope: + def test_uses_token_id_and_ctx_session(self): + token = BzmToken("api-key-id", "secret") + ctx = SimpleNamespace(session_id="mcp-abc") + assert resolve_session_scope(ctx, token) == SessionScope("api-key-id", "mcp-abc") + + def test_defaults_when_missing(self): + assert resolve_session_scope(None, None) == SessionScope("anonymous", "default") + + +class TestToolsManagerDataframesAgainstStorage: + def test_list_query_remove_clear(self, in_memory_session_storage): + token = BzmToken("user-tools", "secret") + ctx = make_ctx(token, "session-tools") + manager = ToolsManager(ctx, in_memory_session_storage, DefaultSessionScopeResolver()) + scope = SessionScope("user-tools", "session-tools") + + meta = run_async( + register_dataframe( + result=[{"id": 1, "name": "a"}, {"id": 2, "name": "b"}], + origin_manager="tests", + origin_action="seed", + json_size_chars=9001, + session_storage=in_memory_session_storage, + scope=scope, + ) + ) + + listed = run_async(manager.dataframes_list()) + assert listed.error is None + assert listed.total == 1 + + queried = run_async( + manager.dataframes_query( + sql=( + f"SELECT id FROM {meta['table_name']} " + f"ORDER BY id LIMIT 100 OFFSET 0" + ) + ) + ) + assert queried.error is None + assert queried.total == 2 + + removed = run_async(manager.dataframes_remove([meta["dataframe_id"]])) + assert removed.error is None + assert run_async(manager.dataframes_list()).total == 0 + + run_async( + register_dataframe( + result=[{"id": 3}], + origin_manager="tests", + origin_action="seed2", + json_size_chars=9001, + session_storage=in_memory_session_storage, + scope=scope, + ) + ) + cleared = run_async(manager.dataframes_clear()) + assert cleared.result[0]["removed_count"] == 1 + assert run_async(manager.dataframes_list()).total == 0 diff --git a/tests/test_tools_manager_tasks.py b/tests/test_tools_manager_tasks.py new file mode 100644 index 0000000..2957488 --- /dev/null +++ b/tests/test_tools_manager_tasks.py @@ -0,0 +1,264 @@ +""" +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 + +from config.storage import DefaultSessionScopeResolver, InMemorySessionStorageProvider, SessionScope +from config.token import BzmToken +from models.result import BaseResult +from tests.conftest import make_ctx +from tools.async_task_manager import ( + STATUS_COMPLETED, + STATUS_WORKING, + TaskRecord, + configure_task_storage, + submit_task, +) +from tools.tools_manager import ToolsManager +import tools.async_task_manager as task_manager + +_TOKEN = BzmToken("local", "secret") +_SESSION_ID = "stdio" +_STDIO_SCOPE = SessionScope(user_id="local", mcp_session_id="stdio") + + +def _configure(): + store = InMemorySessionStorageProvider() + configure_task_storage(store) + return store + + +def _manager(store=None): + store = store or _configure() + ctx = make_ctx(_TOKEN, _SESSION_ID) + return ToolsManager(ctx, store, DefaultSessionScopeResolver()) + + +def test_operation_name_from_manager_method(): + action_payload = { + "manager": "ExecutionManager", + "method": "list", + "params": { + "test_id": 15332595, + "limit": 1, + "offset": 0, + "purpose": "diagnostics", + }, + } + + line = ToolsManager._operation_name(action_payload) + assert line == "execution.list" + + +def test_polling_message_includes_operation_task_and_batch_summary(): + store = _configure() + + async def scenario(): + record = TaskRecord( + task_id="7k2p9m4q", + action={ + "manager": "ExecutionManager", + "method": "list", + "params": {"test_id": 15332595, "limit": 1, "offset": 0}, + }, + created_at=0.0, + last_updated_at=0.0, + time_to_live_ms=None, + status=STATUS_WORKING, + status_message="Task is currently running.", + status_info="", + user_id="local", + mcp_session_id="stdio", + ) + cache = await task_manager._get_or_create_cache(_STDIO_SCOPE) + async with cache.lock: + cache.hydrated = True + cache.tasks[record.task_id] = record + await task_manager._commit_cache(cache, _STDIO_SCOPE) + + manager = _manager(store) + message = await manager._polling_message( + task_record=record, + poll_count=3, + elapsed_seconds=12, + next_poll_seconds=1.0, + window_seconds=30.0, + ) + + assert "Polling 7k2p9m4q[execution.list] (working) attempt=3 elapsed=12s/30s next=1s" in message + assert "batch summary: total=1 completed=0 working=1 parking=0 failed=0" in message + + asyncio.run(scenario()) + + +def test_tasks_list_returns_minimal_snapshot_without_action_payload(): + store = _configure() + + async def scenario(): + record = TaskRecord( + task_id="abc123xy", + action={ + "manager": "TestManager", + "method": "upload_assets", + "params": {"test_id": 1, "file_paths": ["/very/long/path"]}, + }, + created_at=0.0, + last_updated_at=0.0, + time_to_live_ms=None, + status=STATUS_WORKING, + status_message="Task is currently running.", + status_info="", + user_id="local", + mcp_session_id="stdio", + ) + cache = await task_manager._get_or_create_cache(_STDIO_SCOPE) + async with cache.lock: + cache.hydrated = True + cache.tasks[record.task_id] = record + await task_manager._commit_cache(cache, _STDIO_SCOPE) + + manager = _manager(store) + response = await manager.tasks_list() + assert response.result is not None + item = response.result[0] + assert item["task_id"] == "abc123xy" + assert item["operation"] == "test.upload_assets" + assert "action" not in item + + asyncio.run(scenario()) + + +def test_tasks_status_terminal_omits_task_result_payload(): + store = _configure() + + async def scenario(): + record = TaskRecord( + task_id="done1234", + action={"manager": "ExecutionManager", "method": "list", "params": {"limit": 1, "offset": 0}}, + created_at=0.0, + last_updated_at=0.0, + time_to_live_ms=None, + status=STATUS_COMPLETED, + status_message="Task completed.", + status_info="", + result=BaseResult(result=[{"id": 1, "name": "result"}]), + user_id="local", + mcp_session_id="stdio", + ) + cache = await task_manager._get_or_create_cache(_STDIO_SCOPE) + async with cache.lock: + cache.hydrated = True + cache.tasks[record.task_id] = record + await task_manager._commit_cache(cache, _STDIO_SCOPE) + + manager = _manager(store) + response = await manager.tasks_status("done1234") + assert response.result is not None + item = response.result[0] + assert item["task_id"] == "done1234" + assert "task_result" not in item + assert response.info is not None + assert "Use tasks_get to retrieve task_result" in response.info[0] + + asyncio.run(scenario()) + + +def test_tasks_get_terminal_includes_task_result_payload(): + store = _configure() + + async def scenario(): + record = TaskRecord( + task_id="done5678", + action={"manager": "ExecutionManager", "method": "list", "params": {"limit": 1, "offset": 0}}, + created_at=0.0, + last_updated_at=0.0, + time_to_live_ms=None, + status=STATUS_COMPLETED, + status_message="Task completed.", + status_info="", + result=BaseResult(result=[{"id": 2, "name": "final"}]), + user_id="local", + mcp_session_id="stdio", + ) + cache = await task_manager._get_or_create_cache(_STDIO_SCOPE) + async with cache.lock: + cache.hydrated = True + cache.tasks[record.task_id] = record + await task_manager._commit_cache(cache, _STDIO_SCOPE) + + manager = _manager(store) + response = await manager.tasks_get("done5678", remove_on_terminal=False) + assert response.result is not None + item = response.result[0] + assert item["task_id"] == "done5678" + assert item["task_result"]["result"] == [{"id": 2, "name": "final"}] + + asyncio.run(scenario()) + + +def test_execute_with_task_management_fast_path(): + from tools.utils import execute_with_task_management + + _configure() + + async def scenario(): + async def action(): + return BaseResult(result=[{"fast": True}]) + + result = await execute_with_task_management( + action_payload={"manager": "TestManager", "method": "read"}, + coro_factory=action, + fast_response_threshold_seconds=2.0, + scope=_STDIO_SCOPE, + ) + assert result.error is None + assert result.result == [{"fast": True}] + + asyncio.run(scenario()) + + +def test_execute_with_task_management_async_handoff(): + from tools.utils import execute_with_task_management + + store = _configure() + + async def scenario(): + async def action(): + await asyncio.sleep(0.3) + return BaseResult(result=[{"slow": True}]) + + result = await execute_with_task_management( + action_payload={"manager": "TestManager", "method": "read"}, + coro_factory=action, + fast_response_threshold_seconds=0.05, + scope=_STDIO_SCOPE, + ) + assert result.result is not None + assert "task_id" in result.result[0] + assert result.result[0]["status"] in {"parking", "working", "completed"} + assert result.info is not None + assert "tasks_status" in result.info[0] + + task_id = result.result[0]["task_id"] + manager = _manager(store) + for _ in range(50): + status = await manager.tasks_status(task_id) + if status.result and status.result[0]["status"] == STATUS_COMPLETED: + break + await asyncio.sleep(0.02) + got = await manager.tasks_get(task_id, remove_on_terminal=True) + assert got.result[0]["task_result"]["result"] == [{"slow": True}] + + asyncio.run(scenario()) diff --git a/tests/test_utils_normalize_action_args.py b/tests/test_utils_normalize_action_args.py new file mode 100644 index 0000000..8c82452 --- /dev/null +++ b/tests/test_utils_normalize_action_args.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. +""" +from tools.utils import normalize_action_args + + +def test_normalize_action_args_standard_format(): + """Standard format: action + args nested.""" + action, args = normalize_action_args({ + "action": "list", + "args": {"limit": 5, "project_id": 158903, "result_format": "dataframe"}, + }) + assert action == "list" + assert args == {"limit": 5, "project_id": 158903, "result_format": "dataframe"} + + +def test_normalize_action_args_flat_format(): + """Flat format: action + params at top level merged into args.""" + action, args = normalize_action_args({ + "action": "read", + "test_id": 123, + }) + assert action == "read" + assert args == {"test_id": 123} + + +def test_normalize_action_args_double_wrapped(): + """Double-wrapped format: {"arguments": {"action": "x", "args": {...}}}.""" + action, args = normalize_action_args({ + "arguments": { + "action": "list", + "args": { + "limit": 5, + "project_id": 158903, + "result_format": "dataframe", + }, + }, + }) + assert action == "list" + assert args == {"limit": 5, "project_id": 158903, "result_format": "dataframe"} + + +def test_normalize_action_args_double_wrapped_with_action_only(): + """Double-wrapped with only action (no args) still unwraps.""" + action, args = normalize_action_args({ + "arguments": {"action": "list_help_categories"}, + }) + assert action == "list_help_categories" + assert args == {} + + +def test_normalize_action_args_does_not_unwrap_when_extra_keys(): + """When top-level has other keys besides 'arguments', do not unwrap.""" + action, args = normalize_action_args({ + "arguments": {"action": "x", "args": {}}, + "other_key": "value", + }) + assert action == "" + assert "arguments" in args + assert args["other_key"] == "value" diff --git a/tests/test_utils_required_args.py b/tests/test_utils_required_args.py new file mode 100644 index 0000000..94038f0 --- /dev/null +++ b/tests/test_utils_required_args.py @@ -0,0 +1,32 @@ +""" +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 models.result import BaseResult +from tools.utils import validate_required_args, validate_non_empty_str_arg + + +def test_validate_required_args_missing(): + err = validate_required_args("read", {}, ["account_id"]) + assert err is not None + assert "account_id" in err.error + + +def test_validate_required_args_present(): + assert validate_required_args("read", {"account_id": 1}, ["account_id"]) is None + + +def test_validate_non_empty_str_arg(): + assert validate_non_empty_str_arg("x", {"task_id": " "}, "task_id") is not None + assert validate_non_empty_str_arg("x", {"task_id": "abc"}, "task_id") is None diff --git a/tools/account_manager.py b/tools/account_manager.py index 42662e9..f21bcc7 100644 --- a/tools/account_manager.py +++ b/tools/account_manager.py @@ -14,7 +14,7 @@ limitations under the License. """ from typing import Optional, Dict, Any -import httpx + from mcp.server.fastmcp import Context from config.blazemeter import ACCOUNTS_ENDPOINT, TOOLS_PREFIX, SUPPORT_MESSAGE @@ -22,8 +22,8 @@ from formatters.account import format_accounts from models.manager import Manager from models.result import BaseResult -from telemetry import run_tool -from tools.utils import api_request, format_sanitized_traceback +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import api_request, run_as_task class AccountManager(Manager): @@ -38,6 +38,7 @@ def __init__( ): super().__init__(ctx) + @run_as_task() async def read(self, account_id: Optional[int]) -> BaseResult: if not isinstance(account_id, int) or account_id < 1: return BaseResult(error="Missing or invalid required argument 'account_id'. Expected integer.") @@ -59,6 +60,7 @@ async def read(self, account_id: Optional[int]) -> BaseResult: else: return account_result + @run_as_task() async def list(self, limit: int = 50, offset: int = 0) -> BaseResult: if not isinstance(limit, int) or not isinstance(offset, int): return BaseResult(error="Invalid arguments 'limit'/'offset'. Expected integers.") @@ -79,8 +81,24 @@ async def list(self, limit: int = 50, offset: int = 0) -> BaseResult: params=parameters ) -def register(mcp, runtime: AppRuntime) -> None: - @mcp.tool( + +def register(mcp, runtime: AppRuntime): + + async def _dispatch(action, args, token, ctx): + account_manager = AccountManager(ctx) + match action: + case "read": + return await account_manager.read(args.get("account_id")) + case "list": + return await account_manager.list(args.get("limit", 50), args.get("offset", 0)) + case _: + return BaseResult( + error=f"Action {action} not found in account manager tool" + ) + + register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_account", description=""" Operations on account users. @@ -97,30 +115,7 @@ def register(mcp, runtime: AppRuntime) -> None: - If you need to get the default account, use the project id to get the workspace and with that the account. - Use the read operation if AI consent information is needed. The AI Consent it's located at account level. - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + support_message=SUPPORT_MESSAGE, ) - async def account(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - runtime.configure_context(ctx) - account_manager = AccountManager(ctx) - - async def _dispatch(): - match action: - case "read": - return await account_manager.read(args.get("account_id")) - case "list": - return await account_manager.list(args.get("limit", 50), args.get("offset", 0)) - case _: - return BaseResult( - error=f"Action {action} not found in account manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_account", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}" - ) diff --git a/tools/async_task_manager.py b/tools/async_task_manager.py new file mode 100644 index 0000000..9a42206 --- /dev/null +++ b/tools/async_task_manager.py @@ -0,0 +1,621 @@ +""" +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.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, +) + +# Match DefaultSessionScopeResolver fallbacks when token/ctx are absent. +DEFAULT_USER_ID = "anonymous" +DEFAULT_SESSION_ID = "default" +DEFAULT_SCOPE = SessionScope(user_id=DEFAULT_USER_ID, mcp_session_id=DEFAULT_SESSION_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 = 10 + +TASK_ID_MAX_ATTEMPTS = 10 +_MAX_SESSION_CACHES = 256 + +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() + + +@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 + user_id: str = DEFAULT_USER_ID + mcp_session_id: str = DEFAULT_SESSION_ID + + 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() + + +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 DEFAULT_USER_ID), + mcp_session_id=str(payload.get("mcp_session_id") or DEFAULT_SESSION_ID), + ) + + +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]]): + await _set_status_and_persist( + task_record, + STATUS_PARKING, + "Task is waiting for an available execution slot.", + ) + try: + 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.result = BaseResult(error=error_message) + await _set_status_and_persist(task_record, STATUS_FAILED, error_message) + + +async def submit_task( + action: Dict[str, Any], + coro_factory: Callable[[], Awaitable[Any]], + time_to_live_ms: Optional[int] = None, + scope: SessionScope = DEFAULT_SCOPE, +) -> 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 = DEFAULT_SCOPE, +) -> 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 = DEFAULT_SCOPE, +) -> 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 = DEFAULT_SCOPE, +) -> 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 = DEFAULT_SCOPE, +) -> 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: + """Resolve Storage partition keys from a Manager instance (token + ctx).""" + return resolve_session_scope( + getattr(manager, "ctx", None), + token=getattr(manager, "token", None), + scope_resolver=getattr(manager, "scope_resolver", None), + ) diff --git a/tools/billing_manager.py b/tools/billing_manager.py index 6dbf269..656733f 100644 --- a/tools/billing_manager.py +++ b/tools/billing_manager.py @@ -15,7 +15,6 @@ """ from typing import Dict, Any -import httpx from mcp.server.fastmcp import Context from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE @@ -23,8 +22,8 @@ from models.manager import Manager from models.result import BaseResult from tools.billing_utils import calculate_test_cost -from telemetry import run_tool -from tools.utils import format_sanitized_traceback +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import run_as_task class BillingManager(Manager): @@ -35,6 +34,7 @@ def __init__( ): super().__init__(ctx) + @run_as_task() async def calculate_cost_from_config(self, args: Dict) -> BaseResult: result = calculate_test_cost(args) return BaseResult(result=[ @@ -48,9 +48,21 @@ async def calculate_cost_from_config(self, args: Dict) -> BaseResult: } ]) +def register(mcp, runtime: AppRuntime): -def register(mcp, runtime: AppRuntime) -> None: - @mcp.tool( + async def _dispatch(action, args, token, ctx): + billing_manager = BillingManager(ctx) + match action: + case "calculate_cost_from_config": + return await billing_manager.calculate_cost_from_config(args) + case _: + return BaseResult( + error=f"Action {action} not found in billing manager tool" + ) + + register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_billing", description=""" Operations on Billing. @@ -87,28 +99,7 @@ def register(mcp, runtime: AppRuntime) -> None: - Server hours calculation can use provided number_of_servers or estimate based on concurrency (~1000 users per engine). - All calculations are based on official BlazeMeter documentation (blazemeter-usage-billing skill). - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + support_message=SUPPORT_MESSAGE, ) - async def billing(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - runtime.configure_context(ctx) - billing_manager = BillingManager(ctx) - - async def _dispatch(): - match action: - case "calculate_cost_from_config": - return await billing_manager.calculate_cost_from_config(args) - case _: - return BaseResult( - error=f"Action {action} not found in billing manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_billing", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}" - ) diff --git a/tools/dataframe_manager.py b/tools/dataframe_manager.py new file mode 100644 index 0000000..c7cc638 --- /dev/null +++ b/tools/dataframe_manager.py @@ -0,0 +1,1164 @@ +""" +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.token import BzmToken +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, + ) + + +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"], + } + +_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 + + 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, + ) + 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 ( + 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 base_result.result[0].get("dataframe_id") + ): + # Avoid rematerializing a payload that is already a dataframe reference. + return base_result + try: + compact_json = serialize_result_to_compact_json(base_result.result) + 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(base_result.result) + 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)] + base_result.append_info([ + "Large result was stored as a session dataframe. Use blazemeter_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'.", + ]) + 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, + token: Optional[BzmToken] = None, + ctx: Any = 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 token/ctx resolution (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: + scope = resolve_session_scope(ctx, token=token, scope_resolver=scope_resolver) + 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/execution_manager.py b/tools/execution_manager.py index 4f6e427..bea760c 100644 --- a/tools/execution_manager.py +++ b/tools/execution_manager.py @@ -25,8 +25,8 @@ from models.result import BaseResult from tools import bridge, search_utils from tools.report_manager import ReportManager -from telemetry import run_tool -from tools.utils import api_request, timeout, user_agent, format_sanitized_traceback, require_confirmation, Operations +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import api_request, timeout, user_agent, format_sanitized_traceback, require_confirmation, Operations, run_as_task class ExecutionManager(Manager): @@ -103,6 +103,7 @@ def _handle_analyzer_http_error(self, e: httpx.HTTPStatusError, execution_id: in return BaseResult(error=f"HTTP {status_code}: {e.response.text[:200]}") @require_confirmation(operation=Operations.CREATE) + @run_as_task() async def start(self, test_id: Optional[int], delayed_start_ready: bool = True, is_debug_run: bool = False) -> BaseResult: if not isinstance(test_id, int) or test_id < 1: @@ -124,6 +125,7 @@ async def start(self, test_id: Optional[int], delayed_start_ready: bool = True, json=start_body ) + @run_as_task() async def read(self, execution_id: Optional[int]) -> BaseResult: if not isinstance(execution_id, int) or execution_id < 1: return BaseResult(error="Missing or invalid required argument 'execution_id'. Expected integer.") @@ -188,6 +190,7 @@ def _get_execution_status_context() -> str: "When it is archived, it is not possible to read the detailed execution information.\n" ) + @run_as_task() async def list(self, test_id: Optional[int], limit: int = 50, offset: int = 0) -> BaseResult: if not isinstance(test_id, int) or test_id < 1: return BaseResult(error="Missing or invalid required argument 'test_id'. Expected integer.") @@ -213,6 +216,7 @@ async def list(self, test_id: Optional[int], limit: int = 50, offset: int = 0) - params=parameters ) + @run_as_task() async def search(self, args: dict[str, Any]) -> BaseResult: # Check if it's valid or allowed @@ -225,6 +229,7 @@ async def search(self, args: dict[str, Any]) -> BaseResult: return await search_utils.test_execution_search("master", self.token, account_id, args) + @run_as_task() async def search_filter_values(self, account_id: int, filter_names: List[str]) -> BaseResult: # Check if it's valid or allowed @@ -239,6 +244,7 @@ async def search_filter_values(self, account_id: int, filter_names: List[str]) - filter_names, ) + @run_as_task() async def ai_analysis(self, execution_id: Optional[int]) -> BaseResult: if not isinstance(execution_id, int) or execution_id < 1: return BaseResult(error="Missing or invalid required argument 'execution_id'. Expected integer.") @@ -400,6 +406,7 @@ def _build_analysis_result(self, execution_id: int, analysis_state: Dict[str, An } return BaseResult(result=[result]) + @run_as_task() async def read_all_reports(self, execution_id: Optional[int]) -> BaseResult: if not isinstance(execution_id, int) or execution_id < 1: return BaseResult(error="Missing or invalid required argument 'execution_id'. Expected integer.") @@ -438,9 +445,46 @@ def _get_analysis_context_message(is_ready: bool, status_message: str) -> str: "The analysis will be available once processing is complete." ) - def register(mcp, runtime: AppRuntime): - @mcp.tool( + async def _dispatch(action, args, token, ctx): + execution_manager = ExecutionManager(ctx) + report_manager = ReportManager(ctx) + match action: + case "start": + return await execution_manager.start(args.get("test_id")) + case "read": + return await execution_manager.read(args.get("execution_id")) + case "list": + return await execution_manager.list( + args.get("test_id"), + args.get("limit", 50), + args.get("offset", 0), + ) + case "search": + return await execution_manager.search(args) + case "search_filter_values": + return await execution_manager.search_filter_values(args.get("account_id"), + args.get("filter_names", [])) + case "read_summary": + return await report_manager.read_summary(args.get("execution_id")) + case "read_errors": + return await report_manager.read_error(args.get("execution_id")) + case "read_request_stats": + return await report_manager.read_request_stats(args.get("execution_id")) + case "read_all_reports": + return await execution_manager.read_all_reports(args.get("execution_id")) + case "read_anomalies_stats": + return await report_manager.read_anomalies_stats(args.get("execution_id")) + case "ai_analysis": + return await execution_manager.ai_analysis(args.get("execution_id")) + case _: + return BaseResult( + error=f"Action {action} not found in test execution manager tool" + ) + + register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_execution", description=""" Operations on tests executions and results reports. @@ -502,54 +546,7 @@ def register(mcp, runtime: AppRuntime): or create a new analysis entry. It provides dynamic responses indicating whether the analysis is ready or still processing. Hints: - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + support_message=SUPPORT_MESSAGE, ) - async def execution(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - runtime.configure_context(ctx) - execution_manager = ExecutionManager(ctx) - report_manager = ReportManager(ctx) - - async def _dispatch(): - match action: - case "start": - return await execution_manager.start(args.get("test_id")) - case "read": - return await execution_manager.read(args.get("execution_id")) - case "list": - return await execution_manager.list( - args.get("test_id"), - args.get("limit", 50), - args.get("offset", 0), - ) - case "search": - return await execution_manager.search(args) - case "search_filter_values": - return await execution_manager.search_filter_values(args.get("account_id"), - args.get("filter_names", [])) - case "read_summary": - return await report_manager.read_summary(args.get("execution_id")) - case "read_errors": - return await report_manager.read_error(args.get("execution_id")) - case "read_request_stats": - return await report_manager.read_request_stats(args.get("execution_id")) - case "read_all_reports": - return await execution_manager.read_all_reports(args.get("execution_id")) - case "read_anomalies_stats": - return await report_manager.read_anomalies_stats(args.get("execution_id")) - case "ai_analysis": - return await execution_manager.ai_analysis(args.get("execution_id")) - case _: - return BaseResult( - error=f"Action {action} not found in test execution manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_execution", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}" - ) diff --git a/tools/help_manager.py b/tools/help_manager.py index 81b9969..135c1d6 100644 --- a/tools/help_manager.py +++ b/tools/help_manager.py @@ -20,7 +20,6 @@ import httpx from mcp.server.fastmcp import Context -from pydantic import Field from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE, \ HELP_INDEX_URL, HELP_TOC_URL, HELP_BASE_CONTENT_URL @@ -29,8 +28,8 @@ from models.manager import Manager from models.result import BaseResult from tools.help_utils import convert_js_to_py_dict -from telemetry import run_tool -from tools.utils import http_request, format_sanitized_traceback +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import http_request, format_sanitized_traceback, run_as_task class HelpManager(Manager): @@ -140,6 +139,7 @@ async def fetch_chunk(chunk_url: str): help_tree['root_category'] = help_tree.pop('') # Assign a name to the root category HelpManager.help_tree = help_tree + @run_as_task() async def list_help_categories(self) -> BaseResult: if HelpManager.help_tree is None: await self._load_help_tree() @@ -155,6 +155,7 @@ async def list_help_categories(self) -> BaseResult: info=["A list of subcategories is provided for each category"] ) + @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( @@ -233,6 +234,7 @@ async def get_help_object(category_id: str, subcategory_id: str, help_id: str) - return help_object + @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( @@ -257,9 +259,70 @@ async def read_help_info(self, category_id: str, subcategory_id: str, help_id_li }], ) - def register(mcp, runtime: AppRuntime): - @mcp.tool( + help_tool = None + + async def _dispatch(action, args, token, ctx): + args = args or {} + help_manager = HelpManager(ctx) + match action: + case "list_help_categories": + return await help_manager.list_help_categories() + case "list_help_category_content": + return await help_manager.list_help_category_content( + args.get("category_id", "home"), + args.get("subcategory_id_list") + ) + case "read_help_info": + return await help_manager.read_help_info( + args.get("category_id", "home"), + args.get("subcategory_id", ""), + args.get("help_id_list") + ) + case "batch": + # Make sure this initialization doesn't run in parallel + if HelpManager.help_tree is None: + await help_manager._load_help_tree() + + batch_calls = args.get("batch_calls", []) + if not isinstance(batch_calls, list) or not batch_calls: + return BaseResult( + error="batch_calls must be a non-empty list of dicts with 'action' and 'args'") + + semaphore = asyncio.Semaphore(HelpManager.MAX_BATCH_CONCURRENCY) + + async def process_call(call: Dict[str, Any]) -> BaseResult | List[BaseResult]: + sub_action = call.get("action", "") + sub_args = call.get("args", {}) + async with semaphore: + try: + # Recursively call the help function itself + return await help_tool({"action": sub_action, "args": sub_args}, ctx) + except httpx.HTTPStatusError: + return BaseResult( + error=f"HTTP error in sub-action {sub_action}: {format_sanitized_traceback()}" + ) + except Exception: + return BaseResult( + error=f"Error in sub-action {sub_action}: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}") + + # Parallel execution with asyncio.gather + results = await asyncio.gather(*[process_call(call) for call in batch_calls], + return_exceptions=True) + # Handle any exceptions returned + processed_results = [ + r if not isinstance(r, Exception) else BaseResult(error=f"Unhandled exception: {str(r)}") + for r in results + ] + return BaseResult(result=processed_results) + case _: + return BaseResult( + error=f"Action {action} not found in help manager tool" + ) + + help_tool = register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_help", description=""" Operations on documentation and help information. @@ -283,82 +346,8 @@ def register(mcp, runtime: AppRuntime): - Always generates the url attributes as a link in markdown format (like command_url). - **CRITICAL**: For multiple actions, always use the 'batch' action. - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + excluded_actions={'batch'}, + support_message=SUPPORT_MESSAGE, ) - async def help_main( - action: str = Field(description="The action id to execute"), - args: Dict[str, Any] = Field(description="Dictionary with parameters", default=None), - ctx: Context = Field(description="Context object providing access to MCP capabilities") - ) -> BaseResult: - if args is None: - args = {} - - runtime.configure_context(ctx) - help_manager = HelpManager(ctx) - - async def _dispatch(): - match action: - case "list_help_categories": - return await help_manager.list_help_categories() - case "list_help_category_content": - return await help_manager.list_help_category_content( - args.get("category_id", "home"), - args.get("subcategory_id_list") - ) - case "read_help_info": - return await help_manager.read_help_info( - args.get("category_id", "home"), - args.get("subcategory_id", ""), - args.get("help_id_list") - ) - case "batch": - # Make sure this initialization doesn't run in parallel - if HelpManager.help_tree is None: - await help_manager._load_help_tree() - - batch_calls = args.get("batch_calls", []) - if not isinstance(batch_calls, list) or not batch_calls: - return BaseResult( - error="batch_calls must be a non-empty list of dicts with 'action' and 'args'") - - semaphore = asyncio.Semaphore(HelpManager.MAX_BATCH_CONCURRENCY) - - async def process_call(call: Dict[str, Any]) -> BaseResult | List[BaseResult]: - sub_action = call.get("action", "") - sub_args = call.get("args", {}) - async with semaphore: - try: - # Recursively call the help function itself - return await help_main(sub_action, sub_args, ctx) - except httpx.HTTPStatusError: - return BaseResult( - error=f"HTTP error in sub-action {sub_action}: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"Error in sub-action {sub_action}: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}") - - # Parallel execution with asyncio.gather - results = await asyncio.gather(*[process_call(call) for call in batch_calls], - return_exceptions=True) - # Handle any exceptions returned - processed_results = [ - r if not isinstance(r, Exception) else BaseResult(error=f"Unhandled exception: {str(r)}") - for r in results - ] - return BaseResult(result=processed_results) - case _: - return BaseResult( - error=f"Action {action} not found in help manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_help", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}" - ) diff --git a/tools/mcp_entrypoint.py b/tools/mcp_entrypoint.py new file mode 100644 index 0000000..58e1894 --- /dev/null +++ b/tools/mcp_entrypoint.py @@ -0,0 +1,98 @@ +""" +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 typing import Any, Awaitable, Callable, Dict, Optional, Set + +import httpx +from mcp.server.fastmcp import Context + +from config.blazemeter import SUPPORT_MESSAGE +from config.runtime import AppRuntime +from config.token import BzmToken +from models.result import BaseResult +from tools.runtime_tools import run_tool_with_runtime +from tools.utils import ( + format_sanitized_traceback, + normalize_action_args, + tool_result, +) + +ToolDispatch = Callable[ + [str, Dict[str, Any], Optional[BzmToken], Context], + Awaitable[BaseResult], +] + + +def register_managed_tool( + mcp: Any, + runtime: AppRuntime, + *, + name: str, + description: str, + dispatch: ToolDispatch, + excluded_actions: Optional[Set[str]] = None, + disable_materialization: bool = False, + support_message: Optional[str] = SUPPORT_MESSAGE, +) -> Callable[..., Awaitable[BaseResult]]: + """ + Shared MCP tool entrypoint: arguments= normalize → configure_context → + run_tool_with_runtime → @tool_result wrap. + + ``dispatch(action, args, token, ctx)`` owns action routing and validation. + Materialization stays inside ``run_tool_with_runtime`` so tracing includes persist. + Returns the registered tool coroutine (needed for help/skills batch re-entry). + """ + + @mcp.tool(name=name, description=description) + @tool_result( + excluded_actions=excluded_actions, + disable_materialization=True, + ) + async def _tool( + arguments: Dict[str, Any] = None, + ctx: Context = None, + ) -> BaseResult: + action, args = normalize_action_args(arguments) + if not action: + return BaseResult(error="Missing required argument 'action' within tool arguments.") + runtime.configure_context(ctx) + token = runtime.auth.get_token(ctx) + + async def _run() -> BaseResult: + return await dispatch(action, args, token, ctx) + + try: + return await run_tool_with_runtime( + runtime, + name, + action, + ctx, + _run, + token=token, + tool_args=args, + dataframe_excluded_actions=excluded_actions, + disable_dataframe_materialization=disable_materialization, + ) + except httpx.HTTPStatusError: + return BaseResult(error=f"Error: {format_sanitized_traceback()}") + except Exception: + detail = format_sanitized_traceback() + if support_message: + return BaseResult(error=f"Error: {detail}\n{support_message}") + return BaseResult(error=f"Error: {detail}") + + return _tool diff --git a/tools/project_manager.py b/tools/project_manager.py index a8c9d8c..42db71c 100644 --- a/tools/project_manager.py +++ b/tools/project_manager.py @@ -15,17 +15,16 @@ """ from typing import Optional, Dict, Any -import httpx from mcp.server.fastmcp import Context -from config.blazemeter import TOOLS_PREFIX, PROJECTS_ENDPOINT +from config.blazemeter import TOOLS_PREFIX, PROJECTS_ENDPOINT, SUPPORT_MESSAGE from config.runtime import AppRuntime from formatters.project import format_projects from models.manager import Manager from models.result import BaseResult from tools import bridge -from telemetry import run_tool -from tools.utils import api_request, format_sanitized_traceback +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import api_request, run_as_task class ProjectManager(Manager): @@ -36,6 +35,7 @@ def __init__( ): super().__init__(ctx) + @run_as_task() async def read(self, project_id: Optional[int]) -> BaseResult: if not isinstance(project_id, int) or project_id < 1: return BaseResult(error="Missing or invalid required argument 'project_id'. Expected integer.") @@ -60,6 +60,7 @@ async def read(self, project_id: Optional[int]) -> BaseResult: project_element.tests_count = await bridge.count_project_tests(self.token, self.ctx, project_id) return project_result + @run_as_task() async def list(self, workspace_id: Optional[int], limit: int = 50, offset: int = 0) -> BaseResult: if not isinstance(workspace_id, int) or workspace_id < 1: return BaseResult(error="Missing or invalid required argument 'workspace_id'. Expected integer.") @@ -86,8 +87,24 @@ async def list(self, workspace_id: Optional[int], limit: int = 50, offset: int = params=parameters ) + def register(mcp, runtime: AppRuntime): - @mcp.tool( + + async def _dispatch(action, args, token, ctx): + project_manager = ProjectManager(ctx) + match action: + case "read": + return await project_manager.read(args.get("project_id")) + case "list": + return await project_manager.list(args.get("workspace_id"), args.get("limit", 10), args.get("offset", 0)) + case _: + return BaseResult( + error=f"Action {action} not found in project manager tool" + ) + + register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_project", description=""" Operations on projects. @@ -105,31 +122,7 @@ def register(mcp, runtime: AppRuntime): - For a particular project, go directly to the read action (you don't need account or workspace information). - Reading also allows you to obtain the number of tests the project has without having to use a list to count. - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + support_message=SUPPORT_MESSAGE, ) - async def project(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - runtime.configure_context(ctx) - project_manager = ProjectManager(ctx) - - async def _dispatch(): - match action: - case "read": - return await project_manager.read(args.get("project_id")) - case "list": - return await project_manager.list(args.get("workspace_id"), args.get("limit", 10), args.get("offset", 0)) - case _: - return BaseResult( - error=f"Action {action} not found in project manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_project", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"""Error: {format_sanitized_traceback()} - If you think this is a bug, please contact BlazeMeter support or report issue at https://github.com/BlazeMeter/bzm-mcp/issues""" - ) diff --git a/tools/report_manager.py b/tools/report_manager.py index 8b7bec7..7932d10 100644 --- a/tools/report_manager.py +++ b/tools/report_manager.py @@ -27,7 +27,7 @@ from models.manager import Manager from models.result import BaseResult from tools import bridge -from tools.utils import api_request +from tools.utils import api_request, run_as_task EXECUTION_ARCHIVED_MSG = ("Execution report is archived. It is not possible to read execution " "information from an archived execution.") @@ -54,6 +54,7 @@ def _evaluate_archived(execution_result: BaseResult) -> bool: return (execution_result.result and len(execution_result.result) > 0 and execution_result.result[0].get("result").archived) + @run_as_task() async def read_summary(self, master_id: int): execution_result = await bridge.read_execution(self.token, self.ctx, master_id) if execution_result.error: @@ -78,6 +79,7 @@ async def read_summary(self, master_id: int): } ) + @run_as_task() async def read_error(self, master_id: Optional[int]): """ Get error report for a given master_id with formatted, AI-friendly structure. @@ -108,6 +110,7 @@ async def read_error(self, master_id: Optional[int]): } ) + @run_as_task() async def read_request_stats(self, master_id: Optional[int]): """ Get request statistics report for a given master_id with formatted, AI-friendly structure. @@ -139,6 +142,7 @@ async def read_request_stats(self, master_id: Optional[int]): } ) + @run_as_task() async def read_anomalies_stats(self, master_id: Optional[int]): """ Get anomaly statistics for a given master_id (test execution). diff --git a/tools/runtime_tools.py b/tools/runtime_tools.py new file mode 100644 index 0000000..5fd3ca4 --- /dev/null +++ b/tools/runtime_tools.py @@ -0,0 +1,69 @@ +""" +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 typing import Any, Awaitable, Callable, Optional + +from config.runtime import AppRuntime +from telemetry import run_tool +from tools.utils import ( + reset_disable_dataframe_materialization, + set_disable_dataframe_materialization, +) + + +async def run_tool_with_runtime( + runtime: AppRuntime, + tool_name: str, + action: str, + ctx: Any, + dispatch: Callable[[], Awaitable[Any]], + *, + token: Any = None, + tool_args: Any = None, + dataframe_excluded_actions: Optional[set[str]] = None, + disable_dataframe_materialization: bool = False, +) -> Any: + """ + Run a tool action inside telemetry, then materialize large results via SessionStoragePort. + + Managers pass ``runtime`` once; tracing stays unaware of dataframe types. + Materialization runs inside the tool span so duration includes the commit. + """ + resolved_token = token if token is not None else runtime.auth.get_token(ctx) + + 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 + + # Idempotent if the async task runner already materialized the payload. + return await finalize_tool_result( + result, + action=action, + args=tool_args, + origin_manager=tool_name, + session_storage=runtime.storage, + scope_resolver=runtime.scope_resolver, + token=resolved_token, + ctx=ctx, + excluded_actions=dataframe_excluded_actions, + ) + finally: + reset_disable_dataframe_materialization(policy_token) + + return await run_tool(tool_name, action, ctx, _dispatch_and_finalize) diff --git a/tools/skills_manager.py b/tools/skills_manager.py index 993935f..3b42f12 100644 --- a/tools/skills_manager.py +++ b/tools/skills_manager.py @@ -19,14 +19,13 @@ import httpx from mcp.server.fastmcp import Context -from pydantic import Field from config.blazemeter import TOOLS_PREFIX, SUPPORT_MESSAGE from config.runtime import AppRuntime from models.manager import Manager from models.result import BaseResult -from telemetry import run_tool -from tools.utils import format_sanitized_traceback +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import format_sanitized_traceback, run_as_task from tools.skills_utils import list_skills, read_skill_definition, read_skill_file, parse_skill_uri, \ is_skill_uri, list_skill_resources_uri @@ -48,8 +47,8 @@ def __init__( ): super().__init__(ctx) - @staticmethod - async def list_skills() -> BaseResult: + @run_as_task() + async def list_skills(self) -> BaseResult: errors = [] if SkillsManager.skills is None: skills, errors = list_skills() @@ -62,8 +61,8 @@ async def list_skills() -> BaseResult: error=errors[0] if errors and len(errors) > 0 else None # Only the first error ) - @staticmethod - async def read_skill(skill_name: Optional[str]) -> BaseResult: + @run_as_task() + async def read_skill(self, skill_name: Optional[str]) -> BaseResult: if not isinstance(skill_name, str) or not skill_name.strip(): return BaseResult( error="Missing required argument 'skill_name'. Please specify a non-empty skill name." @@ -83,8 +82,8 @@ async def read_skill(skill_name: Optional[str]) -> BaseResult: error=error ) - @staticmethod - async def read_skill_file_path(skill_name: str, file_path: str) -> BaseResult: + @run_as_task() + async def read_skill_file_path(self, skill_name: str, file_path: str) -> BaseResult: skill_content, error = read_skill_file(skill_name, file_path) return BaseResult( result=[{ @@ -97,8 +96,8 @@ async def read_skill_file_path(skill_name: str, file_path: str) -> BaseResult: error=error ) - @staticmethod - async def list_skill_resources(skill_name: Optional[str]) -> BaseResult: + @run_as_task() + async def list_skill_resources(self, skill_name: Optional[str]) -> BaseResult: if not isinstance(skill_name, str) or not skill_name.strip(): return BaseResult( error="Missing required argument 'skill_name'. Please specify a non-empty skill name." @@ -120,8 +119,8 @@ async def list_skill_resources(skill_name: Optional[str]) -> BaseResult: has_more=False, ) - @staticmethod - async def read_skill_resource_uri(skill_uri: Optional[str]) -> BaseResult: + @run_as_task() + async def read_skill_resource_uri(self, skill_uri: Optional[str]) -> BaseResult: if not isinstance(skill_uri, str) or not skill_uri.strip(): return BaseResult( error="Missing required argument 'skill_resource_uri'. Please specify a non-empty skill URI." @@ -145,31 +144,77 @@ async def read_skill_resource_uri(skill_uri: Optional[str]) -> BaseResult: error=f"Invalid Skill URI: {skill_uri}" ) - @staticmethod - async def read_skill_resource_uri_list(skill_uri_list: Optional[List[str]]) -> BaseResult: + @run_as_task() + async def read_skill_resource_uri_list(self, skill_uri_list: Optional[List[str]]) -> BaseResult: if not isinstance(skill_uri_list, list) or not skill_uri_list: return BaseResult( error="Missing required argument 'skill_resource_uri_list'. Please provide a non-empty list of skill URIs." ) results = await asyncio.gather( - *(SkillsManager.read_skill_resource_uri(skill_uri) for skill_uri in skill_uri_list) + *(self.read_skill_resource_uri(skill_uri) for skill_uri in skill_uri_list) ) return BaseResult( result=results, total=len(results), ) - def register(mcp, runtime: AppRuntime): - @mcp.resource("blazemeter-skill-{skill_name}://{path}") - def universal_skills_handler(skill_name: str, path: str) -> str: - path = unquote(path) - content, error = read_skill_file(skill_name, path) - if error: - return error - return content + skills_tool = None - @mcp.tool( + async def _dispatch(action, args, token, ctx): + args = args or {} + skills_manager = SkillsManager(ctx) + match action: + case "list_skills": + return await skills_manager.list_skills() + case "read_skill": + return await skills_manager.read_skill(args.get("skill_name")) + case "list_skill_resources": + return await skills_manager.list_skill_resources(args.get("skill_name")) + case "read_skill_resource_uri": + return await skills_manager.read_skill_resource_uri(args.get("skill_resource_uri")) + case "read_skill_resource_uri_list": + return await skills_manager.read_skill_resource_uri_list(args.get("skill_resource_uri_list")) + case "batch": + batch_calls = args.get("batch_calls", []) + if not isinstance(batch_calls, list) or not batch_calls: + return BaseResult( + error="batch_calls must be a non-empty list of dicts with 'action' and 'args'") + + semaphore = asyncio.Semaphore(SkillsManager.MAX_BATCH_CONCURRENCY) + + async def process_call(call: Dict[str, Any]) -> BaseResult | List[BaseResult]: + sub_action = call.get("action", "") + sub_args = call.get("args", {}) + async with semaphore: + try: + # Recursively call the skills function itself + return await skills_tool({"action": sub_action, "args": sub_args}, ctx) + except httpx.HTTPStatusError: + return BaseResult( + error=f"HTTP error in sub-action {sub_action}: {format_sanitized_traceback()}" + ) + except Exception: + return BaseResult( + error=f"Error in sub-action {sub_action}: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}") + + # Parallel execution with asyncio.gather + results = await asyncio.gather(*[process_call(call) for call in batch_calls], + return_exceptions=True) + # Handle any exceptions returned + processed_results = [ + r if not isinstance(r, Exception) else BaseResult(error=f"Unhandled exception: {str(r)}") + for r in results + ] + return BaseResult(result=processed_results) + case _: + return BaseResult( + error=f"Action {action} not found in skills manager tool" + ) + + skills_tool = register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_skills", description=""" Operations to obtain Skills around BlazeMeter. @@ -196,75 +241,10 @@ def universal_skills_handler(skill_name: str, path: str) -> str: - Always generates the url attributes as a link in markdown format (like command_url). - **CRITICAL**: For multiple actions, always use the 'batch' action. - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + excluded_actions={'batch'}, + # Skills are documents (one content cell), not tabular result sets. + disable_materialization=True, + support_message=SUPPORT_MESSAGE, ) - async def skills( - action: str = Field(description="The action id to execute"), - args: Dict[str, Any] = Field(description="Dictionary with parameters", default=None), - ctx: Context = Field(description="Context object providing access to MCP capabilities") - ) -> BaseResult: - if args is None: - args = {} - - runtime.configure_context(ctx) - skills_manager = SkillsManager(ctx) - - async def _dispatch(): - match action: - case "list_skills": - return await skills_manager.list_skills() - case "read_skill": - return await skills_manager.read_skill(args.get("skill_name")) - case "list_skill_resources": - return await skills_manager.list_skill_resources(args.get("skill_name")) - case "read_skill_resource_uri": - return await skills_manager.read_skill_resource_uri(args.get("skill_resource_uri")) - case "read_skill_resource_uri_list": - return await skills_manager.read_skill_resource_uri_list(args.get("skill_resource_uri_list")) - case "batch": - batch_calls = args.get("batch_calls", []) - if not isinstance(batch_calls, list) or not batch_calls: - return BaseResult( - error="batch_calls must be a non-empty list of dicts with 'action' and 'args'") - - semaphore = asyncio.Semaphore(SkillsManager.MAX_BATCH_CONCURRENCY) - - async def process_call(call: Dict[str, Any]) -> BaseResult | List[BaseResult]: - sub_action = call.get("action", "") - sub_args = call.get("args", {}) - async with semaphore: - try: - # Recursively call the skills function itself - return await skills(sub_action, sub_args, ctx) - except httpx.HTTPStatusError: - return BaseResult( - error=f"HTTP error in sub-action {sub_action}: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"Error in sub-action {sub_action}: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}") - - # Parallel execution with asyncio.gather - results = await asyncio.gather(*[process_call(call) for call in batch_calls], - return_exceptions=True) - # Handle any exceptions returned - processed_results = [ - r if not isinstance(r, Exception) else BaseResult(error=f"Unhandled exception: {str(r)}") - for r in results - ] - return BaseResult(result=processed_results) - case _: - return BaseResult( - error=f"Action {action} not found in skills manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_skills", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}\n{SUPPORT_MESSAGE}" - ) diff --git a/tools/test_manager.py b/tools/test_manager.py index ccc2190..1cf448b 100644 --- a/tools/test_manager.py +++ b/tools/test_manager.py @@ -19,10 +19,9 @@ from typing import Any, Dict from typing import Optional, List -import httpx from mcp.server.fastmcp import Context -from config.blazemeter import TESTS_ENDPOINT, TOOLS_PREFIX +from config.blazemeter import TESTS_ENDPOINT, TOOLS_PREFIX, SUPPORT_MESSAGE from config.file_access import FileAccessPort from config.security import detect_sensitive_upload_path_reason from config.storage import HOSTED_FILE_ACCESS_MESSAGE, SessionScopeResolverPort @@ -37,12 +36,14 @@ from models.performance_test import PerformanceTestObject from models.result import BaseResult from tools import bridge, search_utils -from telemetry import run_tool +from tools.mcp_entrypoint import register_managed_tool from tools.utils import ( api_request, require_confirmation, Operations, format_sanitized_traceback, + run_as_task, + validate_required_args, ) logger = logging.getLogger(__name__) @@ -66,6 +67,7 @@ def __init__( def _current_scope(self): return self.scope_resolver.resolve(self.ctx, self.token) + @run_as_task() async def read(self, test_id: Optional[int]) -> BaseResult: if not isinstance(test_id, int) or test_id < 1: return BaseResult( @@ -91,6 +93,7 @@ async def read(self, test_id: Optional[int]) -> BaseResult: return test_result @require_confirmation(operation=Operations.CREATE) + @run_as_task() async def create( self, test_name: Optional[str], project_id: Optional[int] ) -> BaseResult: @@ -127,6 +130,7 @@ async def create( ) @require_confirmation(operation=Operations.DELETE) + @run_as_task() async def delete(self, test_id: Optional[int]) -> BaseResult: if not isinstance(test_id, int) or test_id < 1: return BaseResult( @@ -207,6 +211,7 @@ def _process_upload_results( successful_uploads.append({"file": valid_files[i], "result": result}) @require_confirmation(operation=Operations.CREATE) + @run_as_task() async def upload_assets( self, test_id: Optional[int], @@ -378,6 +383,7 @@ def _get_script_type(file_name: str) -> str: return script_types.get(extension, "unknown") + @run_as_task() async def list( self, project_id: Optional[int], @@ -415,6 +421,7 @@ async def list( params=parameters, ) + @run_as_task() async def search(self, args: dict[str, Any]) -> BaseResult: # Check if it's valid or allowed account_id = args.get("account_id") @@ -430,6 +437,7 @@ async def search(self, args: dict[str, Any]) -> BaseResult: "test-union", self.token, account_id, args ) + @run_as_task() async def search_filter_values( self, account_id: int, filter_names: List[str] ) -> BaseResult: @@ -498,6 +506,7 @@ def _normalize_configuration_override( return test_data_override @require_confirmation(operation=Operations.UPDATE) + @run_as_task() async def configure(self, performance_test: PerformanceTestObject) -> BaseResult: if not performance_test.is_valid(): raise ValueError("PerformanceTestObject must have a valid test_id") @@ -532,6 +541,7 @@ async def configure(self, performance_test: PerformanceTestObject) -> BaseResult ) @require_confirmation(operation=Operations.UPDATE) + @run_as_task() async def configure_failure_criteria(self, args: Dict[str, Any]) -> BaseResult: """Replace failure criteria for a test via PATCH configuration (preserves plugins.jmeter).""" test_id = args.get("test_id") @@ -562,13 +572,84 @@ async def configure_failure_criteria(self, args: Dict[str, Any]) -> BaseResult: json={"configuration": merged_configuration}, ) + @run_as_task() async def failure_criteria_meta(self, args: Dict[str, Any]) -> BaseResult: """Return the full KPI and condition catalog for building configure_failure_criteria rules (no API call).""" return BaseResult(result=[failure_criteria_meta_payload()]) - def register(mcp, runtime: AppRuntime): - @mcp.tool( + async def _dispatch(action, args, token, ctx): + if runtime.transport == "stdio": + test_manager = TestManager( + ctx, runtime.file_access, runtime.scope_resolver + ) + else: + test_manager = TestManager(ctx) + match action: + case "read": + return await test_manager.read(args.get("test_id")) + case "create": + return await test_manager.create( + args.get("test_name"), args.get("project_id") + ) + case "delete": + return await test_manager.delete(args.get("test_id")) + case "list": + return await test_manager.list( + args.get("project_id"), + args.get("limit", 50), + args.get("offset", 0), + ) + case "search": + return await test_manager.search(args) + case "search_filter_values": + return await test_manager.search_filter_values( + args.get("account_id"), args.get("filter_names", []) + ) + case "configure_load": + performance_test = PerformanceTestObject.from_args(args) + return await test_manager.configure(performance_test) + case "configure_locations": + performance_test = PerformanceTestObject.from_args(args) + return await test_manager.configure(performance_test) + case "upload_assets": + if validation_error := validate_required_args(action, args, ["test_id", "file_paths"]): + return validation_error + upload_result = await test_manager.upload_assets( + args.get("test_id"), + args.get("file_paths"), + args.get("main_script"), + ) + if isinstance(upload_result, BaseResult): + if upload_result.error: + return upload_result + inner = ( + upload_result.result[0] + if upload_result.result and len(upload_result.result) == 1 + else None + ) + if isinstance(inner, dict) and inner.get("error"): + return BaseResult(error=str(inner["error"])) + return upload_result + if isinstance(upload_result, dict) and upload_result.get("error"): + return BaseResult(error=upload_result["error"]) + return BaseResult(result=[upload_result]) + case "configure_failure_criteria": + if validation_error := validate_required_args( + action, args, ["test_id", "enabled", "rules"] + ): + return validation_error + return await test_manager.configure_failure_criteria(args) + case "failure_criteria_meta": + return await test_manager.failure_criteria_meta(args) + case _: + return BaseResult( + error=f"Action {action} not found in tests manager tool" + ) + + register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_tests", description=""" Operations on tests. @@ -669,68 +750,6 @@ def register(mcp, runtime: AppRuntime): - Before configure_failure_criteria, prefer failure_criteria_meta for kpi/condition codes and labels, then read if you must merge with existing rules. - For configure_failure_criteria, call read first and merge client-side if you must keep existing rules; providing rules replaces all criteria rows for that test. """, + dispatch=_dispatch, + support_message=SUPPORT_MESSAGE, ) - async def tests(action: str, args: Dict[str, Any], ctx: Context) -> BaseResult: - runtime.configure_context(ctx) - if runtime.transport == "stdio": - test_manager = TestManager( - ctx, runtime.file_access, runtime.scope_resolver - ) - else: - test_manager = TestManager(ctx) - - async def _dispatch(): - match action: - case "read": - return await test_manager.read(args.get("test_id")) - case "create": - return await test_manager.create( - args.get("test_name"), args.get("project_id") - ) - case "delete": - return await test_manager.delete(args.get("test_id")) - case "list": - return await test_manager.list( - args.get("project_id"), - args.get("limit", 50), - args.get("offset", 0), - ) - case "search": - return await test_manager.search(args) - case "search_filter_values": - return await test_manager.search_filter_values( - args.get("account_id"), args.get("filter_names", []) - ) - case "configure_load": - performance_test = PerformanceTestObject.from_args(args) - return await test_manager.configure(performance_test) - case "configure_locations": - performance_test = PerformanceTestObject.from_args(args) - return await test_manager.configure(performance_test) - case "upload_assets": - upload_result = await test_manager.upload_assets( - args.get("test_id"), - args.get("file_paths"), - args.get("main_script"), - ) - if isinstance(upload_result, dict) and upload_result.get("error"): - return BaseResult(error=upload_result["error"]) - return BaseResult(result=[upload_result]) - case "configure_failure_criteria": - return await test_manager.configure_failure_criteria(args) - case "failure_criteria_meta": - return await test_manager.failure_criteria_meta(args) - case _: - return BaseResult( - error=f"Action {action} not found in tests manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_tests", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult(error=f"Error: {format_sanitized_traceback()}") - except Exception: - return BaseResult( - error=f"""Error: {format_sanitized_traceback()} - If you think this is a bug, please contact BlazeMeter support or report issue at https://github.com/BlazeMeter/bzm-mcp/issues""" - ) diff --git a/tools/tools_manager.py b/tools/tools_manager.py new file mode 100644 index 0000000..765d233 --- /dev/null +++ b/tools/tools_manager.py @@ -0,0 +1,746 @@ +""" +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 re +import time +from typing import Any, Dict, Optional + +from mcp.server.fastmcp import Context + +from config.blazemeter import TOOLS_PREFIX +from config.runtime import AppRuntime +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.mcp_entrypoint import register_managed_tool +from tools.utils import ( + TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME, + run_as_task, + validate_non_empty_str_arg, + validate_required_args, +) + + +class ToolsManager(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(self.ctx, self.token) + + @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="blazemeter_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."], + ) + + +def register(mcp, runtime: AppRuntime): + async def _dispatch(action, args, token, ctx): + manager = ToolsManager(ctx, runtime.storage, runtime.scope_resolver) + args = args or {} + match action: + 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 manager.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 manager.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 manager.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 manager.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 manager.tasks_remove(str(args.get("task_id", "")).strip()) + case "dataframes_list": + return await manager.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 manager.dataframes_get(str(args.get("dataframe_id", "")).strip()) + case "dataframes_schema_groups": + return await manager.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 manager.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 manager.dataframes_remove(args.get("dataframe_id_list")) + case "dataframes_clear": + return await manager.dataframes_clear() + case "dataframes_sql_help": + return await manager.dataframes_sql_help() + case _: + return BaseResult(error=f"Action {action} not found in tools manager tool") + + register_managed_tool( + mcp, + runtime, + name=f"{TOOLS_PREFIX}_tools", + description=""" +Operations for session task and dataframe management (Storage-backed). +Actions: +- 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/queued task on this worker when a local handle exists. + Hosted note: cancel is process-local; other workers may only record cancel in Storage. + 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: +- **CRITICAL**: Always follow the action schema exactly. +- **CRITICAL**: Before writing any dataframe SQL query, call `dataframes_sql_help` first. +- 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. +""", + dispatch=_dispatch, + excluded_actions=set(TOOLS_ACTIONS_SKIP_AUTO_DATAFRAME), + support_message=None, + ) diff --git a/tools/user_manager.py b/tools/user_manager.py index d971559..9d77539 100644 --- a/tools/user_manager.py +++ b/tools/user_manager.py @@ -15,17 +15,15 @@ """ from typing import Any, Dict -import httpx from mcp.server.fastmcp import Context -from pydantic import Field -from config.blazemeter import TOOLS_PREFIX, USER_ENDPOINT +from config.blazemeter import TOOLS_PREFIX, USER_ENDPOINT, SUPPORT_MESSAGE from config.runtime import AppRuntime 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 +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import api_request, run_as_task class UserManager(Manager): @@ -36,6 +34,7 @@ def __init__( ): super().__init__(ctx) + @run_as_task() async def read(self) -> BaseResult: return await api_request( self.token, @@ -46,7 +45,19 @@ async def read(self) -> BaseResult: def register(mcp, runtime: AppRuntime): - @mcp.tool( + async def _dispatch(action, args, token, ctx): + user_manager = UserManager(ctx) + match action: + case "read": + return await user_manager.read() + case _: + return BaseResult( + error=f"Action {action} not found in user manager tool" + ) + + register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_user", description=""" Operations on user information. @@ -54,35 +65,9 @@ def register(mcp, runtime: AppRuntime): - read: Read a current user information from BlazeMeter. Hints: - For default account, workspace and project, use the 'read' action. +- Optional result formatting in args: `result_format` = `auto` (default), `dataframe` (force dataframe), `raw` (disable dataframe materialization). - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + support_message=SUPPORT_MESSAGE, ) - async def user( - action: str = Field(description="The action id to execute"), - args: Dict[str, Any] = Field(description="Dictionary with parameters"), - ctx: Context = Field(description="Context object providing access to MCP capabilities") - ) -> BaseResult: - - runtime.configure_context(ctx) - user_manager = UserManager(ctx) - - async def _dispatch(): - match action: - case "read": - return await user_manager.read() - case _: - return BaseResult( - error=f"Action {action} not found in user manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_user", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"""Error: {format_sanitized_traceback()} - If you think this is a bug, please contact BlazeMeter support or report issue at https://github.com/BlazeMeter/bzm-mcp/issues""" - ) diff --git a/tools/utils.py b/tools/utils.py index 049b8d7..6aa6ff1 100644 --- a/tools/utils.py +++ b/tools/utils.py @@ -16,27 +16,45 @@ """ Simple utilities for BlazeMeter MCP tools. """ +import asyncio +import contextvars import functools +import inspect import os import platform import re +import secrets import sys +import time import traceback from datetime import datetime, timezone from enum import Enum -from typing import Any, Optional, Callable, Awaitable +from typing import Any, Dict, Optional, Callable, Awaitable, Tuple from importlib import resources from pathlib import Path import httpx +from mcp.types import CallToolResult from pydantic import BaseModel from config.blazemeter import BZM_API_BASE_URL -from config.context_resolution import resolve_ctx_user_config +from config.context_resolution import resolve_ctx_token, resolve_ctx_user_config from config.security import validate_http_request_endpoint from config.token import BzmToken from config.version import __version__ -from models.result import BaseResult, HttpBaseResult +from models.result import BaseResult, HttpBaseResult, ToolResult + +SIMPLE_ID_ALPHABET = "0123456789abcdefghjkmnpqrstvwxyz" +SIMPLE_ID_LENGTH = 8 + + +def generate_simple_id() -> str: + return "".join(secrets.choice(SIMPLE_ID_ALPHABET) for _ in range(SIMPLE_ID_LENGTH)) + + +def normalize_simple_id(simple_id: str) -> str: + return str(simple_id).strip().lower() + so = platform.system() # "Windows", "Linux", "Darwin" version = platform.version() # kernel / build version @@ -130,6 +148,15 @@ class ConfirmMode(Enum): DISABLE = "NONE" # No confirmation +_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 +) +_tool_result_depth = contextvars.ContextVar("tool_result_depth", default=0) +_result_debug_enabled = False + + class Operations(Enum): CREATE = "C" # Create READ = "R" # Read @@ -137,6 +164,406 @@ class Operations(Enum): DELETE = "D" # Delete +# MCP tool actions that must stay inline under result_format=auto (no auto-dataframe). +# Shared by @tool_result(excluded_actions=...) and the async task runner so exclusions +# are honored even when @run_as_task materializes before the entrypoint finalize. +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", +}) + + +def set_result_debug_enabled(enabled: bool): + global _result_debug_enabled + _result_debug_enabled = bool(enabled) + + +def is_result_debug_enabled() -> bool: + return _result_debug_enabled + + +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 normalize_action_args(arguments: Optional[Dict[str, Any]] = None) -> tuple[str, Dict[str, Any]]: + """ + Normalize tool arguments to (action, args) format. + Supports: + - {"action": "x", "args": {"key": "value"}} + - {"action": "x", "key": "value"} (params at top level, merged into args) + - {"arguments": {"action": "x", "args": {...}}} (double-wrapped by client) + Top-level keys other than 'action' and 'args' are merged into args. + Use a single 'arguments' param so the full MCP tool call payload is received + (avoids Pydantic dropping extra fields when using action/args separately). + """ + arguments = arguments or {} + # Unwrap double-nested format: {"arguments": {"action": "x", "args": {...}}} + inner = arguments.get("arguments") + if ( + isinstance(inner, dict) + and len(arguments) == 1 + and ("action" in inner or "args" in inner) + ): + arguments = inner + action = str(arguments.get("action") or "").strip() or "" + args = dict(arguments.get("args") or {}) + for key, value in arguments.items(): + if key not in ("action", "args"): + args[key] = value + return action, args + + +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 _resolve_tool_token(ctx: Any) -> Optional[BzmToken]: + if ctx is None: + return None + return resolve_ctx_token(ctx) + + +def _resolve_invocation( + args: tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Tuple[str, Dict[str, Any], Any]: + """Resolve (action, tool_args, ctx) from arguments= or legacy action/args shapes.""" + ctx = kwargs.get("ctx") + arguments = kwargs.get("arguments") + + if arguments is None and args: + if isinstance(args[0], dict): + arguments = args[0] + if ctx is None and len(args) >= 2: + ctx = args[1] + elif isinstance(args[0], str): + action = args[0] + tool_args = args[1] if len(args) > 1 else (kwargs.get("args") or {}) + if ctx is None and len(args) >= 3: + ctx = args[2] + if not isinstance(tool_args, dict): + tool_args = {} + return action, tool_args, ctx + + if isinstance(arguments, dict): + action, tool_args = normalize_action_args(arguments) + return action, tool_args, ctx + + action = kwargs.get("action") or "" + tool_args = kwargs.get("args") or {} + if not isinstance(tool_args, dict): + tool_args = {} + return str(action), tool_args, ctx + + +def _set_tool_call_timing( + result: BaseResult, + started_monotonic: float, + started_wall_clock: float, + extra_timing: Optional[Dict[str, int]] = None, +): + finished_wall_clock = time.time() + duration_ms = int((time.monotonic() - started_monotonic) * 1000) + result.tool_call_started_at = datetime.fromtimestamp(started_wall_clock, tz=timezone.utc).isoformat() + result.tool_call_finished_at = datetime.fromtimestamp(finished_wall_clock, tz=timezone.utc).isoformat() + result.tool_call_duration_ms = duration_ms + if not _result_debug_enabled: + return + debug = result.debug if isinstance(result.debug, dict) else {} + timing = {"total_ms": duration_ms} + if extra_timing: + timing.update({k: int(v) for k, v in extra_timing.items()}) + debug["timing"] = timing + result.debug = debug + + +def tool_result( + excluded_actions: Optional[set[str]] = None, + *, + disable_materialization: bool = False, +): + """ + MCP entrypoint wrapper: set result_format context, attach timing, and return ToolResult. + + Nested calls (e.g. help/skills batch sub-actions) return BaseResult to avoid wrapping. + Materialization is owned by run_tool_with_runtime / the async task runner unless + ``disable_materialization`` is False (legacy/direct finalize path). + """ + excluded = excluded_actions or set() + + def decorator(func: Callable[..., Awaitable[Any]]): + @functools.wraps(func) + async def wrapper(*args, **kwargs) -> ToolResult | CallToolResult | BaseResult: + depth = _tool_result_depth.get() + depth_token = _tool_result_depth.set(depth + 1) + action, tool_args, ctx = _resolve_invocation(args, kwargs) + + result_format = "auto" + if isinstance(tool_args, dict) and "result_format" in tool_args: + raw_format = str(tool_args.get("result_format", "auto")).strip().lower() + if raw_format in {"auto", "dataframe", "raw"}: + result_format = raw_format + else: + result_format = "invalid" + if isinstance(action, str) and action == "batch": + result_format = "raw" + + format_token = _result_format_context.set( + result_format if result_format != "invalid" else "auto" + ) + started_monotonic = time.monotonic() + started_wall_clock = time.time() + try: + if result_format == "invalid": + result: Any = BaseResult( + error="Invalid result_format value. Allowed values: auto, dataframe, raw." + ) + after_func_monotonic = started_monotonic + else: + result = await func(*args, **kwargs) + after_func_monotonic = time.monotonic() + + postprocess_ms = 0 + if ( + not disable_materialization + and isinstance(result, BaseResult) + and not result.error + and result.result is not None + ): + from tools.dataframe_manager import finalize_tool_result + + post_started = time.monotonic() + result = await finalize_tool_result( + result, + action=action, + args=tool_args, + origin_manager=func.__name__, + token=_resolve_tool_token(ctx), + ctx=ctx, + excluded_actions=excluded, + ) + postprocess_ms = int((time.monotonic() - post_started) * 1000) + + if isinstance(result, BaseResult): + _set_tool_call_timing( + result, + started_monotonic, + started_wall_clock, + extra_timing={ + "manager_logic_ms": int((after_func_monotonic - started_monotonic) * 1000), + "postprocess_ms": postprocess_ms, + }, + ) + + if depth > 0: + return result + if isinstance(result, (ToolResult, CallToolResult)): + return result + if isinstance(result, BaseResult): + return ToolResult.from_base_result(result) + return ToolResult.from_base_result(BaseResult(result=[result])) + finally: + _result_format_context.reset(format_token) + _tool_result_depth.reset(depth_token) + + return wrapper + + return decorator + + +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 ( + DEFAULT_SCOPE, + submit_task, + get_task_record, + remove_task, + task_snapshot, + ) + + resolved_scope = scope if isinstance(scope, SessionScope) else DEFAULT_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 + _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) + 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 blazemeter_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) + 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 tools.async_task_manager import session_scope_from_manager + + scope = session_scope_from_manager(self) + 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 + + async def api_request(token: Optional[BzmToken], method: str, endpoint: str, result_formatter: Callable = None, result_formatter_params: Optional[dict] = None, diff --git a/tools/workspace_manager.py b/tools/workspace_manager.py index 84da19a..3ec4f48 100644 --- a/tools/workspace_manager.py +++ b/tools/workspace_manager.py @@ -15,18 +15,16 @@ """ from typing import Any, Dict, Optional -import httpx from mcp.server.fastmcp import Context -from pydantic import Field -from config.blazemeter import WORKSPACES_ENDPOINT, TOOLS_PREFIX +from config.blazemeter import WORKSPACES_ENDPOINT, TOOLS_PREFIX, SUPPORT_MESSAGE from config.runtime import AppRuntime from formatters.workspace import format_workspaces, format_workspaces_detailed, format_workspaces_locations from models.manager import Manager from models.result import BaseResult from tools import bridge -from telemetry import run_tool -from tools.utils import api_request, format_sanitized_traceback +from tools.mcp_entrypoint import register_managed_tool +from tools.utils import api_request, run_as_task class WorkspaceManager(Manager): @@ -41,6 +39,7 @@ def __init__( ): super().__init__(ctx) + @run_as_task() async def read(self, workspace_id: Optional[int]) -> BaseResult: if not isinstance(workspace_id, int) or workspace_id < 1: return BaseResult(error="Missing or invalid required argument 'workspace_id'. Expected integer.") @@ -62,6 +61,7 @@ async def read(self, workspace_id: Optional[int]) -> BaseResult: else: return workspace_result + @run_as_task() async def list(self, account_id: Optional[int], limit: int = 50, offset: int = 0) -> BaseResult: if not isinstance(account_id, int) or account_id < 1: return BaseResult(error="Missing or invalid required argument 'account_id'. Expected integer.") @@ -88,6 +88,7 @@ async def list(self, account_id: Optional[int], limit: int = 50, offset: int = 0 params=parameters ) + @run_as_task() async def read_locations(self, workspace_id: Optional[int], purpose: str = "load") -> BaseResult: if not isinstance(workspace_id, int) or workspace_id < 1: return BaseResult(error="Missing or invalid required argument 'workspace_id'. Expected integer.") @@ -111,9 +112,27 @@ async def read_locations(self, workspace_id: Optional[int], purpose: str = "load return account_result else: return locations_result - def register(mcp, runtime: AppRuntime): - @mcp.tool( + + async def _dispatch(action, args, token, ctx): + workspace_manager = WorkspaceManager(ctx) + match action: + case "read": + return await workspace_manager.read(args.get("workspace_id")) + case "list": + return await workspace_manager.list( + args.get("account_id"), args.get("limit", 50), args.get("offset", 0) + ) + case "read_locations": + return await workspace_manager.read_locations(args.get("workspace_id"), args.get("purpose", "load")) + case _: + return BaseResult( + error=f"Action {action} not found in workspace manager tool" + ) + + register_managed_tool( + mcp, + runtime, name=f"{TOOLS_PREFIX}_workspaces", description=""" Operations on workspaces. @@ -133,40 +152,7 @@ def register(mcp, runtime: AppRuntime): Hints: - For available locations and available billing usage use the 'read' action for a particular workspace. - **CRITICAL**: Always follow the action schema exactly. If args are required, include args with exact names/types. -""" +""", + dispatch=_dispatch, + support_message=SUPPORT_MESSAGE, ) - async def workspace( - action: str = Field(description="The action id to execute"), - args: Dict[str, Any] = Field(description="Dictionary with parameters"), - ctx: Context = Field(description="Context object providing access to MCP capabilities") - ) -> BaseResult: - - runtime.configure_context(ctx) - workspace_manager = WorkspaceManager(ctx) - - async def _dispatch(): - match action: - case "read": - return await workspace_manager.read(args.get("workspace_id")) - case "list": - return await workspace_manager.list( - args.get("account_id"), args.get("limit", 50), args.get("offset", 0) - ) - case "read_locations": - return await workspace_manager.read_locations(args.get("workspace_id"), args.get("purpose", "load")) - case _: - return BaseResult( - error=f"Action {action} not found in workspace manager tool" - ) - - try: - return await run_tool(f"{TOOLS_PREFIX}_workspaces", action, ctx, _dispatch) - except httpx.HTTPStatusError: - return BaseResult( - error=f"Error: {format_sanitized_traceback()}" - ) - except Exception: - return BaseResult( - error=f"""Error: {format_sanitized_traceback()} - If you think this is a bug, please contact BlazeMeter support or report issue at https://github.com/BlazeMeter/bzm-mcp/issues""" - ) diff --git a/uv.lock b/uv.lock index 26a3439..bbc3ef6 100644 --- a/uv.lock +++ b/uv.lock @@ -67,6 +67,7 @@ dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-sdk" }, + { name = "polars" }, { name = "pydantic" }, { name = "pydantic-core" }, { name = "pydantic-settings" }, @@ -88,6 +89,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" }, @@ -842,6 +844,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 = "1.43.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "polars-runtime-32" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/89/13/3873f213304bcbaaf39e63c8b905ceb460a0524448d57f86a829f6d4d0fd/polars-1.43.2.tar.gz", hash = "sha256:c699671b99eb71ff53334d237917aaa3db5ad4dda480abcb6c80e0eaee7b677b", size = 750312, upload-time = "2026-08-01T06:28:30.872Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/fe/0888040a24e4504098b85d8ad486b14cb01cf6b030bbe479dfc2dcffc2ac/polars-1.43.2-py3-none-any.whl", hash = "sha256:22aa0cb92a1ee2d60d6a15a638b2e8e0dd99aea21ac0cd8fb29da8e382e075a9", size = 847150, upload-time = "2026-08-01T06:27:15.543Z" }, +] + +[[package]] +name = "polars-runtime-32" +version = "1.43.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/06/11b578eeef05f867e3ee31b2a2fdd8e7684c2aa47822c49935d1be789c38/polars_runtime_32-1.43.2.tar.gz", hash = "sha256:d7b7c486bccee75a6af0158b87077da3d054657e3c60036b28644f4e1c7fdbf7", size = 3095669, upload-time = "2026-08-01T06:28:32.315Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/fc/12e6d4ca34d820297651134cfa35f86c33e898539fc6629cbb35d0089697/polars_runtime_32-1.43.2-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:91abf205d4ec93f92ba95386b7f8776559ae3dfce425ed2e527efa75d117d04a", size = 53088908, upload-time = "2026-08-01T06:27:18.268Z" }, + { url = "https://files.pythonhosted.org/packages/32/81/833b0853551deb810854f96b43dea342b6e6c9b0ea1afcccf774157d519d/polars_runtime_32-1.43.2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:2cc3ff96fd44789b02eb5c15b98dfcb000101636b177d3034fae2feec19b118f", size = 47540529, upload-time = "2026-08-01T06:27:21.391Z" }, + { url = "https://files.pythonhosted.org/packages/83/55/7b2a75af14c9294d97f3bec132dd3018ddcd988bef32b5d28322150b8c11/polars_runtime_32-1.43.2-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:10ed36e615ab362feb7406e6d084e124b445ad284caa73bd93ae7e65745ed894", size = 51366340, upload-time = "2026-08-01T06:27:24.776Z" }, + { url = "https://files.pythonhosted.org/packages/62/60/64deacb3abc70c52e2d88a808a052d1621c86a48fe9194f2c065579ab1cd/polars_runtime_32-1.43.2-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6d5a7ae004a2723ebf4427f6d6a639f30f86af4cf077075f6b35d04711154fc3", size = 57304599, upload-time = "2026-08-01T06:27:27.875Z" }, + { url = "https://files.pythonhosted.org/packages/52/95/d6e3a236d7630e17c40d0ddee839bf2be9acf548fdc0e5ad65ed9ff0cac6/polars_runtime_32-1.43.2-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:09339eacc6d392206e78aabbaaa37d7276eb969b798f46cb1f367fd718798c60", size = 51520580, upload-time = "2026-08-01T06:27:30.913Z" }, + { url = "https://files.pythonhosted.org/packages/b6/5a/2deb8eac70e9a2ac26d88a66ae7cf52612865026f4f4a5e7ab11ad9d52bf/polars_runtime_32-1.43.2-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:452b400e59e7f56e4c6437f435e796903272a9388feee12de2bea049ae87025e", size = 55204471, upload-time = "2026-08-01T06:27:33.886Z" }, + { url = "https://files.pythonhosted.org/packages/29/9e/647401ae8a607bc0cc40ed7b8592d5b1be90ded0dc9b9d6d3aeb03f9524b/polars_runtime_32-1.43.2-cp310-abi3-win_amd64.whl", hash = "sha256:00e33c28e321410c8d66e814a90043101e3bdd9ed2c6dabda07565aa8adbbdf1", size = 52572176, upload-time = "2026-08-01T06:27:37.048Z" }, + { url = "https://files.pythonhosted.org/packages/96/8d/60a50c3f36c85218a7ffcb48c6fe2ce1f7bec799152d68b8658ebed2179c/polars_runtime_32-1.43.2-cp310-abi3-win_arm64.whl", hash = "sha256:350a4868cae85bf8b3f81b33ba47927c15256bd9264dfc8c0753f1b927eac9d3", size = 46582513, upload-time = "2026-08-01T06:27:40.025Z" }, +] + [[package]] name = "protobuf" version = "6.33.6"