From edc7b77a428069538e5153d46cef8ca36890fa19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 24 Jun 2026 13:36:21 +0200 Subject: [PATCH 01/89] feat(PSGO-261): exchange programmatic tokens for Storage tokens via auth-bridge Implements Part A of the PAT-support RFC: accept Keboola programmatic bearer tokens (kbc_at_* / kbc_pat_*) and exchange them at Connection's auth-bridge resolver for a legacy Storage token, used downstream unchanged. Legacy X-StorageAPI-Token traffic is untouched. - clients/auth_bridge.py: is_programmatic_token() + StorageTokenResolver. Authenticates to POST /manage/internal/auth-bridge/resolve-storage-token with the projected SA JWT (X-Kubernetes-Authorization, read per request), the user token as X-Subject-Token, and projectId in the body. Maps resolver 400/401/403 through; 5xx/timeout/network -> 502. No token material logged. - config.py: project_id field (KBC_PROJECT_ID env / X-KBC-ProjectId header). - mcp.py: SessionStateMiddleware exchanges a programmatic token before building KeboolaClient; KBC_KUBERNETES_TOKEN_PATH is process-env only. - Unit tests for detection, exchange, error mapping, no-token-leak, and the middleware wiring. PKCE login (Part B) and the OAuth->PAT exchange are follow-up PRs. Linear: https://linear.app/keboola/issue/PSGO-261 Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/clients/auth_bridge.py | 140 ++++++++++++++++++ src/keboola_mcp_server/config.py | 6 + src/keboola_mcp_server/mcp.py | 45 +++++- tests/clients/test_auth_bridge.py | 123 +++++++++++++++ tests/test_config.py | 2 +- tests/test_mcp.py | 43 ++++++ 6 files changed, 356 insertions(+), 3 deletions(-) create mode 100644 src/keboola_mcp_server/clients/auth_bridge.py create mode 100644 tests/clients/test_auth_bridge.py diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py new file mode 100644 index 000000000..e6491b59c --- /dev/null +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -0,0 +1,140 @@ +"""Exchange Keboola programmatic tokens for legacy Storage tokens (PSGO-261). + +Implements the decentralized auth-bridge exchange: a programmatic bearer token +(`kbc_at_*` access token or `kbc_pat_*` personal access token) presented to the MCP +server is exchanged at Connection for a legacy Storage token, which is then used for +all downstream Storage-token APIs exactly as before. + +The MCP server authenticates to the resolver with its own projected Kubernetes +ServiceAccount JWT (`X-Kubernetes-Authorization`); the user's token travels as +`X-Subject-Token`. The SA token file is read per call so kubelet rotation is honored. +No token material is ever logged or placed in exception messages. +""" + +import logging +from http import HTTPStatus +from pathlib import Path +from typing import cast +from urllib.parse import urlparse, urlunparse + +import httpx + +LOG = logging.getLogger(__name__) + +_ACCESS_TOKEN_PREFIX = 'kbc_at_' +_PAT_PREFIX = 'kbc_pat_' +_RESOLVE_ENDPOINT = 'manage/internal/auth-bridge/resolve-storage-token' +# Resolver statuses passed through to the client verbatim; anything else (incl. 5xx, +# timeouts, network failures) is mapped to 502 Bad Gateway. +_PASS_THROUGH_STATUSES = frozenset( + {int(HTTPStatus.BAD_REQUEST), int(HTTPStatus.UNAUTHORIZED), int(HTTPStatus.FORBIDDEN)} +) + + +def _strip_bearer(token: str) -> str: + """Removes a leading case-insensitive ``Bearer `` scheme from a token, if present.""" + if token[:7].lower() == 'bearer ': + return token[7:].strip() + return token + + +def is_programmatic_token(token: str | None) -> bool: + """True if ``token`` is a Keboola programmatic bearer token (``kbc_at_`` / ``kbc_pat_``).""" + if not token: + return False + bare = _strip_bearer(token) + return bare.startswith(_ACCESS_TOKEN_PREFIX) or bare.startswith(_PAT_PREFIX) + + +class StorageTokenExchangeError(RuntimeError): + """Raised when the auth-bridge resolver fails to exchange a programmatic token. + + :ivar status_code: The client-facing HTTP status (resolver 400/401/403 pass through; + 5xx/timeout/network map to 502). + """ + + def __init__(self, message: str, status_code: int) -> None: + super().__init__(message, status_code) + self.status_code = status_code + + def __str__(self) -> str: + return self.args[0] + + +class StorageTokenResolver: + """Exchanges a programmatic token for a legacy Storage token via the Connection resolver.""" + + def __init__( + self, + *, + storage_api_url: str, + kubernetes_token_path: str, + timeout: httpx.Timeout | None = None, + transport: httpx.AsyncBaseTransport | None = None, + ) -> None: + """ + :param storage_api_url: Connection Storage API URL (``https://connection.``). + :param kubernetes_token_path: Path to the projected ServiceAccount token file. + :param timeout: Optional HTTP timeout override. + :param transport: Optional httpx transport (for testing). + """ + parsed = urlparse(storage_api_url) + if not parsed.hostname or not parsed.hostname.startswith('connection.'): + raise ValueError(f'Invalid Keboola Storage API URL: {storage_api_url}') + self._base_url = urlunparse(('https', parsed.hostname, '', '', '', '')) + self._kubernetes_token_path = kubernetes_token_path + self._timeout = timeout or httpx.Timeout(connect=5.0, read=30.0, write=10.0, pool=5.0) + self._transport = transport + + def _read_sa_jwt(self) -> str: + # Read per call — the kubelet rotates the projected token in place. + jwt = Path(self._kubernetes_token_path).read_text().strip() + if not jwt: + raise ValueError(f'Kubernetes ServiceAccount token file is empty: {self._kubernetes_token_path}') + return jwt + + async def resolve(self, *, subject_token: str, project_id: int) -> str: + """ + Exchanges ``subject_token`` for the legacy Storage token of ``project_id``. + + :return: The legacy Storage token. + :raises StorageTokenExchangeError: On any resolver failure (status carried on the error). + """ + headers = { + 'Content-Type': 'application/json', + 'Accept': 'application/json', + 'X-Kubernetes-Authorization': f'Bearer {self._read_sa_jwt()}', + 'X-Subject-Token': f'Bearer {_strip_bearer(subject_token)}', + } + try: + async with httpx.AsyncClient(timeout=self._timeout, transport=self._transport) as client: + response = await client.post( + f'{self._base_url}/{_RESOLVE_ENDPOINT}', + headers=headers, + json={'projectId': project_id}, + ) + except httpx.HTTPError as e: + # Network / timeout failure. Raise without chaining so no request (and thus no + # token material) can surface in a traceback. + raise StorageTokenExchangeError( + f'Auth-bridge token exchange could not reach Connection ({type(e).__name__}).', + status_code=int(HTTPStatus.BAD_GATEWAY), + ) from None + + if response.status_code != HTTPStatus.OK: + status = response.status_code + mapped = status if status in _PASS_THROUGH_STATUSES else int(HTTPStatus.BAD_GATEWAY) + LOG.error(f'Auth-bridge token exchange failed: resolver status {status}, mapped to {mapped}.') + raise StorageTokenExchangeError( + f'Auth-bridge token exchange was rejected (resolver status {status}).', + status_code=mapped, + ) + + body = cast(dict, response.json()) + storage_token = body.get('storageToken') + if not storage_token: + raise StorageTokenExchangeError( + 'Auth-bridge token exchange returned no storageToken.', + status_code=int(HTTPStatus.BAD_GATEWAY), + ) + return cast(str, storage_token) diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 6168aeb55..22768403b 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -43,6 +43,12 @@ class Config: """The access-token issued by Keboola OAuth server to be sent in 'Authorization: Bearer ' header.""" conversation_id: str | None = None """The ID of the ongoing conversation with the MCP server. This is supplied only by the HTTP header.""" + project_id: Optional[str] = field(default=None, metadata={'aliases': ['kbc_project_id']}) + """Project id used to scope a programmatic-token (kbc_at_/kbc_pat_) exchange. + + Maps the `X-KBC-ProjectId` HTTP header (via the alias) and the `KBC_PROJECT_ID` env var. + Only consulted when the inbound Storage token is a Keboola programmatic token; the legacy + project-bound Storage token derives its project from the token itself.""" def __post_init__(self) -> None: for f in dataclasses.fields(self): diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index c419e0fab..c8e5b0a66 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -29,6 +29,7 @@ from starlette.requests import Request from starlette.types import ASGIApp, Receive, Scope, Send +from keboola_mcp_server.clients.auth_bridge import StorageTokenResolver, is_programmatic_token from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config, ServerRuntimeInfo, is_same_stack @@ -311,6 +312,38 @@ def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_sto return config + @classmethod + async def _exchange_programmatic_token(cls, config: Config) -> str: + """ + Exchanges a programmatic token (kbc_at_/kbc_pat_) for the project's legacy Storage token. + + The resolver is reached only on the deployed MCP server, which has a projected + ServiceAccount token at ``KBC_KUBERNETES_TOKEN_PATH`` (read from the process + environment only, never from per-request config). A project id is required because + a programmatic token is not project-bound. + """ + kubernetes_token_path = os.environ.get('KBC_KUBERNETES_TOKEN_PATH') + if not kubernetes_token_path: + raise ValueError( + 'Received a Keboola programmatic token (kbc_at_/kbc_pat_) but KBC_KUBERNETES_TOKEN_PATH ' + 'is not configured. Programmatic-token exchange is available only on the deployed MCP server.' + ) + if not config.project_id: + raise ValueError( + 'A project id is required to exchange a programmatic token. ' + 'Set the KBC_PROJECT_ID env var or the X-KBC-ProjectId header.' + ) + try: + project_id = int(config.project_id) + except (TypeError, ValueError): + raise ValueError(f'Invalid project id for programmatic-token exchange: {config.project_id!r}') + + resolver = StorageTokenResolver( + storage_api_url=config.storage_api_url, + kubernetes_token_path=kubernetes_token_path, + ) + return await resolver.resolve(subject_token=config.storage_token, project_id=project_id) + @classmethod async def create_session_state( cls, @@ -343,10 +376,18 @@ async def create_session_state( if not config.storage_api_url: raise ValueError('Storage API URL is not provided.') + storage_token = config.storage_token + bearer_token = config.bearer_token + if is_programmatic_token(storage_token): + # A Keboola programmatic token (kbc_at_/kbc_pat_) is not a Storage token; exchange + # it for the project's legacy Storage token and use that downstream unchanged. + storage_token = await cls._exchange_programmatic_token(config) + bearer_token = None + client = await KeboolaClient( storage_api_url=config.storage_api_url, - storage_api_token=config.storage_token, - bearer_token=config.bearer_token, + storage_api_token=storage_token, + bearer_token=bearer_token, headers=cls._get_headers(runtime_info), readonly=readonly, own_stack_storage_api_url=own_stack_storage_api_url, diff --git a/tests/clients/test_auth_bridge.py b/tests/clients/test_auth_bridge.py new file mode 100644 index 000000000..001b37cd6 --- /dev/null +++ b/tests/clients/test_auth_bridge.py @@ -0,0 +1,123 @@ +"""Tests for the auth-bridge programmatic-token exchange (PSGO-261).""" + +from http import HTTPStatus +from pathlib import Path + +import httpx +import pytest + +from keboola_mcp_server.clients.auth_bridge import ( + StorageTokenExchangeError, + StorageTokenResolver, + is_programmatic_token, +) + +STORAGE_API_URL = 'https://connection.keboola.com' + + +@pytest.mark.parametrize( + ('token', 'expected'), + [ + ('kbc_at_019ef801_abc', True), + ('kbc_pat_019ef801_abc', True), + ('Bearer kbc_at_019ef801_abc', True), + ('bearer kbc_pat_019ef801_abc', True), + ('123-legacy-storage-token', False), + ('kbc_rt_019ef801_abc', False), # refresh token is not a Storage-token bearer + ('', False), + (None, False), + ], +) +def test_is_programmatic_token(token: str | None, expected: bool) -> None: + assert is_programmatic_token(token) is expected + + +@pytest.fixture +def sa_token_file(tmp_path: Path) -> Path: + path = tmp_path / 'sa-token' + path.write_text(' sa-jwt-value\n') # surrounding whitespace must be stripped + return path + + +def _resolver(sa_token_file: Path, handler) -> StorageTokenResolver: + return StorageTokenResolver( + storage_api_url=STORAGE_API_URL, + kubernetes_token_path=str(sa_token_file), + transport=httpx.MockTransport(handler), + ) + + +@pytest.mark.asyncio +async def test_resolve_success_sends_expected_request(sa_token_file: Path) -> None: + captured: dict[str, httpx.Request] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured['request'] = request + return httpx.Response(HTTPStatus.OK, json={'storageToken': 'legacy-token', 'projectId': 42}) + + resolver = _resolver(sa_token_file, handler) + token = await resolver.resolve(subject_token='Bearer kbc_pat_abc', project_id=42) + + assert token == 'legacy-token' + rq = captured['request'] + assert rq.url.path == '/manage/internal/auth-bridge/resolve-storage-token' + assert rq.headers['X-Kubernetes-Authorization'] == 'Bearer sa-jwt-value' + # Subject token is normalized to a single Bearer scheme regardless of inbound form. + assert rq.headers['X-Subject-Token'] == 'Bearer kbc_pat_abc' + + +@pytest.mark.parametrize('status', [HTTPStatus.BAD_REQUEST, HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN]) +@pytest.mark.asyncio +async def test_resolve_passes_through_client_errors(sa_token_file: Path, status: HTTPStatus) -> None: + resolver = _resolver(sa_token_file, lambda rq: httpx.Response(status, json={'error': 'nope'})) + with pytest.raises(StorageTokenExchangeError) as exc: + await resolver.resolve(subject_token='kbc_at_abc', project_id=1) + assert exc.value.status_code == int(status) + + +@pytest.mark.parametrize('status', [HTTPStatus.INTERNAL_SERVER_ERROR, HTTPStatus.BAD_GATEWAY, HTTPStatus.NOT_FOUND]) +@pytest.mark.asyncio +async def test_resolve_maps_other_statuses_to_502(sa_token_file: Path, status: HTTPStatus) -> None: + resolver = _resolver(sa_token_file, lambda rq: httpx.Response(status)) + with pytest.raises(StorageTokenExchangeError) as exc: + await resolver.resolve(subject_token='kbc_at_abc', project_id=1) + assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) + + +@pytest.mark.asyncio +async def test_resolve_maps_network_error_to_502(sa_token_file: Path) -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError('boom', request=request) + + resolver = _resolver(sa_token_file, handler) + with pytest.raises(StorageTokenExchangeError) as exc: + await resolver.resolve(subject_token='kbc_at_abc', project_id=1) + assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) + # No token material leaks into the message. + assert 'kbc_at_abc' not in str(exc.value) + + +@pytest.mark.asyncio +async def test_resolve_missing_storage_token_maps_to_502(sa_token_file: Path) -> None: + resolver = _resolver(sa_token_file, lambda rq: httpx.Response(HTTPStatus.OK, json={'projectId': 1})) + with pytest.raises(StorageTokenExchangeError) as exc: + await resolver.resolve(subject_token='kbc_at_abc', project_id=1) + assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) + + +@pytest.mark.asyncio +async def test_resolve_empty_sa_token_file_fails_loudly(tmp_path: Path) -> None: + empty = tmp_path / 'empty' + empty.write_text(' ') + resolver = StorageTokenResolver( + storage_api_url=STORAGE_API_URL, + kubernetes_token_path=str(empty), + transport=httpx.MockTransport(lambda rq: httpx.Response(HTTPStatus.OK, json={'storageToken': 'x'})), + ) + with pytest.raises(ValueError, match='empty'): + await resolver.resolve(subject_token='kbc_at_abc', project_id=1) + + +def test_invalid_storage_api_url_rejected(sa_token_file: Path) -> None: + with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): + StorageTokenResolver(storage_api_url='https://example.com', kubernetes_token_path=str(sa_token_file)) diff --git a/tests/test_config.py b/tests/test_config.py index 22bbf9e9e..199d63850 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -83,7 +83,7 @@ def test_no_token_password_in_repr(self) -> None: "Config(storage_api_url=None, storage_token='****', branch_id=None, workspace_schema=None, " 'oauth_client_id=None, oauth_client_secret=None, ' 'oauth_server_url=None, oauth_scope=None, mcp_server_url=None, ' - 'jwt_secret=None, bearer_token=None, conversation_id=None)' + 'jwt_secret=None, bearer_token=None, conversation_id=None, project_id=None)' ) @pytest.mark.parametrize( diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8525e5c0c..63736235b 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -797,3 +797,46 @@ def test_apply_request_config_pins_storage_api_url( # Only the Storage API URL is pinned; the other per-request headers keep working. assert applied.storage_token == headers.get('X-Storage-Api-Token', 'server-token') assert applied.branch_id == headers.get('X-Branch-Id') + + +class TestProgrammaticTokenExchange: + """SessionStateMiddleware exchanges programmatic tokens via the auth-bridge resolver (PSGO-261).""" + + @pytest.mark.asyncio + async def test_missing_kubernetes_token_path_raises(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_pat_abc', project_id='1') + with pytest.raises(ValueError, match='KBC_KUBERNETES_TOKEN_PATH'): + await SessionStateMiddleware._exchange_programmatic_token(config) + + @pytest.mark.asyncio + async def test_missing_project_id_raises(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_pat_abc') + with pytest.raises(ValueError, match='project id is required'): + await SessionStateMiddleware._exchange_programmatic_token(config) + + @pytest.mark.asyncio + async def test_invalid_project_id_raises(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config( + storage_api_url='https://connection.keboola.com', storage_token='kbc_pat_abc', project_id='not-an-int' + ) + with pytest.raises(ValueError, match='Invalid project id'): + await SessionStateMiddleware._exchange_programmatic_token(config) + + @pytest.mark.asyncio + async def test_happy_path_calls_resolver(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_abc', project_id='42') + + resolver = MagicMock() + resolver.resolve = AsyncMock(return_value='legacy-storage-token') + with patch('keboola_mcp_server.mcp.StorageTokenResolver', return_value=resolver) as resolver_cls: + token = await SessionStateMiddleware._exchange_programmatic_token(config) + + assert token == 'legacy-storage-token' + resolver_cls.assert_called_once_with( + storage_api_url='https://connection.keboola.com', kubernetes_token_path='/var/run/secrets/token' + ) + resolver.resolve.assert_awaited_once_with(subject_token='kbc_at_abc', project_id=42) From 57f6bd44d933cf0c4ccff6516b40aa63e7ee185b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 24 Jun 2026 22:07:01 +0200 Subject: [PATCH 02/89] feat(PSGO-261): local browser PKCE login (stack-only) with token refresh Adds Part B: the locally-run (stdio) MCP server authenticates with only the stack URL. A browser PKCE login leases a whole-stack session (access + refresh token), stored to a mode-600 file and refreshed during usage; the leased kbc_at_* token is forwarded downstream as the bearer credential. - auth_login.py: PKCE login (authorize -> loopback callback -> /v1/auth/pkce/token), mode-600 credential store keyed by stack host, get_access_token() that refreshes via /v1/auth/token/refresh near expiry and re-logs-in on a dead token. - cli.py: `login` subcommand (interactive); stdio startup loads stored tokens when no KBC_STORAGE_TOKEN is set, so local needs only KBC_STORAGE_API_URL. - mcp.py: a programmatic token with no KBC_KUBERNETES_TOKEN_PATH (local) is now forwarded downstream as a Bearer (+ X-KBC-ProjectId when a project is selected) instead of erroring; the resolver exchange still runs on the deployed server. - Unit tests for PKCE crypto, code exchange, refresh, credential store, and the refresh/dead-token paths. Live browser+stack flow must be verified on a dev stack. The OAuth->PAT exchange remains a separate PR. Linear: https://linear.app/keboola/issue/PSGO-261 Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/auth_login.py | 251 +++++++++++++++++++++++++++ src/keboola_mcp_server/cli.py | 39 +++++ src/keboola_mcp_server/mcp.py | 19 +- tests/test_auth_login.py | 134 ++++++++++++++ 4 files changed, 438 insertions(+), 5 deletions(-) create mode 100644 src/keboola_mcp_server/auth_login.py create mode 100644 tests/test_auth_login.py diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py new file mode 100644 index 000000000..645e0519d --- /dev/null +++ b/src/keboola_mcp_server/auth_login.py @@ -0,0 +1,251 @@ +"""Local browser PKCE login for the MCP server (PSGO-261, Part B). + +Lets a user authenticate the locally-run (stdio) MCP server with only the stack URL: +a browser PKCE flow leases a whole-stack session (access + refresh token), which is +stored to a mode-600 file and refreshed during usage. The leased ``kbc_at_*`` access +token is then forwarded downstream as the bearer credential. + +The interactive browser/loopback orchestration lives in ``perform_login``; the HTTP +calls (``exchange_code``, ``refresh_tokens``) are split out so they can be tested with +an injected httpx transport. +""" + +import base64 +import hashlib +import json +import logging +import os +import secrets +import time +import urllib.parse +import webbrowser +from dataclasses import asdict, dataclass +from http.server import BaseHTTPRequestHandler, HTTPServer +from pathlib import Path +from typing import cast +from urllib.parse import urlparse, urlunparse + +import httpx + +LOG = logging.getLogger(__name__) + +DEFAULT_CLIENT_ID = 'keboola-cli-demo' +_AUTHORIZE_PATH = 'admin/auth/pkce/authorize' +_TOKEN_PATH = 'v1/auth/pkce/token' +_REFRESH_PATH = 'v1/auth/token/refresh' +_REFRESH_SKEW_SECONDS = 60 +_CREDENTIALS_PATH = Path.home() / '.keboola' / 'mcp' / 'credentials.json' + + +def _client_id() -> str: + # Configurable so the real MCP client id can replace the demo value via a secret. + return os.environ.get('KBC_PKCE_CLIENT_ID') or DEFAULT_CLIENT_ID + + +def _base_url(storage_api_url: str) -> str: + parsed = urlparse(storage_api_url) + if not parsed.hostname or not parsed.hostname.startswith('connection.'): + raise ValueError(f'Invalid Keboola Storage API URL: {storage_api_url}') + return urlunparse(('https', parsed.hostname, '', '', '', '')) + + +def _b64url(data: bytes) -> str: + return base64.urlsafe_b64encode(data).decode('ascii').rstrip('=') + + +@dataclass(frozen=True) +class TokenSet: + """A leased session: the access token plus what's needed to refresh it.""" + + access_token: str + refresh_token: str + expires_at: float # epoch seconds + session_id: str | None = None + + @property + def is_near_expiry(self) -> bool: + return time.time() >= (self.expires_at - _REFRESH_SKEW_SECONDS) + + +def _parse_token_response(body: dict, *, now: float | None = None) -> TokenSet: + now = time.time() if now is None else now + return TokenSet( + access_token=cast(str, body['accessToken']), + refresh_token=cast(str, body['refreshToken']), + expires_at=now + float(body.get('expiresIn') or 0), + session_id=cast('str | None', body.get('sessionId')), + ) + + +async def exchange_code( + storage_api_url: str, + *, + code: str, + state: str, + code_verifier: str, + redirect_uri: str, + transport: httpx.AsyncBaseTransport | None = None, +) -> TokenSet: + """Exchanges a PKCE authorization code for a session token set.""" + payload = { + 'clientId': _client_id(), + 'code': code, + 'state': state, + 'redirectUri': redirect_uri, + 'codeVerifier': code_verifier, + } + async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + response = await client.post(f'{_base_url(storage_api_url)}/{_TOKEN_PATH}', json=payload) + response.raise_for_status() + return _parse_token_response(cast(dict, response.json())) + + +async def refresh_tokens( + storage_api_url: str, + *, + refresh_token: str, + transport: httpx.AsyncBaseTransport | None = None, +) -> TokenSet: + """Exchanges a refresh token for a new (rotated) session token set.""" + async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + response = await client.post( + f'{_base_url(storage_api_url)}/{_REFRESH_PATH}', json={'refreshToken': refresh_token} + ) + response.raise_for_status() + return _parse_token_response(cast(dict, response.json())) + + +# --- credential storage (mode-600 file, keyed by stack host) --- + + +def _store_key(storage_api_url: str) -> str: + return cast(str, urlparse(storage_api_url).hostname) + + +def _read_store() -> dict: + if not _CREDENTIALS_PATH.is_file(): + return {} + try: + return cast(dict, json.loads(_CREDENTIALS_PATH.read_text())) + except (ValueError, OSError): + LOG.warning('Could not read MCP credentials file; treating as empty.') + return {} + + +def _write_store(store: dict) -> None: + """Writes the credential store with restrictive permissions, never widening them. + + The file is created 0600 atomically (no world-readable window between create and + chmod) and its parent directory 0700. + """ + _CREDENTIALS_PATH.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + fd = os.open(_CREDENTIALS_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, 'w') as f: + json.dump(store, f, indent=2) + # O_CREAT honors the mode only when creating; chmod covers a pre-existing file. + _CREDENTIALS_PATH.chmod(0o600) + + +def load_tokens(storage_api_url: str) -> TokenSet | None: + entry = _read_store().get(_store_key(storage_api_url)) + if not entry: + return None + return TokenSet(**entry) + + +def save_tokens(storage_api_url: str, tokens: TokenSet) -> None: + store = _read_store() + store[_store_key(storage_api_url)] = asdict(tokens) + _write_store(store) + + +async def get_access_token( + storage_api_url: str, + *, + transport: httpx.AsyncBaseTransport | None = None, +) -> str: + """ + Returns a valid access token for the stack, refreshing (and persisting the rotated + pair) when near expiry. Raises if there are no stored credentials (run ``login``). + """ + tokens = load_tokens(storage_api_url) + if not tokens: + raise RuntimeError( + f'No stored credentials for {storage_api_url}. Run "keboola-mcp-server login --api-url " first.' + ) + if tokens.is_near_expiry: + try: + tokens = await refresh_tokens(storage_api_url, refresh_token=tokens.refresh_token, transport=transport) + except httpx.HTTPStatusError as e: + # Dead token (refresh rejected). Drop the stale credentials and force a re-login. + _forget(storage_api_url) + raise RuntimeError( + f'Session for {storage_api_url} has expired; run "keboola-mcp-server login --api-url " again.' + ) from e + save_tokens(storage_api_url, tokens) + return tokens.access_token + + +def _forget(storage_api_url: str) -> None: + store = _read_store() + if store.pop(_store_key(storage_api_url), None) is not None: + _write_store(store) + + +# --- interactive browser login (not unit-tested; exercises a real browser + loopback) --- + + +class _CallbackHandler(BaseHTTPRequestHandler): + result: dict = {} + + def do_GET(self) -> None: # noqa: N802 (BaseHTTPRequestHandler API) + query = urllib.parse.parse_qs(urlparse(self.path).query) + type(self).result = {k: v[0] for k, v in query.items()} + self.send_response(200) + self.send_header('Content-Type', 'text/plain') + self.end_headers() + self.wfile.write(b'Login complete. You can close this tab and return to the terminal.') + + def log_message(self, *args) -> None: # silence the default stderr logging + pass + + +async def perform_login(storage_api_url: str, *, open_browser=webbrowser.open) -> TokenSet: + """Runs the interactive PKCE browser login and persists the resulting tokens.""" + verifier = _b64url(secrets.token_bytes(48)) # 64 url-safe chars + challenge = _b64url(hashlib.sha256(verifier.encode('ascii')).digest()) + state = _b64url(secrets.token_bytes(32)) + + server = HTTPServer(('127.0.0.1', 0), _CallbackHandler) + redirect_uri = f'http://127.0.0.1:{server.server_address[1]}/callback' + params = { + 'responseType': 'code', + 'clientId': _client_id(), + 'redirectUri': redirect_uri, + 'codeChallenge': challenge, + 'codeChallengeMethod': 'S256', + 'state': state, + } + authorize_url = f'{_base_url(storage_api_url)}/{_AUTHORIZE_PATH}?{urllib.parse.urlencode(params)}' + print(f'Open this URL in your browser to authenticate:\n\n {authorize_url}\n', flush=True) + open_browser(authorize_url) + + _CallbackHandler.result = {} + server.handle_request() # blocks until the browser hits /callback + server.server_close() + result = _CallbackHandler.result + + if result.get('error'): + raise RuntimeError(f'Authorization failed: {result.get("error")} {result.get("errorDescription", "")}'.strip()) + if not secrets.compare_digest(result.get('state', ''), state): + raise RuntimeError('Authorization state mismatch; aborting login.') + code = result.get('code') + if not code: + raise RuntimeError('Authorization callback did not return a code.') + + print('Exchanging authorization code for tokens…', flush=True) + tokens = await exchange_code( + storage_api_url, code=code, state=state, code_verifier=verifier, redirect_uri=redirect_uri + ) + save_tokens(storage_api_url, tokens) + return tokens diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 21f56be9c..46222d390 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -3,11 +3,13 @@ import argparse import asyncio import contextlib +import dataclasses import json import logging.config import os import pathlib import sys +import time import traceback import pydantic @@ -57,6 +59,18 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: parser.add_argument('--port', type=int, default=8000, metavar='INT', help='The port to listen on.') parser.add_argument('--log-config', type=pathlib.Path, metavar='PATH', help='Logging config file.') + subparsers = parser.add_subparsers(dest='command') + login_parser = subparsers.add_parser( + 'login', + help='Authenticate the local MCP server via a browser PKCE login and store the leased tokens.', + ) + login_parser.add_argument( + '--api-url', + metavar='URL', + help='Keboola Storage API URL (e.g. https://connection..keboola.com). ' + 'Falls back to KBC_STORAGE_API_URL.', + ) + return parser.parse_args(args) @@ -98,6 +112,18 @@ async def _http_exception_handler(request: Request, exc: HTTPException): } +async def _run_login(api_url: str | None) -> None: + """Runs the interactive browser PKCE login and stores the leased tokens.""" + from keboola_mcp_server.auth_login import perform_login + + storage_api_url = api_url or os.environ.get('KBC_STORAGE_API_URL') + if not storage_api_url: + raise RuntimeError('A Storage API URL is required for login: pass --api-url or set KBC_STORAGE_API_URL.') + tokens = await perform_login(storage_api_url) + remaining = max(0, int(tokens.expires_at - time.time())) + print(f'\n✓ Session stored for {storage_api_url} (access token expires in ~{remaining}s).') + + async def run_server(args: list[str] | None = None) -> None: """Runs the MCP server in async mode.""" parsed_args = parse_args(args) @@ -124,6 +150,10 @@ async def run_server(args: list[str] | None = None) -> None: stream=sys.stderr, ) + if parsed_args.command == 'login': + await _run_login(getattr(parsed_args, 'api_url', None)) + return + # Create config from the CLI arguments config = Config( storage_api_url=parsed_args.api_url, @@ -134,6 +164,15 @@ async def run_server(args: list[str] | None = None) -> None: try: # Create and run the server if parsed_args.transport == 'stdio': + # Local/stdio needs only the stack URL: with no token configured, use the tokens + # leased by a prior browser `login` (refreshing them as needed). + config = config.replace_by(os.environ) + if not config.storage_token and config.storage_api_url: + from keboola_mcp_server.auth_login import get_access_token + + access_token = await get_access_token(config.storage_api_url) + config = dataclasses.replace(config, storage_token=access_token) + runtime_config = ServerRuntimeInfo(transport=parsed_args.transport) keboola_mcp_server: FastMCP = create_server(config, runtime_info=runtime_config) if config.oauth_client_id or config.oauth_client_secret: diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index c8e5b0a66..31a70eb0c 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -378,17 +378,26 @@ async def create_session_state( storage_token = config.storage_token bearer_token = config.bearer_token + extra_headers: dict[str, Any] = {} if is_programmatic_token(storage_token): - # A Keboola programmatic token (kbc_at_/kbc_pat_) is not a Storage token; exchange - # it for the project's legacy Storage token and use that downstream unchanged. - storage_token = await cls._exchange_programmatic_token(config) - bearer_token = None + if os.environ.get('KBC_KUBERNETES_TOKEN_PATH'): + # Deployed: exchange the programmatic token (kbc_at_/kbc_pat_) for the project's + # legacy Storage token via the auth-bridge resolver, then use it downstream unchanged. + storage_token = await cls._exchange_programmatic_token(config) + bearer_token = None + else: + # Local: no projected SA token to reach the resolver. Forward the programmatic token + # downstream as a Bearer and let PAT-aware services exchange it; name the target + # project when one has been selected. + bearer_token = storage_token + if config.project_id: + extra_headers['X-KBC-ProjectId'] = config.project_id client = await KeboolaClient( storage_api_url=config.storage_api_url, storage_api_token=storage_token, bearer_token=bearer_token, - headers=cls._get_headers(runtime_info), + headers={**cls._get_headers(runtime_info), **extra_headers}, readonly=readonly, own_stack_storage_api_url=own_stack_storage_api_url, ).with_branch_id(config.branch_id) diff --git a/tests/test_auth_login.py b/tests/test_auth_login.py new file mode 100644 index 000000000..e22d761d5 --- /dev/null +++ b/tests/test_auth_login.py @@ -0,0 +1,134 @@ +"""Tests for the local browser PKCE login + credential store (PSGO-261, Part B).""" + +import base64 +import hashlib +import stat +import time +from pathlib import Path + +import httpx +import pytest + +from keboola_mcp_server import auth_login +from keboola_mcp_server.auth_login import ( + TokenSet, + exchange_code, + get_access_token, + load_tokens, + refresh_tokens, + save_tokens, +) + +STACK = 'https://connection.keboola.com' + + +@pytest.fixture +def creds_file(tmp_path: Path, monkeypatch) -> Path: + path = tmp_path / 'creds' / 'credentials.json' + monkeypatch.setattr(auth_login, '_CREDENTIALS_PATH', path) + return path + + +def _token_response(handler_status: int = 200): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + handler_status, + json={ + 'accessToken': 'kbc_at_new', + 'refreshToken': 'kbc_rt_new', + 'tokenType': 'Bearer', + 'expiresIn': 3600, + 'sessionId': 'sess-1', + }, + ) + + return httpx.MockTransport(handler) + + +@pytest.mark.asyncio +async def test_exchange_code_parses_token_set() -> None: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured['url'] = str(request.url) + return httpx.Response(200, json={'accessToken': 'kbc_at_x', 'refreshToken': 'kbc_rt_x', 'expiresIn': 3600}) + + tokens = await exchange_code( + STACK, + code='c', + state='s', + code_verifier='v', + redirect_uri='http://127.0.0.1:1/callback', + transport=httpx.MockTransport(handler), + ) + assert tokens.access_token == 'kbc_at_x' + assert tokens.refresh_token == 'kbc_rt_x' + assert tokens.expires_at > time.time() + assert captured['url'] == 'https://connection.keboola.com/v1/auth/pkce/token' + + +@pytest.mark.asyncio +async def test_refresh_tokens_rotates_pair() -> None: + tokens = await refresh_tokens(STACK, refresh_token='kbc_rt_old', transport=_token_response()) + assert tokens.access_token == 'kbc_at_new' + assert tokens.refresh_token == 'kbc_rt_new' + + +def test_save_and_load_round_trip_mode_600(creds_file: Path) -> None: + ts = TokenSet(access_token='kbc_at_1', refresh_token='kbc_rt_1', expires_at=time.time() + 3600, session_id='s') + save_tokens(STACK, ts) + + assert stat.S_IMODE(creds_file.stat().st_mode) == 0o600 + loaded = load_tokens(STACK) + assert loaded == ts + # A different stack has no credentials. + assert load_tokens('https://connection.other.keboola.com') is None + + +def test_is_near_expiry() -> None: + assert TokenSet('a', 'r', expires_at=time.time() + 10).is_near_expiry is True + assert TokenSet('a', 'r', expires_at=time.time() + 3600).is_near_expiry is False + + +@pytest.mark.asyncio +async def test_get_access_token_without_credentials_raises(creds_file: Path) -> None: + with pytest.raises(RuntimeError, match='Run "keboola-mcp-server login'): + await get_access_token(STACK) + + +@pytest.mark.asyncio +async def test_get_access_token_returns_valid_token_without_refresh(creds_file: Path) -> None: + ts = TokenSet('kbc_at_valid', 'kbc_rt_1', expires_at=time.time() + 3600) + save_tokens(STACK, ts) + # Transport would 500 if called — proves no refresh happens for a fresh token. + token = await get_access_token(STACK, transport=_token_response(500)) + assert token == 'kbc_at_valid' + + +@pytest.mark.asyncio +async def test_get_access_token_refreshes_near_expiry_and_persists(creds_file: Path) -> None: + save_tokens(STACK, TokenSet('kbc_at_old', 'kbc_rt_old', expires_at=time.time() + 5)) + token = await get_access_token(STACK, transport=_token_response()) + assert token == 'kbc_at_new' + # Rotated pair persisted. + assert load_tokens(STACK).refresh_token == 'kbc_rt_new' + + +@pytest.mark.asyncio +async def test_get_access_token_dead_token_forgets_and_raises(creds_file: Path) -> None: + save_tokens(STACK, TokenSet('kbc_at_old', 'kbc_rt_dead', expires_at=time.time() + 5)) + with pytest.raises(RuntimeError, match='has expired'): + await get_access_token(STACK, transport=_token_response(401)) + # Stale credentials dropped so the next start triggers a fresh login. + assert load_tokens(STACK) is None + + +def test_pkce_challenge_is_sha256_of_verifier() -> None: + verifier = auth_login._b64url(b'0123456789abcdef0123456789abcdef0123456789ab') + expected = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode('ascii')).digest()).decode().rstrip('=') + assert auth_login._b64url(hashlib.sha256(verifier.encode('ascii')).digest()) == expected + + +def test_invalid_stack_url_rejected() -> None: + with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): + auth_login._base_url('https://example.com') From 3c242cecaead89b67e6062b72b49613dd40211ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 07:03:52 +0200 Subject: [PATCH 03/89] docs(PSGO-261): extend RFC with multi-project scope Add increment-2 design: token introspection to enumerate accessible projects, scoped-token exchange (/v1/auth/pat/exchange), transparent per-project fan-out for read tools, and the ask-first scope model. Records decisions D1(revised), D2(extended), D6-D8 and resolutions. Co-Authored-By: Claude Opus 4.8 (1M context) --- TOOLS.md | 74 ++++++++++ feature_spec/pat_token_support/RFC.md | 205 ++++++++++++++++++++++++++ 2 files changed, 279 insertions(+) diff --git a/TOOLS.md b/TOOLS.md index fb0861cf4..65fd56ea7 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -50,9 +50,11 @@ providing their configuration IDs. - [modify_streamlit_data_app](#modify_streamlit_data_app): Creates or updates a Streamlit data app. ### Project Tools +- [get_accessible_projects](#get_accessible_projects): Lists the Keboola projects the current login can access across the stack. - [get_project_info](#get_project_info): Retrieves structured information about the current project, including essential context and base instructions for working with it (e. +- [set_project_scope](#set_project_scope): Scopes the current session to a set of Keboola projects. - [update_project_description](#update_project_description): Updates the description of the current Keboola project. ### SQL Tools @@ -3552,6 +3554,31 @@ configuration is created e.g. keboola.ex-google-analytics-v4 and keboola.ex-gmai --- # Project Tools + +## get_accessible_projects +**Annotations**: `read-only` + +**Tags**: `project` + +**Description**: + +Lists the Keboola projects the current login can access across the stack. + +Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), +present the projects, and ask whether they want to work across all of them or a subset. Then call +`set_project_scope` with their choice. + + +**Input JSON Schema**: +```json +{ + "additionalProperties": false, + "properties": {}, + "type": "object" +} +``` + +--- ## get_project_info **Annotations**: `read-only` @@ -3577,6 +3604,53 @@ to establish the project context before using other tools. } ``` +--- + +## set_project_scope +**Annotations**: `read-only` + +**Tags**: `project` + +**Description**: + +Scopes the current session to a set of Keboola projects. + +Mints a scoped access token (narrowed to `project_ids`, optionally read-only) that is used for the +rest of the conversation. Read-only tools then run against every scoped project in a single call; +write operations target the active (first) project only. Call this when the user states which +projects to work on; it can be called again any time to re-scope. + + +**Input JSON Schema**: +```json +{ + "additionalProperties": false, + "properties": { + "project_ids": { + "anyOf": [ + { + "items": { + "type": "integer" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "description": "The project ids to scope the session to. Omit or pass null to scope to ALL accessible projects." + }, + "read_only": { + "default": false, + "description": "If true, mint a read-only scoped token (no write operations in any scoped project).", + "type": "boolean" + } + }, + "type": "object" +} +``` + --- ## update_project_description diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index fc4834a93..e17e8ffe1 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -206,3 +206,208 @@ OAuth is **not** removed (the MCP protocol needs it for HTTP transport). The OAu 3. **Refresh token** — treated as an opaque string (no prefix assumptions). 4. **SA token path env var** — align with the workspace step-up var (`b971146f`) and the Go services' `*_KUBERNETES_TOKEN_PATH` convention; share one file-read helper. 5. **Refresh + dead token** — the server **always refreshes during usage** when it holds the token pair; when the token is dead (refresh fails), it clears stored credentials and **enforces re-login**. + +--- + +# Extension: Multi-project scope via introspect + scoped exchange (PSGO-261, increment 2) + +> This section extends the RFC above. Parts A/B (programmatic-token exchange, PKCE login) are unchanged +> and are the substrate this builds on. It revises decisions **D1** and **D2** (see below). + +## New problem + +Parts A/B give the server a programmatic token and a single `project_id`. A whole-stack PAT/AT can +actually reach **many** projects, and the Kai multi-project workflows (the parent driver, PAT-1838) +need one agent session to act across several of them. Two gaps remain: + +1. **Discovery.** The server has no way to enumerate which projects the inbound token can reach. + (Brainstorm left this as an open question: "the exact stack-level endpoint to enumerate + PAT-accessible projects.") +2. **Scope.** `project_id` is a single value. There is no way to (a) operate over a set of projects, + nor (b) *narrow* a whole-stack token down to a reviewed subset for the rest of the session. + +## Token-contract additions (authoritative, from connection auth API) + +### Introspect — enumerate accessible projects + +``` +GET {connection}/v1/auth/token/introspect +Headers: Authorization: Bearer + +200: +{ "sessionId": "...", "user": { "id", "email", "name" }, + "grantType": "authorization_code", "expiresAt": "", + "projects": [ { "id": , "name": "...", "role": "admin|..." }, ... ] } +``` + +This is the discovery endpoint. It works for any programmatic token and is the source of truth for +"which projects can this session touch." + +### Scoped exchange — mint a token narrowed to chosen projects + +``` +POST {connection}/v1/auth/pat/exchange +Headers: Authorization: Bearer +Body: { "expiresIn": null|, "scope": { "projects": [,...]|null, "readOnly": true|null } } + +201: +{ "accessToken": "", "tokenType": "Bearer", "expiresIn": , + "scope": {...}, "readOnly": , "parentTokenId": "...", "parentTokenType": "session", + "expiresAt": "", "pat": { "id", "name", "scope", "projects": [...], "readOnly", ... } } +``` + +`scope.projects = null` → all projects (whole-stack). A non-null list mints a token that can reach +**only** those projects. `readOnly: true` mints a read-only token. The returned `accessToken` becomes +the session's subject token for all downstream exchange/forwarding. + +## Required behavior + +1. **Discovery tool.** `get_accessible_projects()` calls introspect and returns the projects list + (id, name, role) plus the user identity. Read-only, no side effects. +2. **Scope-selection tool.** `set_project_scope(project_ids: list[int] | "all", read_only: bool=false)`: + - `"all"` → scope = every introspected project id; keep the current (whole-stack) token. + - a subset → call `/v1/auth/pat/exchange` with `scope.projects=project_ids` (+ `readOnly`), store + the returned scoped `accessToken` as the **session subject token**, set scope = `project_ids`, + and clear the per-project client cache so it rebuilds against the scoped token. +3. **Conversation-start nudge.** Server instructions tell the agent: on first interaction call + `get_accessible_projects`, present them, and ask the user **"work across all of these, or a + subset?"**; call `set_project_scope` with the answer. (MCP has no protocol-level startup prompt — + this is the idiomatic discovery-tool + instructions pattern, same shape as `get_project_info`.) +4. **Transparent multi-project execution.** Once scope is set, existing tools run **once per project + in scope** with no per-tool `projects[]` argument. Mechanism (decision D6): + - Session state holds `scope` (ordered project_ids), `active_project_id`, and a lazy + `project_id -> KeboolaClient` cache. `KeboolaClient.from_state(state)` returns the client for + `active_project_id`. **All 43 existing call sites are unchanged.** + - A dispatch-layer wrapper (`SessionStateMiddleware.on_call_tool`) reads the scope: + - **1 project** (or legacy single-project session): set `active_project_id`, call the tool once, + return its result **raw** — byte-for-byte today's behavior. + - **N projects, read tool:** loop the scope, set `active_project_id` per iteration, collect into a + per-project envelope `[{ "project_id": , "result": }, ...]`. No semantic + merge — the envelope preserves each tool's native return shape (this is the answer to the + "merging arbitrary shapes is lossy" risk: we wrap, we don't merge). + - **N projects, write tool:** do **not** fan out. Require a single target project; if scope has + >1 and no explicit single target was confirmed, return a clear error instructing the agent to + confirm with the user and target one project (decision D8). + - Per-project clients build lazily: deployed → resolver exchange per project (Part A) using the + scoped subject token; local → forward bearer + `X-KBC-ProjectId: `. + - Fan-out is sequential in v1. `# ponytail: sequential fan-out; asyncio.gather if N-project latency bites.` +5. **Read/write classification.** An explicit set of mutating tool names (or a registration-time flag) + drives the write-policy branch. Explicit list over magic — there are few write tools. + +## Mode / availability matrix (additions) + +| Inbound credential | Introspect / scope tooling | Multi-project | +| --- | --- | --- | +| programmatic (`kbc_at_*`/`kbc_pat_*`, PKCE or Bearer) | available | yes | +| legacy `KBC_STORAGE_TOKEN` | n/a (project-bound token) | no — single project, unchanged | +| OAuth `SimpleOAuthProvider` (current SAPI mint) | n/a until OAuth→PAT PR lands | no (interim) | + +## Revised decisions + +- **D1 (revised) — scope narrowing is now token-enforced, not advisory.** The original D1 said + "no minting; narrowing is runtime-only session state." With the user choosing the `pat/exchange` + path, narrowing to a subset **mints a scoped token** (in-memory, session-lived, never persisted). + The *stored* (on-disk PKCE) credential is still whole-stack — D1's storage stance holds — but the + *active* session token is the scoped one, so a tool can no longer reach an out-of-scope project even + by bug. Strictly stronger than the original advisory model. +- **D2 (extended) — `project_id` → project scope (a set).** Still explicit session state, never + silently derived. Default scope = `[KBC_PROJECT_ID]` / `X-KBC-ProjectId` (today's single-project + behavior, backward compatible). `get_accessible_projects` + `set_project_scope` replace the + previously-hypothetical "select-project tool"; introspect closes the open enumeration question. +- **D6 (new) — transparent fan-out via active-project indirection.** Tools take no `projects[]` arg; + the dispatch wrapper swaps `active_project_id` and the per-project client cache. Multi-project + results use a per-project envelope, never a semantic merge. Zero changes to the 43 `from_state` + sites. (Chosen over a per-tool `projects[]` param.) +- **D7 (new) — scoped exchange uses `/v1/auth/pat/exchange`.** Not `/v1/auth/pat` (PAT create). The + exchange yields a child token (`parentTokenType: session`) tied to the current session, which is the + right lifetime for a session-scoped narrowing. +- **D8 (new) — multi-project writes are user-driven only.** Read/query/search tools fan out freely. + Mutating tools (create/update/run config, flow, job, data-app, transformation) never fan out + automatically: with >1 project in scope they require a single confirmed target project. Server + instructions state: **the agent must never write to more than one project without explicit user + guidance or confirmation.** Bulk multi-project writes are possible but only on that explicit signal. + +## Scope changes (relative to the base RFC) + +**Moved into scope:** introspect-based project enumeration; `/v1/auth/pat/exchange` scoped token +minting; multi-project read fan-out; user-confirmed multi-project writes; +`get_accessible_projects` + `set_project_scope` tools; per-project client cache. + +**Still out of scope:** OAuth→PAT exchange (separate PR); caching of resolver results; keyring/DB +credential storage; `/v1/auth/pat` PAT lifecycle management (create/list/revoke) tools; parallel +fan-out (sequential in v1). + +## Delivery plan (phased, compact) + +**Phase 1 — Discovery (low risk, read-only).** +- `clients/auth_bridge.py` (or a new `clients/auth.py`): `introspect(subject_token) -> Introspection` + (user + projects[]). GET `/v1/auth/token/introspect`, token redaction same as the resolver. +- `tools/project.py`: `get_accessible_projects()` tool. +- Server instructions: add the "ask all-vs-subset at start" nudge. +- Tests: introspect success/parse, 401/timeout mapping, no-token-in-logs; tool returns projects. + +**Phase 2 — Scoped exchange + session scope state (revises D1).** +- `exchange_scope(subject_token, project_ids|None, read_only, expires_in) -> scoped token` (POST + `/v1/auth/pat/exchange`). +- Session scope state: `scope: list[int]`, `read_only: bool`, scoped subject token; default scope from + `project_id`. `set_project_scope` tool wires exchange → state → cache invalidation. +- Tests: subset → exchange called with right body, scoped token stored; "all" → no exchange; read_only + propagates; scope defaults to single project when unset. + +**Phase 3 — Transparent fan-out (the core refactor, D6/D8).** +- Session state: `active_project_id` + lazy `project_id -> KeboolaClient` cache; `from_state` returns + the active client (indirection only — call sites unchanged). +- `SessionStateMiddleware.on_call_tool` wrapper: 1-project raw passthrough; N-project read envelope; + N-project write guard. +- Explicit write-tool name set. +- Tests: single-project unchanged (regression); 2-project read returns enveloped per-project results; + write tool with N-project scope refuses without a confirmed target; per-project client built with the + right token/header. + +**Cross-cutting:** version bump (minor — new capability), `uv.lock`, `TOOLS.md` regen (new tools + +the per-project envelope shape change the docs), integration tests on a dev stack with a real +`kbc_pat_*` across ≥2 projects. + +## Open questions (new) + +- [ ] **Envelope vs raw for exactly-1-in-scope-but-explicitly-multi.** Confirm: a scope of exactly one + project returns raw (not a 1-element envelope) so single-project UX never regresses. (Assumed yes.) +- [ ] **`expiresIn` for the scoped exchange.** Use `null` (inherit parent/default) in v1, or pin to the + remaining parent lifetime? Affects mid-session expiry of the scoped token. +- [ ] **Scope change mid-session re-introspect.** After `set_project_scope`, do we re-introspect to + validate the subset is still reachable, or trust the prior introspect? (Lean: trust; resolver/exchange + will reject an out-of-scope project anyway.) +- [ ] **Write-target confirmation mechanism.** Is the "confirmed single target" a tool argument + (`project_id` on the write tool), a separate `set_write_target` call, or purely instruction-driven? + (Lean: explicit `project_id` arg on write tools, honored only when scope >1.) + +## Resolutions (2026-06-30) — answers to the increment-2 open questions + +- **Scoping requires a dedicated tool (`set_project_scope`); it is the only mechanism.** The MCP + server receives nothing from the conversation except tool calls — plain chat text never reaches the + server. Scope is server-side state (scoped token + per-project client cache + active project), so the + user's in-conversation intent can only change scope by the agent invoking the tool. The tool is + callable **at any point mid-conversation**, not just at start; the conversation-start nudge is an + instruction-level suggestion, not a gate. The user drives scope changes by saying so; the agent + translates that into the tool call. (Resolves the recurring "do I need a tool / is it user-driven" + question: yes, a tool; driven by the user via conversation, any time.) +- **The tool does not swap a single client — it invalidates the cache (D6).** `set_project_scope` + stores the new scope + scoped token and **clears the per-project client cache**. `from_state` then + lazily rebuilds each project's client against the new scoped token. Cleaner than replacing one + `KeboolaClient` object in state. +- **Q2 (scoped-token lifetime) — resolved: the child token is re-minted, not independently refreshed.** + The `pat/exchange` response carries `accessToken` + `expiresAt` but **no `refreshToken`** — the child + (scoped) token is not refreshable on its own. The refreshable credential is the **parent** PKCE + session token (Part B, `/v1/auth/token/refresh`). MCP **remembers the scope selection** + (`project_ids`, `readOnly`); when the scoped child nears expiry it **re-runs `pat/exchange`** against + the still-valid (refresh-backed) parent token to lease a fresh scoped token. `expiresIn: null` at + exchange time (inherit server default) is fine because we re-mint on demand. *Flag: confirm against + the auth API that the child token genuinely has no own refresh token.* +- **Q3 (re-introspect on scope change) — resolved: trust prior, let exchange reject.** When the user + picks a subset, MCP goes straight to `pat/exchange` without re-calling `introspect`. The exchange + endpoint itself rejects any project the token can't reach, so a pre-check is redundant — one fewer + round-trip, and the exchange is the authority. +- **Q1 (exactly-1-in-scope) — confirmed: a single-project scope returns the raw result, not a + 1-element envelope.** Single-project UX is byte-for-byte unchanged. +- **Q4 (write-target confirmation) — lean: explicit `project_id` arg on write tools, honored only when + scope > 1.** (Still open; not blocking.) From 8c9f13cc2d5c27692db227af8771800cf262e398 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 07:04:11 +0200 Subject: [PATCH 04/89] feat(PSGO-261): multi-project scope via introspect + scoped exchange Add token introspection and scoped-token exchange to the local PKCE token lifecycle, two tools (get_accessible_projects, set_project_scope), and transparent per-project fan-out for read-only tools. - auth_login: introspect_token() + exchange_scoped_token() (project ids sent as strings, as the exchange API requires) - mcp: SessionScope persisted across the per-request state rebuild; parent token refreshed during usage and scoped token re-minted near expiry; auto-lease ALL accessible projects as the default scope (MPA); MultiProjectMiddleware fans read tools out per project and deep-merges the structured outputs so results still validate each tool's schema; ask-first gate blocks data tools until set_project_scope confirms scope - tools/project: get_accessible_projects + set_project_scope - server: register MultiProjectMiddleware; MPA/ask-first instructions Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/auth_login.py | 99 +++++++++ src/keboola_mcp_server/mcp.py | 269 ++++++++++++++++++++++++ src/keboola_mcp_server/server.py | 14 ++ src/keboola_mcp_server/tools/project.py | 158 +++++++++++++- tests/test_auth_login.py | 61 ++++++ tests/test_mcp.py | 211 +++++++++++++++++++ tests/test_server.py | 4 +- tests/tools/test_project.py | 88 ++++++++ 8 files changed, 902 insertions(+), 2 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 645e0519d..64b34b2e3 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -33,6 +33,8 @@ _AUTHORIZE_PATH = 'admin/auth/pkce/authorize' _TOKEN_PATH = 'v1/auth/pkce/token' _REFRESH_PATH = 'v1/auth/token/refresh' +_INTROSPECT_PATH = 'v1/auth/token/introspect' +_EXCHANGE_PATH = 'v1/auth/pat/exchange' _REFRESH_SKEW_SECONDS = 60 _CREDENTIALS_PATH = Path.home() / '.keboola' / 'mcp' / 'credentials.json' @@ -77,6 +79,103 @@ def _parse_token_response(body: dict, *, now: float | None = None) -> TokenSet: ) +@dataclass(frozen=True) +class ProjectAccess: + """A project the introspected token can reach.""" + + id: int + name: str | None = None + role: str | None = None + + +@dataclass(frozen=True) +class Introspection: + """Identity + the set of projects a programmatic token can reach (token introspection).""" + + user_id: int | None + user_email: str | None + user_name: str | None + projects: list[ProjectAccess] + + +@dataclass(frozen=True) +class ScopedToken: + """A child access token minted by /v1/auth/pat/exchange, narrowed to a set of projects.""" + + access_token: str + expires_at: float # epoch seconds + project_ids: list[int] + read_only: bool + + @property + def is_near_expiry(self) -> bool: + return time.time() >= (self.expires_at - _REFRESH_SKEW_SECONDS) + + +async def introspect_token( + storage_api_url: str, + *, + subject_token: str, + transport: httpx.AsyncBaseTransport | None = None, +) -> Introspection: + """Enumerates the projects a programmatic token can reach via /v1/auth/token/introspect.""" + async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + response = await client.get( + f'{_base_url(storage_api_url)}/{_INTROSPECT_PATH}', + headers={'Authorization': f'Bearer {subject_token}'}, + ) + response.raise_for_status() + body = cast(dict, response.json()) + user = body.get('user') or {} + projects = [ + ProjectAccess(id=int(p['id']), name=p.get('name'), role=p.get('role')) + for p in body.get('projects', []) + if p.get('id') is not None + ] + return Introspection( + user_id=user.get('id'), + user_email=user.get('email'), + user_name=user.get('name'), + projects=projects, + ) + + +async def exchange_scoped_token( + storage_api_url: str, + *, + subject_token: str, + project_ids: list[int], + read_only: bool = False, + expires_in: int | None = None, + transport: httpx.AsyncBaseTransport | None = None, +) -> ScopedToken: + """ + Mints a child access token scoped to ``project_ids`` via /v1/auth/pat/exchange. + + The child token has no refresh token of its own; it is re-minted from the (refreshable) + parent token when it nears expiry. ``read_only=True`` mints a read-only token. + """ + # The exchange API expects project ids as strings (and rejects integers with a 400). + payload = { + 'expiresIn': expires_in, + 'scope': {'projects': [str(p) for p in project_ids], 'readOnly': read_only or None}, + } + async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + response = await client.post( + f'{_base_url(storage_api_url)}/{_EXCHANGE_PATH}', + headers={'Authorization': f'Bearer {subject_token}'}, + json=payload, + ) + response.raise_for_status() + body = cast(dict, response.json()) + return ScopedToken( + access_token=cast(str, body['accessToken']), + expires_at=time.time() + float(body.get('expiresIn') or 0), + project_ids=list(project_ids), + read_only=bool(body.get('readOnly')), + ) + + async def exchange_code( storage_api_url: str, *, diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 31a70eb0c..31c38c8f5 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -10,6 +10,7 @@ import logging import os import textwrap +import time from collections.abc import Awaitable, Callable, Iterable from typing import Any, TypeVar from unittest.mock import MagicMock @@ -21,6 +22,7 @@ from fastmcp.server.dependencies import get_http_request from fastmcp.server.middleware import CallNext, MiddlewareContext from fastmcp.tools import Tool +from fastmcp.tools.tool import ToolResult from mcp import types as mt from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from pydantic import BaseModel @@ -29,6 +31,7 @@ from starlette.requests import Request from starlette.types import ASGIApp, Receive, Scope, Send +from keboola_mcp_server.auth_login import exchange_scoped_token, get_access_token, introspect_token from keboola_mcp_server.clients.auth_bridge import StorageTokenResolver, is_programmatic_token from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient @@ -39,6 +42,49 @@ LOG = logging.getLogger(__name__) CONVERSATION_ID = 'conversation_id' +SCOPE_KEY = 'project_scope' + +# Tools that must not be fanned out across multiple projects, even when a multi-project scope is +# active and they are read-only: the scope/auth tools operate on the whole-stack token (not a single +# project), and query_data resolves through a per-project WorkspaceManager that the fan-out does not +# swap (it would otherwise run against the active project's workspace N times). +# ponytail: explicit deny-set; per-project WorkspaceManager fan-out is a follow-up, not v1. +# get_project_info and query_data resolve through the active project's WorkspaceManager (workspace id +# / sql dialect), which the fan-out does not swap, so they report the active project only. +_NO_FANOUT_TOOLS = {'get_accessible_projects', 'set_project_scope', 'get_project_info', 'query_data'} + +# Tools allowed before the user has confirmed a project scope. Everything else is blocked with a +# message telling the assistant to ask the user which projects to work on first (ask-first UX). +_BOOTSTRAP_TOOLS = {'get_accessible_projects', 'set_project_scope'} + + +@dataclasses.dataclass(frozen=True) +class SessionScope: + """In-conversation multi-project scope (PSGO-261 increment 2). + + Persisted on the session across the per-request state rebuild. ``project_ids`` is the + user-selected set; ``scoped_token`` is the child access token minted by /v1/auth/pat/exchange + and narrowed to those projects (re-minted from the parent when near expiry). + """ + + project_ids: list[int] + read_only: bool = False + scoped_token: str | None = None + scoped_expires_at: float | None = None + confirmed: bool = False + """True once the user has explicitly chosen a scope via ``set_project_scope``. The default + auto-leased scope is unconfirmed, which gates data tools until the user decides.""" + + @property + def active_project_id(self) -> int | None: + return self.project_ids[0] if self.project_ids else None + + @property + def is_near_expiry(self) -> bool: + if self.scoped_expires_at is None: + return False + return time.time() >= (self.scoped_expires_at - 60) + R = TypeVar('R') T = TypeVar('T') @@ -226,6 +272,16 @@ async def on_request( if http_rq := get_http_request_or_none(): config = self.apply_request_config(http_rq, config, own_stack_storage_api_url=own_stack_storage_api_url) + # In-conversation multi-project scope persists on the session across this per-request + # state rebuild. Read it before the state is overwritten, refresh the stored token during + # usage, and re-mint the scoped token when it nears expiry. With no scope and no preset + # project, auto-lease ALL accessible projects (multi-project mode) so the session + # bootstraps without a hand-picked project and read tools fan out across everything. + scope = self._read_persisted_scope(ctx.session) + if scope is None and not config.project_id: + scope = await self._autolease_default_scope(config) + config, scope = await self._resolve_local_tokens(config, scope) + # TODO: We could probably get rid of the 'state' attribute set on ctx.session and just # pass KeboolaClient and WorkspaceManager instances to a tool as extra parameters. @@ -241,6 +297,8 @@ async def on_request( state = await self.create_session_state( config, runtime_info, own_stack_storage_api_url=own_stack_storage_api_url ) + if scope is not None: + state[SCOPE_KEY] = scope ctx.session.state = state try: @@ -312,6 +370,103 @@ def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_sto return config + @staticmethod + def _read_persisted_scope(session: Any) -> 'SessionScope | None': + """Reads the multi-project scope stashed in the prior request's session state, if any.""" + prior = getattr(session, 'state', None) + if isinstance(prior, dict): + scope = prior.get(SCOPE_KEY) + if isinstance(scope, SessionScope): + return scope + return None + + @classmethod + def _is_local_programmatic(cls, config: Config) -> bool: + """True for a local (non-deployed) session carrying a Keboola programmatic token.""" + return ( + not os.environ.get('KBC_KUBERNETES_TOKEN_PATH') + and bool(config.storage_token) + and bool(config.storage_api_url) + and is_programmatic_token(config.storage_token) + ) + + @classmethod + async def _autolease_default_scope(cls, config: Config) -> 'SessionScope | None': + """ + Default to multi-project mode: scope the session to ALL accessible projects. + + Introspects the programmatic token once to enumerate the projects it can reach and returns a + scope covering all of them (no minted token — the whole-stack parent token is used, narrowed + per request only by the ``X-KBC-ProjectId`` header). Returns None when introspection is + unavailable (deployed server, legacy token, or no reachable projects) so the caller falls + back to the existing single-project behavior. + """ + if not cls._is_local_programmatic(config): + return None + try: + parent = await get_access_token(config.storage_api_url) + except RuntimeError: + parent = config.storage_token + try: + introspection = await introspect_token(config.storage_api_url, subject_token=parent) + except Exception as e: + LOG.warning(f'Could not auto-lease projects from token introspection: {e}') + return None + project_ids = [p.id for p in introspection.projects] + if not project_ids: + return None + LOG.info(f'Multi-project mode: auto-leased {len(project_ids)} accessible project(s) as the default scope.') + return SessionScope(project_ids=project_ids) + + @classmethod + async def _resolve_local_tokens( + cls, config: Config, scope: 'SessionScope | None' + ) -> 'tuple[Config, SessionScope | None]': + """ + For local (non-deployed) programmatic-token sessions, keep tokens fresh during usage. + + Refreshes the stored whole-stack (parent) token via the PKCE credential store. When the user + has explicitly narrowed scope (a minted scoped token is present), that token is re-minted from + the parent when it nears expiry. The default (auto-leased) multi-project scope carries no + minted token and simply uses the parent token, narrowed per request by ``X-KBC-ProjectId``. + On the deployed server (``KBC_KUBERNETES_TOKEN_PATH`` set) the per-request resolver exchange + already handles freshness, so this is a no-op there. + """ + if not cls._is_local_programmatic(config): + return config, scope + + parent = config.storage_token + try: + # Refreshes (and persists the rotated pair) when near expiry; raises if no stored creds. + parent = await get_access_token(config.storage_api_url) + except RuntimeError: + pass # token supplied directly (no PKCE login) — use it as-is + + token = parent + project_id = config.project_id + if scope and scope.project_ids: + project_id = str(scope.active_project_id) + if scope.scoped_token is not None: + if scope.is_near_expiry: + try: + minted = await exchange_scoped_token( + config.storage_api_url, + subject_token=parent, + project_ids=scope.project_ids, + read_only=scope.read_only, + ) + scope = dataclasses.replace( + scope, scoped_token=minted.access_token, scoped_expires_at=minted.expires_at + ) + except Exception as e: + # Don't break the session if re-minting fails; fall back to the parent token. + LOG.warning(f'Could not refresh the scoped token; using the parent token: {e}') + scope = dataclasses.replace(scope, scoped_token=None, scoped_expires_at=None) + token = scope.scoped_token or parent + + config = dataclasses.replace(config, storage_token=token, project_id=project_id) + return config, scope + @classmethod async def _exchange_programmatic_token(cls, config: Config) -> str: """ @@ -628,6 +783,120 @@ async def on_call_tool( return await call_next(context) +class MultiProjectMiddleware(fmw.Middleware): + """Fans a read-only tool call out across every project in the active multi-project scope. + + Single-project (or no) scope is an unchanged passthrough. With >1 project selected, a read-only + tool runs once per project — the active ``KeboolaClient`` in session state is swapped to each + project's client and the per-project results are labelled and concatenated (no structured-content + merge, so each tool keeps its native output shape). Write tools never fan out: they target the + active project only, so the agent can never write to multiple projects without the user explicitly + re-scoping (PSGO-261 decision D8). + """ + + async def on_call_tool( + self, + context: MiddlewareContext[mt.CallToolRequestParams], + call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult], + ) -> mt.CallToolResult: + ctx = context.fastmcp_context + state = ctx.session.state + scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None + + # Ask-first gate: until the user confirms a scope via set_project_scope, block data tools and + # tell the assistant to ask the user which projects to work on. Only applies when a scope has + # been auto-leased (local programmatic session); deployed/legacy sessions have no scope. + if isinstance(scope, SessionScope) and not scope.confirmed and context.message.name not in _BOOTSTRAP_TOOLS: + raise ToolError( + f'This session can access {len(scope.project_ids)} Keboola project(s), but no scope has ' + 'been confirmed yet. Call "get_accessible_projects", show the user their projects, and ask ' + 'whether to work across ALL of them or a subset. Then call "set_project_scope" ' + '(no arguments = all projects, or pass the chosen project ids, optionally read_only=true). ' + 'This confirmation is required once per session.' + ) + + if not isinstance(scope, SessionScope) or len(scope.project_ids) <= 1: + return await call_next(context) + if context.message.name in _NO_FANOUT_TOOLS: + return await call_next(context) + + tool = await ctx.fastmcp.get_tool(context.message.name) + if not is_read_only_tool(tool): + # Write tools target the active project only — never an automatic multi-project write. + return await call_next(context) + + server_state = ServerState.from_context(ctx) + original_client = state.get(KeboolaClient.STATE_KEY) + # Default (auto-leased) scope carries no minted token; fall back to the active client's token. + base_token = scope.scoped_token or (original_client.token if isinstance(original_client, KeboolaClient) else '') + results: list[tuple[int, ToolResult]] = [] + try: + for project_id in scope.project_ids: + state[KeboolaClient.STATE_KEY] = await self._client_for_project( + server_state, base_token, project_id, scope.read_only + ) + results.append((project_id, await call_next(context))) + finally: + state[KeboolaClient.STATE_KEY] = original_client + + return self._merge(results) + + @staticmethod + async def _client_for_project( + server_state: ServerState, token: str, project_id: int, read_only: bool + ) -> KeboolaClient: + return await KeboolaClient( + storage_api_url=server_state.config.storage_api_url, + storage_api_token=token, + bearer_token=token, + headers={ + **SessionStateMiddleware._get_headers(server_state.runtime_info), + 'X-KBC-ProjectId': str(project_id), + }, + readonly=read_only or None, + ).with_branch_id(None) + + @staticmethod + def _deep_merge(a: Any, b: Any) -> Any: + """Merges two per-project structured outputs so the result still validates the tool's schema. + + Lists are concatenated (the combined slice across projects), nested objects merged key by key, + and numeric counters summed; any other scalar keeps the first project's value. This keeps every + required field present with its declared type, so the merged object validates against the + single-project output schema. + """ + if isinstance(a, list) and isinstance(b, list): + return a + b + if isinstance(a, dict) and isinstance(b, dict): + merged = dict(a) + for key, value in b.items(): + merged[key] = MultiProjectMiddleware._deep_merge(a[key], value) if key in a else value + return merged + if isinstance(a, bool) or isinstance(b, bool): + return a + if isinstance(a, (int, float)) and isinstance(b, (int, float)): + return a + b # counters like search "total" + return a + + @staticmethod + def _merge(results: list[tuple[int, 'ToolResult']]) -> 'ToolResult': + # Label each project's output in the text content (attribution the model can read) and, when the + # tool declares structured output, deep-merge the per-project structured payloads into a single + # schema-valid object (lists concatenated across projects) so the client's output-schema + # validation passes. + content: list[Any] = [] + merged_structured: Any = None + for project_id, result in results: + content.append(mt.TextContent(type='text', text=f'=== project {project_id} ===')) + content.extend(result.content or []) + sc = result.structured_content + if sc is not None: + merged_structured = ( + sc if merged_structured is None else MultiProjectMiddleware._deep_merge(merged_structured, sc) + ) + return ToolResult(content=content, structured_content=merged_structured) + + def _to_python(data: Any, exclude_none: bool = True) -> Any | None: if isinstance(data, BaseModel): return data.model_dump(exclude_none=exclude_none, by_alias=False) diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index 286d8bea9..322d31c44 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -20,6 +20,7 @@ from keboola_mcp_server.errors import ValidationErrorMiddleware from keboola_mcp_server.mcp import ( KeboolaMcpServer, + MultiProjectMiddleware, ServerState, SessionStateMiddleware, ToolsFilteringMiddleware, @@ -228,6 +229,18 @@ def create_server( server_state = ServerState(config=config, runtime_info=runtime_info) mcp = KeboolaMcpServer( name='Keboola MCP Server', + instructions=( + 'This server runs in multi-project mode. When the user logs in with a stack-wide Keboola ' + 'token, data tools are BLOCKED until a project scope is confirmed. So at the very START of ' + 'the conversation, before doing anything else: call "get_accessible_projects", show the user ' + 'their projects, and ASK whether to work across ALL of them or a subset. Do not decide for ' + 'them. Then call "set_project_scope" with their answer (no arguments = all projects, or the ' + 'chosen project ids, optionally read_only=true). After that, read-only tools return results ' + 'per project. Never write to more than one project without explicit user confirmation — ' + 'write operations target the active (first-scoped) project only. Note: outside the Storage ' + 'API, some tools may need per-project token support not yet available on every stack; ' + 'surface such errors plainly rather than retrying.' + ), lifespan=create_keboola_lifespan(server_state), auth=oauth_provider, middleware=[ @@ -235,6 +248,7 @@ def create_server( SessionStateMiddleware(), ToolAuthorizationMiddleware(), ToolsFilteringMiddleware(), + MultiProjectMiddleware(), ValidationErrorMiddleware(), ], ) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index f9dc56a58..3f851d166 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -1,16 +1,18 @@ import logging -from typing import Annotated, cast +from typing import Annotated, Optional, cast from fastmcp import Context, FastMCP from fastmcp.tools import FunctionTool from mcp.types import ToolAnnotations from pydantic import BaseModel, Field +from keboola_mcp_server.auth_login import exchange_scoped_token, get_access_token, introspect_token from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager +from keboola_mcp_server.mcp import SCOPE_KEY, SessionScope from keboola_mcp_server.resources.prompts import get_project_system_prompt from keboola_mcp_server.workspace import WorkspaceManager @@ -40,9 +42,46 @@ def add_project_tools(mcp: FastMCP) -> None: ) ) + LOG.info(f'Adding tool {get_accessible_projects.__name__} to the MCP server.') + mcp.add_tool( + FunctionTool.from_function( + get_accessible_projects, + annotations=ToolAnnotations(readOnlyHint=True), + tags={PROJECT_TOOLS_TAG}, + ) + ) + + LOG.info(f'Adding tool {set_project_scope.__name__} to the MCP server.') + mcp.add_tool( + FunctionTool.from_function( + set_project_scope, + annotations=ToolAnnotations(readOnlyHint=True), + tags={PROJECT_TOOLS_TAG}, + ) + ) + LOG.info('Project tools initialized.') +async def _parent_subject_token(client: KeboolaClient) -> str: + """ + Resolves the whole-stack (parent) programmatic token used to introspect/scope. + + Prefers the refreshable token from the local PKCE credential store (so re-scoping always starts + from the parent, never from an already-narrowed scoped token); falls back to whatever bearer the + client currently carries (a directly-supplied PAT, or an HTTP bearer). + """ + if client.bearer_token is None: + raise ValueError( + 'Project scoping requires a Keboola programmatic token (kbc_at_/kbc_pat_). ' + 'Run "keboola-mcp-server login --api-url " first, or supply such a token.' + ) + try: + return await get_access_token(client.storage_api_url) + except RuntimeError: + return client.bearer_token + + async def _resolve_branch_context(client: KeboolaClient) -> tuple[str | int, str, bool]: """ Resolves the current branch's id, name, and dev-branch flag from the storage API. @@ -224,3 +263,120 @@ async def get_project_info( LOG.info('Returning unified project info.') return project_info + + +class AccessibleProject(BaseModel): + id: int = Field(description='The project id.') + name: str | None = Field(default=None, description='The project name.') + role: str | None = Field(default=None, description='The user role in this project (e.g. "admin").') + + +class AccessibleProjects(BaseModel): + user_email: str | None = Field(default=None, description='The email of the authenticated user.') + projects: list[AccessibleProject] = Field(description='The projects the current token can reach across the stack.') + llm_instruction: str = Field( + description='Guidance for the assistant on how to use this result.', + ) + + +class ProjectScope(BaseModel): + project_ids: list[int] = Field(description='The projects the session is now scoped to.') + read_only: bool = Field(description='Whether the scoped token is read-only.') + active_project_id: int | None = Field( + default=None, + description='The project that write operations and single-project tools target.', + ) + llm_instruction: str = Field(description='Guidance for the assistant on the new scope.') + + +@tool_errors() +async def get_accessible_projects(ctx: Context) -> AccessibleProjects: + """ + Lists the Keboola projects the current login can access across the stack. + + Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), + present the projects, and ask whether they want to work across all of them or a subset. Then call + `set_project_scope` with their choice. + """ + client = KeboolaClient.from_state(ctx.session.state) + subject_token = await _parent_subject_token(client) + introspection = await introspect_token(client.storage_api_url, subject_token=subject_token) + projects = [AccessibleProject(id=p.id, name=p.name, role=p.role) for p in introspection.projects] + return AccessibleProjects( + user_email=introspection.user_email, + projects=projects, + llm_instruction=( + 'Ask the user whether to operate across all these projects or a subset, then call ' + '"set_project_scope" with the chosen project ids. Never write to more than one project ' + 'without explicit user confirmation.' + ), + ) + + +@tool_errors() +async def set_project_scope( + ctx: Context, + project_ids: Annotated[ + Optional[list[int]], + Field( + description='The project ids to scope the session to. ' + 'Omit or pass null to scope to ALL accessible projects.' + ), + ] = None, + read_only: Annotated[ + bool, + Field(description='If true, mint a read-only scoped token (no write operations in any scoped project).'), + ] = False, +) -> ProjectScope: + """ + Scopes the current session to a set of Keboola projects. + + Mints a scoped access token (narrowed to `project_ids`, optionally read-only) that is used for the + rest of the conversation. Read-only tools then run against every scoped project in a single call; + write operations target the active (first) project only. Call this when the user states which + projects to work on; it can be called again any time to re-scope. + """ + client = KeboolaClient.from_state(ctx.session.state) + parent_token = await _parent_subject_token(client) + + ids = list(project_ids or []) + if not ids: + introspection = await introspect_token(client.storage_api_url, subject_token=parent_token) + ids = [p.id for p in introspection.projects] + if not ids: + raise ValueError('No accessible projects to scope to.') + + # Mint a token narrowed to the chosen projects. If the exchange endpoint is unavailable on the + # stack, fall back to the whole-stack parent token (still narrowed per request by X-KBC-ProjectId) + # so scoping/fan-out keeps working without the security narrowing. + try: + minted = await exchange_scoped_token( + client.storage_api_url, subject_token=parent_token, project_ids=ids, read_only=read_only + ) + scope = SessionScope( + project_ids=ids, + read_only=minted.read_only, + scoped_token=minted.access_token, + scoped_expires_at=minted.expires_at, + confirmed=True, + ) + except Exception as e: + LOG.warning(f'Scoped-token exchange failed ({e}); scoping with the whole-stack token instead.') + scope = SessionScope(project_ids=ids, read_only=read_only, confirmed=True) + ctx.session.state[SCOPE_KEY] = scope + + multi = len(ids) > 1 + return ProjectScope( + project_ids=ids, + read_only=minted.read_only, + active_project_id=scope.active_project_id, + llm_instruction=( + ( + f'Session scoped to {len(ids)} projects. Read-only tools now return results per project. ' + f'Write operations target the active project {scope.active_project_id} only; to write to a ' + 'different project, re-scope or confirm with the user first.' + ) + if multi + else f'Session scoped to project {ids[0]}.' + ), + ) diff --git a/tests/test_auth_login.py b/tests/test_auth_login.py index e22d761d5..6ad8eb136 100644 --- a/tests/test_auth_login.py +++ b/tests/test_auth_login.py @@ -2,6 +2,7 @@ import base64 import hashlib +import json import stat import time from pathlib import Path @@ -13,7 +14,9 @@ from keboola_mcp_server.auth_login import ( TokenSet, exchange_code, + exchange_scoped_token, get_access_token, + introspect_token, load_tokens, refresh_tokens, save_tokens, @@ -132,3 +135,61 @@ def test_pkce_challenge_is_sha256_of_verifier() -> None: def test_invalid_stack_url_rejected() -> None: with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): auth_login._base_url('https://example.com') + + +# --- introspection + scoped exchange (PSGO-261 increment 2) --- + + +@pytest.mark.asyncio +async def test_introspect_token_parses_projects() -> None: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured['url'] = str(request.url) + captured['auth'] = request.headers['Authorization'] + return httpx.Response( + 200, + json={ + 'user': {'id': 60, 'email': 'm@k.com', 'name': 'M'}, + 'projects': [ + {'id': 18, 'name': 'A', 'role': 'admin'}, + {'id': 83, 'name': 'B', 'role': 'admin'}, + ], + }, + ) + + intro = await introspect_token(STACK, subject_token='kbc_at_x', transport=httpx.MockTransport(handler)) + + assert captured['url'] == 'https://connection.keboola.com/v1/auth/token/introspect' + assert captured['auth'] == 'Bearer kbc_at_x' + assert intro.user_email == 'm@k.com' + assert [(p.id, p.name, p.role) for p in intro.projects] == [(18, 'A', 'admin'), (83, 'B', 'admin')] + + +@pytest.mark.asyncio +async def test_exchange_scoped_token_sends_scope_and_parses() -> None: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured['url'] = str(request.url) + captured['auth'] = request.headers['Authorization'] + captured['body'] = json.loads(request.content) + return httpx.Response(201, json={'accessToken': 'kbc_at_scoped', 'expiresIn': 3600, 'readOnly': True}) + + scoped = await exchange_scoped_token( + STACK, + subject_token='kbc_at_parent', + project_ids=[18, 83], + read_only=True, + transport=httpx.MockTransport(handler), + ) + + assert captured['url'] == 'https://connection.keboola.com/v1/auth/pat/exchange' + assert captured['auth'] == 'Bearer kbc_at_parent' + # the exchange API requires project ids as strings + assert captured['body'] == {'expiresIn': None, 'scope': {'projects': ['18', '83'], 'readOnly': True}} + assert scoped.access_token == 'kbc_at_scoped' + assert scoped.read_only is True + assert scoped.project_ids == [18, 83] + assert scoped.expires_at > time.time() + assert not scoped.is_near_expiry diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 63736235b..c819a1310 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,5 @@ import asyncio +import time from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -6,14 +7,19 @@ import pytest from fastmcp import Context from fastmcp.exceptions import ToolError +from fastmcp.tools.tool import ToolResult +from mcp import types as mt from pydantic import BaseModel, Field from starlette.requests import Request from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config, ServerRuntimeInfo from keboola_mcp_server.mcp import ( + SCOPE_KEY, AggregateError, + MultiProjectMiddleware, ServerState, + SessionScope, SessionStateMiddleware, ToolsFilteringMiddleware, _exclude_none_serializer, @@ -840,3 +846,208 @@ async def test_happy_path_calls_resolver(self, monkeypatch) -> None: storage_api_url='https://connection.keboola.com', kubernetes_token_path='/var/run/secrets/token' ) resolver.resolve.assert_awaited_once_with(subject_token='kbc_at_abc', project_id=42) + + +class TestResolveLocalTokens: + """SessionStateMiddleware keeps local tokens fresh and re-mints the scoped token (PSGO-261).""" + + @pytest.mark.asyncio + async def test_deployed_is_noop(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x') + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, None) + assert out_config is config + assert out_scope is None + + @pytest.mark.asyncio + async def test_legacy_token_is_noop(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='legacy-sapi-token') + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, None) + assert out_config is config + assert out_scope is None + + @pytest.mark.asyncio + async def test_programmatic_no_scope_refreshes_parent(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_old') + with patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_fresh')): + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, None) + assert out_config.storage_token == 'kbc_at_fresh' + assert out_scope is None + + @pytest.mark.asyncio + async def test_default_scope_uses_parent_token_without_minting(self, monkeypatch) -> None: + # The default (auto-leased) multi-project scope carries no minted token: it uses the parent + # token and just sets the active project — no exchange call. + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_old') + scope = SessionScope(project_ids=[11, 22], read_only=False) + with ( + patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_parent')), + patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock()) as exch, + ): + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + exch.assert_not_awaited() + assert out_config.storage_token == 'kbc_at_parent' + assert out_config.project_id == '11' # active project = first in scope + assert out_scope.scoped_token is None + + @pytest.mark.asyncio + async def test_fresh_scoped_token_is_not_reminted(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_old') + scope = SessionScope(project_ids=[11], scoped_token='kbc_at_live', scoped_expires_at=time.time() + 3600) + with ( + patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_parent')), + patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock()) as exch, + ): + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + exch.assert_not_awaited() + assert out_config.storage_token == 'kbc_at_live' + + @pytest.mark.asyncio + async def test_near_expiry_scoped_token_is_reminted(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_old') + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_stale', scoped_expires_at=time.time() - 1) + minted = SimpleNamespace(access_token='kbc_at_fresh_scoped', expires_at=time.time() + 900) + with ( + patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_parent')), + patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock(return_value=minted)) as exch, + ): + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + exch.assert_awaited_once() + assert out_scope.scoped_token == 'kbc_at_fresh_scoped' + assert out_config.storage_token == 'kbc_at_fresh_scoped' + + @pytest.mark.asyncio + async def test_autolease_scopes_all_accessible_projects(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x') + introspection = SimpleNamespace( + projects=[SimpleNamespace(id=11), SimpleNamespace(id=22), SimpleNamespace(id=33)] + ) + with ( + patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_parent')), + patch('keboola_mcp_server.mcp.introspect_token', AsyncMock(return_value=introspection)), + ): + scope = await SessionStateMiddleware._autolease_default_scope(config) + assert scope.project_ids == [11, 22, 33] + assert scope.scoped_token is None # default scope uses the parent token + + @pytest.mark.asyncio + async def test_autolease_noop_when_deployed(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x') + assert await SessionStateMiddleware._autolease_default_scope(config) is None + + +class TestMultiProjectMiddleware: + """Read tools fan out across the scoped projects; writes and single-project scope do not.""" + + @staticmethod + def _ctx(scope: SessionScope | None, tool_name: str, read_only: bool): + state: dict = {KeboolaClient.STATE_KEY: 'orig-client'} + if scope is not None: + state[SCOPE_KEY] = scope + ctx = MagicMock(spec=Context) + ctx.session = SimpleNamespace(state=state) + ctx.request_context.lifespan_context = ServerState( + config=Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x'), + runtime_info=ServerRuntimeInfo(transport='stdio'), + ) + tool = MagicMock() + tool.name = tool_name + if read_only: + tool.annotations.readOnlyHint = True + else: + tool.annotations = None + ctx.fastmcp.get_tool = AsyncMock(return_value=tool) + context = SimpleNamespace(message=SimpleNamespace(name=tool_name), fastmcp_context=ctx) + return context, state + + @staticmethod + def _result(text: str) -> ToolResult: + return ToolResult( + content=[mt.TextContent(type='text', text=text)], + structured_content={'rows': [text]}, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ('scope', 'tool_name', 'read_only'), + [ + (None, 'get_tables', True), + (SessionScope(project_ids=[11], confirmed=True), 'get_tables', True), + (SessionScope(project_ids=[11, 22], confirmed=True), 'update_config', False), # write: no fan-out + (SessionScope(project_ids=[11, 22], confirmed=True), 'query_data', True), # excluded tool + ], + ids=['no_scope', 'single_project', 'write_tool', 'excluded_tool'], + ) + async def test_passthrough_calls_once(self, scope, tool_name, read_only) -> None: + context, _ = self._ctx(scope, tool_name, read_only) + calls = [] + + async def call_next(_): + calls.append(1) + return self._result('single') + + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + assert len(calls) == 1 + assert result.content[0].text == 'single' + + @pytest.mark.asyncio + async def test_unconfirmed_scope_blocks_data_tools(self) -> None: + # Default (auto-leased, unconfirmed) scope: data tools are gated with an ask-first message. + scope = SessionScope(project_ids=[11, 22], confirmed=False) + context, _ = self._ctx(scope, 'get_tables', read_only=True) + + async def call_next(_): + raise AssertionError('call_next must not run for a gated tool') + + with pytest.raises(ToolError, match='no scope has been confirmed'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + @pytest.mark.asyncio + async def test_unconfirmed_scope_allows_bootstrap_tools(self) -> None: + scope = SessionScope(project_ids=[11, 22], confirmed=False) + context, _ = self._ctx(scope, 'get_accessible_projects', read_only=True) + calls = [] + + async def call_next(_): + calls.append(1) + return self._result('projects') + + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + assert len(calls) == 1 + assert result.content[0].text == 'projects' + + @pytest.mark.asyncio + async def test_read_tool_fans_out_per_project(self) -> None: + scope = SessionScope( + project_ids=[11, 22], scoped_token='kbc_at_s', scoped_expires_at=time.time() + 3600, confirmed=True + ) + context, state = self._ctx(scope, 'get_tables', read_only=True) + active_clients: list = [] + + async def call_next(_): + active_clients.append(state[KeboolaClient.STATE_KEY]) + return self._result('rows') + + with patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + # Ran once per project, each against that project's client. + assert active_clients == ['client-11', 'client-22'] + # Active client restored afterwards. + assert state[KeboolaClient.STATE_KEY] == 'orig-client' + # Per-project results are labelled in the text content. + texts = [c.text for c in result.content] + assert texts == ['=== project 11 ===', 'rows', '=== project 22 ===', 'rows'] + # Structured output is deep-merged (list fields concatenated) so it still validates the schema. + assert result.structured_content == {'rows': ['rows', 'rows']} diff --git a/tests/test_server.py b/tests/test_server.py index 02b025a64..fced0c834 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -60,6 +60,7 @@ async def test_list_tools(self): 'deploy_data_app', 'docs_query', 'find_component_id', + 'get_accessible_projects', 'get_buckets', 'get_components', 'get_config_examples', @@ -83,6 +84,7 @@ async def test_list_tools(self): 'run_sync_action', 'search', 'search_semantic_context', + 'set_project_scope', 'update_config', 'update_config_row', 'update_descriptions', @@ -150,7 +152,7 @@ async def test_tools_input_schema(self): missing_default.append(f'{tool.name}.{prop_name}') missing_properties.sort() - assert missing_properties == ['get_project_info'] + assert missing_properties == ['get_accessible_projects', 'get_project_info'] missing_type.sort() assert not missing_type, f'These tool params have no "type" info: {missing_type}' missing_default.sort() diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 3e0ea2497..ea3c4e799 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -1,3 +1,6 @@ +import time +from types import SimpleNamespace + import pytest from mcp.server.fastmcp import Context from pytest_mock import MockerFixture @@ -5,15 +8,20 @@ from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.links import Link +from keboola_mcp_server.mcp import SCOPE_KEY from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, _resolve_branch_context, + get_accessible_projects, get_project_info, + set_project_scope, update_project_description, ) from keboola_mcp_server.workspace import WorkspaceManager +STACK = 'https://connection.test.keboola.com' + @pytest.mark.parametrize( ('role', 'expected_substring', 'expect_none'), @@ -246,3 +254,83 @@ async def test_update_project_description( keboola_client.storage_client.branch_metadata_update.assert_called_once_with( {MetadataField.PROJECT_DESCRIPTION: description} ) + + +# --- multi-project scope tools (PSGO-261 increment 2) --- + + +def _prep_client(mcp_context_client: Context, mocker: MockerFixture, *, bearer: str | None = 'kbc_at_parent'): + client = KeboolaClient.from_state(mcp_context_client.session.state) + client.bearer_token = bearer + client.storage_api_url = STACK + mocker.patch( + 'keboola_mcp_server.tools.project.get_access_token', + new=mocker.AsyncMock(return_value='kbc_at_parent'), + ) + return client + + +@pytest.mark.asyncio +async def test_get_accessible_projects(mcp_context_client: Context, mocker: MockerFixture) -> None: + _prep_client(mcp_context_client, mocker) + introspection = SimpleNamespace( + user_email='m@k.com', + projects=[SimpleNamespace(id=18, name='A', role='admin'), SimpleNamespace(id=83, name='B', role='admin')], + ) + introspect = mocker.patch( + 'keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection) + ) + + result = await get_accessible_projects(mcp_context_client) + + introspect.assert_awaited_once_with(STACK, subject_token='kbc_at_parent') + assert result.user_email == 'm@k.com' + assert [(p.id, p.name, p.role) for p in result.projects] == [(18, 'A', 'admin'), (83, 'B', 'admin')] + + +@pytest.mark.asyncio +async def test_set_project_scope_subset_exchanges_and_stores( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + _prep_client(mcp_context_client, mocker) + minted = SimpleNamespace(access_token='kbc_at_scoped', expires_at=time.time() + 3600, read_only=False) + exch = mocker.patch( + 'keboola_mcp_server.tools.project.exchange_scoped_token', new=mocker.AsyncMock(return_value=minted) + ) + + result = await set_project_scope(mcp_context_client, project_ids=[18, 83]) + + exch.assert_awaited_once_with(STACK, subject_token='kbc_at_parent', project_ids=[18, 83], read_only=False) + assert result.project_ids == [18, 83] + assert result.active_project_id == 18 + scope = mcp_context_client.session.state[SCOPE_KEY] + assert scope.scoped_token == 'kbc_at_scoped' + assert scope.project_ids == [18, 83] + + +@pytest.mark.asyncio +async def test_set_project_scope_all_introspects_then_exchanges( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + _prep_client(mcp_context_client, mocker) + introspection = SimpleNamespace( + user_email=None, + projects=[SimpleNamespace(id=18, name='A', role='admin'), SimpleNamespace(id=83, name='B', role='x')], + ) + mocker.patch('keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection)) + minted = SimpleNamespace(access_token='kbc_at_all', expires_at=time.time() + 3600, read_only=False) + exch = mocker.patch( + 'keboola_mcp_server.tools.project.exchange_scoped_token', new=mocker.AsyncMock(return_value=minted) + ) + + result = await set_project_scope(mcp_context_client, project_ids=None) + + exch.assert_awaited_once_with(STACK, subject_token='kbc_at_parent', project_ids=[18, 83], read_only=False) + assert result.project_ids == [18, 83] + + +@pytest.mark.asyncio +async def test_scope_requires_programmatic_token(mcp_context_client: Context, mocker: MockerFixture) -> None: + _prep_client(mcp_context_client, mocker, bearer=None) + with pytest.raises(ValueError, match='programmatic token'): + await get_accessible_projects(mcp_context_client) From 9dc5550fbb33dc4b7cb8a2f01da7270345ccb486 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 07:10:26 +0200 Subject: [PATCH 05/89] feat(PSGO-261): per-call project filter for read tools Add an optional project_ids argument to fan-out-eligible read tools so a single call can target a subset of the scoped projects without changing the session scope. The argument is advertised at list time (on_list_tools) only while a multi-project scope is active, and consumed + stripped in on_call_tool before the tool runs. - one target -> single call against that project's X-KBC-ProjectId (no per-project envelope); a subset -> fan out over just those projects - out-of-scope project ids are rejected with a clear error - bootstrap tools keep their own project_ids arg (not stripped); write and workspace-bound tools ignore the filter and target the active project Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/mcp.py | 91 ++++++++++++++++++++++++++++++++--- tests/test_mcp.py | 83 +++++++++++++++++++++++++++++++- 2 files changed, 165 insertions(+), 9 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 31c38c8f5..5e30a730e 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -57,6 +57,10 @@ # message telling the assistant to ask the user which projects to work on first (ask-first UX). _BOOTSTRAP_TOOLS = {'get_accessible_projects', 'set_project_scope'} +# Optional per-call argument injected on fan-out-eligible read tools to restrict a single call to a +# subset of the scoped projects (consumed and stripped by MultiProjectMiddleware.on_call_tool). +_PROJECT_FILTER_ARG = 'project_ids' + @dataclasses.dataclass(frozen=True) class SessionScope: @@ -802,11 +806,12 @@ async def on_call_tool( ctx = context.fastmcp_context state = ctx.session.state scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None + name = context.message.name # Ask-first gate: until the user confirms a scope via set_project_scope, block data tools and # tell the assistant to ask the user which projects to work on. Only applies when a scope has # been auto-leased (local programmatic session); deployed/legacy sessions have no scope. - if isinstance(scope, SessionScope) and not scope.confirmed and context.message.name not in _BOOTSTRAP_TOOLS: + if isinstance(scope, SessionScope) and not scope.confirmed and name not in _BOOTSTRAP_TOOLS: raise ToolError( f'This session can access {len(scope.project_ids)} Keboola project(s), but no scope has ' 'been confirmed yet. Call "get_accessible_projects", show the user their projects, and ask ' @@ -815,23 +820,58 @@ async def on_call_tool( 'This confirmation is required once per session.' ) - if not isinstance(scope, SessionScope) or len(scope.project_ids) <= 1: + # No auto-leased scope (deployed / legacy) or a bootstrap/scope tool: pass through untouched. + # Bootstrap tools own a real `project_ids` argument, so we must not strip it. + if not isinstance(scope, SessionScope) or name in _BOOTSTRAP_TOOLS: return await call_next(context) - if context.message.name in _NO_FANOUT_TOOLS: + # Workspace-bound and write tools always target the active project (no fan-out, filter ignored). + if name in _NO_FANOUT_TOOLS: return await call_next(context) - - tool = await ctx.fastmcp.get_tool(context.message.name) + tool = await ctx.fastmcp.get_tool(name) if not is_read_only_tool(tool): - # Write tools target the active project only — never an automatic multi-project write. + return await call_next(context) + + # Read tool: consume the optional per-call project filter (advertised via on_list_tools) so the + # tool never receives it, then narrow this call's target projects to the requested subset. + requested = None + args = getattr(context.message, 'arguments', None) + if isinstance(args, dict): + requested = args.pop(_PROJECT_FILTER_ARG, None) + + targets = list(scope.project_ids) + if requested: + outside = [p for p in requested if p not in scope.project_ids] + if outside: + raise ToolError( + f'Project(s) {outside} are outside the current scope {scope.project_ids}. ' + 'Call "set_project_scope" to change the scope first.' + ) + targets = [p for p in scope.project_ids if p in requested] + if not targets: return await call_next(context) server_state = ServerState.from_context(ctx) original_client = state.get(KeboolaClient.STATE_KEY) # Default (auto-leased) scope carries no minted token; fall back to the active client's token. base_token = scope.scoped_token or (original_client.token if isinstance(original_client, KeboolaClient) else '') + + # A single target (scope of one, or narrowed to one via the filter) runs once against that + # project only — one call, that project's X-KBC-ProjectId, no per-project envelope. + if len(targets) == 1: + target = targets[0] + if target == scope.active_project_id: + return await call_next(context) + try: + state[KeboolaClient.STATE_KEY] = await self._client_for_project( + server_state, base_token, target, scope.read_only + ) + return await call_next(context) + finally: + state[KeboolaClient.STATE_KEY] = original_client + results: list[tuple[int, ToolResult]] = [] try: - for project_id in scope.project_ids: + for project_id in targets: state[KeboolaClient.STATE_KEY] = await self._client_for_project( server_state, base_token, project_id, scope.read_only ) @@ -841,6 +881,43 @@ async def on_call_tool( return self._merge(results) + async def on_list_tools( + self, + context: MiddlewareContext[mt.ListToolsRequest], + call_next: CallNext[mt.ListToolsRequest, list[Tool]], + ) -> list[Tool]: + # Advertise the optional per-call `project_ids` filter on fan-out-eligible read tools while a + # multi-project scope is active, so the assistant can target a subset (e.g. a single project) + # without changing the session scope. The value is consumed and stripped in on_call_tool. + tools = await call_next(context) + ctx = context.fastmcp_context + state = getattr(ctx.session, 'state', None) + scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None + if not (isinstance(scope, SessionScope) and len(scope.project_ids) > 1): + return tools + + patched: list[Tool] = [] + for tool in tools: + if tool.name in _BOOTSTRAP_TOOLS or tool.name in _NO_FANOUT_TOOLS or not is_read_only_tool(tool): + patched.append(tool) + continue + params = dict(tool.parameters or {}) + props = dict(params.get('properties') or {}) + if _PROJECT_FILTER_ARG in props: + patched.append(tool) + continue + props[_PROJECT_FILTER_ARG] = { + 'type': 'array', + 'items': {'type': 'integer'}, + 'description': ( + 'Optional. Restrict this call to these project ids (a subset of the confirmed ' + 'multi-project scope). Omit to run across all scoped projects.' + ), + } + params['properties'] = props + patched.append(tool.model_copy(update={'parameters': params})) + return patched + @staticmethod async def _client_for_project( server_state: ServerState, token: str, project_id: int, read_only: bool diff --git a/tests/test_mcp.py b/tests/test_mcp.py index c819a1310..e72f4fbbc 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -947,7 +947,7 @@ class TestMultiProjectMiddleware: """Read tools fan out across the scoped projects; writes and single-project scope do not.""" @staticmethod - def _ctx(scope: SessionScope | None, tool_name: str, read_only: bool): + def _ctx(scope: SessionScope | None, tool_name: str, read_only: bool, arguments: dict | None = None): state: dict = {KeboolaClient.STATE_KEY: 'orig-client'} if scope is not None: state[SCOPE_KEY] = scope @@ -964,7 +964,8 @@ def _ctx(scope: SessionScope | None, tool_name: str, read_only: bool): else: tool.annotations = None ctx.fastmcp.get_tool = AsyncMock(return_value=tool) - context = SimpleNamespace(message=SimpleNamespace(name=tool_name), fastmcp_context=ctx) + message = SimpleNamespace(name=tool_name, arguments=arguments if arguments is not None else {}) + context = SimpleNamespace(message=message, fastmcp_context=ctx) return context, state @staticmethod @@ -1051,3 +1052,81 @@ async def call_next(_): assert texts == ['=== project 11 ===', 'rows', '=== project 22 ===', 'rows'] # Structured output is deep-merged (list fields concatenated) so it still validates the schema. assert result.structured_content == {'rows': ['rows', 'rows']} + + @pytest.mark.asyncio + async def test_project_filter_single_target_runs_once(self) -> None: + # project_ids filter narrows a multi-project scope to one project: one call, that project's + # client, and the filter is stripped from the arguments the tool receives. + scope = SessionScope(project_ids=[11, 22, 33], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [22]}) + seen_clients: list = [] + + async def call_next(_): + seen_clients.append(state[KeboolaClient.STATE_KEY]) + return self._result('t') + + with patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_clients == ['client-22'] # ran once, against project 22 only + assert state[KeboolaClient.STATE_KEY] == 'orig-client' # restored + assert 'project_ids' not in context.message.arguments # stripped before the tool + assert result.content[0].text == 't' # raw single-project result, not an envelope + + @pytest.mark.asyncio + async def test_project_filter_subset_fans_out(self) -> None: + scope = SessionScope(project_ids=[11, 22, 33], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [11, 33]}) + seen_clients: list = [] + + async def call_next(_): + seen_clients.append(state[KeboolaClient.STATE_KEY]) + return self._result('t') + + with patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_clients == ['client-11', 'client-33'] # only the requested subset, in scope order + + @pytest.mark.asyncio + async def test_project_filter_outside_scope_raises(self) -> None: + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, _ = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [99]}) + + async def call_next(_): + raise AssertionError('must not run for an out-of-scope filter') + + with pytest.raises(ToolError, match='outside the current scope'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + @pytest.mark.asyncio + async def test_on_list_tools_injects_project_filter(self) -> None: + scope = SessionScope(project_ids=[11, 22], confirmed=True) + # a read fan-out tool, an excluded tool, and a write tool + read_tool = _tool('get_tables', read_only=True) + read_tool.parameters = {'type': 'object', 'properties': {'bucket_ids': {'type': 'array'}}} + excluded = _tool('query_data', read_only=True) + excluded.parameters = {'type': 'object', 'properties': {}} + write_tool = _tool('update_config', read_only=False) + write_tool.parameters = {'type': 'object', 'properties': {}} + for t in (read_tool, excluded, write_tool): + t.model_copy = lambda update, _t=t: SimpleNamespace(name=_t.name, parameters=update['parameters']) + + context, _ = self._ctx(scope, 'x', read_only=True) + + async def call_next(_): + return [read_tool, excluded, write_tool] + + tools = await MultiProjectMiddleware().on_list_tools(context, call_next) + by_name = {t.name: t for t in tools} + assert 'project_ids' in by_name['get_tables'].parameters['properties'] + assert 'project_ids' not in by_name['query_data'].parameters['properties'] + assert 'project_ids' not in by_name['update_config'].parameters['properties'] From 546cac0c4c08c71bfde1057c69f1f4ad9cf2931a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 07:11:14 +0200 Subject: [PATCH 06/89] feat(PSGO-261): surface current scope in get_accessible_projects Report which projects are in the confirmed scope and which is active, plus top-level scoped_project_ids / active_project_id / read_only, and tailor the llm_instruction to whether a scope has been confirmed yet. Lets the assistant answer "which project am I on / scoped to" from one call. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/tools/project.py | 48 +++++++++++++++++++++---- tests/tools/test_project.py | 15 +++++++- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 3f851d166..a9222330a 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -269,11 +269,23 @@ class AccessibleProject(BaseModel): id: int = Field(description='The project id.') name: str | None = Field(default=None, description='The project name.') role: str | None = Field(default=None, description='The user role in this project (e.g. "admin").') + in_scope: bool = Field(default=False, description='Whether the session is currently scoped to this project.') + is_active: bool = Field( + default=False, description='Whether this is the active project (write / single-project tool target).' + ) class AccessibleProjects(BaseModel): user_email: str | None = Field(default=None, description='The email of the authenticated user.') projects: list[AccessibleProject] = Field(description='The projects the current token can reach across the stack.') + scoped_project_ids: list[int] | None = Field( + default=None, + description='The projects the session is currently scoped to, or null if no scope has been confirmed yet.', + ) + active_project_id: int | None = Field( + default=None, description='The active project that write operations and single-project tools target.' + ) + read_only: bool | None = Field(default=None, description='Whether the current scoped token is read-only.') llm_instruction: str = Field( description='Guidance for the assistant on how to use this result.', ) @@ -301,15 +313,39 @@ async def get_accessible_projects(ctx: Context) -> AccessibleProjects: client = KeboolaClient.from_state(ctx.session.state) subject_token = await _parent_subject_token(client) introspection = await introspect_token(client.storage_api_url, subject_token=subject_token) - projects = [AccessibleProject(id=p.id, name=p.name, role=p.role) for p in introspection.projects] + + scope = ctx.session.state.get(SCOPE_KEY) + scoped_ids = scope.project_ids if isinstance(scope, SessionScope) and scope.confirmed else None + active_id = scope.active_project_id if scoped_ids else None + + projects = [ + AccessibleProject( + id=p.id, + name=p.name, + role=p.role, + in_scope=scoped_ids is not None and p.id in scoped_ids, + is_active=p.id == active_id, + ) + for p in introspection.projects + ] + if scoped_ids is None: + instruction = ( + 'No project scope has been confirmed yet. Ask the user whether to operate across all these ' + 'projects or a subset, then call "set_project_scope" with the chosen project ids. Never write ' + 'to more than one project without explicit user confirmation.' + ) + else: + instruction = ( + f'Session is currently scoped to {len(scoped_ids)} project(s); the active project is {active_id}. ' + 'Call "set_project_scope" to change the scope.' + ) return AccessibleProjects( user_email=introspection.user_email, projects=projects, - llm_instruction=( - 'Ask the user whether to operate across all these projects or a subset, then call ' - '"set_project_scope" with the chosen project ids. Never write to more than one project ' - 'without explicit user confirmation.' - ), + scoped_project_ids=scoped_ids, + active_project_id=active_id, + read_only=scope.read_only if scoped_ids is not None else None, + llm_instruction=instruction, ) diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index ea3c4e799..2e0e9cb23 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -8,7 +8,7 @@ from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.links import Link -from keboola_mcp_server.mcp import SCOPE_KEY +from keboola_mcp_server.mcp import SCOPE_KEY, SessionScope from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, @@ -281,11 +281,24 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock 'keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection) ) + # No scope confirmed yet. result = await get_accessible_projects(mcp_context_client) introspect.assert_awaited_once_with(STACK, subject_token='kbc_at_parent') assert result.user_email == 'm@k.com' assert [(p.id, p.name, p.role) for p in result.projects] == [(18, 'A', 'admin'), (83, 'B', 'admin')] + assert result.scoped_project_ids is None + assert result.active_project_id is None + assert result.read_only is None + assert all(not p.in_scope and not p.is_active for p in result.projects) + + # Once scoped, the current scope is surfaced on the projects and at the top level. + mcp_context_client.session.state[SCOPE_KEY] = SessionScope(project_ids=[83], read_only=True, confirmed=True) + result = await get_accessible_projects(mcp_context_client) + assert result.scoped_project_ids == [83] + assert result.active_project_id == 83 + assert result.read_only is True + assert [(p.id, p.in_scope, p.is_active) for p in result.projects] == [(18, False, False), (83, True, True)] @pytest.mark.asyncio From 5200b63ccfcb5e9c6b577ce0cb15a31b19317152 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 07:28:04 +0200 Subject: [PATCH 07/89] feat(PSGO-261): fan out query_data across scoped projects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit query_data was pinned to the active project's workspace, so a user could not query tables from any other scoped project without re-scoping the session. Make query_data a normal fan-out read tool: MultiProjectMiddleware._swap_project now swaps both the KeboolaClient AND a per-project WorkspaceManager into session state for the duration of a fanned-out call, so the SQL runs in the targeted project's own workspace. Narrow to one project with the project_ids filter, or run across all scoped projects. Cross-project / cross-backend SQL in a single statement remains unsupported (BigQuery has no cross-project access; Snowflake only via materialized linked aliases) — accepted edge case, to be documented in the MPA RFC. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/mcp.py | 44 ++++++++++++++------ tests/test_mcp.py | 76 +++++++++++++++++++++++++++-------- 2 files changed, 91 insertions(+), 29 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 5e30a730e..0aae46104 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -46,12 +46,12 @@ # Tools that must not be fanned out across multiple projects, even when a multi-project scope is # active and they are read-only: the scope/auth tools operate on the whole-stack token (not a single -# project), and query_data resolves through a per-project WorkspaceManager that the fan-out does not -# swap (it would otherwise run against the active project's workspace N times). -# ponytail: explicit deny-set; per-project WorkspaceManager fan-out is a follow-up, not v1. -# get_project_info and query_data resolve through the active project's WorkspaceManager (workspace id -# / sql dialect), which the fan-out does not swap, so they report the active project only. -_NO_FANOUT_TOOLS = {'get_accessible_projects', 'set_project_scope', 'get_project_info', 'query_data'} +# project), and get_project_info resolves through the active project's WorkspaceManager (workspace id +# / sql dialect), so it reports the active project only. +# query_data is intentionally NOT here: the fan-out swaps a per-project WorkspaceManager (see +# MultiProjectMiddleware._swap_project) so a query runs against the workspace of each targeted +# project — narrow to one with the project_ids filter, or run across all scoped projects. +_NO_FANOUT_TOOLS = {'get_accessible_projects', 'set_project_scope', 'get_project_info'} # Tools allowed before the user has confirmed a project scope. Everything else is blocked with a # message telling the assistant to ask the user which projects to work on first (ask-first UX). @@ -852,6 +852,7 @@ async def on_call_tool( server_state = ServerState.from_context(ctx) original_client = state.get(KeboolaClient.STATE_KEY) + original_workspace = state.get(WorkspaceManager.STATE_KEY) # Default (auto-leased) scope carries no minted token; fall back to the active client's token. base_token = scope.scoped_token or (original_client.token if isinstance(original_client, KeboolaClient) else '') @@ -862,22 +863,20 @@ async def on_call_tool( if target == scope.active_project_id: return await call_next(context) try: - state[KeboolaClient.STATE_KEY] = await self._client_for_project( - server_state, base_token, target, scope.read_only - ) + await self._swap_project(state, server_state, base_token, target, scope.read_only) return await call_next(context) finally: state[KeboolaClient.STATE_KEY] = original_client + state[WorkspaceManager.STATE_KEY] = original_workspace results: list[tuple[int, ToolResult]] = [] try: for project_id in targets: - state[KeboolaClient.STATE_KEY] = await self._client_for_project( - server_state, base_token, project_id, scope.read_only - ) + await self._swap_project(state, server_state, base_token, project_id, scope.read_only) results.append((project_id, await call_next(context))) finally: state[KeboolaClient.STATE_KEY] = original_client + state[WorkspaceManager.STATE_KEY] = original_workspace return self._merge(results) @@ -918,6 +917,27 @@ async def on_list_tools( patched.append(tool.model_copy(update={'parameters': params})) return patched + @classmethod + async def _swap_project( + cls, + state: dict[str, Any], + server_state: ServerState, + base_token: str, + project_id: int, + read_only: bool, + ) -> None: + """Points the session state at `project_id` for the duration of one fanned-out tool call. + + Swaps in a per-project `KeboolaClient` AND a `WorkspaceManager` built on it, so + workspace-bound reads (query_data) run against *this* project's workspace rather than the + active project's. The workspace is provisioned lazily on first use per project. + ponytail: rebuilt per call; caching across calls would need a store that survives the + per-request state rebuild — add if provisioning latency shows up in practice. + """ + client = await cls._client_for_project(server_state, base_token, project_id, read_only) + state[KeboolaClient.STATE_KEY] = client + state[WorkspaceManager.STATE_KEY] = await WorkspaceManager.create(client, server_state.config.workspace_schema) + @staticmethod async def _client_for_project( server_state: ServerState, token: str, project_id: int, read_only: bool diff --git a/tests/test_mcp.py b/tests/test_mcp.py index e72f4fbbc..5f89b6e0d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -28,6 +28,7 @@ toon_serializer, unwrap_results, ) +from keboola_mcp_server.workspace import WorkspaceManager class SimpleModel(BaseModel): @@ -982,7 +983,7 @@ def _result(text: str) -> ToolResult: (None, 'get_tables', True), (SessionScope(project_ids=[11], confirmed=True), 'get_tables', True), (SessionScope(project_ids=[11, 22], confirmed=True), 'update_config', False), # write: no fan-out - (SessionScope(project_ids=[11, 22], confirmed=True), 'query_data', True), # excluded tool + (SessionScope(project_ids=[11, 22], confirmed=True), 'get_project_info', True), # excluded tool ], ids=['no_scope', 'single_project', 'write_tool', 'excluded_tool'], ) @@ -1031,28 +1032,60 @@ async def test_read_tool_fans_out_per_project(self) -> None: ) context, state = self._ctx(scope, 'get_tables', read_only=True) active_clients: list = [] + active_workspaces: list = [] async def call_next(_): active_clients.append(state[KeboolaClient.STATE_KEY]) + active_workspaces.append(state[WorkspaceManager.STATE_KEY]) return self._result('rows') - with patch.object( - MultiProjectMiddleware, - '_client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + with ( + patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): result = await MultiProjectMiddleware().on_call_tool(context, call_next) - # Ran once per project, each against that project's client. + # Ran once per project, each against that project's client AND workspace. assert active_clients == ['client-11', 'client-22'] - # Active client restored afterwards. + assert active_workspaces == ['wsm-client-11', 'wsm-client-22'] + # Active client and workspace restored afterwards. assert state[KeboolaClient.STATE_KEY] == 'orig-client' + assert state.get(WorkspaceManager.STATE_KEY) is None # Per-project results are labelled in the text content. texts = [c.text for c in result.content] assert texts == ['=== project 11 ===', 'rows', '=== project 22 ===', 'rows'] # Structured output is deep-merged (list fields concatenated) so it still validates the schema. assert result.structured_content == {'rows': ['rows', 'rows']} + @pytest.mark.asyncio + async def test_query_data_targets_single_project_workspace(self) -> None: + # query_data is no longer excluded: with the project_ids filter it runs once against that + # project's own workspace, so the user can query any scoped project without re-scoping. + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'query_data', read_only=True, arguments={'project_ids': [22]}) + seen_workspaces: list = [] + + async def call_next(_): + seen_workspaces.append(state[WorkspaceManager.STATE_KEY]) + return self._result('csv') + + with ( + patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_workspaces == ['wsm-client-22'] # ran against project 22's workspace + assert result.content[0].text == 'csv' + @pytest.mark.asyncio async def test_project_filter_single_target_runs_once(self) -> None: # project_ids filter narrows a multi-project scope to one project: one call, that project's @@ -1060,19 +1093,25 @@ async def test_project_filter_single_target_runs_once(self) -> None: scope = SessionScope(project_ids=[11, 22, 33], scoped_token='kbc_at_s', confirmed=True) context, state = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [22]}) seen_clients: list = [] + seen_workspaces: list = [] async def call_next(_): seen_clients.append(state[KeboolaClient.STATE_KEY]) + seen_workspaces.append(state[WorkspaceManager.STATE_KEY]) return self._result('t') - with patch.object( - MultiProjectMiddleware, - '_client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + with ( + patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): result = await MultiProjectMiddleware().on_call_tool(context, call_next) assert seen_clients == ['client-22'] # ran once, against project 22 only + assert seen_workspaces == ['wsm-client-22'] # its own workspace assert state[KeboolaClient.STATE_KEY] == 'orig-client' # restored assert 'project_ids' not in context.message.arguments # stripped before the tool assert result.content[0].text == 't' # raw single-project result, not an envelope @@ -1087,10 +1126,13 @@ async def call_next(_): seen_clients.append(state[KeboolaClient.STATE_KEY]) return self._result('t') - with patch.object( - MultiProjectMiddleware, - '_client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + with ( + patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): await MultiProjectMiddleware().on_call_tool(context, call_next) @@ -1113,7 +1155,7 @@ async def test_on_list_tools_injects_project_filter(self) -> None: # a read fan-out tool, an excluded tool, and a write tool read_tool = _tool('get_tables', read_only=True) read_tool.parameters = {'type': 'object', 'properties': {'bucket_ids': {'type': 'array'}}} - excluded = _tool('query_data', read_only=True) + excluded = _tool('get_project_info', read_only=True) excluded.parameters = {'type': 'object', 'properties': {}} write_tool = _tool('update_config', read_only=False) write_tool.parameters = {'type': 'object', 'properties': {}} @@ -1128,5 +1170,5 @@ async def call_next(_): tools = await MultiProjectMiddleware().on_list_tools(context, call_next) by_name = {t.name: t for t in tools} assert 'project_ids' in by_name['get_tables'].parameters['properties'] - assert 'project_ids' not in by_name['query_data'].parameters['properties'] + assert 'project_ids' not in by_name['get_project_info'].parameters['properties'] assert 'project_ids' not in by_name['update_config'].parameters['properties'] From 5d54a8994bb98480bbf8e34f6108a89db1872eed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 08:09:24 +0200 Subject: [PATCH 08/89] feat(PSGO-261): enrich get_accessible_projects with per-project SQL dialect get_accessible_projects now compacts several API calls into one bootstrap result so the assistant doesn't need a get_project_info per project: - Per-project sql_dialect, derived from the token's owner.defaultBackend via a token verify narrowed with X-KBC-ProjectId (concurrent, no workspace needed). - Optional with_llm_instruction flag returns the base working instructions as a top-level array grouped by SQL dialect (deduplicated, not copied per project), so mixed BigQuery/Snowflake scopes each get correct dialect guidance. workspace_id is omitted (not needed for bootstrap); branch context stays in get_project_info. Metastore and data-science already use the bearer/PAT token plus X-KBC-ProjectId, so no client change was needed there; Queue/AI/SyncActions still pass the raw storage token and remain follow-ups for PAT support. Co-Authored-By: Claude Opus 4.8 --- TOOLS.md | 17 +++- src/keboola_mcp_server/tools/project.py | 100 +++++++++++++++++++++++- tests/test_server.py | 2 +- tests/tools/test_project.py | 43 +++++++++- 4 files changed, 152 insertions(+), 10 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index 65fd56ea7..9b78363c6 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -50,7 +50,7 @@ providing their configuration IDs. - [modify_streamlit_data_app](#modify_streamlit_data_app): Creates or updates a Streamlit data app. ### Project Tools -- [get_accessible_projects](#get_accessible_projects): Lists the Keboola projects the current login can access across the stack. +- [get_accessible_projects](#get_accessible_projects): Lists the Keboola projects the current login can access across the stack, each with its SQL dialect. - [get_project_info](#get_project_info): Retrieves structured information about the current project, including essential context and base instructions for working with it (e. @@ -3562,18 +3562,27 @@ configuration is created e.g. keboola.ex-google-analytics-v4 and keboola.ex-gmai **Description**: -Lists the Keboola projects the current login can access across the stack. +Lists the Keboola projects the current login can access across the stack, each with its SQL dialect. Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), present the projects, and ask whether they want to work across all of them or a subset. Then call -`set_project_scope` with their choice. +`set_project_scope` with their choice. This tool compacts several API calls (token introspection +plus a per-project token verify for the SQL dialect) into one result, so the assistant does not +need a separate get_project_info call per project. Pass with_llm_instruction=true on the first +call to also receive the base working instructions grouped by dialect. **Input JSON Schema**: ```json { "additionalProperties": false, - "properties": {}, + "properties": { + "with_llm_instruction": { + "default": false, + "description": "If true, include the base working instructions (llm_instructions), grouped by SQL dialect. Request this once at the very start of a conversation; omit it on later calls.", + "type": "boolean" + } + }, "type": "object" } ``` diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index a9222330a..20b526265 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -12,7 +12,7 @@ from keboola_mcp_server.config import MetadataField from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager -from keboola_mcp_server.mcp import SCOPE_KEY, SessionScope +from keboola_mcp_server.mcp import SCOPE_KEY, MultiProjectMiddleware, ServerState, SessionScope, process_concurrently from keboola_mcp_server.resources.prompts import get_project_system_prompt from keboola_mcp_server.workspace import WorkspaceManager @@ -265,6 +265,16 @@ async def get_project_info( return project_info +def _sql_dialect_from_token(token_data: JsonDict) -> str | None: + """Derives the project's SQL dialect from the token's owner.defaultBackend, without a workspace.""" + backend = cast(JsonDict, token_data.get('owner', {})).get('defaultBackend') + if backend == 'snowflake': + return 'Snowflake' + if backend == 'bigquery': + return 'BigQuery' + return None + + class AccessibleProject(BaseModel): id: int = Field(description='The project id.') name: str | None = Field(default=None, description='The project name.') @@ -273,6 +283,17 @@ class AccessibleProject(BaseModel): is_active: bool = Field( default=False, description='Whether this is the active project (write / single-project tool target).' ) + sql_dialect: str | None = Field( + default=None, description='The SQL dialect of the project ("Snowflake" or "BigQuery").' + ) + + +class LlmInstructionGroup(BaseModel): + """Base working instructions shared by all projects of one SQL dialect (dialect-specific system prompt).""" + + project_ids: list[int] = Field(description='The scoped projects this instruction applies to.') + sql_dialect: str | None = Field(default=None, description='The SQL dialect these projects share.') + llm_instruction: str = Field(description='The base working instructions for projects of this dialect.') class AccessibleProjects(BaseModel): @@ -286,6 +307,14 @@ class AccessibleProjects(BaseModel): default=None, description='The active project that write operations and single-project tools target.' ) read_only: bool | None = Field(default=None, description='Whether the current scoped token is read-only.') + llm_instructions: list[LlmInstructionGroup] | None = Field( + default=None, + description=( + 'The base working instructions, grouped by SQL dialect (deduplicated across projects). ' + 'Only present when the tool is called with with_llm_instruction=true; request this once at the ' + 'start of a conversation.' + ), + ) llm_instruction: str = Field( description='Guidance for the assistant on how to use this result.', ) @@ -301,14 +330,43 @@ class ProjectScope(BaseModel): llm_instruction: str = Field(description='Guidance for the assistant on the new scope.') +async def _project_sql_dialect( + server_state: ServerState, subject_token: str, project_id: int +) -> tuple[int, str | None]: + """Fetches one project's SQL dialect by verifying the parent token narrowed with X-KBC-ProjectId. + + No workspace is provisioned — the dialect comes from the token's owner.defaultBackend, so this is + a single cheap Storage API call per project. + """ + per_client = await MultiProjectMiddleware._client_for_project( + server_state, subject_token, project_id, read_only=True + ) + token_data = await per_client.storage_client.verify_token() + return project_id, _sql_dialect_from_token(token_data) + + @tool_errors() -async def get_accessible_projects(ctx: Context) -> AccessibleProjects: +async def get_accessible_projects( + ctx: Context, + with_llm_instruction: Annotated[ + bool, + Field( + description=( + 'If true, include the base working instructions (llm_instructions), grouped by SQL dialect. ' + 'Request this once at the very start of a conversation; omit it on later calls.' + ) + ), + ] = False, +) -> AccessibleProjects: """ - Lists the Keboola projects the current login can access across the stack. + Lists the Keboola projects the current login can access across the stack, each with its SQL dialect. Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), present the projects, and ask whether they want to work across all of them or a subset. Then call - `set_project_scope` with their choice. + `set_project_scope` with their choice. This tool compacts several API calls (token introspection + plus a per-project token verify for the SQL dialect) into one result, so the assistant does not + need a separate get_project_info call per project. Pass with_llm_instruction=true on the first + call to also receive the base working instructions grouped by dialect. """ client = KeboolaClient.from_state(ctx.session.state) subject_token = await _parent_subject_token(client) @@ -318,6 +376,21 @@ async def get_accessible_projects(ctx: Context) -> AccessibleProjects: scoped_ids = scope.project_ids if isinstance(scope, SessionScope) and scope.confirmed else None active_id = scope.active_project_id if scoped_ids else None + # Enrich each project with its SQL dialect (concurrently). Best-effort: a project whose verify + # fails simply keeps sql_dialect=None rather than failing the whole listing. + server_state = ServerState.from_context(ctx) + dialects: dict[int, str | None] = {} + results = await process_concurrently( + [p.id for p in introspection.projects], + lambda pid: _project_sql_dialect(server_state, subject_token, pid), + ) + for result in results: + if isinstance(result, BaseException): + LOG.warning(f'Could not resolve SQL dialect for a project: {result}') + continue + pid, dialect = result + dialects[pid] = dialect + projects = [ AccessibleProject( id=p.id, @@ -325,9 +398,27 @@ async def get_accessible_projects(ctx: Context) -> AccessibleProjects: role=p.role, in_scope=scoped_ids is not None and p.id in scoped_ids, is_active=p.id == active_id, + sql_dialect=dialects.get(p.id), ) for p in introspection.projects ] + + # Optionally attach the base working instructions, grouped by dialect so the (large) prompt is + # sent once per distinct dialect rather than duplicated per project. + llm_instructions: list[LlmInstructionGroup] | None = None + if with_llm_instruction: + by_dialect: dict[str | None, list[int]] = {} + for p in projects: + by_dialect.setdefault(p.sql_dialect, []).append(p.id) + llm_instructions = [ + LlmInstructionGroup( + project_ids=ids, + sql_dialect=dialect, + llm_instruction=get_project_system_prompt(dialect or 'Snowflake'), + ) + for dialect, ids in by_dialect.items() + ] + if scoped_ids is None: instruction = ( 'No project scope has been confirmed yet. Ask the user whether to operate across all these ' @@ -345,6 +436,7 @@ async def get_accessible_projects(ctx: Context) -> AccessibleProjects: scoped_project_ids=scoped_ids, active_project_id=active_id, read_only=scope.read_only if scoped_ids is not None else None, + llm_instructions=llm_instructions, llm_instruction=instruction, ) diff --git a/tests/test_server.py b/tests/test_server.py index fced0c834..bea20d81f 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -152,7 +152,7 @@ async def test_tools_input_schema(self): missing_default.append(f'{tool.name}.{prop_name}') missing_properties.sort() - assert missing_properties == ['get_accessible_projects', 'get_project_info'] + assert missing_properties == ['get_project_info'] missing_type.sort() assert not missing_type, f'These tool params have no "type" info: {missing_type}' missing_default.sort() diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 2e0e9cb23..a446883ef 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -280,16 +280,27 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock introspect = mocker.patch( 'keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection) ) + # Per-project SQL dialect is resolved via a token verify narrowed by X-KBC-ProjectId; mock that. + mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) + dialects = {18: 'BigQuery', 83: 'Snowflake'} + mocker.patch( + 'keboola_mcp_server.tools.project._project_sql_dialect', + new=mocker.AsyncMock(side_effect=lambda _ss, _tok, pid: (pid, dialects[pid])), + ) # No scope confirmed yet. result = await get_accessible_projects(mcp_context_client) introspect.assert_awaited_once_with(STACK, subject_token='kbc_at_parent') assert result.user_email == 'm@k.com' - assert [(p.id, p.name, p.role) for p in result.projects] == [(18, 'A', 'admin'), (83, 'B', 'admin')] + assert [(p.id, p.name, p.role, p.sql_dialect) for p in result.projects] == [ + (18, 'A', 'admin', 'BigQuery'), + (83, 'B', 'admin', 'Snowflake'), + ] assert result.scoped_project_ids is None assert result.active_project_id is None assert result.read_only is None + assert result.llm_instructions is None # not requested assert all(not p.in_scope and not p.is_active for p in result.projects) # Once scoped, the current scope is surfaced on the projects and at the top level. @@ -301,6 +312,36 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock assert [(p.id, p.in_scope, p.is_active) for p in result.projects] == [(18, False, False), (83, True, True)] +@pytest.mark.asyncio +async def test_get_accessible_projects_llm_instructions_grouped_by_dialect( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + _prep_client(mcp_context_client, mocker) + introspection = SimpleNamespace( + user_email='m@k.com', + projects=[ + SimpleNamespace(id=18, name='A', role='admin'), + SimpleNamespace(id=86, name='B', role='admin'), + SimpleNamespace(id=95, name='C', role='admin'), + ], + ) + mocker.patch('keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection)) + mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) + dialects = {18: 'BigQuery', 86: 'BigQuery', 95: 'Snowflake'} + mocker.patch( + 'keboola_mcp_server.tools.project._project_sql_dialect', + new=mocker.AsyncMock(side_effect=lambda _ss, _tok, pid: (pid, dialects[pid])), + ) + + result = await get_accessible_projects(mcp_context_client, with_llm_instruction=True) + + assert result.llm_instructions is not None + # One group per distinct dialect, projects deduplicated into their dialect group (no per-project copies). + groups = {g.sql_dialect: g.project_ids for g in result.llm_instructions} + assert groups == {'BigQuery': [18, 86], 'Snowflake': [95]} + assert all(g.llm_instruction for g in result.llm_instructions) + + @pytest.mark.asyncio async def test_set_project_scope_subset_exchanges_and_stores( mcp_context_client: Context, mocker: MockerFixture From 21b68979a0489a7190b1bad37a31bce2420ebfbe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 08:19:16 +0200 Subject: [PATCH 09/89] refactor(PSGO-261): rename get_accessible_projects llm_instructions -> base_instructions The output had both llm_instruction (singular: how-to-use-this-result guidance, a codebase-wide convention) and llm_instructions (plural: the base working system prompts). The near-identical names were confusing. Rename the plural to base_instructions and its inner field to `instructions` so the two concepts read as distinct. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/tools/project.py | 20 ++++++++++---------- tests/tools/test_project.py | 8 ++++---- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 20b526265..b30ae865e 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -288,12 +288,12 @@ class AccessibleProject(BaseModel): ) -class LlmInstructionGroup(BaseModel): +class BaseInstructionGroup(BaseModel): """Base working instructions shared by all projects of one SQL dialect (dialect-specific system prompt).""" - project_ids: list[int] = Field(description='The scoped projects this instruction applies to.') + project_ids: list[int] = Field(description='The scoped projects these instructions apply to.') sql_dialect: str | None = Field(default=None, description='The SQL dialect these projects share.') - llm_instruction: str = Field(description='The base working instructions for projects of this dialect.') + instructions: str = Field(description='The base working instructions for projects of this dialect.') class AccessibleProjects(BaseModel): @@ -307,7 +307,7 @@ class AccessibleProjects(BaseModel): default=None, description='The active project that write operations and single-project tools target.' ) read_only: bool | None = Field(default=None, description='Whether the current scoped token is read-only.') - llm_instructions: list[LlmInstructionGroup] | None = Field( + base_instructions: list[BaseInstructionGroup] | None = Field( default=None, description=( 'The base working instructions, grouped by SQL dialect (deduplicated across projects). ' @@ -316,7 +316,7 @@ class AccessibleProjects(BaseModel): ), ) llm_instruction: str = Field( - description='Guidance for the assistant on how to use this result.', + description='Guidance for the assistant on how to use this result (distinct from base_instructions).', ) @@ -405,16 +405,16 @@ async def get_accessible_projects( # Optionally attach the base working instructions, grouped by dialect so the (large) prompt is # sent once per distinct dialect rather than duplicated per project. - llm_instructions: list[LlmInstructionGroup] | None = None + base_instructions: list[BaseInstructionGroup] | None = None if with_llm_instruction: by_dialect: dict[str | None, list[int]] = {} for p in projects: by_dialect.setdefault(p.sql_dialect, []).append(p.id) - llm_instructions = [ - LlmInstructionGroup( + base_instructions = [ + BaseInstructionGroup( project_ids=ids, sql_dialect=dialect, - llm_instruction=get_project_system_prompt(dialect or 'Snowflake'), + instructions=get_project_system_prompt(dialect or 'Snowflake'), ) for dialect, ids in by_dialect.items() ] @@ -436,7 +436,7 @@ async def get_accessible_projects( scoped_project_ids=scoped_ids, active_project_id=active_id, read_only=scope.read_only if scoped_ids is not None else None, - llm_instructions=llm_instructions, + base_instructions=base_instructions, llm_instruction=instruction, ) diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index a446883ef..a233f26f3 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -300,7 +300,7 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock assert result.scoped_project_ids is None assert result.active_project_id is None assert result.read_only is None - assert result.llm_instructions is None # not requested + assert result.base_instructions is None # not requested assert all(not p.in_scope and not p.is_active for p in result.projects) # Once scoped, the current scope is surfaced on the projects and at the top level. @@ -335,11 +335,11 @@ async def test_get_accessible_projects_llm_instructions_grouped_by_dialect( result = await get_accessible_projects(mcp_context_client, with_llm_instruction=True) - assert result.llm_instructions is not None + assert result.base_instructions is not None # One group per distinct dialect, projects deduplicated into their dialect group (no per-project copies). - groups = {g.sql_dialect: g.project_ids for g in result.llm_instructions} + groups = {g.sql_dialect: g.project_ids for g in result.base_instructions} assert groups == {'BigQuery': [18, 86], 'Snowflake': [95]} - assert all(g.llm_instruction for g in result.llm_instructions) + assert all(g.instructions for g in result.base_instructions) @pytest.mark.asyncio From ac9513767803e5cb5de12f7adb3532f38a9af537 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 08:23:59 +0200 Subject: [PATCH 10/89] =?UTF-8?q?docs(PSGO-261):=20RFC=20increment=203=20?= =?UTF-8?q?=E2=80=94=20query=20fan-out,=20dialect=20bootstrap,=20service-t?= =?UTF-8?q?oken=20gaps?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Document the increment-3 work: query_data per-project workspace fan-out (with the read-only-provisioning and cross-project-SQL limitations), get_accessible_projects dialect-aware bootstrap (per-project sql_dialect via token verify, base_instructions grouped by dialect), and the per-service PAT token status — Queue/AI/SyncActions still pass the raw storage token (future PAT work), metastore/data-science already use the bearer/PAT + X-KBC-ProjectId path. Co-Authored-By: Claude Opus 4.8 --- feature_spec/pat_token_support/RFC.md | 97 +++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index e17e8ffe1..9c7a1ea96 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -411,3 +411,100 @@ the per-project envelope shape change the docs), integration tests on a dev stac 1-element envelope.** Single-project UX is byte-for-byte unchanged. - **Q4 (write-target confirmation) — lean: explicit `project_id` arg on write tools, honored only when scope > 1.** (Still open; not blocking.) + +# Extension: query fan-out, dialect-aware bootstrap, per-service token gaps (PSGO-261, increment 3) + +## Context + +Increment 2 delivered read fan-out + scope tools but left three rough edges: `query_data` was +pinned to the active project's workspace, `get_accessible_projects` returned only id/name/role +(forcing a `get_project_info` per project for the SQL dialect), and only Storage + the Query +Service actually honor the multi-project token narrowing. This increment addresses the first two +and documents the third. + +## `query_data` fan-out + per-project workspace + +`query_data` is now a normal fan-out read tool — it was removed from `_NO_FANOUT_TOOLS`. The fan-out +swaps **both** the `KeboolaClient` **and** a `WorkspaceManager` built on it into session state for +the duration of a call (`MultiProjectMiddleware._swap_project`), so the SQL runs inside the targeted +project's own read-only workspace (its BigQuery dataset / Snowflake schema), not the active +project's. + +- Narrow to one project with the `project_ids` filter (`query_data(project_ids=[86])`), or run + across all scoped projects. +- Per-project workspaces are provisioned lazily on first use. `ponytail:` the manager is rebuilt + per call; a cache surviving the per-request state rebuild is a follow-up if provisioning latency + shows up. +- Merged `structured_content` for a fanned-out query keeps the first project's `csv_data` (scalar + deep-merge); every project's full result is present in the per-project text envelopes. Structured + multi-CSV merge is deferred. + +### Known limitations (accepted, not solved this increment) +- **Read-only scope can't provision a first-time workspace** (workspace creation is a POST). The + first `query_data` into a project without an existing MCP workspace needs a non-read-only scope. + Verified: read-only scope → `Forbidden POST operation on a readonly client` on workspace create. +- **No cross-project SQL in a single statement.** BigQuery has no cross-project data access; + Snowflake reaches another project only via a *materialized* linked-bucket alias. A single + `query_data` call always executes inside exactly one project's workspace. This is a backend + constraint, not an MCP limitation — a future increment could add FQN-aware routing, but the join + itself is impossible in one statement regardless. + +## `get_accessible_projects` as the dialect-aware bootstrap call + +`get_accessible_projects` now compacts several API calls into one bootstrap result so the assistant +does not need a `get_project_info` per project: + +- **Introspection** → reachable projects (id, name, role). +- **Per-project token verify** (parent token narrowed with `X-KBC-ProjectId`, run concurrently) → + each project's `sql_dialect`, derived from `owner.defaultBackend` — **no workspace provisioned**. +- **Current scope surfaced** → `scoped_project_ids`, `active_project_id`, `read_only`, and + per-project `in_scope` / `is_active` flags. (There is no separate scope-introspection tool; this + is the read side of scope state without mutating the token.) +- **Optional base instructions** → `with_llm_instruction=true` returns `base_instructions`: a + top-level array grouped by SQL dialect (deduplicated, **not** copied per project), e.g. + `[{project_ids:[18,86], sql_dialect:"BigQuery", instructions:"…"}, {project_ids:[95], + sql_dialect:"Snowflake", instructions:"…"}]`. Request once at the start of a conversation. + +`workspace_id` is intentionally omitted here (not needed for bootstrap). The result keeps the +codebase-wide singular `llm_instruction` field (how-to-use-this-result guidance) distinct from the +plural `base_instructions` (the working system prompts). + +### `get_project_info` caveat +`get_project_info` stays in `_NO_FANOUT_TOOLS` and reports only the active project. In a +mixed-dialect scope its single `sql_dialect` / dialect-specific `llm_instruction` is misleading for +the other projects. Prefer `get_accessible_projects` for multi-project bootstrap. Follow-up: fan out +`get_project_info`, or split its static prompt from the per-project dialect/branch/workspace facts. + +## Per-service token support under multi-project scope + +Fan-out narrows a call to one project via the **`X-KBC-ProjectId` header** on a shared token. Only +services that read that header work under header-narrowing. Current wiring (`clients/client.py`): + +| Service | Token today | PAT / multi-project status | +|---|---|---| +| Storage (`connection`) | `bearer_or_sapi_token` + `X-KBC-ProjectId` | ✅ works | +| Query Service | workspace bearer | ✅ per-project workspace | +| Metastore (semantic) | `bearer_or_sapi_token` | ✅ PAT/bearer-first, SAPI fallback (guarded); feature-gated, untested on stacks without `mcp-semantic-tooling` | +| Data Science (sandboxes) | `bearer_or_sapi_token` + `X-KBC-ProjectId` | ✅ PAT + project header (verified: data-app create + deploy) | +| Scheduler | `bearer_or_sapi_token` | ✅ bearer-first (writes only) | +| **Jobs Queue** | `self._token` (raw → `X-StorageAPI-Token`) | ❌ 401 under PAT fan-out — **future PAT work** | +| **AI Service** | `self._token` | ⚠️ raw token — **future PAT work** | +| **Sync Actions** | `self._token` | ⚠️ raw token — **future PAT work** | + +### Future PAT work (not in this increment) +`jobs_queue`, `ai_service`, and `sync_actions` still pass the raw `self._token`, so under a +PAT/multi-project session the satellite service rejects it (verified: `get_jobs` → 401 "Invalid +access token" from the Queue API). To support PAT they must use the bearer/PAT token (like +metastore/data-science/scheduler) or be handed a per-project minted token. Metastore and +data-science already use the bearer/PAT path, so no change was needed there. + +## Decisions (increment 3) + +- **`query_data` fans out with a per-project workspace** rather than being pinned to the active + project. Single-project targeting via the `project_ids` filter; cross-project SQL stays out of + scope (backend-impossible in one statement). +- **`get_accessible_projects` is the multi-project bootstrap**: per-project dialect via token verify + (no workspace), current scope surfaced, base instructions grouped by dialect behind + `with_llm_instruction`. +- **Queue / AI / SyncActions PAT support is deferred** and documented above; metastore + data-science + already satisfy the PAT/bearer + `X-KBC-ProjectId` contract. From 5bc49e4dca155b4909e6f753b856e60cac5c353a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 10:36:26 +0200 Subject: [PATCH 11/89] test(PSGO-261): add get_accessible_projects + set_project_scope to integtest tool set _assert_basic_setup hardcodes the expected tool list and fails on any unlisted tool. The two multi-project bootstrap tools are always registered, so add them to expected_tools (they are not feature-gated, so they don't belong in the exclude set). Co-Authored-By: Claude Opus 4.8 --- integtests/test_mcp_server.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/integtests/test_mcp_server.py b/integtests/test_mcp_server.py index 8581da255..2569681b4 100644 --- a/integtests/test_mcp_server.py +++ b/integtests/test_mcp_server.py @@ -178,6 +178,7 @@ async def _assert_basic_setup(client: Client): 'find_component_id', 'get_buckets', 'get_components', + 'get_accessible_projects', 'get_config_examples', 'get_configs', 'get_data_apps', @@ -189,6 +190,7 @@ async def _assert_basic_setup(client: Client): 'get_shared_buckets', 'get_tables', 'link_shared_bucket', + 'set_project_scope', 'modify_flow', 'modify_python_js_data_app', 'modify_streamlit_data_app', From e54bc49d58cdf3e4ad309c84baa778ea0cae4a16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 10:41:07 +0200 Subject: [PATCH 12/89] test(PSGO-261): CI integration coverage for the multi-project PAT flow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds integtests that exercise the real introspect + scoped-exchange + fan-out path using a Keboola programmatic token (kbc_pat_/kbc_at_): - conftest: `programmatic_token` (from INTEGTEST_STORAGE_PAT) and `programmatic_token_url` (INTEGTEST_STORAGE_PAT_URL, falling back to the pool URL) fixtures. Both skip cleanly when the secret is absent, so CI stays green until the token is wired in. - test_pat_multiproject.py: get_accessible_projects (per-project sql_dialect + dialect-grouped base_instructions) and set_project_scope → fanned-out get_buckets, driven through the real middleware pipeline via an in-process Client. - ci.yml: pass INTEGTEST_STORAGE_PAT (secret) and INTEGTEST_STORAGE_PAT_URL (var) to the integtests env. Co-Authored-By: Claude Opus 4.8 --- .github/workflows/ci.yml | 4 ++ integtests/conftest.py | 27 +++++++++++ integtests/test_pat_multiproject.py | 71 +++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+) create mode 100644 integtests/test_pat_multiproject.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d341d7555..f0195655f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -138,6 +138,10 @@ jobs: INTEGTEST_STORAGE_TOKENS: ${{ secrets.INTEGTEST_STORAGE_TOKENS }} INTEGTEST_POOL_STORAGE_API_URL: ${{ vars.INTEGTEST_POOL_STORAGE_API_URL }} INTEGTEST_STORAGE_TOKEN_STORAGE_BRANCHES: ${{ secrets.INTEGTEST_STORAGE_TOKEN_STORAGE_BRANCHES }} + # Programmatic token (kbc_pat_/kbc_at_) for the multi-project PAT flow tests. When unset, + # integtests/test_pat_multiproject.py skips. INTEGTEST_STORAGE_PAT_URL falls back to the pool URL. + INTEGTEST_STORAGE_PAT: ${{ secrets.INTEGTEST_STORAGE_PAT }} + INTEGTEST_STORAGE_PAT_URL: ${{ vars.INTEGTEST_STORAGE_PAT_URL }} run: | uv run tox -e integtests diff --git a/integtests/conftest.py b/integtests/conftest.py index 133d0406e..14f9f6e0d 100644 --- a/integtests/conftest.py +++ b/integtests/conftest.py @@ -45,6 +45,11 @@ # The second pair of token/schema for testing simultaneous access to two different projects. STORAGE_API_TOKEN_ENV_VAR_2 = 'INTEGTEST_STORAGE_TOKEN_PRJ2' WORKSPACE_SCHEMA_ENV_VAR_2 = 'INTEGTEST_WORKSPACE_SCHEMA_PRJ2' +# A Keboola programmatic token (kbc_pat_/kbc_at_) used to exercise the multi-project PAT flow +# (introspect + scoped exchange + fan-out). Optional: tests skip when it is not set. +STORAGE_PAT_ENV_VAR = 'INTEGTEST_STORAGE_PAT' +# Stack URL the PAT belongs to; falls back to the pool URL when not set. +STORAGE_PAT_URL_ENV_VAR = 'INTEGTEST_STORAGE_PAT_URL' # We reset dev environment variables to integtest values to ensure tests run locally using .env settings. DEV_STORAGE_API_URL_ENV_VAR = 'STORAGE_API_URL' DEV_STORAGE_TOKEN_ENV_VAR = 'KBC_STORAGE_TOKEN' @@ -205,6 +210,28 @@ def workspace_schema(_clean_project: None, storage_api_token: str, storage_api_u LOG.exception(f'Failed to delete test-session workspace {workspace_id}') +@pytest.fixture(scope='session') +def programmatic_token(env_file_loaded: bool) -> str: + """A Keboola programmatic token (kbc_pat_/kbc_at_) for the multi-project PAT flow. + + Skips the test when INTEGTEST_STORAGE_PAT is not configured, so the suite stays green + on stacks/CI runs where no PAT is provided. + """ + token = os.getenv(STORAGE_PAT_ENV_VAR) + if not token: + pytest.skip(f'{STORAGE_PAT_ENV_VAR} not set; skipping PAT multi-project integration tests.') + return token + + +@pytest.fixture(scope='session') +def programmatic_token_url(env_file_loaded: bool) -> str: + """Stack URL for the programmatic token; falls back to the pool URL.""" + url = os.getenv(STORAGE_PAT_URL_ENV_VAR) or os.getenv(POOL_STORAGE_API_URL_ENV_VAR) + if not url: + pytest.skip(f'Neither {STORAGE_PAT_URL_ENV_VAR} nor {POOL_STORAGE_API_URL_ENV_VAR} set.') + return url + + @pytest.fixture(scope='session') def storage_api_token_2(env_file_loaded: bool) -> str | None: return os.getenv(STORAGE_API_TOKEN_ENV_VAR_2) diff --git a/integtests/test_pat_multiproject.py b/integtests/test_pat_multiproject.py new file mode 100644 index 000000000..b3c89711a --- /dev/null +++ b/integtests/test_pat_multiproject.py @@ -0,0 +1,71 @@ +"""Integration tests for the multi-project PAT flow (PSGO-261). + +These exercise the real introspect + scoped-exchange + fan-out path, which only activates for a +Keboola programmatic token (kbc_pat_/kbc_at_) in local (non-deployed) mode. They are gated on the +INTEGTEST_STORAGE_PAT secret via the `programmatic_token` fixture and skip when it is absent, so CI +stays green until the token is wired in. +""" + +import logging + +import pytest +import pytest_asyncio +from fastmcp import Client, FastMCP + +from integtests.conftest import INTEGTEST_CLIENT_INFO +from keboola_mcp_server.config import Config, ServerRuntimeInfo +from keboola_mcp_server.server import create_server +from keboola_mcp_server.tools.project import AccessibleProjects, ProjectScope + +LOG = logging.getLogger(__name__) + + +@pytest.fixture +def pat_mcp_server(programmatic_token: str, programmatic_token_url: str) -> FastMCP: + # No workspace_schema and no project_id → the server auto-leases a multi-project scope from the + # PAT introspection (local programmatic mode). + config = Config(storage_api_url=programmatic_token_url, storage_token=programmatic_token) + server = create_server(config, runtime_info=ServerRuntimeInfo(transport='stdio')) + assert isinstance(server, FastMCP) + return server + + +@pytest_asyncio.fixture +async def pat_mcp_client(pat_mcp_server: FastMCP): + async with Client(pat_mcp_server, client_info=INTEGTEST_CLIENT_INFO) as client: + yield client + + +@pytest.mark.asyncio +async def test_pat_get_accessible_projects(pat_mcp_client: Client): + # get_accessible_projects is a bootstrap tool: callable before any scope is confirmed. + response = await pat_mcp_client.call_tool('get_accessible_projects', {'with_llm_instruction': True}) + result = AccessibleProjects.model_validate(response.structured_content) + + assert result.projects, 'the PAT should reach at least one project' + # Every project is enriched with its SQL dialect (derived from the per-project token verify). + assert all(p.sql_dialect in ('Snowflake', 'BigQuery') for p in result.projects) + # with_llm_instruction=True → base instructions grouped by dialect (one group per distinct dialect). + assert result.base_instructions, 'base_instructions expected when with_llm_instruction=true' + dialects_in_groups = {g.sql_dialect for g in result.base_instructions} + dialects_in_projects = {p.sql_dialect for p in result.projects} + assert dialects_in_groups == dialects_in_projects + + +@pytest.mark.asyncio +async def test_pat_scope_then_read_fans_out(pat_mcp_client: Client): + # Discover reachable projects, then confirm a scope over all of them. + accessible = AccessibleProjects.model_validate( + (await pat_mcp_client.call_tool('get_accessible_projects', {})).structured_content + ) + project_ids = [p.id for p in accessible.projects] + + scope_response = await pat_mcp_client.call_tool('set_project_scope', {'read_only': True}) + scope = ProjectScope.model_validate(scope_response.structured_content) + assert set(scope.project_ids) == set(project_ids) + assert scope.read_only is True + assert scope.active_project_id == project_ids[0] + + # A read tool now runs against the confirmed scope instead of raising the ask-first gate error. + buckets_response = await pat_mcp_client.call_tool('get_buckets', {}) + assert buckets_response.structured_content is not None From a55e015b1a0707ccb59ccb7e9820c5c73a3716ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 1 Jul 2026 10:58:19 +0200 Subject: [PATCH 13/89] test(PSGO-261): run integtests across sapi/pat_single/pat_mpa auth over the same pool Reuse the existing project pool + per-project lock and vary only authentication, instead of a separate PAT pool: - A single INTEGTEST_STORAGE_PAT (user member of all pool projects) drives the PAT modes against the same locked pool projects; the pool's SAPI token still handles lock + cleanup, so the PAT only needs read access. - New auth_client fixture parametrizes [sapi, pat_single, pat_mpa]; PAT modes build the server with the PAT and set_project_scope (read-only) to the locked project(s). - mpa_second_project acquires one more pool project (non-blocking, best-effort) so MPA fan-out spans >1 project; degrades to a single-project scope when none is free. - test_read_buckets_across_auth_modes runs the same read question under every mode; test_pat_accessible_projects_enrichment checks per-project dialect + grouped instructions. - Drop INTEGTEST_STORAGE_PAT_URL (PAT uses the pool URL). Co-Authored-By: Claude Opus 4.8 --- .github/workflows/ci.yml | 6 +- integtests/conftest.py | 52 +++++++++--- integtests/test_pat_multiproject.py | 122 ++++++++++++++++++---------- 3 files changed, 125 insertions(+), 55 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f0195655f..cc5db3c6a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -138,10 +138,10 @@ jobs: INTEGTEST_STORAGE_TOKENS: ${{ secrets.INTEGTEST_STORAGE_TOKENS }} INTEGTEST_POOL_STORAGE_API_URL: ${{ vars.INTEGTEST_POOL_STORAGE_API_URL }} INTEGTEST_STORAGE_TOKEN_STORAGE_BRANCHES: ${{ secrets.INTEGTEST_STORAGE_TOKEN_STORAGE_BRANCHES }} - # Programmatic token (kbc_pat_/kbc_at_) for the multi-project PAT flow tests. When unset, - # integtests/test_pat_multiproject.py skips. INTEGTEST_STORAGE_PAT_URL falls back to the pool URL. + # A single programmatic token (kbc_pat_/kbc_at_) whose user is a member of all pool projects. + # Drives the PAT/MPA auth modes against the SAME pool projects (uses the pool URL). PAT-auth + # tests in integtests/test_pat_multiproject.py skip when it is unset. INTEGTEST_STORAGE_PAT: ${{ secrets.INTEGTEST_STORAGE_PAT }} - INTEGTEST_STORAGE_PAT_URL: ${{ vars.INTEGTEST_STORAGE_PAT_URL }} run: | uv run tox -e integtests diff --git a/integtests/conftest.py b/integtests/conftest.py index 14f9f6e0d..84c47264f 100644 --- a/integtests/conftest.py +++ b/integtests/conftest.py @@ -45,11 +45,11 @@ # The second pair of token/schema for testing simultaneous access to two different projects. STORAGE_API_TOKEN_ENV_VAR_2 = 'INTEGTEST_STORAGE_TOKEN_PRJ2' WORKSPACE_SCHEMA_ENV_VAR_2 = 'INTEGTEST_WORKSPACE_SCHEMA_PRJ2' -# A Keboola programmatic token (kbc_pat_/kbc_at_) used to exercise the multi-project PAT flow -# (introspect + scoped exchange + fan-out). Optional: tests skip when it is not set. +# A single Keboola programmatic token (kbc_pat_/kbc_at_) whose user is a member of ALL pool +# projects. It exercises the multi-project PAT flow (introspect + scoped exchange + fan-out) +# against the SAME pool projects, just with different authentication — no separate PAT pool. +# Optional: PAT-auth tests skip when it is not set. The PAT uses the pool storage_api_url. STORAGE_PAT_ENV_VAR = 'INTEGTEST_STORAGE_PAT' -# Stack URL the PAT belongs to; falls back to the pool URL when not set. -STORAGE_PAT_URL_ENV_VAR = 'INTEGTEST_STORAGE_PAT_URL' # We reset dev environment variables to integtest values to ensure tests run locally using .env settings. DEV_STORAGE_API_URL_ENV_VAR = 'STORAGE_API_URL' DEV_STORAGE_TOKEN_ENV_VAR = 'KBC_STORAGE_TOKEN' @@ -223,13 +223,45 @@ def programmatic_token(env_file_loaded: bool) -> str: return token +def _try_acquire_additional_project(exclude_project_ids: set[str]) -> AcquiredProject | None: + """Non-blocking single pass over the pool to lock one more project (for MPA breadth). + + Reuses the same per-project lock as the primary acquisition, but does NOT block/retry: if no + other pool project is free right now, returns None and the caller degrades to a single-project + MPA scope. Never waits, so it can't hang MPA setup when the pool has only one project or the + others are busy with concurrent runs. + """ + if _project_pool is None: + return None + for endpoint in _project_pool._endpoints: + if endpoint.project_id in exclude_project_ids: + continue + lock_info = _project_pool._make_lock(endpoint)._try_acquire_once() + if lock_info is not None: + return AcquiredProject(endpoint=endpoint, lock_info=lock_info) + return None + + @pytest.fixture(scope='session') -def programmatic_token_url(env_file_loaded: bool) -> str: - """Stack URL for the programmatic token; falls back to the pool URL.""" - url = os.getenv(STORAGE_PAT_URL_ENV_VAR) or os.getenv(POOL_STORAGE_API_URL_ENV_VAR) - if not url: - pytest.skip(f'Neither {STORAGE_PAT_URL_ENV_VAR} nor {POOL_STORAGE_API_URL_ENV_VAR} set.') - return url +def mpa_second_project(project_lock: AcquiredProject) -> Generator[AcquiredProject | None, Any, None]: + """A second, exclusively-locked pool project so MPA fan-out spans >1 project. + + Best-effort: yields None when the pool has no other free project (MPA tests then run against a + single-project scope). Released and cleaned on teardown. + """ + second = _try_acquire_additional_project(exclude_project_ids={project_lock.endpoint.project_id}) + if second is None: + yield None + return + try: + yield second + finally: + if _project_pool is not None: + try: + _clean_project(second.endpoint.storage_api_token, second.endpoint.storage_api_url) + except Exception: + LOG.exception(f'Failed to clean second MPA project {second.endpoint.project_id}') + _project_pool.release(second) @pytest.fixture(scope='session') diff --git a/integtests/test_pat_multiproject.py b/integtests/test_pat_multiproject.py index b3c89711a..8380ba9a2 100644 --- a/integtests/test_pat_multiproject.py +++ b/integtests/test_pat_multiproject.py @@ -1,71 +1,109 @@ -"""Integration tests for the multi-project PAT flow (PSGO-261). +"""Integration tests that run the SAME questions under different authentication (PSGO-261). -These exercise the real introspect + scoped-exchange + fan-out path, which only activates for a -Keboola programmatic token (kbc_pat_/kbc_at_) in local (non-deployed) mode. They are gated on the -INTEGTEST_STORAGE_PAT secret via the `programmatic_token` fixture and skip when it is absent, so CI -stays green until the token is wired in. +Instead of a separate PAT project pool, these reuse the existing project pool + per-project lock and +vary only the auth used to drive the MCP tools: + +- ``sapi`` — the pool project's legacy Storage API token (today's single-project path). +- ``pat_single`` — a programmatic token (kbc_pat_/kbc_at_) scoped to the one locked pool project. +- ``pat_mpa`` — the same PAT scoped to two locked pool projects (fan-out), or one when the pool + has no second free project. + +The PAT (INTEGTEST_STORAGE_PAT) must be a member of the pool projects; PAT modes skip when it is not +set. Lock + cleanup always use the project's SAPI token, so the PAT only needs read access here. """ import logging +from collections.abc import AsyncGenerator import pytest import pytest_asyncio from fastmcp import Client, FastMCP from integtests.conftest import INTEGTEST_CLIENT_INFO +from integtests.project_lock import AcquiredProject from keboola_mcp_server.config import Config, ServerRuntimeInfo from keboola_mcp_server.server import create_server from keboola_mcp_server.tools.project import AccessibleProjects, ProjectScope LOG = logging.getLogger(__name__) +AUTH_MODES = ['sapi', 'pat_single', 'pat_mpa'] + -@pytest.fixture -def pat_mcp_server(programmatic_token: str, programmatic_token_url: str) -> FastMCP: - # No workspace_schema and no project_id → the server auto-leases a multi-project scope from the - # PAT introspection (local programmatic mode). - config = Config(storage_api_url=programmatic_token_url, storage_token=programmatic_token) +def _make_server(config: Config) -> FastMCP: server = create_server(config, runtime_info=ServerRuntimeInfo(transport='stdio')) assert isinstance(server, FastMCP) return server -@pytest_asyncio.fixture -async def pat_mcp_client(pat_mcp_server: FastMCP): - async with Client(pat_mcp_server, client_info=INTEGTEST_CLIENT_INFO) as client: - yield client +@pytest_asyncio.fixture(params=AUTH_MODES) +async def auth_client( + request: pytest.FixtureRequest, + storage_api_url: str, + storage_api_token: str, + workspace_schema: str, + project_lock: AcquiredProject, +) -> AsyncGenerator[tuple[str, Client, list[int]], None]: + """Yields (auth_mode, ready-to-use Client, scoped_project_ids) for each auth mode. + + For PAT modes the client is already scoped (read-only) to the locked pool project(s), so the test + body is identical across modes. + """ + auth_mode = request.param + primary_id = int(project_lock.endpoint.project_id) + + if auth_mode == 'sapi': + config = Config( + storage_api_url=storage_api_url, storage_token=storage_api_token, workspace_schema=workspace_schema + ) + async with Client(_make_server(config), client_info=INTEGTEST_CLIENT_INFO) as client: + yield auth_mode, client, [primary_id] + return + + # PAT modes: skip when no programmatic token is configured. + pat = request.getfixturevalue('programmatic_token') + target_ids = [primary_id] + if auth_mode == 'pat_mpa': + second = request.getfixturevalue('mpa_second_project') + if second is not None: + target_ids.append(int(second.endpoint.project_id)) + + config = Config(storage_api_url=storage_api_url, storage_token=pat) + async with Client(_make_server(config), client_info=INTEGTEST_CLIENT_INFO) as client: + scope = ProjectScope.model_validate( + ( + await client.call_tool('set_project_scope', {'project_ids': target_ids, 'read_only': True}) + ).structured_content + ) + assert set(scope.project_ids) == set(target_ids) + yield auth_mode, client, target_ids @pytest.mark.asyncio -async def test_pat_get_accessible_projects(pat_mcp_client: Client): - # get_accessible_projects is a bootstrap tool: callable before any scope is confirmed. - response = await pat_mcp_client.call_tool('get_accessible_projects', {'with_llm_instruction': True}) - result = AccessibleProjects.model_validate(response.structured_content) +async def test_read_buckets_across_auth_modes(auth_client: tuple[str, Client, list[int]]): + """The same read question (list buckets) works under sapi / pat_single / pat_mpa.""" + auth_mode, client, target_ids = auth_client + response = await client.call_tool('get_buckets', {}) + sc = response.structured_content + assert sc is not None, f'{auth_mode}: expected structured content' + # Both single-project (raw) and MPA (merged fan-out) shapes expose a top-level "buckets" list. + assert 'buckets' in sc, f'{auth_mode}: expected a buckets list, got keys {list(sc)}' + LOG.info(f'{auth_mode}: get_buckets returned {len(sc["buckets"])} bucket(s) across {target_ids}') + + +@pytest.mark.asyncio +async def test_pat_accessible_projects_enrichment( + programmatic_token: str, + storage_api_url: str, +): + """PAT bootstrap: per-project SQL dialect + dialect-grouped base instructions.""" + config = Config(storage_api_url=storage_api_url, storage_token=programmatic_token) + async with Client(_make_server(config), client_info=INTEGTEST_CLIENT_INFO) as client: + result = AccessibleProjects.model_validate( + (await client.call_tool('get_accessible_projects', {'with_llm_instruction': True})).structured_content + ) assert result.projects, 'the PAT should reach at least one project' - # Every project is enriched with its SQL dialect (derived from the per-project token verify). assert all(p.sql_dialect in ('Snowflake', 'BigQuery') for p in result.projects) - # with_llm_instruction=True → base instructions grouped by dialect (one group per distinct dialect). assert result.base_instructions, 'base_instructions expected when with_llm_instruction=true' - dialects_in_groups = {g.sql_dialect for g in result.base_instructions} - dialects_in_projects = {p.sql_dialect for p in result.projects} - assert dialects_in_groups == dialects_in_projects - - -@pytest.mark.asyncio -async def test_pat_scope_then_read_fans_out(pat_mcp_client: Client): - # Discover reachable projects, then confirm a scope over all of them. - accessible = AccessibleProjects.model_validate( - (await pat_mcp_client.call_tool('get_accessible_projects', {})).structured_content - ) - project_ids = [p.id for p in accessible.projects] - - scope_response = await pat_mcp_client.call_tool('set_project_scope', {'read_only': True}) - scope = ProjectScope.model_validate(scope_response.structured_content) - assert set(scope.project_ids) == set(project_ids) - assert scope.read_only is True - assert scope.active_project_id == project_ids[0] - - # A read tool now runs against the confirmed scope instead of raising the ask-first gate error. - buckets_response = await pat_mcp_client.call_tool('get_buckets', {}) - assert buckets_response.structured_content is not None + assert {g.sql_dialect for g in result.base_instructions} == {p.sql_dialect for p in result.projects} From 475624af978a24e1aabd8b7ac680015556ef56d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 2 Jul 2026 07:35:06 +0200 Subject: [PATCH 14/89] feat(PSGO-261): auto-login on local server start when no/dead session The locally-run (stdio) server no longer needs a separate `login` step: when it is started with only the stack URL and no token, and there is no stored session (or the stored one can no longer be refreshed), it now runs the browser PKCE login on the spot via new ensure_access_token(), stores the session, and continues. Only applies to local stdio (browser + loopback reachable on the same machine); remote /deployed servers still use client-driven OAuth. Mid-session dead refresh is unchanged (surfaces the re-login guidance). Also add INTEGTEST_STORAGE_PAT to the tox integtests pass_env allowlist so the PAT multi-project tests actually receive the secret in CI. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/auth_login.py | 37 +++++++++++++++++++++++-- src/keboola_mcp_server/cli.py | 16 ++++++++--- tests/test_auth_login.py | 40 ++++++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 6 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 64b34b2e3..725e9fedf 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -16,6 +16,7 @@ import logging import os import secrets +import sys import time import urllib.parse import webbrowser @@ -285,6 +286,37 @@ async def get_access_token( return tokens.access_token +async def ensure_access_token( + storage_api_url: str, + *, + allow_interactive: bool = True, + open_browser=webbrowser.open, + transport: httpx.AsyncBaseTransport | None = None, +) -> str: + """Return a valid access token, running the browser PKCE login when needed and allowed. + + Convenience for the locally-run (stdio) server so it can be started with only the stack URL + and no separate ``login`` step: if no session is stored, or the stored one can no longer be + refreshed, this logs in interactively (opens a browser + loopback callback), persists the + session, and returns the fresh token. + + ``allow_interactive`` MUST be false unless a real terminal is attached. When the stdio server + is launched by an MCP client its stdout is the JSON-RPC channel and there is no TTY, so an + interactive login would both corrupt the protocol stream and block the initialize handshake + (the loopback wait has no timeout). In that case this raises the same "run login" guidance as + ``get_access_token`` instead of attempting a browser login. Remote/deployed servers must use + client-driven OAuth regardless. + """ + try: + return await get_access_token(storage_api_url, transport=transport) + except RuntimeError as exc: + if not allow_interactive: + raise + LOG.info(f'No usable stored session for {storage_api_url} ({exc}); starting browser login.') + await perform_login(storage_api_url, open_browser=open_browser) + return await get_access_token(storage_api_url, transport=transport) + + def _forget(storage_api_url: str) -> None: store = _read_store() if store.pop(_store_key(storage_api_url), None) is not None: @@ -326,7 +358,8 @@ async def perform_login(storage_api_url: str, *, open_browser=webbrowser.open) - 'state': state, } authorize_url = f'{_base_url(storage_api_url)}/{_AUTHORIZE_PATH}?{urllib.parse.urlencode(params)}' - print(f'Open this URL in your browser to authenticate:\n\n {authorize_url}\n', flush=True) + # Never write to stdout: under the stdio transport stdout is the JSON-RPC channel. Use stderr. + print(f'Open this URL in your browser to authenticate:\n\n {authorize_url}\n', file=sys.stderr, flush=True) open_browser(authorize_url) _CallbackHandler.result = {} @@ -342,7 +375,7 @@ async def perform_login(storage_api_url: str, *, open_browser=webbrowser.open) - if not code: raise RuntimeError('Authorization callback did not return a code.') - print('Exchanging authorization code for tokens…', flush=True) + print('Exchanging authorization code for tokens…', file=sys.stderr, flush=True) tokens = await exchange_code( storage_api_url, code=code, state=state, code_verifier=verifier, redirect_uri=redirect_uri ) diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 46222d390..4e2a8794e 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -165,12 +165,20 @@ async def run_server(args: list[str] | None = None) -> None: # Create and run the server if parsed_args.transport == 'stdio': # Local/stdio needs only the stack URL: with no token configured, use the tokens - # leased by a prior browser `login` (refreshing them as needed). + # leased by a prior browser `login` (refreshing them as needed). When no session is + # stored or it can no longer be refreshed, log in interactively on the spot — so the + # server can be started without a separate `login` step. config = config.replace_by(os.environ) if not config.storage_token and config.storage_api_url: - from keboola_mcp_server.auth_login import get_access_token - - access_token = await get_access_token(config.storage_api_url) + from keboola_mcp_server.auth_login import ensure_access_token + + # Only run the interactive browser login when a real terminal is attached. When an + # MCP client launches this stdio server, stdin/stdout are pipes (no TTY) and stdout + # is the JSON-RPC channel — an interactive login there would corrupt the protocol + # and block the initialize handshake. In that case require a prior `login` (or a + # configured token) and fail fast with guidance instead. + allow_interactive = sys.stdin.isatty() and sys.stderr.isatty() + access_token = await ensure_access_token(config.storage_api_url, allow_interactive=allow_interactive) config = dataclasses.replace(config, storage_token=access_token) runtime_config = ServerRuntimeInfo(transport=parsed_args.transport) diff --git a/tests/test_auth_login.py b/tests/test_auth_login.py index 6ad8eb136..fec841a4c 100644 --- a/tests/test_auth_login.py +++ b/tests/test_auth_login.py @@ -13,6 +13,7 @@ from keboola_mcp_server import auth_login from keboola_mcp_server.auth_login import ( TokenSet, + ensure_access_token, exchange_code, exchange_scoped_token, get_access_token, @@ -126,6 +127,45 @@ async def test_get_access_token_dead_token_forgets_and_raises(creds_file: Path) assert load_tokens(STACK) is None +@pytest.mark.asyncio +async def test_ensure_access_token_returns_stored_without_login(creds_file: Path, monkeypatch) -> None: + save_tokens(STACK, TokenSet('kbc_at_valid', 'kbc_rt_1', expires_at=time.time() + 3600)) + + async def _must_not_login(*_a, **_k): + raise AssertionError('perform_login must not run when a valid session is stored') + + monkeypatch.setattr(auth_login, 'perform_login', _must_not_login) + token = await ensure_access_token(STACK, transport=_token_response(500)) + assert token == 'kbc_at_valid' + + +@pytest.mark.asyncio +async def test_ensure_access_token_logs_in_when_no_session(creds_file: Path, monkeypatch) -> None: + # No stored session → ensure_access_token runs the browser login, then returns the fresh token. + calls: list[str] = [] + + async def _fake_login(storage_api_url: str, **_k): + calls.append(storage_api_url) + save_tokens(storage_api_url, TokenSet('kbc_at_fresh', 'kbc_rt_fresh', expires_at=time.time() + 3600)) + + monkeypatch.setattr(auth_login, 'perform_login', _fake_login) + token = await ensure_access_token(STACK, transport=_token_response(500)) + assert token == 'kbc_at_fresh' + assert calls == [STACK] + + +@pytest.mark.asyncio +async def test_ensure_access_token_non_interactive_raises_without_login(creds_file: Path, monkeypatch) -> None: + # No TTY (e.g. launched by an MCP client): must NOT attempt a browser login (it would corrupt + # the stdio protocol / hang the handshake); raise the clear guidance instead. + async def _must_not_login(*_a, **_k): + raise AssertionError('perform_login must not run when interactive login is disallowed') + + monkeypatch.setattr(auth_login, 'perform_login', _must_not_login) + with pytest.raises(RuntimeError, match='Run "keboola-mcp-server login'): + await ensure_access_token(STACK, allow_interactive=False) + + def test_pkce_challenge_is_sha256_of_verifier() -> None: verifier = auth_login._b64url(b'0123456789abcdef0123456789abcdef0123456789ab') expected = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode('ascii')).digest()).decode().rstrip('=') From 207004b5ae6d34ab916f7b5d0daf515a0db2876e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 2 Jul 2026 09:02:57 +0200 Subject: [PATCH 15/89] feat(PSGO-261): login --pat leases a PAT over all projects (sudo MFA + create) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `keboola-mcp-server login --pat --totp ` (or --recovery): after the browser PKCE login returns the session token, lease a Personal Access Token over ALL accessible projects and print it. The flow is introspect (all project ids) → POST /v1/auth/sudo (MFA elevation; totp/recovery mutually exclusive, type omitted) → POST /v1/auth/pat (name + ~1 month expiry + projects as strings). New auth_login helpers: elevate_session, create_pat, lease_pat, unit-tested with mocked transports. NOTE: the /v1/auth/sudo and /v1/auth/pat request/response field names are assumed to mirror the existing exchange endpoint's camelCase conventions (defensive response parsing) and must be confirmed against the auth API once reachable. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/auth_login.py | 108 +++++++++++++++++++++++++++ src/keboola_mcp_server/cli.py | 55 +++++++++++++- tests/test_auth_login.py | 85 +++++++++++++++++++++ 3 files changed, 244 insertions(+), 4 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 725e9fedf..9e80a538c 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -36,7 +36,10 @@ _REFRESH_PATH = 'v1/auth/token/refresh' _INTROSPECT_PATH = 'v1/auth/token/introspect' _EXCHANGE_PATH = 'v1/auth/pat/exchange' +_SUDO_PATH = 'v1/auth/sudo' +_PAT_PATH = 'v1/auth/pat' _REFRESH_SKEW_SECONDS = 60 +_PAT_DEFAULT_EXPIRES_SECONDS = 30 * 24 * 60 * 60 # ~1 month _CREDENTIALS_PATH = Path.home() / '.keboola' / 'mcp' / 'credentials.json' @@ -177,6 +180,111 @@ async def exchange_scoped_token( ) +async def elevate_session( + storage_api_url: str, + *, + subject_token: str, + totp_code: str | None = None, + recovery_code: str | None = None, + transport: httpx.AsyncBaseTransport | None = None, +) -> str: + """Elevates (``sudo``) the session with an MFA code via POST /v1/auth/sudo. + + ``totp_code`` and ``recovery_code`` are mutually exclusive — pass exactly one. The ``type`` + field is intentionally omitted (empty). Returns the elevated bearer token to authorize + sensitive operations such as PAT creation; if the endpoint elevates the session in place and + returns no token, falls back to the original ``subject_token``. + + NOTE: response field name assumed (``token``/``accessToken``) — confirm against the auth API. + """ + if bool(totp_code) == bool(recovery_code): + raise ValueError('Provide exactly one of totp_code or recovery_code.') + payload = {'totpCode': totp_code} if totp_code else {'recoveryCode': recovery_code} + async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + response = await client.post( + f'{_base_url(storage_api_url)}/{_SUDO_PATH}', + headers={'Authorization': f'Bearer {subject_token}'}, + json=payload, + ) + if response.is_error: + raise RuntimeError(f'POST /{_SUDO_PATH} failed ({response.status_code}): {response.text}') + body = cast(dict, response.json()) if response.content else {} + return cast(str, body.get('token') or body.get('accessToken') or subject_token) + + +async def create_pat( + storage_api_url: str, + *, + subject_token: str, + project_ids: list[int], + name: str, + expires_in: int = _PAT_DEFAULT_EXPIRES_SECONDS, + transport: httpx.AsyncBaseTransport | None = None, +) -> str: + """Creates a Personal Access Token (``kbc_pat_*``) via POST /v1/auth/pat. + + ``subject_token`` must be an elevated (sudo) bearer. ``project_ids`` are sent as strings (the + auth service rejects integers, per the exchange endpoint). Requires a prior ``elevate_session``. + + Projects are nested under ``scope`` (mirroring /v1/auth/pat/exchange); a top-level ``projects`` + field is rejected by the API. Response token field assumed (``token``/``pat``/``accessToken``). + """ + payload = { + 'name': name, + 'expiresIn': expires_in, + 'scope': {'projects': [str(p) for p in project_ids]}, + } + async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + response = await client.post( + f'{_base_url(storage_api_url)}/{_PAT_PATH}', + headers={'Authorization': f'Bearer {subject_token}'}, + json=payload, + ) + if response.is_error: + # Surface the validation body so a wrong/missing field is visible (the schema is assumed). + raise RuntimeError(f'POST /{_PAT_PATH} failed ({response.status_code}) with {payload=}: {response.text}') + body = cast(dict, response.json()) + pat = body.get('token') or body.get('pat') or body.get('accessToken') + if not pat: + raise RuntimeError(f'PAT creation response did not contain a token: keys={sorted(body)}') + return cast(str, pat) + + +async def lease_pat( + storage_api_url: str, + *, + subject_token: str, + totp_code: str | None = None, + recovery_code: str | None = None, + name: str = 'keboola-mcp-server', + expires_in: int = _PAT_DEFAULT_EXPIRES_SECONDS, + transport: httpx.AsyncBaseTransport | None = None, +) -> str: + """Leases a PAT over ALL accessible projects: introspect → sudo (MFA) → create PAT. + + ``subject_token`` is the whole-stack session access token (``kbc_at_*``) from the PKCE login. + """ + introspection = await introspect_token(storage_api_url, subject_token=subject_token, transport=transport) + project_ids = [p.id for p in introspection.projects] + if not project_ids: + raise RuntimeError('The session token can not reach any projects; cannot create a PAT.') + elevated = await elevate_session( + storage_api_url, + subject_token=subject_token, + totp_code=totp_code, + recovery_code=recovery_code, + transport=transport, + ) + return await create_pat( + storage_api_url, + subject_token=elevated, + project_ids=project_ids, + name=name, + expires_in=expires_in, + transport=transport, + ) + + async def exchange_code( storage_api_url: str, *, diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 4e2a8794e..8a6af6703 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -70,6 +70,22 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: help='Keboola Storage API URL (e.g. https://connection..keboola.com). ' 'Falls back to KBC_STORAGE_API_URL.', ) + login_parser.add_argument( + '--pat', + action='store_true', + help='After the browser login, lease a Personal Access Token (kbc_pat_) over all accessible ' + 'projects and print it. Requires an MFA code (--totp or --recovery).', + ) + login_parser.add_argument('--totp', metavar='CODE', help='TOTP MFA code for the sudo elevation (--pat).') + login_parser.add_argument( + '--recovery', metavar='CODE', help='Recovery MFA code for the sudo elevation (--pat); alternative to --totp.' + ) + login_parser.add_argument( + '--pat-name', + metavar='STR', + default='keboola-mcp-server', + help='Name for the leased PAT (--pat).', + ) return parser.parse_args(args) @@ -112,17 +128,42 @@ async def _http_exception_handler(request: Request, exc: HTTPException): } -async def _run_login(api_url: str | None) -> None: - """Runs the interactive browser PKCE login and stores the leased tokens.""" - from keboola_mcp_server.auth_login import perform_login +async def _run_login( + api_url: str | None, + *, + pat: bool = False, + totp: str | None = None, + recovery: str | None = None, + pat_name: str = 'keboola-mcp-server', +) -> None: + """Runs the interactive browser PKCE login and stores the leased tokens. + + With ``pat=True``, additionally leases a Personal Access Token over all accessible projects + (introspect → sudo with the MFA code → create PAT) and prints it. + """ + from keboola_mcp_server.auth_login import lease_pat, perform_login storage_api_url = api_url or os.environ.get('KBC_STORAGE_API_URL') if not storage_api_url: raise RuntimeError('A Storage API URL is required for login: pass --api-url or set KBC_STORAGE_API_URL.') + + if pat and bool(totp) == bool(recovery): + raise RuntimeError('Leasing a PAT (--pat) requires exactly one MFA code: pass --totp or --recovery.') + tokens = await perform_login(storage_api_url) remaining = max(0, int(tokens.expires_at - time.time())) print(f'\n✓ Session stored for {storage_api_url} (access token expires in ~{remaining}s).') + if pat: + pat_token = await lease_pat( + storage_api_url, + subject_token=tokens.access_token, + totp_code=totp, + recovery_code=recovery, + name=pat_name, + ) + print(f'\n✓ Personal Access Token (valid ~1 month, all accessible projects):\n\n {pat_token}\n') + async def run_server(args: list[str] | None = None) -> None: """Runs the MCP server in async mode.""" @@ -151,7 +192,13 @@ async def run_server(args: list[str] | None = None) -> None: ) if parsed_args.command == 'login': - await _run_login(getattr(parsed_args, 'api_url', None)) + await _run_login( + getattr(parsed_args, 'api_url', None), + pat=getattr(parsed_args, 'pat', False), + totp=getattr(parsed_args, 'totp', None), + recovery=getattr(parsed_args, 'recovery', None), + pat_name=getattr(parsed_args, 'pat_name', 'keboola-mcp-server'), + ) return # Create config from the CLI arguments diff --git a/tests/test_auth_login.py b/tests/test_auth_login.py index fec841a4c..87ac7151c 100644 --- a/tests/test_auth_login.py +++ b/tests/test_auth_login.py @@ -13,11 +13,14 @@ from keboola_mcp_server import auth_login from keboola_mcp_server.auth_login import ( TokenSet, + create_pat, + elevate_session, ensure_access_token, exchange_code, exchange_scoped_token, get_access_token, introspect_token, + lease_pat, load_tokens, refresh_tokens, save_tokens, @@ -233,3 +236,85 @@ def handler(request: httpx.Request) -> httpx.Response: assert scoped.project_ids == [18, 83] assert scoped.expires_at > time.time() assert not scoped.is_near_expiry + + +# --- sudo elevation + PAT creation (PSGO-261) --- + + +@pytest.mark.asyncio +async def test_elevate_session_sends_totp_and_returns_token() -> None: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured['url'] = str(request.url) + captured['auth'] = request.headers['Authorization'] + captured['body'] = json.loads(request.content) + return httpx.Response(200, json={'token': 'kbc_sudo_1'}) + + token = await elevate_session( + STACK, subject_token='kbc_at_x', totp_code='123456', transport=httpx.MockTransport(handler) + ) + assert captured['url'] == 'https://connection.keboola.com/v1/auth/sudo' + assert captured['auth'] == 'Bearer kbc_at_x' + assert captured['body'] == {'totpCode': '123456'} # recoveryCode/type omitted + assert token == 'kbc_sudo_1' + + +@pytest.mark.asyncio +async def test_elevate_session_requires_exactly_one_code() -> None: + with pytest.raises(ValueError, match='exactly one'): + await elevate_session(STACK, subject_token='kbc_at_x', totp_code='1', recovery_code='2') + with pytest.raises(ValueError, match='exactly one'): + await elevate_session(STACK, subject_token='kbc_at_x') + + +@pytest.mark.asyncio +async def test_create_pat_sends_projects_and_parses_token() -> None: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured['url'] = str(request.url) + captured['auth'] = request.headers['Authorization'] + captured['body'] = json.loads(request.content) + return httpx.Response(201, json={'token': 'kbc_pat_new'}) + + pat = await create_pat( + STACK, + subject_token='kbc_sudo_1', + project_ids=[18, 83], + name='demo', + expires_in=2592000, + transport=httpx.MockTransport(handler), + ) + assert captured['url'] == 'https://connection.keboola.com/v1/auth/pat' + assert captured['auth'] == 'Bearer kbc_sudo_1' + # project ids serialized as strings and nested under scope, like the exchange endpoint + assert captured['body'] == {'name': 'demo', 'expiresIn': 2592000, 'scope': {'projects': ['18', '83']}} + assert pat == 'kbc_pat_new' + + +@pytest.mark.asyncio +async def test_lease_pat_introspects_then_sudo_then_creates() -> None: + # One routing transport across the three endpoints the flow hits, asserting the sudo token is + # what authorizes PAT creation and that all introspected projects are included. + seen: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + path = request.url.path + seen.append(path) + if path.endswith('/token/introspect'): + return httpx.Response(200, json={'projects': [{'id': 18}, {'id': 83}, {'id': 95}]}) + if path.endswith('/auth/sudo'): + assert json.loads(request.content) == {'recoveryCode': 'rec-9'} + return httpx.Response(200, json={'token': 'kbc_sudo_1'}) + if path.endswith('/auth/pat'): + assert request.headers['Authorization'] == 'Bearer kbc_sudo_1' + assert json.loads(request.content)['scope']['projects'] == ['18', '83', '95'] + return httpx.Response(201, json={'token': 'kbc_pat_leased'}) + raise AssertionError(f'unexpected path {path}') + + pat = await lease_pat( + STACK, subject_token='kbc_at_parent', recovery_code='rec-9', transport=httpx.MockTransport(handler) + ) + assert pat == 'kbc_pat_leased' + assert [p.split('/')[-1] for p in seen] == ['introspect', 'sudo', 'pat'] From ee7677bd445c56a801a22c06c3d254fc58bb9340 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 07:51:29 +0200 Subject: [PATCH 16/89] feat(PSGO-261): scope-first tool visibility for programmatic sessions Replace the "auto-lease a phantom active project + block data tools at call-time" UX with a list-time gate: while a programmatic (kbc_*) session has an unconfirmed scope, on_list_tools advertises only the scoping tools (get_accessible_projects, set_project_scope). set_project_scope confirms the scope and emits notifications/tools/list_changed so the client reveals the full tool set. Backward compatible: the gate keys on SessionScope presence, which only programmatic/ PKCE sessions get. Legacy Storage-token sessions have no scope -> unchanged passthrough (all tools advertised at connect, no gate). Call-time ask-first gate kept as defense. Also fixes a latent bug where set_project_scope read minted.read_only on the exchange-failure path (minted unbound) -> use the stored scope's read_only. RFC increment-4 documents this plus the analysis vs PR #451 (fan-out relevance/latency conceded/answered; attribution + per-project error isolation tracked as follow-ups). Co-Authored-By: Claude Opus 4.8 --- feature_spec/pat_token_support/RFC.md | 73 +++++++++++++++++++++++++ src/keboola_mcp_server/mcp.py | 9 +++ src/keboola_mcp_server/tools/project.py | 10 +++- tests/test_mcp.py | 29 ++++++++++ 4 files changed, 120 insertions(+), 1 deletion(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 9c7a1ea96..438a6b44b 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -508,3 +508,76 @@ data-science already use the bearer/PAT path, so no change was needed there. `with_llm_instruction`. - **Queue / AI / SyncActions PAT support is deferred** and documented above; metastore + data-science already satisfy the PAT/bearer + `X-KBC-ProjectId` contract. + +# Extension: scope-first tool visibility + reviewer feedback (PSGO-261, increment 4) + +## Context — reviewer feedback vs. PR #451 + +An earlier MPA attempt (PR #451, `davidesner`) took a different shape: static numbered SAPI tokens +(`KBC_STORAGE_TOKEN_1..N`) in `.mcp.json`, a middleware that injects a `project_id`/`branch_id` +parameter into every tool schema, and the **agent** passing `project_id` per call (so covering N +projects means the agent calls the tool N times). Two critiques of our fan-out/scope model were +raised against that backdrop. Verdict after analysis: + +- **"Fan-out is worse than N explicit calls."** Partly conceded, partly not: + - *Relevance / context bloat* — not a real differentiator: the user can scope the token or use the + `project_ids` filter to target one project, and a genuine all-projects request bloats context in + either design. + - *Latency* — fixable: the fan-out loop should run **concurrently** (it is currently sequential). + - *Attribution & error isolation* — the one real gap (see below). Kept as follow-up. +- **"Active project while unscoped feels weird; expect tools to load after the first scope."** — + Accepted. Implemented as scope-first tool visibility (below). + +## Attribution & error isolation (the remaining fan-out gap) + +Concrete, with a 2-project read: + +- **Attribution.** `get_buckets` fan-out concatenates both projects' `buckets` lists via + `_deep_merge`; each bucket has `source_project: null`, so the merged `structured_content` cannot + say which project a bucket came from (only the `=== project N ===` text envelope can, which + structured-output clients don't parse). Two explicit calls each carry their project by construction. +- **Error isolation.** The fan-out loop is `for p in targets: results.append(await call_next())` — + if one project raises (e.g. `get_jobs` → Queue 401 on project 95 while 86 succeeds), the exception + propagates and the **whole** call fails, discarding project 86's good result. Two explicit calls + isolate the failure (86 returns jobs, 95 returns its 401). + +Follow-up (not in this increment): make fan-out concurrent, catch per-project errors into a +per-project `{project_id, ok|error}` envelope, and stamp `source_project` on merged rows. + +## Scope-first tool visibility (implemented) + +Replaces the "auto-lease a phantom active project + block data tools at call-time" UX with a +list-time gate: + +- A **programmatic (`kbc_*`) session** auto-leases an *unconfirmed* scope. While unconfirmed, + `on_list_tools` advertises **only** the scoping tools (`get_accessible_projects`, + `set_project_scope`). Data tools are not shown. +- `set_project_scope` confirms the scope and emits `notifications/tools/list_changed`, so the client + re-fetches and the full tool set appears. (Best-effort: the data tools are also no longer gated at + call-time once scope is confirmed, so a client that ignores `list_changed` still works.) +- The call-time ask-first gate is retained as defense in depth. + +### Backward compatibility — the gate is `kbc_*`-only +This is safe precisely because the gate keys on the presence of an (unconfirmed) `SessionScope`, +which **only** programmatic/PKCE sessions ever get: + +- **Legacy Storage-token sessions** (`1234-…`, not `kbc_*`) have no `SessionScope` → `on_list_tools` + is an unchanged passthrough: every tool is advertised from connect, no gate, single project. Zero + BC impact — no existing integration relying on "all tools at connect" breaks. +- **Deployed programmatic** sessions resolve to a single configured project (no auto-lease) → no + gate, single-project behavior. +- Only **local multi-project `kbc_*`** sessions (a new capability with no prior integrations) see the + hide-then-reveal flow, so there is nothing to break. + +The remaining risk is a client that ignores `tools/list_changed`; those still function (post-scope +calls are ungated) but won't visually refresh the tool list until reconnect. + +## Decisions (increment 4) + +- **Scope-first tool visibility** for `kbc_*` sessions; legacy Storage-token sessions unchanged + (gate keyed on `SessionScope` presence, which is equivalent to the `kbc_*` prefix). +- **No phantom active project** surfaced before a scope is confirmed (tools that need one are hidden). +- **Fan-out stays**, with the relevance/latency critiques answered by the `project_ids` filter and a + (follow-up) concurrent loop; attribution + per-project error isolation are the tracked follow-ups. +- Fixed a latent bug: `set_project_scope` referenced `minted.read_only` on the exchange-failure path + where `minted` is unbound — now uses the stored scope's `read_only`. diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 0aae46104..a9365d790 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -892,6 +892,15 @@ async def on_list_tools( ctx = context.fastmcp_context state = getattr(ctx.session, 'state', None) scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None + + # Scope-first UX (programmatic / kbc_* sessions only — those are the ones that get an + # auto-leased scope): until the user confirms a scope via set_project_scope, expose ONLY the + # scoping tools. The full tool set is revealed after confirmation (set_project_scope emits + # notifications/tools/list_changed). Legacy Storage-token sessions have no SessionScope, so + # they keep advertising every tool unchanged — no BC impact. + if isinstance(scope, SessionScope) and not scope.confirmed: + return [t for t in tools if t.name in _BOOTSTRAP_TOOLS] + if not (isinstance(scope, SessionScope) and len(scope.project_ids) > 1): return tools diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index b30ae865e..3bc36cb67 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -493,10 +493,18 @@ async def set_project_scope( scope = SessionScope(project_ids=ids, read_only=read_only, confirmed=True) ctx.session.state[SCOPE_KEY] = scope + # Scope-first UX: the tool list is filtered to scoping-only until a scope is confirmed. Now that + # it is, tell the client to re-fetch so the full tool set appears. Best-effort — clients that + # don't act on list_changed still work (the data tools are no longer gated once scope is set). + try: + await ctx.session.send_tool_list_changed() + except Exception as e: + LOG.debug(f'Could not send tools/list_changed after scoping: {e}') + multi = len(ids) > 1 return ProjectScope( project_ids=ids, - read_only=minted.read_only, + read_only=scope.read_only, active_project_id=scope.active_project_id, llm_instruction=( ( diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 5f89b6e0d..608da514f 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1172,3 +1172,32 @@ async def call_next(_): assert 'project_ids' in by_name['get_tables'].parameters['properties'] assert 'project_ids' not in by_name['get_project_info'].parameters['properties'] assert 'project_ids' not in by_name['update_config'].parameters['properties'] + + @pytest.mark.asyncio + async def test_on_list_tools_scope_first_hides_data_tools_until_confirmed(self) -> None: + # Unconfirmed (auto-leased) scope: only the scoping tools are advertised; data tools appear + # after set_project_scope confirms. Legacy sessions (no scope) are covered by the passthrough. + scope = SessionScope(project_ids=[11, 22], confirmed=False) + context, _ = self._ctx(scope, 'x', read_only=True) + + async def call_next(_): + return [ + _tool('get_accessible_projects', read_only=True), + _tool('set_project_scope', read_only=True), + _tool('get_tables', read_only=True), + _tool('update_config', read_only=False), + ] + + tools = await MultiProjectMiddleware().on_list_tools(context, call_next) + assert {t.name for t in tools} == {'get_accessible_projects', 'set_project_scope'} + + @pytest.mark.asyncio + async def test_on_list_tools_no_scope_is_passthrough(self) -> None: + # Legacy Storage-token session (no SessionScope): every tool stays advertised, unchanged. + context, _ = self._ctx(None, 'x', read_only=True) + + async def call_next(_): + return [_tool('get_tables', read_only=True), _tool('update_config', read_only=False)] + + tools = await MultiProjectMiddleware().on_list_tools(context, call_next) + assert {t.name for t in tools} == {'get_tables', 'update_config'} From b561c171b2023ee832be35290cf78388bea4a2fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 09:10:54 +0200 Subject: [PATCH 17/89] feat(PSGO-261): count-first fan-out with a safety cap for big projects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fan-out saves fixed structural overhead, not data volume, so on big projects an unbounded get_buckets/get_tables fanned across N projects can return hundreds of thousands of tokens in one result and overflow the context window. MultiProjectMiddleware._merge now degrades to count-first past _FANOUT_MAX_ITEMS (200 total items): under the cap it keeps the full per-project detail; over it, it returns a single note with per-project counts + a truncated (schema-safe) sample and steers to project_ids / search. Counters (bucket_counts, search total) stay summed so true totals survive truncation. RFC increment-4 gains the scaling analysis (fixed-overhead, %→0) and this fix. Follow-ups: real limit/offset pagination on the enumerators, and concurrent fan-out. Co-Authored-By: Claude Opus 4.8 --- feature_spec/pat_token_support/RFC.md | 29 +++++++++++++ src/keboola_mcp_server/mcp.py | 59 +++++++++++++++++++++++---- tests/test_mcp.py | 27 ++++++++++++ 3 files changed, 107 insertions(+), 8 deletions(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 438a6b44b..9d5f51e6d 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -581,3 +581,32 @@ calls are ungated) but won't visually refresh the tool list until reconnect. (follow-up) concurrent loop; attribution + per-project error isolation are the tracked follow-ups. - Fixed a latent bug: `set_project_scope` referenced `minted.read_only` on the exchange-failure path where `minted` is unbound — now uses the stored scope's `read_only`. + +## Scale: count-first fan-out with a safety cap + +Fan-out's saving is a *fixed* structural overhead (deduped envelopes/wrappers/turns, ~a few hundred +tokens across N projects) — it does **not** compress data. So as projects grow, the percentage cut +trends to zero and the binding cost becomes raw **data volume**: + +| buckets/proj (×6) | data | fan-out | explicit | cut % | +|---|--:|--:|--:|--:| +| 5 | 3,375 tok | 3,658 | 3,978 | 8.0% | +| 50 | 33,750 | 34,033 | 34,353 | 0.9% | +| 500 | 337,500 | 337,783 | 338,103 | 0.1% | + +An unbounded enumerator (`get_buckets`/`get_tables` have no `limit`/`offset`) fanned out across N +big projects returns hundreds of thousands of tokens in one tool result — overflowing the context +window in *either* model. Fan-out is a round-trip/turn optimizer, not a data-volume one. + +**Fix — `MultiProjectMiddleware._merge` degrades to count-first past a cap** (`_FANOUT_MAX_ITEMS`, +default 200 total items across projects): +- Under the cap: unchanged — per-project text envelopes + fully merged lists. +- Over the cap: return a single guidance note with **per-project item counts**, a **truncated sample** + (first `_FANOUT_MAX_ITEMS`, schema-safe — a shorter list still validates), and steer the agent to + **narrow with `project_ids`** or **use `search`**. Counters (e.g. `bucket_counts`, search `total`) + are summed by `_deep_merge`, so they keep reflecting the true totals even when the item lists are + truncated. The per-project full text dumps are dropped in this path (that is the context saving). + +This makes the multi-project path safe on humongous projects: it can never wedge the session, and it +nudges toward the scalable access patterns (search / per-project drill-down) instead of bulk-listing. +Follow-up: real `limit`/`offset` pagination on the enumerators, and concurrent fan-out. diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index a9365d790..6d1246b2e 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -984,23 +984,66 @@ def _deep_merge(a: Any, b: Any) -> Any: return a + b # counters like search "total" return a + # Total list items across projects before a fanned-out result degrades to count-first: instead of + # dumping every project's full listing (which, on big projects, overflows the context window in a + # single tool result), return per-project counts + a truncated sample + guidance to narrow. Small + # multi-project results stay fully detailed. Class attribute so tests can lower it. + _FANOUT_MAX_ITEMS = 200 + + @staticmethod + def _largest_list_len(sc: Any) -> int: + """Item count of a structured payload = the length of its largest top-level list (buckets/tables/hits).""" + if isinstance(sc, dict): + return max((len(v) for v in sc.values() if isinstance(v, list)), default=0) + if isinstance(sc, list): + return len(sc) + return 0 + + @staticmethod + def _truncate_lists(sc: Any, limit: int) -> Any: + """Truncate every top-level list to `limit` (schema-safe: a shorter list still validates).""" + if isinstance(sc, dict): + return {k: (v[:limit] if isinstance(v, list) else v) for k, v in sc.items()} + if isinstance(sc, list): + return sc[:limit] + return sc + @staticmethod def _merge(results: list[tuple[int, 'ToolResult']]) -> 'ToolResult': - # Label each project's output in the text content (attribution the model can read) and, when the - # tool declares structured output, deep-merge the per-project structured payloads into a single - # schema-valid object (lists concatenated across projects) so the client's output-schema - # validation passes. - content: list[Any] = [] + # Deep-merge the per-project structured payloads into one schema-valid object (lists concatenated + # across projects). Counters (e.g. bucket_counts, search total) are summed by _deep_merge, so they + # keep reflecting the true totals even if the item lists get truncated below. merged_structured: Any = None + per_project_counts: list[tuple[int, int]] = [] + total_items = 0 for project_id, result in results: - content.append(mt.TextContent(type='text', text=f'=== project {project_id} ===')) - content.extend(result.content or []) sc = result.structured_content + per_project_counts.append((project_id, MultiProjectMiddleware._largest_list_len(sc))) + total_items += MultiProjectMiddleware._largest_list_len(sc) if sc is not None: merged_structured = ( sc if merged_structured is None else MultiProjectMiddleware._deep_merge(merged_structured, sc) ) - return ToolResult(content=content, structured_content=merged_structured) + + # Small enough: full detail with per-project text envelopes (attribution the model can read). + if total_items <= MultiProjectMiddleware._FANOUT_MAX_ITEMS: + content: list[Any] = [] + for project_id, result in results: + content.append(mt.TextContent(type='text', text=f'=== project {project_id} ===')) + content.extend(result.content or []) + return ToolResult(content=content, structured_content=merged_structured) + + # Count-first: the combined listing is too large for one result. Return per-project counts, a + # truncated sample (first _FANOUT_MAX_ITEMS), and guidance — instead of every project's full dump. + summary = ', '.join(f'project {pid}: {n}' for pid, n in per_project_counts) + note = ( + f'Multi-project result is large — {total_items} items across {len(results)} project(s) ' + f'({summary}). Showing the first {MultiProjectMiddleware._FANOUT_MAX_ITEMS} in structured_content; ' + f'counters reflect the true totals. Narrow with project_ids=[...] on this tool, or use the ' + f'search tool to find specific items.' + ) + truncated = MultiProjectMiddleware._truncate_lists(merged_structured, MultiProjectMiddleware._FANOUT_MAX_ITEMS) + return ToolResult(content=[mt.TextContent(type='text', text=note)], structured_content=truncated) def _to_python(data: Any, exclude_none: bool = True) -> Any | None: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 608da514f..84d3d6196 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1201,3 +1201,30 @@ async def call_next(_): tools = await MultiProjectMiddleware().on_list_tools(context, call_next) assert {t.name for t in tools} == {'get_tables', 'update_config'} + + @staticmethod + def _items_result(n: int) -> ToolResult: + return ToolResult( + content=[mt.TextContent(type='text', text=f'{n} items')], + structured_content={'buckets': list(range(n)), 'total': n}, + ) + + def test_merge_small_keeps_full_detail(self) -> None: + merged = MultiProjectMiddleware._merge([(11, self._items_result(2)), (22, self._items_result(3))]) + # Under the cap: per-project text envelopes + fully merged lists; counters summed. + assert merged.structured_content == {'buckets': [0, 1, 0, 1, 2], 'total': 5} + assert [c.text for c in merged.content] == ['=== project 11 ===', '2 items', '=== project 22 ===', '3 items'] + + def test_merge_large_degrades_to_count_first(self, monkeypatch) -> None: + # Lower the cap so a modest result trips the count-first path. + monkeypatch.setattr(MultiProjectMiddleware, '_FANOUT_MAX_ITEMS', 3) + merged = MultiProjectMiddleware._merge([(11, self._items_result(2)), (22, self._items_result(3))]) + # Single guidance note (no per-project full dump), lists truncated, counters preserved. + assert len(merged.content) == 1 + note = merged.content[0].text + assert 'project 11: 2' in note + assert 'project 22: 3' in note + assert 'search tool' in note + assert 'project_ids' in note + assert len(merged.structured_content['buckets']) == 3 # truncated to the cap + assert merged.structured_content['total'] == 5 # true total preserved From 196b6a693e06c0dde44895b5ff9a5320ec7d4605 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 10:57:30 +0200 Subject: [PATCH 18/89] feat(PSGO-261): refresh-first login + partial-result fan-out on per-project errors login is now refresh-first: it reuses a stored session and refreshes the access token when the refresh token is still valid (no browser); the browser PKCE flow runs only when there is no stored session or the refresh token itself is dead. Re-running `login` an hour later just leases a fresh access token. Uses ensure_access_token(allow_interactive=True). Fan-out no longer fails wholesale when one project errors: MultiProjectMiddleware catches per-project exceptions, returns the successful projects' merged result plus a per-project retry hint ("retry with project_ids=[N]") in the text content, and only raises an aggregate error when *every* scoped project failed. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/cli.py | 18 +++++++++----- src/keboola_mcp_server/mcp.py | 33 ++++++++++++++++++++++---- tests/test_mcp.py | 44 +++++++++++++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 11 deletions(-) diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 8a6af6703..63fddd4b7 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -136,12 +136,16 @@ async def _run_login( recovery: str | None = None, pat_name: str = 'keboola-mcp-server', ) -> None: - """Runs the interactive browser PKCE login and stores the leased tokens. + """Establishes a stored session and, with ``pat=True``, leases a PAT. + + Refresh-first: if a stored session exists and its refresh token is still valid, this refreshes + (no browser) — so re-running `login` an hour later just leases a fresh access token. A browser + PKCE login runs only when there is no stored session or the refresh token itself is dead. With ``pat=True``, additionally leases a Personal Access Token over all accessible projects (introspect → sudo with the MFA code → create PAT) and prints it. """ - from keboola_mcp_server.auth_login import lease_pat, perform_login + from keboola_mcp_server.auth_login import ensure_access_token, lease_pat, load_tokens storage_api_url = api_url or os.environ.get('KBC_STORAGE_API_URL') if not storage_api_url: @@ -150,14 +154,16 @@ async def _run_login( if pat and bool(totp) == bool(recovery): raise RuntimeError('Leasing a PAT (--pat) requires exactly one MFA code: pass --totp or --recovery.') - tokens = await perform_login(storage_api_url) - remaining = max(0, int(tokens.expires_at - time.time())) - print(f'\n✓ Session stored for {storage_api_url} (access token expires in ~{remaining}s).') + # Refresh-first, browser only when dead (interactive: this is the terminal `login` command). + access_token = await ensure_access_token(storage_api_url, allow_interactive=True) + tokens = load_tokens(storage_api_url) + remaining = max(0, int(tokens.expires_at - time.time())) if tokens else 0 + print(f'\n✓ Session ready for {storage_api_url} (access token expires in ~{remaining}s).') if pat: pat_token = await lease_pat( storage_api_url, - subject_token=tokens.access_token, + subject_token=access_token, totp_code=totp, recovery_code=recovery, name=pat_name, diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 6d1246b2e..5fdd8944d 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -870,15 +870,29 @@ async def on_call_tool( state[WorkspaceManager.STATE_KEY] = original_workspace results: list[tuple[int, ToolResult]] = [] + errors: list[tuple[int, str]] = [] try: for project_id in targets: await self._swap_project(state, server_state, base_token, project_id, scope.read_only) - results.append((project_id, await call_next(context))) + # Isolate per-project failures: one project's error (e.g. Queue 401, a transient 5xx) + # must not discard the other projects' good results. Collect it and keep going, so the + # agent gets a partial response plus a retry hint. CancelledError is BaseException, so + # `except Exception` lets client cancellation propagate. + try: + results.append((project_id, await call_next(context))) + except Exception as e: + LOG.warning(f'Fan-out call failed for project {project_id}: {e}') + errors.append((project_id, str(e))) finally: state[KeboolaClient.STATE_KEY] = original_client state[WorkspaceManager.STATE_KEY] = original_workspace - return self._merge(results) + # Every project failed → nothing partial to return; surface a single aggregate error. + if not results and errors: + detail = '; '.join(f'project {pid}: {msg}' for pid, msg in errors) + raise ToolError(f'The tool failed for all {len(errors)} scoped project(s): {detail}') + + return self._merge(results, errors) async def on_list_tools( self, @@ -1009,10 +1023,19 @@ def _truncate_lists(sc: Any, limit: int) -> Any: return sc @staticmethod - def _merge(results: list[tuple[int, 'ToolResult']]) -> 'ToolResult': + def _merge(results: list[tuple[int, 'ToolResult']], errors: 'list[tuple[int, str]] | None' = None) -> 'ToolResult': # Deep-merge the per-project structured payloads into one schema-valid object (lists concatenated # across projects). Counters (e.g. bucket_counts, search total) are summed by _deep_merge, so they # keep reflecting the true totals even if the item lists get truncated below. + # Per-project failures (partial success) are surfaced as retry-hint notes in the text content, + # so the model can re-run just the failed project(s) via the project_ids filter. + error_notes = [ + mt.TextContent( + type='text', + text=f'project {pid} failed (retry with project_ids=[{pid}]): {msg}', + ) + for pid, msg in (errors or []) + ] merged_structured: Any = None per_project_counts: list[tuple[int, int]] = [] total_items = 0 @@ -1027,7 +1050,7 @@ def _merge(results: list[tuple[int, 'ToolResult']]) -> 'ToolResult': # Small enough: full detail with per-project text envelopes (attribution the model can read). if total_items <= MultiProjectMiddleware._FANOUT_MAX_ITEMS: - content: list[Any] = [] + content: list[Any] = list(error_notes) for project_id, result in results: content.append(mt.TextContent(type='text', text=f'=== project {project_id} ===')) content.extend(result.content or []) @@ -1043,7 +1066,7 @@ def _merge(results: list[tuple[int, 'ToolResult']]) -> 'ToolResult': f'search tool to find specific items.' ) truncated = MultiProjectMiddleware._truncate_lists(merged_structured, MultiProjectMiddleware._FANOUT_MAX_ITEMS) - return ToolResult(content=[mt.TextContent(type='text', text=note)], structured_content=truncated) + return ToolResult(content=error_notes + [mt.TextContent(type='text', text=note)], structured_content=truncated) def _to_python(data: Any, exclude_none: bool = True) -> Any | None: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 84d3d6196..48c2f3768 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1215,6 +1215,50 @@ def test_merge_small_keeps_full_detail(self) -> None: assert merged.structured_content == {'buckets': [0, 1, 0, 1, 2], 'total': 5} assert [c.text for c in merged.content] == ['=== project 11 ===', '2 items', '=== project 22 ===', '3 items'] + @pytest.mark.asyncio + async def test_fan_out_partial_failure_returns_successes_with_retry_hint(self) -> None: + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True) + + async def call_next(_): + if state[KeboolaClient.STATE_KEY] == 'client-22': + raise RuntimeError('boom-22') + return self._result('rows') + + with ( + patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + # Project 11 succeeded; project 22's failure is a retry hint, not a total failure. + assert result.structured_content == {'rows': ['rows']} + texts = [c.text for c in result.content] + assert any('project 22 failed' in t and 'project_ids=[22]' in t for t in texts) + + @pytest.mark.asyncio + async def test_fan_out_all_failed_raises_aggregate(self) -> None: + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, _ = self._ctx(scope, 'get_tables', read_only=True) + + async def call_next(_): + raise RuntimeError('down') + + with ( + patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), + ): + with pytest.raises(ToolError, match='failed for all 2 scoped'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + def test_merge_large_degrades_to_count_first(self, monkeypatch) -> None: # Lower the cap so a modest result trips the count-first path. monkeypatch.setattr(MultiProjectMiddleware, '_FANOUT_MAX_ITEMS', 3) From 9fc4aa6ae7e22152b54a6dc69aaa472b30dcaa76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 11:04:21 +0200 Subject: [PATCH 19/89] feat(PSGO-261): login --show-token to print the session access token Adds `login --show-token`, which prints the kbc_at_ session access token (in addition to storing it), so it can be passed as a header to a locally-run streamable-HTTP server for MPA over HTTP. Warns it expires in ~1 hour. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/cli.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 63fddd4b7..15fc11e34 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -76,6 +76,12 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: help='After the browser login, lease a Personal Access Token (kbc_pat_) over all accessible ' 'projects and print it. Requires an MFA code (--totp or --recovery).', ) + login_parser.add_argument( + '--show-token', + action='store_true', + help='Also print the session access token (kbc_at_) to stdout — e.g. to pass as a header to a ' + 'locally-run streamable-HTTP server. Note: it expires in ~1 hour.', + ) login_parser.add_argument('--totp', metavar='CODE', help='TOTP MFA code for the sudo elevation (--pat).') login_parser.add_argument( '--recovery', metavar='CODE', help='Recovery MFA code for the sudo elevation (--pat); alternative to --totp.' @@ -135,6 +141,7 @@ async def _run_login( totp: str | None = None, recovery: str | None = None, pat_name: str = 'keboola-mcp-server', + show_token: bool = False, ) -> None: """Establishes a stored session and, with ``pat=True``, leases a PAT. @@ -160,6 +167,10 @@ async def _run_login( remaining = max(0, int(tokens.expires_at - time.time())) if tokens else 0 print(f'\n✓ Session ready for {storage_api_url} (access token expires in ~{remaining}s).') + if show_token: + # Explicitly requested (e.g. to pass as a header to a local streamable-HTTP server). + print(f'\nAccess token (kbc_at_, expires in ~{remaining}s):\n\n {access_token}\n') + if pat: pat_token = await lease_pat( storage_api_url, @@ -204,6 +215,7 @@ async def run_server(args: list[str] | None = None) -> None: totp=getattr(parsed_args, 'totp', None), recovery=getattr(parsed_args, 'recovery', None), pat_name=getattr(parsed_args, 'pat_name', 'keboola-mcp-server'), + show_token=getattr(parsed_args, 'show_token', False), ) return From bb85cc40d56b9ab1080204e89b426d0b9733087a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 11:19:04 +0200 Subject: [PATCH 20/89] feat(PSGO-261): local HTTP server self-refreshes from the stored PKCE session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For a local streamable-HTTP request with no token in headers/env, SessionStateMiddleware now falls back to the stored login session (~/.keboola/mcp/credentials.json) via get_access_token — reading, refreshing, and persisting the rotated pair per request. This gives self-refreshing MPA over HTTP with a single `login`: run the HTTP server and connect with NO auth header; the server keeps the access token fresh and handles refresh- token rotation (which a static refresh-token header could not). No-op when a token is provided or on the deployed server (KBC_KUBERNETES_TOKEN_PATH set). Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/mcp.py | 26 ++++++++++++++++++++++++++ tests/test_mcp.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 5fdd8944d..f38c53fb8 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -276,6 +276,14 @@ async def on_request( if http_rq := get_http_request_or_none(): config = self.apply_request_config(http_rq, config, own_stack_storage_api_url=own_stack_storage_api_url) + # Local streamable-HTTP with no token supplied (no header / env): fall back to the stored + # PKCE session and keep it fresh. Lets you `login` once, run the HTTP server, and connect + # with NO auth header — the server reads ~/.keboola/mcp/credentials.json, refreshes the + # access token when it nears expiry, and persists the rotated pair (so refresh-token + # rotation is handled, which a static header could never do). No-op when a token is already + # provided or on the deployed server (KBC_KUBERNETES_TOKEN_PATH set). + config = await self._maybe_use_stored_session(config) + # In-conversation multi-project scope persists on the session across this per-request # state rebuild. Read it before the state is overwritten, refresh the stored token during # usage, and re-mint the scoped token when it nears expiry. With no scope and no preset @@ -394,6 +402,24 @@ def _is_local_programmatic(cls, config: Config) -> bool: and is_programmatic_token(config.storage_token) ) + @classmethod + async def _maybe_use_stored_session(cls, config: Config) -> Config: + """Populate the token from the stored PKCE session for a local, tokenless request. + + Only when: no token is set, a stack URL is known, and this is not the deployed server. Reads + (and refreshes + persists) the session leased by ``keboola-mcp-server login``. If there is no + stored session, leaves the config unchanged (a clear "no token" error surfaces downstream). + """ + if config.storage_token or not config.storage_api_url: + return config + if os.environ.get('KBC_KUBERNETES_TOKEN_PATH'): + return config + try: + access_token = await get_access_token(config.storage_api_url) + except RuntimeError: + return config + return dataclasses.replace(config, storage_token=access_token) + @classmethod async def _autolease_default_scope(cls, config: Config) -> 'SessionScope | None': """ diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 48c2f3768..4e9cf5079 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -849,6 +849,41 @@ async def test_happy_path_calls_resolver(self, monkeypatch) -> None: resolver.resolve.assert_awaited_once_with(subject_token='kbc_at_abc', project_id=42) +class TestMaybeUseStoredSession: + """Local HTTP with no token falls back to the stored PKCE session (PSGO-261).""" + + @pytest.mark.asyncio + async def test_no_token_loads_stored_session(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com') # no token + with patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_stored')): + out = await SessionStateMiddleware._maybe_use_stored_session(config) + assert out.storage_token == 'kbc_at_stored' + + @pytest.mark.asyncio + async def test_existing_token_is_noop(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_hdr') + with patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(side_effect=AssertionError('must not read'))): + out = await SessionStateMiddleware._maybe_use_stored_session(config) + assert out is config + + @pytest.mark.asyncio + async def test_deployed_is_noop(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com') + out = await SessionStateMiddleware._maybe_use_stored_session(config) + assert out is config + + @pytest.mark.asyncio + async def test_no_stored_session_is_noop(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com') + with patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(side_effect=RuntimeError('no creds'))): + out = await SessionStateMiddleware._maybe_use_stored_session(config) + assert out.storage_token is None + + class TestResolveLocalTokens: """SessionStateMiddleware keeps local tokens fresh and re-mints the scoped token (PSGO-261).""" From 09a92360fafa44e1eb9f343ee21eee44ef946580 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 12:21:23 +0200 Subject: [PATCH 21/89] fix(PSGO-261): call-time tool gating + drop surfaced active_project_id Two MPA UX corrections from review: 1) Don't hide tools before scope. Scope-first hiding relied on the client re-fetching the tool list after notifications/tools/list_changed, which Claude Code does not do mid-session, so unlocked tools never appeared without a reconnect. All tools now stay listed; the call-time ask-first gate blocks data tools until set_project_scope. The project_ids filter is injected only for a confirmed multi-project scope. 2) Remove active_project_id / is_active from the tool outputs. It always equals project_ids[0] and read as a chosen "primary" when scoped to many projects. The internal SessionScope.active_project_id (write / get_project_info / base-client default target) stays; it is just no longer surfaced. set_project_scope guidance now says writes target the first scoped project without naming an "active project". TOOLS.md regenerated; RFC increment-4 updated. Co-Authored-By: Claude Opus 4.8 --- TOOLS.md | 12 ++++-- feature_spec/pat_token_support/RFC.md | 57 +++++++++++-------------- src/keboola_mcp_server/mcp.py | 15 +++---- src/keboola_mcp_server/tools/project.py | 22 ++-------- tests/test_mcp.py | 14 ++++-- tests/tools/test_project.py | 7 +-- 6 files changed, 56 insertions(+), 71 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index 9b78363c6..a6f02efba 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -832,7 +832,8 @@ WORKFLOW: "type": "string" }, "value": { - "description": "Value to append to the list" + "description": "Value to append to the list", + "title": "Value" } }, "required": [ @@ -900,7 +901,8 @@ WORKFLOW: "type": "string" }, "value": { - "description": "New value to set" + "description": "New value to set", + "title": "Value" } }, "required": [ @@ -1152,7 +1154,8 @@ WORKFLOW: "type": "string" }, "value": { - "description": "Value to append to the list" + "description": "Value to append to the list", + "title": "Value" } }, "required": [ @@ -1220,7 +1223,8 @@ WORKFLOW: "type": "string" }, "value": { - "description": "New value to set" + "description": "New value to set", + "title": "Value" } }, "required": [ diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 9d5f51e6d..8f374d16d 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -544,41 +544,36 @@ Concrete, with a 2-project read: Follow-up (not in this increment): make fan-out concurrent, catch per-project errors into a per-project `{project_id, ok|error}` envelope, and stamp `source_project` on merged rows. -## Scope-first tool visibility (implemented) - -Replaces the "auto-lease a phantom active project + block data tools at call-time" UX with a -list-time gate: - -- A **programmatic (`kbc_*`) session** auto-leases an *unconfirmed* scope. While unconfirmed, - `on_list_tools` advertises **only** the scoping tools (`get_accessible_projects`, - `set_project_scope`). Data tools are not shown. -- `set_project_scope` confirms the scope and emits `notifications/tools/list_changed`, so the client - re-fetches and the full tool set appears. (Best-effort: the data tools are also no longer gated at - call-time once scope is confirmed, so a client that ignores `list_changed` still works.) -- The call-time ask-first gate is retained as defense in depth. - -### Backward compatibility — the gate is `kbc_*`-only -This is safe precisely because the gate keys on the presence of an (unconfirmed) `SessionScope`, -which **only** programmatic/PKCE sessions ever get: - -- **Legacy Storage-token sessions** (`1234-…`, not `kbc_*`) have no `SessionScope` → `on_list_tools` - is an unchanged passthrough: every tool is advertised from connect, no gate, single project. Zero - BC impact — no existing integration relying on "all tools at connect" breaks. -- **Deployed programmatic** sessions resolve to a single configured project (no auto-lease) → no - gate, single-project behavior. -- Only **local multi-project `kbc_*`** sessions (a new capability with no prior integrations) see the - hide-then-reveal flow, so there is nothing to break. - -The remaining risk is a client that ignores `tools/list_changed`; those still function (post-scope -calls are ungated) but won't visually refresh the tool list until reconnect. +## Tool gating: call-time, not list-time (why hide-then-reveal was reverted) + +We first tried **scope-first tool visibility**: while a programmatic session's scope was unconfirmed, +`on_list_tools` advertised only the scoping tools, and `set_project_scope` emitted +`notifications/tools/list_changed` to reveal the rest. **This does not work on Claude Code** (and +likely other clients): the client does **not re-fetch the tool list** after `list_changed` +mid-session, so the newly-unlocked tools never enter its inventory (and `ToolSearch` can't find them) +until a reconnect. Hiding therefore left the session stuck with only two tools. + +**Reverted to call-time gating** (robust on every client, no reconnect): +- **All tools stay listed** from connect. No hide. +- The **call-time ask-first gate** (`on_call_tool`) blocks data tools with a "confirm a scope first" + error until `set_project_scope` is called. After scoping, the already-listed tools just work. +- `set_project_scope` still emits `notifications/tools/list_changed` — now only meaningful because a + **confirmed multi-project scope adds the `project_ids` filter param** to read tools (a real schema + change); clients that honor it refresh, clients that don't still work (the param is optional). +- The `project_ids` filter is injected only for a **confirmed** scope of >1 project. + +This keeps the reviewer's other win (no phantom active project *before* a scope exists) without +depending on a client capability that isn't there. The "tools appear after scope" ideal is only +achievable on clients that re-fetch on `list_changed`; we don't rely on it. ## Decisions (increment 4) -- **Scope-first tool visibility** for `kbc_*` sessions; legacy Storage-token sessions unchanged - (gate keyed on `SessionScope` presence, which is equivalent to the `kbc_*` prefix). -- **No phantom active project** surfaced before a scope is confirmed (tools that need one are hidden). +- **Call-time gate, not list-time hiding** — hide-then-reveal needs client `list_changed` re-fetch + (absent in Claude Code mid-session), so all tools stay listed and data tools are gated at call time. +- **No phantom active project before a scope is confirmed**; after `set_project_scope` the + `active_project_id` is the write / `query_data`-default target and is surfaced intentionally. - **Fan-out stays**, with the relevance/latency critiques answered by the `project_ids` filter and a - (follow-up) concurrent loop; attribution + per-project error isolation are the tracked follow-ups. + (follow-up) concurrent loop; per-project error isolation is now implemented (partial results). - Fixed a latent bug: `set_project_scope` referenced `minted.read_only` on the exchange-failure path where `minted` is unbound — now uses the stored scope's `read_only`. diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index f38c53fb8..7c7002bec 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -933,15 +933,12 @@ async def on_list_tools( state = getattr(ctx.session, 'state', None) scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None - # Scope-first UX (programmatic / kbc_* sessions only — those are the ones that get an - # auto-leased scope): until the user confirms a scope via set_project_scope, expose ONLY the - # scoping tools. The full tool set is revealed after confirmation (set_project_scope emits - # notifications/tools/list_changed). Legacy Storage-token sessions have no SessionScope, so - # they keep advertising every tool unchanged — no BC impact. - if isinstance(scope, SessionScope) and not scope.confirmed: - return [t for t in tools if t.name in _BOOTSTRAP_TOOLS] - - if not (isinstance(scope, SessionScope) and len(scope.project_ids) > 1): + # NOTE: we intentionally do NOT hide data tools before a scope is confirmed. Hiding relied on + # the client re-fetching the tool list after notifications/tools/list_changed, which Claude Code + # (and others) don't do mid-session — that left the newly-unlocked tools invisible until a + # reconnect. Instead every tool stays listed and the call-time ask-first gate (on_call_tool) + # steers the user to set_project_scope first; once scoped, the already-listed tools just work. + if not (isinstance(scope, SessionScope) and scope.confirmed and len(scope.project_ids) > 1): return tools patched: list[Tool] = [] diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 3bc36cb67..69ef794b8 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -280,9 +280,6 @@ class AccessibleProject(BaseModel): name: str | None = Field(default=None, description='The project name.') role: str | None = Field(default=None, description='The user role in this project (e.g. "admin").') in_scope: bool = Field(default=False, description='Whether the session is currently scoped to this project.') - is_active: bool = Field( - default=False, description='Whether this is the active project (write / single-project tool target).' - ) sql_dialect: str | None = Field( default=None, description='The SQL dialect of the project ("Snowflake" or "BigQuery").' ) @@ -303,9 +300,6 @@ class AccessibleProjects(BaseModel): default=None, description='The projects the session is currently scoped to, or null if no scope has been confirmed yet.', ) - active_project_id: int | None = Field( - default=None, description='The active project that write operations and single-project tools target.' - ) read_only: bool | None = Field(default=None, description='Whether the current scoped token is read-only.') base_instructions: list[BaseInstructionGroup] | None = Field( default=None, @@ -323,10 +317,6 @@ class AccessibleProjects(BaseModel): class ProjectScope(BaseModel): project_ids: list[int] = Field(description='The projects the session is now scoped to.') read_only: bool = Field(description='Whether the scoped token is read-only.') - active_project_id: int | None = Field( - default=None, - description='The project that write operations and single-project tools target.', - ) llm_instruction: str = Field(description='Guidance for the assistant on the new scope.') @@ -374,7 +364,6 @@ async def get_accessible_projects( scope = ctx.session.state.get(SCOPE_KEY) scoped_ids = scope.project_ids if isinstance(scope, SessionScope) and scope.confirmed else None - active_id = scope.active_project_id if scoped_ids else None # Enrich each project with its SQL dialect (concurrently). Best-effort: a project whose verify # fails simply keeps sql_dialect=None rather than failing the whole listing. @@ -397,7 +386,6 @@ async def get_accessible_projects( name=p.name, role=p.role, in_scope=scoped_ids is not None and p.id in scoped_ids, - is_active=p.id == active_id, sql_dialect=dialects.get(p.id), ) for p in introspection.projects @@ -427,14 +415,13 @@ async def get_accessible_projects( ) else: instruction = ( - f'Session is currently scoped to {len(scoped_ids)} project(s); the active project is {active_id}. ' + f'Session is currently scoped to {len(scoped_ids)} project(s). ' 'Call "set_project_scope" to change the scope.' ) return AccessibleProjects( user_email=introspection.user_email, projects=projects, scoped_project_ids=scoped_ids, - active_project_id=active_id, read_only=scope.read_only if scoped_ids is not None else None, base_instructions=base_instructions, llm_instruction=instruction, @@ -505,12 +492,11 @@ async def set_project_scope( return ProjectScope( project_ids=ids, read_only=scope.read_only, - active_project_id=scope.active_project_id, llm_instruction=( ( - f'Session scoped to {len(ids)} projects. Read-only tools now return results per project. ' - f'Write operations target the active project {scope.active_project_id} only; to write to a ' - 'different project, re-scope or confirm with the user first.' + f'Session scoped to {len(ids)} projects. Read-only tools return results per project. ' + 'Write operations are not fanned out — they target the first scoped project; to write ' + 'elsewhere, re-scope to that project first (confirm with the user).' ) if multi else f'Session scoped to project {ids[0]}.' diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 4e9cf5079..ac0a92493 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1209,9 +1209,10 @@ async def call_next(_): assert 'project_ids' not in by_name['update_config'].parameters['properties'] @pytest.mark.asyncio - async def test_on_list_tools_scope_first_hides_data_tools_until_confirmed(self) -> None: - # Unconfirmed (auto-leased) scope: only the scoping tools are advertised; data tools appear - # after set_project_scope confirms. Legacy sessions (no scope) are covered by the passthrough. + async def test_on_list_tools_unconfirmed_scope_lists_all_tools(self) -> None: + # Data tools are NOT hidden before scope is confirmed: hiding relied on the client re-fetching + # after tools/list_changed, which Claude Code doesn't do mid-session. All tools stay listed; + # the call-time ask-first gate steers to set_project_scope instead. scope = SessionScope(project_ids=[11, 22], confirmed=False) context, _ = self._ctx(scope, 'x', read_only=True) @@ -1224,7 +1225,12 @@ async def call_next(_): ] tools = await MultiProjectMiddleware().on_list_tools(context, call_next) - assert {t.name for t in tools} == {'get_accessible_projects', 'set_project_scope'} + assert {t.name for t in tools} == { + 'get_accessible_projects', + 'set_project_scope', + 'get_tables', + 'update_config', + } @pytest.mark.asyncio async def test_on_list_tools_no_scope_is_passthrough(self) -> None: diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index a233f26f3..40c872b39 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -298,18 +298,16 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock (83, 'B', 'admin', 'Snowflake'), ] assert result.scoped_project_ids is None - assert result.active_project_id is None assert result.read_only is None assert result.base_instructions is None # not requested - assert all(not p.in_scope and not p.is_active for p in result.projects) + assert all(not p.in_scope for p in result.projects) # Once scoped, the current scope is surfaced on the projects and at the top level. mcp_context_client.session.state[SCOPE_KEY] = SessionScope(project_ids=[83], read_only=True, confirmed=True) result = await get_accessible_projects(mcp_context_client) assert result.scoped_project_ids == [83] - assert result.active_project_id == 83 assert result.read_only is True - assert [(p.id, p.in_scope, p.is_active) for p in result.projects] == [(18, False, False), (83, True, True)] + assert [(p.id, p.in_scope) for p in result.projects] == [(18, False), (83, True)] @pytest.mark.asyncio @@ -356,7 +354,6 @@ async def test_set_project_scope_subset_exchanges_and_stores( exch.assert_awaited_once_with(STACK, subject_token='kbc_at_parent', project_ids=[18, 83], read_only=False) assert result.project_ids == [18, 83] - assert result.active_project_id == 18 scope = mcp_context_client.session.state[SCOPE_KEY] assert scope.scoped_token == 'kbc_at_scoped' assert scope.project_ids == [18, 83] From d95e7009b9665aadb793ccc0bf05de9093507f00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 13:50:14 +0200 Subject: [PATCH 22/89] fix(PSGO-261): don't introspect on /list requests (tools/list 30s timeout) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Capability-discovery requests (tools/list, prompts/list, resources/list) ran the multi-project auto-lease, which calls token introspect against Connection. When that network call is slow, every /list hung until the client's 30s timeout (observed on all three list methods). Skip auto-lease for /list — listing capabilities must not depend on a Connection round trip. The scope is still established on the first real (non-list) tool call, so gating and fan-out are unaffected. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/auth_login.py | 16 ++++--- src/keboola_mcp_server/mcp.py | 69 +++++++++++++++++++--------- tests/test_mcp.py | 30 ++++++++++++ 3 files changed, 87 insertions(+), 28 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 9e80a538c..5e86d91b9 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -41,6 +41,10 @@ _REFRESH_SKEW_SECONDS = 60 _PAT_DEFAULT_EXPIRES_SECONDS = 30 * 24 * 60 * 60 # ~1 month _CREDENTIALS_PATH = Path.home() / '.keboola' / 'mcp' / 'credentials.json' +# Short connect timeout so an unreachable stack (e.g. VPN off — internal `.dev` stacks resolve to a +# private 10.x IP) fails in a few seconds with a clear ConnectTimeout instead of blocking the full +# window. A longer read timeout still tolerates a slow-but-reachable Connection. +_AUTH_TIMEOUT = httpx.Timeout(connect=5.0, read=30.0, write=10.0, pool=5.0) def _client_id() -> str: @@ -123,7 +127,7 @@ async def introspect_token( transport: httpx.AsyncBaseTransport | None = None, ) -> Introspection: """Enumerates the projects a programmatic token can reach via /v1/auth/token/introspect.""" - async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + async with httpx.AsyncClient(timeout=_AUTH_TIMEOUT, transport=transport) as client: response = await client.get( f'{_base_url(storage_api_url)}/{_INTROSPECT_PATH}', headers={'Authorization': f'Bearer {subject_token}'}, @@ -164,7 +168,7 @@ async def exchange_scoped_token( 'expiresIn': expires_in, 'scope': {'projects': [str(p) for p in project_ids], 'readOnly': read_only or None}, } - async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + async with httpx.AsyncClient(timeout=_AUTH_TIMEOUT, transport=transport) as client: response = await client.post( f'{_base_url(storage_api_url)}/{_EXCHANGE_PATH}', headers={'Authorization': f'Bearer {subject_token}'}, @@ -200,7 +204,7 @@ async def elevate_session( if bool(totp_code) == bool(recovery_code): raise ValueError('Provide exactly one of totp_code or recovery_code.') payload = {'totpCode': totp_code} if totp_code else {'recoveryCode': recovery_code} - async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + async with httpx.AsyncClient(timeout=_AUTH_TIMEOUT, transport=transport) as client: response = await client.post( f'{_base_url(storage_api_url)}/{_SUDO_PATH}', headers={'Authorization': f'Bearer {subject_token}'}, @@ -234,7 +238,7 @@ async def create_pat( 'expiresIn': expires_in, 'scope': {'projects': [str(p) for p in project_ids]}, } - async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + async with httpx.AsyncClient(timeout=_AUTH_TIMEOUT, transport=transport) as client: response = await client.post( f'{_base_url(storage_api_url)}/{_PAT_PATH}', headers={'Authorization': f'Bearer {subject_token}'}, @@ -302,7 +306,7 @@ async def exchange_code( 'redirectUri': redirect_uri, 'codeVerifier': code_verifier, } - async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + async with httpx.AsyncClient(timeout=_AUTH_TIMEOUT, transport=transport) as client: response = await client.post(f'{_base_url(storage_api_url)}/{_TOKEN_PATH}', json=payload) response.raise_for_status() return _parse_token_response(cast(dict, response.json())) @@ -315,7 +319,7 @@ async def refresh_tokens( transport: httpx.AsyncBaseTransport | None = None, ) -> TokenSet: """Exchanges a refresh token for a new (rotated) session token set.""" - async with httpx.AsyncClient(timeout=30.0, transport=transport) as client: + async with httpx.AsyncClient(timeout=_AUTH_TIMEOUT, transport=transport) as client: response = await client.post( f'{_base_url(storage_api_url)}/{_REFRESH_PATH}', json={'refreshToken': refresh_token} ) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 7c7002bec..cf3e4456d 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -31,7 +31,7 @@ from starlette.requests import Request from starlette.types import ASGIApp, Receive, Scope, Send -from keboola_mcp_server.auth_login import exchange_scoped_token, get_access_token, introspect_token +from keboola_mcp_server.auth_login import exchange_scoped_token, get_access_token, introspect_token, load_tokens from keboola_mcp_server.clients.auth_bridge import StorageTokenResolver, is_programmatic_token from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient @@ -276,23 +276,28 @@ async def on_request( if http_rq := get_http_request_or_none(): config = self.apply_request_config(http_rq, config, own_stack_storage_api_url=own_stack_storage_api_url) + # Capability-discovery requests (tools/list, prompts/list, resources/list) MUST be fast and + # network-free: a client fetches all three on connect, so any Connection round-trip here + # (token introspect, refresh, or scoped-exchange) makes connecting hang until the client's + # 30s timeout. For /list we do zero network in on_request — no auto-lease, no token refresh, + # no scoped re-mint — and use the stored session token as-is (no refresh). The scope and + # fresh tokens are established on the first real (non-list) tool call. + is_list = context.method.endswith('/list') + # Local streamable-HTTP with no token supplied (no header / env): fall back to the stored - # PKCE session and keep it fresh. Lets you `login` once, run the HTTP server, and connect - # with NO auth header — the server reads ~/.keboola/mcp/credentials.json, refreshes the - # access token when it nears expiry, and persists the rotated pair (so refresh-token - # rotation is handled, which a static header could never do). No-op when a token is already - # provided or on the deployed server (KBC_KUBERNETES_TOKEN_PATH set). - config = await self._maybe_use_stored_session(config) + # PKCE session. For non-list requests keep it fresh (refresh + persist rotation); for /list + # read it without a network refresh. No-op when a token is provided or on the deployed + # server (KBC_KUBERNETES_TOKEN_PATH set). + config = await self._maybe_use_stored_session(config, refresh=not is_list) # In-conversation multi-project scope persists on the session across this per-request - # state rebuild. Read it before the state is overwritten, refresh the stored token during - # usage, and re-mint the scoped token when it nears expiry. With no scope and no preset - # project, auto-lease ALL accessible projects (multi-project mode) so the session - # bootstraps without a hand-picked project and read tools fan out across everything. + # state rebuild. With no scope and no preset project, auto-lease ALL accessible projects + # (multi-project mode) so read tools fan out across everything — but never on /list. scope = self._read_persisted_scope(ctx.session) - if scope is None and not config.project_id: + if scope is None and not config.project_id and not is_list: scope = await self._autolease_default_scope(config) - config, scope = await self._resolve_local_tokens(config, scope) + if not is_list: + config, scope = await self._resolve_local_tokens(config, scope) # TODO: We could probably get rid of the 'state' attribute set on ctx.session and just # pass KeboolaClient and WorkspaceManager instances to a tool as extra parameters. @@ -301,7 +306,7 @@ async def on_request( # so that clients can discover available tools even when the configured branch ID doesn't # exist yet. For these requests the client is created without a branch ID. Otherwise, the branch is # validated via a SAPI call. - if context.method.endswith('/list'): + if is_list: if config.branch_id: LOG.info(f'Skipping branch validation for {context.method} request.') config = dataclasses.replace(config, branch_id=None) @@ -403,21 +408,29 @@ def _is_local_programmatic(cls, config: Config) -> bool: ) @classmethod - async def _maybe_use_stored_session(cls, config: Config) -> Config: + async def _maybe_use_stored_session(cls, config: Config, *, refresh: bool = True) -> Config: """Populate the token from the stored PKCE session for a local, tokenless request. - Only when: no token is set, a stack URL is known, and this is not the deployed server. Reads - (and refreshes + persists) the session leased by ``keboola-mcp-server login``. If there is no - stored session, leaves the config unchanged (a clear "no token" error surfaces downstream). + Only when: no token is set, a stack URL is known, and this is not the deployed server. With + ``refresh=True`` reads (and refreshes + persists) the session leased by + ``keboola-mcp-server login``. With ``refresh=False`` (capability-discovery /list requests) + reads the stored token WITHOUT any network refresh, so listing never blocks on Connection. + If there is no stored session, leaves the config unchanged. """ if config.storage_token or not config.storage_api_url: return config if os.environ.get('KBC_KUBERNETES_TOKEN_PATH'): return config - try: - access_token = await get_access_token(config.storage_api_url) - except RuntimeError: - return config + if refresh: + try: + access_token = await get_access_token(config.storage_api_url) + except RuntimeError: + return config + else: + tokens = load_tokens(config.storage_api_url) + if not tokens: + return config + access_token = tokens.access_token return dataclasses.replace(config, storage_token=access_token) @classmethod @@ -682,6 +695,18 @@ async def on_list_tools( self, context: MiddlewareContext[mt.ListToolsRequest], call_next: CallNext[mt.ListToolsRequest, list[Tool]] ) -> list[Tool]: tools = await call_next(context) + + # Feature/role filtering needs verify_token (a Connection round-trip). For a programmatic + # (kbc_*) session that has NOT confirmed a scope yet there is no single project to evaluate + # features against, and doing the call on tools/list would block connecting on a slow stack. + # Skip list-time filtering there and advertise the superset; the on_call_tool guards still + # enforce every feature/role/branch rule per project when a tool is actually invoked. + client = KeboolaClient.from_state(context.fastmcp_context.session.state) + scope = context.fastmcp_context.session.state.get(SCOPE_KEY) + scope_confirmed = isinstance(scope, SessionScope) and scope.confirmed + if is_programmatic_token(client.token) and not scope_confirmed: + return tools + token_info = await self.get_token_info(context.fastmcp_context) features = self.get_project_features(token_info) token_role = self.get_token_role(token_info).lower() diff --git a/tests/test_mcp.py b/tests/test_mcp.py index ac0a92493..8fe7d4141 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -416,6 +416,24 @@ async def call_next(_): assert name in result_names assert 'other_tool' in result_names + @pytest.mark.asyncio + async def test_list_tools_programmatic_pre_scope_skips_verify(self, mcp_context_client) -> None: + # Programmatic session with no confirmed scope: tools/list must not call verify_token (it would + # block connecting on a slow stack). Advertise the superset; on_call_tool still enforces. + client = KeboolaClient.from_state(mcp_context_client.session.state) + client.storage_client.verify_token = AsyncMock(side_effect=AssertionError('verify must not run pre-scope')) + + tools = [_tool('get_tables', read_only=True), _tool('create_flow'), _tool('get_semantic_context')] + + async def call_next(_): + return tools + + context = SimpleNamespace(fastmcp_context=mcp_context_client) + # programmatic token + no confirmed scope → filtering skipped, verify_token not called + with patch('keboola_mcp_server.mcp.is_programmatic_token', return_value=True): + result = await ToolsFilteringMiddleware().on_list_tools(context, call_next) + assert {t.name for t in result} == {'get_tables', 'create_flow', 'get_semantic_context'} + @pytest.mark.asyncio @pytest.mark.parametrize( ('token_role', 'bearer_token', 'hidden_tools', 'visible_tools'), @@ -868,6 +886,18 @@ async def test_existing_token_is_noop(self, monkeypatch) -> None: out = await SessionStateMiddleware._maybe_use_stored_session(config) assert out is config + @pytest.mark.asyncio + async def test_list_request_uses_stored_token_without_network_refresh(self, monkeypatch) -> None: + # /list must not do a network refresh: read the stored token as-is via load_tokens. + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com') + with ( + patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(side_effect=AssertionError('no network'))), + patch('keboola_mcp_server.mcp.load_tokens', return_value=SimpleNamespace(access_token='kbc_at_file')), + ): + out = await SessionStateMiddleware._maybe_use_stored_session(config, refresh=False) + assert out.storage_token == 'kbc_at_file' + @pytest.mark.asyncio async def test_deployed_is_noop(self, monkeypatch) -> None: monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') From 8d3ed0e5ee4c1f47a3427bf3ce7dbf1f0ff5af0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 15:02:37 +0200 Subject: [PATCH 23/89] fix(PSGO-261): add --no-stateless-http so multi-project scope persists over HTTP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-project scope lives in the MCP session state, but the streamable-HTTP server ran with stateless_http=True — every request was an independent session (Terminating session: None), so set_project_scope's state never reached the next call and data tools kept reporting "no scope confirmed". Over stdio the session is long-lived, which is why it worked there. Add a --stateless-http/--no-stateless-http flag (default stateless, preserving the deployed scaled behavior). Run a local MPA-over-HTTP server with --no-stateless-http so the session (and its scope) persists across requests. RFC gets a transport note. Co-Authored-By: Claude Opus 4.8 --- feature_spec/pat_token_support/RFC.md | 16 ++++++++++++++++ src/keboola_mcp_server/cli.py | 10 +++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 8f374d16d..49c53117a 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -605,3 +605,19 @@ default 200 total items across projects): This makes the multi-project path safe on humongous projects: it can never wedge the session, and it nudges toward the scalable access patterns (search / per-project drill-down) instead of bulk-listing. Follow-up: real `limit`/`offset` pagination on the enumerators, and concurrent fan-out. + +## Transport note: multi-project scope needs a stateful session + +Multi-project scope lives in the MCP **session** state (`ctx.session.state[SCOPE_KEY]`), read back on +each request. This only persists when the transport keeps the session alive across requests: + +- **stdio** — one long-lived session per process → scope persists (this is how MPA was developed/tested). +- **streamable-HTTP, stateless** (`stateless_http=True`, the deployed default for horizontal scaling) + → every request is an independent session (`Terminating session: None`), so `set_project_scope`'s + state never reaches the next call and data tools keep reporting "no scope confirmed". +- **streamable-HTTP, stateful** (`--no-stateless-http`) → the server issues an `Mcp-Session-Id` the + client echoes back, the `ServerSession` is reused, and scope persists. + +So to run MPA locally over HTTP: `keboola-mcp-server --transport streamable-http --no-stateless-http`. +Deployed multi-replica MPA over stateless HTTP would need a shared/sticky scope store (out of scope +here; deployed sessions today are single-project via the resolver + `KBC_PROJECT_ID`). diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 15fc11e34..66740899e 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -57,6 +57,14 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: parser.add_argument('--workspace-schema', metavar='STR', help='Keboola Storage API workspace schema.') parser.add_argument('--host', default='localhost', metavar='STR', help='The host to listen on.') parser.add_argument('--port', type=int, default=8000, metavar='INT', help='The port to listen on.') + parser.add_argument( + '--stateless-http', + action=argparse.BooleanOptionalAction, + default=True, + help='Streamable-HTTP session mode. Stateless (default) suits scaled/deployed servers where ' + 'any replica handles any request. Use --no-stateless-http for a local server so in-session ' + 'state — notably multi-project scope from set_project_scope — persists across requests.', + ) parser.add_argument('--log-config', type=pathlib.Path, metavar='PATH', help='Logging config file.') subparsers = parser.add_subparsers(dest='command') @@ -274,7 +282,7 @@ async def run_server(args: list[str] | None = None) -> None: http_app: StarletteWithLifespan = mcp_server.http_app( path='/', transport='streamable-http', - stateless_http=True, + stateless_http=parsed_args.stateless_http, ) mount_paths['/mcp'] = http_app transports.append('Streamable-HTTP') From a50671dbd9458bce36b388a35d963ac4b8db78d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 15:16:52 +0200 Subject: [PATCH 24/89] fix(PSGO-261): skip list-time feature filtering for all programmatic sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Post-scope, tools/list ran verify_token with the multi-project scoped token but no X-KBC-ProjectId, so Connection returned 401 on every list — the client retried and then disconnected. Feature/role filtering at list time is ill-defined for a multi-project token anyway (which project's features?), and tool CALLS already verify per-project with the active project_id. Skip list-time filtering for ALL programmatic (kbc_*) sessions (was: only pre-scope) and advertise the superset; on_call_tool still enforces feature/role/branch rules per project. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/mcp.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index cf3e4456d..acc730486 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -696,15 +696,15 @@ async def on_list_tools( ) -> list[Tool]: tools = await call_next(context) - # Feature/role filtering needs verify_token (a Connection round-trip). For a programmatic - # (kbc_*) session that has NOT confirmed a scope yet there is no single project to evaluate - # features against, and doing the call on tools/list would block connecting on a slow stack. - # Skip list-time filtering there and advertise the superset; the on_call_tool guards still - # enforce every feature/role/branch rule per project when a tool is actually invoked. + # Feature/role filtering needs verify_token (a Connection round-trip with a single project + # context). For a programmatic (kbc_*) session this doesn't work at list time: pre-scope there + # is no project (and the call would block connecting on a slow stack); post-scope the session + # holds a multi-project scoped token and verify without an X-KBC-ProjectId returns 401 — which + # made every tools/list fail and the client disconnect. So skip list-time filtering for ALL + # programmatic sessions and advertise the superset; the on_call_tool guards still enforce every + # feature/role/branch rule per project (with the right project_id) when a tool is invoked. client = KeboolaClient.from_state(context.fastmcp_context.session.state) - scope = context.fastmcp_context.session.state.get(SCOPE_KEY) - scope_confirmed = isinstance(scope, SessionScope) and scope.confirmed - if is_programmatic_token(client.token) and not scope_confirmed: + if is_programmatic_token(client.token): return tools token_info = await self.get_token_info(context.fastmcp_context) From 3de5d0d217f957b587e49777a9561929277b42df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 15:26:17 +0200 Subject: [PATCH 25/89] feat(PSGO-261): add logout command and login --force to switch tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refresh-first login reuses the stored session, so there was no way to switch user/token. Add: - `keboola-mcp-server logout [--api-url URL | --all]` — delete the stored PKCE session (one stack or all) so the next login starts fresh. Backed by new forget_tokens(). - `login --force` — drop any stored session and always run the browser flow. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/auth_login.py | 18 ++++++++++ src/keboola_mcp_server/cli.py | 51 ++++++++++++++++++++++++++-- tests/test_auth_login.py | 17 ++++++++++ 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 5e86d91b9..9582dc94d 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -435,6 +435,24 @@ def _forget(storage_api_url: str) -> None: _write_store(store) +def forget_tokens(storage_api_url: str | None = None) -> bool: + """Deletes the stored PKCE session — for one stack, or all when ``storage_api_url`` is None. + + Returns True if anything was removed. Used by the ``logout`` command so the next ``login`` starts + a fresh browser flow (e.g. to switch user/token) instead of refreshing the old session. + """ + store = _read_store() + if not store: + return False + if storage_api_url is None: + _write_store({}) + return True + if store.pop(_store_key(storage_api_url), None) is not None: + _write_store(store) + return True + return False + + # --- interactive browser login (not unit-tested; exercises a real browser + loopback) --- diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 66740899e..ecef2f988 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -100,6 +100,23 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: default='keboola-mcp-server', help='Name for the leased PAT (--pat).', ) + login_parser.add_argument( + '--force', + action='store_true', + help='Force a fresh browser login even if a valid stored session exists (e.g. to switch ' + 'user/token). Without it, login refreshes the existing session.', + ) + + logout_parser = subparsers.add_parser( + 'logout', + help='Delete the stored PKCE session so the next login starts fresh (switch user/token).', + ) + logout_parser.add_argument( + '--api-url', + metavar='URL', + help='Stack to log out of (default: KBC_STORAGE_API_URL). Use --all to clear every stack.', + ) + logout_parser.add_argument('--all', action='store_true', help='Delete stored sessions for all stacks.') return parser.parse_args(args) @@ -150,6 +167,7 @@ async def _run_login( recovery: str | None = None, pat_name: str = 'keboola-mcp-server', show_token: bool = False, + force: bool = False, ) -> None: """Establishes a stored session and, with ``pat=True``, leases a PAT. @@ -160,7 +178,7 @@ async def _run_login( With ``pat=True``, additionally leases a Personal Access Token over all accessible projects (introspect → sudo with the MFA code → create PAT) and prints it. """ - from keboola_mcp_server.auth_login import ensure_access_token, lease_pat, load_tokens + from keboola_mcp_server.auth_login import ensure_access_token, forget_tokens, lease_pat, load_tokens, perform_login storage_api_url = api_url or os.environ.get('KBC_STORAGE_API_URL') if not storage_api_url: @@ -169,8 +187,13 @@ async def _run_login( if pat and bool(totp) == bool(recovery): raise RuntimeError('Leasing a PAT (--pat) requires exactly one MFA code: pass --totp or --recovery.') - # Refresh-first, browser only when dead (interactive: this is the terminal `login` command). - access_token = await ensure_access_token(storage_api_url, allow_interactive=True) + if force: + # Drop any stored session and always run the browser flow (e.g. to switch user/token). + forget_tokens(storage_api_url) + access_token = (await perform_login(storage_api_url)).access_token + else: + # Refresh-first, browser only when dead (interactive: this is the terminal `login` command). + access_token = await ensure_access_token(storage_api_url, allow_interactive=True) tokens = load_tokens(storage_api_url) remaining = max(0, int(tokens.expires_at - time.time())) if tokens else 0 print(f'\n✓ Session ready for {storage_api_url} (access token expires in ~{remaining}s).') @@ -190,6 +213,23 @@ async def _run_login( print(f'\n✓ Personal Access Token (valid ~1 month, all accessible projects):\n\n {pat_token}\n') +async def _run_logout(api_url: str | None, *, all_stacks: bool = False) -> None: + """Deletes the stored PKCE session so the next login starts fresh.""" + from keboola_mcp_server.auth_login import forget_tokens + + if all_stacks: + removed = forget_tokens(None) + print('✓ Logged out of all stacks.' if removed else 'No stored sessions to remove.') + return + storage_api_url = api_url or os.environ.get('KBC_STORAGE_API_URL') + if not storage_api_url: + raise RuntimeError( + 'A Storage API URL is required for logout: pass --api-url, set KBC_STORAGE_API_URL, or use --all.' + ) + removed = forget_tokens(storage_api_url) + print(f'✓ Logged out of {storage_api_url}.' if removed else f'No stored session for {storage_api_url}.') + + async def run_server(args: list[str] | None = None) -> None: """Runs the MCP server in async mode.""" parsed_args = parse_args(args) @@ -224,9 +264,14 @@ async def run_server(args: list[str] | None = None) -> None: recovery=getattr(parsed_args, 'recovery', None), pat_name=getattr(parsed_args, 'pat_name', 'keboola-mcp-server'), show_token=getattr(parsed_args, 'show_token', False), + force=getattr(parsed_args, 'force', False), ) return + if parsed_args.command == 'logout': + await _run_logout(getattr(parsed_args, 'api_url', None), all_stacks=getattr(parsed_args, 'all', False)) + return + # Create config from the CLI arguments config = Config( storage_api_url=parsed_args.api_url, diff --git a/tests/test_auth_login.py b/tests/test_auth_login.py index 87ac7151c..88f573a97 100644 --- a/tests/test_auth_login.py +++ b/tests/test_auth_login.py @@ -18,6 +18,7 @@ ensure_access_token, exchange_code, exchange_scoped_token, + forget_tokens, get_access_token, introspect_token, lease_pat, @@ -81,6 +82,22 @@ async def test_refresh_tokens_rotates_pair() -> None: assert tokens.refresh_token == 'kbc_rt_new' +def test_forget_tokens_one_stack_and_all(creds_file: Path) -> None: + save_tokens(STACK, TokenSet('kbc_at_a', 'kbc_rt_a', expires_at=time.time() + 3600)) + other = 'https://connection.other.keboola.com' + save_tokens(other, TokenSet('kbc_at_b', 'kbc_rt_b', expires_at=time.time() + 3600)) + + # forget one stack leaves the other intact + assert forget_tokens(STACK) is True + assert load_tokens(STACK) is None + assert load_tokens(other) is not None + assert forget_tokens(STACK) is False # already gone + + # forget all clears everything + assert forget_tokens(None) is True + assert load_tokens(other) is None + + def test_save_and_load_round_trip_mode_600(creds_file: Path) -> None: ts = TokenSet(access_token='kbc_at_1', refresh_token='kbc_rt_1', expires_at=time.time() + 3600, session_id='s') save_tokens(STACK, ts) From 0d12f27e7437c13d195f5db065b1b2ad6d99687d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 15:38:49 +0200 Subject: [PATCH 26/89] fix(PSGO-261): stop fanning out argument-validation errors across projects A bad tool call (e.g. get_components with no component_ids) fails identically in every scoped project, so fan-out emitted N copies of the pydantic error plus a confusing "failed for all N projects" aggregate. Re-raise validation errors immediately on the first project so a single clean error surfaces. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/mcp.py | 7 +++++++ tests/test_mcp.py | 27 +++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index acc730486..0fa146ed1 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -18,6 +18,7 @@ import toon_format from fastmcp import Context, FastMCP from fastmcp.exceptions import ToolError +from fastmcp.exceptions import ValidationError as FastMCPValidationError from fastmcp.server import middleware as fmw from fastmcp.server.dependencies import get_http_request from fastmcp.server.middleware import CallNext, MiddlewareContext @@ -26,6 +27,7 @@ from mcp import types as mt from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from pydantic import BaseModel +from pydantic import ValidationError as PydanticValidationError from pydantic_core import to_json from starlette.applications import Starlette from starlette.requests import Request @@ -931,6 +933,11 @@ async def on_call_tool( # `except Exception` lets client cancellation propagate. try: results.append((project_id, await call_next(context))) + except (FastMCPValidationError, PydanticValidationError): + # Argument-level validation error: the same bad arguments fail identically in + # every project, so fanning out would emit N identical copies plus a confusing + # "failed for all N projects" aggregate. Abort and surface the single clean error. + raise except Exception as e: LOG.warning(f'Fan-out call failed for project {project_id}: {e}') errors.append((project_id, str(e))) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8fe7d4141..f4a020d0d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -10,6 +10,7 @@ from fastmcp.tools.tool import ToolResult from mcp import types as mt from pydantic import BaseModel, Field +from pydantic import ValidationError as PydanticValidationError from starlette.requests import Request from keboola_mcp_server.clients.client import KeboolaClient @@ -1330,6 +1331,32 @@ async def call_next(_): with pytest.raises(ToolError, match='failed for all 2 scoped'): await MultiProjectMiddleware().on_call_tool(context, call_next) + @pytest.mark.asyncio + async def test_fan_out_validation_error_raised_once_not_per_project(self) -> None: + # A bad argument (e.g. get_components with no component_ids) fails identically in every + # project, so it must surface as ONE clean validation error, not N copies + an aggregate. + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True) + calls = [] + + async def call_next(_): + calls.append(state[KeboolaClient.STATE_KEY]) + raise PydanticValidationError.from_exception_data('get_tables', []) + + with ( + patch.object( + MultiProjectMiddleware, + '_client_for_project', + AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), + ): + with pytest.raises(PydanticValidationError): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + # Aborted after the first project; not retried across the rest. + assert calls == ['client-11'] + def test_merge_large_degrades_to_count_first(self, monkeypatch) -> None: # Lower the cap so a modest result trips the count-first path. monkeypatch.setattr(MultiProjectMiddleware, '_FANOUT_MAX_ITEMS', 3) From 4029bd3286ed930821c584c167afb6b53cfee5b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 3 Jul 2026 15:41:51 +0200 Subject: [PATCH 27/89] fix(PSGO-261): stop polling workspace-creation job after terminal failure The workspace poll loop only treated 'success' as terminal; an 'error' (or warning/terminated/cancelled) status fell through to sleep-and-retry, so a job that failed at 6s kept being polled until the 300s timeout. Return None immediately on any terminal failure status so the caller surfaces the error right away. Co-Authored-By: Claude Opus 4.8 --- src/keboola_mcp_server/workspace.py | 6 ++++++ tests/test_workspace.py | 21 +++++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/src/keboola_mcp_server/workspace.py b/src/keboola_mcp_server/workspace.py index 0ec4ac7d3..1524133b2 100644 --- a/src/keboola_mcp_server/workspace.py +++ b/src/keboola_mcp_server/workspace.py @@ -786,6 +786,12 @@ async def _create_ws(self, *, timeout_sec: float = 300.0) -> _WspInfo | None: LOG.info(f'Created workspace: {workspace_id}') return await self._find_ws_by_id(workspace_id) + elif job_status in ('error', 'warning', 'terminated', 'cancelled', 'canceled'): + # Terminal failure states: the job will never reach 'success', so stop polling + # immediately instead of spinning until the timeout. + LOG.warning(f'Workspace creation job failed: job_id={job_id}, status={job_status}') + return None + elif duration > timeout_sec: LOG.info(f'Workspace creation timed out after {duration:.2f} seconds.') return None diff --git a/tests/test_workspace.py b/tests/test_workspace.py index 766fee096..cfcbb7251 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -171,6 +171,27 @@ async def test_workspace_creation_cleans_up_config_on_failure(): ) +@pytest.mark.asyncio +async def test_workspace_creation_stops_on_terminal_error_status(): + """A job that reaches a terminal failure status must stop polling at once, not spin to timeout.""" + mock_client = Mock(spec=KeboolaClient) + mock_client.branch_id = None + mock_storage_client = AsyncMock() + mock_client.storage_client = mock_storage_client + + mock_storage_client.verify_token.return_value = {'owner': {'defaultBackend': 'snowflake'}} + mock_storage_client.configuration_create.return_value = {'id': 'cfg-1', 'name': 'test'} + mock_storage_client.workspace_create_for_config.return_value = {'id': 999} + mock_storage_client.job_detail.return_value = {'status': 'error'} + + manager = WorkspaceManager(mock_client) + result = await manager._create_ws() + + assert result is None + # Polled exactly once — the terminal 'error' status short-circuits the loop. + mock_storage_client.job_detail.assert_awaited_once() + + @pytest.mark.asyncio @pytest.mark.parametrize( ('input_branch_id', 'has_sb_feature', 'workspace_schema', 'expected_bound_branch_id'), From abbd61c84ec8b7e38a7c30b6ed8da06b3a8e53c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 07:34:32 +0200 Subject: [PATCH 28/89] fix(PSGO-261): send PAT as Bearer token on jobs-queue, ai-service and sync-actions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The queue, ai-service and sync-actions clients were wired with the raw storage token, so a programmatic PAT/AT session sent the token as X-StorageAPI-Token — which those services reject. The queue accepts `Authorization: Bearer kbc_at_...` + X-KBC-ProjectId instead. Switch these three clients onto bearer_or_sapi_token (already used by storage/scheduler/ data-science/metastore), which falls back to the raw token for plain SAPI sessions, so run_job/get_jobs work end-to-end under PAT login. Consolidates the three near-identical per-client token-selection tests into one parametrized test covering all six bearer-capable clients. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/clients/client.py | 10 ++- tests/clients/test_client.py | 101 ++++++----------------- 2 files changed, 30 insertions(+), 81 deletions(-) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index 7cdd4acae..683205b89 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -193,10 +193,14 @@ def __init__( encryption_client=self._encryption_client, ) self._jobs_queue_client = JobsQueueClient.create( - root_url=queue_api_url, token=self._token, branch_id=branch_id, headers=self._headers, readonly=readonly + root_url=queue_api_url, + token=bearer_or_sapi_token, + branch_id=branch_id, + headers=self._headers, + readonly=readonly, ) self._ai_service_client = AIServiceClient.create( - root_url=ai_service_api_url, token=self._token, headers=self._headers, readonly=readonly + root_url=ai_service_api_url, token=bearer_or_sapi_token, headers=self._headers, readonly=readonly ) # Data-science (sandboxes-service) git-repo credential endpoints require an admin-context # token (CanManageAppRepoCredentials -> StorageApiToken::isAdminToken()). The OAuth bearer @@ -215,7 +219,7 @@ def __init__( ) self._sync_actions_client = SyncActionsClient.create( root_url=sync_actions_api_url, - token=self._token, + token=bearer_or_sapi_token, branch_id=branch_id, headers=self._headers, readonly=readonly, diff --git a/tests/clients/test_client.py b/tests/clients/test_client.py index 644eca11d..9c2284d29 100644 --- a/tests/clients/test_client.py +++ b/tests/clients/test_client.py @@ -407,41 +407,6 @@ async def test_with_branch_id_http_error( await keboola_client.with_branch_id('non-existent-branch') mock_client.get.assert_called_once() - @pytest.mark.parametrize( - ('bearer_token', 'storage_token', 'expected_scheduler_token'), - [ - ('oauth_bearer_123', 'sapi_token_456', 'Bearer oauth_bearer_123'), - (None, 'sapi_token_456', 'sapi_token_456'), - ('', 'sapi_token_456', 'sapi_token_456'), - ], - ids=['with_bearer_token', 'without_bearer_token', 'empty_bearer_token'], - ) - def test_scheduler_client_token_selection( - self, bearer_token: str | None, storage_token: str, expected_scheduler_token: str - ): - """Test SchedulerClient uses bearer token when available, falls back to storage token.""" - # Create KeboolaClient with different token configurations - client = KeboolaClient( - storage_api_url='https://connection.keboola.com', - storage_api_token=storage_token, - bearer_token=bearer_token, - ) - - # Verify scheduler client was initialized with correct token - # Check the headers of the underlying RawKeboolaClient - scheduler_headers = client.scheduler_client.raw_client.headers - - if expected_scheduler_token.startswith('Bearer '): - # Should use Authorization header for bearer token - assert 'Authorization' in scheduler_headers - assert scheduler_headers['Authorization'] == expected_scheduler_token - assert 'X-StorageAPI-Token' not in scheduler_headers - else: - # Should use X-StorageAPI-Token header for storage token - assert 'X-StorageAPI-Token' in scheduler_headers - assert scheduler_headers['X-StorageAPI-Token'] == expected_scheduler_token - assert 'Authorization' not in scheduler_headers - def test_metastore_client_url_derivation(self) -> None: client = KeboolaClient( storage_api_url='https://connection.canary-orion.keboola.dev', @@ -451,38 +416,25 @@ def test_metastore_client_url_derivation(self) -> None: assert client.metastore_client.raw_client.base_api_url == 'https://metastore.canary-orion.keboola.dev' assert client.metastore_client.raw_client.headers['X-StorageAPI-Token'] == 'sapi_token_456' + # All clients below use the bearer token (Authorization header) when one is available and fall + # back to the raw storage token (X-StorageAPI-Token) otherwise. The jobs-queue/ai-service/ + # sync-actions clients were switched onto the bearer token so PAT (kbc_at_/kbc_pat_) sessions + # work end-to-end — the queue accepts `Authorization: Bearer kbc_at_...` + X-KBC-ProjectId but + # rejects the PAT sent as X-StorageAPI-Token (PSGO-261). Data-science needs it for admin-context + # git-credential endpoints (AI-3398). @pytest.mark.parametrize( - ('bearer_token', 'storage_token', 'expected_metastore_token'), + 'client_attr', [ - ('oauth_bearer_123', 'sapi_token_456', 'Bearer oauth_bearer_123'), - (None, 'sapi_token_456', 'sapi_token_456'), - ('', 'sapi_token_456', 'sapi_token_456'), + 'scheduler_client', + 'metastore_client', + 'data_science_client', + 'jobs_queue_client', + 'ai_service_client', + 'sync_actions_client', ], - ids=['with_bearer_token', 'without_bearer_token', 'empty_bearer_token'], ) - def test_metastore_client_token_selection( - self, bearer_token: str | None, storage_token: str, expected_metastore_token: str - ): - """Test MetastoreClient uses bearer token when available, falls back to storage token.""" - client = KeboolaClient( - storage_api_url='https://connection.keboola.com', - storage_api_token=storage_token, - bearer_token=bearer_token, - ) - - metastore_headers = client.metastore_client.raw_client.headers - - if expected_metastore_token.startswith('Bearer '): - assert 'Authorization' in metastore_headers - assert metastore_headers['Authorization'] == expected_metastore_token - assert 'X-StorageAPI-Token' not in metastore_headers - else: - assert 'X-StorageAPI-Token' in metastore_headers - assert metastore_headers['X-StorageAPI-Token'] == expected_metastore_token - assert 'Authorization' not in metastore_headers - @pytest.mark.parametrize( - ('bearer_token', 'storage_token', 'expected_data_science_token'), + ('bearer_token', 'storage_token', 'expected_token'), [ ('oauth_bearer_123', 'sapi_token_456', 'Bearer oauth_bearer_123'), (None, 'sapi_token_456', 'sapi_token_456'), @@ -490,31 +442,24 @@ def test_metastore_client_token_selection( ], ids=['with_bearer_token', 'without_bearer_token', 'empty_bearer_token'], ) - def test_data_science_client_token_selection( - self, bearer_token: str | None, storage_token: str, expected_data_science_token: str + def test_client_bearer_token_selection( + self, client_attr: str, bearer_token: str | None, storage_token: str, expected_token: str ): - """DataScienceClient uses the bearer token when available, falls back to the storage token. - - The sandboxes-service git-repo credential endpoints require an admin-context token - (CanManageAppRepoCredentials -> isAdminToken()); the OAuth bearer token carries it while the - minted SAPI token does not (AI-3398). - """ + """Clients use the bearer token when available, falling back to the storage token.""" client = KeboolaClient( storage_api_url='https://connection.keboola.com', storage_api_token=storage_token, bearer_token=bearer_token, ) - data_science_headers = client.data_science_client.raw_client.headers + headers = getattr(client, client_attr).raw_client.headers - if expected_data_science_token.startswith('Bearer '): - assert 'Authorization' in data_science_headers - assert data_science_headers['Authorization'] == expected_data_science_token - assert 'X-StorageAPI-Token' not in data_science_headers + if expected_token.startswith('Bearer '): + assert headers.get('Authorization') == expected_token + assert 'X-StorageAPI-Token' not in headers else: - assert 'X-StorageAPI-Token' in data_science_headers - assert data_science_headers['X-StorageAPI-Token'] == expected_data_science_token - assert 'Authorization' not in data_science_headers + assert headers.get('X-StorageAPI-Token') == expected_token + assert 'Authorization' not in headers def test_flow_schema_cache_roundtrip(): From 3dfad9899216ab92b2fd5482d5d7bf4d402e36d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 07:52:10 +0200 Subject: [PATCH 29/89] docs(PSGO-261): reconcile as-built RFC with shipped per-service token wiring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The increment-3 per-service table and "Future PAT work" section still claimed jobs_queue/ai_service/sync_actions used the raw self._token and were 401 under PAT fan-out. Commit 5b8c65ed wired all three onto bearer_or_sapi_token, so mark them ✅ and rewrite the deferral note as resolved. Also correct the scoped- exchange body (project ids are sent as strings, not ints) and the increment-2 default-scope decision (auto-lease of all reachable projects, not single-project). Co-Authored-By: Claude Opus 4.8 (1M context) --- feature_spec/pat_token_support/RFC.md | 39 ++++++++++++++++----------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 49c53117a..88244a9ec 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -248,7 +248,8 @@ This is the discovery endpoint. It works for any programmatic token and is the s ``` POST {connection}/v1/auth/pat/exchange Headers: Authorization: Bearer -Body: { "expiresIn": null|, "scope": { "projects": [,...]|null, "readOnly": true|null } } +Body: { "expiresIn": null|, "scope": { "projects": ["",...]|null, "readOnly": true|null } } + # NB: project ids are sent as STRINGS — the exchange API 400s on integers (auth_login.py:166) 201: { "accessToken": "", "tokenType": "Bearer", "expiresIn": , @@ -311,9 +312,13 @@ the session's subject token for all downstream exchange/forwarding. *active* session token is the scoped one, so a tool can no longer reach an out-of-scope project even by bug. Strictly stronger than the original advisory model. - **D2 (extended) — `project_id` → project scope (a set).** Still explicit session state, never - silently derived. Default scope = `[KBC_PROJECT_ID]` / `X-KBC-ProjectId` (today's single-project - behavior, backward compatible). `get_accessible_projects` + `set_project_scope` replace the - previously-hypothetical "select-project tool"; introspect closes the open enumeration question. + silently derived. An explicit `KBC_PROJECT_ID` / `X-KBC-ProjectId` pins a single project + (backward compatible). _(as-built: when a local programmatic session sets **no** explicit project, + `SessionStateMiddleware._autolease_default_scope` introspects and defaults to **all** reachable + projects — multi-project by default — gated by an ask-first confirmation (`SessionScope.confirmed`, + `_BOOTSTRAP_TOOLS`); it is not single-project-by-default.)_ `get_accessible_projects` + + `set_project_scope` replace the previously-hypothetical "select-project tool"; introspect closes + the open enumeration question. - **D6 (new) — transparent fan-out via active-project indirection.** Tools take no `projects[]` arg; the dispatch wrapper swaps `active_project_id` and the per-project client cache. Multi-project results use a per-project envelope, never a semantic merge. Zero changes to the 43 `from_state` @@ -487,16 +492,18 @@ services that read that header work under header-narrowing. Current wiring (`cli | Metastore (semantic) | `bearer_or_sapi_token` | ✅ PAT/bearer-first, SAPI fallback (guarded); feature-gated, untested on stacks without `mcp-semantic-tooling` | | Data Science (sandboxes) | `bearer_or_sapi_token` + `X-KBC-ProjectId` | ✅ PAT + project header (verified: data-app create + deploy) | | Scheduler | `bearer_or_sapi_token` | ✅ bearer-first (writes only) | -| **Jobs Queue** | `self._token` (raw → `X-StorageAPI-Token`) | ❌ 401 under PAT fan-out — **future PAT work** | -| **AI Service** | `self._token` | ⚠️ raw token — **future PAT work** | -| **Sync Actions** | `self._token` | ⚠️ raw token — **future PAT work** | - -### Future PAT work (not in this increment) -`jobs_queue`, `ai_service`, and `sync_actions` still pass the raw `self._token`, so under a -PAT/multi-project session the satellite service rejects it (verified: `get_jobs` → 401 "Invalid -access token" from the Queue API). To support PAT they must use the bearer/PAT token (like -metastore/data-science/scheduler) or be handed a per-project minted token. Metastore and -data-science already use the bearer/PAT path, so no change was needed there. +| **Jobs Queue** | `bearer_or_sapi_token` + `X-KBC-ProjectId` | ✅ bearer/PAT-first, SAPI fallback | +| **AI Service** | `bearer_or_sapi_token` | ✅ bearer/PAT-first, SAPI fallback | +| **Sync Actions** | `bearer_or_sapi_token` + `X-KBC-ProjectId` | ✅ bearer/PAT-first, SAPI fallback | + +### Resolved: Queue / AI / Sync-Actions now speak bearer/PAT +`jobs_queue`, `ai_service`, and `sync_actions` originally passed the raw `self._token`, so under a +PAT/multi-project session the satellite service rejected it (`get_jobs` → 401 "Invalid access +token" from the Queue API). Fixed in commit `5b8c65ed`: all three are now wired with +`bearer_or_sapi_token` (`clients/client.py:169,184,190,209`), which forwards `Authorization: +Bearer ` for programmatic sessions and falls back to `X-StorageAPI-Token` for legacy SAPI — +matching metastore/data-science/scheduler. The queue accepts `Authorization: Bearer kbc_at_…` + +`X-KBC-ProjectId` (verified by hand against the Queue API). ## Decisions (increment 3) @@ -506,8 +513,8 @@ data-science already use the bearer/PAT path, so no change was needed there. - **`get_accessible_projects` is the multi-project bootstrap**: per-project dialect via token verify (no workspace), current scope surfaced, base instructions grouped by dialect behind `with_llm_instruction`. -- **Queue / AI / SyncActions PAT support is deferred** and documented above; metastore + data-science - already satisfy the PAT/bearer + `X-KBC-ProjectId` contract. +- **Queue / AI / SyncActions now use the bearer/PAT path** (commit `5b8c65ed`), joining + metastore + data-science in satisfying the PAT/bearer + `X-KBC-ProjectId` contract. # Extension: scope-first tool visibility + reviewer feedback (PSGO-261, increment 4) From 353c0a2b984cf8df41c2c90fd6096d2f04e66337 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 07:53:30 +0200 Subject: [PATCH 30/89] chore(PSGO-261): regenerate TOOLS.md with locked pydantic 2.13.4 check-tools-docs regenerates TOOLS.md from the locked deps; pydantic 2.13.4 drops redundant JSON-schema title fields, so the committed doc must match. Co-Authored-By: Claude Opus 4.8 (1M context) --- TOOLS.md | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index a6f02efba..9b78363c6 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -832,8 +832,7 @@ WORKFLOW: "type": "string" }, "value": { - "description": "Value to append to the list", - "title": "Value" + "description": "Value to append to the list" } }, "required": [ @@ -901,8 +900,7 @@ WORKFLOW: "type": "string" }, "value": { - "description": "New value to set", - "title": "Value" + "description": "New value to set" } }, "required": [ @@ -1154,8 +1152,7 @@ WORKFLOW: "type": "string" }, "value": { - "description": "Value to append to the list", - "title": "Value" + "description": "Value to append to the list" } }, "required": [ @@ -1223,8 +1220,7 @@ WORKFLOW: "type": "string" }, "value": { - "description": "New value to set", - "title": "Value" + "description": "New value to set" } }, "required": [ From c36fb8fdd0ea2fc11e1d05b004222bf60c2f7745 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 08:27:08 +0200 Subject: [PATCH 31/89] refactor(PSGO-261): apply post-rebase review cleanups (simplify/reuse/altitude) From the /simplify + ponytail + thermonuclear review passes, the safe, behavior-preserving wins: - Extract shared clients/base.py helpers normalize_storage_api_url() and read_service_account_jwt(), replacing the 3x duplicated URL validation and the 2x duplicated SA-JWT read across client.py, auth_bridge.py, auth_login.py. - Centralize the KBC_KUBERNETES_TOKEN_PATH am-I-deployed check into config.deployed_sa_token_path(), used by mcp.py (5 sites) and errors.py. - Compute MultiProjectMiddleware._largest_list_len once per project in _merge. - Harden the auth-bridge resolver: map a non-JSON / non-dict 200 body to a 502 StorageTokenExchangeError instead of letting it bubble as a 500 (Copilot #605). - Add KBC_PROJECT_ID / X-KBC-ProjectId config parsing regression cases (Copilot). Larger structural findings deferred to a focused follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/auth_login.py | 9 +++--- src/keboola_mcp_server/clients/auth_bridge.py | 24 ++++++++-------- src/keboola_mcp_server/clients/base.py | 28 +++++++++++++++++++ src/keboola_mcp_server/clients/client.py | 16 ++++------- src/keboola_mcp_server/config.py | 11 ++++++++ src/keboola_mcp_server/errors.py | 4 +-- src/keboola_mcp_server/mcp.py | 18 ++++++------ tests/test_config.py | 8 ++++++ 8 files changed, 79 insertions(+), 39 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 9582dc94d..19ab9273f 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -24,10 +24,12 @@ from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path from typing import cast -from urllib.parse import urlparse, urlunparse +from urllib.parse import urlparse import httpx +from keboola_mcp_server.clients.base import normalize_storage_api_url + LOG = logging.getLogger(__name__) DEFAULT_CLIENT_ID = 'keboola-cli-demo' @@ -53,10 +55,7 @@ def _client_id() -> str: def _base_url(storage_api_url: str) -> str: - parsed = urlparse(storage_api_url) - if not parsed.hostname or not parsed.hostname.startswith('connection.'): - raise ValueError(f'Invalid Keboola Storage API URL: {storage_api_url}') - return urlunparse(('https', parsed.hostname, '', '', '', '')) + return normalize_storage_api_url(storage_api_url) def _b64url(data: bytes) -> str: diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index e6491b59c..4ed730067 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -13,12 +13,12 @@ import logging from http import HTTPStatus -from pathlib import Path from typing import cast -from urllib.parse import urlparse, urlunparse import httpx +from keboola_mcp_server.clients.base import normalize_storage_api_url, read_service_account_jwt + LOG = logging.getLogger(__name__) _ACCESS_TOKEN_PREFIX = 'kbc_at_' @@ -78,20 +78,14 @@ def __init__( :param timeout: Optional HTTP timeout override. :param transport: Optional httpx transport (for testing). """ - parsed = urlparse(storage_api_url) - if not parsed.hostname or not parsed.hostname.startswith('connection.'): - raise ValueError(f'Invalid Keboola Storage API URL: {storage_api_url}') - self._base_url = urlunparse(('https', parsed.hostname, '', '', '', '')) + self._base_url = normalize_storage_api_url(storage_api_url) self._kubernetes_token_path = kubernetes_token_path self._timeout = timeout or httpx.Timeout(connect=5.0, read=30.0, write=10.0, pool=5.0) self._transport = transport def _read_sa_jwt(self) -> str: # Read per call — the kubelet rotates the projected token in place. - jwt = Path(self._kubernetes_token_path).read_text().strip() - if not jwt: - raise ValueError(f'Kubernetes ServiceAccount token file is empty: {self._kubernetes_token_path}') - return jwt + return read_service_account_jwt(self._kubernetes_token_path) async def resolve(self, *, subject_token: str, project_id: int) -> str: """ @@ -130,8 +124,14 @@ async def resolve(self, *, subject_token: str, project_id: int) -> str: status_code=mapped, ) - body = cast(dict, response.json()) - storage_token = body.get('storageToken') + try: + body = response.json() + except ValueError: + raise StorageTokenExchangeError( + 'Auth-bridge token exchange returned a non-JSON body.', + status_code=int(HTTPStatus.BAD_GATEWAY), + ) from None + storage_token = body.get('storageToken') if isinstance(body, dict) else None if not storage_token: raise StorageTokenExchangeError( 'Auth-bridge token exchange returned no storageToken.', diff --git a/src/keboola_mcp_server/clients/base.py b/src/keboola_mcp_server/clients/base.py index 1d1c14341..8975b3cce 100644 --- a/src/keboola_mcp_server/clients/base.py +++ b/src/keboola_mcp_server/clients/base.py @@ -1,7 +1,9 @@ import json import logging from http import HTTPStatus +from pathlib import Path from typing import Any, Union, cast +from urllib.parse import urlparse, urlunparse import httpx from httpx_retries import Retry, RetryTransport @@ -17,6 +19,32 @@ LOG = logging.getLogger(__name__) +def normalize_storage_api_url(storage_api_url: str) -> str: + """ + Validates a Keboola Storage API URL and returns its canonical ``https://connection.`` base. + + :raises ValueError: if the host is missing or is not a ``connection.*`` host. + """ + parsed = urlparse(storage_api_url) + if not parsed.hostname or not parsed.hostname.startswith('connection.'): + raise ValueError(f'Invalid Keboola Storage API URL: {storage_api_url}') + return urlunparse(('https', parsed.hostname, '', '', '', '')) + + +def read_service_account_jwt(path: str) -> str: + """ + Reads the projected Kubernetes ServiceAccount JWT from ``path``. + + Read per call so kubelet rotation of the projected token is honored. + + :raises ValueError: if the token file is empty. + """ + jwt = Path(path).read_text().strip() + if not jwt: + raise ValueError(f'Kubernetes ServiceAccount token file is empty: {path}') + return jwt + + class RawKeboolaClient: """ Raw async client for Keboola services. diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index 683205b89..673c1a08a 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -2,13 +2,13 @@ import logging from collections.abc import Mapping, Sequence -from pathlib import Path -from typing import Any, Literal, TypeVar +from typing import Any, Literal, TypeVar, cast from urllib.parse import urlparse, urlunparse import httpx from keboola_mcp_server.clients.ai_service import AIServiceClient +from keboola_mcp_server.clients.base import normalize_storage_api_url, read_service_account_jwt from keboola_mcp_server.clients.data_science import DataScienceClient from keboola_mcp_server.clients.encryption import EncryptionClient from keboola_mcp_server.clients.jobs_queue import JobsQueueClient @@ -164,12 +164,8 @@ def __init__( # Mirrors _features_cache: fetched once per session so it is never stale across runs. self._flow_schema_cache: dict[str, JsonDict] = {} - sapi_url_parsed = urlparse(storage_api_url) - if not sapi_url_parsed.hostname or not sapi_url_parsed.hostname.startswith('connection.'): - raise ValueError(f'Invalid Keboola Storage API URL: {storage_api_url}') - - self._hostname_suffix = sapi_url_parsed.hostname.split('connection.')[1] - self._storage_api_url = urlunparse(('https', f'connection.{self._hostname_suffix}', '', '', '', '')) + self._storage_api_url = normalize_storage_api_url(storage_api_url) + self._hostname_suffix = cast(str, urlparse(self._storage_api_url).hostname).split('connection.')[1] metastore_api_url = urlunparse(('https', f'metastore.{self._hostname_suffix}', '', '', '', '')) queue_api_url = urlunparse(('https', f'queue.{self._hostname_suffix}', '', '', '', '')) ai_service_api_url = urlunparse(('https', f'ai.{self._hostname_suffix}', '', '', '', '')) @@ -315,9 +311,7 @@ def step_up_storage_client(self, kubernetes_token_path: str) -> 'AsyncStorageCli ) return self._storage_client - jwt = Path(kubernetes_token_path).read_text().strip() - if not jwt: - raise ValueError(f'Kubernetes ServiceAccount token file is empty: {kubernetes_token_path}') + jwt = read_service_account_jwt(kubernetes_token_path) headers = dict(self._headers or {}) headers['X-Kubernetes-Authorization'] = f'Bearer {jwt}' diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 22768403b..0c36580e3 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -15,6 +15,17 @@ Transport = Literal['stdio', 'streamable-http', 'http-compat/streamable-http'] +def deployed_sa_token_path() -> str | None: + """ + Path to the deployed server's projected Kubernetes ServiceAccount token, or None when running locally. + + The presence of the ``KBC_KUBERNETES_TOKEN_PATH`` env var is the single signal that this process is + the Keboola-deployed MCP server (able to reach the auth-bridge resolver) rather than a local session. + Read from the process environment only, never from per-request config. + """ + return os.environ.get('KBC_KUBERNETES_TOKEN_PATH') + + @dataclass(frozen=True) class Config: """Server configuration.""" diff --git a/src/keboola_mcp_server/errors.py b/src/keboola_mcp_server/errors.py index bb6b64cac..b048f6705 100644 --- a/src/keboola_mcp_server/errors.py +++ b/src/keboola_mcp_server/errors.py @@ -1,7 +1,6 @@ import inspect import json import logging -import os import time from collections.abc import Callable, Mapping from functools import wraps @@ -21,6 +20,7 @@ from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.clients.storage import StorageEventType +from keboola_mcp_server.config import deployed_sa_token_path from keboola_mcp_server.mcp import CONVERSATION_ID, ServerState, get_http_request_or_none LOG = logging.getLogger(__name__) @@ -132,7 +132,7 @@ async def _trigger_event( # the user's own client otherwise — this code runs in a `finally:` block that swallows its # errors, so it must not depend on anything failing loudly. storage_client = client.storage_client - if kubernetes_token_path := os.environ.get('KBC_KUBERNETES_TOKEN_PATH'): + if kubernetes_token_path := deployed_sa_token_path(): storage_client = client.step_up_storage_client(kubernetes_token_path) resp = await storage_client.trigger_event( message=message, diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 0fa146ed1..61f1ea0c7 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -8,7 +8,6 @@ import asyncio import dataclasses import logging -import os import textwrap import time from collections.abc import Awaitable, Callable, Iterable @@ -37,7 +36,7 @@ from keboola_mcp_server.clients.auth_bridge import StorageTokenResolver, is_programmatic_token from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient -from keboola_mcp_server.config import Config, ServerRuntimeInfo, is_same_stack +from keboola_mcp_server.config import Config, ServerRuntimeInfo, deployed_sa_token_path, is_same_stack from keboola_mcp_server.oauth import ProxyAccessToken from keboola_mcp_server.tools.constants import MODIFY_FLOW_TOOL_NAME, SEMANTIC_TOOLS_TAG, UPDATE_FLOW_TOOL_NAME from keboola_mcp_server.workspace import WorkspaceManager @@ -403,7 +402,7 @@ def _read_persisted_scope(session: Any) -> 'SessionScope | None': def _is_local_programmatic(cls, config: Config) -> bool: """True for a local (non-deployed) session carrying a Keboola programmatic token.""" return ( - not os.environ.get('KBC_KUBERNETES_TOKEN_PATH') + not deployed_sa_token_path() and bool(config.storage_token) and bool(config.storage_api_url) and is_programmatic_token(config.storage_token) @@ -421,7 +420,7 @@ async def _maybe_use_stored_session(cls, config: Config, *, refresh: bool = True """ if config.storage_token or not config.storage_api_url: return config - if os.environ.get('KBC_KUBERNETES_TOKEN_PATH'): + if deployed_sa_token_path(): return config if refresh: try: @@ -522,7 +521,7 @@ async def _exchange_programmatic_token(cls, config: Config) -> str: environment only, never from per-request config). A project id is required because a programmatic token is not project-bound. """ - kubernetes_token_path = os.environ.get('KBC_KUBERNETES_TOKEN_PATH') + kubernetes_token_path = deployed_sa_token_path() if not kubernetes_token_path: raise ValueError( 'Received a Keboola programmatic token (kbc_at_/kbc_pat_) but KBC_KUBERNETES_TOKEN_PATH ' @@ -580,7 +579,7 @@ async def create_session_state( bearer_token = config.bearer_token extra_headers: dict[str, Any] = {} if is_programmatic_token(storage_token): - if os.environ.get('KBC_KUBERNETES_TOKEN_PATH'): + if deployed_sa_token_path(): # Deployed: exchange the programmatic token (kbc_at_/kbc_pat_) for the project's # legacy Storage token via the auth-bridge resolver, then use it downstream unchanged. storage_token = await cls._exchange_programmatic_token(config) @@ -615,7 +614,7 @@ async def create_session_state( # overridable per request. An unforgeable path is not enough on its own, because the # destination can come from a header — `KeboolaClient.step_up_storage_client()` # therefore attaches the JWT only when the target is this server's own stack. - kubernetes_token_path = os.environ.get('KBC_KUBERNETES_TOKEN_PATH') + kubernetes_token_path = deployed_sa_token_path() workspace_manager = await WorkspaceManager.create( client, config.workspace_schema, kubernetes_token_path=kubernetes_token_path ) @@ -1096,8 +1095,9 @@ def _merge(results: list[tuple[int, 'ToolResult']], errors: 'list[tuple[int, str total_items = 0 for project_id, result in results: sc = result.structured_content - per_project_counts.append((project_id, MultiProjectMiddleware._largest_list_len(sc))) - total_items += MultiProjectMiddleware._largest_list_len(sc) + item_count = MultiProjectMiddleware._largest_list_len(sc) + per_project_counts.append((project_id, item_count)) + total_items += item_count if sc is not None: merged_structured = ( sc if merged_structured is None else MultiProjectMiddleware._deep_merge(merged_structured, sc) diff --git a/tests/test_config.py b/tests/test_config.py index 199d63850..249d5ec1e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -34,6 +34,14 @@ class TestConfig: {'X-Conversation-ID': '1234'}, Config(conversation_id='1234'), ), + ( + {'KBC_PROJECT_ID': '1888'}, + Config(project_id='1888'), + ), + ( + {'X-KBC-ProjectId': '1888'}, + Config(project_id='1888'), + ), ], ) def test_from_dict(self, d: Mapping[str, str], expected: Config) -> None: From 44a3a8f38e0f354a2a070ceeabb854ce8bd9090f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 10:25:13 +0200 Subject: [PATCH 32/89] docs(PSGO-261): document browser login for local setup; drop legacy token paste Rewrite the Local Setup section around the one-time browser `login` flow (no token to create or paste): `login`/`logout`/`--force`/`--show-token`, the ~/.keboola/mcp/credentials.json session store, and starting the server with only KBC_STORAGE_API_URL. Remove the KBC_STORAGE_TOKEN / KBC_WORKSPACE_SCHEMA paste-a-token instructions from every local config example (Claude, Cursor, WSL, local-dev, manual CLI). Keep a token path only for headless containers/CI (access/personal access token + KBC_PROJECT_ID, or per-request headers). Also refresh CLAUDE.md: tox runs five environments (add isort, which was missing from the list). Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 66 ++++++++++++++++++++++++++----------------------------- 1 file changed, 31 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index bee7d0788..bdb4dcf9f 100644 --- a/README.md +++ b/README.md @@ -124,40 +124,41 @@ For detailed documentation, see [developers.keboola.com/integrate/mcp/#tool-auth ## Local MCP Server Setup (Custom or Dev Way) -Run the MCP server on your own machine for full control and easy development. Choose this when you want to customize tools, debug locally, or iterate quickly. You’ll clone the repo, set Keboola credentials via environment variables or headers depending on the server transport, install dependencies, and start the server. This approach offers maximum flexibility (custom tools, local logging, offline iteration) but requires manual setup and you manage updates and secrets yourself. +Run the MCP server on your own machine for full control and easy development. Choose this when you want to customize tools, debug locally, or iterate quickly. You’ll install the server, authenticate (a one-time browser login — no token to paste), and start it. This approach offers maximum flexibility (custom tools, local logging, offline iteration) but requires manual setup and you manage updates and secrets yourself. The server supports multiple **transport** options, which can be selected by providing the `--transport ` argument when starting the server: - `stdio` - Default when `--transport` is not specified. Standard input/output, typically used for local deployment with a single client. - `streamable-http` - Runs the server remotely over HTTP with a bidirectional streaming channel, allowing the client and server to continuously exchange messages. Connect via /mcp (e.g., http://localhost:8000/mcp). - `http-compat` - An alias for `streamable-http`, kept for backwards compatibility. -For client–server communication, Keboola credentials must be provided to enable working with your project in your Keboola Region. The following are required: `KBC_STORAGE_TOKEN`, `KBC_STORAGE_API_URL`, `KBC_WORKSPACE_SCHEMA` and optionally `KBC_BRANCH_ID`. You can provide these in two ways: -- For personal use (mainly with stdio transport): set the environment variables before starting the server. All requests will reuse these predefined credentials. -- For multi-user use: include the variables in the request headers so that each request uses the credentials provided with it. +To work with your Keboola project the server needs two things: your **Keboola Region** (`KBC_STORAGE_API_URL`) and a way to **authenticate**. The recommended way is a one-time browser **login** — you never create, copy, or paste a token. Optionally set `KBC_BRANCH_ID` to work inside a development branch. Two of the variables are not taken from the request headers: - `KBC_STORAGE_API_URL`: a server that was started with its own Storage API URL (the `--api-url` parameter or the `KBC_STORAGE_API_URL` environment variable) only serves that one Keboola stack. An `X-Storage-Api-Url` header asking for a different host is ignored (a warning is logged) — the server keeps its own URL for the request. Start the server without a Storage API URL of its own if you want each request to choose its stack. - `KBC_KUBERNETES_TOKEN_PATH` (deployed servers only, see [docs/kubernetes-sa-auth.md](docs/kubernetes-sa-auth.md)): read from the environment only, never from a header. +### Logging in -### KBC_STORAGE_TOKEN +Sign in once with your browser; the server stores the session and refreshes it automatically, so there are no tokens to manage: -This is your authentication token for Keboola: - -For instructions on how to create and manage Storage API tokens, refer to the [official Keboola documentation](https://help.keboola.com/management/project/tokens/). - -**Note**: If you want the MCP server to have limited access, use custom storage token, if you want the MCP to access everything in your project, use the master token. +```bash +uvx keboola_mcp_server login --api-url https://connection.YOUR_REGION.keboola.com +``` -### KBC_WORKSPACE_SCHEMA +This opens your browser to sign in to Keboola and pick a project, then saves the session to `~/.keboola/mcp/credentials.json` (readable only by you, one entry per stack). Afterwards, start the server with only `KBC_STORAGE_API_URL` set — no token required. -This identifies your workspace in Keboola and is used for SQL queries. However, this is **only required if you're using a custom storage token** instead of the Master Token: +| Command | What it does | +|---------|--------------| +| `login --api-url ` | Sign in to a stack | +| `login --force` | Sign in again / switch account | +| `login --show-token` | Print the current session token (debugging) | +| `logout [--api-url ] [--all]` | Remove the stored session for a stack (or all stacks) | -- If using [Master Token](https://help.keboola.com/management/project/tokens/#master-tokens): The workspace is created automatically behind the scenes -- If using [custom storage token](https://help.keboola.com/management/project/tokens/#limited-tokens): Follow this [Keboola guide](https://help.keboola.com/tutorial/manipulate/workspace/) to get your KBC_WORKSPACE_SCHEMA +When you start the server over **stdio in an interactive terminal** with no stored session, it runs this browser login automatically on first start. MCP clients (Claude, Cursor, …) launch the server in the background where a browser can't open, so run `login` once yourself first. -**Note**: When creating a workspace manually, check Grant read-only access to all Project data option +#### Authenticating without a browser -**Note**: KBC_WORKSPACE_SCHEMA is called Dataset Name in BigQuery workspaces, you simply click connect and copy the Dataset Name +For containers or CI where a browser login isn't possible, provide a Keboola [access or personal access token](https://help.keboola.com/management/project/tokens/) directly — set `KBC_STORAGE_TOKEN` (env var) or send the `X-StorageAPI-Token` header — together with `KBC_PROJECT_ID` (or the `X-KBC-ProjectId` header) to select the project. On HTTP transports these can be supplied per request as headers, so each request carries its own credentials. ### KBC_STORAGE_API_URL (Keboola Region) @@ -223,10 +224,14 @@ There are four ways to use the Keboola MCP Server, depending on your needs: ### Option A: Integrated Mode (Recommended) -In this mode, Claude or Cursor automatically starts the MCP server for you. **You do not need to run any commands in your terminal**. +In this mode, Claude or Cursor automatically starts the MCP server for you. -1. Configure your MCP client (Claude/Cursor) with the appropriate settings -2. The client will automatically launch the MCP server when needed +1. **Log in once** in a terminal so a session is stored (the client launches the server in the background, where a browser can't open): + ```bash + uvx keboola_mcp_server login --api-url https://connection.YOUR_REGION.keboola.com + ``` +2. Configure your MCP client (Claude/Cursor) with the settings below — only `KBC_STORAGE_API_URL` is needed. +3. The client will automatically launch the MCP server when needed. #### Claude Desktop Configuration @@ -242,8 +247,6 @@ In this mode, Claude or Cursor automatically starts the MCP server for you. **Yo "args": ["keboola_mcp_server --transport "], "env": { "KBC_STORAGE_API_URL": "https://connection.YOUR_REGION.keboola.com", - "KBC_STORAGE_TOKEN": "your_keboola_storage_token", - "KBC_WORKSPACE_SCHEMA": "your_workspace_schema", "KBC_BRANCH_ID": "your_branch_id_optional" } } @@ -270,8 +273,6 @@ Config file locations: "args": ["keboola_mcp_server --transport "], "env": { "KBC_STORAGE_API_URL": "https://connection.YOUR_REGION.keboola.com", - "KBC_STORAGE_TOKEN": "your_keboola_storage_token", - "KBC_WORKSPACE_SCHEMA": "your_workspace_schema", "KBC_BRANCH_ID": "your_branch_id_optional" } } @@ -295,8 +296,6 @@ When running the MCP server from Windows Subsystem for Linux with Cursor AI, use "bash", "-c '", "export KBC_STORAGE_API_URL=https://connection.YOUR_REGION.keboola.com &&", - "export KBC_STORAGE_TOKEN=your_keboola_storage_token &&", - "export KBC_WORKSPACE_SCHEMA=your_workspace_schema &&", "export KBC_BRANCH_ID=your_branch_id_optional &&", "/snap/bin/uvx keboola_mcp_server --transport ", "'" @@ -324,8 +323,6 @@ For developers working on the MCP server code itself: ], "env": { "KBC_STORAGE_API_URL": "https://connection.YOUR_REGION.keboola.com", - "KBC_STORAGE_TOKEN": "your_keboola_storage_token", - "KBC_WORKSPACE_SCHEMA": "your_workspace_schema", "KBC_BRANCH_ID": "your_branch_id_optional" } } @@ -338,11 +335,9 @@ For developers working on the MCP server code itself: You can run the server manually in a terminal for testing or debugging: ```bash -# Set environment variables +# Sign in once (stores a session under ~/.keboola/mcp), then start the server. export KBC_STORAGE_API_URL=https://connection.YOUR_REGION.keboola.com -export KBC_STORAGE_TOKEN=your_keboola_storage_token -export KBC_WORKSPACE_SCHEMA=your_workspace_schema -export KBC_BRANCH_ID=your_branch_id_optional +uvx keboola_mcp_server login --api-url "$KBC_STORAGE_API_URL" uvx keboola_mcp_server --transport streamable-http ``` @@ -355,6 +350,8 @@ uvx keboola_mcp_server --transport streamable-http ### Option D: Using Docker +A container can't open a browser, so authenticate with a token (see [Authenticating without a browser](#authenticating-without-a-browser)): set `KBC_STORAGE_TOKEN` to a Keboola access/personal access token and `KBC_PROJECT_ID` to the target project. (Over HTTP you can instead pass `X-StorageAPI-Token` / `X-KBC-ProjectId` headers per request and omit these.) + ```shell docker pull keboola/mcp-server:latest @@ -364,8 +361,8 @@ docker run \ -it \ -p 127.0.0.1:8000:8000 \ -e KBC_STORAGE_API_URL="https://connection.YOUR_REGION.keboola.com" \ - -e KBC_STORAGE_TOKEN="YOUR_KEBOOLA_STORAGE_TOKEN" \ - -e KBC_WORKSPACE_SCHEMA="YOUR_WORKSPACE_SCHEMA" \ + -e KBC_STORAGE_TOKEN="YOUR_KEBOOLA_TOKEN" \ + -e KBC_PROJECT_ID="YOUR_PROJECT_ID" \ -e KBC_BRANCH_ID="YOUR_BRANCH_ID_OPTIONAL" \ keboola/mcp-server:latest \ --transport streamable-http \ @@ -437,8 +434,7 @@ For a complete list of available tools with detailed descriptions, parameters, a | Issue | Solution | |-------|----------| -| **Authentication Errors** | Verify `KBC_STORAGE_TOKEN` is valid | -| **Workspace Issues** | Confirm `KBC_WORKSPACE_SCHEMA` is correct | +| **Authentication Errors** | Re-run `keboola_mcp_server login` (or, if authenticating with a token, verify the token and `KBC_PROJECT_ID`) | | **Connection Timeout** | Check network connectivity | ## Development From d88d949c430aaec5c173d23e1aa514c13fb78b83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 10:42:17 +0200 Subject: [PATCH 33/89] fix(PSGO-261): normalize inbound Bearer scheme to avoid Bearer Bearer headers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A programmatic token may arrive already prefixed with `Bearer ` (the codebase tolerates that on input). Downstream, introspect/exchange helpers and KeboolaClient re-add the `Bearer ` scheme, so a pre-prefixed token produced `Authorization: Bearer Bearer …` and broke auth. Strip the scheme (public `strip_bearer`) at the four sites that forward the token: _autolease_default_scope, _resolve_local_tokens, create_session_state (local branch), and _client_for_project. Adds a regression test asserting the exchange subject token is bare. Also fixes the MultiProjectMiddleware docstring, which claimed 'no structured-content merge' while _merge deep-merges structured content (lists concatenated, counters summed, count-first past the cap). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/clients/auth_bridge.py | 6 ++--- src/keboola_mcp_server/mcp.py | 25 ++++++++++++------- tests/test_mcp.py | 19 ++++++++++++++ 3 files changed, 38 insertions(+), 12 deletions(-) diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index 4ed730067..315d1a32b 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -31,7 +31,7 @@ ) -def _strip_bearer(token: str) -> str: +def strip_bearer(token: str) -> str: """Removes a leading case-insensitive ``Bearer `` scheme from a token, if present.""" if token[:7].lower() == 'bearer ': return token[7:].strip() @@ -42,7 +42,7 @@ def is_programmatic_token(token: str | None) -> bool: """True if ``token`` is a Keboola programmatic bearer token (``kbc_at_`` / ``kbc_pat_``).""" if not token: return False - bare = _strip_bearer(token) + bare = strip_bearer(token) return bare.startswith(_ACCESS_TOKEN_PREFIX) or bare.startswith(_PAT_PREFIX) @@ -98,7 +98,7 @@ async def resolve(self, *, subject_token: str, project_id: int) -> str: 'Content-Type': 'application/json', 'Accept': 'application/json', 'X-Kubernetes-Authorization': f'Bearer {self._read_sa_jwt()}', - 'X-Subject-Token': f'Bearer {_strip_bearer(subject_token)}', + 'X-Subject-Token': f'Bearer {strip_bearer(subject_token)}', } try: async with httpx.AsyncClient(timeout=self._timeout, transport=self._transport) as client: diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 61f1ea0c7..a76f7478b 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -33,7 +33,7 @@ from starlette.types import ASGIApp, Receive, Scope, Send from keboola_mcp_server.auth_login import exchange_scoped_token, get_access_token, introspect_token, load_tokens -from keboola_mcp_server.clients.auth_bridge import StorageTokenResolver, is_programmatic_token +from keboola_mcp_server.clients.auth_bridge import StorageTokenResolver, is_programmatic_token, strip_bearer from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config, ServerRuntimeInfo, deployed_sa_token_path, is_same_stack @@ -450,7 +450,7 @@ async def _autolease_default_scope(cls, config: Config) -> 'SessionScope | None' try: parent = await get_access_token(config.storage_api_url) except RuntimeError: - parent = config.storage_token + parent = strip_bearer(config.storage_token) try: introspection = await introspect_token(config.storage_api_url, subject_token=parent) except Exception as e: @@ -479,7 +479,9 @@ async def _resolve_local_tokens( if not cls._is_local_programmatic(config): return config, scope - parent = config.storage_token + # Strip any inbound `Bearer ` scheme; introspect/exchange helpers add the scheme themselves, + # so a pre-prefixed token would produce an `Authorization: Bearer Bearer …` header. + parent = strip_bearer(config.storage_token) try: # Refreshes (and persists the rotated pair) when near expiry; raises if no stored creds. parent = await get_access_token(config.storage_api_url) @@ -587,8 +589,9 @@ async def create_session_state( else: # Local: no projected SA token to reach the resolver. Forward the programmatic token # downstream as a Bearer and let PAT-aware services exchange it; name the target - # project when one has been selected. - bearer_token = storage_token + # project when one has been selected. Strip any inbound `Bearer ` scheme so the + # client's own `Bearer ` prefixing can't produce `Authorization: Bearer Bearer …`. + bearer_token = strip_bearer(storage_token) if config.project_id: extra_headers['X-KBC-ProjectId'] = config.project_id @@ -844,10 +847,11 @@ class MultiProjectMiddleware(fmw.Middleware): Single-project (or no) scope is an unchanged passthrough. With >1 project selected, a read-only tool runs once per project — the active ``KeboolaClient`` in session state is swapped to each - project's client and the per-project results are labelled and concatenated (no structured-content - merge, so each tool keeps its native output shape). Write tools never fan out: they target the - active project only, so the agent can never write to multiple projects without the user explicitly - re-scoping (PSGO-261 decision D8). + project's client and the per-project results are labelled with a per-project text envelope. Their + structured content is deep-merged (lists concatenated across projects, counters summed) into one + schema-valid object, degrading to count-first with a truncated sample past ``_FANOUT_MAX_ITEMS``. + Write tools never fan out: they target the active project only, so the agent can never write to + multiple projects without the user explicitly re-scoping (PSGO-261 decision D8). """ async def on_call_tool( @@ -1019,6 +1023,9 @@ async def _swap_project( async def _client_for_project( server_state: ServerState, token: str, project_id: int, read_only: bool ) -> KeboolaClient: + # Normalize any inbound `Bearer ` scheme; KeboolaClient adds it back for bearer tokens, + # so a pre-prefixed value would otherwise become `Authorization: Bearer Bearer …`. + token = strip_bearer(token) return await KeboolaClient( storage_api_url=server_state.config.storage_api_url, storage_api_token=token, diff --git a/tests/test_mcp.py b/tests/test_mcp.py index f4a020d0d..c1e16b245 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -988,6 +988,25 @@ async def test_near_expiry_scoped_token_is_reminted(self, monkeypatch) -> None: assert out_scope.scoped_token == 'kbc_at_fresh_scoped' assert out_config.storage_token == 'kbc_at_fresh_scoped' + @pytest.mark.asyncio + async def test_bearer_prefixed_token_is_stripped_before_exchange(self, monkeypatch) -> None: + # A programmatic token supplied with an explicit `Bearer ` scheme (tolerated on input) must be + # normalized to bare form; the exchange/introspect helpers add the scheme themselves, so a + # pre-prefixed value would otherwise become `Authorization: Bearer Bearer …` (PSGO-261). + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='Bearer kbc_pat_x') + scope = SessionScope(project_ids=[11], scoped_token='kbc_at_stale', scoped_expires_at=time.time() - 1) + minted = SimpleNamespace(access_token='kbc_at_fresh_scoped', expires_at=time.time() + 900) + exch = AsyncMock(return_value=minted) + with ( + # No stored PKCE session → falls back to the directly-supplied config token. + patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(side_effect=RuntimeError)), + patch('keboola_mcp_server.mcp.exchange_scoped_token', exch), + ): + await SessionStateMiddleware._resolve_local_tokens(config, scope) + exch.assert_awaited_once() + assert exch.await_args.kwargs['subject_token'] == 'kbc_pat_x' # bare, no `Bearer ` prefix + @pytest.mark.asyncio async def test_autolease_scopes_all_accessible_projects(self, monkeypatch) -> None: monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) From 8f5903ea26a2ab7b607940b3e90e3216de4035c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 11:15:13 +0200 Subject: [PATCH 34/89] docs(PSGO-261): use conventional Note: marker in _swap_project docstring Replace the nonstandard 'ponytail:' marker with 'Note:' so the caching-follow-up remark reads clearly to maintainers (Copilot review, #605). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/mcp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index a76f7478b..c11b49a89 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -1012,7 +1012,7 @@ async def _swap_project( Swaps in a per-project `KeboolaClient` AND a `WorkspaceManager` built on it, so workspace-bound reads (query_data) run against *this* project's workspace rather than the active project's. The workspace is provisioned lazily on first use per project. - ponytail: rebuilt per call; caching across calls would need a store that survives the + Note: rebuilt per call; caching across calls would need a store that survives the per-request state rebuild — add if provisioning latency shows up in practice. """ client = await cls._client_for_project(server_state, base_token, project_id, read_only) From 68095e7a101f70c91199e1071fc2f98ae883e110 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 11:26:13 +0200 Subject: [PATCH 35/89] fix(PSGO-261): safer dialect fallback, login timeout, traceback logging - get_accessible_projects: an unresolved SQL dialect (None) now yields the no-dialect prompt instead of defaulting to Snowflake, so a BigQuery/unknown project isn't given Snowflake-specific SQL guidance (+ regression test). - perform_login: bound the loopback callback wait with a 300s server timeout and raise a clear error, so a closed tab / blocked browser can't hang login (and stdio auto-login). - Log swallowed scoped-exchange and per-project fan-out failures with exc_info=True so the tracebacks reach Datadog (matching the codebase convention). Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/auth_login.py | 13 +++++++++++- src/keboola_mcp_server/mcp.py | 2 +- src/keboola_mcp_server/tools/project.py | 8 +++++--- tests/tools/test_project.py | 27 +++++++++++++++++++++++++ 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 19ab9273f..0f6e0c003 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -41,6 +41,9 @@ _SUDO_PATH = 'v1/auth/sudo' _PAT_PATH = 'v1/auth/pat' _REFRESH_SKEW_SECONDS = 60 +# Max time to wait for the browser to hit the loopback /callback. Generous enough for SSO/MFA, but +# bounded so a closed tab or blocked browser can't hang `login` (and stdio auto-login) forever. +_LOGIN_CALLBACK_TIMEOUT_SECONDS = 300 _PAT_DEFAULT_EXPIRES_SECONDS = 30 * 24 * 60 * 60 # ~1 month _CREDENTIALS_PATH = Path.home() / '.keboola' / 'mcp' / 'credentials.json' # Short connect timeout so an unreachable stack (e.g. VPN off — internal `.dev` stacks resolve to a @@ -492,10 +495,18 @@ async def perform_login(storage_api_url: str, *, open_browser=webbrowser.open) - open_browser(authorize_url) _CallbackHandler.result = {} - server.handle_request() # blocks until the browser hits /callback + # Bound the wait: handle_request() returns after `timeout` seconds even if no callback arrives, + # so a closed tab / blocked browser fails with a clear error instead of hanging indefinitely. + server.timeout = _LOGIN_CALLBACK_TIMEOUT_SECONDS + server.handle_request() # blocks until the browser hits /callback or the timeout elapses server.server_close() result = _CallbackHandler.result + if not result: + raise RuntimeError( + f'Timed out after {_LOGIN_CALLBACK_TIMEOUT_SECONDS}s waiting for the browser sign-in callback. ' + 'Re-run the login and complete authentication in the opened browser window.' + ) if result.get('error'): raise RuntimeError(f'Authorization failed: {result.get("error")} {result.get("errorDescription", "")}'.strip()) if not secrets.compare_digest(result.get('state', ''), state): diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index c11b49a89..705584056 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -942,7 +942,7 @@ async def on_call_tool( # "failed for all N projects" aggregate. Abort and surface the single clean error. raise except Exception as e: - LOG.warning(f'Fan-out call failed for project {project_id}: {e}') + LOG.warning(f'Fan-out call failed for project {project_id}: {e}', exc_info=True) errors.append((project_id, str(e))) finally: state[KeboolaClient.STATE_KEY] = original_client diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 69ef794b8..bc647b9f8 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -402,7 +402,9 @@ async def get_accessible_projects( BaseInstructionGroup( project_ids=ids, sql_dialect=dialect, - instructions=get_project_system_prompt(dialect or 'Snowflake'), + # No/unknown dialect -> pass '' so the prompt omits dialect-specific guidance rather + # than defaulting to Snowflake (which would mislead a BigQuery/unknown project). + instructions=get_project_system_prompt(dialect or ''), ) for dialect, ids in by_dialect.items() ] @@ -475,8 +477,8 @@ async def set_project_scope( scoped_expires_at=minted.expires_at, confirmed=True, ) - except Exception as e: - LOG.warning(f'Scoped-token exchange failed ({e}); scoping with the whole-stack token instead.') + except Exception: + LOG.warning('Scoped-token exchange failed; scoping with the whole-stack token instead.', exc_info=True) scope = SessionScope(project_ids=ids, read_only=read_only, confirmed=True) ctx.session.state[SCOPE_KEY] = scope diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 40c872b39..08f88e530 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -340,6 +340,33 @@ async def test_get_accessible_projects_llm_instructions_grouped_by_dialect( assert all(g.instructions for g in result.base_instructions) +@pytest.mark.asyncio +async def test_get_accessible_projects_unknown_dialect_omits_snowflake_guidance( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + # A project whose dialect can't be resolved (None) must NOT fall back to Snowflake guidance — + # that would mislead the assistant into Snowflake-specific SQL for a non-Snowflake project. + from keboola_mcp_server.resources.prompts import get_project_system_prompt + + _prep_client(mcp_context_client, mocker) + introspection = SimpleNamespace(user_email='m@k.com', projects=[SimpleNamespace(id=42, name='X', role='admin')]) + mocker.patch('keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection)) + mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) + mocker.patch( + 'keboola_mcp_server.tools.project._project_sql_dialect', + new=mocker.AsyncMock(side_effect=lambda _ss, _tok, pid: (pid, None)), + ) + + result = await get_accessible_projects(mcp_context_client, with_llm_instruction=True) + + assert result.base_instructions is not None + (group,) = result.base_instructions + assert group.sql_dialect is None + # The unknown-dialect group gets the no-dialect prompt, not the Snowflake one. + assert group.instructions == get_project_system_prompt('') + assert group.instructions != get_project_system_prompt('Snowflake') + + @pytest.mark.asyncio async def test_set_project_scope_subset_exchanges_and_stores( mcp_context_client: Context, mocker: MockerFixture From 03b54fa5c43a5be45808a78415716e742f9d4566 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 11:40:23 +0200 Subject: [PATCH 36/89] refactor(PSGO-261): public client_for_project, legacy-token instructions, store polish - Promote MultiProjectMiddleware._client_for_project to public client_for_project so the project tool no longer reaches into a protected middleware member (also a thermonuclear finding); update its one external caller and the tests. - Reword the server instructions so they don't mislead legacy Storage-API-token sessions: multi-project gating applies to programmatic (kbc_at_/kbc_pat_) tokens; a legacy token is already single-project and uses tools directly. - credentials.json: json.dump(ensure_ascii=False) per the project JSON guideline. - Fix ensure_access_token docstring: the loopback wait is now bounded by _LOGIN_CALLBACK_TIMEOUT_SECONDS. Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/auth_login.py | 5 +++-- src/keboola_mcp_server/mcp.py | 4 ++-- src/keboola_mcp_server/server.py | 24 ++++++++++++++---------- src/keboola_mcp_server/tools/project.py | 2 +- tests/test_mcp.py | 14 +++++++------- 5 files changed, 27 insertions(+), 22 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 0f6e0c003..c34b3eba3 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -355,7 +355,7 @@ def _write_store(store: dict) -> None: _CREDENTIALS_PATH.parent.mkdir(parents=True, exist_ok=True, mode=0o700) fd = os.open(_CREDENTIALS_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) with os.fdopen(fd, 'w') as f: - json.dump(store, f, indent=2) + json.dump(store, f, indent=2, ensure_ascii=False) # O_CREAT honors the mode only when creating; chmod covers a pre-existing file. _CREDENTIALS_PATH.chmod(0o600) @@ -417,7 +417,8 @@ async def ensure_access_token( ``allow_interactive`` MUST be false unless a real terminal is attached. When the stdio server is launched by an MCP client its stdout is the JSON-RPC channel and there is no TTY, so an interactive login would both corrupt the protocol stream and block the initialize handshake - (the loopback wait has no timeout). In that case this raises the same "run login" guidance as + (and the loopback wait, though bounded by ``_LOGIN_CALLBACK_TIMEOUT_SECONDS``, would still stall + the handshake for its duration). In that case this raises the same "run login" guidance as ``get_access_token`` instead of attempting a browser login. Remote/deployed servers must use client-driven OAuth regardless. """ diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 705584056..f84a63748 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -1015,12 +1015,12 @@ async def _swap_project( Note: rebuilt per call; caching across calls would need a store that survives the per-request state rebuild — add if provisioning latency shows up in practice. """ - client = await cls._client_for_project(server_state, base_token, project_id, read_only) + client = await cls.client_for_project(server_state, base_token, project_id, read_only) state[KeboolaClient.STATE_KEY] = client state[WorkspaceManager.STATE_KEY] = await WorkspaceManager.create(client, server_state.config.workspace_schema) @staticmethod - async def _client_for_project( + async def client_for_project( server_state: ServerState, token: str, project_id: int, read_only: bool ) -> KeboolaClient: # Normalize any inbound `Bearer ` scheme; KeboolaClient adds it back for bearer tokens, diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index 322d31c44..eb9d01236 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -230,16 +230,20 @@ def create_server( mcp = KeboolaMcpServer( name='Keboola MCP Server', instructions=( - 'This server runs in multi-project mode. When the user logs in with a stack-wide Keboola ' - 'token, data tools are BLOCKED until a project scope is confirmed. So at the very START of ' - 'the conversation, before doing anything else: call "get_accessible_projects", show the user ' - 'their projects, and ASK whether to work across ALL of them or a subset. Do not decide for ' - 'them. Then call "set_project_scope" with their answer (no arguments = all projects, or the ' - 'chosen project ids, optionally read_only=true). After that, read-only tools return results ' - 'per project. Never write to more than one project without explicit user confirmation — ' - 'write operations target the active (first-scoped) project only. Note: outside the Storage ' - 'API, some tools may need per-project token support not yet available on every stack; ' - 'surface such errors plainly rather than retrying.' + 'This server supports multi-project mode for stack-wide Keboola programmatic tokens ' + '(kbc_at_/kbc_pat_). When the session uses such a token, data tools are BLOCKED until a ' + 'project scope is confirmed. So at the very START of the conversation, before doing anything ' + 'else: call "get_accessible_projects", show the user their projects, and ASK whether to work ' + 'across ALL of them or a subset. Do not decide for them. Then call "set_project_scope" with ' + 'their answer (no arguments = all projects, or the chosen project ids, optionally ' + 'read_only=true). After that, read-only tools return results per project. Never write to ' + 'more than one project without explicit user confirmation — write operations target the ' + 'active (first-scoped) project only. If instead the session uses a legacy project-scoped ' + 'Storage API token, it is already bound to a single project: use the tools directly — ' + '"get_accessible_projects" / "set_project_scope" do not apply (they will report that no ' + 'programmatic token is present). Note: outside the Storage API, some tools may need ' + 'per-project token support not yet available on every stack; surface such errors plainly ' + 'rather than retrying.' ), lifespan=create_keboola_lifespan(server_state), auth=oauth_provider, diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index bc647b9f8..196a8ccf4 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -328,7 +328,7 @@ async def _project_sql_dialect( No workspace is provisioned — the dialect comes from the token's owner.defaultBackend, so this is a single cheap Storage API call per project. """ - per_client = await MultiProjectMiddleware._client_for_project( + per_client = await MultiProjectMiddleware.client_for_project( server_state, subject_token, project_id, read_only=True ) token_data = await per_client.storage_client.verify_token() diff --git a/tests/test_mcp.py b/tests/test_mcp.py index c1e16b245..e8c777fc7 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1127,7 +1127,7 @@ async def call_next(_): with ( patch.object( MultiProjectMiddleware, - '_client_for_project', + 'client_for_project', AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), ), patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), @@ -1161,7 +1161,7 @@ async def call_next(_): with ( patch.object( MultiProjectMiddleware, - '_client_for_project', + 'client_for_project', AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), ), patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), @@ -1188,7 +1188,7 @@ async def call_next(_): with ( patch.object( MultiProjectMiddleware, - '_client_for_project', + 'client_for_project', AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), ), patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), @@ -1214,7 +1214,7 @@ async def call_next(_): with ( patch.object( MultiProjectMiddleware, - '_client_for_project', + 'client_for_project', AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), ), patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), @@ -1319,7 +1319,7 @@ async def call_next(_): with ( patch.object( MultiProjectMiddleware, - '_client_for_project', + 'client_for_project', AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), ), patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), @@ -1342,7 +1342,7 @@ async def call_next(_): with ( patch.object( MultiProjectMiddleware, - '_client_for_project', + 'client_for_project', AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), ), patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), @@ -1365,7 +1365,7 @@ async def call_next(_): with ( patch.object( MultiProjectMiddleware, - '_client_for_project', + 'client_for_project', AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), ), patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), From b05b4d3a88f85587f6fb8ebd647c02d03c9c14e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 11:51:53 +0200 Subject: [PATCH 37/89] fix(PSGO-261): propagate cancellation in dialect fan-out; correct /list comment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - get_accessible_projects: re-raise asyncio.CancelledError from the concurrent dialect lookups instead of treating it as a best-effort per-project failure, so a cancelled request stops promptly rather than continuing work. - Reword the /list on_request comment: it skips the extra auth round-trips (introspect/ refresh/scoped-exchange), but create_session_state below may still make ordinary Storage calls (WorkspaceManager.create) — the old 'zero network' claim was misleading. Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/mcp.py | 14 ++++++++------ src/keboola_mcp_server/tools/project.py | 3 +++ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index f84a63748..8348bea2a 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -277,12 +277,14 @@ async def on_request( if http_rq := get_http_request_or_none(): config = self.apply_request_config(http_rq, config, own_stack_storage_api_url=own_stack_storage_api_url) - # Capability-discovery requests (tools/list, prompts/list, resources/list) MUST be fast and - # network-free: a client fetches all three on connect, so any Connection round-trip here - # (token introspect, refresh, or scoped-exchange) makes connecting hang until the client's - # 30s timeout. For /list we do zero network in on_request — no auto-lease, no token refresh, - # no scoped re-mint — and use the stored session token as-is (no refresh). The scope and - # fresh tokens are established on the first real (non-list) tool call. + # Capability-discovery requests (tools/list, prompts/list, resources/list) MUST be fast: a + # client fetches all three on connect, so any Connection AUTH round-trip here (token + # introspect, refresh, or scoped-exchange) makes connecting hang until the client's 30s + # timeout. For /list we skip all that extra auth work — no auto-lease, no token refresh, no + # scoped re-mint — and use the stored session token as-is. (create_session_state below may + # still make ordinary Storage calls, e.g. WorkspaceManager.create; the point is /list adds + # none of the introspect/refresh/exchange round-trips.) Scope and fresh tokens are + # established on the first real (non-list) tool call. is_list = context.method.endswith('/list') # Local streamable-HTTP with no token supplied (no header / env): fall back to the stored diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 196a8ccf4..a4d95d66a 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -1,3 +1,4 @@ +import asyncio import logging from typing import Annotated, Optional, cast @@ -374,6 +375,8 @@ async def get_accessible_projects( lambda pid: _project_sql_dialect(server_state, subject_token, pid), ) for result in results: + if isinstance(result, asyncio.CancelledError): + raise result # never swallow cancellation — let it propagate if isinstance(result, BaseException): LOG.warning(f'Could not resolve SQL dialect for a project: {result}') continue From 5b923b889831faa28f92952abd201b1d95811c78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 13:18:10 +0200 Subject: [PATCH 38/89] fix(PSGO-261): keep the workspace when creation ends with a 'warning' status A workspace-creation job that finishes with status 'warning' (a child job failed) can still have produced the workspace (results.id present). The polling loop treated 'warning' as an outright failure and returned None, discarding a live workspace and forcing callers into a retry/timeout. Now: 'warning' with a results.id returns that workspace (like 'success'); 'warning' without one stays a terminal failure. Adds regression tests for both. Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/workspace.py | 17 +++++++++++-- tests/test_workspace.py | 39 ++++++++++++++++++++++++++--- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/src/keboola_mcp_server/workspace.py b/src/keboola_mcp_server/workspace.py index 1524133b2..bc890273d 100644 --- a/src/keboola_mcp_server/workspace.py +++ b/src/keboola_mcp_server/workspace.py @@ -786,9 +786,22 @@ async def _create_ws(self, *, timeout_sec: float = 300.0) -> _WspInfo | None: LOG.info(f'Created workspace: {workspace_id}') return await self._find_ws_by_id(workspace_id) + elif ( + job_status == 'warning' + and isinstance(job_info.get('results'), dict) + and isinstance(job_info['results'].get('id'), int) + ): + # 'warning' = the job finished but a child job failed; the workspace itself may still + # have been created (results.id present). Use it instead of discarding a live workspace. + workspace_id = job_info['results']['id'] + LOG.warning( + f'Workspace creation finished with warning; using workspace {workspace_id}: job_id={job_id}' + ) + return await self._find_ws_by_id(workspace_id) + elif job_status in ('error', 'warning', 'terminated', 'cancelled', 'canceled'): - # Terminal failure states: the job will never reach 'success', so stop polling - # immediately instead of spinning until the timeout. + # Terminal failure states (incl. 'warning' with no workspace id): the job will never + # reach 'success', so stop polling immediately instead of spinning until the timeout. LOG.warning(f'Workspace creation job failed: job_id={job_id}, status={job_status}') return None diff --git a/tests/test_workspace.py b/tests/test_workspace.py index cfcbb7251..91ec1a5a5 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -172,7 +172,18 @@ async def test_workspace_creation_cleans_up_config_on_failure(): @pytest.mark.asyncio -async def test_workspace_creation_stops_on_terminal_error_status(): +@pytest.mark.parametrize( + 'job_detail', + [ + {'status': 'error'}, + {'status': 'terminated'}, + # 'warning' with no workspace id is still a failure (nothing usable was produced). + {'status': 'warning'}, + {'status': 'warning', 'results': {}}, + ], + ids=['error', 'terminated', 'warning_no_results', 'warning_no_id'], +) +async def test_workspace_creation_stops_on_terminal_error_status(job_detail: dict): """A job that reaches a terminal failure status must stop polling at once, not spin to timeout.""" mock_client = Mock(spec=KeboolaClient) mock_client.branch_id = None @@ -182,16 +193,38 @@ async def test_workspace_creation_stops_on_terminal_error_status(): mock_storage_client.verify_token.return_value = {'owner': {'defaultBackend': 'snowflake'}} mock_storage_client.configuration_create.return_value = {'id': 'cfg-1', 'name': 'test'} mock_storage_client.workspace_create_for_config.return_value = {'id': 999} - mock_storage_client.job_detail.return_value = {'status': 'error'} + mock_storage_client.job_detail.return_value = job_detail manager = WorkspaceManager(mock_client) result = await manager._create_ws() assert result is None - # Polled exactly once — the terminal 'error' status short-circuits the loop. + # Polled exactly once — the terminal status short-circuits the loop. mock_storage_client.job_detail.assert_awaited_once() +@pytest.mark.asyncio +async def test_workspace_creation_warning_with_id_uses_workspace(mocker): + """A 'warning' job that still created a workspace (results.id present) must not be discarded.""" + mock_client = Mock(spec=KeboolaClient) + mock_client.branch_id = None + mock_storage_client = AsyncMock() + mock_client.storage_client = mock_storage_client + + mock_storage_client.verify_token.return_value = {'owner': {'defaultBackend': 'snowflake'}} + mock_storage_client.configuration_create.return_value = {'id': 'cfg-1', 'name': 'test'} + mock_storage_client.workspace_create_for_config.return_value = {'id': 999} + mock_storage_client.job_detail.return_value = {'status': 'warning', 'results': {'id': 999}} + + manager = WorkspaceManager(mock_client) + sentinel = object() + mocker.patch.object(manager, '_find_ws_by_id', AsyncMock(return_value=sentinel)) + result = await manager._create_ws() + + assert result is sentinel # the created workspace is used despite the warning + manager._find_ws_by_id.assert_awaited_once_with(999) + + @pytest.mark.asyncio @pytest.mark.parametrize( ('input_branch_id', 'has_sb_feature', 'workspace_schema', 'expected_bound_branch_id'), From c1916fd2f5ccb0a1b2bcb80c6278373daf5037f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 13:30:12 +0200 Subject: [PATCH 39/89] fix(PSGO-261): reject empty project_ids; chmod credentials before writing tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - set_project_scope and the per-call project_ids fan-out filter now reject an explicit empty list (ValueError / ToolError) instead of silently treating [] like null and broadening to all projects — an accidental [] must not widen scope. Adds a regression test. - _write_store now fchmod(0600) the fd BEFORE writing token material (after O_TRUNC empties it), closing the window where a pre-existing over-permissioned credentials file could hold freshly-written tokens before the trailing chmod. Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/auth_login.py | 6 ++++-- src/keboola_mcp_server/mcp.py | 9 ++++++++- src/keboola_mcp_server/tools/project.py | 5 +++++ tests/tools/test_project.py | 10 ++++++++++ 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index c34b3eba3..32685327e 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -355,9 +355,11 @@ def _write_store(store: dict) -> None: _CREDENTIALS_PATH.parent.mkdir(parents=True, exist_ok=True, mode=0o700) fd = os.open(_CREDENTIALS_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) with os.fdopen(fd, 'w') as f: + # O_CREAT only applies the mode when creating; a pre-existing file could be world-readable. + # fchmod BEFORE writing any token material so there is no exposure window (O_TRUNC already + # emptied the file, so nothing sensitive exists until json.dump runs after this). + os.fchmod(f.fileno(), 0o600) json.dump(store, f, indent=2, ensure_ascii=False) - # O_CREAT honors the mode only when creating; chmod covers a pre-existing file. - _CREDENTIALS_PATH.chmod(0o600) def load_tokens(storage_api_url: str) -> TokenSet | None: diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 8348bea2a..55eee7241 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -897,7 +897,14 @@ async def on_call_tool( requested = args.pop(_PROJECT_FILTER_ARG, None) targets = list(scope.project_ids) - if requested: + if requested is not None: + # Omit the filter to run across the full scope; an explicit empty list is a caller mistake + # (it must not silently fall through to the whole scope). + if not requested: + raise ToolError( + f'"{_PROJECT_FILTER_ARG}" must be a non-empty list of project ids, ' + 'or omitted to run across the full scope.' + ) outside = [p for p in requested if p not in scope.project_ids] if outside: raise ToolError( diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index a4d95d66a..cd1d5c550 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -459,6 +459,11 @@ async def set_project_scope( client = KeboolaClient.from_state(ctx.session.state) parent_token = await _parent_subject_token(client) + # Distinguish "omit/null" (scope to all) from an explicit empty list, which is almost certainly a + # caller mistake and must not silently broaden the scope to every project. + if project_ids is not None and len(project_ids) == 0: + raise ValueError('project_ids must be a non-empty list of project ids, or omitted/null to scope to all.') + ids = list(project_ids or []) if not ids: introspection = await introspect_token(client.storage_api_url, subject_token=parent_token) diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 08f88e530..d75857232 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -412,3 +412,13 @@ async def test_scope_requires_programmatic_token(mcp_context_client: Context, mo _prep_client(mcp_context_client, mocker, bearer=None) with pytest.raises(ValueError, match='programmatic token'): await get_accessible_projects(mcp_context_client) + + +@pytest.mark.asyncio +async def test_set_project_scope_rejects_explicit_empty_list( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + # An explicit [] must NOT be treated like null (all projects) — it's almost certainly a mistake. + _prep_client(mcp_context_client, mocker) + with pytest.raises(ValueError, match='non-empty'): + await set_project_scope(mcp_context_client, project_ids=[]) From be5a848456a272f81a8f89a78b40fe8e9ba068af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 13:38:22 +0200 Subject: [PATCH 40/89] fix(PSGO-261): refresh an expired stored token on /list (keep valid tokens network-free) /list read the stored PKCE token as-is; if it was already expired, create_session_state's Storage calls (WorkspaceManager.create) would fail, forcing a needless re-login after ~1h. Now /list refreshes only a locally-known-expired token (checked via TokenSet.is_near_expiry, no network for valid tokens), preserving the 'never block /list on connect' property. Adds tests for the valid (no refresh) and expired (one refresh) cases. Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/keboola_mcp_server/mcp.py | 11 ++++++++++- tests/test_mcp.py | 25 ++++++++++++++++++++++--- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 55eee7241..f31f3a053 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -433,7 +433,16 @@ async def _maybe_use_stored_session(cls, config: Config, *, refresh: bool = True tokens = load_tokens(config.storage_api_url) if not tokens: return config - access_token = tokens.access_token + if tokens.is_near_expiry: + # The stored access token is (near) expired; using it as-is would make the /list + # session-state build fail its Storage calls. Refresh only this case via the network — + # valid tokens still take the network-free fast path so /list never blocks on connect. + try: + access_token = await get_access_token(config.storage_api_url) + except RuntimeError: + return config + else: + access_token = tokens.access_token return dataclasses.replace(config, storage_token=access_token) @classmethod diff --git a/tests/test_mcp.py b/tests/test_mcp.py index e8c777fc7..bc4302cfb 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -888,17 +888,36 @@ async def test_existing_token_is_noop(self, monkeypatch) -> None: assert out is config @pytest.mark.asyncio - async def test_list_request_uses_stored_token_without_network_refresh(self, monkeypatch) -> None: - # /list must not do a network refresh: read the stored token as-is via load_tokens. + async def test_list_request_uses_valid_stored_token_without_network_refresh(self, monkeypatch) -> None: + # /list with a still-valid stored token must not do a network refresh: read it as-is. monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) config = Config(storage_api_url='https://connection.keboola.com') with ( patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(side_effect=AssertionError('no network'))), - patch('keboola_mcp_server.mcp.load_tokens', return_value=SimpleNamespace(access_token='kbc_at_file')), + patch( + 'keboola_mcp_server.mcp.load_tokens', + return_value=SimpleNamespace(access_token='kbc_at_file', is_near_expiry=False), + ), ): out = await SessionStateMiddleware._maybe_use_stored_session(config, refresh=False) assert out.storage_token == 'kbc_at_file' + @pytest.mark.asyncio + async def test_list_request_refreshes_only_an_expired_stored_token(self, monkeypatch) -> None: + # /list with an EXPIRED stored token refreshes (once) so session-state Storage calls don't fail. + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com') + with ( + patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_fresh')) as gat, + patch( + 'keboola_mcp_server.mcp.load_tokens', + return_value=SimpleNamespace(access_token='kbc_at_stale', is_near_expiry=True), + ), + ): + out = await SessionStateMiddleware._maybe_use_stored_session(config, refresh=False) + gat.assert_awaited_once() + assert out.storage_token == 'kbc_at_fresh' + @pytest.mark.asyncio async def test_deployed_is_noop(self, monkeypatch) -> None: monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') From 8fdb8dd0f1dd907ff14f87724a0e6d8fc6eda72a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 21 Jul 2026 14:45:32 +0200 Subject: [PATCH 41/89] fix(PSGO-261): log swallowed auth-refresh errors with traceback; don't mask exchange client errors - _autolease_default_scope and _resolve_local_tokens now log their swallowed introspect/ scoped-exchange failures with exc_info=True, so the underlying httpx/status error is visible in Datadog instead of just the message. - set_project_scope now re-raises a 400/401/403 from /v1/auth/pat/exchange (bad project_ids, invalid/insufficient token) instead of silently downgrading to an unscoped whole-stack token. Only endpoint-unavailable errors (5xx, network/timeout) still fall back. Adds regression tests for both branches. - README: the browser login leases a stack-wide session; it doesn't pick a project. Project selection happens afterwards via get_accessible_projects/set_project_scope. Code review follow-up, #605. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 2 +- src/keboola_mcp_server/mcp.py | 4 +- src/keboola_mcp_server/tools/project.py | 11 +++++ tests/tools/test_project.py | 56 +++++++++++++++++++++++++ 4 files changed, 70 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index bdb4dcf9f..ab9341a2a 100644 --- a/README.md +++ b/README.md @@ -145,7 +145,7 @@ Sign in once with your browser; the server stores the session and refreshes it a uvx keboola_mcp_server login --api-url https://connection.YOUR_REGION.keboola.com ``` -This opens your browser to sign in to Keboola and pick a project, then saves the session to `~/.keboola/mcp/credentials.json` (readable only by you, one entry per stack). Afterwards, start the server with only `KBC_STORAGE_API_URL` set — no token required. +This opens your browser to sign in to Keboola, then saves the stack-wide session to `~/.keboola/mcp/credentials.json` (readable only by you, one entry per stack). Afterwards, start the server with only `KBC_STORAGE_API_URL` set — no token required. Which project(s) to work on is chosen afterwards, in the conversation (`get_accessible_projects` / `set_project_scope`), not during login. | Command | What it does | |---------|--------------| diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index f31f3a053..3dd95432d 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -465,7 +465,7 @@ async def _autolease_default_scope(cls, config: Config) -> 'SessionScope | None' try: introspection = await introspect_token(config.storage_api_url, subject_token=parent) except Exception as e: - LOG.warning(f'Could not auto-lease projects from token introspection: {e}') + LOG.warning(f'Could not auto-lease projects from token introspection: {e}', exc_info=True) return None project_ids = [p.id for p in introspection.projects] if not project_ids: @@ -517,7 +517,7 @@ async def _resolve_local_tokens( ) except Exception as e: # Don't break the session if re-minting fails; fall back to the parent token. - LOG.warning(f'Could not refresh the scoped token; using the parent token: {e}') + LOG.warning(f'Could not refresh the scoped token; using the parent token: {e}', exc_info=True) scope = dataclasses.replace(scope, scoped_token=None, scoped_expires_at=None) token = scope.scoped_token or parent diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index cd1d5c550..0c71a0c04 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -2,6 +2,7 @@ import logging from typing import Annotated, Optional, cast +import httpx from fastmcp import Context, FastMCP from fastmcp.tools import FunctionTool from mcp.types import ToolAnnotations @@ -485,7 +486,17 @@ async def set_project_scope( scoped_expires_at=minted.expires_at, confirmed=True, ) + except httpx.HTTPStatusError as e: + if e.response.status_code in (400, 401, 403): + # Client error (bad project_ids, invalid/insufficient token): the input or auth is wrong, + # not the exchange endpoint — surface it instead of silently downgrading to an unscoped + # whole-stack token, which would mislead the caller about what was actually scoped. + raise + LOG.warning('Scoped-token exchange failed; scoping with the whole-stack token instead.', exc_info=True) + scope = SessionScope(project_ids=ids, read_only=read_only, confirmed=True) except Exception: + # Network/timeout/unavailable exchange endpoint: fall back so scoping still works, narrowed + # per request by X-KBC-ProjectId, without the extra token-scoping security narrowing. LOG.warning('Scoped-token exchange failed; scoping with the whole-stack token instead.', exc_info=True) scope = SessionScope(project_ids=ids, read_only=read_only, confirmed=True) ctx.session.state[SCOPE_KEY] = scope diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index d75857232..f2f4dce0a 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -1,6 +1,7 @@ import time from types import SimpleNamespace +import httpx import pytest from mcp.server.fastmcp import Context from pytest_mock import MockerFixture @@ -407,6 +408,61 @@ async def test_set_project_scope_all_introspects_then_exchanges( assert result.project_ids == [18, 83] +@pytest.mark.asyncio +@pytest.mark.parametrize('status_code', [400, 401, 403]) +async def test_set_project_scope_reraises_client_error( + mcp_context_client: Context, mocker: MockerFixture, status_code: int +) -> None: + # A 400/401/403 from the exchange means bad input/auth, not an unavailable endpoint — it must + # surface to the caller rather than silently downgrading to an unscoped whole-stack token. + _prep_client(mcp_context_client, mocker) + response = httpx.Response(status_code, request=httpx.Request('POST', 'https://x/v1/auth/pat/exchange')) + mocker.patch( + 'keboola_mcp_server.tools.project.exchange_scoped_token', + new=mocker.AsyncMock(side_effect=httpx.HTTPStatusError('bad', request=response.request, response=response)), + ) + + with pytest.raises(httpx.HTTPStatusError): + await set_project_scope(mcp_context_client, project_ids=[18]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('status_code', [500, 502, 503]) +async def test_set_project_scope_falls_back_on_server_error( + mcp_context_client: Context, mocker: MockerFixture, status_code: int +) -> None: + # A 5xx (endpoint unavailable) still falls back to the whole-stack token so scoping keeps working. + _prep_client(mcp_context_client, mocker) + response = httpx.Response(status_code, request=httpx.Request('POST', 'https://x/v1/auth/pat/exchange')) + mocker.patch( + 'keboola_mcp_server.tools.project.exchange_scoped_token', + new=mocker.AsyncMock(side_effect=httpx.HTTPStatusError('down', request=response.request, response=response)), + ) + + result = await set_project_scope(mcp_context_client, project_ids=[18]) + + assert result.project_ids == [18] + scope = mcp_context_client.session.state[SCOPE_KEY] + assert scope.scoped_token is None + + +@pytest.mark.asyncio +async def test_set_project_scope_falls_back_on_network_error( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + _prep_client(mcp_context_client, mocker) + mocker.patch( + 'keboola_mcp_server.tools.project.exchange_scoped_token', + new=mocker.AsyncMock(side_effect=httpx.ConnectTimeout('timed out')), + ) + + result = await set_project_scope(mcp_context_client, project_ids=[18]) + + assert result.project_ids == [18] + scope = mcp_context_client.session.state[SCOPE_KEY] + assert scope.scoped_token is None + + @pytest.mark.asyncio async def test_scope_requires_programmatic_token(mcp_context_client: Context, mocker: MockerFixture) -> None: _prep_client(mcp_context_client, mocker, bearer=None) From 200cd5c2eeb6a6767e8938221320cdcc1de46ea7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 22 Jul 2026 14:04:51 +0200 Subject: [PATCH 42/89] docs(PSGO-261): RFC for OAuth login exchanging into a programmatic session Connection (keboola/connection#7836) is removing /oauth/authorize outright, in favor of /oauth/consent + a new internal auth-bridge endpoint (manage/internal/auth-bridge/exchange-oauth-token) that exchanges the league OAuth claudai/projectless-scoped token for a whole-stack kbc_at_* programmatic session, discarding the league token immediately after. This closes the exact 'OAuth->PAT exchange is a separate PR' gap the original PSGO-261 RFC deferred: once wired through apply_request_config, OAuth sessions get the full multi-project scoping machinery (get_accessible_projects/set_project_scope) for free. Verified against Connection's own E2E test and TokenRefreshProcessor.php: projectless = whole-stack (not project-pinned), and the exchanged session refreshes independently forever after via the existing /v1/auth/token/refresh - no dependency on refreshing the original league OAuth session. Co-Authored-By: Claude Opus 4.8 (1M context) --- feature_spec/oauth_session_exchange/RFC.md | 103 +++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 feature_spec/oauth_session_exchange/RFC.md diff --git a/feature_spec/oauth_session_exchange/RFC.md b/feature_spec/oauth_session_exchange/RFC.md new file mode 100644 index 000000000..c1db19df9 --- /dev/null +++ b/feature_spec/oauth_session_exchange/RFC.md @@ -0,0 +1,103 @@ +# RFC: OAuth login exchanges for a programmatic session, replacing the project-bound SAPI mint + +Linear: [PSGO-261](https://linear.app/keboola/issue/PSGO-261/support-pat-tokens-in-mcp-server-mcp-server) +Parent: PSGO-261 (multi-project PAT support) — this closes the "OAuth→PAT exchange is a separate PR" carve-out that RFC explicitly deferred. +Related: [keboola/connection#7836](https://github.com/keboola/connection/pull/7836) — the new Connection-side internal endpoint this RFC integrates with. + +--- + +## Problem + +The MCP server's public/remote OAuth login (`SimpleOAuthProvider`, `oauth.py`) currently: + +1. Redirects to `{server_url}/oauth/authorize` (`oauth.py:160,226-237`) — no `scope` sent on purpose (`# send no scopes ... let it use its own default scope`). +2. Exchanges the resulting code at `{server_url}/oauth/token` for a league OAuth access/refresh token pair (`oauth.py:264-297`). +3. Mints a **project-bound legacy Storage API token** from that OAuth access token via `POST {storage_api_url}/v2/storage/tokens` (`_create_sapi_token`, `oauth.py:626-658`), because "AI Service and Jobs Queue... do not support bearer tokens yet" (`ProxyAccessToken.sapi_token` docstring, `oauth.py:114-118`). +4. Stores that legacy token as `config.storage_token` (`mcp.py` `apply_request_config`) — a session pinned to whichever single project was implicit at authorize time, entirely outside the PSGO-261 multi-project architecture (`get_accessible_projects`/`set_project_scope`/fan-out never apply to OAuth sessions today). + +**This is changing, unconditionally and immediately** (per keboola/connection#7836): +- `/oauth/authorize` is being **removed outright** — no back-compat, no old-client fallback, no deprecation window, no rollout coordination needed on our side (verify locally, ship when ready). +- The front-channel authorize step moves to a new endpoint, `/oauth/consent`, requesting scope `claudai projectless` (see Decisions §1). +- After the standard code→token exchange (still against Connection, still yielding a league OAuth access token — now `claudai`+`projectless`-scoped), the MCP server must call a **new internal auth-bridge endpoint** to turn that OAuth token into a real Keboola session: + + ``` + POST {manage-host}/internal/auth-bridge/exchange-oauth-token + Headers: + X-Kubernetes-Authorization: Bearer # same mechanism as resolve-storage-token + X-KBC-ManageApiToken: + X-Subject-Token: Bearer + Auth: caller's own Manage token must be TYPE_SUPER or carry scope + SCOPE_INTERNAL_AUTH_BRIDGE_EXCHANGE_OAUTH_TOKEN; caller must be + Kubernetes-authenticated (validated KubernetesClaims) + + Response 200 (CliTokenResponse): + { accessToken, refreshToken, tokenType: "Bearer", expiresIn, sessionId, user: {email, ...} } + 401: subject token missing/invalid, or not bound to an active admin + (AuthBridgeAuthenticationException | OAuthExchangeUnauthorizedException) + 403: caller not Kubernetes-authenticated / Manage-access-denied, or subject + token missing the claudai scope (ManageAccessDeniedException | MissingClaudaiScopeException) + ``` + +- **Confirmed by keboola/connection's own E2E test** (`AuthBridgeOAuthExchangeTest.php`): + - A normal `claudai`-scoped subject token exchanges to a **project-pinned** session (`testExchangeIssuesPinnedManagelessSession`). + - A subject token additionally carrying the **`projectless`** scope (league `user_identifier` = `admin:{id}`, not project-bound) exchanges to an **unrestricted, whole-stack** session — the test explicitly asserts it can reach a project via live membership alone, with no pin (`testExchangeProjectLess...`). + - The exchanged session has **no Manage API access** (`testExchangedSessionHasNoManageApiAccess` — 401 on `/manage/projects`): it's a pure Storage-scoped programmatic session, not a Manage token. + - **Exchange-only enforcement:** a `projectless` league token cannot be used directly as a Storage bearer (401) nor via `resolve-storage-token` (401) — this new endpoint is the *only* redemption path for it. +- The `CliTokenResponse` shape is **identical** to a PKCE `login` session (`auth_login.py`'s `TokenSet`: `accessToken`/`refreshToken`/`expiresIn`/`sessionId`). This is the same `kbc_at_*`-style programmatic token the rest of PSGO-261 already knows how to handle. +- **The original league OAuth access/refresh token pair is used exactly once** (for this exchange call) **and then permanently discarded** — never stored, never sent to any other Keboola service, and — per `TokenRefreshProcessor.php` — never touched again even on refresh (see Decisions §4). + +**Symptom if unaddressed:** the moment Connection removes `/oauth/authorize`, every public/remote MCP OAuth login (Claude.ai, Cursor, any HTTP client using this server's OAuth flow) breaks outright, with no fallback. + +## Required Behavior + +### Token contract + +| Token | Source | Sent as | Lifetime | Fate | +|---|---|---|---|---| +| League OAuth access token (`claudai`+`projectless` scope) | Connection `/oauth/token` (front-channel via `/oauth/consent`) | `X-Subject-Token: Bearer ` — **only** to the new internal exchange | existing league OAuth TTL | Used once, then discarded permanently | +| Exchanged session (`accessToken`/`refreshToken`) | `POST manage/internal/auth-bridge/exchange-oauth-token` | `kbc_at_*` — same shape as a PKCE `login` `TokenSet`, whole-stack (projectless) | per `CliTokenResponse.expiresIn` | **Becomes the session's only credential; refreshed independently forever after (§4)** | + +### Flow + +1. `authorize()` redirects to `{server_url}/oauth/consent` (was `/oauth/authorize`), requesting scope `claudai projectless`. +2. Code→token exchange is **unchanged**: still `POST {server_url}/oauth/token`, still returns a league OAuth access/refresh pair (now carrying both scopes). +3. **New step, replacing `_create_sapi_token`:** exchange the league OAuth access token for a Keboola programmatic session via `manage/internal/auth-bridge/exchange-oauth-token`, reusing the exact SA-JWT + `X-Subject-Token` mechanism already implemented for `resolve-storage-token` (`clients/auth_bridge.py`). +4. Parse the `CliTokenResponse` into the same shape `auth_login.py._parse_token_response` already builds from a PKCE response. +5. From here on this session is **indistinguishable, downstream, from a directly-supplied `kbc_at_*` token**: `is_programmatic_token()` detects it, `create_session_state`'s existing deployed-path branch (`_exchange_programmatic_token` → `StorageTokenResolver.resolve`) runs unchanged, and the full PSGO-261 multi-project machinery (`get_accessible_projects`, `set_project_scope`, read fan-out) becomes available to every OAuth client for the first time, starting whole-stack/unconfirmed exactly like a fresh PKCE login. +6. The league OAuth token pair from step 2 is discarded after step 3 completes — never persisted, never refreshed. +7. `ProxyAccessToken.sapi_token` (currently required, justified only by "Jobs Queue/AI Service don't support bearer tokens yet") is **obsolete**: those clients now speak bearer via `bearer_or_sapi_token` (PSGO-261, commit `5b8c65ed`). Remove the field; repurpose `ProxyAccessToken` to carry the new `kbc_at_` token and its own `refresh_token`/`session_id` instead. +8. **Refresh is fully decoupled from the league OAuth session** (confirmed, §4): `exchange_refresh_token()` calls `POST /v1/auth/token/refresh` (already implemented, `auth_login.py.refresh_tokens`) directly against the previously-exchanged refresh token. The league OAuth `/oauth/token` refresh grant is never invoked again after step 2. + +## Resolution Strategy + +- **`oauth.py`:** + - `_oauth_server_auth_url` → `/oauth/consent`. + - Add `'scope': 'claudai projectless'` to `authorize()`'s `url_params` (`oauth.py:226-232`). + - Replace `_create_sapi_token()` with a new method (e.g. `_exchange_oauth_for_session`) that POSTs to the new internal endpoint. Build it as a sibling to `StorageTokenResolver` in `clients/auth_bridge.py` — **reuse**, don't reimplement, `read_service_account_jwt`/`normalize_storage_api_url` (`clients/base.py`) and the existing error-mapping convention, adapted to this endpoint's exception set (401 for `AuthBridgeAuthenticationException`/`OAuthExchangeUnauthorizedException`, 403 for `ManageAccessDeniedException`/`MissingClaudaiScopeException`). Reuse `auth_login.py._parse_token_response` to build the resulting `TokenSet`. + - `ProxyAccessToken`: drop `sapi_token: str`; the `delegate` (league OAuth) token is kept only long enough to complete the exchange, then not referenced again — no ongoing refresh dependency on it (§4). + - `exchange_authorization_code()`: call the new exchange method instead of `_create_sapi_token`; store the exchanged `refreshToken`/`sessionId` needed for the *independent* refresh path. + - `exchange_refresh_token()`: **simplify** — drop the `POST {server_url}/oauth/token` (`grant_type=refresh_token`) call to Connection's league OAuth server entirely; call `refresh_tokens()` (`auth_login.py`) directly against the previously-exchanged `kbc_at_` refresh token. +- **`mcp.py`, `apply_request_config`:** set `config.storage_token` to the new `kbc_at_` token; the separate `bearer_token=user.access_token.delegate.token` assignment goes away (the league OAuth delegate token is discarded per step 6, never used downstream). `is_programmatic_token()` then does the rest unchanged. +- **No changes** to the PSGO-261 scoping tools, the local PKCE `login` CLI flow, or the local-stdio auth-bridge path — this RFC only touches the remote/HTTP OAuth flow. + +## Scope + +**In scope:** `oauth.py` flow change (consent endpoint + scope, new exchange call + client, `ProxyAccessToken` shape change, simplified refresh), `mcp.py` `apply_request_config` change, unit + integration tests for the new bridge. + +**Out of scope:** the multi-project scoping tools themselves (reused unchanged); the CLI PKCE `login` flow (unaffected); the local-stdio auth-bridge/deployed-resolver path (unaffected — this only changes how the *OAuth* front door feeds a token into the *same* downstream pipe). + +## Testing / Verification + +**Unit** — mock the new internal endpoint: `authorize()` builds the `/oauth/consent` URL with `scope=claudai projectless`; `exchange_authorization_code` calls the new exchange instead of `_create_sapi_token`; error mapping matches the PHP action's declared exceptions (401/403); `apply_request_config` sets `storage_token` to the new `kbc_at_` token and it round-trips through `is_programmatic_token()` as `True`; `exchange_refresh_token` calls `refresh_tokens()` and never calls Connection's league OAuth refresh grant. + +**Integration** — full `authorize→consent→callback→token→internal-exchange` cycle against a real (dev) stack; confirm `get_accessible_projects`/`set_project_scope` work immediately after OAuth login with no project pre-selected; confirm a refresh cycle works purely via `/v1/auth/token/refresh` with no call back to Connection's OAuth server. + +**Manual** — connect a real OAuth MCP client (Claude.ai, Cursor) to a server running this change; confirm login completes and the scoping tools appear/work as expected. This is the practical local test @martin.vasko is running before finalizing implementation. + +## Decisions + +1. **Scope requested at `/oauth/consent` is `claudai projectless`** (space-separated, standard OAuth2 multi-scope) — `claudai` satisfies the exchange endpoint's `MissingClaudaiScopeException` guard; `projectless` is what makes the league token's `user_identifier` claim `admin:{id}` (not project-bound), which is what makes the *exchanged* session whole-stack. **Verify the exact literal string via the local test** — inferred from Connection's E2E test fixture comments, not from an explicit request example. +2. **Projectless = whole-stack, confirmed.** Both you and Connection's own E2E test agree: a `projectless`-scoped exchange yields an unrestricted session equivalent to a PKCE `login` lease — starts unconfirmed/whole-stack, `get_accessible_projects`/`set_project_scope` apply exactly as they do for a directly-supplied PAT today. +3. **`X-Subject-Token` confirmed** as the header name (shared constant with `resolve-storage-token`, both defined as `SUBJECT_TOKEN_HEADER = 'X-Subject-Token'` in Connection's source). Connection's E2E test also sends `X-KBC-ManageApiToken` alongside `X-Kubernetes-Authorization` — our current `StorageTokenResolver.resolve()` does **not** send that header and works today against the sibling `resolve-storage-token` endpoint, so it may be satisfied automatically for a validated k8s caller. **Verify empirically in the local test**: if the exchange 401s/403s without it, add `X-KBC-ManageApiToken` (same JWT value) to the request. +4. **No dual refresh — confirmed, not assumed.** `TokenRefreshProcessor.php` (Connection) operates on `ProgrammaticSession`/`ProgrammaticSessionRepository` — the same entity and `/v1/auth/token/refresh` mechanism already used by PAT/PKCE-login sessions, fully independent of the league OAuth session. The exchanged session refreshes on its own, forever, via the existing `refresh_tokens()`; the league OAuth token/refresh-token pair is used exactly once (at initial exchange) and never touched again, including on refresh. This **simplifies** `exchange_refresh_token()` relative to today's implementation (which currently re-negotiates with Connection's OAuth server on every refresh) rather than adding a second refresh call. +5. **No rollout coordination.** `/oauth/authorize` removal is immediate with no old-client fallback and no deploy-order dependency communicated from Connection's side. Verify locally against a real stack before shipping; no special deploy sequencing planned. From f893055e2cc78ec281f264da94497299aa55513f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 22 Jul 2026 14:42:16 +0200 Subject: [PATCH 43/89] feat(PSGO-261): implement OAuth login exchange into a programmatic session Per feature_spec/oauth_session_exchange/RFC.md: /oauth/authorize is being removed outright, so authorize() now redirects to /oauth/consent with scope "claudai projectless". After the code->token exchange, the league OAuth access token is exchanged exactly once via the new manage/internal/auth-bridge/exchange-oauth-token endpoint (OAuthSessionExchanger, sibling of StorageTokenResolver) for a whole-stack kbc_at_ programmatic session, which then feeds into the existing PSGO-261 multi-project pipe unchanged (is_programmatic_token/get_accessible_projects/set_project_scope all apply). ProxyAccessToken/ProxyRefreshToken drop the obsolete sapi_token/delegate fields (AI Service and Jobs Queue now speak bearer via bearer_or_sapi_token) and instead carry the exchanged session's access/refresh tokens directly. Refresh is decoupled from the league OAuth session per Connection's TokenRefreshProcessor: exchange_refresh_token() now calls refresh_tokens() directly instead of re-negotiating with Connection's OAuth server. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/auth_login.py | 6 +- src/keboola_mcp_server/clients/auth_bridge.py | 115 +++++++++++- src/keboola_mcp_server/mcp.py | 10 +- src/keboola_mcp_server/oauth.py | 170 +++++++----------- tests/clients/test_auth_bridge.py | 79 ++++++++ tests/test_mcp.py | 24 +++ tests/test_oauth.py | 144 ++++++++++++++- 7 files changed, 427 insertions(+), 121 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 32685327e..ac76446b8 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -79,7 +79,7 @@ def is_near_expiry(self) -> bool: return time.time() >= (self.expires_at - _REFRESH_SKEW_SECONDS) -def _parse_token_response(body: dict, *, now: float | None = None) -> TokenSet: +def parse_token_response(body: dict, *, now: float | None = None) -> TokenSet: now = time.time() if now is None else now return TokenSet( access_token=cast(str, body['accessToken']), @@ -311,7 +311,7 @@ async def exchange_code( async with httpx.AsyncClient(timeout=_AUTH_TIMEOUT, transport=transport) as client: response = await client.post(f'{_base_url(storage_api_url)}/{_TOKEN_PATH}', json=payload) response.raise_for_status() - return _parse_token_response(cast(dict, response.json())) + return parse_token_response(cast(dict, response.json())) async def refresh_tokens( @@ -326,7 +326,7 @@ async def refresh_tokens( f'{_base_url(storage_api_url)}/{_REFRESH_PATH}', json={'refreshToken': refresh_token} ) response.raise_for_status() - return _parse_token_response(cast(dict, response.json())) + return parse_token_response(cast(dict, response.json())) # --- credential storage (mode-600 file, keyed by stack host) --- diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index 315d1a32b..e918ab526 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -1,14 +1,18 @@ -"""Exchange Keboola programmatic tokens for legacy Storage tokens (PSGO-261). +"""Auth-bridge exchanges against Connection's internal endpoints (PSGO-261). -Implements the decentralized auth-bridge exchange: a programmatic bearer token -(`kbc_at_*` access token or `kbc_pat_*` personal access token) presented to the MCP -server is exchanged at Connection for a legacy Storage token, which is then used for -all downstream Storage-token APIs exactly as before. +Two exchanges live here, both authenticating to Connection with the MCP server's own +projected Kubernetes ServiceAccount JWT (`X-Kubernetes-Authorization`): -The MCP server authenticates to the resolver with its own projected Kubernetes -ServiceAccount JWT (`X-Kubernetes-Authorization`); the user's token travels as -`X-Subject-Token`. The SA token file is read per call so kubelet rotation is honored. -No token material is ever logged or placed in exception messages. +- `StorageTokenResolver`: a programmatic bearer token (`kbc_at_*`/`kbc_pat_*`) presented + to the MCP server is exchanged for a legacy Storage token, used for all downstream + Storage-token APIs exactly as before. +- `OAuthSessionExchanger`: a league OAuth access token from the remote/HTTP OAuth login + flow (`oauth.py`) is exchanged for a whole-stack Keboola programmatic session, which + then feeds into the same downstream pipe as a directly-supplied `kbc_at_*` token. + +In both cases the user's token travels as `X-Subject-Token`. The SA token file is read +per call so kubelet rotation is honored. No token material is ever logged or placed in +exception messages. """ import logging @@ -24,6 +28,7 @@ _ACCESS_TOKEN_PREFIX = 'kbc_at_' _PAT_PREFIX = 'kbc_pat_' _RESOLVE_ENDPOINT = 'manage/internal/auth-bridge/resolve-storage-token' +_EXCHANGE_OAUTH_ENDPOINT = 'manage/internal/auth-bridge/exchange-oauth-token' # Resolver statuses passed through to the client verbatim; anything else (incl. 5xx, # timeouts, network failures) is mapped to 502 Bad Gateway. _PASS_THROUGH_STATUSES = frozenset( @@ -138,3 +143,95 @@ async def resolve(self, *, subject_token: str, project_id: int) -> str: status_code=int(HTTPStatus.BAD_GATEWAY), ) return cast(str, storage_token) + + +class OAuthTokenExchangeError(RuntimeError): + """Raised when the auth-bridge fails to exchange a league OAuth token for a programmatic session. + + :ivar status_code: The client-facing HTTP status (resolver 401/403 pass through; 5xx/timeout/ + network map to 502). + """ + + def __init__(self, message: str, status_code: int) -> None: + super().__init__(message, status_code) + self.status_code = status_code + + def __str__(self) -> str: + return self.args[0] + + +class OAuthSessionExchanger: + """Exchanges a league OAuth access token (``claudai projectless`` scope) for a whole-stack + Keboola programmatic session (PSGO-261 oauth_session_exchange RFC). + + Sibling of `StorageTokenResolver`, reusing the same SA-JWT / ``X-Subject-Token`` mechanism. + """ + + def __init__( + self, + *, + storage_api_url: str, + kubernetes_token_path: str, + timeout: httpx.Timeout | None = None, + transport: httpx.AsyncBaseTransport | None = None, + ) -> None: + """ + :param storage_api_url: Connection Storage API URL (``https://connection.``). + :param kubernetes_token_path: Path to the projected ServiceAccount token file. + :param timeout: Optional HTTP timeout override. + :param transport: Optional httpx transport (for testing). + """ + self._base_url = normalize_storage_api_url(storage_api_url) + self._kubernetes_token_path = kubernetes_token_path + self._timeout = timeout or httpx.Timeout(connect=5.0, read=30.0, write=10.0, pool=5.0) + self._transport = transport + + async def exchange(self, *, oauth_access_token: str) -> dict: + """ + Exchanges ``oauth_access_token`` for a ``CliTokenResponse`` (same shape as a PKCE login). + + :return: The raw response body (``accessToken``/``refreshToken``/``expiresIn``/``sessionId``). + :raises OAuthTokenExchangeError: On any exchange failure (status carried on the error). + """ + # Connection's E2E test for this endpoint sends X-KBC-ManageApiToken alongside + # X-Kubernetes-Authorization; send both since the sibling resolve-storage-token + # endpoint only needs the latter (unconfirmed whether this one needs both too). + sa_jwt = read_service_account_jwt(self._kubernetes_token_path) + headers = { + 'Content-Type': 'application/json', + 'Accept': 'application/json', + 'X-Kubernetes-Authorization': f'Bearer {sa_jwt}', + 'X-KBC-ManageApiToken': sa_jwt, + 'X-Subject-Token': f'Bearer {strip_bearer(oauth_access_token)}', + } + try: + async with httpx.AsyncClient(timeout=self._timeout, transport=self._transport) as client: + response = await client.post(f'{self._base_url}/{_EXCHANGE_OAUTH_ENDPOINT}', headers=headers, json={}) + except httpx.HTTPError as e: + raise OAuthTokenExchangeError( + f'OAuth-token exchange could not reach Connection ({type(e).__name__}).', + status_code=int(HTTPStatus.BAD_GATEWAY), + ) from None + + if response.status_code != HTTPStatus.OK: + status = response.status_code + mapped = status if status in _PASS_THROUGH_STATUSES else int(HTTPStatus.BAD_GATEWAY) + LOG.error(f'OAuth-token exchange failed: resolver status {status}, mapped to {mapped}.') + raise OAuthTokenExchangeError( + f'OAuth-token exchange was rejected (resolver status {status}).', + status_code=mapped, + ) + + try: + body = response.json() + except ValueError: + raise OAuthTokenExchangeError( + 'OAuth-token exchange returned a non-JSON body.', + status_code=int(HTTPStatus.BAD_GATEWAY), + ) from None + if not isinstance(body, dict) or not body.get('accessToken') or not body.get('refreshToken'): + raise OAuthTokenExchangeError( + 'OAuth-token exchange returned an incomplete response.', + status_code=int(HTTPStatus.BAD_GATEWAY), + ) + return cast(dict, body) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 3dd95432d..03d698215 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -377,16 +377,14 @@ def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_sto config = dataclasses.replace(config, storage_api_url=own_stack_storage_api_url) if user := http_rq.scope.get('user'): - LOG.debug(f'Injecting bearer and SAPI tokens: user={user}, access_token={user.access_token}') + LOG.debug(f'Injecting exchanged session token: user={user}, access_token={user.access_token}') assert isinstance(user, AuthenticatedUser), f'Expecting AuthenticatedUser, got: {type(user)}' assert isinstance(user.access_token, ProxyAccessToken), ( f'Expecting ProxyAccessToken, got: {type(user.access_token)}' ) - config = dataclasses.replace( - config, - storage_token=user.access_token.sapi_token, - bearer_token=user.access_token.delegate.token, - ) + # The exchanged kbc_at_ token is a Keboola programmatic token; is_programmatic_token() + # detects it downstream and the full PSGO-261 multi-project machinery applies unchanged. + config = dataclasses.replace(config, storage_token=user.access_token.kbc_access_token) return config diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index a17871e87..db3275f5d 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -25,6 +25,10 @@ from mcp.shared.auth import InvalidRedirectUriError, OAuthClientInformationFull, OAuthToken from pydantic import AnyHttpUrl, AnyUrl +from keboola_mcp_server.auth_login import TokenSet, parse_token_response, refresh_tokens +from keboola_mcp_server.clients.auth_bridge import OAuthSessionExchanger, OAuthTokenExchangeError +from keboola_mcp_server.config import deployed_sa_token_path + LOG = logging.getLogger(__name__) _OAUTH_LOG_ALL = bool(os.getenv('KEBOOLA_MCP_SERVER_OAUTH_LOG_ALL')) _RE_LOCALHOST = re.compile(r'^(localhost|127\.0\.0\.1|\[::1]|::1)$', re.IGNORECASE) @@ -113,14 +117,18 @@ class _ExtendedAuthorizationCode(AuthorizationCode): class ProxyAccessToken(AccessToken): - delegate: AccessToken - # This token is created by the MCP server and used for calling AI Service and Jobs Queue, - # which do not support 'Authorization: Bearer ' header yet. - sapi_token: str + # The whole-stack Keboola programmatic session obtained by exchanging the league OAuth + # access token (`oauth_session_exchange` RFC). `kbc_access_token` is forwarded downstream + # as `config.storage_token`, exactly like a directly-supplied `kbc_at_*` token. + kbc_access_token: str + kbc_refresh_token: str + session_id: str | None = None class ProxyRefreshToken(RefreshToken): - delegate: RefreshToken + # The refresh side of the same exchanged session; used to refresh independently of the + # (single-use, discarded) league OAuth token pair. + kbc_refresh_token: str class SimpleOAuthProvider(OAuthProvider): @@ -153,11 +161,11 @@ def __init__( client_registration_options=ClientRegistrationOptions(enabled=True), ) - self._sapi_tokens_url = urljoin(storage_api_url, '/v2/storage/tokens') + self._storage_api_url = storage_api_url self._mcp_callback_url = urljoin(mcp_server_url, callback_endpoint) self._oauth_client_id = client_id self._oauth_client_secret = client_secret - self._oauth_server_auth_url = urljoin(server_url, '/oauth/authorize') + self._oauth_server_auth_url = urljoin(server_url, '/oauth/consent') self._oauth_server_token_url = urljoin(server_url, '/oauth/token') self._oauth_scope = scope self._jwt_secret = jwt_secret or secrets.token_hex(32) @@ -228,7 +236,9 @@ async def authorize(self, client: OAuthClientInformationFull, params: Authorizat 'response_type': 'code', 'redirect_uri': self._mcp_callback_url, 'state': state_jwt, - # send no scopes to Keboola OAuth server and let it use its own default scope + # 'claudai' satisfies the exchange endpoint's MissingClaudaiScopeException guard; + # 'projectless' makes the exchanged session whole-stack instead of project-pinned. + 'scope': 'claudai projectless', } auth_url = construct_redirect_uri(self._oauth_server_auth_url, **url_params) @@ -379,30 +389,27 @@ async def exchange_authorization_code( # Check that we get the instance loaded by load_authorization_code() function. assert isinstance(authorization_code, _ExtendedAuthorizationCode) - expires_in = max(0, int(authorization_code.oauth_access_token.expires_at - time.time())) # seconds - sapi_token = await self._create_sapi_token( - oauth_access_token=authorization_code.oauth_access_token.token, - expires_in=self._ceil_to_hour(expires_in * 2), # twice as much as the access token's time out - ) + # Exchange the league OAuth access token for a whole-stack Keboola programmatic session. + # The league token is used exactly once, here, and then never referenced again. + token_set = await self._exchange_oauth_for_session(authorization_code.oauth_access_token.token) - # wrap the access_token from the OAuth into our own access_token access_token = ProxyAccessToken( token=f'mcp_{secrets.token_hex(32)}', client_id=client.client_id, scopes=authorization_code.scopes, - expires_at=authorization_code.oauth_access_token.expires_at, - delegate=authorization_code.oauth_access_token, - sapi_token=sapi_token, + expires_at=int(token_set.expires_at), + kbc_access_token=token_set.access_token, + kbc_refresh_token=token_set.refresh_token, + session_id=token_set.session_id, ) access_token_jwt = self._encode(access_token.model_dump()) - # wrap the refresh_token from the OAuth into our own refresh_token refresh_token = ProxyRefreshToken( token=f'mcp_{secrets.token_hex(32)}', client_id=client.client_id, scopes=authorization_code.scopes, - expires_at=authorization_code.oauth_refresh_token.expires_at, - delegate=authorization_code.oauth_refresh_token, + expires_at=int(token_set.expires_at), + kbc_refresh_token=token_set.refresh_token, ) refresh_token_jwt = self._encode(refresh_token.model_dump()) @@ -410,7 +417,7 @@ async def exchange_authorization_code( access_token=access_token_jwt, refresh_token=refresh_token_jwt, token_type='Bearer', - expires_in=expires_in, + expires_in=max(0, int(token_set.expires_at - time.time())), scope=' '.join(access_token.scopes), ) @@ -486,8 +493,9 @@ async def exchange_refresh_token( scopes: list[str], ) -> OAuthToken: """ - Swaps the refresh token for a new access and refresh tokens from the OAuth server. The function also creates - a new Storage API token for accessing the AI Service and Jobs Queue APIs. + Refreshes the exchanged Keboola programmatic session directly (PSGO-261 + oauth_session_exchange RFC) — no round-trip to the league OAuth server: that token pair + was used once, at initial exchange, and is never touched again. :param client: The OAuth client details. :param refresh_token: The refresh token to use for renewing the tokens. @@ -496,7 +504,7 @@ async def exchange_refresh_token( :return: A new OAuthToken containing the access and refresh tokens. - :raises HTTPException: If the OAuth server response indicates an error. + :raises HTTPException: If the session-refresh call indicates an error. """ _log_debug( f'[exchange_refresh_token] client_id={client.client_id}, refresh_token={refresh_token}, scopes={scopes}' @@ -504,62 +512,32 @@ async def exchange_refresh_token( assert isinstance(refresh_token, ProxyRefreshToken), f'Expected ProxyRefreshToken, got {type(refresh_token)}' - # get new access and refresh tokens from the OAuth server - async with self._create_http_client() as http_client: - response = await http_client.post( - self._oauth_server_token_url, - data={ - 'client_id': self._oauth_client_id, - 'client_secret': self._oauth_client_secret, - 'grant_type': 'refresh_token', - 'refresh_token': refresh_token.delegate.token, - }, - headers={'Accept': 'application/json'}, - ) - - if response.status_code != 200: - LOG.exception( - '[exchange_refresh_token] Failed to refresh token, ' - f'OAuth server response: status={response.status_code}, text={response.text}' - ) - raise HTTPException( - 400, f'Failed to refresh token: status={response.status_code}, text={response.text}' - ) - - data = response.json() - _log_debug(f'[exchange_refresh_token] OAuth server response: {data}') - - if 'error' in data: - LOG.exception(f'[exchange_refresh_token] Error when refreshing token: data={data}') - raise HTTPException(400, data.get('error_description', data['error'])) - - oauth_access_token, oauth_refresh_token = self._read_oauth_tokens(data, scopes or refresh_token.scopes) - expires_in = max(0, int(oauth_access_token.expires_at - time.time())) # seconds - sapi_token = await self._create_sapi_token( - oauth_access_token=oauth_access_token.token, - expires_in=self._ceil_to_hour(expires_in * 2), # twice as much as the access token's time out - ) + try: + token_set = await refresh_tokens(self._storage_api_url, refresh_token=refresh_token.kbc_refresh_token) + except httpx.HTTPStatusError as e: + LOG.exception(f'[exchange_refresh_token] Failed to refresh session: status={e.response.status_code}') + raise HTTPException(400, f'Failed to refresh token: status={e.response.status_code}') from e - # wrap the access_token from the OAuth into our own access_token + new_scopes = scopes or refresh_token.scopes access_token = ProxyAccessToken( token=f'mcp_{secrets.token_hex(32)}', client_id=client.client_id, - scopes=oauth_access_token.scopes, - expires_at=oauth_access_token.expires_at, - delegate=oauth_access_token, - sapi_token=sapi_token, + scopes=new_scopes, + expires_at=int(token_set.expires_at), + kbc_access_token=token_set.access_token, + kbc_refresh_token=token_set.refresh_token, + session_id=token_set.session_id, ) access_token_jwt = self._encode(access_token.model_dump()) - # wrap the refresh_token from the OAuth into our own refresh_token - refresh_token = ProxyRefreshToken( + new_refresh_token = ProxyRefreshToken( token=f'mcp_{secrets.token_hex(32)}', client_id=client.client_id, - scopes=oauth_refresh_token.scopes, - expires_at=oauth_refresh_token.expires_at, - delegate=oauth_refresh_token, + scopes=new_scopes, + expires_at=int(token_set.expires_at), + kbc_refresh_token=token_set.refresh_token, ) - refresh_token_jwt = self._encode(refresh_token.model_dump()) + refresh_token_jwt = self._encode(new_refresh_token.model_dump()) oauth_token = OAuthToken( access_token=access_token_jwt, @@ -570,7 +548,7 @@ async def exchange_refresh_token( ) _log_debug( - f'[exchange_refresh_token] access_token={access_token}, refresh_token={refresh_token}, ' + f'[exchange_refresh_token] access_token={access_token}, refresh_token={new_refresh_token}, ' f'oauth_token={oauth_token}' ) @@ -622,39 +600,29 @@ def _read_oauth_tokens(self, data: dict[str, Any], scopes: list[str]) -> tuple[A return access_token, refresh_token - async def _create_sapi_token(self, oauth_access_token: str, expires_in: int) -> str: + async def _exchange_oauth_for_session(self, oauth_access_token: str) -> TokenSet: """ - Creates a new Storage API token for accessing AI and Jobs Queue services that do not support bearer tokens yet. + Exchanges a league OAuth access token (``claudai projectless`` scope) for a whole-stack + Keboola programmatic session via ``manage/internal/auth-bridge/exchange-oauth-token``. """ - async with self._create_http_client() as http_client: - response = await http_client.post( - self._sapi_tokens_url, - json={ - 'description': 'Created by the MCP server.', - 'expiresIn': expires_in, - 'canReadAllFileUploads': True, - 'canManageBuckets': True, - }, - headers={ - 'Accept': 'application/json', - 'Authorization': f'Bearer {oauth_access_token}', - }, - ) - - if response.status_code != 200: - LOG.error( - '[_create_sapi_token] Failed to create Storage API token, ' - f'Storage API response: status={response.status_code}, text={response.text}' - ) - raise HTTPException( - response.status_code, - f'Failed to create Storage API token: status={response.status_code}, text={response.text}', - ) - - data = response.json() - _log_debug(f'[_create_sapi_token] Storage API response: {data}') + kubernetes_token_path = deployed_sa_token_path() + if not kubernetes_token_path: + # OAuth login only runs on the deployed server; a missing SA token path means + # KBC_KUBERNETES_TOKEN_PATH isn't set there, which is a deployment misconfiguration. + raise HTTPException(500, 'OAuth login is misconfigured: no Kubernetes ServiceAccount token available.') + + exchanger = OAuthSessionExchanger( + storage_api_url=self._storage_api_url, + kubernetes_token_path=kubernetes_token_path, + ) + try: + body = await exchanger.exchange(oauth_access_token=oauth_access_token) + except OAuthTokenExchangeError as e: + LOG.error(f'[_exchange_oauth_for_session] {e}') + raise HTTPException(e.status_code, str(e)) from e - return data['token'] + _log_debug(f'[_exchange_oauth_for_session] exchange response: {body}') + return parse_token_response(body) @staticmethod def _ceil_to_hour(seconds: int) -> int: diff --git a/tests/clients/test_auth_bridge.py b/tests/clients/test_auth_bridge.py index 001b37cd6..b8eadc1b5 100644 --- a/tests/clients/test_auth_bridge.py +++ b/tests/clients/test_auth_bridge.py @@ -7,6 +7,8 @@ import pytest from keboola_mcp_server.clients.auth_bridge import ( + OAuthSessionExchanger, + OAuthTokenExchangeError, StorageTokenExchangeError, StorageTokenResolver, is_programmatic_token, @@ -121,3 +123,80 @@ async def test_resolve_empty_sa_token_file_fails_loudly(tmp_path: Path) -> None: def test_invalid_storage_api_url_rejected(sa_token_file: Path) -> None: with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): StorageTokenResolver(storage_api_url='https://example.com', kubernetes_token_path=str(sa_token_file)) + + +def _exchanger(sa_token_file: Path, handler) -> OAuthSessionExchanger: + return OAuthSessionExchanger( + storage_api_url=STORAGE_API_URL, + kubernetes_token_path=str(sa_token_file), + transport=httpx.MockTransport(handler), + ) + + +@pytest.mark.asyncio +async def test_exchange_success_sends_expected_request(sa_token_file: Path) -> None: + captured: dict[str, httpx.Request] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured['request'] = request + return httpx.Response( + HTTPStatus.OK, + json={'accessToken': 'kbc_at_new', 'refreshToken': 'kbc_rt_new', 'expiresIn': 3600, 'sessionId': 's1'}, + ) + + exchanger = _exchanger(sa_token_file, handler) + body = await exchanger.exchange(oauth_access_token='Bearer league-oauth-token') + + assert body == {'accessToken': 'kbc_at_new', 'refreshToken': 'kbc_rt_new', 'expiresIn': 3600, 'sessionId': 's1'} + rq = captured['request'] + assert rq.url.path == '/manage/internal/auth-bridge/exchange-oauth-token' + assert rq.headers['X-Kubernetes-Authorization'] == 'Bearer sa-jwt-value' + assert rq.headers['X-KBC-ManageApiToken'] == 'sa-jwt-value' + # Subject token is normalized to a single Bearer scheme regardless of inbound form. + assert rq.headers['X-Subject-Token'] == 'Bearer league-oauth-token' + + +@pytest.mark.parametrize('status', [HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN]) +@pytest.mark.asyncio +async def test_exchange_passes_through_client_errors(sa_token_file: Path, status: HTTPStatus) -> None: + exchanger = _exchanger(sa_token_file, lambda rq: httpx.Response(status, json={'error': 'nope'})) + with pytest.raises(OAuthTokenExchangeError) as exc: + await exchanger.exchange(oauth_access_token='league-oauth-token') + assert exc.value.status_code == int(status) + + +@pytest.mark.parametrize('status', [HTTPStatus.INTERNAL_SERVER_ERROR, HTTPStatus.BAD_GATEWAY, HTTPStatus.NOT_FOUND]) +@pytest.mark.asyncio +async def test_exchange_maps_other_statuses_to_502(sa_token_file: Path, status: HTTPStatus) -> None: + exchanger = _exchanger(sa_token_file, lambda rq: httpx.Response(status)) + with pytest.raises(OAuthTokenExchangeError) as exc: + await exchanger.exchange(oauth_access_token='league-oauth-token') + assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) + + +@pytest.mark.asyncio +async def test_exchange_maps_network_error_to_502(sa_token_file: Path) -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError('boom', request=request) + + exchanger = _exchanger(sa_token_file, handler) + with pytest.raises(OAuthTokenExchangeError) as exc: + await exchanger.exchange(oauth_access_token='league-oauth-token') + assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) + assert 'league-oauth-token' not in str(exc.value) + + +@pytest.mark.parametrize( + 'body', + [ + {'projectId': 1}, # missing both tokens + {'accessToken': 'kbc_at_new'}, # missing refreshToken + {'refreshToken': 'kbc_rt_new'}, # missing accessToken + ], +) +@pytest.mark.asyncio +async def test_exchange_incomplete_response_maps_to_502(sa_token_file: Path, body: dict) -> None: + exchanger = _exchanger(sa_token_file, lambda rq: httpx.Response(HTTPStatus.OK, json=body)) + with pytest.raises(OAuthTokenExchangeError) as exc: + await exchanger.exchange(oauth_access_token='league-oauth-token') + assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index bc4302cfb..961ff228d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -824,6 +824,30 @@ def test_apply_request_config_pins_storage_api_url( assert applied.storage_token == headers.get('X-Storage-Api-Token', 'server-token') assert applied.branch_id == headers.get('X-Branch-Id') + def test_apply_request_config_injects_exchanged_session_token(self): + from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser + from starlette.requests import Request + + from keboola_mcp_server.clients.auth_bridge import is_programmatic_token + from keboola_mcp_server.oauth import ProxyAccessToken + + access_token = ProxyAccessToken( + token='mcp_proxy', + client_id='claude.ai', + scopes=['claudai', 'projectless'], + expires_at=int(time.time() + 3600), + kbc_access_token='kbc_at_exchanged', + kbc_refresh_token='kbc_rt_exchanged', + session_id='session-1', + ) + http_rq = Request({'type': 'http', 'headers': [], 'user': AuthenticatedUser(access_token)}) + config = Config(storage_api_url='https://connection.test.keboola.com') + + out_config = SessionStateMiddleware.apply_request_config(http_rq, config) + + assert out_config.storage_token == 'kbc_at_exchanged' + assert is_programmatic_token(out_config.storage_token) + class TestProgrammaticTokenExchange: """SessionStateMiddleware exchanges programmatic tokens via the auth-bridge resolver (PSGO-261).""" diff --git a/tests/test_oauth.py b/tests/test_oauth.py index f80d305ef..865ec2d95 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -1,13 +1,21 @@ import time from collections.abc import Mapping +from http import HTTPStatus from typing import Any +from urllib.parse import parse_qs, urlparse import pytest -from mcp.server.auth.provider import AccessToken, RefreshToken +from mcp.server.auth.provider import AccessToken, AuthorizationParams, RefreshToken from mcp.shared.auth import InvalidRedirectUriError, OAuthClientInformationFull from pydantic import AnyHttpUrl, AnyUrl -from keboola_mcp_server.oauth import SimpleOAuthProvider, _ExtendedAuthorizationCode, _OAuthClientInformationFull +from keboola_mcp_server.clients.auth_bridge import OAuthTokenExchangeError +from keboola_mcp_server.oauth import ( + ProxyRefreshToken, + SimpleOAuthProvider, + _ExtendedAuthorizationCode, + _OAuthClientInformationFull, +) JWT_KEY = 'secret' @@ -217,3 +225,135 @@ def test_validate_redirect_uri(self, uri: AnyUrl | None, valid: bool): else: with pytest.raises(InvalidRedirectUriError): info.validate_redirect_uri(uri) + + @pytest.mark.asyncio + async def test_authorize_redirects_to_consent_with_claudai_projectless_scope( + self, oauth_provider: SimpleOAuthProvider + ): + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + params = AuthorizationParams( + redirect_uri=AnyUrl('http://foo/callback'), + redirect_uri_provided_explicitly=True, + code_challenge='challenge', + state='client-state', + scopes=None, + ) + auth_url = await oauth_provider.authorize(client, params) + + parsed = urlparse(auth_url) + assert parsed.path == '/oauth/consent' + query = parse_qs(parsed.query) + assert query['scope'] == ['claudai projectless'] + + @pytest.mark.asyncio + async def test_exchange_authorization_code_exchanges_for_session( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + from keboola_mcp_server import oauth as oauth_module + + monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: '/tmp/sa-token') + captured: dict[str, Any] = {} + + class _FakeExchanger: + def __init__(self, **kwargs): + captured['init_kwargs'] = kwargs + + async def exchange(self, *, oauth_access_token: str): + captured['oauth_access_token'] = oauth_access_token + return {'accessToken': 'kbc_at_new', 'refreshToken': 'kbc_rt_new', 'expiresIn': 3600} + + monkeypatch.setattr(oauth_module, 'OAuthSessionExchanger', _FakeExchanger) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) + + oauth_token = await oauth_provider.exchange_authorization_code(client, auth_code) + + assert captured['oauth_access_token'] == 'oauth-access-token' + loaded = await oauth_provider.load_access_token(oauth_token.access_token) + assert loaded is not None + assert loaded.kbc_access_token == 'kbc_at_new' + assert loaded.kbc_refresh_token == 'kbc_rt_new' + + @pytest.mark.asyncio + async def test_exchange_authorization_code_maps_exchange_error( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + from http.client import HTTPException + + from keboola_mcp_server import oauth as oauth_module + + monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: '/tmp/sa-token') + + class _FailingExchanger: + def __init__(self, **kwargs): + pass + + async def exchange(self, *, oauth_access_token: str): + raise OAuthTokenExchangeError('rejected', status_code=int(HTTPStatus.FORBIDDEN)) + + monkeypatch.setattr(oauth_module, 'OAuthSessionExchanger', _FailingExchanger) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) + + with pytest.raises(HTTPException) as exc: + await oauth_provider.exchange_authorization_code(client, auth_code) + assert exc.value.args[0] == int(HTTPStatus.FORBIDDEN) + + @pytest.mark.asyncio + async def test_exchange_authorization_code_missing_sa_token_path( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + from http.client import HTTPException + + from keboola_mcp_server import oauth as oauth_module + + monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: None) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) + + with pytest.raises(HTTPException) as exc: + await oauth_provider.exchange_authorization_code(client, auth_code) + assert exc.value.args[0] == 500 + + @pytest.mark.asyncio + async def test_exchange_refresh_token_calls_refresh_tokens_directly( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + from keboola_mcp_server import oauth as oauth_module + from keboola_mcp_server.auth_login import TokenSet + + captured: dict[str, Any] = {} + + async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, transport=None): + captured['storage_api_url'] = storage_api_url + captured['refresh_token'] = refresh_token + return TokenSet( + access_token='kbc_at_rotated', refresh_token='kbc_rt_rotated', expires_at=time.time() + 3600 + ) + + monkeypatch.setattr(oauth_module, 'refresh_tokens', _fake_refresh_tokens) + # If exchange_refresh_token ever called Connection's league OAuth server, this transport + # would raise, proving the refresh is fully decoupled from it (RFC Decision §4). + oauth_provider._create_http_client = lambda: (_ for _ in ()).throw( # type: ignore[method-assign] + AssertionError('exchange_refresh_token must not call the league OAuth server') + ) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + refresh_token = ProxyRefreshToken( + token='mcp_old', + client_id='foo-client-id', + scopes=['claudai', 'projectless'], + expires_at=int(time.time() + 3600), + kbc_refresh_token='kbc_rt_old', + ) + + oauth_token = await oauth_provider.exchange_refresh_token(client, refresh_token, []) + + assert captured['refresh_token'] == 'kbc_rt_old' + loaded = await oauth_provider.load_access_token(oauth_token.access_token) + assert loaded is not None + assert loaded.kbc_access_token == 'kbc_at_rotated' + assert loaded.kbc_refresh_token == 'kbc_rt_rotated' From 228d1e8fe96efbf9eba4b6dcc80e373489e725d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 22 Jul 2026 14:50:44 +0200 Subject: [PATCH 44/89] refactor(PSGO-261): dedupe StorageTokenResolver/OAuthSessionExchanger boilerplate Both auth-bridge clients shared identical __init__/SA-JWT-read code and identical exception __init__/__str__; factor into _AuthBridgeClient and _AuthBridgeExchangeError base classes, no behavior change. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/clients/auth_bridge.py | 54 ++++++------------- 1 file changed, 16 insertions(+), 38 deletions(-) diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index e918ab526..05ceaeb3e 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -51,8 +51,8 @@ def is_programmatic_token(token: str | None) -> bool: return bare.startswith(_ACCESS_TOKEN_PREFIX) or bare.startswith(_PAT_PREFIX) -class StorageTokenExchangeError(RuntimeError): - """Raised when the auth-bridge resolver fails to exchange a programmatic token. +class _AuthBridgeExchangeError(RuntimeError): + """Base for auth-bridge exchange failures. :ivar status_code: The client-facing HTTP status (resolver 400/401/403 pass through; 5xx/timeout/network map to 502). @@ -66,8 +66,12 @@ def __str__(self) -> str: return self.args[0] -class StorageTokenResolver: - """Exchanges a programmatic token for a legacy Storage token via the Connection resolver.""" +class StorageTokenExchangeError(_AuthBridgeExchangeError): + """Raised when the auth-bridge resolver fails to exchange a programmatic token.""" + + +class _AuthBridgeClient: + """Shared setup for auth-bridge clients: base URL, SA-token path, timeout, transport.""" def __init__( self, @@ -92,6 +96,10 @@ def _read_sa_jwt(self) -> str: # Read per call — the kubelet rotates the projected token in place. return read_service_account_jwt(self._kubernetes_token_path) + +class StorageTokenResolver(_AuthBridgeClient): + """Exchanges a programmatic token for a legacy Storage token via the Connection resolver.""" + async def resolve(self, *, subject_token: str, project_id: int) -> str: """ Exchanges ``subject_token`` for the legacy Storage token of ``project_id``. @@ -145,47 +153,17 @@ async def resolve(self, *, subject_token: str, project_id: int) -> str: return cast(str, storage_token) -class OAuthTokenExchangeError(RuntimeError): - """Raised when the auth-bridge fails to exchange a league OAuth token for a programmatic session. +class OAuthTokenExchangeError(_AuthBridgeExchangeError): + """Raised when the auth-bridge fails to exchange a league OAuth token for a programmatic session.""" - :ivar status_code: The client-facing HTTP status (resolver 401/403 pass through; 5xx/timeout/ - network map to 502). - """ - def __init__(self, message: str, status_code: int) -> None: - super().__init__(message, status_code) - self.status_code = status_code - - def __str__(self) -> str: - return self.args[0] - - -class OAuthSessionExchanger: +class OAuthSessionExchanger(_AuthBridgeClient): """Exchanges a league OAuth access token (``claudai projectless`` scope) for a whole-stack Keboola programmatic session (PSGO-261 oauth_session_exchange RFC). Sibling of `StorageTokenResolver`, reusing the same SA-JWT / ``X-Subject-Token`` mechanism. """ - def __init__( - self, - *, - storage_api_url: str, - kubernetes_token_path: str, - timeout: httpx.Timeout | None = None, - transport: httpx.AsyncBaseTransport | None = None, - ) -> None: - """ - :param storage_api_url: Connection Storage API URL (``https://connection.``). - :param kubernetes_token_path: Path to the projected ServiceAccount token file. - :param timeout: Optional HTTP timeout override. - :param transport: Optional httpx transport (for testing). - """ - self._base_url = normalize_storage_api_url(storage_api_url) - self._kubernetes_token_path = kubernetes_token_path - self._timeout = timeout or httpx.Timeout(connect=5.0, read=30.0, write=10.0, pool=5.0) - self._transport = transport - async def exchange(self, *, oauth_access_token: str) -> dict: """ Exchanges ``oauth_access_token`` for a ``CliTokenResponse`` (same shape as a PKCE login). @@ -196,7 +174,7 @@ async def exchange(self, *, oauth_access_token: str) -> dict: # Connection's E2E test for this endpoint sends X-KBC-ManageApiToken alongside # X-Kubernetes-Authorization; send both since the sibling resolve-storage-token # endpoint only needs the latter (unconfirmed whether this one needs both too). - sa_jwt = read_service_account_jwt(self._kubernetes_token_path) + sa_jwt = self._read_sa_jwt() headers = { 'Content-Type': 'application/json', 'Accept': 'application/json', From aaf6c9798e93a61048c91ecbebe62f0637c10dd3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 22 Jul 2026 14:58:42 +0200 Subject: [PATCH 45/89] fix(PSGO-261): map refresh network errors to HTTPException; dedupe token-wrap exchange_refresh_token only caught httpx.HTTPStatusError, so a network/timeout failure talking to Connection's refresh endpoint propagated as a raw httpx exception instead of the documented HTTPException (502, matching the sibling auth-bridge exchangers' network-error mapping). Also logs the missing KBC_KUBERNETES_TOKEN_PATH misconfiguration before raising, matching every other failure branch in the same method. Extracted _wrap_session_as_oauth_token() to remove the ~35-line duplicated ProxyAccessToken/ProxyRefreshToken/OAuthToken construction shared by exchange_authorization_code and exchange_refresh_token. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/oauth.py | 63 ++++++++++----------------------- tests/test_oauth.py | 29 +++++++++++++++ 2 files changed, 47 insertions(+), 45 deletions(-) diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index db3275f5d..3b641d0c5 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -392,41 +392,7 @@ async def exchange_authorization_code( # Exchange the league OAuth access token for a whole-stack Keboola programmatic session. # The league token is used exactly once, here, and then never referenced again. token_set = await self._exchange_oauth_for_session(authorization_code.oauth_access_token.token) - - access_token = ProxyAccessToken( - token=f'mcp_{secrets.token_hex(32)}', - client_id=client.client_id, - scopes=authorization_code.scopes, - expires_at=int(token_set.expires_at), - kbc_access_token=token_set.access_token, - kbc_refresh_token=token_set.refresh_token, - session_id=token_set.session_id, - ) - access_token_jwt = self._encode(access_token.model_dump()) - - refresh_token = ProxyRefreshToken( - token=f'mcp_{secrets.token_hex(32)}', - client_id=client.client_id, - scopes=authorization_code.scopes, - expires_at=int(token_set.expires_at), - kbc_refresh_token=token_set.refresh_token, - ) - refresh_token_jwt = self._encode(refresh_token.model_dump()) - - oauth_token = OAuthToken( - access_token=access_token_jwt, - refresh_token=refresh_token_jwt, - token_type='Bearer', - expires_in=max(0, int(token_set.expires_at - time.time())), - scope=' '.join(access_token.scopes), - ) - - _log_debug( - f'[exchange_authorization_code] access_token={access_token}, refresh_token={refresh_token},' - f'oauth_token={oauth_token}' - ) - - return oauth_token + return self._wrap_session_as_oauth_token(client, token_set, authorization_code.scopes) async def load_access_token(self, token: str) -> AccessToken | None: """ @@ -517,12 +483,20 @@ async def exchange_refresh_token( except httpx.HTTPStatusError as e: LOG.exception(f'[exchange_refresh_token] Failed to refresh session: status={e.response.status_code}') raise HTTPException(400, f'Failed to refresh token: status={e.response.status_code}') from e + except httpx.HTTPError as e: + LOG.exception(f'[exchange_refresh_token] Could not reach Connection to refresh session: {e}') + raise HTTPException(502, f'Failed to refresh token: could not reach Connection ({e}).') from e + + return self._wrap_session_as_oauth_token(client, token_set, scopes or refresh_token.scopes) - new_scopes = scopes or refresh_token.scopes + def _wrap_session_as_oauth_token( + self, client: OAuthClientInformationFull, token_set: TokenSet, scopes: list[str] + ) -> OAuthToken: + """Wraps an exchanged Keboola session (`TokenSet`) into our own proxy access/refresh tokens.""" access_token = ProxyAccessToken( token=f'mcp_{secrets.token_hex(32)}', client_id=client.client_id, - scopes=new_scopes, + scopes=scopes, expires_at=int(token_set.expires_at), kbc_access_token=token_set.access_token, kbc_refresh_token=token_set.refresh_token, @@ -530,28 +504,26 @@ async def exchange_refresh_token( ) access_token_jwt = self._encode(access_token.model_dump()) - new_refresh_token = ProxyRefreshToken( + refresh_token = ProxyRefreshToken( token=f'mcp_{secrets.token_hex(32)}', client_id=client.client_id, - scopes=new_scopes, + scopes=scopes, expires_at=int(token_set.expires_at), kbc_refresh_token=token_set.refresh_token, ) - refresh_token_jwt = self._encode(new_refresh_token.model_dump()) + refresh_token_jwt = self._encode(refresh_token.model_dump()) oauth_token = OAuthToken( access_token=access_token_jwt, refresh_token=refresh_token_jwt, token_type='Bearer', - expires_in=max(0, int(access_token.expires_at - time.time())), - scope=' '.join(access_token.scopes), + expires_in=max(0, int(token_set.expires_at - time.time())), + scope=' '.join(scopes), ) - _log_debug( - f'[exchange_refresh_token] access_token={access_token}, refresh_token={new_refresh_token}, ' + f'[_wrap_session_as_oauth_token] access_token={access_token}, refresh_token={refresh_token}, ' f'oauth_token={oauth_token}' ) - return oauth_token async def revoke_token(self, token: str, token_type_hint: str | None = None) -> None: @@ -609,6 +581,7 @@ async def _exchange_oauth_for_session(self, oauth_access_token: str) -> TokenSet if not kubernetes_token_path: # OAuth login only runs on the deployed server; a missing SA token path means # KBC_KUBERNETES_TOKEN_PATH isn't set there, which is a deployment misconfiguration. + LOG.error('[_exchange_oauth_for_session] KBC_KUBERNETES_TOKEN_PATH is not set; cannot exchange session.') raise HTTPException(500, 'OAuth login is misconfigured: no Kubernetes ServiceAccount token available.') exchanger = OAuthSessionExchanger( diff --git a/tests/test_oauth.py b/tests/test_oauth.py index 865ec2d95..f0f00040e 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -4,6 +4,7 @@ from typing import Any from urllib.parse import parse_qs, urlparse +import httpx import pytest from mcp.server.auth.provider import AccessToken, AuthorizationParams, RefreshToken from mcp.shared.auth import InvalidRedirectUriError, OAuthClientInformationFull @@ -357,3 +358,31 @@ async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, tran assert loaded is not None assert loaded.kbc_access_token == 'kbc_at_rotated' assert loaded.kbc_refresh_token == 'kbc_rt_rotated' + + @pytest.mark.asyncio + async def test_exchange_refresh_token_maps_network_error_to_http_exception( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + from http.client import HTTPException + + from keboola_mcp_server import oauth as oauth_module + + async def _failing_refresh_tokens(storage_api_url: str, *, refresh_token: str, transport=None): + raise httpx.ConnectError('boom') + + monkeypatch.setattr(oauth_module, 'refresh_tokens', _failing_refresh_tokens) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + refresh_token = ProxyRefreshToken( + token='mcp_old', + client_id='foo-client-id', + scopes=['claudai', 'projectless'], + expires_at=int(time.time() + 3600), + kbc_refresh_token='kbc_rt_old', + ) + + # A network failure talking to Connection must surface as a clean HTTPException, not + # propagate as a raw httpx error (which the caller has no reason to expect/handle). + with pytest.raises(HTTPException) as exc: + await oauth_provider.exchange_refresh_token(client, refresh_token, []) + assert exc.value.args[0] == 502 From e8ff64098828c6b4f2016453b5bbf4cce03f9048 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 22 Jul 2026 15:27:10 +0200 Subject: [PATCH 46/89] fix(PSGO-261): stop logging token material; don't require project_id for OAuth sessions apply_request_config logged the full ProxyAccessToken via %s, whose default pydantic repr includes the raw kbc_access_token/kbc_refresh_token -- log only the client_id/session_id instead. create_session_state unconditionally exchanged a deployed programmatic token via the auth-bridge resolver, which requires config.project_id. A freshly OAuth-exchanged session starts whole-stack/projectless (RFC decision 2), so every OAuth login on the deployed server would fail at session-state init before get_accessible_projects/set_project_scope could ever run. Now the resolver exchange only happens once a project is already known; otherwise the token is forwarded as a Bearer (matching the existing local-session and per-project fan-out behavior) so project scoping works first. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/mcp.py | 23 +++++++++++++++-------- tests/test_mcp.py | 23 +++++++++++++++++++++++ 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 03d698215..dda79985d 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -377,11 +377,16 @@ def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_sto config = dataclasses.replace(config, storage_api_url=own_stack_storage_api_url) if user := http_rq.scope.get('user'): - LOG.debug(f'Injecting exchanged session token: user={user}, access_token={user.access_token}') assert isinstance(user, AuthenticatedUser), f'Expecting AuthenticatedUser, got: {type(user)}' assert isinstance(user.access_token, ProxyAccessToken), ( f'Expecting ProxyAccessToken, got: {type(user.access_token)}' ) + # Log only non-sensitive identifiers; ProxyAccessToken's default repr includes the raw + # kbc_access_token/kbc_refresh_token, which must never be logged. + LOG.debug( + f'Injecting exchanged session token: client_id={user.access_token.client_id}, ' + f'session_id={user.access_token.session_id}' + ) # The exchanged kbc_at_ token is a Keboola programmatic token; is_programmatic_token() # detects it downstream and the full PSGO-261 multi-project machinery applies unchanged. config = dataclasses.replace(config, storage_token=user.access_token.kbc_access_token) @@ -590,16 +595,18 @@ async def create_session_state( bearer_token = config.bearer_token extra_headers: dict[str, Any] = {} if is_programmatic_token(storage_token): - if deployed_sa_token_path(): - # Deployed: exchange the programmatic token (kbc_at_/kbc_pat_) for the project's - # legacy Storage token via the auth-bridge resolver, then use it downstream unchanged. + if deployed_sa_token_path() and config.project_id: + # Deployed, and a project is already known (header, or a prior scope selection): + # exchange the programmatic token for that project's legacy Storage token via the + # auth-bridge resolver, then use it downstream unchanged. storage_token = await cls._exchange_programmatic_token(config) bearer_token = None else: - # Local: no projected SA token to reach the resolver. Forward the programmatic token - # downstream as a Bearer and let PAT-aware services exchange it; name the target - # project when one has been selected. Strip any inbound `Bearer ` scheme so the - # client's own `Bearer ` prefixing can't produce `Authorization: Bearer Bearer …`. + # No projected SA token to reach the resolver (local), OR no project is known yet + # (deployed, e.g. a freshly-exchanged whole-stack OAuth session pre-scoping): forward + # the programmatic token downstream as a Bearer so get_accessible_projects/ + # set_project_scope can introspect/scope it directly. Strip any inbound `Bearer ` + # scheme so the client's own `Bearer ` prefixing can't produce `Bearer Bearer …`. bearer_token = strip_bearer(storage_token) if config.project_id: extra_headers['X-KBC-ProjectId'] = config.project_id diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 961ff228d..df6096d85 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -891,6 +891,29 @@ async def test_happy_path_calls_resolver(self, monkeypatch) -> None: ) resolver.resolve.assert_awaited_once_with(subject_token='kbc_at_abc', project_id=42) + @pytest.mark.asyncio + async def test_deployed_without_project_id_forwards_bearer_instead_of_exchanging(self, monkeypatch) -> None: + """A deployed OAuth session starts whole-stack (no project_id yet, RFC decision §2); it must + forward the programmatic token as a Bearer for get_accessible_projects/set_project_scope to + introspect, not fail by attempting a resolver exchange that requires a project id.""" + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_abc') + runtime_info = ServerRuntimeInfo(transport='http') + + with ( + patch.object( + SessionStateMiddleware, + '_exchange_programmatic_token', + AsyncMock(side_effect=AssertionError('must not exchange without a known project_id')), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')), + ): + state = await SessionStateMiddleware.create_session_state(config, runtime_info) + + client = state[KeboolaClient.STATE_KEY] + assert client.bearer_token == 'kbc_at_abc' + assert client.token == 'kbc_at_abc' + class TestMaybeUseStoredSession: """Local HTTP with no token falls back to the stored PKCE session (PSGO-261).""" From 8ecc889f8206e8addcd8bd29a59de359e3786a6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 22 Jul 2026 15:36:33 +0200 Subject: [PATCH 47/89] fix(PSGO-261): correct with_llm_instruction wording; log dialect failures with traceback get_accessible_projects's with_llm_instruction description referenced a non-existent llm_instructions field; the response field is base_instructions. The best-effort SQL-dialect resolution swallowed the per-project exception without exc_info, losing the traceback needed to debug why verify_token failed for that project. Co-Authored-By: Claude Sonnet 5 --- TOOLS.md | 2 +- src/keboola_mcp_server/tools/project.py | 4 ++-- tests/tools/test_project.py | 23 +++++++++++++++++++++++ 3 files changed, 26 insertions(+), 3 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index 9b78363c6..76f05d827 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -3579,7 +3579,7 @@ call to also receive the base working instructions grouped by dialect. "properties": { "with_llm_instruction": { "default": false, - "description": "If true, include the base working instructions (llm_instructions), grouped by SQL dialect. Request this once at the very start of a conversation; omit it on later calls.", + "description": "If true, include the base working instructions (base_instructions), grouped by SQL dialect. Request this once at the very start of a conversation; omit it on later calls.", "type": "boolean" } }, diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 0c71a0c04..878d4f562 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -344,7 +344,7 @@ async def get_accessible_projects( bool, Field( description=( - 'If true, include the base working instructions (llm_instructions), grouped by SQL dialect. ' + 'If true, include the base working instructions (base_instructions), grouped by SQL dialect. ' 'Request this once at the very start of a conversation; omit it on later calls.' ) ), @@ -379,7 +379,7 @@ async def get_accessible_projects( if isinstance(result, asyncio.CancelledError): raise result # never swallow cancellation — let it propagate if isinstance(result, BaseException): - LOG.warning(f'Could not resolve SQL dialect for a project: {result}') + LOG.warning(f'Could not resolve SQL dialect for a project: {result}', exc_info=result) continue pid, dialect = result dialects[pid] = dialect diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index f2f4dce0a..70670aaf3 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -368,6 +368,29 @@ async def test_get_accessible_projects_unknown_dialect_omits_snowflake_guidance( assert group.instructions != get_project_system_prompt('Snowflake') +@pytest.mark.asyncio +async def test_get_accessible_projects_logs_dialect_failure_with_traceback( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + # A per-project dialect-resolution failure is swallowed (best-effort), but must still log with + # exc_info so the traceback isn't lost. + _prep_client(mcp_context_client, mocker) + introspection = SimpleNamespace(user_email='m@k.com', projects=[SimpleNamespace(id=42, name='X', role='admin')]) + mocker.patch('keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection)) + mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) + mocker.patch( + 'keboola_mcp_server.tools.project._project_sql_dialect', + new=mocker.AsyncMock(side_effect=RuntimeError('verify failed')), + ) + log_warning = mocker.patch('keboola_mcp_server.tools.project.LOG.warning') + + result = await get_accessible_projects(mcp_context_client) + + assert result.projects[0].sql_dialect is None + log_warning.assert_called_once() + assert log_warning.call_args.kwargs.get('exc_info') is not None + + @pytest.mark.asyncio async def test_set_project_scope_subset_exchanges_and_stores( mcp_context_client: Context, mocker: MockerFixture From 5f414c30ba0a2feafaf83c26933f5c25ea631680 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 23 Jul 2026 08:50:03 +0200 Subject: [PATCH 48/89] fix(PSGO-261): use per-request Storage API URL in fan-out; don't duplicate refresh token MultiProjectMiddleware.client_for_project() hard-coded server_state.config.storage_api_url (the startup/lifespan config), not the current request's URL -- wrong for streamable-HTTP setups that supply the URL per request via headers. It now takes storage_api_url explicitly, threaded from the active KeboolaClient's own URL at each call site (on_call_tool's fan-out, and tools/project.py's per-project dialect verify). _swap_project() also built its per-project WorkspaceManager without kubernetes_token_path, unlike create_session_state, breaking the SA-token step-up on the deployed server; it now passes it through. ProxyAccessToken no longer duplicates kbc_refresh_token (only ProxyRefreshToken, which is what exchange_refresh_token actually consumes) -- access tokens are sent/handled far more often, so carrying the longer-lived refresh token on them needlessly widened its exposure. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/mcp.py | 24 +++++-- src/keboola_mcp_server/oauth.py | 7 +- src/keboola_mcp_server/tools/project.py | 6 +- tests/test_mcp.py | 91 +++++++++++++++++++++---- tests/test_oauth.py | 11 ++- tests/tools/test_project.py | 6 +- 6 files changed, 112 insertions(+), 33 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index dda79985d..4747763ef 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -932,8 +932,12 @@ async def on_call_tool( server_state = ServerState.from_context(ctx) original_client = state.get(KeboolaClient.STATE_KEY) original_workspace = state.get(WorkspaceManager.STATE_KEY) + is_real_client = isinstance(original_client, KeboolaClient) # Default (auto-leased) scope carries no minted token; fall back to the active client's token. - base_token = scope.scoped_token or (original_client.token if isinstance(original_client, KeboolaClient) else '') + base_token = scope.scoped_token or (original_client.token if is_real_client else '') + # The active client's own URL — the current request/session's, not the startup config's + # (which can differ or be unset for streamable-HTTP setups that supply it per request). + storage_api_url = original_client.storage_api_url if is_real_client else server_state.config.storage_api_url # A single target (scope of one, or narrowed to one via the filter) runs once against that # project only — one call, that project's X-KBC-ProjectId, no per-project envelope. @@ -942,7 +946,7 @@ async def on_call_tool( if target == scope.active_project_id: return await call_next(context) try: - await self._swap_project(state, server_state, base_token, target, scope.read_only) + await self._swap_project(state, server_state, storage_api_url, base_token, target, scope.read_only) return await call_next(context) finally: state[KeboolaClient.STATE_KEY] = original_client @@ -952,7 +956,7 @@ async def on_call_tool( errors: list[tuple[int, str]] = [] try: for project_id in targets: - await self._swap_project(state, server_state, base_token, project_id, scope.read_only) + await self._swap_project(state, server_state, storage_api_url, base_token, project_id, scope.read_only) # Isolate per-project failures: one project's error (e.g. Queue 401, a transient 5xx) # must not discard the other projects' good results. Collect it and keep going, so the # agent gets a partial response plus a retry hint. CancelledError is BaseException, so @@ -1026,6 +1030,7 @@ async def _swap_project( cls, state: dict[str, Any], server_state: ServerState, + storage_api_url: str, base_token: str, project_id: int, read_only: bool, @@ -1038,19 +1043,24 @@ async def _swap_project( Note: rebuilt per call; caching across calls would need a store that survives the per-request state rebuild — add if provisioning latency shows up in practice. """ - client = await cls.client_for_project(server_state, base_token, project_id, read_only) + client = await cls.client_for_project(server_state, storage_api_url, base_token, project_id, read_only) state[KeboolaClient.STATE_KEY] = client - state[WorkspaceManager.STATE_KEY] = await WorkspaceManager.create(client, server_state.config.workspace_schema) + state[WorkspaceManager.STATE_KEY] = await WorkspaceManager.create( + client, server_state.config.workspace_schema, kubernetes_token_path=deployed_sa_token_path() + ) @staticmethod async def client_for_project( - server_state: ServerState, token: str, project_id: int, read_only: bool + server_state: ServerState, storage_api_url: str, token: str, project_id: int, read_only: bool ) -> KeboolaClient: + # `storage_api_url` is the current request/session URL (e.g. the active `KeboolaClient`'s), + # not `server_state.config.storage_api_url` — that's the startup/lifespan config, which can + # differ (or be unset) for streamable-HTTP setups that supply the URL per request. # Normalize any inbound `Bearer ` scheme; KeboolaClient adds it back for bearer tokens, # so a pre-prefixed value would otherwise become `Authorization: Bearer Bearer …`. token = strip_bearer(token) return await KeboolaClient( - storage_api_url=server_state.config.storage_api_url, + storage_api_url=storage_api_url, storage_api_token=token, bearer_token=token, headers={ diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index 3b641d0c5..eca1f4c8c 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -119,9 +119,11 @@ class _ExtendedAuthorizationCode(AuthorizationCode): class ProxyAccessToken(AccessToken): # The whole-stack Keboola programmatic session obtained by exchanging the league OAuth # access token (`oauth_session_exchange` RFC). `kbc_access_token` is forwarded downstream - # as `config.storage_token`, exactly like a directly-supplied `kbc_at_*` token. + # as `config.storage_token`, exactly like a directly-supplied `kbc_at_*` token. The refresh + # token is deliberately NOT carried here (only on `ProxyRefreshToken`, which is what + # `exchange_refresh_token` actually receives) — access tokens are sent/handled far more often, + # so duplicating the longer-lived refresh token onto them would needlessly widen its exposure. kbc_access_token: str - kbc_refresh_token: str session_id: str | None = None @@ -499,7 +501,6 @@ def _wrap_session_as_oauth_token( scopes=scopes, expires_at=int(token_set.expires_at), kbc_access_token=token_set.access_token, - kbc_refresh_token=token_set.refresh_token, session_id=token_set.session_id, ) access_token_jwt = self._encode(access_token.model_dump()) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 878d4f562..e097b4462 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -323,7 +323,7 @@ class ProjectScope(BaseModel): async def _project_sql_dialect( - server_state: ServerState, subject_token: str, project_id: int + server_state: ServerState, storage_api_url: str, subject_token: str, project_id: int ) -> tuple[int, str | None]: """Fetches one project's SQL dialect by verifying the parent token narrowed with X-KBC-ProjectId. @@ -331,7 +331,7 @@ async def _project_sql_dialect( a single cheap Storage API call per project. """ per_client = await MultiProjectMiddleware.client_for_project( - server_state, subject_token, project_id, read_only=True + server_state, storage_api_url, subject_token, project_id, read_only=True ) token_data = await per_client.storage_client.verify_token() return project_id, _sql_dialect_from_token(token_data) @@ -373,7 +373,7 @@ async def get_accessible_projects( dialects: dict[int, str | None] = {} results = await process_concurrently( [p.id for p in introspection.projects], - lambda pid: _project_sql_dialect(server_state, subject_token, pid), + lambda pid: _project_sql_dialect(server_state, client.storage_api_url, subject_token, pid), ) for result in results: if isinstance(result, asyncio.CancelledError): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index df6096d85..3c93d63b0 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -837,7 +837,6 @@ def test_apply_request_config_injects_exchanged_session_token(self): scopes=['claudai', 'projectless'], expires_at=int(time.time() + 3600), kbc_access_token='kbc_at_exchanged', - kbc_refresh_token='kbc_rt_exchanged', session_id='session-1', ) http_rq = Request({'type': 'http', 'headers': [], 'user': AuthenticatedUser(access_token)}) @@ -1194,9 +1193,13 @@ async def call_next(_): patch.object( MultiProjectMiddleware, 'client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), - patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): result = await MultiProjectMiddleware().on_call_tool(context, call_next) @@ -1212,6 +1215,40 @@ async def call_next(_): # Structured output is deep-merged (list fields concatenated) so it still validates the schema. assert result.structured_content == {'rows': ['rows', 'rows']} + @pytest.mark.asyncio + async def test_swap_project_uses_active_client_url_and_sa_token_path(self, monkeypatch) -> None: + # _swap_project must use the CURRENT request's Storage API URL (the active client's), not + # server_state.config's startup/lifespan URL, and must pass the deployed SA token path + # through to WorkspaceManager.create exactly like create_session_state does. + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True) + # server_state.config carries a different (stale/absent) URL than the active request client. + state[KeboolaClient.STATE_KEY] = KeboolaClient( + storage_api_url='https://connection.request.keboola.com', storage_api_token='kbc_at_s' + ) + seen_calls: list = [] + + async def fake_client_for_project(_ss, storage_api_url, _token, pid, _ro): + seen_calls.append((storage_api_url, pid)) + return f'client-{pid}' + + async def call_next(_): + return self._result('rows') + + with ( + patch.object(MultiProjectMiddleware, 'client_for_project', AsyncMock(side_effect=fake_client_for_project)), + patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')) as ws_create, + ): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_calls == [ + ('https://connection.request.keboola.com', 11), + ('https://connection.request.keboola.com', 22), + ] + for call in ws_create.await_args_list: + assert call.kwargs.get('kubernetes_token_path') == '/var/run/secrets/token' + @pytest.mark.asyncio async def test_query_data_targets_single_project_workspace(self) -> None: # query_data is no longer excluded: with the project_ids filter it runs once against that @@ -1228,9 +1265,13 @@ async def call_next(_): patch.object( MultiProjectMiddleware, 'client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), - patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): result = await MultiProjectMiddleware().on_call_tool(context, call_next) @@ -1255,9 +1296,13 @@ async def call_next(_): patch.object( MultiProjectMiddleware, 'client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), - patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): result = await MultiProjectMiddleware().on_call_tool(context, call_next) @@ -1281,9 +1326,13 @@ async def call_next(_): patch.object( MultiProjectMiddleware, 'client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), - patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): await MultiProjectMiddleware().on_call_tool(context, call_next) @@ -1386,9 +1435,13 @@ async def call_next(_): patch.object( MultiProjectMiddleware, 'client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), - patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): result = await MultiProjectMiddleware().on_call_tool(context, call_next) @@ -1409,9 +1462,13 @@ async def call_next(_): patch.object( MultiProjectMiddleware, 'client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), - patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): with pytest.raises(ToolError, match='failed for all 2 scoped'): await MultiProjectMiddleware().on_call_tool(context, call_next) @@ -1432,9 +1489,13 @@ async def call_next(_): patch.object( MultiProjectMiddleware, 'client_for_project', - AsyncMock(side_effect=lambda _ss, _token, pid, _ro: f'client-{pid}'), + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), - patch.object(WorkspaceManager, 'create', AsyncMock(side_effect=lambda client, _schema: f'wsm-{client}')), ): with pytest.raises(PydanticValidationError): await MultiProjectMiddleware().on_call_tool(context, call_next) diff --git a/tests/test_oauth.py b/tests/test_oauth.py index f0f00040e..0208fdcc1 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -274,7 +274,12 @@ async def exchange(self, *, oauth_access_token: str): loaded = await oauth_provider.load_access_token(oauth_token.access_token) assert loaded is not None assert loaded.kbc_access_token == 'kbc_at_new' - assert loaded.kbc_refresh_token == 'kbc_rt_new' + # The refresh token is carried on ProxyRefreshToken only, not duplicated onto the (more + # frequently sent/handled) access token. + assert not hasattr(loaded, 'kbc_refresh_token') + loaded_refresh = await oauth_provider.load_refresh_token(client, oauth_token.refresh_token) + assert loaded_refresh is not None + assert loaded_refresh.kbc_refresh_token == 'kbc_rt_new' @pytest.mark.asyncio async def test_exchange_authorization_code_maps_exchange_error( @@ -357,7 +362,9 @@ async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, tran loaded = await oauth_provider.load_access_token(oauth_token.access_token) assert loaded is not None assert loaded.kbc_access_token == 'kbc_at_rotated' - assert loaded.kbc_refresh_token == 'kbc_rt_rotated' + loaded_refresh = await oauth_provider.load_refresh_token(client, oauth_token.refresh_token) + assert loaded_refresh is not None + assert loaded_refresh.kbc_refresh_token == 'kbc_rt_rotated' @pytest.mark.asyncio async def test_exchange_refresh_token_maps_network_error_to_http_exception( diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 70670aaf3..420a9cf53 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -286,7 +286,7 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock dialects = {18: 'BigQuery', 83: 'Snowflake'} mocker.patch( 'keboola_mcp_server.tools.project._project_sql_dialect', - new=mocker.AsyncMock(side_effect=lambda _ss, _tok, pid: (pid, dialects[pid])), + new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, dialects[pid])), ) # No scope confirmed yet. @@ -329,7 +329,7 @@ async def test_get_accessible_projects_llm_instructions_grouped_by_dialect( dialects = {18: 'BigQuery', 86: 'BigQuery', 95: 'Snowflake'} mocker.patch( 'keboola_mcp_server.tools.project._project_sql_dialect', - new=mocker.AsyncMock(side_effect=lambda _ss, _tok, pid: (pid, dialects[pid])), + new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, dialects[pid])), ) result = await get_accessible_projects(mcp_context_client, with_llm_instruction=True) @@ -355,7 +355,7 @@ async def test_get_accessible_projects_unknown_dialect_omits_snowflake_guidance( mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) mocker.patch( 'keboola_mcp_server.tools.project._project_sql_dialect', - new=mocker.AsyncMock(side_effect=lambda _ss, _tok, pid: (pid, None)), + new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, None)), ) result = await get_accessible_projects(mcp_context_client, with_llm_instruction=True) From f49c74a6d75449892f7e5b91a0f3c0cd8fc5c1dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 23 Jul 2026 09:06:32 +0200 Subject: [PATCH 49/89] fix(PSGO-261): re-run per-project authorization during fan-out; decouple refresh-token TTL MultiProjectMiddleware ran after ToolsFilteringMiddleware in the middleware list, i.e. inner relative to it -- so per-project feature/role/branch authorization was evaluated once against the pre-fan-out client and never re-checked against each swapped project during multi-project fan-out, contradicting the documented invariant that "the on_call_tool guards still enforce every feature/role/branch rule per project". Swapped the order so MultiProjectMiddleware wraps ToolsFilteringMiddleware. ProxyRefreshToken's expires_at was tied to the exchanged session's (short, ~1h) access-token expiry instead of a longer independent window, which would force a re-login every ~1h even though the underlying Keboola refresh token can keep the session alive indefinitely (RFC Decision 4). Restored the ~7-day-cap derivation the pre-exchange code used for the league OAuth refresh token. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/oauth.py | 9 ++++++++- src/keboola_mcp_server/server.py | 6 +++++- tests/test_oauth.py | 4 ++++ 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index eca1f4c8c..13c8b2e1f 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -505,11 +505,18 @@ def _wrap_session_as_oauth_token( ) access_token_jwt = self._encode(access_token.model_dump()) + # The proxy refresh token's own expiry must NOT be tied to the (short-lived) access token's: + # the underlying Keboola refresh token keeps the session alive indefinitely (RFC Decision §4), + # but the mcp SDK enforces `expires_at` on the object load_refresh_token() returns. Derive a + # longer window the same way the pre-exchange code did for the league OAuth refresh token + # (up to ~7 days), so a client that doesn't refresh for a while isn't forced to re-login. + access_expires_in = max(0, int(token_set.expires_at - time.time())) + refresh_expires_at = int(time.time()) + self._ceil_to_hour(min(168 * access_expires_in, 168 * 3600)) refresh_token = ProxyRefreshToken( token=f'mcp_{secrets.token_hex(32)}', client_id=client.client_id, scopes=scopes, - expires_at=int(token_set.expires_at), + expires_at=refresh_expires_at, kbc_refresh_token=token_set.refresh_token, ) refresh_token_jwt = self._encode(refresh_token.model_dump()) diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index eb9d01236..de9bb8276 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -251,8 +251,12 @@ def create_server( LoggingMiddleware(log_level=logging.DEBUG), SessionStateMiddleware(), ToolAuthorizationMiddleware(), - ToolsFilteringMiddleware(), + # MultiProjectMiddleware must wrap ToolsFilteringMiddleware (run first in this list = + # outer), not the reverse: it swaps the active KeboolaClient per project during fan-out, + # and ToolsFilteringMiddleware's per-project feature/role/branch checks must be + # re-evaluated against each swapped client — not just once against the pre-fan-out client. MultiProjectMiddleware(), + ToolsFilteringMiddleware(), ValidationErrorMiddleware(), ], ) diff --git a/tests/test_oauth.py b/tests/test_oauth.py index 0208fdcc1..6812608be 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -280,6 +280,10 @@ async def exchange(self, *, oauth_access_token: str): loaded_refresh = await oauth_provider.load_refresh_token(client, oauth_token.refresh_token) assert loaded_refresh is not None assert loaded_refresh.kbc_refresh_token == 'kbc_rt_new' + # The refresh token's own expiry must be much longer than the (1h) access token's -- it must + # not be tied to it, or the mcp SDK would force a re-login every ~1h even though the + # underlying Keboola refresh token can keep the session alive indefinitely. + assert loaded_refresh.expires_at - loaded.expires_at > 6 * 24 * 3600 # at least ~6 more days @pytest.mark.asyncio async def test_exchange_authorization_code_maps_exchange_error( From d64779eeee639ca315e204c7cd03a61bbaf2a00e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 23 Jul 2026 09:17:13 +0200 Subject: [PATCH 50/89] fix(PSGO-261): validate _parent_subject_token's bearer is actually programmatic Only checked client.bearer_token is non-None; a non-programmatic bearer or a Bearer-prefixed value would reach introspect_token/exchange_scoped_token directly and fail with an unclear downstream 401/400 instead of the intended explicit validation error. Now uses is_programmatic_token()/strip_bearer(), matching the error message's own claim ("requires a Keboola programmatic token"). Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/tools/project.py | 5 +++-- tests/tools/test_project.py | 28 ++++++++++++++++++++++--- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index e097b4462..41cfc53c7 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -9,6 +9,7 @@ from pydantic import BaseModel, Field from keboola_mcp_server.auth_login import exchange_scoped_token, get_access_token, introspect_token +from keboola_mcp_server.clients.auth_bridge import is_programmatic_token, strip_bearer from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField @@ -73,7 +74,7 @@ async def _parent_subject_token(client: KeboolaClient) -> str: from the parent, never from an already-narrowed scoped token); falls back to whatever bearer the client currently carries (a directly-supplied PAT, or an HTTP bearer). """ - if client.bearer_token is None: + if not is_programmatic_token(client.bearer_token): raise ValueError( 'Project scoping requires a Keboola programmatic token (kbc_at_/kbc_pat_). ' 'Run "keboola-mcp-server login --api-url " first, or supply such a token.' @@ -81,7 +82,7 @@ async def _parent_subject_token(client: KeboolaClient) -> str: try: return await get_access_token(client.storage_api_url) except RuntimeError: - return client.bearer_token + return strip_bearer(cast(str, client.bearer_token)) async def _resolve_branch_context(client: KeboolaClient) -> tuple[str | int, str, bool]: diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 420a9cf53..70cac9bc8 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -487,10 +487,32 @@ async def test_set_project_scope_falls_back_on_network_error( @pytest.mark.asyncio -async def test_scope_requires_programmatic_token(mcp_context_client: Context, mocker: MockerFixture) -> None: - _prep_client(mcp_context_client, mocker, bearer=None) - with pytest.raises(ValueError, match='programmatic token'): +@pytest.mark.parametrize( + 'bearer', + [ + None, # no bearer at all + 'legacy-sapi-token-123', # a non-programmatic bearer must not be accepted either + 'Bearer kbc_at_prefixed', # accepted, but exercises the strip_bearer normalization path + ], + ids=['no_bearer', 'non_programmatic_bearer', 'bearer_prefixed'], +) +async def test_scope_requires_programmatic_token( + mcp_context_client: Context, mocker: MockerFixture, bearer: str | None +) -> None: + _prep_client(mcp_context_client, mocker, bearer=bearer) + mocker.patch('keboola_mcp_server.tools.project.get_access_token', new=mocker.AsyncMock(side_effect=RuntimeError)) + if bearer == 'Bearer kbc_at_prefixed': + introspection = SimpleNamespace(user_email='m@k.com', projects=[]) + introspect = mocker.patch( + 'keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection) + ) + mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) await get_accessible_projects(mcp_context_client) + # The inbound bearer's `Bearer ` scheme must be stripped before use as a subject token. + introspect.assert_awaited_once_with(STACK, subject_token='kbc_at_prefixed') + else: + with pytest.raises(ValueError, match='programmatic token'): + await get_accessible_projects(mcp_context_client) @pytest.mark.asyncio From 4d7a5d6a938801545cb0d9407e365e37c610d49c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 23 Jul 2026 15:18:42 +0200 Subject: [PATCH 51/89] fix(PSGO-261): raise the right exception class for OAuth-visible errors oauth.py raised the stdlib http.client.HTTPException everywhere, but its two callers each expect a different type: - server.py's oauth_callback_handler (custom /oauth/callback route) catches starlette.exceptions.HTTPException specifically to re-raise it; the stdlib class doesn't match, so every intentional error there (invalid state, failed code exchange, expired token) fell through to the generic except-Exception branch and came back as an opaque 500. - exchange_authorization_code/exchange_refresh_token are invoked by the mcp SDK's own /token endpoint handler, which only recognizes its own TokenError and turns it into a spec-compliant TokenErrorResponse body ({"error": ..., "error_description": ...}). A bare HTTPException there bubbles up uncaught and reaches OAuth clients as a non-compliant error shape that fails their own response validation (confirmed live: a real client rejected it with a Zod schema error). Split accordingly: handle_oauth_callback/_read_oauth_tokens now raise starlette.exceptions.HTTPException (matching their only caller's except clause); exchange_refresh_token/_exchange_oauth_for_session now raise TokenError (matching the mcp SDK's /token handler contract). Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/oauth.py | 28 +++++++++++++++++++++++----- tests/test_oauth.py | 28 ++++++++++++---------------- tests/test_server.py | 20 +++++++++++++++++++- 3 files changed, 54 insertions(+), 22 deletions(-) diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index 13c8b2e1f..d261868d3 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -7,7 +7,6 @@ import secrets import time from collections.abc import Mapping -from http.client import HTTPException from typing import Any, cast from urllib.parse import urljoin @@ -19,11 +18,13 @@ AuthorizationCode, AuthorizationParams, RefreshToken, + TokenError, construct_redirect_uri, ) from mcp.server.auth.settings import ClientRegistrationOptions from mcp.shared.auth import InvalidRedirectUriError, OAuthClientInformationFull, OAuthToken from pydantic import AnyHttpUrl, AnyUrl +from starlette.exceptions import HTTPException from keboola_mcp_server.auth_login import TokenSet, parse_token_response, refresh_tokens from keboola_mcp_server.clients.auth_bridge import OAuthSessionExchanger, OAuthTokenExchangeError @@ -480,14 +481,23 @@ async def exchange_refresh_token( assert isinstance(refresh_token, ProxyRefreshToken), f'Expected ProxyRefreshToken, got {type(refresh_token)}' + # Raised as TokenError (not HTTPException): this method is invoked by the mcp SDK's own + # /token endpoint handler, which only recognizes TokenError and formats it into a spec- + # compliant TokenErrorResponse body ({"error": ..., "error_description": ...}) -- an + # HTTPException here would bubble up uncaught and reach the client as an opaque, non-OAuth + # shaped error. try: token_set = await refresh_tokens(self._storage_api_url, refresh_token=refresh_token.kbc_refresh_token) except httpx.HTTPStatusError as e: LOG.exception(f'[exchange_refresh_token] Failed to refresh session: status={e.response.status_code}') - raise HTTPException(400, f'Failed to refresh token: status={e.response.status_code}') from e + raise TokenError( + error='invalid_grant', error_description=f'Failed to refresh token: status={e.response.status_code}' + ) from e except httpx.HTTPError as e: LOG.exception(f'[exchange_refresh_token] Could not reach Connection to refresh session: {e}') - raise HTTPException(502, f'Failed to refresh token: could not reach Connection ({e}).') from e + raise TokenError( + error='invalid_grant', error_description=f'Failed to refresh token: could not reach Connection ({e}).' + ) from e return self._wrap_session_as_oauth_token(client, token_set, scopes or refresh_token.scopes) @@ -584,13 +594,21 @@ async def _exchange_oauth_for_session(self, oauth_access_token: str) -> TokenSet """ Exchanges a league OAuth access token (``claudai projectless`` scope) for a whole-stack Keboola programmatic session via ``manage/internal/auth-bridge/exchange-oauth-token``. + + Raised as ``TokenError`` (not ``HTTPException``): this runs inside ``exchange_authorization_code``, + invoked by the mcp SDK's own ``/token`` endpoint handler, which only recognizes ``TokenError`` + and formats it into a spec-compliant ``TokenErrorResponse`` body. An ``HTTPException`` here + would bubble up uncaught and reach the client as an opaque, non-OAuth-shaped error. """ kubernetes_token_path = deployed_sa_token_path() if not kubernetes_token_path: # OAuth login only runs on the deployed server; a missing SA token path means # KBC_KUBERNETES_TOKEN_PATH isn't set there, which is a deployment misconfiguration. LOG.error('[_exchange_oauth_for_session] KBC_KUBERNETES_TOKEN_PATH is not set; cannot exchange session.') - raise HTTPException(500, 'OAuth login is misconfigured: no Kubernetes ServiceAccount token available.') + raise TokenError( + error='invalid_request', + error_description='OAuth login is misconfigured: no Kubernetes ServiceAccount token available.', + ) exchanger = OAuthSessionExchanger( storage_api_url=self._storage_api_url, @@ -600,7 +618,7 @@ async def _exchange_oauth_for_session(self, oauth_access_token: str) -> TokenSet body = await exchanger.exchange(oauth_access_token=oauth_access_token) except OAuthTokenExchangeError as e: LOG.error(f'[_exchange_oauth_for_session] {e}') - raise HTTPException(e.status_code, str(e)) from e + raise TokenError(error='invalid_grant', error_description=str(e)) from e _log_debug(f'[_exchange_oauth_for_session] exchange response: {body}') return parse_token_response(body) diff --git a/tests/test_oauth.py b/tests/test_oauth.py index 6812608be..c858de107 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -6,7 +6,7 @@ import httpx import pytest -from mcp.server.auth.provider import AccessToken, AuthorizationParams, RefreshToken +from mcp.server.auth.provider import AccessToken, AuthorizationParams, RefreshToken, TokenError from mcp.shared.auth import InvalidRedirectUriError, OAuthClientInformationFull from pydantic import AnyHttpUrl, AnyUrl @@ -289,8 +289,6 @@ async def exchange(self, *, oauth_access_token: str): async def test_exchange_authorization_code_maps_exchange_error( self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch ): - from http.client import HTTPException - from keboola_mcp_server import oauth as oauth_module monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: '/tmp/sa-token') @@ -307,16 +305,16 @@ async def exchange(self, *, oauth_access_token: str): client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) - with pytest.raises(HTTPException) as exc: + # Raised as TokenError (not HTTPException): the mcp SDK's /token handler only recognizes + # TokenError and turns it into a spec-compliant TokenErrorResponse body. + with pytest.raises(TokenError) as exc: await oauth_provider.exchange_authorization_code(client, auth_code) - assert exc.value.args[0] == int(HTTPStatus.FORBIDDEN) + assert exc.value.error == 'invalid_grant' @pytest.mark.asyncio async def test_exchange_authorization_code_missing_sa_token_path( self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch ): - from http.client import HTTPException - from keboola_mcp_server import oauth as oauth_module monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: None) @@ -324,9 +322,9 @@ async def test_exchange_authorization_code_missing_sa_token_path( client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) - with pytest.raises(HTTPException) as exc: + with pytest.raises(TokenError) as exc: await oauth_provider.exchange_authorization_code(client, auth_code) - assert exc.value.args[0] == 500 + assert exc.value.error == 'invalid_request' @pytest.mark.asyncio async def test_exchange_refresh_token_calls_refresh_tokens_directly( @@ -371,11 +369,9 @@ async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, tran assert loaded_refresh.kbc_refresh_token == 'kbc_rt_rotated' @pytest.mark.asyncio - async def test_exchange_refresh_token_maps_network_error_to_http_exception( + async def test_exchange_refresh_token_maps_network_error_to_token_error( self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch ): - from http.client import HTTPException - from keboola_mcp_server import oauth as oauth_module async def _failing_refresh_tokens(storage_api_url: str, *, refresh_token: str, transport=None): @@ -392,8 +388,8 @@ async def _failing_refresh_tokens(storage_api_url: str, *, refresh_token: str, t kbc_refresh_token='kbc_rt_old', ) - # A network failure talking to Connection must surface as a clean HTTPException, not - # propagate as a raw httpx error (which the caller has no reason to expect/handle). - with pytest.raises(HTTPException) as exc: + # A network failure talking to Connection must surface as a clean TokenError, not + # propagate as a raw httpx error (which the mcp SDK's /token handler can't format). + with pytest.raises(TokenError) as exc: await oauth_provider.exchange_refresh_token(client, refresh_token, []) - assert exc.value.args[0] == 502 + assert exc.value.error == 'invalid_grant' diff --git a/tests/test_server.py b/tests/test_server.py index bea20d81f..9e249745b 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -14,6 +14,7 @@ from fastmcp.tools import FunctionTool from mcp.types import TextContent from pydantic import Field +from starlette.exceptions import HTTPException from starlette.requests import Request from keboola_mcp_server import cli @@ -26,7 +27,7 @@ toon_serializer, toon_serializer_compact, ) -from keboola_mcp_server.server import create_server +from keboola_mcp_server.server import CustomRoutes, create_server from keboola_mcp_server.tools.components.tools import COMPONENT_TOOLS_TAG from keboola_mcp_server.tools.constants import CONFIG_DIFF_PREVIEW_TAG from keboola_mcp_server.tools.data_apps import DATA_APP_TOOLS_TAG @@ -599,3 +600,20 @@ async def read_stream(stream, lines_list): missing_top_names = {'fastmcp', 'keboola_mcp_server', 'uvicorn'} - top_names assert not missing_top_names, f'Missing logger names: {missing_top_names}' + + +@pytest.mark.asyncio +async def test_oauth_callback_handler_propagates_http_exception(mocker) -> None: + # handle_oauth_callback() raises starlette.exceptions.HTTPException; oauth_callback_handler must + # re-raise it as-is (so Starlette renders the real status/detail) rather than falling through to + # the generic except-Exception branch, which would mask it as an opaque 500. + server_state = ServerState(config=Config(), runtime_info=ServerRuntimeInfo(transport='streamable-http')) + oauth_provider = mocker.Mock() + oauth_provider.handle_oauth_callback = mocker.AsyncMock(side_effect=HTTPException(400, 'Invalid state parameter')) + routes = CustomRoutes(server_state=server_state, oauth_provider=oauth_provider) + + request = Request({'type': 'http', 'headers': [], 'query_string': b'code=abc&state=xyz'}) + with pytest.raises(HTTPException) as exc: + await routes.oauth_callback_handler(request) + assert exc.value.status_code == 400 + assert exc.value.detail == 'Invalid state parameter' From 7fbeb25b8bf5a851113c380e336d5527f09a47ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 23 Jul 2026 15:35:24 +0200 Subject: [PATCH 52/89] fix(PSGO-261): stop sending X-KBC-ManageApiToken to exchange-oauth-token Confirmed against a real stack (401 Unauthorized) and against Connection's source: X-KBC-ManageApiToken (ManageTokenAuthenticator, a real Manage-token lookup) and X-Kubernetes-Authorization (KubernetesAuthenticator, synthetic token) are mutually exclusive authenticators for this endpoint. AuthBridgeOAuthExchangeProcessor explicitly rejects a non-synthetic token, so sending both headers -- even with a merely-reused SA JWT value for the Manage header -- causes a hard 401. Only X-Kubernetes-Authorization is sent now, exactly like the sibling resolve-storage-token endpoint. Co-Authored-By: Claude Sonnet 5 --- feature_spec/oauth_session_exchange/RFC.md | 2 +- src/keboola_mcp_server/clients/auth_bridge.py | 8 ++++---- tests/clients/test_auth_bridge.py | 4 +++- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/feature_spec/oauth_session_exchange/RFC.md b/feature_spec/oauth_session_exchange/RFC.md index c1db19df9..c961b9ab9 100644 --- a/feature_spec/oauth_session_exchange/RFC.md +++ b/feature_spec/oauth_session_exchange/RFC.md @@ -98,6 +98,6 @@ The MCP server's public/remote OAuth login (`SimpleOAuthProvider`, `oauth.py`) c 1. **Scope requested at `/oauth/consent` is `claudai projectless`** (space-separated, standard OAuth2 multi-scope) — `claudai` satisfies the exchange endpoint's `MissingClaudaiScopeException` guard; `projectless` is what makes the league token's `user_identifier` claim `admin:{id}` (not project-bound), which is what makes the *exchanged* session whole-stack. **Verify the exact literal string via the local test** — inferred from Connection's E2E test fixture comments, not from an explicit request example. 2. **Projectless = whole-stack, confirmed.** Both you and Connection's own E2E test agree: a `projectless`-scoped exchange yields an unrestricted session equivalent to a PKCE `login` lease — starts unconfirmed/whole-stack, `get_accessible_projects`/`set_project_scope` apply exactly as they do for a directly-supplied PAT today. -3. **`X-Subject-Token` confirmed** as the header name (shared constant with `resolve-storage-token`, both defined as `SUBJECT_TOKEN_HEADER = 'X-Subject-Token'` in Connection's source). Connection's E2E test also sends `X-KBC-ManageApiToken` alongside `X-Kubernetes-Authorization` — our current `StorageTokenResolver.resolve()` does **not** send that header and works today against the sibling `resolve-storage-token` endpoint, so it may be satisfied automatically for a validated k8s caller. **Verify empirically in the local test**: if the exchange 401s/403s without it, add `X-KBC-ManageApiToken` (same JWT value) to the request. +3. **`X-Subject-Token` confirmed** as the header name (shared constant with `resolve-storage-token`, both defined as `SUBJECT_TOKEN_HEADER = 'X-Subject-Token'` in Connection's source). **`X-KBC-ManageApiToken` confirmed NOT sent, verified against a real stack.** It's a separate, mutually-exclusive authenticator (`ManageTokenAuthenticator`, a real Manage-token lookup) from `X-Kubernetes-Authorization` (`KubernetesAuthenticator`, synthetic-token path) — sending both caused a live 401, because `AuthBridgeOAuthExchangeProcessor` explicitly rejects a non-synthetic (i.e. not Kubernetes-authenticated) token even if it also happens to carry the right Manage scope. Only `X-Kubernetes-Authorization` is sent, exactly like the sibling `resolve-storage-token` endpoint. 4. **No dual refresh — confirmed, not assumed.** `TokenRefreshProcessor.php` (Connection) operates on `ProgrammaticSession`/`ProgrammaticSessionRepository` — the same entity and `/v1/auth/token/refresh` mechanism already used by PAT/PKCE-login sessions, fully independent of the league OAuth session. The exchanged session refreshes on its own, forever, via the existing `refresh_tokens()`; the league OAuth token/refresh-token pair is used exactly once (at initial exchange) and never touched again, including on refresh. This **simplifies** `exchange_refresh_token()` relative to today's implementation (which currently re-negotiates with Connection's OAuth server on every refresh) rather than adding a second refresh call. 5. **No rollout coordination.** `/oauth/authorize` removal is immediate with no old-client fallback and no deploy-order dependency communicated from Connection's side. Verify locally against a real stack before shipping; no special deploy sequencing planned. diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index 05ceaeb3e..f31ef2c6a 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -171,15 +171,15 @@ async def exchange(self, *, oauth_access_token: str) -> dict: :return: The raw response body (``accessToken``/``refreshToken``/``expiresIn``/``sessionId``). :raises OAuthTokenExchangeError: On any exchange failure (status carried on the error). """ - # Connection's E2E test for this endpoint sends X-KBC-ManageApiToken alongside - # X-Kubernetes-Authorization; send both since the sibling resolve-storage-token - # endpoint only needs the latter (unconfirmed whether this one needs both too). + # X-KBC-ManageApiToken is a DIFFERENT, mutually-exclusive authenticator (a real Manage + # token lookup) -- confirmed against Connection's source that it must never be sent + # alongside X-Kubernetes-Authorization; the k8s JWT alone authorizes this endpoint, + # exactly like the sibling resolve-storage-token endpoint. sa_jwt = self._read_sa_jwt() headers = { 'Content-Type': 'application/json', 'Accept': 'application/json', 'X-Kubernetes-Authorization': f'Bearer {sa_jwt}', - 'X-KBC-ManageApiToken': sa_jwt, 'X-Subject-Token': f'Bearer {strip_bearer(oauth_access_token)}', } try: diff --git a/tests/clients/test_auth_bridge.py b/tests/clients/test_auth_bridge.py index b8eadc1b5..ba831c750 100644 --- a/tests/clients/test_auth_bridge.py +++ b/tests/clients/test_auth_bridge.py @@ -151,7 +151,9 @@ def handler(request: httpx.Request) -> httpx.Response: rq = captured['request'] assert rq.url.path == '/manage/internal/auth-bridge/exchange-oauth-token' assert rq.headers['X-Kubernetes-Authorization'] == 'Bearer sa-jwt-value' - assert rq.headers['X-KBC-ManageApiToken'] == 'sa-jwt-value' + # X-KBC-ManageApiToken is a distinct, mutually-exclusive authenticator -- must never be sent + # alongside X-Kubernetes-Authorization (confirmed against Connection's source). + assert 'X-KBC-ManageApiToken' not in rq.headers # Subject token is normalized to a single Bearer scheme regardless of inbound form. assert rq.headers['X-Subject-Token'] == 'Bearer league-oauth-token' From 660aab91554ee06b3dedddb457a26c008306e540 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 28 Jul 2026 08:29:11 +0200 Subject: [PATCH 53/89] fix(PSGO-261): never consult the local PKCE store on the deployed server _parent_subject_token() tried the local PKCE credential store (get_access_token(), which reads/writes ~/.keboola/mcp/credentials.json) before falling back to the request's own bearer token. That's correct for a single-user local stdio session, but the deployed server is multi-tenant: it holds no legitimate local session, and get_access_token() calls save_tokens() on every refresh -- writing one tenant's session to a file a different concurrent request could then read back instead of its own bearer token. Confirmed live: a deployed OAuth session's get_accessible_projects intermittently failed tokens/verify with a token that wasn't the one actually presented. Now gated on deployed_sa_token_path(): the deployed server always uses the current request's own bearer token; only local sessions try the PKCE store. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/tools/project.py | 14 ++++++++++---- tests/tools/test_project.py | 21 +++++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 41cfc53c7..a6e3d140a 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -12,7 +12,7 @@ from keboola_mcp_server.clients.auth_bridge import is_programmatic_token, strip_bearer from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient -from keboola_mcp_server.config import MetadataField +from keboola_mcp_server.config import MetadataField, deployed_sa_token_path from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager from keboola_mcp_server.mcp import SCOPE_KEY, MultiProjectMiddleware, ServerState, SessionScope, process_concurrently @@ -70,15 +70,21 @@ async def _parent_subject_token(client: KeboolaClient) -> str: """ Resolves the whole-stack (parent) programmatic token used to introspect/scope. - Prefers the refreshable token from the local PKCE credential store (so re-scoping always starts - from the parent, never from an already-narrowed scoped token); falls back to whatever bearer the - client currently carries (a directly-supplied PAT, or an HTTP bearer). + On a local (non-deployed) session, prefers the refreshable token from the local PKCE + credential store (so re-scoping always starts from the parent, never from an already-narrowed + scoped token). On the deployed server the local store is never consulted: it holds no session + for the current request's caller, and -- since it's shared across every concurrent session on + the pod -- reading (or refreshing-and-writing) it here would risk leaking one tenant's session + into another's request. Falls back to whatever bearer the client currently carries (a + directly-supplied PAT, an HTTP bearer, or an OAuth-exchanged session token). """ if not is_programmatic_token(client.bearer_token): raise ValueError( 'Project scoping requires a Keboola programmatic token (kbc_at_/kbc_pat_). ' 'Run "keboola-mcp-server login --api-url " first, or supply such a token.' ) + if deployed_sa_token_path(): + return strip_bearer(cast(str, client.bearer_token)) try: return await get_access_token(client.storage_api_url) except RuntimeError: diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 70cac9bc8..03bf4748c 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -13,6 +13,7 @@ from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, + _parent_subject_token, _resolve_branch_context, get_accessible_projects, get_project_info, @@ -271,6 +272,26 @@ def _prep_client(mcp_context_client: Context, mocker: MockerFixture, *, bearer: return client +@pytest.mark.asyncio +async def test_parent_subject_token_ignores_local_store_when_deployed(mocker: MockerFixture) -> None: + # On the deployed (multi-tenant) server, the local PKCE credential store must never be consulted + # -- it holds no session for this request's caller, and since it's shared across every concurrent + # request on the pod, reading (or refresh-writing) it here would risk leaking one tenant's session + # into another's. Only the request's own bearer token may be used. + mocker.patch('keboola_mcp_server.tools.project.deployed_sa_token_path', return_value='/var/run/secrets/token') + get_access_token = mocker.patch( + 'keboola_mcp_server.tools.project.get_access_token', + new=mocker.AsyncMock(return_value='kbc_at_wrong_tenant'), + ) + client = mocker.Mock() + client.bearer_token = 'Bearer kbc_at_this_request' + + token = await _parent_subject_token(client) + + assert token == 'kbc_at_this_request' + get_access_token.assert_not_called() + + @pytest.mark.asyncio async def test_get_accessible_projects(mcp_context_client: Context, mocker: MockerFixture) -> None: _prep_client(mcp_context_client, mocker) From eba47d9c882a4573274acf5cf875e8368bc61a87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 28 Jul 2026 08:51:21 +0200 Subject: [PATCH 54/89] fix(PSGO-261): exempt bootstrap tools from ToolsFilteringMiddleware's verify on_call_tool unconditionally called verify_token() (needs a single-project X-KBC-ProjectId context) before the tool's own body ran -- including for get_accessible_projects/set_project_scope, whose entire purpose is to establish that context in the first place. On a fresh, unscoped OAuth or PAT session this 401'd before introspect_token() was ever reached, contradicting on_list_tools' own comment that "the on_call_tool guards still enforce... per project (with the right project_id)" -- confirmed live via Datadog logs showing session-state creation followed directly by a verify call, no introspect at all. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/mcp.py | 9 +++++++++ tests/test_mcp.py | 22 ++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 4747763ef..06791eb84 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -835,6 +835,15 @@ async def on_call_tool( call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult], ) -> mt.CallToolResult: tool = await context.fastmcp_context.fastmcp.get_tool(context.message.name) + + # Bootstrap tools (get_accessible_projects/set_project_scope) must work before any project + # is chosen -- that's their entire purpose. verify_token() needs a single-project context + # (X-KBC-ProjectId); calling it here pre-scope would 401 before the tool's own body (which + # establishes that context, e.g. via introspect_token) ever runs. Mirrors the same exemption + # in on_list_tools and MultiProjectMiddleware. + if tool.name in _BOOTSTRAP_TOOLS: + return await call_next(context) + token_info = await self.get_token_info(context.fastmcp_context) has_semantic_models = False diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 3c93d63b0..8adbf4e5f 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -537,6 +537,28 @@ async def call_next(_): result = await middleware.on_call_tool(context, call_next) assert result is expected + @pytest.mark.asyncio + @pytest.mark.parametrize('tool_name', ['get_accessible_projects', 'set_project_scope']) + async def test_call_tool_bootstrap_tools_skip_verify( + self, mcp_context_client, keboola_client, tool_name: str + ) -> None: + # Bootstrap tools must work before any project is chosen (that's their purpose): calling + # verify_token() here -- which needs a single-project context (X-KBC-ProjectId) -- would 401 + # before the tool's own body (which establishes that context) ever runs. + keboola_client.storage_client.verify_token = AsyncMock(side_effect=AssertionError('verify must not run')) + + tool = _tool(tool_name) + mcp_context_client.fastmcp = SimpleNamespace(get_tool=AsyncMock(return_value=tool)) + context = SimpleNamespace(fastmcp_context=mcp_context_client, message=SimpleNamespace(name=tool_name)) + + expected = MagicMock() + + async def call_next(_): + return expected + + result = await ToolsFilteringMiddleware().on_call_tool(context, call_next) + assert result is expected + @pytest.mark.asyncio @pytest.mark.parametrize( 'tool_name', From 1c2bd3664e4a66fe6e99aa07302952fef041dd78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 28 Jul 2026 09:13:26 +0200 Subject: [PATCH 55/89] fix(PSGO-261): apply confirmed scope's project id on deployed sessions; fix step-up bearer _resolve_local_tokens applied a confirmed SessionScope's active project id to config only for local sessions; deployed sessions skipped it entirely, so every call after set_project_scope kept building the active client from the unscoped whole-stack token (no X-KBC-ProjectId) and 401'd, even though scoping itself succeeded. Now the active project id is threaded through for deployed sessions too; create_session_state's existing resolver-exchange path (keyed off project_id) then narrows the token correctly, unchanged. step_up_storage_client (used for best-effort tool-call event logging) built its client from the raw storage_api_token directly, bypassing the bearer_or_sapi_token selection done everywhere else -- so a programmatic (kbc_at_/kbc_pat_) session's token was sent as X-StorageAPI-Token, which Storage API rejects outright, silently dropping event telemetry for every programmatic session. Reuses the same bearer-preferring token now. Co-Authored-By: Claude Sonnet 5 --- src/keboola_mcp_server/clients/client.py | 7 +++++-- src/keboola_mcp_server/mcp.py | 10 +++++++++- tests/clients/test_client.py | 17 ++++++++++++++++ tests/test_mcp.py | 25 ++++++++++++++++++++++++ 4 files changed, 56 insertions(+), 3 deletions(-) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index 673c1a08a..337c8837f 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -175,7 +175,7 @@ def __init__( sync_actions_api_url = urlunparse(('https', f'sync-actions.{self._hostname_suffix}', '', '', '', '')) # Initialize clients for individual services - bearer_or_sapi_token = f'Bearer {bearer_token}' if bearer_token else self._token + bearer_or_sapi_token = self._bearer_or_sapi_token = f'Bearer {bearer_token}' if bearer_token else self._token # The encryption service does not require an authorization header, so we pass None as the token self._encryption_client = EncryptionClient.create( root_url=encryption_api_url, token=None, headers=self._headers @@ -317,7 +317,10 @@ def step_up_storage_client(self, kubernetes_token_path: str) -> 'AsyncStorageCli headers['X-Kubernetes-Authorization'] = f'Bearer {jwt}' return AsyncStorageClient.create( root_url=self._storage_api_url, - token=self._token, + # Bearer, not the raw storage_api_token: for a programmatic (kbc_at_/kbc_pat_) session + # the raw token would be sent as X-StorageAPI-Token, which Storage API rejects outright + # -- it only accepts a programmatic token via Authorization: Bearer. + token=self._bearer_or_sapi_token, branch_id=self._branch_id, headers=headers, readonly=self._storage_client.raw_client.readonly, diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 06791eb84..6d3154497 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -488,9 +488,17 @@ async def _resolve_local_tokens( the parent when it nears expiry. The default (auto-leased) multi-project scope carries no minted token and simply uses the parent token, narrowed per request by ``X-KBC-ProjectId``. On the deployed server (``KBC_KUBERNETES_TOKEN_PATH`` set) the per-request resolver exchange - already handles freshness, so this is a no-op there. + already handles token freshness/narrowing once ``project_id`` is known -- but it only runs + once ``project_id`` is known, and nothing else threads a confirmed scope's active project id + into ``config`` for a deployed session. Without that, ``create_session_state`` keeps building + the active client from the unscoped whole-stack token with no ``X-KBC-ProjectId``, so every + call after ``set_project_scope`` 401s even though scoping itself succeeded. So still apply + just the active project id here for deployed sessions; the token itself is left alone since + the resolver-exchange path (keyed off that project id) handles narrowing it correctly. """ if not cls._is_local_programmatic(config): + if scope and scope.project_ids and not config.project_id: + config = dataclasses.replace(config, project_id=str(scope.active_project_id)) return config, scope # Strip any inbound `Bearer ` scheme; introspect/exchange helpers add the scheme themselves, diff --git a/tests/clients/test_client.py b/tests/clients/test_client.py index 9c2284d29..06aa255aa 100644 --- a/tests/clients/test_client.py +++ b/tests/clients/test_client.py @@ -747,6 +747,23 @@ def test_attaches_step_up_header_and_keeps_user_token(self, tmp_path, own_stack_ # ... and pre-existing headers are preserved. assert headers['User-Agent'] == 'test' + def test_uses_bearer_for_programmatic_token(self, tmp_path): + # A programmatic (kbc_at_/kbc_pat_) session's token must ride as Authorization: Bearer, not + # X-StorageAPI-Token, which Storage API rejects outright for that token shape. + token_file = tmp_path / 'token' + token_file.write_text('sa-jwt') + client = KeboolaClient( + storage_api_url='https://connection.keboola.com', + storage_api_token='kbc_at_abc', + bearer_token='kbc_at_abc', + ) + + stepped = client.step_up_storage_client(str(token_file)) + + headers = stepped.raw_client.headers + assert headers['Authorization'] == 'Bearer kbc_at_abc' + assert 'X-StorageAPI-Token' not in headers + @pytest.mark.parametrize('readonly', [None, True, False]) def test_propagates_readonly_guard(self, tmp_path, readonly): token_file = tmp_path / 'token' diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8adbf4e5f..7f518fa69 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1013,6 +1013,31 @@ async def test_deployed_is_noop(self, monkeypatch) -> None: assert out_config is config assert out_scope is None + @pytest.mark.asyncio + async def test_deployed_with_confirmed_scope_applies_active_project_id(self, monkeypatch) -> None: + # Deployed sessions skip token refresh/re-minting (the resolver-exchange path in + # create_session_state handles that once project_id is known) -- but a confirmed scope's + # active project id must still be threaded through, or every call after set_project_scope + # keeps building the active client from the unscoped whole-stack token (PSGO-261 regression: + # get_accessible_projects worked, every subsequent scoped call 401'd). + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x') + scope = SessionScope(project_ids=[18], scoped_token='kbc_at_scoped', confirmed=True) + with patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(side_effect=AssertionError('no PKCE store'))): + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + assert out_config.project_id == '18' + assert out_config.storage_token == 'kbc_at_x' # untouched; resolver-exchange narrows it + assert out_scope is scope # untouched + + @pytest.mark.asyncio + async def test_deployed_no_scope_or_already_set_project_id_is_noop(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x', project_id='7') + scope = SessionScope(project_ids=[18], confirmed=True) + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + assert out_config is config # project_id already set -- not overwritten + assert out_scope is scope + @pytest.mark.asyncio async def test_legacy_token_is_noop(self, monkeypatch) -> None: monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) From b79e2988e6ed864d5240a3b2f9cf9d81af1a8474 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 28 Jul 2026 10:11:33 +0200 Subject: [PATCH 56/89] fix(PSGO-261): carry multi-project scope as a signed scope_token instead of session state Datadog trace evidence showed three separate POST /mcp/ requests sharing the same runtime-id (single process, not a replica issue): ctx.session.state is rebuilt empty on every request under this server's default stateless-HTTP transport, so a scope confirmed via set_project_scope never survived to the next tool call. The mcp 2026-07-28 RC formalizes this across the spec (drops Mcp-Session-Id/session pinning entirely), so the fix follows the same self-contained-JWT pattern this server already uses for OAuth tokens: set_project_scope/get_accessible_projects now return a signed scope_token that the caller resends as a tool argument on every subsequent call, unified across local stdio, local OAuth, and deployed OAuth sessions. Co-Authored-By: Claude Sonnet 5 --- TOOLS.md | 3 + src/keboola_mcp_server/jwt_utils.py | 24 ++++++ src/keboola_mcp_server/mcp.py | 107 +++++++++++++++++++++--- src/keboola_mcp_server/oauth.py | 17 +--- src/keboola_mcp_server/server.py | 18 ++-- src/keboola_mcp_server/tools/project.py | 43 ++++++++-- tests/test_mcp.py | 84 +++++++++++++++++++ tests/tools/test_project.py | 19 ++++- 8 files changed, 272 insertions(+), 43 deletions(-) create mode 100644 src/keboola_mcp_server/jwt_utils.py diff --git a/TOOLS.md b/TOOLS.md index 76f05d827..d3948de18 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -3629,6 +3629,9 @@ rest of the conversation. Read-only tools then run against every scoped project write operations target the active (first) project only. Call this when the user states which projects to work on; it can be called again any time to re-scope. +The server does not remember this scope between calls: pass the returned `scope_token` as the +`scope_token` argument on every subsequent tool call in this conversation to keep it in effect. + **Input JSON Schema**: ```json diff --git a/src/keboola_mcp_server/jwt_utils.py b/src/keboola_mcp_server/jwt_utils.py new file mode 100644 index 000000000..5562cc5ac --- /dev/null +++ b/src/keboola_mcp_server/jwt_utils.py @@ -0,0 +1,24 @@ +""" +Shared helpers for signing small JSON payloads into opaque, self-contained JWTs. + +Used wherever this server hands a client something to carry and resend later instead of keeping +it in server-side memory (OAuth state/access/refresh tokens in ``oauth.py``; the multi-project +``scope_token`` in ``mcp.py``) -- gzip-compressed JSON, HMAC-signed so any process holding the same +secret can verify it without a shared store. +""" + +import gzip +import json +from typing import Any, Mapping + +import jwt.api_jws + + +def encode_jwt(data: Mapping[str, Any], secret: str) -> str: + json_gzip = gzip.compress(json.dumps(data).encode('utf-8')) + return jwt.api_jws.encode(json_gzip, secret) + + +def decode_jwt(token: str, secret: str) -> dict[str, Any]: + json_gzip = jwt.api_jws.decode(token, secret, algorithms=['HS256']) + return json.loads(gzip.decompress(json_gzip).decode('utf-8')) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 6d3154497..34ff3970e 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -8,6 +8,7 @@ import asyncio import dataclasses import logging +import secrets import textwrap import time from collections.abc import Awaitable, Callable, Iterable @@ -37,6 +38,7 @@ from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config, ServerRuntimeInfo, deployed_sa_token_path, is_same_stack +from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt from keboola_mcp_server.oauth import ProxyAccessToken from keboola_mcp_server.tools.constants import MODIFY_FLOW_TOOL_NAME, SEMANTIC_TOOLS_TAG, UPDATE_FLOW_TOOL_NAME from keboola_mcp_server.workspace import WorkspaceManager @@ -62,6 +64,27 @@ # subset of the scoped projects (consumed and stripped by MultiProjectMiddleware.on_call_tool). _PROJECT_FILTER_ARG = 'project_ids' +# Per-call argument that carries the confirmed multi-project scope forward (consumed and stripped +# by SessionStateMiddleware.on_request). See SessionScope.to_token/from_token: under the server's +# default stateless-HTTP transport a fresh, empty session is built for every request (the mcp +# 2026-07-28 RC formalizes this across the spec, dropping Mcp-Session-Id/session pinning entirely), +# so nothing survives in ctx.session.state between one tool call and the next -- on one replica or +# many, even within a single process. A scope set via "set_project_scope" only persists if the +# caller resends the token it returned. +_SCOPE_TOKEN_ARG = 'scope_token' + +# Process-local fallback signing key for scope_token, used when no shared KBC_JWT_SECRET is +# configured (e.g. local stdio/login sessions). A per-process secret is enough there since a stdio +# process serves exactly one conversation end-to-end; deployed multi-replica setups already require +# a shared jwt_secret for the OAuth-provider JWTs (see oauth.py), which this reuses. +_FALLBACK_SCOPE_SECRET = secrets.token_hex(32) + + +def resolve_scope_secret(config: Config) -> str: + """The HMAC key used to sign/verify ``scope_token`` -- shared across replicas when + ``config.jwt_secret`` (``KBC_JWT_SECRET``) is configured, otherwise a process-local fallback.""" + return config.jwt_secret or _FALLBACK_SCOPE_SECRET + @dataclasses.dataclass(frozen=True) class SessionScope: @@ -90,6 +113,16 @@ def is_near_expiry(self) -> bool: return False return time.time() >= (self.scoped_expires_at - 60) + def to_token(self, secret: str) -> str: + """Signs this scope into the opaque ``scope_token`` a caller resends on later calls.""" + return encode_jwt(dataclasses.asdict(self), secret) + + @classmethod + def from_token(cls, token: str, secret: str) -> 'SessionScope': + """Inverse of ``to_token``. Raises on a missing/invalid/tampered token -- callers should + treat any exception as "no scope" rather than fail the request.""" + return cls(**decode_jwt(token, secret)) + R = TypeVar('R') T = TypeVar('T') @@ -293,10 +326,12 @@ async def on_request( # server (KBC_KUBERNETES_TOKEN_PATH set). config = await self._maybe_use_stored_session(config, refresh=not is_list) - # In-conversation multi-project scope persists on the session across this per-request - # state rebuild. With no scope and no preset project, auto-lease ALL accessible projects - # (multi-project mode) so read tools fan out across everything — but never on /list. - scope = self._read_persisted_scope(ctx.session) + # In-conversation multi-project scope is carried by the caller as the `scope_token` tool + # argument (see SessionScope.to_token/from_token) rather than read back from + # ctx.session.state, which is rebuilt empty on every request under this server's default + # stateless-HTTP transport. With no scope and no preset project, auto-lease ALL accessible + # projects (multi-project mode) so read tools fan out across everything — but never on /list. + scope = self._read_scope_from_request(context, config) if scope is None and not config.project_id and not is_list: scope = await self._autolease_default_scope(config) if not is_list: @@ -328,6 +363,40 @@ async def on_request( # ctx.session.state = {} pass + async def on_list_tools( + self, + context: fmw.MiddlewareContext[mt.ListToolsRequest], + call_next: fmw.CallNext[mt.ListToolsRequest, list[Tool]], + ) -> list[Tool]: + """Advertises the optional `scope_token` argument on every tool. + + Unconditional (unlike MultiProjectMiddleware's `_PROJECT_FILTER_ARG` patch, which is gated on + an active multi-project scope): a `tools/list` request cannot itself carry `scope_token`, so + whether a scope is currently confirmed can't be known while building this response. Showing + the parameter always costs nothing when unused and is what lets the caller learn about it + before ever calling `set_project_scope`. + """ + tools = await call_next(context) + patched: list[Tool] = [] + for tool in tools: + params = dict(tool.parameters or {}) + props = dict(params.get('properties') or {}) + if _SCOPE_TOKEN_ARG in props: + patched.append(tool) + continue + props[_SCOPE_TOKEN_ARG] = { + 'type': 'string', + 'description': ( + 'Opaque token returned by "set_project_scope" (also echoed by ' + '"get_accessible_projects" once a scope is confirmed). The server does not ' + 'remember the scope between calls -- resend this value on every tool call in ' + 'this conversation once you have it.' + ), + } + params['properties'] = props + patched.append(tool.model_copy(update={'parameters': params})) + return patched + @classmethod def _get_headers(cls, runtime_info: ServerRuntimeInfo) -> dict[str, Any]: """ @@ -393,15 +462,27 @@ def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_sto return config - @staticmethod - def _read_persisted_scope(session: Any) -> 'SessionScope | None': - """Reads the multi-project scope stashed in the prior request's session state, if any.""" - prior = getattr(session, 'state', None) - if isinstance(prior, dict): - scope = prior.get(SCOPE_KEY) - if isinstance(scope, SessionScope): - return scope - return None + @classmethod + def _read_scope_from_request(cls, context: fmw.MiddlewareContext[Any], config: Config) -> 'SessionScope | None': + """Decodes the ``scope_token`` tool-call argument (if any) back into a SessionScope. + + Pops the argument so it never reaches the tool function, matching how + MultiProjectMiddleware.on_call_tool consumes _PROJECT_FILTER_ARG. Absent, malformed, or + expired tokens are treated as "no scope yet" rather than an error -- the ask-first gate in + MultiProjectMiddleware then steers the caller back through get_accessible_projects / + set_project_scope. + """ + args = getattr(getattr(context, 'message', None), 'arguments', None) + if not isinstance(args, dict): + return None + token = args.pop(_SCOPE_TOKEN_ARG, None) + if not token: + return None + try: + return SessionScope.from_token(token, resolve_scope_secret(config)) + except Exception: + LOG.warning('Ignoring invalid or expired scope_token.', exc_info=True) + return None @classmethod def _is_local_programmatic(cls, config: Config) -> bool: diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index d261868d3..9a431d4ac 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -1,5 +1,3 @@ -import gzip -import json import logging import math import os @@ -11,7 +9,7 @@ from urllib.parse import urljoin import httpx -import jwt.api_jws +import jwt from fastmcp.server.auth.auth import OAuthProvider from mcp.server.auth.provider import ( AccessToken, @@ -29,6 +27,7 @@ from keboola_mcp_server.auth_login import TokenSet, parse_token_response, refresh_tokens from keboola_mcp_server.clients.auth_bridge import OAuthSessionExchanger, OAuthTokenExchangeError from keboola_mcp_server.config import deployed_sa_token_path +from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt LOG = logging.getLogger(__name__) _OAUTH_LOG_ALL = bool(os.getenv('KEBOOLA_MCP_SERVER_OAUTH_LOG_ALL')) @@ -632,15 +631,7 @@ def _create_http_client(): return httpx.AsyncClient(follow_redirects=True, timeout=httpx.Timeout(30.0)) def _encode(self, data: Mapping[str, Any], *, key: str | None = None) -> str: - json_str = json.dumps(data) - json_bytes = json_str.encode('utf-8') - json_gzip = gzip.compress(json_bytes) - json_encrypted = jwt.api_jws.encode(json_gzip, key or self._jwt_secret) - return json_encrypted + return encode_jwt(data, key or self._jwt_secret) def _decode(self, data: str, *, key: str | None = None) -> dict[str, Any]: - json_gzip = jwt.api_jws.decode(data, key or self._jwt_secret, algorithms=['HS256']) - json_bytes = gzip.decompress(json_gzip) - json_str = json_bytes.decode('utf-8') - json_obj = json.loads(json_str) - return json_obj + return decode_jwt(data, key or self._jwt_secret) diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index de9bb8276..24a5c16f0 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -236,14 +236,16 @@ def create_server( 'else: call "get_accessible_projects", show the user their projects, and ASK whether to work ' 'across ALL of them or a subset. Do not decide for them. Then call "set_project_scope" with ' 'their answer (no arguments = all projects, or the chosen project ids, optionally ' - 'read_only=true). After that, read-only tools return results per project. Never write to ' - 'more than one project without explicit user confirmation — write operations target the ' - 'active (first-scoped) project only. If instead the session uses a legacy project-scoped ' - 'Storage API token, it is already bound to a single project: use the tools directly — ' - '"get_accessible_projects" / "set_project_scope" do not apply (they will report that no ' - 'programmatic token is present). Note: outside the Storage API, some tools may need ' - 'per-project token support not yet available on every stack; surface such errors plainly ' - 'rather than retrying.' + 'read_only=true). Both tools return a "scope_token" -- the server does not remember the ' + 'scope between calls, so resend that value as the "scope_token" argument on every ' + 'subsequent tool call in this conversation. After that, read-only tools return results per ' + 'project. Never write to more than one project without explicit user confirmation — write ' + 'operations target the active (first-scoped) project only. If instead the session uses a ' + 'legacy project-scoped Storage API token, it is already bound to a single project: use the ' + 'tools directly — "get_accessible_projects" / "set_project_scope" do not apply (they will ' + 'report that no programmatic token is present). Note: outside the Storage API, some tools ' + 'may need per-project token support not yet available on every stack; surface such errors ' + 'plainly rather than retrying.' ), lifespan=create_keboola_lifespan(server_state), auth=oauth_provider, diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index a6e3d140a..512a4ab52 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -15,7 +15,14 @@ from keboola_mcp_server.config import MetadataField, deployed_sa_token_path from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager -from keboola_mcp_server.mcp import SCOPE_KEY, MultiProjectMiddleware, ServerState, SessionScope, process_concurrently +from keboola_mcp_server.mcp import ( + SCOPE_KEY, + MultiProjectMiddleware, + ServerState, + SessionScope, + process_concurrently, + resolve_scope_secret, +) from keboola_mcp_server.resources.prompts import get_project_system_prompt from keboola_mcp_server.workspace import WorkspaceManager @@ -310,6 +317,14 @@ class AccessibleProjects(BaseModel): description='The projects the session is currently scoped to, or null if no scope has been confirmed yet.', ) read_only: bool | None = Field(default=None, description='Whether the current scoped token is read-only.') + scope_token: str | None = Field( + default=None, + description=( + 'Opaque token encoding the confirmed scope, or null if none is confirmed yet. The server ' + 'does not remember the scope between calls -- pass this value as the "scope_token" ' + 'argument on every subsequent tool call in this conversation.' + ), + ) base_instructions: list[BaseInstructionGroup] | None = Field( default=None, description=( @@ -326,6 +341,12 @@ class AccessibleProjects(BaseModel): class ProjectScope(BaseModel): project_ids: list[int] = Field(description='The projects the session is now scoped to.') read_only: bool = Field(description='Whether the scoped token is read-only.') + scope_token: str = Field( + description=( + 'Opaque token encoding this scope. The server does not remember it between calls -- pass ' + 'this value as the "scope_token" argument on every subsequent tool call in this conversation.' + ), + ) llm_instruction: str = Field(description='Guidance for the assistant on the new scope.') @@ -428,14 +449,15 @@ async def get_accessible_projects( ) else: instruction = ( - f'Session is currently scoped to {len(scoped_ids)} project(s). ' - 'Call "set_project_scope" to change the scope.' + f'Session is currently scoped to {len(scoped_ids)} project(s). Resend "scope_token" on every ' + 'subsequent tool call to keep it in effect; call "set_project_scope" to change the scope.' ) return AccessibleProjects( user_email=introspection.user_email, projects=projects, scoped_project_ids=scoped_ids, read_only=scope.read_only if scoped_ids is not None else None, + scope_token=scope.to_token(resolve_scope_secret(server_state.config)) if scoped_ids is not None else None, base_instructions=base_instructions, llm_instruction=instruction, ) @@ -463,6 +485,9 @@ async def set_project_scope( rest of the conversation. Read-only tools then run against every scoped project in a single call; write operations target the active (first) project only. Call this when the user states which projects to work on; it can be called again any time to re-scope. + + The server does not remember this scope between calls: pass the returned `scope_token` as the + `scope_token` argument on every subsequent tool call in this conversation to keep it in effect. """ client = KeboolaClient.from_state(ctx.session.state) parent_token = await _parent_subject_token(client) @@ -517,16 +542,24 @@ async def set_project_scope( LOG.debug(f'Could not send tools/list_changed after scoping: {e}') multi = len(ids) > 1 + scope_token = scope.to_token(resolve_scope_secret(ServerState.from_context(ctx).config)) return ProjectScope( project_ids=ids, read_only=scope.read_only, + scope_token=scope_token, llm_instruction=( ( f'Session scoped to {len(ids)} projects. Read-only tools return results per project. ' 'Write operations are not fanned out — they target the first scoped project; to write ' - 'elsewhere, re-scope to that project first (confirm with the user).' + 'elsewhere, re-scope to that project first (confirm with the user). The server does not ' + 'remember this scope between calls -- pass "scope_token" as an argument on every ' + 'subsequent tool call in this conversation.' ) if multi - else f'Session scoped to project {ids[0]}.' + else ( + f'Session scoped to project {ids[0]}. The server does not remember this scope between ' + 'calls -- pass "scope_token" as an argument on every subsequent tool call in this ' + 'conversation.' + ) ), ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 7f518fa69..6480d4484 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -26,6 +26,7 @@ _exclude_none_serializer, _filter_toon_nulls, process_concurrently, + resolve_scope_secret, toon_serializer, unwrap_results, ) @@ -1141,6 +1142,89 @@ async def test_autolease_noop_when_deployed(self, monkeypatch) -> None: assert await SessionStateMiddleware._autolease_default_scope(config) is None +class TestScopeToken: + """The multi-project scope is carried by the caller as the `scope_token` tool argument, not read + back from ctx.session.state -- which is rebuilt empty on every request under this server's + default stateless-HTTP transport, so nothing survives there between one tool call and the next. + """ + + def test_round_trip(self) -> None: + scope = SessionScope( + project_ids=[11, 22], read_only=True, scoped_token='kbc_at_s', scoped_expires_at=1234.0, confirmed=True + ) + token = scope.to_token('secret') + assert SessionScope.from_token(token, 'secret') == scope + + def test_wrong_secret_rejected(self) -> None: + token = SessionScope(project_ids=[11], confirmed=True).to_token('secret-a') + with pytest.raises(Exception, match='.+'): + SessionScope.from_token(token, 'secret-b') + + def test_resolve_scope_secret_prefers_configured_jwt_secret(self) -> None: + assert resolve_scope_secret(Config(jwt_secret='shared-secret')) == 'shared-secret' + + def test_resolve_scope_secret_fallback_is_stable_within_process(self) -> None: + config = Config() + assert resolve_scope_secret(config) == resolve_scope_secret(config) + + @staticmethod + def _call_tool_context(arguments: dict) -> SimpleNamespace: + message = SimpleNamespace(name='get_tables', arguments=arguments) + return SimpleNamespace(message=message, method='tools/call') + + def test_read_scope_from_request_decodes_and_pops_token(self) -> None: + config = Config(jwt_secret='shared-secret') + scope = SessionScope(project_ids=[11, 22], confirmed=True) + arguments = {'scope_token': scope.to_token('shared-secret'), 'other_arg': 1} + + context = self._call_tool_context(arguments) + result = SessionStateMiddleware._read_scope_from_request(context, config) + + assert result == scope + # Popped so the underlying tool function never sees it as an unexpected argument. + assert 'scope_token' not in arguments + assert arguments == {'other_arg': 1} + + @pytest.mark.parametrize( + 'arguments', + [{}, {'scope_token': None}, {'scope_token': ''}, {'scope_token': 'not-a-valid-jwt'}], + ids=['missing', 'none', 'empty', 'malformed'], + ) + def test_read_scope_from_request_returns_none_when_absent_or_invalid(self, arguments: dict) -> None: + config = Config(jwt_secret='shared-secret') + context = self._call_tool_context(dict(arguments)) + assert SessionStateMiddleware._read_scope_from_request(context, config) is None + + def test_read_scope_from_request_ignores_non_call_tool_requests(self) -> None: + # tools/list (and other non-call requests) carry no `.arguments` at all. + context = SimpleNamespace(method='tools/list', fastmcp_context=None) + assert SessionStateMiddleware._read_scope_from_request(context, Config()) is None + + def test_wrong_secret_falls_back_to_no_scope_via_read_scope_from_request(self) -> None: + # A token minted with a different secret (e.g. a replica whose fallback secret differs) must + # degrade to "no scope" rather than raise -- the ask-first gate then re-prompts the caller. + token = SessionScope(project_ids=[11], confirmed=True).to_token('secret-a') + context = self._call_tool_context({'scope_token': token}) + assert SessionStateMiddleware._read_scope_from_request(context, Config(jwt_secret='secret-b')) is None + + @pytest.mark.asyncio + async def test_on_list_tools_advertises_scope_token_unconditionally(self) -> None: + # Unlike MultiProjectMiddleware's `project_ids` filter, this must show up even with no scope + # confirmed yet (indeed, even before get_accessible_projects has ever been called) -- a + # tools/list request can't itself carry a scope_token, so scope state can't gate this. + tool = _tool('get_tables', read_only=True) + tool.parameters = {'type': 'object', 'properties': {}} + tool.model_copy = lambda update, _t=tool: SimpleNamespace(name=_t.name, parameters=update['parameters']) + context = SimpleNamespace(method='tools/list') + + async def call_next(_): + return [tool] + + tools = await SessionStateMiddleware().on_list_tools(context, call_next) + + assert 'scope_token' in tools[0].parameters['properties'] + + class TestMultiProjectMiddleware: """Read tools fan out across the scoped projects; writes and single-project scope do not.""" diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 03bf4748c..4750165f0 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -7,9 +7,9 @@ from pytest_mock import MockerFixture from keboola_mcp_server.clients.client import KeboolaClient -from keboola_mcp_server.config import MetadataField +from keboola_mcp_server.config import Config, MetadataField from keboola_mcp_server.links import Link -from keboola_mcp_server.mcp import SCOPE_KEY, SessionScope +from keboola_mcp_server.mcp import SCOPE_KEY, SessionScope, resolve_scope_secret from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, @@ -303,7 +303,9 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock 'keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection) ) # Per-project SQL dialect is resolved via a token verify narrowed by X-KBC-ProjectId; mock that. - mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) + mocker.patch( + 'keboola_mcp_server.tools.project.ServerState.from_context', return_value=SimpleNamespace(config=Config()) + ) dialects = {18: 'BigQuery', 83: 'Snowflake'} mocker.patch( 'keboola_mcp_server.tools.project._project_sql_dialect', @@ -322,14 +324,20 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock assert result.scoped_project_ids is None assert result.read_only is None assert result.base_instructions is None # not requested + assert result.scope_token is None assert all(not p.in_scope for p in result.projects) - # Once scoped, the current scope is surfaced on the projects and at the top level. + # Once scoped, the current scope is surfaced on the projects and at the top level, and echoed + # back as a scope_token the caller must resend on later calls (the server does not remember it). mcp_context_client.session.state[SCOPE_KEY] = SessionScope(project_ids=[83], read_only=True, confirmed=True) result = await get_accessible_projects(mcp_context_client) assert result.scoped_project_ids == [83] assert result.read_only is True assert [(p.id, p.in_scope) for p in result.projects] == [(18, False), (83, True)] + assert result.scope_token is not None + assert SessionScope.from_token(result.scope_token, resolve_scope_secret(Config())) == SessionScope( + project_ids=[83], read_only=True, confirmed=True + ) @pytest.mark.asyncio @@ -429,6 +437,9 @@ async def test_set_project_scope_subset_exchanges_and_stores( scope = mcp_context_client.session.state[SCOPE_KEY] assert scope.scoped_token == 'kbc_at_scoped' assert scope.project_ids == [18, 83] + # The server does not remember this scope between calls; the caller must resend scope_token. + assert result.scope_token is not None + assert SessionScope.from_token(result.scope_token, resolve_scope_secret(Config())) == scope @pytest.mark.asyncio From 78a0bf6efd89065ae9c115b9a704cace05ad3a6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 28 Jul 2026 10:34:23 +0200 Subject: [PATCH 57/89] fix(PSGO-261): drop the auth-bridge resolve-storage-token exchange for programmatic tokens Live testing against a real dev stack hit a 403 on resolve-storage-token once the scope_token fix correctly started populating project_id for deployed OAuth sessions -- a separate Manage scope from exchange-oauth-token's, never granted on this stack since this code path had never actually fired before. Rather than chase another infra scope grant, drop the resolver call: KeboolaClient already forwards a Bearer token to every service it wraps (Storage, Queue, AI, Data Science, Scheduler, Sync Actions, Metastore), so a programmatic token (kbc_at_/kbc_pat_) -- OAuth-exchanged or directly supplied -- is now always sent as Authorization: Bearer, narrowed by X-KBC-ProjectId once a project is known, on both local and deployed sessions. No legacy Storage token is minted for it anywhere in this flow anymore; the old X-StorageAPI-Token header path is untouched and only applies to a genuinely legacy, non-programmatic token supplied directly. Co-Authored-By: Claude Sonnet 5 --- feature_spec/oauth_session_exchange/RFC.md | 1 + src/keboola_mcp_server/clients/auth_bridge.py | 97 +++-------------- src/keboola_mcp_server/mcp.py | 59 ++-------- tests/clients/test_auth_bridge.py | 102 +++--------------- tests/test_mcp.py | 71 ++++-------- 5 files changed, 61 insertions(+), 269 deletions(-) diff --git a/feature_spec/oauth_session_exchange/RFC.md b/feature_spec/oauth_session_exchange/RFC.md index c961b9ab9..80060d507 100644 --- a/feature_spec/oauth_session_exchange/RFC.md +++ b/feature_spec/oauth_session_exchange/RFC.md @@ -101,3 +101,4 @@ The MCP server's public/remote OAuth login (`SimpleOAuthProvider`, `oauth.py`) c 3. **`X-Subject-Token` confirmed** as the header name (shared constant with `resolve-storage-token`, both defined as `SUBJECT_TOKEN_HEADER = 'X-Subject-Token'` in Connection's source). **`X-KBC-ManageApiToken` confirmed NOT sent, verified against a real stack.** It's a separate, mutually-exclusive authenticator (`ManageTokenAuthenticator`, a real Manage-token lookup) from `X-Kubernetes-Authorization` (`KubernetesAuthenticator`, synthetic-token path) — sending both caused a live 401, because `AuthBridgeOAuthExchangeProcessor` explicitly rejects a non-synthetic (i.e. not Kubernetes-authenticated) token even if it also happens to carry the right Manage scope. Only `X-Kubernetes-Authorization` is sent, exactly like the sibling `resolve-storage-token` endpoint. 4. **No dual refresh — confirmed, not assumed.** `TokenRefreshProcessor.php` (Connection) operates on `ProgrammaticSession`/`ProgrammaticSessionRepository` — the same entity and `/v1/auth/token/refresh` mechanism already used by PAT/PKCE-login sessions, fully independent of the league OAuth session. The exchanged session refreshes on its own, forever, via the existing `refresh_tokens()`; the league OAuth token/refresh-token pair is used exactly once (at initial exchange) and never touched again, including on refresh. This **simplifies** `exchange_refresh_token()` relative to today's implementation (which currently re-negotiates with Connection's OAuth server on every refresh) rather than adding a second refresh call. 5. **No rollout coordination.** `/oauth/authorize` removal is immediate with no old-client fallback and no deploy-order dependency communicated from Connection's side. Verify locally against a real stack before shipping; no special deploy sequencing planned. +6. **As-built deviation: `resolve-storage-token`/`StorageTokenResolver` removed entirely, not reused.** §64/§66 originally assumed `create_session_state`'s deployed-path branch would keep converting a programmatic token into a legacy per-project Storage token via the auth-bridge resolver once a project id is known. Live testing against a real dev stack surfaced a 403 on that resolver call (a separate, independently-provisioned Manage scope, `SCOPE_INTERNAL_AUTH_BRIDGE_RESOLVE_STORAGE_TOKEN`, from `exchange-oauth-token`'s) — and confirmed `KeboolaClient` already forwards `bearer_or_sapi_token` (`Authorization: Bearer`) to every service it wraps (Storage, Queue, AI, Data Science, Scheduler, Sync Actions, Metastore) whenever a bearer token is set. So a programmatic token — OAuth-exchanged or a directly-supplied `kbc_pat_*` — is now **always** forwarded as a Bearer, narrowed to a project via `X-KBC-ProjectId` once known, on both local and deployed sessions. No further exchange into a legacy Storage token is performed anywhere in this flow; that resolver/endpoint is no longer called by this codebase. The old `X-StorageAPI-Token` legacy-token header path is unaffected and untouched — it only applies to a genuinely old, non-programmatic token supplied directly, and is expected to be deprecated separately in the future. diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index f31ef2c6a..ac5cd304c 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -1,18 +1,16 @@ -"""Auth-bridge exchanges against Connection's internal endpoints (PSGO-261). - -Two exchanges live here, both authenticating to Connection with the MCP server's own -projected Kubernetes ServiceAccount JWT (`X-Kubernetes-Authorization`): - -- `StorageTokenResolver`: a programmatic bearer token (`kbc_at_*`/`kbc_pat_*`) presented - to the MCP server is exchanged for a legacy Storage token, used for all downstream - Storage-token APIs exactly as before. -- `OAuthSessionExchanger`: a league OAuth access token from the remote/HTTP OAuth login - flow (`oauth.py`) is exchanged for a whole-stack Keboola programmatic session, which - then feeds into the same downstream pipe as a directly-supplied `kbc_at_*` token. - -In both cases the user's token travels as `X-Subject-Token`. The SA token file is read -per call so kubelet rotation is honored. No token material is ever logged or placed in -exception messages. +"""Auth-bridge exchange against Connection's internal endpoint (PSGO-261). + +`OAuthSessionExchanger` exchanges a league OAuth access token from the remote/HTTP OAuth +login flow (`oauth.py`) for a whole-stack Keboola programmatic session (`kbc_at_*`), +authenticating to Connection with the MCP server's own projected Kubernetes ServiceAccount +JWT (`X-Kubernetes-Authorization`); the user's token travels as `X-Subject-Token`. The +resulting `kbc_at_*` session feeds into the same downstream pipe as a directly-supplied +one -- forwarded as a Bearer to every service `KeboolaClient` wraps (Storage, Queue, AI, +etc.), narrowed to a project via `X-KBC-ProjectId` once known. No further exchange into a +legacy per-project Storage token is needed or performed. + +The SA token file is read per call so kubelet rotation is honored. No token material is +ever logged or placed in exception messages. """ import logging @@ -27,7 +25,6 @@ _ACCESS_TOKEN_PREFIX = 'kbc_at_' _PAT_PREFIX = 'kbc_pat_' -_RESOLVE_ENDPOINT = 'manage/internal/auth-bridge/resolve-storage-token' _EXCHANGE_OAUTH_ENDPOINT = 'manage/internal/auth-bridge/exchange-oauth-token' # Resolver statuses passed through to the client verbatim; anything else (incl. 5xx, # timeouts, network failures) is mapped to 502 Bad Gateway. @@ -66,10 +63,6 @@ def __str__(self) -> str: return self.args[0] -class StorageTokenExchangeError(_AuthBridgeExchangeError): - """Raised when the auth-bridge resolver fails to exchange a programmatic token.""" - - class _AuthBridgeClient: """Shared setup for auth-bridge clients: base URL, SA-token path, timeout, transport.""" @@ -97,72 +90,13 @@ def _read_sa_jwt(self) -> str: return read_service_account_jwt(self._kubernetes_token_path) -class StorageTokenResolver(_AuthBridgeClient): - """Exchanges a programmatic token for a legacy Storage token via the Connection resolver.""" - - async def resolve(self, *, subject_token: str, project_id: int) -> str: - """ - Exchanges ``subject_token`` for the legacy Storage token of ``project_id``. - - :return: The legacy Storage token. - :raises StorageTokenExchangeError: On any resolver failure (status carried on the error). - """ - headers = { - 'Content-Type': 'application/json', - 'Accept': 'application/json', - 'X-Kubernetes-Authorization': f'Bearer {self._read_sa_jwt()}', - 'X-Subject-Token': f'Bearer {strip_bearer(subject_token)}', - } - try: - async with httpx.AsyncClient(timeout=self._timeout, transport=self._transport) as client: - response = await client.post( - f'{self._base_url}/{_RESOLVE_ENDPOINT}', - headers=headers, - json={'projectId': project_id}, - ) - except httpx.HTTPError as e: - # Network / timeout failure. Raise without chaining so no request (and thus no - # token material) can surface in a traceback. - raise StorageTokenExchangeError( - f'Auth-bridge token exchange could not reach Connection ({type(e).__name__}).', - status_code=int(HTTPStatus.BAD_GATEWAY), - ) from None - - if response.status_code != HTTPStatus.OK: - status = response.status_code - mapped = status if status in _PASS_THROUGH_STATUSES else int(HTTPStatus.BAD_GATEWAY) - LOG.error(f'Auth-bridge token exchange failed: resolver status {status}, mapped to {mapped}.') - raise StorageTokenExchangeError( - f'Auth-bridge token exchange was rejected (resolver status {status}).', - status_code=mapped, - ) - - try: - body = response.json() - except ValueError: - raise StorageTokenExchangeError( - 'Auth-bridge token exchange returned a non-JSON body.', - status_code=int(HTTPStatus.BAD_GATEWAY), - ) from None - storage_token = body.get('storageToken') if isinstance(body, dict) else None - if not storage_token: - raise StorageTokenExchangeError( - 'Auth-bridge token exchange returned no storageToken.', - status_code=int(HTTPStatus.BAD_GATEWAY), - ) - return cast(str, storage_token) - - class OAuthTokenExchangeError(_AuthBridgeExchangeError): """Raised when the auth-bridge fails to exchange a league OAuth token for a programmatic session.""" class OAuthSessionExchanger(_AuthBridgeClient): """Exchanges a league OAuth access token (``claudai projectless`` scope) for a whole-stack - Keboola programmatic session (PSGO-261 oauth_session_exchange RFC). - - Sibling of `StorageTokenResolver`, reusing the same SA-JWT / ``X-Subject-Token`` mechanism. - """ + Keboola programmatic session (PSGO-261 oauth_session_exchange RFC).""" async def exchange(self, *, oauth_access_token: str) -> dict: """ @@ -173,8 +107,7 @@ async def exchange(self, *, oauth_access_token: str) -> dict: """ # X-KBC-ManageApiToken is a DIFFERENT, mutually-exclusive authenticator (a real Manage # token lookup) -- confirmed against Connection's source that it must never be sent - # alongside X-Kubernetes-Authorization; the k8s JWT alone authorizes this endpoint, - # exactly like the sibling resolve-storage-token endpoint. + # alongside X-Kubernetes-Authorization; the k8s JWT alone authorizes this endpoint. sa_jwt = self._read_sa_jwt() headers = { 'Content-Type': 'application/json', diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 34ff3970e..0766575a8 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -34,7 +34,7 @@ from starlette.types import ASGIApp, Receive, Scope, Send from keboola_mcp_server.auth_login import exchange_scoped_token, get_access_token, introspect_token, load_tokens -from keboola_mcp_server.clients.auth_bridge import StorageTokenResolver, is_programmatic_token, strip_bearer +from keboola_mcp_server.clients.auth_bridge import is_programmatic_token, strip_bearer from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config, ServerRuntimeInfo, deployed_sa_token_path, is_same_stack @@ -616,38 +616,6 @@ async def _resolve_local_tokens( config = dataclasses.replace(config, storage_token=token, project_id=project_id) return config, scope - @classmethod - async def _exchange_programmatic_token(cls, config: Config) -> str: - """ - Exchanges a programmatic token (kbc_at_/kbc_pat_) for the project's legacy Storage token. - - The resolver is reached only on the deployed MCP server, which has a projected - ServiceAccount token at ``KBC_KUBERNETES_TOKEN_PATH`` (read from the process - environment only, never from per-request config). A project id is required because - a programmatic token is not project-bound. - """ - kubernetes_token_path = deployed_sa_token_path() - if not kubernetes_token_path: - raise ValueError( - 'Received a Keboola programmatic token (kbc_at_/kbc_pat_) but KBC_KUBERNETES_TOKEN_PATH ' - 'is not configured. Programmatic-token exchange is available only on the deployed MCP server.' - ) - if not config.project_id: - raise ValueError( - 'A project id is required to exchange a programmatic token. ' - 'Set the KBC_PROJECT_ID env var or the X-KBC-ProjectId header.' - ) - try: - project_id = int(config.project_id) - except (TypeError, ValueError): - raise ValueError(f'Invalid project id for programmatic-token exchange: {config.project_id!r}') - - resolver = StorageTokenResolver( - storage_api_url=config.storage_api_url, - kubernetes_token_path=kubernetes_token_path, - ) - return await resolver.resolve(subject_token=config.storage_token, project_id=project_id) - @classmethod async def create_session_state( cls, @@ -684,21 +652,16 @@ async def create_session_state( bearer_token = config.bearer_token extra_headers: dict[str, Any] = {} if is_programmatic_token(storage_token): - if deployed_sa_token_path() and config.project_id: - # Deployed, and a project is already known (header, or a prior scope selection): - # exchange the programmatic token for that project's legacy Storage token via the - # auth-bridge resolver, then use it downstream unchanged. - storage_token = await cls._exchange_programmatic_token(config) - bearer_token = None - else: - # No projected SA token to reach the resolver (local), OR no project is known yet - # (deployed, e.g. a freshly-exchanged whole-stack OAuth session pre-scoping): forward - # the programmatic token downstream as a Bearer so get_accessible_projects/ - # set_project_scope can introspect/scope it directly. Strip any inbound `Bearer ` - # scheme so the client's own `Bearer ` prefixing can't produce `Bearer Bearer …`. - bearer_token = strip_bearer(storage_token) - if config.project_id: - extra_headers['X-KBC-ProjectId'] = config.project_id + # A programmatic token (kbc_at_/kbc_pat_) is forwarded downstream as a Bearer -- + # KeboolaClient already sends it that way to every service it wraps (Storage, Queue, + # AI, etc.), so no legacy per-project Storage token needs to be minted for it. Strip + # any inbound `Bearer ` scheme so the client's own `Bearer ` prefixing can't produce + # `Bearer Bearer …`. Narrow to a specific project via X-KBC-ProjectId when known + # (header, or a prior scope selection) -- unset (whole-stack) is exactly what + # get_accessible_projects/set_project_scope need before a project is chosen. + bearer_token = strip_bearer(storage_token) + if config.project_id: + extra_headers['X-KBC-ProjectId'] = config.project_id client = await KeboolaClient( storage_api_url=config.storage_api_url, diff --git a/tests/clients/test_auth_bridge.py b/tests/clients/test_auth_bridge.py index ba831c750..037259346 100644 --- a/tests/clients/test_auth_bridge.py +++ b/tests/clients/test_auth_bridge.py @@ -6,13 +6,7 @@ import httpx import pytest -from keboola_mcp_server.clients.auth_bridge import ( - OAuthSessionExchanger, - OAuthTokenExchangeError, - StorageTokenExchangeError, - StorageTokenResolver, - is_programmatic_token, -) +from keboola_mcp_server.clients.auth_bridge import OAuthSessionExchanger, OAuthTokenExchangeError, is_programmatic_token STORAGE_API_URL = 'https://connection.keboola.com' @@ -41,88 +35,9 @@ def sa_token_file(tmp_path: Path) -> Path: return path -def _resolver(sa_token_file: Path, handler) -> StorageTokenResolver: - return StorageTokenResolver( - storage_api_url=STORAGE_API_URL, - kubernetes_token_path=str(sa_token_file), - transport=httpx.MockTransport(handler), - ) - - -@pytest.mark.asyncio -async def test_resolve_success_sends_expected_request(sa_token_file: Path) -> None: - captured: dict[str, httpx.Request] = {} - - def handler(request: httpx.Request) -> httpx.Response: - captured['request'] = request - return httpx.Response(HTTPStatus.OK, json={'storageToken': 'legacy-token', 'projectId': 42}) - - resolver = _resolver(sa_token_file, handler) - token = await resolver.resolve(subject_token='Bearer kbc_pat_abc', project_id=42) - - assert token == 'legacy-token' - rq = captured['request'] - assert rq.url.path == '/manage/internal/auth-bridge/resolve-storage-token' - assert rq.headers['X-Kubernetes-Authorization'] == 'Bearer sa-jwt-value' - # Subject token is normalized to a single Bearer scheme regardless of inbound form. - assert rq.headers['X-Subject-Token'] == 'Bearer kbc_pat_abc' - - -@pytest.mark.parametrize('status', [HTTPStatus.BAD_REQUEST, HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN]) -@pytest.mark.asyncio -async def test_resolve_passes_through_client_errors(sa_token_file: Path, status: HTTPStatus) -> None: - resolver = _resolver(sa_token_file, lambda rq: httpx.Response(status, json={'error': 'nope'})) - with pytest.raises(StorageTokenExchangeError) as exc: - await resolver.resolve(subject_token='kbc_at_abc', project_id=1) - assert exc.value.status_code == int(status) - - -@pytest.mark.parametrize('status', [HTTPStatus.INTERNAL_SERVER_ERROR, HTTPStatus.BAD_GATEWAY, HTTPStatus.NOT_FOUND]) -@pytest.mark.asyncio -async def test_resolve_maps_other_statuses_to_502(sa_token_file: Path, status: HTTPStatus) -> None: - resolver = _resolver(sa_token_file, lambda rq: httpx.Response(status)) - with pytest.raises(StorageTokenExchangeError) as exc: - await resolver.resolve(subject_token='kbc_at_abc', project_id=1) - assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) - - -@pytest.mark.asyncio -async def test_resolve_maps_network_error_to_502(sa_token_file: Path) -> None: - def handler(request: httpx.Request) -> httpx.Response: - raise httpx.ConnectError('boom', request=request) - - resolver = _resolver(sa_token_file, handler) - with pytest.raises(StorageTokenExchangeError) as exc: - await resolver.resolve(subject_token='kbc_at_abc', project_id=1) - assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) - # No token material leaks into the message. - assert 'kbc_at_abc' not in str(exc.value) - - -@pytest.mark.asyncio -async def test_resolve_missing_storage_token_maps_to_502(sa_token_file: Path) -> None: - resolver = _resolver(sa_token_file, lambda rq: httpx.Response(HTTPStatus.OK, json={'projectId': 1})) - with pytest.raises(StorageTokenExchangeError) as exc: - await resolver.resolve(subject_token='kbc_at_abc', project_id=1) - assert exc.value.status_code == int(HTTPStatus.BAD_GATEWAY) - - -@pytest.mark.asyncio -async def test_resolve_empty_sa_token_file_fails_loudly(tmp_path: Path) -> None: - empty = tmp_path / 'empty' - empty.write_text(' ') - resolver = StorageTokenResolver( - storage_api_url=STORAGE_API_URL, - kubernetes_token_path=str(empty), - transport=httpx.MockTransport(lambda rq: httpx.Response(HTTPStatus.OK, json={'storageToken': 'x'})), - ) - with pytest.raises(ValueError, match='empty'): - await resolver.resolve(subject_token='kbc_at_abc', project_id=1) - - def test_invalid_storage_api_url_rejected(sa_token_file: Path) -> None: with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): - StorageTokenResolver(storage_api_url='https://example.com', kubernetes_token_path=str(sa_token_file)) + OAuthSessionExchanger(storage_api_url='https://example.com', kubernetes_token_path=str(sa_token_file)) def _exchanger(sa_token_file: Path, handler) -> OAuthSessionExchanger: @@ -158,6 +73,19 @@ def handler(request: httpx.Request) -> httpx.Response: assert rq.headers['X-Subject-Token'] == 'Bearer league-oauth-token' +@pytest.mark.asyncio +async def test_exchange_empty_sa_token_file_fails_loudly(tmp_path: Path) -> None: + empty = tmp_path / 'empty' + empty.write_text(' ') + exchanger = OAuthSessionExchanger( + storage_api_url=STORAGE_API_URL, + kubernetes_token_path=str(empty), + transport=httpx.MockTransport(lambda rq: httpx.Response(HTTPStatus.OK, json={'accessToken': 'x'})), + ) + with pytest.raises(ValueError, match='empty'): + await exchanger.exchange(oauth_access_token='league-oauth-token') + + @pytest.mark.parametrize('status', [HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN]) @pytest.mark.asyncio async def test_exchange_passes_through_client_errors(sa_token_file: Path, status: HTTPStatus) -> None: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 6480d4484..fab85977d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -871,70 +871,37 @@ def test_apply_request_config_injects_exchanged_session_token(self): assert is_programmatic_token(out_config.storage_token) -class TestProgrammaticTokenExchange: - """SessionStateMiddleware exchanges programmatic tokens via the auth-bridge resolver (PSGO-261).""" +class TestProgrammaticTokenForwarding: + """A programmatic token (kbc_at_/kbc_pat_) is always forwarded downstream as a Bearer (PSGO-261). - @pytest.mark.asyncio - async def test_missing_kubernetes_token_path_raises(self, monkeypatch) -> None: - monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) - config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_pat_abc', project_id='1') - with pytest.raises(ValueError, match='KBC_KUBERNETES_TOKEN_PATH'): - await SessionStateMiddleware._exchange_programmatic_token(config) - - @pytest.mark.asyncio - async def test_missing_project_id_raises(self, monkeypatch) -> None: - monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') - config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_pat_abc') - with pytest.raises(ValueError, match='project id is required'): - await SessionStateMiddleware._exchange_programmatic_token(config) + KeboolaClient already sends a Bearer token to every service it wraps (Storage, Queue, AI, + etc.), so no legacy per-project Storage token needs to be minted via the auth-bridge resolver + -- that resolver call was removed entirely; see git history for the prior + `_exchange_programmatic_token`/`StorageTokenResolver` code this replaced. + """ @pytest.mark.asyncio - async def test_invalid_project_id_raises(self, monkeypatch) -> None: - monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + @pytest.mark.parametrize('kubernetes_token_path', [None, '/var/run/secrets/token'], ids=['local', 'deployed']) + @pytest.mark.parametrize('project_id', [None, '42'], ids=['no_project_id', 'with_project_id']) + async def test_forwards_bearer_regardless_of_deployment_or_project_id( + self, monkeypatch, kubernetes_token_path: str | None, project_id: str | None + ) -> None: + if kubernetes_token_path: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', kubernetes_token_path) + else: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) config = Config( - storage_api_url='https://connection.keboola.com', storage_token='kbc_pat_abc', project_id='not-an-int' - ) - with pytest.raises(ValueError, match='Invalid project id'): - await SessionStateMiddleware._exchange_programmatic_token(config) - - @pytest.mark.asyncio - async def test_happy_path_calls_resolver(self, monkeypatch) -> None: - monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') - config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_abc', project_id='42') - - resolver = MagicMock() - resolver.resolve = AsyncMock(return_value='legacy-storage-token') - with patch('keboola_mcp_server.mcp.StorageTokenResolver', return_value=resolver) as resolver_cls: - token = await SessionStateMiddleware._exchange_programmatic_token(config) - - assert token == 'legacy-storage-token' - resolver_cls.assert_called_once_with( - storage_api_url='https://connection.keboola.com', kubernetes_token_path='/var/run/secrets/token' + storage_api_url='https://connection.keboola.com', storage_token='kbc_at_abc', project_id=project_id ) - resolver.resolve.assert_awaited_once_with(subject_token='kbc_at_abc', project_id=42) - - @pytest.mark.asyncio - async def test_deployed_without_project_id_forwards_bearer_instead_of_exchanging(self, monkeypatch) -> None: - """A deployed OAuth session starts whole-stack (no project_id yet, RFC decision §2); it must - forward the programmatic token as a Bearer for get_accessible_projects/set_project_scope to - introspect, not fail by attempting a resolver exchange that requires a project id.""" - monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') - config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_abc') runtime_info = ServerRuntimeInfo(transport='http') - with ( - patch.object( - SessionStateMiddleware, - '_exchange_programmatic_token', - AsyncMock(side_effect=AssertionError('must not exchange without a known project_id')), - ), - patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')), - ): + with patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')): state = await SessionStateMiddleware.create_session_state(config, runtime_info) client = state[KeboolaClient.STATE_KEY] assert client.bearer_token == 'kbc_at_abc' assert client.token == 'kbc_at_abc' + assert client.headers.get('X-KBC-ProjectId') == project_id class TestMaybeUseStoredSession: From 7dff1afbf05db6181de15032c3556fff100866de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 29 Jul 2026 10:00:38 +0200 Subject: [PATCH 58/89] docs(PSGO-261): add Postgres-backed OAuth session store RFC; reconcile prior RFCs with as-built code New RFC (oauth_session_persistence) proposes replacing the self-contained OAuth access/refresh JWTs with an opaque, Postgres-backed session reference -- enables revocation and server-managed refresh that the current JWT design structurally can't provide, at the cost of a new stateful dependency for the OAuth path only (local login and header/PAT tokens are explicitly unaffected). Also updates pat_token_support/RFC.md and oauth_session_exchange/RFC.md, whose "stateful session"/"resolver exchange" sections had drifted from the scope_token and resolve-storage-token-removal fixes shipped earlier in this branch. Co-Authored-By: Claude Sonnet 5 --- feature_spec/oauth_session_exchange/RFC.md | 2 +- feature_spec/oauth_session_persistence/RFC.md | 237 ++++++++++++++++++ feature_spec/pat_token_support/RFC.md | 45 ++-- 3 files changed, 268 insertions(+), 16 deletions(-) create mode 100644 feature_spec/oauth_session_persistence/RFC.md diff --git a/feature_spec/oauth_session_exchange/RFC.md b/feature_spec/oauth_session_exchange/RFC.md index 80060d507..111f52a00 100644 --- a/feature_spec/oauth_session_exchange/RFC.md +++ b/feature_spec/oauth_session_exchange/RFC.md @@ -63,7 +63,7 @@ The MCP server's public/remote OAuth login (`SimpleOAuthProvider`, `oauth.py`) c 2. Code→token exchange is **unchanged**: still `POST {server_url}/oauth/token`, still returns a league OAuth access/refresh pair (now carrying both scopes). 3. **New step, replacing `_create_sapi_token`:** exchange the league OAuth access token for a Keboola programmatic session via `manage/internal/auth-bridge/exchange-oauth-token`, reusing the exact SA-JWT + `X-Subject-Token` mechanism already implemented for `resolve-storage-token` (`clients/auth_bridge.py`). 4. Parse the `CliTokenResponse` into the same shape `auth_login.py._parse_token_response` already builds from a PKCE response. -5. From here on this session is **indistinguishable, downstream, from a directly-supplied `kbc_at_*` token**: `is_programmatic_token()` detects it, `create_session_state`'s existing deployed-path branch (`_exchange_programmatic_token` → `StorageTokenResolver.resolve`) runs unchanged, and the full PSGO-261 multi-project machinery (`get_accessible_projects`, `set_project_scope`, read fan-out) becomes available to every OAuth client for the first time, starting whole-stack/unconfirmed exactly like a fresh PKCE login. +5. From here on this session is **indistinguishable, downstream, from a directly-supplied `kbc_at_*` token**: `is_programmatic_token()` detects it, `create_session_state` forwards it as `Authorization: Bearer` narrowed by `X-KBC-ProjectId` once a project is known (see Decision §6 — no legacy-token resolver exchange is performed), and the full PSGO-261 multi-project machinery (`get_accessible_projects`, `set_project_scope`, read fan-out) becomes available to every OAuth client for the first time, starting whole-stack/unconfirmed exactly like a fresh PKCE login. 6. The league OAuth token pair from step 2 is discarded after step 3 completes — never persisted, never refreshed. 7. `ProxyAccessToken.sapi_token` (currently required, justified only by "Jobs Queue/AI Service don't support bearer tokens yet") is **obsolete**: those clients now speak bearer via `bearer_or_sapi_token` (PSGO-261, commit `5b8c65ed`). Remove the field; repurpose `ProxyAccessToken` to carry the new `kbc_at_` token and its own `refresh_token`/`session_id` instead. 8. **Refresh is fully decoupled from the league OAuth session** (confirmed, §4): `exchange_refresh_token()` calls `POST /v1/auth/token/refresh` (already implemented, `auth_login.py.refresh_tokens`) directly against the previously-exchanged refresh token. The league OAuth `/oauth/token` refresh grant is never invoked again after step 2. diff --git a/feature_spec/oauth_session_persistence/RFC.md b/feature_spec/oauth_session_persistence/RFC.md new file mode 100644 index 000000000..fd9a059c4 --- /dev/null +++ b/feature_spec/oauth_session_persistence/RFC.md @@ -0,0 +1,237 @@ +# RFC: Postgres-backed OAuth session store (replaces self-contained JWT session) + +Linear: [PSGO-261](https://linear.app/keboola/issue/PSGO-261/support-pat-tokens-in-mcp-server-mcp-server) +Parent: PSGO-261. Closes the "keyring/DB credential storage" carve-out `pat_token_support/RFC.md` +explicitly deferred in increment 1 ("Still out of scope: ... keyring/DB credential storage"). +Related: `oauth_session_exchange/RFC.md` (the exchange this RFC changes how the result is stored), +`pat_token_support/RFC.md` §"Transport note" (the `scope_token` mechanism this RFC's scope columns +can absorb for OAuth sessions specifically). + +## Problem + +Today, the deployed OAuth login flow (`oauth.py`, `SimpleOAuthProvider`) stores **nothing** +server-side. Every piece of session state — the OAuth authorize-state, the authorization code, the +access token, the refresh token, and (as of the `scope_token` fix) the confirmed multi-project scope +— is self-encoded into a signed, gzip-compressed JWT (`jwt_utils.py`) and handed to the client, which +resends it on every subsequent request. This was a deliberate choice (see `oauth.py`'s own comment: +*"We don't store the authentication states... instead we encode them to JWT"*) and it is fully +stateless: any replica can decode any token with the shared `KBC_JWT_SECRET`, no shared datastore +needed, correct under the MCP 2026-07-28 RC's stateless-transport direction. + +That design has three real costs, all inherent to "the client holds the truth, signed": + +1. **No revocation.** A leaked or compromised `ProxyAccessToken`/`ProxyRefreshToken`/`scope_token` is + valid until its embedded expiry, full stop — there is no server-side list to delete from. Ending a + session early (logout, incident response, revoking a compromised token) is not possible today. +2. **Client-visible plumbing.** `scope_token` must be threaded through every tool call as an explicit + argument (see `pat_token_support/RFC.md` "Transport note") because there is nowhere else for the + confirmed scope to live between requests. This works, but it is visible surface area the calling + agent has to carry correctly every single call. +3. **Refresh is the client's problem.** The MCP client must notice its access token is nearing + expiry and call this server's `/oauth/token` with `grant_type=refresh_token` — today's + `exchange_refresh_token()` only runs when the client initiates that call. There's no way for this + server to refresh the underlying `kbc_access_token`/`kbc_refresh_token` proactively or transparently. + +**Proposed change:** store the OAuth session server-side, in Postgres, encrypted at rest. The token +the MCP client holds becomes a short, opaque, random reference (not a JWT carrying real credentials) +that this server looks up, decrypts, and — if the underlying Keboola token is near expiry — refreshes +transparently before using. This is a deliberate, scoped trade: give up "zero shared infra" for the +OAuth path specifically, in exchange for revocation, a smaller/opaque client-facing token, and +server-managed refresh. It does **not** change the local PKCE `login` flow (`~/.keboola/mcp/credentials.json`, +unaffected) or the header/PAT-supplied-token flow (still fully stateless, unaffected) — see Scope. + +## Required Behavior + +### Token model change + +| | Today (JWT, self-contained) | Proposed (Postgres, opaque reference) | +|---|---|---| +| What the MCP client holds | A JWT with the real `kbc_access_token`/`kbc_refresh_token`/`scope` embedded, HMAC-signed | A random opaque string (e.g. 256 bits, base64url) that is *only* a lookup key | +| How the server validates it | Verify HMAC signature, decode payload | Look up by the opaque string (hashed) in Postgres; row must exist, not be revoked, not be expired | +| Where the real Keboola credentials live | Inside the JWT, in the client's possession | Encrypted (AES-256-GCM) in Postgres only; never sent to the client | +| Revocation | Not possible before natural expiry | `DELETE`/soft-revoke the row; token is dead on the next lookup | +| Refresh | Client-initiated, via `/oauth/token` `grant_type=refresh_token` | Server-initiated, lazily: on lookup, if `kbc_access_token` is near expiry, refresh via `refresh_tokens()` and update the row in place — the client's opaque token does not need to change | +| Multi-project scope (`scope_token`) | Separate signed JWT, resent as a tool argument every call | Columns on the same session row; no `scope_token` argument needed for OAuth sessions at all | + +### Schema (new `oauth_sessions` table, one row per logged-in session) + +| Column | Type | Notes | +|---|---|---| +| `id` | `uuid`, PK | Internal row id | +| `access_token_hash` | `bytea`, unique, indexed | `sha256` of the opaque access token the client holds — store the hash, not the token, so a DB read alone can't leak a live bearer credential | +| `refresh_token_hash` | `bytea`, unique, indexed, nullable | Same, for the opaque refresh token | +| `client_id` | `text` | The OAuth client (`claude.ai`, etc.) — audit/introspection only | +| `user_email` | `text`, nullable | From the exchange response, for audit/introspection | +| `kbc_access_token_enc` | `bytea` | AES-256-GCM ciphertext of the real `kbc_access_token` | +| `kbc_refresh_token_enc` | `bytea` | AES-256-GCM ciphertext of the real `kbc_refresh_token` | +| `kbc_access_expires_at` | `timestamptz` | Drives the lazy-refresh check | +| `scope_project_ids` | `int[]`, nullable | Confirmed multi-project scope — absorbs `scope_token`'s job for OAuth sessions | +| `scope_read_only` | `boolean`, default `false` | | +| `scope_confirmed` | `boolean`, default `false` | | +| `scope_scoped_token_enc` | `bytea`, nullable | AES-256-GCM ciphertext of the minted scoped token (`/v1/auth/pat/exchange` result) | +| `scope_scoped_expires_at` | `timestamptz`, nullable | | +| `created_at` / `updated_at` / `last_used_at` | `timestamptz` | | +| `revoked_at` | `timestamptz`, nullable | Soft-revoke; a non-null value makes lookup fail as if the row didn't exist | + +Every encrypted column uses **AES-256-GCM** (authenticated encryption — tamper-evident, not just +confidential) via the `cryptography` package, already a pinned dependency (`pyproject.toml:22`, +`~= 49.0`) — no new crypto library needed, just a new small module (`session_store/crypto.py`) wrapping +`cryptography.hazmat.primitives.ciphers.aead.AESGCM`. + +### New env var + +`KBC_SESSION_ENCRYPTION_KEY` — 32 raw bytes, base64-encoded for env-var transport (`base64.b64decode` +on load, fail loudly at startup if it doesn't decode to exactly 32 bytes). Single static key for v1 +(see Open Questions on rotation). Mirrors how `KBC_JWT_SECRET` is already handled today +(`config.jwt_secret`) — same "required in production, generate an ephemeral one locally if unset" +posture, so local dev/tests work with zero setup. + +### Refresh strategy: lazy, on lookup — not a background job + +Every session lookup (equivalent to today's `load_access_token()`) checks `kbc_access_expires_at` +against `is_near_expiry`-style logic (reuse the exact same 60-second-early check `SessionScope` +already uses) and refreshes in place via the existing `refresh_tokens()` (`auth_login.py`) before +returning the decrypted token — the same "check-then-refresh-then-use" shape already implemented for +`scope_token`'s `scoped_token` re-mint in `_resolve_local_tokens`. **No new scheduler, no background +worker, no cron** for v1 — refresh only happens on an actual request, which is simpler to reason +about and test, at the cost of the first request after a long idle period paying one extra refresh +round-trip (acceptable; this is the same trade the current design already makes for `scope_token`). + +### What the MCP client actually sees + +Nothing changes about the OAuth *dance* — `/oauth/consent`, code exchange, redirect — only what comes +back at the end. `SimpleOAuthProvider.exchange_authorization_code()` mints a session row instead of a +JWT and returns the row's opaque access/refresh token pair as today's `AccessToken`/`RefreshToken` +Pydantic models (same wire shape, different contents — a random string instead of a JWT). Client code +requires zero changes; this is entirely a server-internal storage swap from the MCP client's point of +view. `set_project_scope`/`get_accessible_projects` **stop returning `scope_token`** for OAuth +sessions (scope now lives on the row, found via the same access-token lookup already required on +every request) — the two tools' models keep `scope_token: str | None` for backward compat with +header/PAT sessions (see Scope), just always `None` when the caller authenticated via OAuth. + +## Resolution Strategy + +- **New package `src/keboola_mcp_server/session_store/`**: + - `crypto.py` — `encrypt(plaintext: bytes, key: bytes) -> bytes` / `decrypt(...)`, thin AES-256-GCM + wrapper (nonce prepended to ciphertext, standard practice). + - `repository.py` — `SessionStore` protocol (`create`, `get_by_access_token`, `get_by_refresh_token`, + `update_kbc_tokens`, `update_scope`, `revoke`) + a `PostgresSessionStore` implementation using + `asyncpg` (async-native, matches this codebase's existing all-`httpx`-async style; **no ORM** — + a single-table store doesn't earn SQLAlchemy's overhead, and the rest of the codebase has zero + ORM precedent to extend). The protocol exists so `oauth.py`/tests can mock the store without a + real database (unit tests) while integration tests exercise `PostgresSessionStore` against a real + one (see Testing). + - `migrations/0001_oauth_sessions.sql` (+ a ~15-line runner: a `schema_migrations` tracking table, + apply un-applied numbered `.sql` files in order at startup or via a `keboola-mcp-server migrate` + CLI subcommand — deliberately not `alembic`; one table doesn't need a migration framework, a + numbered-SQL-files-plus-tracking-table is the whole mechanism and is trivially testable). +- **`config.py`**: add `postgres_dsn: Optional[str]` and `session_encryption_key: Optional[str]` + fields, same env-var-mapping mechanism as every other `Config` field. +- **`oauth.py`**: `SimpleOAuthProvider` gains a `session_store: SessionStore` constructor param. + `exchange_authorization_code`/`exchange_refresh_token`/`load_access_token`/`load_refresh_token` + are rewritten against the store instead of `self._encode`/`self._decode`. The authorize-state JWT + (5-minute TTL, `authorize()`) and the authorization-code JWT (`_ExtendedAuthorizationCode`) are + **unchanged** — they're short-lived, single-use, pre-authentication artifacts with no real + credentials embedded, and encoding them as JWTs today is already fine; only the *long-lived, + real-credential-carrying* access/refresh/scope tokens move to Postgres. +- **`mcp.py`**: `SessionStateMiddleware`/`MultiProjectMiddleware` gain a scope-store lookup path for + OAuth sessions (keyed by the same `AuthenticatedUser.access_token` already resolved by the MCP SDK's + auth layer) instead of decoding `scope_token` from arguments — `set_project_scope` writes + `scope_*` columns on the row instead of minting a JWT. The `_read_scope_from_request`/`_SCOPE_TOKEN_ARG` + path stays exactly as-is for non-OAuth sessions (see Scope). +- **`server.py`**: construct the `SessionStore` (real `PostgresSessionStore` if `postgres_dsn` is set, + otherwise refuse to start an OAuth-enabled server without one — no silent in-memory fallback for a + production auth path) and pass it into `SimpleOAuthProvider`. +- **`docker-compose.yml`** (new, repo root): a single `postgres:16` service for local dev/integration + tests — named volume, healthcheck, default credentials for local use only (never used in any real + deployment, which gets its own managed Postgres instance via kbc-stacks, out of scope here). + +## Scope + +**In scope:** the deployed, OAuth-authenticated session path only — `exchange-oauth-token` result +storage, refresh, and (optionally, see Open Questions) the multi-project scope for OAuth sessions. +Postgres schema + migrations + local docker-compose + AES-256-GCM encryption + unit/integration tests. + +**Explicitly out of scope, unchanged by this RFC:** +- **Local PKCE `login` flow** (`auth_login.py`, `~/.keboola/mcp/credentials.json`) — keeps using its + existing mode-600 file. It has no multi-replica concern (one stdio process, one conversation) and + no revocation need proportionate to the complexity of adding a DB dependency to a local CLI tool. +- **Header/PAT-supplied tokens** (`is_programmatic_token()` path) — stays fully stateless, `scope_token` + keeps working exactly as today for this path. There is no stable per-conversation identifier to key + a DB row on for a bare supplied token the way `session_id`/the OAuth access token gives us for free. +- Postgres HA/backup/monitoring in the actual deployed stacks — that's a kbc-stacks-side concern + (separate repo), same carve-out pattern used for the k8s ServiceAccount scope grants in + `oauth_session_exchange/RFC.md`. +- Encryption-key rotation (see Open Questions) — single static key for v1. +- Background/proactive refresh — lazy-on-lookup only for v1 (see Required Behavior). + +## Delivery plan (phased, compact) + +**Phase 1 — Schema, crypto, store, local dev infra (no behavior change yet).** +- `session_store/` package: `crypto.py`, `repository.py` (`SessionStore` protocol + + `PostgresSessionStore`), `migrations/0001_oauth_sessions.sql` + runner. +- `docker-compose.yml`; `config.py` additions; `KBC_SESSION_ENCRYPTION_KEY` handling (generate + ephemeral locally if unset, same posture as `KBC_JWT_SECRET`). +- Tests: crypto round-trip (encrypt/decrypt, tamper detection via GCM auth tag), migration runner + applies-once idempotency, `PostgresSessionStore` CRUD against a real docker-compose Postgres. + +**Phase 2 — Wire `SimpleOAuthProvider` to the store (the core swap).** +- Replace `_encode`/`_decode` calls for access/refresh tokens with store lookups; lazy-refresh-on-lookup. +- `exchange_authorization_code`: mint a row instead of a JWT pair. +- Tests: exchange creates a row with encrypted tokens; `load_access_token` decrypts + returns; expired + access token triggers exactly one `refresh_tokens()` call and updates the row in place; a revoked + row fails lookup; a tampered ciphertext (flipped bit) fails GCM auth and is treated as invalid, not + silently decrypted wrong. + +**Phase 3 — Move multi-project scope onto the row for OAuth sessions.** +- `set_project_scope`/`get_accessible_projects`: write/read `scope_*` columns via the store when the + session is OAuth-authenticated; `scope_token` stays `None` in their output for this case. +- `mcp.py`: scope resolution branches on session type — OAuth → store lookup, everything else → + existing `_read_scope_from_request`/`scope_token` path, unchanged. +- Tests: OAuth session's `set_project_scope` never returns a `scope_token`; a subsequent call with no + `scope_token` argument still resolves the previously-confirmed scope correctly via the store. + +**Cross-cutting:** version bump (minor — new capability + new required infra for OAuth deployments), +`uv.lock`, new `asyncpg` dependency, CI: a `postgres` service container for the integration-test tox +env, `TOOLS.md` regen (the `scope_token` field's description changes to note it's OAuth-session-conditional). + +## Testing / Verification + +**Unit** — `PostgresSessionStore` mocked out via the `SessionStore` protocol wherever `oauth.py`/`mcp.py` +logic is under test (no real DB needed for these); crypto module tested in full isolation (round-trip, +wrong-key failure, tampered-ciphertext failure) with no DB at all. + +**Integration** — a real Postgres via docker-compose (`docker compose up -d postgres` in the +integration-test tox env, matching how `integtests/` already needs real external services): full +`authorize → consent → callback → token → tool call → refresh-after-forced-expiry → revoke → 401` +cycle against it. + +**Manual** — the same real-dev-stack OAuth login test used throughout this PR's live debugging, +confirming: no `scope_token` in `set_project_scope`'s output for an OAuth session; killing/restarting +the MCP server process mid-conversation and confirming the session survives (this is the concrete, +demonstrable win over the JWT design — a process restart today does *not* invalidate a signed JWT +either, so this specific test doesn't distinguish them; the real differentiator is **revocation**: +manually deleting the row and confirming the *next* request 401s, which a signed JWT cannot do before +its embedded expiry). + +## Open Questions + +1. **Does multi-project scope move to Postgres for OAuth sessions in the same delivery, or later?** + Phase 3 above assumes yes (it's a small addition once the row exists for the access/refresh tokens + anyway) — confirm before starting Phase 3, since it's the part of this RFC that changes tool-facing + output shape (`scope_token` becomes conditionally absent), not just internal storage. +2. **Encryption key rotation.** V1 ships a single static `KBC_SESSION_ENCRYPTION_KEY`. Rotating it + invalidates every stored session (can't decrypt with the old key). Acceptable for v1 (forces + re-login, not data loss — no Keboola data lives in this table, only session credentials), but worth + a documented runbook step before this ships, and a versioned-key-prefix scheme (`v1:`) + would make future rotation non-disruptive if we want to add it later — flagging now so the column + format (prefix the ciphertext with a key-version byte) is decided before Phase 1's migration ships, + not retrofitted after real rows exist. +3. **Session expiry / cleanup.** Rows never get deleted automatically today's plan — need a retention + policy (e.g. delete rows with `revoked_at` set or `last_used_at` older than N days) — a cron/cleanup + job, explicitly out of scope for v1 but should be tracked as an immediate follow-up, not forgotten. +4. **Does Postgres downtime take down OAuth login entirely, or degrade gracefully?** With no DB, no + OAuth session can be created or validated — this is a new hard dependency for the OAuth path (by + design, per Scope: "no silent in-memory fallback for a production auth path"). Confirm this is + acceptable given kbc-stacks' Postgres HA posture before shipping, since it changes the failure mode + of OAuth login from "always works" (today, stateless) to "works iff Postgres is reachable." diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 88244a9ec..98405047c 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -613,18 +613,33 @@ This makes the multi-project path safe on humongous projects: it can never wedge nudges toward the scalable access patterns (search / per-project drill-down) instead of bulk-listing. Follow-up: real `limit`/`offset` pagination on the enumerators, and concurrent fan-out. -## Transport note: multi-project scope needs a stateful session - -Multi-project scope lives in the MCP **session** state (`ctx.session.state[SCOPE_KEY]`), read back on -each request. This only persists when the transport keeps the session alive across requests: - -- **stdio** — one long-lived session per process → scope persists (this is how MPA was developed/tested). -- **streamable-HTTP, stateless** (`stateless_http=True`, the deployed default for horizontal scaling) - → every request is an independent session (`Terminating session: None`), so `set_project_scope`'s - state never reaches the next call and data tools keep reporting "no scope confirmed". -- **streamable-HTTP, stateful** (`--no-stateless-http`) → the server issues an `Mcp-Session-Id` the - client echoes back, the `ServerSession` is reused, and scope persists. - -So to run MPA locally over HTTP: `keboola-mcp-server --transport streamable-http --no-stateless-http`. -Deployed multi-replica MPA over stateless HTTP would need a shared/sticky scope store (out of scope -here; deployed sessions today are single-project via the resolver + `KBC_PROJECT_ID`). +## Transport note: multi-project scope is carried by the caller, not the session (superseded) + +**Superseded.** This section originally assumed multi-project scope had to live in the MCP +**session** state (`ctx.session.state[SCOPE_KEY]`), read back on each request, and that this only +persists when the transport keeps the session alive across requests — fine on stdio (one long-lived +process), broken on the deployed default (`stateless_http=True`, a fresh empty session per request, +confirmed live via Datadog trace evidence: three separate `POST /mcp/` requests sharing one +process/`runtime-id` yet never seeing each other's session state), and only working around that with +`--no-stateless-http` (a single-replica-only workaround, and itself in tension with the direction the +MCP spec is taking: the 2026-07-28 RC removes `Mcp-Session-Id`/session pinning from the protocol +entirely, in favor of stateless-by-default operation). + +**As built:** `set_project_scope`/`get_accessible_projects` sign the confirmed `SessionScope` into an +opaque `scope_token` (`SessionScope.to_token`/`from_token`, `mcp.py`; HMAC-JWT, the same +gzip+`jwt.api_jws` mechanism `SimpleOAuthProvider` already uses for OAuth tokens, extracted into +`jwt_utils.py`) and return it to the caller, who resends it as a tool-call argument on every +subsequent call. `SessionStateMiddleware` decodes it fresh from the request each time +(`_read_scope_from_request`) instead of reading `ctx.session.state` from a prior request. This is +stateless by construction: it works identically on stdio, one HTTP replica, or many, with no shared +store, no sticky routing, and no `--no-stateless-http` workaround needed. The signing secret is +`config.jwt_secret` (`KBC_JWT_SECRET`) when set — required to be shared across replicas for the +existing OAuth JWTs already, so scope tokens ride along for free — or a process-local fallback +(fine for stdio, since one process serves exactly one conversation). + +Separately, the deployed session no longer needs to be single-project via a resolver exchange at +all: `create_session_state` forwards any programmatic token (`kbc_at_*`/`kbc_pat_*`) as +`Authorization: Bearer`, narrowed by `X-KBC-ProjectId` once a project is known — the +`resolve-storage-token` auth-bridge exchange this section referenced has been removed (see +`oauth_session_exchange/RFC.md` Decision §6). Full multi-project scope now works the same way on +the deployed server as it does locally. From 9abc2a53a29f0fec68a8e9d73081628728025dab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 29 Jul 2026 10:16:13 +0200 Subject: [PATCH 59/89] feat(PSGO-261): add Postgres-backed session_store package (Phase 1, no behavior change yet) Foundational layer for oauth_session_persistence RFC: AES-256-GCM encryption (session_store/crypto.py), a SessionStore protocol + PostgresSessionStore implementation over asyncpg (repository.py), and a minimal numbered-SQL-file migration runner (migrator.py) with the oauth_sessions table's first migration. Local dev/test Postgres via docker-compose.yml. New Config fields (postgres_dsn, session_encryption_key) redacted in __repr__ like every other secret field. Nothing wires into oauth.py/mcp.py yet -- SimpleOAuthProvider still uses the existing self-contained JWT session. That's Phase 2. Co-Authored-By: Claude Sonnet 5 --- docker-compose.yml | 19 ++ pyproject.toml | 3 +- src/keboola_mcp_server/config.py | 9 +- .../session_store/__init__.py | 6 + .../session_store/crypto.py | 63 ++++++ .../migrations/0001_oauth_sessions.sql | 29 +++ .../session_store/migrator.py | 50 +++++ .../session_store/repository.py | 206 ++++++++++++++++++ tests/session_store/__init__.py | 0 tests/session_store/conftest.py | 42 ++++ tests/session_store/test_crypto.py | 69 ++++++ tests/session_store/test_migrator.py | 44 ++++ tests/session_store/test_repository.py | 139 ++++++++++++ tests/test_config.py | 5 +- uv.lock | 72 +++++- 15 files changed, 750 insertions(+), 6 deletions(-) create mode 100644 docker-compose.yml create mode 100644 src/keboola_mcp_server/session_store/__init__.py create mode 100644 src/keboola_mcp_server/session_store/crypto.py create mode 100644 src/keboola_mcp_server/session_store/migrations/0001_oauth_sessions.sql create mode 100644 src/keboola_mcp_server/session_store/migrator.py create mode 100644 src/keboola_mcp_server/session_store/repository.py create mode 100644 tests/session_store/__init__.py create mode 100644 tests/session_store/conftest.py create mode 100644 tests/session_store/test_crypto.py create mode 100644 tests/session_store/test_migrator.py create mode 100644 tests/session_store/test_repository.py diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 000000000..91742a871 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,19 @@ +services: + postgres: + image: postgres:16 + environment: + POSTGRES_DB: keboola_mcp + POSTGRES_USER: keboola_mcp + POSTGRES_PASSWORD: keboola_mcp + ports: + - "5432:5432" + volumes: + - postgres_data:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U keboola_mcp"] + interval: 2s + timeout: 2s + retries: 30 + +volumes: + postgres_data: diff --git a/pyproject.toml b/pyproject.toml index 5e8c8429d..22548bf12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.75.3" +version = "1.76.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" @@ -20,6 +20,7 @@ dependencies = [ "pyjwt ~= 2.13", "json-log-formatter ~= 1.1", "cryptography ~= 49.0", + "asyncpg ~= 0.31", "pydantic ~= 2.13.0", "sqlglot ~= 30.0", "toon-format ~= 0.9.0b1", diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 0c36580e3..48606daa2 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -50,11 +50,16 @@ class Config: """The URL where the MCP server si reachable.""" jwt_secret: str | None = None """The secret key for encoding and decoding JWT tokens.""" + postgres_dsn: str | None = None + """Connection string for the Postgres-backed OAuth session store (oauth_session_persistence RFC). + Required to enable OAuth login when oauth_client_id/oauth_client_secret are set.""" + session_encryption_key: str | None = None + """Base64-encoded 32-byte AES-256 key used to encrypt OAuth session credentials at rest.""" bearer_token: str | None = None """The access-token issued by Keboola OAuth server to be sent in 'Authorization: Bearer ' header.""" conversation_id: str | None = None """The ID of the ongoing conversation with the MCP server. This is supplied only by the HTTP header.""" - project_id: Optional[str] = field(default=None, metadata={'aliases': ['kbc_project_id']}) + project_id: str | None = field(default=None, metadata={'aliases': ['kbc_project_id']}) """Project id used to scope a programmatic-token (kbc_at_/kbc_pat_) exchange. Maps the `X-KBC-ProjectId` HTTP header (via the alias) and the `KBC_PROJECT_ID` env var. @@ -143,7 +148,7 @@ def __repr__(self) -> str: for f in dataclasses.fields(self): value = getattr(self, f.name) if value: - if 'token' in f.name or 'password' in f.name or 'secret' in f.name: + if any(kw in f.name for kw in ('token', 'password', 'secret', 'key', 'dsn')): params.append(f"{f.name}='****'") else: if isinstance(value, str): diff --git a/src/keboola_mcp_server/session_store/__init__.py b/src/keboola_mcp_server/session_store/__init__.py new file mode 100644 index 000000000..0ede3402c --- /dev/null +++ b/src/keboola_mcp_server/session_store/__init__.py @@ -0,0 +1,6 @@ +"""Postgres-backed OAuth session storage (PSGO-261, oauth_session_persistence RFC). + +Replaces the self-contained OAuth access/refresh JWTs with an opaque, server-side session +reference: the MCP client holds only a random lookup key, never the real Keboola credentials. +See ``feature_spec/oauth_session_persistence/RFC.md`` for the design. +""" diff --git a/src/keboola_mcp_server/session_store/crypto.py b/src/keboola_mcp_server/session_store/crypto.py new file mode 100644 index 000000000..6696600e2 --- /dev/null +++ b/src/keboola_mcp_server/session_store/crypto.py @@ -0,0 +1,63 @@ +"""AES-256-GCM encryption for session data at rest (oauth_session_persistence RFC). + +GCM is authenticated encryption: tampering with the ciphertext (or decrypting with the wrong +key) raises ``InvalidTag`` rather than silently returning garbage plaintext. + +Ciphertext layout: ````. The +key-version prefix exists so a future key rotation can be introduced without a data migration +(RFC Open Question #2) -- v1 ships with exactly one supported version. +""" + +import base64 +import os +import secrets + +from cryptography.exceptions import InvalidTag +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +KEY_SIZE = 32 # AES-256 +_NONCE_SIZE = 12 # 96-bit GCM nonce, standard choice +_KEY_VERSION = 1 + +# Process-local fallback key for local dev/tests when KBC_SESSION_ENCRYPTION_KEY is unset. +# Mirrors mcp.py's _FALLBACK_SCOPE_SECRET: fine for a throwaway local Postgres, useless across a +# process restart -- a real deployment must set the env var. +_FALLBACK_KEY = secrets.token_bytes(KEY_SIZE) + + +class DecryptionError(Exception): + """Raised when ciphertext fails authentication (wrong key, corruption, or tampering).""" + + +def encrypt(plaintext: bytes, key: bytes) -> bytes: + if len(key) != KEY_SIZE: + raise ValueError(f'Encryption key must be {KEY_SIZE} bytes, got {len(key)}.') + nonce = os.urandom(_NONCE_SIZE) + ciphertext = AESGCM(key).encrypt(nonce, plaintext, None) + return bytes([_KEY_VERSION]) + nonce + ciphertext + + +def decrypt(blob: bytes, key: bytes) -> bytes: + if len(key) != KEY_SIZE: + raise ValueError(f'Encryption key must be {KEY_SIZE} bytes, got {len(key)}.') + if not blob or blob[0] != _KEY_VERSION: + raise DecryptionError(f'Unsupported or missing key version in ciphertext: {blob[:1]!r}.') + nonce, ciphertext = blob[1 : 1 + _NONCE_SIZE], blob[1 + _NONCE_SIZE :] + try: + return AESGCM(key).decrypt(nonce, ciphertext, None) + except InvalidTag as e: + raise DecryptionError('Ciphertext failed authentication (wrong key or tampered data).') from e + + +def resolve_encryption_key(session_encryption_key: str | None) -> bytes: + """Decodes the base64-encoded ``KBC_SESSION_ENCRYPTION_KEY``, or falls back to a process-local + key when unset (local dev/tests only -- see module docstring).""" + if not session_encryption_key: + return _FALLBACK_KEY + try: + key = base64.b64decode(session_encryption_key, validate=True) + except Exception as e: + raise ValueError('KBC_SESSION_ENCRYPTION_KEY is not valid base64.') from e + if len(key) != KEY_SIZE: + raise ValueError(f'KBC_SESSION_ENCRYPTION_KEY must decode to {KEY_SIZE} bytes, got {len(key)}.') + return key diff --git a/src/keboola_mcp_server/session_store/migrations/0001_oauth_sessions.sql b/src/keboola_mcp_server/session_store/migrations/0001_oauth_sessions.sql new file mode 100644 index 000000000..8686e365f --- /dev/null +++ b/src/keboola_mcp_server/session_store/migrations/0001_oauth_sessions.sql @@ -0,0 +1,29 @@ +-- oauth_session_persistence RFC: one row per OAuth-authenticated MCP session. +-- Real Keboola credentials are stored only as AES-256-GCM ciphertext (session_store/crypto.py); +-- the MCP client holds only the opaque access/refresh token, whose SHA-256 hash is looked up here. + +CREATE TABLE oauth_sessions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + access_token_hash BYTEA NOT NULL, + refresh_token_hash BYTEA, + client_id TEXT NOT NULL, + user_email TEXT, + kbc_access_token_enc BYTEA NOT NULL, + kbc_refresh_token_enc BYTEA NOT NULL, + kbc_access_expires_at TIMESTAMPTZ NOT NULL, + scope_project_ids INTEGER[], + scope_read_only BOOLEAN NOT NULL DEFAULT FALSE, + scope_confirmed BOOLEAN NOT NULL DEFAULT FALSE, + scope_scoped_token_enc BYTEA, + scope_scoped_expires_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + last_used_at TIMESTAMPTZ NOT NULL DEFAULT now(), + revoked_at TIMESTAMPTZ +); + +-- Unique so a hash collision (practically impossible with SHA-256) can't silently pick the +-- wrong session; also serves as the lookup index for the two access paths. +CREATE UNIQUE INDEX oauth_sessions_access_token_hash_idx ON oauth_sessions (access_token_hash); +CREATE UNIQUE INDEX oauth_sessions_refresh_token_hash_idx ON oauth_sessions (refresh_token_hash) + WHERE refresh_token_hash IS NOT NULL; diff --git a/src/keboola_mcp_server/session_store/migrator.py b/src/keboola_mcp_server/session_store/migrator.py new file mode 100644 index 000000000..07271f230 --- /dev/null +++ b/src/keboola_mcp_server/session_store/migrator.py @@ -0,0 +1,50 @@ +"""Tiny numbered-SQL-file migration runner. + +One table doesn't earn a migration framework (alembic, etc.) -- this is the whole mechanism: +numbered ``.sql`` files applied in order, tracked in ``schema_migrations`` so re-running is a +no-op. Not general-purpose (no down-migrations, no branching) by design. +""" + +import logging +from importlib import resources +from typing import cast + +import asyncpg + +LOG = logging.getLogger(__name__) + +_CREATE_TRACKING_TABLE = """ +CREATE TABLE IF NOT EXISTS schema_migrations ( + filename TEXT PRIMARY KEY, + applied_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +""" + + +def _migration_files() -> list[tuple[str, str]]: + """Returns (filename, sql) pairs for every ``*.sql`` file in this package's migrations/ dir, + sorted by filename -- the numeric prefix (``0001_...``) is what defines application order.""" + migrations_dir = resources.files(__package__) / 'migrations' + files = sorted(p for p in migrations_dir.iterdir() if p.name.endswith('.sql')) + return [(p.name, p.read_text()) for p in files] + + +async def apply_migrations(pool: asyncpg.Pool) -> list[str]: + """Applies every not-yet-applied migration file, in order, each in its own transaction. + + :return: filenames actually applied (empty if the schema was already up to date). + """ + applied: list[str] = [] + async with pool.acquire() as conn: + conn = cast(asyncpg.Connection, conn) + await conn.execute(_CREATE_TRACKING_TABLE) + already_applied = {r['filename'] for r in await conn.fetch('SELECT filename FROM schema_migrations')} + for filename, sql in _migration_files(): + if filename in already_applied: + continue + async with conn.transaction(): + await conn.execute(sql) + await conn.execute('INSERT INTO schema_migrations (filename) VALUES ($1)', filename) + LOG.info(f'Applied migration: {filename}') + applied.append(filename) + return applied diff --git a/src/keboola_mcp_server/session_store/repository.py b/src/keboola_mcp_server/session_store/repository.py new file mode 100644 index 000000000..033004ed3 --- /dev/null +++ b/src/keboola_mcp_server/session_store/repository.py @@ -0,0 +1,206 @@ +"""``SessionStore`` protocol + Postgres implementation (oauth_session_persistence RFC). + +The protocol exists so `oauth.py`/`mcp.py` logic can be unit-tested against an in-memory fake +without a real database; `PostgresSessionStore` is the only production implementation. +""" + +import dataclasses +import hashlib +import secrets +from datetime import datetime +from typing import Protocol + +import asyncpg + +from keboola_mcp_server.session_store import crypto +from keboola_mcp_server.session_store.migrator import apply_migrations + +# Length of the opaque, randomly-generated access/refresh tokens handed to the MCP client. 256 +# bits: not guessable, and this is the *entire* security check for these tokens (no signature to +# verify) -- see repository/RFC for why that's sufficient once the real credential lives server-side. +_TOKEN_BYTES = 32 + + +def generate_opaque_token() -> str: + return secrets.token_urlsafe(_TOKEN_BYTES) + + +def _hash_token(token: str) -> bytes: + return hashlib.sha256(token.encode('utf-8')).digest() + + +@dataclasses.dataclass(frozen=True) +class OAuthSession: + id: str + client_id: str + user_email: str | None + kbc_access_token: str + kbc_refresh_token: str + kbc_access_expires_at: datetime + scope_project_ids: list[int] | None + scope_read_only: bool + scope_confirmed: bool + scope_scoped_token: str | None + scope_scoped_expires_at: datetime | None + + +class SessionStore(Protocol): + async def create( + self, + *, + client_id: str, + user_email: str | None, + kbc_access_token: str, + kbc_refresh_token: str, + kbc_access_expires_at: datetime, + ) -> tuple[str, str, OAuthSession]: + """Creates a session row. Returns (opaque_access_token, opaque_refresh_token, session).""" + ... + + async def get_by_access_token(self, access_token: str) -> OAuthSession | None: + """None if the token doesn't exist, is revoked, or its underlying row is gone.""" + ... + + async def get_by_refresh_token(self, refresh_token: str) -> OAuthSession | None: ... + + async def rotate_kbc_tokens( + self, session_id: str, *, kbc_access_token: str, kbc_refresh_token: str, kbc_access_expires_at: datetime + ) -> None: + """Replaces the encrypted Keboola credentials in place (server-managed refresh).""" + ... + + async def update_scope( + self, + session_id: str, + *, + project_ids: list[int], + read_only: bool, + confirmed: bool, + scoped_token: str | None, + scoped_expires_at: datetime | None, + ) -> None: ... + + async def revoke(self, session_id: str) -> None: ... + + +class PostgresSessionStore: + def __init__(self, pool: asyncpg.Pool, encryption_key: bytes) -> None: + self._pool = pool + self._key = encryption_key + + @classmethod + async def connect(cls, dsn: str, encryption_key: bytes) -> 'PostgresSessionStore': + pool = await asyncpg.create_pool(dsn) + await apply_migrations(pool) + return cls(pool, encryption_key) + + async def close(self) -> None: + await self._pool.close() + + def _to_session(self, row: asyncpg.Record) -> OAuthSession: + return OAuthSession( + id=str(row['id']), + client_id=row['client_id'], + user_email=row['user_email'], + kbc_access_token=crypto.decrypt(row['kbc_access_token_enc'], self._key).decode('utf-8'), + kbc_refresh_token=crypto.decrypt(row['kbc_refresh_token_enc'], self._key).decode('utf-8'), + kbc_access_expires_at=row['kbc_access_expires_at'], + scope_project_ids=list(row['scope_project_ids']) if row['scope_project_ids'] is not None else None, + scope_read_only=row['scope_read_only'], + scope_confirmed=row['scope_confirmed'], + scope_scoped_token=( + crypto.decrypt(row['scope_scoped_token_enc'], self._key).decode('utf-8') + if row['scope_scoped_token_enc'] is not None + else None + ), + scope_scoped_expires_at=row['scope_scoped_expires_at'], + ) + + async def create( + self, + *, + client_id: str, + user_email: str | None, + kbc_access_token: str, + kbc_refresh_token: str, + kbc_access_expires_at: datetime, + ) -> tuple[str, str, OAuthSession]: + access_token = generate_opaque_token() + refresh_token = generate_opaque_token() + row = await self._pool.fetchrow( + """ + INSERT INTO oauth_sessions ( + access_token_hash, refresh_token_hash, client_id, user_email, + kbc_access_token_enc, kbc_refresh_token_enc, kbc_access_expires_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7) + RETURNING * + """, + _hash_token(access_token), + _hash_token(refresh_token), + client_id, + user_email, + crypto.encrypt(kbc_access_token.encode('utf-8'), self._key), + crypto.encrypt(kbc_refresh_token.encode('utf-8'), self._key), + kbc_access_expires_at, + ) + assert row is not None + return access_token, refresh_token, self._to_session(row) + + async def get_by_access_token(self, access_token: str) -> OAuthSession | None: + row = await self._pool.fetchrow( + 'UPDATE oauth_sessions SET last_used_at = now() ' + 'WHERE access_token_hash = $1 AND revoked_at IS NULL RETURNING *', + _hash_token(access_token), + ) + return self._to_session(row) if row is not None else None + + async def get_by_refresh_token(self, refresh_token: str) -> OAuthSession | None: + row = await self._pool.fetchrow( + 'SELECT * FROM oauth_sessions WHERE refresh_token_hash = $1 AND revoked_at IS NULL', + _hash_token(refresh_token), + ) + return self._to_session(row) if row is not None else None + + async def rotate_kbc_tokens( + self, session_id: str, *, kbc_access_token: str, kbc_refresh_token: str, kbc_access_expires_at: datetime + ) -> None: + await self._pool.execute( + """ + UPDATE oauth_sessions + SET kbc_access_token_enc = $2, kbc_refresh_token_enc = $3, kbc_access_expires_at = $4, + updated_at = now() + WHERE id = $1 + """, + session_id, + crypto.encrypt(kbc_access_token.encode('utf-8'), self._key), + crypto.encrypt(kbc_refresh_token.encode('utf-8'), self._key), + kbc_access_expires_at, + ) + + async def update_scope( + self, + session_id: str, + *, + project_ids: list[int], + read_only: bool, + confirmed: bool, + scoped_token: str | None, + scoped_expires_at: datetime | None, + ) -> None: + await self._pool.execute( + """ + UPDATE oauth_sessions + SET scope_project_ids = $2, scope_read_only = $3, scope_confirmed = $4, + scope_scoped_token_enc = $5, scope_scoped_expires_at = $6, updated_at = now() + WHERE id = $1 + """, + session_id, + project_ids, + read_only, + confirmed, + crypto.encrypt(scoped_token.encode('utf-8'), self._key) if scoped_token is not None else None, + scoped_expires_at, + ) + + async def revoke(self, session_id: str) -> None: + await self._pool.execute('UPDATE oauth_sessions SET revoked_at = now() WHERE id = $1', session_id) diff --git a/tests/session_store/__init__.py b/tests/session_store/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/session_store/conftest.py b/tests/session_store/conftest.py new file mode 100644 index 000000000..5b5b946e4 --- /dev/null +++ b/tests/session_store/conftest.py @@ -0,0 +1,42 @@ +import os + +import asyncpg +import pytest +import pytest_asyncio + +from keboola_mcp_server.session_store.crypto import KEY_SIZE +from keboola_mcp_server.session_store.repository import PostgresSessionStore + +TEST_DSN = os.environ.get('KBC_TEST_POSTGRES_DSN', 'postgresql://keboola_mcp:keboola_mcp@localhost:5432/keboola_mcp') + + +def _postgres_available() -> bool: + import socket + from urllib.parse import urlparse + + parsed = urlparse(TEST_DSN) + try: + with socket.create_connection((parsed.hostname, parsed.port or 5432), timeout=0.5): + return True + except OSError: + return False + + +requires_postgres = pytest.mark.skipif( + not _postgres_available(), reason=f'No Postgres reachable at {TEST_DSN} (see docker-compose.yml)' +) + + +@pytest_asyncio.fixture +async def store(): + pool = await asyncpg.create_pool(TEST_DSN) + try: + # Clean slate per test: drop and let apply_migrations (inside connect) recreate. + await pool.execute('DROP TABLE IF EXISTS oauth_sessions, schema_migrations CASCADE') + finally: + await pool.close() + s = await PostgresSessionStore.connect(TEST_DSN, encryption_key=bytes([1] * KEY_SIZE)) + try: + yield s + finally: + await s.close() diff --git a/tests/session_store/test_crypto.py b/tests/session_store/test_crypto.py new file mode 100644 index 000000000..17dbcd021 --- /dev/null +++ b/tests/session_store/test_crypto.py @@ -0,0 +1,69 @@ +import pytest + +from keboola_mcp_server.session_store.crypto import ( + KEY_SIZE, + DecryptionError, + decrypt, + encrypt, + resolve_encryption_key, +) + + +def _key(byte: int = 1) -> bytes: + return bytes([byte]) * KEY_SIZE + + +def test_round_trip() -> None: + ciphertext = encrypt(b'kbc_at_secret', _key()) + assert decrypt(ciphertext, _key()) == b'kbc_at_secret' + + +def test_ciphertext_differs_each_call() -> None: + # Random nonce per call -- same plaintext must not produce identical ciphertext. + assert encrypt(b'same plaintext', _key()) != encrypt(b'same plaintext', _key()) + + +def test_wrong_key_fails() -> None: + ciphertext = encrypt(b'kbc_at_secret', _key(1)) + with pytest.raises(DecryptionError): + decrypt(ciphertext, _key(2)) + + +def test_tampered_ciphertext_fails() -> None: + ciphertext = bytearray(encrypt(b'kbc_at_secret', _key())) + ciphertext[-1] ^= 0xFF # flip a bit in the GCM tag/ciphertext + with pytest.raises(DecryptionError): + decrypt(bytes(ciphertext), _key()) + + +def test_wrong_key_version_fails() -> None: + ciphertext = bytearray(encrypt(b'kbc_at_secret', _key())) + ciphertext[0] = 99 + with pytest.raises(DecryptionError, match='key version'): + decrypt(bytes(ciphertext), _key()) + + +@pytest.mark.parametrize('key_len', [16, 31, 33]) +def test_rejects_wrong_key_length(key_len: int) -> None: + with pytest.raises(ValueError, match='32 bytes'): + encrypt(b'x', bytes(key_len)) + + +def test_resolve_encryption_key_decodes_base64() -> None: + import base64 + + raw = _key(7) + assert resolve_encryption_key(base64.b64encode(raw).decode()) == raw + + +def test_resolve_encryption_key_falls_back_when_unset() -> None: + key = resolve_encryption_key(None) + assert len(key) == KEY_SIZE + # Stable within the process (same fallback reused, not regenerated per call). + assert resolve_encryption_key(None) == key + + +@pytest.mark.parametrize('bad_value', ['not-base64!!!', 'aGVsbG8=']) # valid base64, wrong length +def test_resolve_encryption_key_rejects_invalid_input(bad_value: str) -> None: + with pytest.raises(ValueError, match='.+'): + resolve_encryption_key(bad_value) diff --git a/tests/session_store/test_migrator.py b/tests/session_store/test_migrator.py new file mode 100644 index 000000000..d8018cd2d --- /dev/null +++ b/tests/session_store/test_migrator.py @@ -0,0 +1,44 @@ +import asyncpg +import pytest +import pytest_asyncio + +from keboola_mcp_server.session_store.migrator import apply_migrations +from tests.session_store.conftest import TEST_DSN, requires_postgres + +pytestmark = [pytest.mark.asyncio, requires_postgres] + + +@pytest_asyncio.fixture(autouse=True) +async def _clean_slate(): + pool = await asyncpg.create_pool(TEST_DSN) + try: + await pool.execute('DROP TABLE IF EXISTS oauth_sessions, schema_migrations CASCADE') + finally: + await pool.close() + + +async def test_applies_migrations_once() -> None: + pool = await asyncpg.create_pool(TEST_DSN) + try: + applied = await apply_migrations(pool) + assert applied == ['0001_oauth_sessions.sql'] + + # Re-running is a no-op -- the table already exists, so re-applying the DDL would fail + # if the tracking table didn't correctly skip it. + applied_again = await apply_migrations(pool) + assert applied_again == [] + finally: + await pool.close() + + +async def test_creates_oauth_sessions_table() -> None: + pool = await asyncpg.create_pool(TEST_DSN) + try: + await apply_migrations(pool) + columns = await pool.fetch( + "SELECT column_name FROM information_schema.columns WHERE table_name = 'oauth_sessions'" + ) + names = {r['column_name'] for r in columns} + assert {'access_token_hash', 'kbc_access_token_enc', 'scope_project_ids', 'revoked_at'} <= names + finally: + await pool.close() diff --git a/tests/session_store/test_repository.py b/tests/session_store/test_repository.py new file mode 100644 index 000000000..50afa56ae --- /dev/null +++ b/tests/session_store/test_repository.py @@ -0,0 +1,139 @@ +from datetime import datetime, timedelta, timezone + +import pytest + +from tests.session_store.conftest import requires_postgres + +pytestmark = [pytest.mark.asyncio, requires_postgres] + + +async def test_create_and_get_by_access_token(store) -> None: + expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + access_token, refresh_token, session = await store.create( + client_id='claude.ai', + user_email='m@k.com', + kbc_access_token='kbc_at_secret', + kbc_refresh_token='kbc_rt_secret', + kbc_access_expires_at=expires_at, + ) + + assert access_token + assert refresh_token + assert access_token != refresh_token + assert session.client_id == 'claude.ai' + assert session.kbc_access_token == 'kbc_at_secret' + assert session.kbc_refresh_token == 'kbc_rt_secret' + assert session.scope_confirmed is False + assert session.scope_project_ids is None + + fetched = await store.get_by_access_token(access_token) + assert fetched is not None + assert fetched.id == session.id + assert fetched.kbc_access_token == 'kbc_at_secret' + + +async def test_get_by_access_token_unknown_returns_none(store) -> None: + assert await store.get_by_access_token('does-not-exist') is None + + +async def test_get_by_refresh_token(store) -> None: + _, refresh_token, session = await store.create( + client_id='claude.ai', + user_email=None, + kbc_access_token='kbc_at_x', + kbc_refresh_token='kbc_rt_x', + kbc_access_expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + fetched = await store.get_by_refresh_token(refresh_token) + assert fetched is not None + assert fetched.id == session.id + + +async def test_rotate_kbc_tokens_replaces_credentials(store) -> None: + access_token, _, session = await store.create( + client_id='claude.ai', + user_email=None, + kbc_access_token='kbc_at_old', + kbc_refresh_token='kbc_rt_old', + kbc_access_expires_at=datetime.now(timezone.utc) - timedelta(seconds=1), + ) + new_expiry = datetime.now(timezone.utc) + timedelta(hours=1) + await store.rotate_kbc_tokens( + session.id, kbc_access_token='kbc_at_new', kbc_refresh_token='kbc_rt_new', kbc_access_expires_at=new_expiry + ) + + fetched = await store.get_by_access_token(access_token) + assert fetched is not None + assert fetched.kbc_access_token == 'kbc_at_new' + assert fetched.kbc_refresh_token == 'kbc_rt_new' + + +async def test_update_scope_confirms_project_selection(store) -> None: + access_token, _, session = await store.create( + client_id='claude.ai', + user_email=None, + kbc_access_token='kbc_at_x', + kbc_refresh_token='kbc_rt_x', + kbc_access_expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + scoped_expiry = datetime.now(timezone.utc) + timedelta(minutes=30) + await store.update_scope( + session.id, + project_ids=[18, 83], + read_only=True, + confirmed=True, + scoped_token='kbc_pat_scoped', + scoped_expires_at=scoped_expiry, + ) + + fetched = await store.get_by_access_token(access_token) + assert fetched is not None + assert fetched.scope_project_ids == [18, 83] + assert fetched.scope_read_only is True + assert fetched.scope_confirmed is True + assert fetched.scope_scoped_token == 'kbc_pat_scoped' + + +async def test_update_scope_without_scoped_token(store) -> None: + # The whole-stack fallback path (resolver exchange unavailable) confirms scope with no minted + # token -- must not choke on a None scoped_token. + access_token, _, session = await store.create( + client_id='claude.ai', + user_email=None, + kbc_access_token='kbc_at_x', + kbc_refresh_token='kbc_rt_x', + kbc_access_expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + await store.update_scope( + session.id, project_ids=[18], read_only=False, confirmed=True, scoped_token=None, scoped_expires_at=None + ) + fetched = await store.get_by_access_token(access_token) + assert fetched is not None + assert fetched.scope_scoped_token is None + + +async def test_revoke_makes_session_unreachable(store) -> None: + access_token, refresh_token, session = await store.create( + client_id='claude.ai', + user_email=None, + kbc_access_token='kbc_at_x', + kbc_refresh_token='kbc_rt_x', + kbc_access_expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + await store.revoke(session.id) + + assert await store.get_by_access_token(access_token) is None + assert await store.get_by_refresh_token(refresh_token) is None + + +async def test_credentials_are_encrypted_at_rest(store) -> None: + # Read the raw row directly -- the plaintext secret must never appear in storage. + _, _, session = await store.create( + client_id='claude.ai', + user_email=None, + kbc_access_token='kbc_at_should_not_appear_in_plaintext', + kbc_refresh_token='kbc_rt_x', + kbc_access_expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + raw = await store._pool.fetchrow('SELECT kbc_access_token_enc FROM oauth_sessions WHERE id = $1', session.id) + assert b'kbc_at_should_not_appear_in_plaintext' not in raw['kbc_access_token_enc'] diff --git a/tests/test_config.py b/tests/test_config.py index 249d5ec1e..fb738d151 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -86,12 +86,13 @@ def test_defaults(self) -> None: assert getattr(config, f.name) is None, f'Expected default value for {f.name} to be None' def test_no_token_password_in_repr(self) -> None: - config = Config(storage_token='foo') + config = Config(storage_token='foo', postgres_dsn='postgresql://u:p@host/db', session_encryption_key='abc') assert str(config) == ( "Config(storage_api_url=None, storage_token='****', branch_id=None, workspace_schema=None, " 'oauth_client_id=None, oauth_client_secret=None, ' 'oauth_server_url=None, oauth_scope=None, mcp_server_url=None, ' - 'jwt_secret=None, bearer_token=None, conversation_id=None, project_id=None)' + "jwt_secret=None, postgres_dsn='****', session_encryption_key='****', " + 'bearer_token=None, conversation_id=None, project_id=None)' ) @pytest.mark.parametrize( diff --git a/uv.lock b/uv.lock index 2c80fe7ec..33e2a2f87 100644 --- a/uv.lock +++ b/uv.lock @@ -71,6 +71,74 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, ] +[[package]] +name = "async-timeout" +version = "5.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/ae/136395dfbfe00dfc94da3f3e136d0b13f394cba8f4841120e34226265780/async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3", size = 9274, upload-time = "2024-11-06T16:41:39.6Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/ba/e2081de779ca30d473f21f5b30e0e737c438205440784c7dfc81efc2b029/async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c", size = 6233, upload-time = "2024-11-06T16:41:37.9Z" }, +] + +[[package]] +name = "asyncpg" +version = "0.31.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "async-timeout", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fe/cc/d18065ce2380d80b1bcce927c24a2642efd38918e33fd724bc4bca904877/asyncpg-0.31.0.tar.gz", hash = "sha256:c989386c83940bfbd787180f2b1519415e2d3d6277a70d9d0f0145ac73500735", size = 993667, upload-time = "2025-11-24T23:27:00.812Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/d9/507c80bdac2e95e5a525644af94b03fa7f9a44596a84bd48a6e80f854f92/asyncpg-0.31.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:831712dd3cf117eec68575a9b50da711893fd63ebe277fc155ecae1c6c9f0f61", size = 644865, upload-time = "2025-11-24T23:25:23.527Z" }, + { url = "https://files.pythonhosted.org/packages/ea/03/f93b5e543f65c5f504e91405e8d21bb9e600548be95032951a754781a41d/asyncpg-0.31.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0b17c89312c2f4ccea222a3a6571f7df65d4ba2c0e803339bfc7bed46a96d3be", size = 639297, upload-time = "2025-11-24T23:25:25.192Z" }, + { url = "https://files.pythonhosted.org/packages/e5/1e/de2177e57e03a06e697f6c1ddf2a9a7fcfdc236ce69966f54ffc830fd481/asyncpg-0.31.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3faa62f997db0c9add34504a68ac2c342cfee4d57a0c3062fcf0d86c7f9cb1e8", size = 2816679, upload-time = "2025-11-24T23:25:26.718Z" }, + { url = "https://files.pythonhosted.org/packages/d0/98/1a853f6870ac7ad48383a948c8ff3c85dc278066a4d69fc9af7d3d4b1106/asyncpg-0.31.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ea599d45c361dfbf398cb67da7fd052affa556a401482d3ff1ee99bd68808a1", size = 2867087, upload-time = "2025-11-24T23:25:28.399Z" }, + { url = "https://files.pythonhosted.org/packages/11/29/7e76f2a51f2360a7c90d2cf6d0d9b210c8bb0ae342edebd16173611a55c2/asyncpg-0.31.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:795416369c3d284e1837461909f58418ad22b305f955e625a4b3a2521d80a5f3", size = 2747631, upload-time = "2025-11-24T23:25:30.154Z" }, + { url = "https://files.pythonhosted.org/packages/5d/3f/716e10cb57c4f388248db46555e9226901688fbfabd0afb85b5e1d65d5a7/asyncpg-0.31.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:a8d758dac9d2e723e173d286ef5e574f0b350ec00e9186fce84d0fc5f6a8e6b8", size = 2855107, upload-time = "2025-11-24T23:25:31.888Z" }, + { url = "https://files.pythonhosted.org/packages/7e/ec/3ebae9dfb23a1bd3f68acfd4f795983b65b413291c0e2b0d982d6ae6c920/asyncpg-0.31.0-cp310-cp310-win32.whl", hash = "sha256:2d076d42eb583601179efa246c5d7ae44614b4144bc1c7a683ad1222814ed095", size = 521990, upload-time = "2025-11-24T23:25:33.402Z" }, + { url = "https://files.pythonhosted.org/packages/20/b4/9fbb4b0af4e36d96a61d026dd37acab3cf521a70290a09640b215da5ab7c/asyncpg-0.31.0-cp310-cp310-win_amd64.whl", hash = "sha256:9ea33213ac044171f4cac23740bed9a3805abae10e7025314cfbd725ec670540", size = 581629, upload-time = "2025-11-24T23:25:34.846Z" }, + { url = "https://files.pythonhosted.org/packages/08/17/cc02bc49bc350623d050fa139e34ea512cd6e020562f2a7312a7bcae4bc9/asyncpg-0.31.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:eee690960e8ab85063ba93af2ce128c0f52fd655fdff9fdb1a28df01329f031d", size = 643159, upload-time = "2025-11-24T23:25:36.443Z" }, + { url = "https://files.pythonhosted.org/packages/a4/62/4ded7d400a7b651adf06f49ea8f73100cca07c6df012119594d1e3447aa6/asyncpg-0.31.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2657204552b75f8288de08ca60faf4a99a65deef3a71d1467454123205a88fab", size = 638157, upload-time = "2025-11-24T23:25:37.89Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5b/4179538a9a72166a0bf60ad783b1ef16efb7960e4d7b9afe9f77a5551680/asyncpg-0.31.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a429e842a3a4b4ea240ea52d7fe3f82d5149853249306f7ff166cb9948faa46c", size = 2918051, upload-time = "2025-11-24T23:25:39.461Z" }, + { url = "https://files.pythonhosted.org/packages/e6/35/c27719ae0536c5b6e61e4701391ffe435ef59539e9360959240d6e47c8c8/asyncpg-0.31.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c0807be46c32c963ae40d329b3a686356e417f674c976c07fa49f1b30303f109", size = 2972640, upload-time = "2025-11-24T23:25:41.512Z" }, + { url = "https://files.pythonhosted.org/packages/43/f4/01ebb9207f29e645a64699b9ce0eefeff8e7a33494e1d29bb53736f7766b/asyncpg-0.31.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e5d5098f63beeae93512ee513d4c0c53dc12e9aa2b7a1af5a81cddf93fe4e4da", size = 2851050, upload-time = "2025-11-24T23:25:43.153Z" }, + { url = "https://files.pythonhosted.org/packages/3e/f4/03ff1426acc87be0f4e8d40fa2bff5c3952bef0080062af9efc2212e3be8/asyncpg-0.31.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:37fc6c00a814e18eef51833545d1891cac9aa69140598bb076b4cd29b3e010b9", size = 2962574, upload-time = "2025-11-24T23:25:44.942Z" }, + { url = "https://files.pythonhosted.org/packages/c7/39/cc788dfca3d4060f9d93e67be396ceec458dfc429e26139059e58c2c244d/asyncpg-0.31.0-cp311-cp311-win32.whl", hash = "sha256:5a4af56edf82a701aece93190cc4e094d2df7d33f6e915c222fb09efbb5afc24", size = 521076, upload-time = "2025-11-24T23:25:46.486Z" }, + { url = "https://files.pythonhosted.org/packages/28/fc/735af5384c029eb7f1ca60ccb8fa95521dbdaeef788edf4cecfc604c3cab/asyncpg-0.31.0-cp311-cp311-win_amd64.whl", hash = "sha256:480c4befbdf079c14c9ca43c8c5e1fe8b6296c96f1f927158d4f1e750aacc047", size = 584980, upload-time = "2025-11-24T23:25:47.938Z" }, + { url = "https://files.pythonhosted.org/packages/2a/a6/59d0a146e61d20e18db7396583242e32e0f120693b67a8de43f1557033e2/asyncpg-0.31.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b44c31e1efc1c15188ef183f287c728e2046abb1d26af4d20858215d50d91fad", size = 662042, upload-time = "2025-11-24T23:25:49.578Z" }, + { url = "https://files.pythonhosted.org/packages/36/01/ffaa189dcb63a2471720615e60185c3f6327716fdc0fc04334436fbb7c65/asyncpg-0.31.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0c89ccf741c067614c9b5fc7f1fc6f3b61ab05ae4aaa966e6fd6b93097c7d20d", size = 638504, upload-time = "2025-11-24T23:25:51.501Z" }, + { url = "https://files.pythonhosted.org/packages/9f/62/3f699ba45d8bd24c5d65392190d19656d74ff0185f42e19d0bbd973bb371/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:12b3b2e39dc5470abd5e98c8d3373e4b1d1234d9fbdedf538798b2c13c64460a", size = 3426241, upload-time = "2025-11-24T23:25:53.278Z" }, + { url = "https://files.pythonhosted.org/packages/8c/d1/a867c2150f9c6e7af6462637f613ba67f78a314b00db220cd26ff559d532/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:aad7a33913fb8bcb5454313377cc330fbb19a0cd5faa7272407d8a0c4257b671", size = 3520321, upload-time = "2025-11-24T23:25:54.982Z" }, + { url = "https://files.pythonhosted.org/packages/7a/1a/cce4c3f246805ecd285a3591222a2611141f1669d002163abef999b60f98/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3df118d94f46d85b2e434fd62c84cb66d5834d5a890725fe625f498e72e4d5ec", size = 3316685, upload-time = "2025-11-24T23:25:57.43Z" }, + { url = "https://files.pythonhosted.org/packages/40/ae/0fc961179e78cc579e138fad6eb580448ecae64908f95b8cb8ee2f241f67/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bd5b6efff3c17c3202d4b37189969acf8927438a238c6257f66be3c426beba20", size = 3471858, upload-time = "2025-11-24T23:25:59.636Z" }, + { url = "https://files.pythonhosted.org/packages/52/b2/b20e09670be031afa4cbfabd645caece7f85ec62d69c312239de568e058e/asyncpg-0.31.0-cp312-cp312-win32.whl", hash = "sha256:027eaa61361ec735926566f995d959ade4796f6a49d3bde17e5134b9964f9ba8", size = 527852, upload-time = "2025-11-24T23:26:01.084Z" }, + { url = "https://files.pythonhosted.org/packages/b5/f0/f2ed1de154e15b107dc692262395b3c17fc34eafe2a78fc2115931561730/asyncpg-0.31.0-cp312-cp312-win_amd64.whl", hash = "sha256:72d6bdcbc93d608a1158f17932de2321f68b1a967a13e014998db87a72ed3186", size = 597175, upload-time = "2025-11-24T23:26:02.564Z" }, + { url = "https://files.pythonhosted.org/packages/95/11/97b5c2af72a5d0b9bc3fa30cd4b9ce22284a9a943a150fdc768763caf035/asyncpg-0.31.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c204fab1b91e08b0f47e90a75d1b3c62174dab21f670ad6c5d0f243a228f015b", size = 661111, upload-time = "2025-11-24T23:26:04.467Z" }, + { url = "https://files.pythonhosted.org/packages/1b/71/157d611c791a5e2d0423f09f027bd499935f0906e0c2a416ce712ba51ef3/asyncpg-0.31.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:54a64f91839ba59008eccf7aad2e93d6e3de688d796f35803235ea1c4898ae1e", size = 636928, upload-time = "2025-11-24T23:26:05.944Z" }, + { url = "https://files.pythonhosted.org/packages/2e/fc/9e3486fb2bbe69d4a867c0b76d68542650a7ff1574ca40e84c3111bb0c6e/asyncpg-0.31.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0e0822b1038dc7253b337b0f3f676cadc4ac31b126c5d42691c39691962e403", size = 3424067, upload-time = "2025-11-24T23:26:07.957Z" }, + { url = "https://files.pythonhosted.org/packages/12/c6/8c9d076f73f07f995013c791e018a1cd5f31823c2a3187fc8581706aa00f/asyncpg-0.31.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bef056aa502ee34204c161c72ca1f3c274917596877f825968368b2c33f585f4", size = 3518156, upload-time = "2025-11-24T23:26:09.591Z" }, + { url = "https://files.pythonhosted.org/packages/ae/3b/60683a0baf50fbc546499cfb53132cb6835b92b529a05f6a81471ab60d0c/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:0bfbcc5b7ffcd9b75ab1558f00db2ae07db9c80637ad1b2469c43df79d7a5ae2", size = 3319636, upload-time = "2025-11-24T23:26:11.168Z" }, + { url = "https://files.pythonhosted.org/packages/50/dc/8487df0f69bd398a61e1792b3cba0e47477f214eff085ba0efa7eac9ce87/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:22bc525ebbdc24d1261ecbf6f504998244d4e3be1721784b5f64664d61fbe602", size = 3472079, upload-time = "2025-11-24T23:26:13.164Z" }, + { url = "https://files.pythonhosted.org/packages/13/a1/c5bbeeb8531c05c89135cb8b28575ac2fac618bcb60119ee9696c3faf71c/asyncpg-0.31.0-cp313-cp313-win32.whl", hash = "sha256:f890de5e1e4f7e14023619399a471ce4b71f5418cd67a51853b9910fdfa73696", size = 527606, upload-time = "2025-11-24T23:26:14.78Z" }, + { url = "https://files.pythonhosted.org/packages/91/66/b25ccb84a246b470eb943b0107c07edcae51804912b824054b3413995a10/asyncpg-0.31.0-cp313-cp313-win_amd64.whl", hash = "sha256:dc5f2fa9916f292e5c5c8b2ac2813763bcd7f58e130055b4ad8a0531314201ab", size = 596569, upload-time = "2025-11-24T23:26:16.189Z" }, + { url = "https://files.pythonhosted.org/packages/3c/36/e9450d62e84a13aea6580c83a47a437f26c7ca6fa0f0fd40b6670793ea30/asyncpg-0.31.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:f6b56b91bb0ffc328c4e3ed113136cddd9deefdf5f79ab448598b9772831df44", size = 660867, upload-time = "2025-11-24T23:26:17.631Z" }, + { url = "https://files.pythonhosted.org/packages/82/4b/1d0a2b33b3102d210439338e1beea616a6122267c0df459ff0265cd5807a/asyncpg-0.31.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:334dec28cf20d7f5bb9e45b39546ddf247f8042a690bff9b9573d00086e69cb5", size = 638349, upload-time = "2025-11-24T23:26:19.689Z" }, + { url = "https://files.pythonhosted.org/packages/41/aa/e7f7ac9a7974f08eff9183e392b2d62516f90412686532d27e196c0f0eeb/asyncpg-0.31.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:98cc158c53f46de7bb677fd20c417e264fc02b36d901cc2a43bd6cb0dc6dbfd2", size = 3410428, upload-time = "2025-11-24T23:26:21.275Z" }, + { url = "https://files.pythonhosted.org/packages/6f/de/bf1b60de3dede5c2731e6788617a512bc0ebd9693eac297ee74086f101d7/asyncpg-0.31.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9322b563e2661a52e3cdbc93eed3be7748b289f792e0011cb2720d278b366ce2", size = 3471678, upload-time = "2025-11-24T23:26:23.627Z" }, + { url = "https://files.pythonhosted.org/packages/46/78/fc3ade003e22d8bd53aaf8f75f4be48f0b460fa73738f0391b9c856a9147/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19857a358fc811d82227449b7ca40afb46e75b33eb8897240c3839dd8b744218", size = 3313505, upload-time = "2025-11-24T23:26:25.235Z" }, + { url = "https://files.pythonhosted.org/packages/bf/e9/73eb8a6789e927816f4705291be21f2225687bfa97321e40cd23055e903a/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ba5f8886e850882ff2c2ace5732300e99193823e8107e2c53ef01c1ebfa1e85d", size = 3434744, upload-time = "2025-11-24T23:26:26.944Z" }, + { url = "https://files.pythonhosted.org/packages/08/4b/f10b880534413c65c5b5862f79b8e81553a8f364e5238832ad4c0af71b7f/asyncpg-0.31.0-cp314-cp314-win32.whl", hash = "sha256:cea3a0b2a14f95834cee29432e4ddc399b95700eb1d51bbc5bfee8f31fa07b2b", size = 532251, upload-time = "2025-11-24T23:26:28.404Z" }, + { url = "https://files.pythonhosted.org/packages/d3/2d/7aa40750b7a19efa5d66e67fc06008ca0f27ba1bd082e457ad82f59aba49/asyncpg-0.31.0-cp314-cp314-win_amd64.whl", hash = "sha256:04d19392716af6b029411a0264d92093b6e5e8285ae97a39957b9a9c14ea72be", size = 604901, upload-time = "2025-11-24T23:26:30.34Z" }, + { url = "https://files.pythonhosted.org/packages/ce/fe/b9dfe349b83b9dee28cc42360d2c86b2cdce4cb551a2c2d27e156bcac84d/asyncpg-0.31.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:bdb957706da132e982cc6856bb2f7b740603472b54c3ebc77fe60ea3e57e1bd2", size = 702280, upload-time = "2025-11-24T23:26:32Z" }, + { url = "https://files.pythonhosted.org/packages/6a/81/e6be6e37e560bd91e6c23ea8a6138a04fd057b08cf63d3c5055c98e81c1d/asyncpg-0.31.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6d11b198111a72f47154fa03b85799f9be63701e068b43f84ac25da0bda9cb31", size = 682931, upload-time = "2025-11-24T23:26:33.572Z" }, + { url = "https://files.pythonhosted.org/packages/a6/45/6009040da85a1648dd5bc75b3b0a062081c483e75a1a29041ae63a0bf0dc/asyncpg-0.31.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18c83b03bc0d1b23e6230f5bf8d4f217dc9bc08644ce0502a9d91dc9e634a9c7", size = 3581608, upload-time = "2025-11-24T23:26:35.638Z" }, + { url = "https://files.pythonhosted.org/packages/7e/06/2e3d4d7608b0b2b3adbee0d0bd6a2d29ca0fc4d8a78f8277df04e2d1fd7b/asyncpg-0.31.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e009abc333464ff18b8f6fd146addffd9aaf63e79aa3bb40ab7a4c332d0c5e9e", size = 3498738, upload-time = "2025-11-24T23:26:37.275Z" }, + { url = "https://files.pythonhosted.org/packages/7d/aa/7d75ede780033141c51d83577ea23236ba7d3a23593929b32b49db8ed36e/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3b1fbcb0e396a5ca435a8826a87e5c2c2cc0c8c68eb6fadf82168056b0e53a8c", size = 3401026, upload-time = "2025-11-24T23:26:39.423Z" }, + { url = "https://files.pythonhosted.org/packages/ba/7a/15e37d45e7f7c94facc1e9148c0e455e8f33c08f0b8a0b1deb2c5171771b/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:8df714dba348efcc162d2adf02d213e5fab1bd9f557e1305633e851a61814a7a", size = 3429426, upload-time = "2025-11-24T23:26:41.032Z" }, + { url = "https://files.pythonhosted.org/packages/13/d5/71437c5f6ae5f307828710efbe62163974e71237d5d46ebd2869ea052d10/asyncpg-0.31.0-cp314-cp314t-win32.whl", hash = "sha256:1b41f1afb1033f2b44f3234993b15096ddc9cd71b21a42dbd87fc6a57b43d65d", size = 614495, upload-time = "2025-11-24T23:26:42.659Z" }, + { url = "https://files.pythonhosted.org/packages/3c/d7/8fb3044eaef08a310acfe23dae9a8e2e07d305edc29a53497e52bc76eca7/asyncpg-0.31.0-cp314-cp314t-win_amd64.whl", hash = "sha256:bd4107bb7cdd0e9e65fae66a62afd3a249663b844fa34d479f6d5b3bef9c04c3", size = 706062, upload-time = "2025-11-24T23:26:44.086Z" }, +] + [[package]] name = "attrs" version = "26.1.0" @@ -1172,9 +1240,10 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.75.3" +version = "1.76.0" source = { editable = "." } dependencies = [ + { name = "asyncpg" }, { name = "cryptography" }, { name = "fastmcp" }, { name = "httpx" }, @@ -1213,6 +1282,7 @@ tests = [ [package.metadata] requires-dist = [ + { name = "asyncpg", specifier = "~=0.31" }, { name = "cryptography", specifier = "~=49.0" }, { name = "fastmcp", specifier = "==3.4.4" }, { name = "httpx", specifier = "~=0.28" }, From f214cade7811908d3905c48c2fcd901842433b7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 29 Jul 2026 10:25:36 +0200 Subject: [PATCH 60/89] fix(PSGO-261): alias postgres_dsn to MCP_DB_URL, matching this infra's naming convention Keeps the field named postgres_dsn internally (so __repr__'s dsn-keyword redaction still applies) while accepting MCP_DB_URL / KBC_MCP_DB_URL / KBC_POSTGRES_DSN as the env var, consistent with how other Config fields (e.g. storage_token/storage_api_token) already alias multiple accepted names. Co-Authored-By: Claude Sonnet 5 --- feature_spec/oauth_session_persistence/RFC.md | 10 +++++++++- src/keboola_mcp_server/config.py | 6 ++++-- tests/test_config.py | 12 ++++++++++++ 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/feature_spec/oauth_session_persistence/RFC.md b/feature_spec/oauth_session_persistence/RFC.md index fd9a059c4..c634bfed9 100644 --- a/feature_spec/oauth_session_persistence/RFC.md +++ b/feature_spec/oauth_session_persistence/RFC.md @@ -78,7 +78,7 @@ confidential) via the `cryptography` package, already a pinned dependency (`pypr `~= 49.0`) — no new crypto library needed, just a new small module (`session_store/crypto.py`) wrapping `cryptography.hazmat.primitives.ciphers.aead.AESGCM`. -### New env var +### New env vars `KBC_SESSION_ENCRYPTION_KEY` — 32 raw bytes, base64-encoded for env-var transport (`base64.b64decode` on load, fail loudly at startup if it doesn't decode to exactly 32 bytes). Single static key for v1 @@ -86,6 +86,14 @@ on load, fail loudly at startup if it doesn't decode to exactly 32 bytes). Singl (`config.jwt_secret`) — same "required in production, generate an ephemeral one locally if unset" posture, so local dev/tests work with zero setup. +`config.postgres_dsn` — the infra-facing env var is **`MCP_DB_URL`** (aliased, matching the naming +already used for other freshly-provisioned Postgres instances in this infra, e.g. +`mcp_docs_database_init[0].postgresql_url`); `KBC_MCP_DB_URL`/`KBC_POSTGRES_DSN` also work via the +same alias/prefix mechanism every other `Config` field already supports. A single connection-string +value (`postgresql://user:pass@host:port/dbname`), not split host/port/user/password env vars — +matches every other `Config` field (one env var, one value) and is exactly what +`asyncpg.create_pool(dsn)` wants directly, no glue needed. + ### Refresh strategy: lazy, on lookup — not a background job Every session lookup (equivalent to today's `load_access_token()`) checks `kbc_access_expires_at` diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 48606daa2..5724442ec 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -50,9 +50,11 @@ class Config: """The URL where the MCP server si reachable.""" jwt_secret: str | None = None """The secret key for encoding and decoding JWT tokens.""" - postgres_dsn: str | None = None + postgres_dsn: str | None = field(default=None, metadata={'aliases': ['mcp_db_url']}) """Connection string for the Postgres-backed OAuth session store (oauth_session_persistence RFC). - Required to enable OAuth login when oauth_client_id/oauth_client_secret are set.""" + Required to enable OAuth login when oauth_client_id/oauth_client_secret are set. + + Maps the `MCP_DB_URL` / `KBC_MCP_DB_URL` env var (via the alias) as well as `KBC_POSTGRES_DSN`.""" session_encryption_key: str | None = None """Base64-encoded 32-byte AES-256 key used to encrypt OAuth session credentials at rest.""" bearer_token: str | None = None diff --git a/tests/test_config.py b/tests/test_config.py index fb738d151..f78cec800 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -42,6 +42,18 @@ class TestConfig: {'X-KBC-ProjectId': '1888'}, Config(project_id='1888'), ), + ( + {'MCP_DB_URL': 'postgresql://u:p@host/db'}, + Config(postgres_dsn='postgresql://u:p@host/db'), + ), + ( + {'KBC_MCP_DB_URL': 'postgresql://u:p@host/db'}, + Config(postgres_dsn='postgresql://u:p@host/db'), + ), + ( + {'KBC_POSTGRES_DSN': 'postgresql://u:p@host/db'}, + Config(postgres_dsn='postgresql://u:p@host/db'), + ), ], ) def test_from_dict(self, d: Mapping[str, str], expected: Config) -> None: From 7e3e6e9bd3b9390f01163238cdeb31dd5f91154c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 30 Jul 2026 13:09:09 +0200 Subject: [PATCH 61/89] feat(PSGO-261): wire SimpleOAuthProvider onto the Postgres session store (Phase 2) OAuth access/refresh tokens are no longer self-contained JWTs -- they're opaque, randomly-generated references to a row in Postgres (oauth_sessions), carrying the real (encrypted) Keboola access/refresh token server-side only. This buys two things the JWT design structurally couldn't: revocation (revoke_token now actually deletes the session instead of being a no-op) and transparent server-managed refresh (load_access_token refreshes the underlying Keboola credential in place when it's near expiry, so the client's own opaque token never needs to change). exchange_refresh_token rotates both the underlying Keboola credential and the client-facing opaque token pair (OAuth 2.1 refresh-token-rotation). Neither opaque token carries a client-visible expiry -- validity is revocation-based (row deleted/soft-revoked), not TTL-based. server.py now refuses to enable OAuth without a Postgres DSN configured (no silent in-memory fallback for a production auth path) and constructs the PostgresSessionStore, whose connection pool is lazily created on first use so create_server() stays a plain sync function -- forcing its many call sites (including a dozen-plus synchronous tests) to become async would have been a much bigger, unrelated change. Also adds the `keboola-mcp-server migrate` CLI subcommand: a thin wrapper over the migration runner added in the prior commit, meant to run as a one-shot Job before the server deployment rolls out (schema migrations are deliberately NOT auto-applied by the app itself). Multi-project scope-on-DB for OAuth sessions (RFC Phase 3, replacing scope_token for this session type specifically) is not part of this commit -- scope_token keeps working unchanged for every session type, OAuth included. Co-Authored-By: Claude Sonnet 5 --- pyproject.toml | 2 +- src/keboola_mcp_server/cli.py | 36 ++++ src/keboola_mcp_server/oauth.py | 191 +++++++++-------- src/keboola_mcp_server/server.py | 16 ++ .../session_store/repository.py | 74 +++++-- tests/session_store/conftest.py | 7 +- tests/test_cli.py | 69 ++++++ tests/test_oauth.py | 196 +++++++++++++++++- tests/test_server.py | 37 ++++ uv.lock | 2 +- 10 files changed, 519 insertions(+), 111 deletions(-) create mode 100644 tests/test_cli.py diff --git a/pyproject.toml b/pyproject.toml index 22548bf12..58ed50874 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.76.0" +version = "1.77.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index ecef2f988..93cb6c720 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -118,6 +118,12 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: ) logout_parser.add_argument('--all', action='store_true', help='Delete stored sessions for all stacks.') + subparsers.add_parser( + 'migrate', + help='Applies pending Postgres schema migrations for the OAuth session store, then exits. ' + 'Intended to run as a one-shot job before the server deployment rolls out.', + ) + return parser.parse_args(args) @@ -230,6 +236,32 @@ async def _run_logout(api_url: str | None, *, all_stacks: bool = False) -> None: print(f'✓ Logged out of {storage_api_url}.' if removed else f'No stored session for {storage_api_url}.') +async def _run_migrate() -> None: + """Applies pending Postgres schema migrations for the OAuth session store, then exits. + + Reads the DSN from the same env vars the server itself uses (MCP_DB_URL / KBC_MCP_DB_URL / + KBC_POSTGRES_DSN) so a migration Job can share the exact same envFrom secret as the deployment. + """ + import asyncpg + + from keboola_mcp_server.session_store.migrator import apply_migrations + + config = Config().replace_by(os.environ) + if not config.postgres_dsn: + raise RuntimeError('A Postgres DSN is required to run migrations: set MCP_DB_URL (or KBC_POSTGRES_DSN).') + + pool = await asyncpg.create_pool(config.postgres_dsn) + try: + applied = await apply_migrations(pool) + finally: + await pool.close() + + if applied: + print(f"✓ Applied {len(applied)} migration(s): {', '.join(applied)}") + else: + print('✓ Schema already up to date -- no migrations applied.') + + async def run_server(args: list[str] | None = None) -> None: """Runs the MCP server in async mode.""" parsed_args = parse_args(args) @@ -272,6 +304,10 @@ async def run_server(args: list[str] | None = None) -> None: await _run_logout(getattr(parsed_args, 'api_url', None), all_stacks=getattr(parsed_args, 'all', False)) return + if parsed_args.command == 'migrate': + await _run_migrate() + return + # Create config from the CLI arguments config = Config( storage_api_url=parsed_args.api_url, diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index 9a431d4ac..d3591ee51 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -1,3 +1,4 @@ +import dataclasses import logging import math import os @@ -5,6 +6,7 @@ import secrets import time from collections.abc import Mapping +from datetime import datetime, timezone from typing import Any, cast from urllib.parse import urljoin @@ -28,6 +30,12 @@ from keboola_mcp_server.clients.auth_bridge import OAuthSessionExchanger, OAuthTokenExchangeError from keboola_mcp_server.config import deployed_sa_token_path from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt +from keboola_mcp_server.session_store.repository import SessionStore + +# The OAuth scope this server always requests at /oauth/consent (see authorize()) -- fixed for +# this flow, so the opaque access/refresh tokens don't need their own scopes column; carried here +# only to satisfy the mcp SDK's AccessToken/RefreshToken (scopes: list[str], required). +_OAUTH_SCOPES = ['claudai', 'projectless'] LOG = logging.getLogger(__name__) _OAUTH_LOG_ALL = bool(os.getenv('KEBOOLA_MCP_SERVER_OAUTH_LOG_ALL')) @@ -131,6 +139,7 @@ class ProxyRefreshToken(RefreshToken): # The refresh side of the same exchanged session; used to refresh independently of the # (single-use, discarded) league OAuth token pair. kbc_refresh_token: str + session_id: str | None = None class SimpleOAuthProvider(OAuthProvider): @@ -144,6 +153,7 @@ def __init__( client_secret: str, server_url: str, scope: str, + session_store: SessionStore, jwt_secret: str | None = None, ) -> None: """ @@ -156,12 +166,17 @@ def __init__( :param client_secret: The client secret registered with the OAuth server :param server_url: The URL of the OAuth server that the MCP server should authenticate to. :param scope: The scope of access to request from the OAuth server. - :param jwt_secret: The secret key for encoding and decoding JWT tokens. + :param session_store: Postgres-backed store for the exchanged Keboola session (access/refresh + token + multi-project scope) -- see oauth_session_persistence RFC. The short-lived, + pre-authentication artifacts (authorize-state, authorization code) still use `jwt_secret` + below; only the long-lived, real-credential-carrying tokens live in the store. + :param jwt_secret: The secret key for encoding and decoding the pre-auth JWT artifacts. """ super().__init__( base_url=mcp_server_url, client_registration_options=ClientRegistrationOptions(enabled=True), ) + self._session_store = session_store self._storage_api_url = storage_api_url self._mcp_callback_url = urljoin(mcp_server_url, callback_endpoint) @@ -394,64 +409,85 @@ async def exchange_authorization_code( # Exchange the league OAuth access token for a whole-stack Keboola programmatic session. # The league token is used exactly once, here, and then never referenced again. token_set = await self._exchange_oauth_for_session(authorization_code.oauth_access_token.token) - return self._wrap_session_as_oauth_token(client, token_set, authorization_code.scopes) + access_token, refresh_token, _session = await self._session_store.create( + client_id=client.client_id, + user_email=None, + kbc_access_token=token_set.access_token, + kbc_refresh_token=token_set.refresh_token, + kbc_access_expires_at=datetime.fromtimestamp(token_set.expires_at, tz=timezone.utc), + ) + return self._oauth_token(access_token, refresh_token, authorization_code.scopes) async def load_access_token(self, token: str) -> AccessToken | None: """ - Loads and validates an access token. - The method decrypts a JWT access token, validates its content, and returns a `ProxyAccessToken` object - if the token is valid and not expired. Returns `None` if the token is invalid or expired. + Loads an access token by looking up the opaque, randomly-generated token in the Postgres + session store (oauth_session_persistence RFC) -- no signature to verify, the DB row's mere + existence (and not being revoked) is the entire validity check. + + Refreshes the underlying Keboola credential transparently if it's near expiry, so a client + that never proactively refreshes its own (non-expiring) opaque token still always gets a + live Keboola session underneath. - :param token: The JWT access token to be loaded and validated. - :return: A `ProxyAccessToken` instance if the token is valid and not expired, otherwise `None`. + :param token: The opaque access token to look up. + :return: A `ProxyAccessToken` carrying the (possibly just-refreshed) Keboola access token, + or `None` if the token doesn't exist or was revoked. """ - try: - access_token_raw = self._decode(token) - except jwt.InvalidTokenError: - LOG.debug(f'[load_access_token] Invalid token: {token}', exc_info=True) + session = await self._session_store.get_by_access_token(token) + if session is None: + _log_debug(f'[load_access_token] Unknown or revoked token: {token}') return None - proxy_token = ProxyAccessToken.model_validate(access_token_raw) - _log_debug(f'[load_access_token] token={token}, proxy_token={proxy_token}') - - # Log the expired authorization code. - # The mcp library itself performs the check and returns a proper response, but no logs. - now = time.time() - if proxy_token.expires_at and proxy_token.expires_at < now: - LOG.info( - f'[load_access_token] Expired access token: proxy_token.expires_at={proxy_token.expires_at}, now={now}' - ) + if session.kbc_access_expires_at.timestamp() <= time.time() + 60: + try: + token_set = await refresh_tokens(self._storage_api_url, refresh_token=session.kbc_refresh_token) + except httpx.HTTPError as e: + # Don't fail the request over a refresh hiccup -- the (soon-to-expire) credential we + # already have may still work for the next little while; the *next* lookup retries. + LOG.warning(f'[load_access_token] Could not refresh near-expiry Keboola session: {e}', exc_info=True) + else: + await self._session_store.rotate_kbc_tokens( + session.id, + kbc_access_token=token_set.access_token, + kbc_refresh_token=token_set.refresh_token, + kbc_access_expires_at=datetime.fromtimestamp(token_set.expires_at, tz=timezone.utc), + ) + session = dataclasses.replace( + session, kbc_access_token=token_set.access_token, kbc_refresh_token=token_set.refresh_token + ) + proxy_token = ProxyAccessToken( + token=token, + client_id=session.client_id, + scopes=_OAUTH_SCOPES, + expires_at=None, # no client-visible expiry -- see load_access_token docstring + kbc_access_token=session.kbc_access_token, + session_id=session.id, + ) + _log_debug(f'[load_access_token] token={token}, session_id={session.id}') return proxy_token async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshToken | None: """ - Loads and validates a refresh token. - The method decrypts a JWT refresh token, validates its content, and returns a `RefreshToken` object - if the token is valid and not expired. Returns `None` if the token is invalid or expired. + Loads a refresh token by looking up the opaque token in the Postgres session store. :param client: The OAuth client details. - :param refresh_token: A string representing the refresh token in JWT format. - :return: A `ProxyRefreshToken` instance if the token is valid and not expired, otherwise `None`. + :param refresh_token: The opaque refresh token to look up. + :return: A `ProxyRefreshToken`, or `None` if the token doesn't exist or was revoked. """ - try: - refresh_token_raw = self._decode(refresh_token) - except jwt.InvalidTokenError: - LOG.debug(f'[load_refresh_token] Invalid token: {refresh_token}', exc_info=True) + session = await self._session_store.get_by_refresh_token(refresh_token) + if session is None: + _log_debug(f'[load_refresh_token] Unknown or revoked token: {refresh_token}') return None - proxy_token = ProxyRefreshToken.model_validate(refresh_token_raw) - _log_debug(f'[load_refresh_token] token={refresh_token}, proxy_token={proxy_token}') - - # Log the expired authorization code. - # The mcp library itself performs the check and returns a proper response, but no logs. - now = time.time() - if proxy_token.expires_at and proxy_token.expires_at < now: - LOG.info( - f'[load_refresh_token] Expired refresh token: proxy_token.expires_at={proxy_token.expires_at}, ' - f'now={now}' - ) - + proxy_token = ProxyRefreshToken( + token=refresh_token, + client_id=session.client_id, + scopes=_OAUTH_SCOPES, + expires_at=None, + kbc_refresh_token=session.kbc_refresh_token, + session_id=session.id, + ) + _log_debug(f'[load_refresh_token] token={refresh_token}, session_id={session.id}') return proxy_token async def exchange_refresh_token( @@ -465,6 +501,9 @@ async def exchange_refresh_token( oauth_session_exchange RFC) — no round-trip to the league OAuth server: that token pair was used once, at initial exchange, and is never touched again. + Also rotates the client-facing opaque access/refresh token pair (OAuth 2.1's refresh-token- + rotation recommendation) -- the old pair stops resolving to this session immediately after. + :param client: The OAuth client details. :param refresh_token: The refresh token to use for renewing the tokens. :param scopes: List of scopes to associate with the new tokens. If not provided, the scopes @@ -472,13 +511,14 @@ async def exchange_refresh_token( :return: A new OAuthToken containing the access and refresh tokens. - :raises HTTPException: If the session-refresh call indicates an error. + :raises TokenError: If the session-refresh call indicates an error. """ _log_debug( f'[exchange_refresh_token] client_id={client.client_id}, refresh_token={refresh_token}, scopes={scopes}' ) assert isinstance(refresh_token, ProxyRefreshToken), f'Expected ProxyRefreshToken, got {type(refresh_token)}' + assert refresh_token.session_id is not None # Raised as TokenError (not HTTPException): this method is invoked by the mcp SDK's own # /token endpoint handler, which only recognizes TokenError and formats it into a spec- @@ -498,63 +538,42 @@ async def exchange_refresh_token( error='invalid_grant', error_description=f'Failed to refresh token: could not reach Connection ({e}).' ) from e - return self._wrap_session_as_oauth_token(client, token_set, scopes or refresh_token.scopes) - - def _wrap_session_as_oauth_token( - self, client: OAuthClientInformationFull, token_set: TokenSet, scopes: list[str] - ) -> OAuthToken: - """Wraps an exchanged Keboola session (`TokenSet`) into our own proxy access/refresh tokens.""" - access_token = ProxyAccessToken( - token=f'mcp_{secrets.token_hex(32)}', - client_id=client.client_id, - scopes=scopes, - expires_at=int(token_set.expires_at), + await self._session_store.rotate_kbc_tokens( + refresh_token.session_id, kbc_access_token=token_set.access_token, - session_id=token_set.session_id, - ) - access_token_jwt = self._encode(access_token.model_dump()) - - # The proxy refresh token's own expiry must NOT be tied to the (short-lived) access token's: - # the underlying Keboola refresh token keeps the session alive indefinitely (RFC Decision §4), - # but the mcp SDK enforces `expires_at` on the object load_refresh_token() returns. Derive a - # longer window the same way the pre-exchange code did for the league OAuth refresh token - # (up to ~7 days), so a client that doesn't refresh for a while isn't forced to re-login. - access_expires_in = max(0, int(token_set.expires_at - time.time())) - refresh_expires_at = int(time.time()) + self._ceil_to_hour(min(168 * access_expires_in, 168 * 3600)) - refresh_token = ProxyRefreshToken( - token=f'mcp_{secrets.token_hex(32)}', - client_id=client.client_id, - scopes=scopes, - expires_at=refresh_expires_at, kbc_refresh_token=token_set.refresh_token, + kbc_access_expires_at=datetime.fromtimestamp(token_set.expires_at, tz=timezone.utc), ) - refresh_token_jwt = self._encode(refresh_token.model_dump()) + new_access_token, new_refresh_token = await self._session_store.rotate_opaque_tokens(refresh_token.session_id) + return self._oauth_token(new_access_token, new_refresh_token, scopes or refresh_token.scopes) - oauth_token = OAuthToken( - access_token=access_token_jwt, - refresh_token=refresh_token_jwt, + @staticmethod + def _oauth_token(access_token: str, refresh_token: str, scopes: list[str]) -> OAuthToken: + # expires_in=None: these opaque tokens don't carry a client-visible expiry (see + # load_access_token) -- the server refreshes the underlying Keboola credential + # transparently, so the client never needs to proactively refresh either. + return OAuthToken( + access_token=access_token, + refresh_token=refresh_token, token_type='Bearer', - expires_in=max(0, int(token_set.expires_at - time.time())), + expires_in=None, scope=' '.join(scopes), ) - _log_debug( - f'[_wrap_session_as_oauth_token] access_token={access_token}, refresh_token={refresh_token}, ' - f'oauth_token={oauth_token}' - ) - return oauth_token async def revoke_token(self, token: str, token_type_hint: str | None = None) -> None: """ - Revokes a token. - - This is a no-op function as the tokens are not stored and so there is no way to revoke tokens that have already - been issued. + Revokes a token by deleting its session from the Postgres store (soft-delete via + `revoked_at`) -- both the access and refresh token immediately stop resolving. - :param token: The token to be revoked. + :param token: The token to be revoked (access or refresh; `token_type_hint` is advisory). :param token_type_hint: An optional hint about the type of the token. """ _log_debug(f'[revoke_token] token={token}, token_type_hint={token_type_hint}') - # This is no-op as we don't store the tokens. + session = await self._session_store.get_by_access_token( + token + ) or await self._session_store.get_by_refresh_token(token) + if session is not None: + await self._session_store.revoke(session.id) def _read_oauth_tokens(self, data: dict[str, Any], scopes: list[str]) -> tuple[AccessToken, RefreshToken]: """ diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index 24a5c16f0..8254482c4 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -28,6 +28,8 @@ from keboola_mcp_server.oauth import SimpleOAuthProvider from keboola_mcp_server.preview import preview_config_diff from keboola_mcp_server.prompts.add_prompts import add_keboola_prompts +from keboola_mcp_server.session_store.crypto import resolve_encryption_key +from keboola_mcp_server.session_store.repository import PostgresSessionStore from keboola_mcp_server.tools.components.tools import add_component_tools from keboola_mcp_server.tools.data_apps import add_data_app_tools from keboola_mcp_server.tools.doc import add_doc_tools @@ -209,6 +211,19 @@ def create_server( if not config.oauth_scope: config = dataclasses.replace(config, oauth_scope='email') + # OAuth sessions (the real Keboola access/refresh tokens) live in Postgres, not in a + # self-contained JWT (oauth_session_persistence RFC) -- revocation and server-managed + # refresh both need a durable, deletable row. No silent in-memory fallback for this + # production auth path: refuse to start rather than accept OAuth logins nothing can revoke. + if not config.postgres_dsn: + raise RuntimeError( + 'OAuth is configured (oauth_client_id/oauth_client_secret) but no Postgres DSN is set. ' + 'Set MCP_DB_URL (or KBC_POSTGRES_DSN) so OAuth sessions can be stored.' + ) + session_store = PostgresSessionStore( + config.postgres_dsn, encryption_key=resolve_encryption_key(config.session_encryption_key) + ) + oauth_provider = SimpleOAuthProvider( storage_api_url=config.storage_api_url, client_id=config.oauth_client_id, @@ -220,6 +235,7 @@ def create_server( # The path corresponds to oauth_callback_handler() set up below. callback_endpoint='/oauth/callback', jwt_secret=config.jwt_secret, + session_store=session_store, ) else: oauth_provider = None diff --git a/src/keboola_mcp_server/session_store/repository.py b/src/keboola_mcp_server/session_store/repository.py index 033004ed3..a7eed78b2 100644 --- a/src/keboola_mcp_server/session_store/repository.py +++ b/src/keboola_mcp_server/session_store/repository.py @@ -4,6 +4,7 @@ without a real database; `PostgresSessionStore` is the only production implementation. """ +import asyncio import dataclasses import hashlib import secrets @@ -13,7 +14,6 @@ import asyncpg from keboola_mcp_server.session_store import crypto -from keboola_mcp_server.session_store.migrator import apply_migrations # Length of the opaque, randomly-generated access/refresh tokens handed to the MCP client. 256 # bits: not guessable, and this is the *entire* security check for these tokens (no signature to @@ -69,6 +69,14 @@ async def rotate_kbc_tokens( """Replaces the encrypted Keboola credentials in place (server-managed refresh).""" ... + async def rotate_opaque_tokens(self, session_id: str) -> tuple[str, str]: + """Issues a fresh opaque access/refresh token pair for an existing session (OAuth refresh + grant, per OAuth 2.1's refresh-token-rotation recommendation), invalidating the old pair. + + :return: (new_opaque_access_token, new_opaque_refresh_token) + """ + ... + async def update_scope( self, session_id: str, @@ -84,18 +92,32 @@ async def revoke(self, session_id: str) -> None: ... class PostgresSessionStore: - def __init__(self, pool: asyncpg.Pool, encryption_key: bytes) -> None: - self._pool = pool + """Schema migrations are NOT applied here -- that's the `keboola-mcp-server migrate` CLI/Job's + job, run once per deployment before this app starts (oauth_session_persistence RFC). This class + only ever reads/writes rows, assuming the schema is already in place. + + The connection pool is created lazily, on first use, so construction stays synchronous (no + event loop required) -- `server.py`'s `create_server()` is a plain sync function, and forcing + every one of its many call sites (including a dozen-plus sync tests) to become async just to + accommodate this would be a much bigger, unrelated change. + """ + + def __init__(self, dsn: str, encryption_key: bytes) -> None: + self._dsn = dsn self._key = encryption_key + self._pool: asyncpg.Pool | None = None + self._pool_lock = asyncio.Lock() - @classmethod - async def connect(cls, dsn: str, encryption_key: bytes) -> 'PostgresSessionStore': - pool = await asyncpg.create_pool(dsn) - await apply_migrations(pool) - return cls(pool, encryption_key) + async def _get_pool(self) -> asyncpg.Pool: + if self._pool is None: + async with self._pool_lock: + if self._pool is None: # re-check: another task may have won the lock race first + self._pool = await asyncpg.create_pool(self._dsn) + return self._pool async def close(self) -> None: - await self._pool.close() + if self._pool is not None: + await self._pool.close() def _to_session(self, row: asyncpg.Record) -> OAuthSession: return OAuthSession( @@ -127,7 +149,8 @@ async def create( ) -> tuple[str, str, OAuthSession]: access_token = generate_opaque_token() refresh_token = generate_opaque_token() - row = await self._pool.fetchrow( + pool = await self._get_pool() + row = await pool.fetchrow( """ INSERT INTO oauth_sessions ( access_token_hash, refresh_token_hash, client_id, user_email, @@ -147,7 +170,8 @@ async def create( return access_token, refresh_token, self._to_session(row) async def get_by_access_token(self, access_token: str) -> OAuthSession | None: - row = await self._pool.fetchrow( + pool = await self._get_pool() + row = await pool.fetchrow( 'UPDATE oauth_sessions SET last_used_at = now() ' 'WHERE access_token_hash = $1 AND revoked_at IS NULL RETURNING *', _hash_token(access_token), @@ -155,7 +179,8 @@ async def get_by_access_token(self, access_token: str) -> OAuthSession | None: return self._to_session(row) if row is not None else None async def get_by_refresh_token(self, refresh_token: str) -> OAuthSession | None: - row = await self._pool.fetchrow( + pool = await self._get_pool() + row = await pool.fetchrow( 'SELECT * FROM oauth_sessions WHERE refresh_token_hash = $1 AND revoked_at IS NULL', _hash_token(refresh_token), ) @@ -164,7 +189,8 @@ async def get_by_refresh_token(self, refresh_token: str) -> OAuthSession | None: async def rotate_kbc_tokens( self, session_id: str, *, kbc_access_token: str, kbc_refresh_token: str, kbc_access_expires_at: datetime ) -> None: - await self._pool.execute( + pool = await self._get_pool() + await pool.execute( """ UPDATE oauth_sessions SET kbc_access_token_enc = $2, kbc_refresh_token_enc = $3, kbc_access_expires_at = $4, @@ -177,6 +203,22 @@ async def rotate_kbc_tokens( kbc_access_expires_at, ) + async def rotate_opaque_tokens(self, session_id: str) -> tuple[str, str]: + access_token = generate_opaque_token() + refresh_token = generate_opaque_token() + pool = await self._get_pool() + await pool.execute( + """ + UPDATE oauth_sessions + SET access_token_hash = $2, refresh_token_hash = $3, updated_at = now() + WHERE id = $1 + """, + session_id, + _hash_token(access_token), + _hash_token(refresh_token), + ) + return access_token, refresh_token + async def update_scope( self, session_id: str, @@ -187,7 +229,8 @@ async def update_scope( scoped_token: str | None, scoped_expires_at: datetime | None, ) -> None: - await self._pool.execute( + pool = await self._get_pool() + await pool.execute( """ UPDATE oauth_sessions SET scope_project_ids = $2, scope_read_only = $3, scope_confirmed = $4, @@ -203,4 +246,5 @@ async def update_scope( ) async def revoke(self, session_id: str) -> None: - await self._pool.execute('UPDATE oauth_sessions SET revoked_at = now() WHERE id = $1', session_id) + pool = await self._get_pool() + await pool.execute('UPDATE oauth_sessions SET revoked_at = now() WHERE id = $1', session_id) diff --git a/tests/session_store/conftest.py b/tests/session_store/conftest.py index 5b5b946e4..5aa70e940 100644 --- a/tests/session_store/conftest.py +++ b/tests/session_store/conftest.py @@ -5,6 +5,7 @@ import pytest_asyncio from keboola_mcp_server.session_store.crypto import KEY_SIZE +from keboola_mcp_server.session_store.migrator import apply_migrations from keboola_mcp_server.session_store.repository import PostgresSessionStore TEST_DSN = os.environ.get('KBC_TEST_POSTGRES_DSN', 'postgresql://keboola_mcp:keboola_mcp@localhost:5432/keboola_mcp') @@ -31,11 +32,13 @@ def _postgres_available() -> bool: async def store(): pool = await asyncpg.create_pool(TEST_DSN) try: - # Clean slate per test: drop and let apply_migrations (inside connect) recreate. + # Clean slate per test: drop, then re-apply migrations -- standing in for the migration + # Job that would normally run once, ahead of the app, in a real deployment. await pool.execute('DROP TABLE IF EXISTS oauth_sessions, schema_migrations CASCADE') + await apply_migrations(pool) finally: await pool.close() - s = await PostgresSessionStore.connect(TEST_DSN, encryption_key=bytes([1] * KEY_SIZE)) + s = PostgresSessionStore(TEST_DSN, encryption_key=bytes([1] * KEY_SIZE)) try: yield s finally: diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 000000000..a56beec9c --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,69 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from keboola_mcp_server.cli import _run_migrate, parse_args + + +def test_parse_args_migrate() -> None: + args = parse_args(['migrate']) + assert args.command == 'migrate' + + +class TestRunMigrate: + @pytest.mark.asyncio + async def test_requires_postgres_dsn(self, monkeypatch) -> None: + monkeypatch.delenv('MCP_DB_URL', raising=False) + monkeypatch.delenv('KBC_POSTGRES_DSN', raising=False) + monkeypatch.delenv('KBC_MCP_DB_URL', raising=False) + with pytest.raises(RuntimeError, match='Postgres DSN'): + await _run_migrate() + + @pytest.mark.asyncio + async def test_applies_migrations_and_closes_pool(self, monkeypatch, capsys) -> None: + monkeypatch.setenv('MCP_DB_URL', 'postgresql://u:p@host/db') + pool = MagicMock() + pool.close = AsyncMock() + with ( + patch('asyncpg.create_pool', AsyncMock(return_value=pool)) as create_pool, + patch( + 'keboola_mcp_server.session_store.migrator.apply_migrations', + AsyncMock(return_value=['0001_oauth_sessions.sql']), + ), + ): + await _run_migrate() + + create_pool.assert_awaited_once_with('postgresql://u:p@host/db') + pool.close.assert_awaited_once() + assert '0001_oauth_sessions.sql' in capsys.readouterr().out + + @pytest.mark.asyncio + async def test_no_pending_migrations_still_closes_pool(self, monkeypatch, capsys) -> None: + monkeypatch.setenv('MCP_DB_URL', 'postgresql://u:p@host/db') + pool = MagicMock() + pool.close = AsyncMock() + with ( + patch('asyncpg.create_pool', AsyncMock(return_value=pool)), + patch('keboola_mcp_server.session_store.migrator.apply_migrations', AsyncMock(return_value=[])), + ): + await _run_migrate() + + pool.close.assert_awaited_once() + assert 'up to date' in capsys.readouterr().out + + @pytest.mark.asyncio + async def test_closes_pool_even_if_migration_fails(self, monkeypatch) -> None: + monkeypatch.setenv('MCP_DB_URL', 'postgresql://u:p@host/db') + pool = MagicMock() + pool.close = AsyncMock() + with ( + patch('asyncpg.create_pool', AsyncMock(return_value=pool)), + patch( + 'keboola_mcp_server.session_store.migrator.apply_migrations', + AsyncMock(side_effect=RuntimeError('boom')), + ), + ): + with pytest.raises(RuntimeError, match='boom'): + await _run_migrate() + + pool.close.assert_awaited_once() diff --git a/tests/test_oauth.py b/tests/test_oauth.py index c858de107..9d5ea924c 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -1,5 +1,8 @@ +import dataclasses +import secrets import time from collections.abc import Mapping +from datetime import datetime, timedelta, timezone from http import HTTPStatus from typing import Any from urllib.parse import parse_qs, urlparse @@ -17,10 +20,90 @@ _ExtendedAuthorizationCode, _OAuthClientInformationFull, ) +from keboola_mcp_server.session_store.repository import OAuthSession JWT_KEY = 'secret' +class FakeSessionStore: + """In-memory `SessionStore` (no real Postgres) for exercising `SimpleOAuthProvider` in isolation.""" + + def __init__(self) -> None: + self._sessions: dict[str, OAuthSession] = {} + self._access_tokens: dict[str, str] = {} + self._refresh_tokens: dict[str, str] = {} + self._next_id = 0 + + def _new_token_pair(self, session_id: str) -> tuple[str, str]: + access_token = f'at_{session_id}_{secrets.token_hex(4)}' + refresh_token = f'rt_{session_id}_{secrets.token_hex(4)}' + self._access_tokens[access_token] = session_id + self._refresh_tokens[refresh_token] = session_id + return access_token, refresh_token + + async def create( + self, *, client_id, user_email, kbc_access_token, kbc_refresh_token, kbc_access_expires_at + ) -> tuple[str, str, OAuthSession]: + self._next_id += 1 + session_id = str(self._next_id) + session = OAuthSession( + id=session_id, + client_id=client_id, + user_email=user_email, + kbc_access_token=kbc_access_token, + kbc_refresh_token=kbc_refresh_token, + kbc_access_expires_at=kbc_access_expires_at, + scope_project_ids=None, + scope_read_only=False, + scope_confirmed=False, + scope_scoped_token=None, + scope_scoped_expires_at=None, + ) + self._sessions[session_id] = session + access_token, refresh_token = self._new_token_pair(session_id) + return access_token, refresh_token, session + + async def get_by_access_token(self, access_token: str) -> OAuthSession | None: + session_id = self._access_tokens.get(access_token) + return self._sessions.get(session_id) if session_id else None + + async def get_by_refresh_token(self, refresh_token: str) -> OAuthSession | None: + session_id = self._refresh_tokens.get(refresh_token) + return self._sessions.get(session_id) if session_id else None + + async def rotate_kbc_tokens( + self, session_id: str, *, kbc_access_token: str, kbc_refresh_token: str, kbc_access_expires_at: datetime + ) -> None: + session = self._sessions[session_id] + self._sessions[session_id] = dataclasses.replace( + session, + kbc_access_token=kbc_access_token, + kbc_refresh_token=kbc_refresh_token, + kbc_access_expires_at=kbc_access_expires_at, + ) + + async def rotate_opaque_tokens(self, session_id: str) -> tuple[str, str]: + self._access_tokens = {k: v for k, v in self._access_tokens.items() if v != session_id} + self._refresh_tokens = {k: v for k, v in self._refresh_tokens.items() if v != session_id} + return self._new_token_pair(session_id) + + async def update_scope( + self, session_id: str, *, project_ids, read_only, confirmed, scoped_token, scoped_expires_at + ) -> None: + session = self._sessions[session_id] + self._sessions[session_id] = dataclasses.replace( + session, + scope_project_ids=project_ids, + scope_read_only=read_only, + scope_confirmed=confirmed, + scope_scoped_token=scoped_token, + scope_scoped_expires_at=scoped_expires_at, + ) + + async def revoke(self, session_id: str) -> None: + self._sessions.pop(session_id, None) + + class TestSimpleOAuthProvider: @pytest.fixture def oauth_provider(self) -> SimpleOAuthProvider: @@ -33,6 +116,7 @@ def oauth_provider(self) -> SimpleOAuthProvider: server_url='https://oauth', scope='scope', jwt_secret=JWT_KEY, + session_store=FakeSessionStore(), ) @staticmethod @@ -280,10 +364,11 @@ async def exchange(self, *, oauth_access_token: str): loaded_refresh = await oauth_provider.load_refresh_token(client, oauth_token.refresh_token) assert loaded_refresh is not None assert loaded_refresh.kbc_refresh_token == 'kbc_rt_new' - # The refresh token's own expiry must be much longer than the (1h) access token's -- it must - # not be tied to it, or the mcp SDK would force a re-login every ~1h even though the - # underlying Keboola refresh token can keep the session alive indefinitely. - assert loaded_refresh.expires_at - loaded.expires_at > 6 * 24 * 3600 # at least ~6 more days + # Neither opaque token carries a client-visible expiry (oauth_session_persistence RFC): the + # server refreshes the underlying Keboola credential transparently on lookup, so there's no + # forced-relogin window tied to the (1h) Keboola access token's lifetime. + assert loaded.expires_at is None + assert loaded_refresh.expires_at is None @pytest.mark.asyncio async def test_exchange_authorization_code_maps_exchange_error( @@ -350,12 +435,20 @@ async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, tran ) client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + _at, _rt, session = await oauth_provider._session_store.create( + client_id='foo-client-id', + user_email=None, + kbc_access_token='kbc_at_old', + kbc_refresh_token='kbc_rt_old', + kbc_access_expires_at=datetime.now(timezone.utc), + ) refresh_token = ProxyRefreshToken( token='mcp_old', client_id='foo-client-id', scopes=['claudai', 'projectless'], - expires_at=int(time.time() + 3600), + expires_at=None, kbc_refresh_token='kbc_rt_old', + session_id=session.id, ) oauth_token = await oauth_provider.exchange_refresh_token(client, refresh_token, []) @@ -380,12 +473,20 @@ async def _failing_refresh_tokens(storage_api_url: str, *, refresh_token: str, t monkeypatch.setattr(oauth_module, 'refresh_tokens', _failing_refresh_tokens) client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + _at, _rt, session = await oauth_provider._session_store.create( + client_id='foo-client-id', + user_email=None, + kbc_access_token='kbc_at_old', + kbc_refresh_token='kbc_rt_old', + kbc_access_expires_at=datetime.now(timezone.utc), + ) refresh_token = ProxyRefreshToken( token='mcp_old', client_id='foo-client-id', scopes=['claudai', 'projectless'], - expires_at=int(time.time() + 3600), + expires_at=None, kbc_refresh_token='kbc_rt_old', + session_id=session.id, ) # A network failure talking to Connection must surface as a clean TokenError, not @@ -393,3 +494,86 @@ async def _failing_refresh_tokens(storage_api_url: str, *, refresh_token: str, t with pytest.raises(TokenError) as exc: await oauth_provider.exchange_refresh_token(client, refresh_token, []) assert exc.value.error == 'invalid_grant' + + @pytest.mark.asyncio + async def test_load_access_token_refreshes_near_expiry_session_transparently( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + from keboola_mcp_server import oauth as oauth_module + from keboola_mcp_server.auth_login import TokenSet + + access_token, _rt, session = await oauth_provider._session_store.create( + client_id='foo-client-id', + user_email=None, + kbc_access_token='kbc_at_stale', + kbc_refresh_token='kbc_rt_stale', + kbc_access_expires_at=datetime.now(timezone.utc), # already at/past expiry + ) + + async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, transport=None): + assert refresh_token == 'kbc_rt_stale' + return TokenSet(access_token='kbc_at_fresh', refresh_token='kbc_rt_fresh', expires_at=time.time() + 3600) + + monkeypatch.setattr(oauth_module, 'refresh_tokens', _fake_refresh_tokens) + + loaded = await oauth_provider.load_access_token(access_token) + + assert loaded is not None + assert loaded.kbc_access_token == 'kbc_at_fresh' + # The refresh is persisted, not just returned once -- a second lookup sees it too. + stored = await oauth_provider._session_store.get_by_access_token(access_token) + assert stored is not None + assert stored.kbc_access_token == 'kbc_at_fresh' + assert stored.kbc_refresh_token == 'kbc_rt_fresh' + + @pytest.mark.asyncio + async def test_load_access_token_tolerates_refresh_failure( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + # A refresh hiccup must not break the current request -- the (soon-to-expire) credential + # already on the session may still work; the next lookup retries the refresh. + from keboola_mcp_server import oauth as oauth_module + + access_token, _rt, _session = await oauth_provider._session_store.create( + client_id='foo-client-id', + user_email=None, + kbc_access_token='kbc_at_stale', + kbc_refresh_token='kbc_rt_stale', + kbc_access_expires_at=datetime.now(timezone.utc), + ) + + async def _failing_refresh_tokens(storage_api_url: str, *, refresh_token: str, transport=None): + raise httpx.ConnectError('boom') + + monkeypatch.setattr(oauth_module, 'refresh_tokens', _failing_refresh_tokens) + + loaded = await oauth_provider.load_access_token(access_token) + + assert loaded is not None + assert loaded.kbc_access_token == 'kbc_at_stale' # unchanged, refresh failed but didn't raise + + @pytest.mark.asyncio + async def test_load_access_token_unknown_token_returns_none(self, oauth_provider: SimpleOAuthProvider) -> None: + assert await oauth_provider.load_access_token('never-issued') is None + + @pytest.mark.asyncio + async def test_revoke_token_invalidates_both_access_and_refresh_token( + self, oauth_provider: SimpleOAuthProvider + ) -> None: + access_token, refresh_token, _session = await oauth_provider._session_store.create( + client_id='foo-client-id', + user_email=None, + kbc_access_token='kbc_at_x', + kbc_refresh_token='kbc_rt_x', + kbc_access_expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + + await oauth_provider.revoke_token(access_token) + + assert await oauth_provider.load_access_token(access_token) is None + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + assert await oauth_provider.load_refresh_token(client, refresh_token) is None + + @pytest.mark.asyncio + async def test_revoke_token_unknown_token_is_a_noop(self, oauth_provider: SimpleOAuthProvider) -> None: + await oauth_provider.revoke_token('never-issued') # must not raise diff --git a/tests/test_server.py b/tests/test_server.py index 9e249745b..4d42aa480 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -617,3 +617,40 @@ async def test_oauth_callback_handler_propagates_http_exception(mocker) -> None: await routes.oauth_callback_handler(request) assert exc.value.status_code == 400 assert exc.value.detail == 'Invalid state parameter' + + +class TestCreateServerOAuthSessionStore: + """OAuth sessions live in Postgres (oauth_session_persistence RFC) -- create_server() must + refuse to enable OAuth without a DSN rather than silently falling back to something unrevoked.""" + + @staticmethod + def _oauth_config(**overrides) -> Config: + return Config( + storage_api_url='https://connection.keboola.com', + oauth_client_id='client-id', + oauth_client_secret='client-secret', + oauth_server_url='https://connection.keboola.com', + mcp_server_url='https://mcp.keboola.com', + **overrides, + ) + + def test_raises_without_postgres_dsn(self) -> None: + with pytest.raises(RuntimeError, match='MCP_DB_URL'): + create_server(self._oauth_config(), runtime_info=ServerRuntimeInfo(transport='streamable-http')) + + def test_constructs_session_store_when_dsn_is_set(self) -> None: + from keboola_mcp_server.session_store.repository import PostgresSessionStore + + server = create_server( + self._oauth_config(postgres_dsn='postgresql://u:p@host/db'), + runtime_info=ServerRuntimeInfo(transport='streamable-http'), + ) + assert isinstance(server, FastMCP) + assert isinstance(server.auth._session_store, PostgresSessionStore) + + def test_no_oauth_configured_needs_no_postgres_dsn(self) -> None: + # The vast majority of create_server() call sites (local stdio, header/PAT-token sessions) + # have no OAuth at all -- this must keep working with zero Postgres setup. + server = create_server(Config(), runtime_info=ServerRuntimeInfo(transport='stdio')) + assert isinstance(server, FastMCP) + assert server.auth is None diff --git a/uv.lock b/uv.lock index 33e2a2f87..3fe81c914 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.76.0" +version = "1.77.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From a406aec3088121bab7ec1fa1207ebf4bd5b48203 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 08:46:29 +0200 Subject: [PATCH 62/89] feat(PSGO-261): persist multi-project scope on the OAuth session row OAuth-authenticated sessions no longer need to resend scope_token: the opaque OAuth access token already round-trips through the Postgres session store on every call (load_access_token), so set_project_scope now persists the confirmed scope there too (SessionStore.update_scope) instead of only minting a scope_token. PAT/header-token sessions have no session row to persist against and keep relying on scope_token as before. Bump to 1.78.0 (new capability, not just a fix). --- pyproject.toml | 2 +- src/keboola_mcp_server/mcp.py | 52 +++++++++++++++++++ src/keboola_mcp_server/oauth.py | 15 ++++++ src/keboola_mcp_server/server.py | 3 +- src/keboola_mcp_server/tools/project.py | 66 ++++++++++++++++++++----- tests/test_mcp.py | 58 ++++++++++++++++++++++ tests/tools/test_project.py | 34 ++++++++++++- uv.lock | 2 +- 8 files changed, 214 insertions(+), 18 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 58ed50874..b60ed8ebc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.77.0" +version = "1.78.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 0766575a8..d19d0ae94 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -40,6 +40,7 @@ from keboola_mcp_server.config import Config, ServerRuntimeInfo, deployed_sa_token_path, is_same_stack from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt from keboola_mcp_server.oauth import ProxyAccessToken +from keboola_mcp_server.session_store.repository import SessionStore from keboola_mcp_server.tools.constants import MODIFY_FLOW_TOOL_NAME, SEMANTIC_TOOLS_TAG, UPDATE_FLOW_TOOL_NAME from keboola_mcp_server.workspace import WorkspaceManager @@ -47,6 +48,12 @@ CONVERSATION_ID = 'conversation_id' SCOPE_KEY = 'project_scope' +# The OAuth session's DB row id (see session_store.repository.OAuthSession), stashed on +# ctx.session.state so set_project_scope can persist a newly-confirmed scope back to Postgres +# instead of only returning a scope_token. Absent for non-OAuth (PAT/header-token) sessions, which +# have no session row to persist against -- those keep relying on scope_token. +OAUTH_SESSION_ID_KEY = 'oauth_session_id' + # Tools that must not be fanned out across multiple projects, even when a multi-project scope is # active and they are read-only: the scope/auth tools operate on the whole-stack token (not a single # project), and get_project_info resolves through the active project's WorkspaceManager (workspace id @@ -183,6 +190,7 @@ async def project_has_semantic_models(client: KeboolaClient) -> bool: class ServerState: config: Config runtime_info: ServerRuntimeInfo + session_store: SessionStore | None = None @property def own_stack_storage_api_url(self) -> str | None: @@ -332,6 +340,12 @@ async def on_request( # stateless-HTTP transport. With no scope and no preset project, auto-lease ALL accessible # projects (multi-project mode) so read tools fan out across everything — but never on /list. scope = self._read_scope_from_request(context, config) + # OAuth-authenticated sessions don't need scope_token at all: the opaque OAuth access + # token is already resent on every call and resolves through the Postgres session store + # (load_access_token), so a confirmed scope persisted there (via set_project_scope -> + # SessionStore.update_scope) is read back here instead of round-tripping it as an argument. + if scope is None: + scope = self._read_persisted_oauth_scope(http_rq) if scope is None and not config.project_id and not is_list: scope = await self._autolease_default_scope(config) if not is_list: @@ -354,6 +368,8 @@ async def on_request( ) if scope is not None: state[SCOPE_KEY] = scope + if oauth_session_id := self._read_oauth_session_id(http_rq): + state[OAUTH_SESSION_ID_KEY] = oauth_session_id ctx.session.state = state try: @@ -484,6 +500,42 @@ def _read_scope_from_request(cls, context: fmw.MiddlewareContext[Any], config: C LOG.warning('Ignoring invalid or expired scope_token.', exc_info=True) return None + @staticmethod + def _oauth_access_token(http_rq: Request | None) -> ProxyAccessToken | None: + if http_rq is None: + return None + user = http_rq.scope.get('user') + if not isinstance(user, AuthenticatedUser) or not isinstance(user.access_token, ProxyAccessToken): + return None + return user.access_token + + @classmethod + def _read_persisted_oauth_scope(cls, http_rq: Request | None) -> 'SessionScope | None': + """The multi-project scope persisted on the OAuth session row, if any. + + Only used as a fallback when the caller sent no ``scope_token`` -- an explicit scope_token + (e.g. a fresher re-scope from the same request) always takes precedence. + """ + access_token = cls._oauth_access_token(http_rq) + if access_token is None or not access_token.scope_confirmed or access_token.scope_project_ids is None: + return None + return SessionScope( + project_ids=access_token.scope_project_ids, + read_only=access_token.scope_read_only, + scoped_token=access_token.scope_scoped_token, + scoped_expires_at=( + access_token.scope_scoped_expires_at.timestamp() + if access_token.scope_scoped_expires_at is not None + else None + ), + confirmed=True, + ) + + @classmethod + def _read_oauth_session_id(cls, http_rq: Request | None) -> str | None: + access_token = cls._oauth_access_token(http_rq) + return access_token.session_id if access_token is not None else None + @classmethod def _is_local_programmatic(cls, config: Config) -> bool: """True for a local (non-deployed) session carrying a Keboola programmatic token.""" diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index d3591ee51..e322ca6bb 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -134,6 +134,16 @@ class ProxyAccessToken(AccessToken): kbc_access_token: str session_id: str | None = None + # The multi-project scope persisted on the oauth_sessions row (see SessionStore.update_scope), + # carried here so mcp.py can rebuild a SessionScope without a second DB round-trip -- the row is + # already fetched in load_access_token below. Same exposure-minimization reasoning as + # kbc_access_token: only the fields mcp.py actually needs, not the whole OAuthSession. + scope_project_ids: list[int] | None = None + scope_read_only: bool = False + scope_confirmed: bool = False + scope_scoped_token: str | None = None + scope_scoped_expires_at: datetime | None = None + class ProxyRefreshToken(RefreshToken): # The refresh side of the same exchanged session; used to refresh independently of the @@ -462,6 +472,11 @@ async def load_access_token(self, token: str) -> AccessToken | None: expires_at=None, # no client-visible expiry -- see load_access_token docstring kbc_access_token=session.kbc_access_token, session_id=session.id, + scope_project_ids=session.scope_project_ids, + scope_read_only=session.scope_read_only, + scope_confirmed=session.scope_confirmed, + scope_scoped_token=session.scope_scoped_token, + scope_scoped_expires_at=session.scope_scoped_expires_at, ) _log_debug(f'[load_access_token] token={token}, session_id={session.id}') return proxy_token diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index 8254482c4..c3acf7f2d 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -239,10 +239,11 @@ def create_server( ) else: oauth_provider = None + session_store = None # Initialize FastMCP server with system lifespan LOG.info(f'Creating server with config: {config}') - server_state = ServerState(config=config, runtime_info=runtime_info) + server_state = ServerState(config=config, runtime_info=runtime_info, session_store=session_store) mcp = KeboolaMcpServer( name='Keboola MCP Server', instructions=( diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 512a4ab52..496c9638c 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -1,5 +1,6 @@ import asyncio import logging +from datetime import datetime, timezone from typing import Annotated, Optional, cast import httpx @@ -16,6 +17,7 @@ from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager from keboola_mcp_server.mcp import ( + OAUTH_SESSION_ID_KEY, SCOPE_KEY, MultiProjectMiddleware, ServerState, @@ -341,15 +343,43 @@ class AccessibleProjects(BaseModel): class ProjectScope(BaseModel): project_ids: list[int] = Field(description='The projects the session is now scoped to.') read_only: bool = Field(description='Whether the scoped token is read-only.') - scope_token: str = Field( + scope_token: str | None = Field( + default=None, description=( - 'Opaque token encoding this scope. The server does not remember it between calls -- pass ' + 'Opaque token encoding this scope, or null for an OAuth-authenticated session (the ' + 'server persists the scope itself in that case -- no need to resend it). Otherwise, pass ' 'this value as the "scope_token" argument on every subsequent tool call in this conversation.' ), ) llm_instruction: str = Field(description='Guidance for the assistant on the new scope.') +async def _persist_oauth_scope(ctx: Context, scope: SessionScope) -> bool: + """Persists ``scope`` on the caller's OAuth session row, if this is an OAuth-authenticated + session (see mcp.OAUTH_SESSION_ID_KEY). No-op (returns False) for PAT/header-token sessions, + which have no session row to persist against -- those keep relying on scope_token. + """ + session_id = ctx.session.state.get(OAUTH_SESSION_ID_KEY) + if not session_id: + return False + session_store = ServerState.from_context(ctx).session_store + if session_store is None: + return False + await session_store.update_scope( + session_id, + project_ids=scope.project_ids, + read_only=scope.read_only, + confirmed=scope.confirmed, + scoped_token=scope.scoped_token, + scoped_expires_at=( + datetime.fromtimestamp(scope.scoped_expires_at, tz=timezone.utc) + if scope.scoped_expires_at is not None + else None + ), + ) + return True + + async def _project_sql_dialect( server_state: ServerState, storage_api_url: str, subject_token: str, project_id: int ) -> tuple[int, str | None]: @@ -447,9 +477,14 @@ async def get_accessible_projects( 'projects or a subset, then call "set_project_scope" with the chosen project ids. Never write ' 'to more than one project without explicit user confirmation.' ) + is_oauth_persisted = False else: + is_oauth_persisted = bool(ctx.session.state.get(OAUTH_SESSION_ID_KEY)) instruction = ( - f'Session is currently scoped to {len(scoped_ids)} project(s). Resend "scope_token" on every ' + f'Session is currently scoped to {len(scoped_ids)} project(s). Call "set_project_scope" to ' + 'change the scope.' + if is_oauth_persisted + else f'Session is currently scoped to {len(scoped_ids)} project(s). Resend "scope_token" on every ' 'subsequent tool call to keep it in effect; call "set_project_scope" to change the scope.' ) return AccessibleProjects( @@ -457,7 +492,11 @@ async def get_accessible_projects( projects=projects, scoped_project_ids=scoped_ids, read_only=scope.read_only if scoped_ids is not None else None, - scope_token=scope.to_token(resolve_scope_secret(server_state.config)) if scoped_ids is not None else None, + scope_token=( + scope.to_token(resolve_scope_secret(server_state.config)) + if scoped_ids is not None and not is_oauth_persisted + else None + ), base_instructions=base_instructions, llm_instruction=instruction, ) @@ -542,7 +581,14 @@ async def set_project_scope( LOG.debug(f'Could not send tools/list_changed after scoping: {e}') multi = len(ids) > 1 - scope_token = scope.to_token(resolve_scope_secret(ServerState.from_context(ctx).config)) + persisted = await _persist_oauth_scope(ctx, scope) + scope_token = None if persisted else scope.to_token(resolve_scope_secret(ServerState.from_context(ctx).config)) + resend_instruction = ( + 'The server persists this scope server-side for the rest of the conversation -- no need to resend it.' + if persisted + else 'The server does not remember this scope between calls -- pass "scope_token" as an argument on ' + 'every subsequent tool call in this conversation.' + ) return ProjectScope( project_ids=ids, read_only=scope.read_only, @@ -551,15 +597,9 @@ async def set_project_scope( ( f'Session scoped to {len(ids)} projects. Read-only tools return results per project. ' 'Write operations are not fanned out — they target the first scoped project; to write ' - 'elsewhere, re-scope to that project first (confirm with the user). The server does not ' - 'remember this scope between calls -- pass "scope_token" as an argument on every ' - 'subsequent tool call in this conversation.' + f'elsewhere, re-scope to that project first (confirm with the user). {resend_instruction}' ) if multi - else ( - f'Session scoped to project {ids[0]}. The server does not remember this scope between ' - 'calls -- pass "scope_token" as an argument on every subsequent tool call in this ' - 'conversation.' - ) + else f'Session scoped to project {ids[0]}. {resend_instruction}' ), ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index fab85977d..2c951c3eb 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -9,6 +9,7 @@ from fastmcp.exceptions import ToolError from fastmcp.tools.tool import ToolResult from mcp import types as mt +from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from pydantic import BaseModel, Field from pydantic import ValidationError as PydanticValidationError from starlette.requests import Request @@ -1174,6 +1175,63 @@ def test_wrong_secret_falls_back_to_no_scope_via_read_scope_from_request(self) - context = self._call_tool_context({'scope_token': token}) assert SessionStateMiddleware._read_scope_from_request(context, Config(jwt_secret='secret-b')) is None + @staticmethod + def _http_rq_with_oauth_user(**access_token_kwargs) -> SimpleNamespace: + from keboola_mcp_server.oauth import ProxyAccessToken + + access_token = ProxyAccessToken( + token='opaque-access-token', + client_id='client-1', + scopes=[], + expires_at=None, + kbc_access_token='kbc_at_x', + **access_token_kwargs, + ) + user = AuthenticatedUser(access_token) + return SimpleNamespace(scope={'user': user}) + + def test_read_persisted_oauth_scope_builds_scope_when_confirmed(self) -> None: + http_rq = self._http_rq_with_oauth_user( + session_id='session-1', + scope_project_ids=[11, 22], + scope_read_only=True, + scope_confirmed=True, + scope_scoped_token='kbc_at_scoped', + scope_scoped_expires_at=datetime.fromtimestamp(1234.0, tz=timezone.utc), + ) + scope = SessionStateMiddleware._read_persisted_oauth_scope(http_rq) + assert scope == SessionScope( + project_ids=[11, 22], + read_only=True, + scoped_token='kbc_at_scoped', + scoped_expires_at=1234.0, + confirmed=True, + ) + + @pytest.mark.parametrize( + 'access_token_kwargs', + [ + {'scope_confirmed': False, 'scope_project_ids': [11]}, + {'scope_confirmed': True, 'scope_project_ids': None}, + ], + ids=['unconfirmed', 'no_project_ids'], + ) + def test_read_persisted_oauth_scope_returns_none_when_not_confirmed(self, access_token_kwargs: dict) -> None: + http_rq = self._http_rq_with_oauth_user(**access_token_kwargs) + assert SessionStateMiddleware._read_persisted_oauth_scope(http_rq) is None + + def test_read_persisted_oauth_scope_returns_none_for_non_oauth_request(self) -> None: + assert SessionStateMiddleware._read_persisted_oauth_scope(None) is None + assert SessionStateMiddleware._read_persisted_oauth_scope(SimpleNamespace(scope={})) is None + + def test_read_oauth_session_id_returns_session_id_for_oauth_request(self) -> None: + http_rq = self._http_rq_with_oauth_user(session_id='session-1') + assert SessionStateMiddleware._read_oauth_session_id(http_rq) == 'session-1' + + def test_read_oauth_session_id_returns_none_for_non_oauth_request(self) -> None: + assert SessionStateMiddleware._read_oauth_session_id(None) is None + assert SessionStateMiddleware._read_oauth_session_id(SimpleNamespace(scope={})) is None + @pytest.mark.asyncio async def test_on_list_tools_advertises_scope_token_unconditionally(self) -> None: # Unlike MultiProjectMiddleware's `project_ids` filter, this must show up even with no scope diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 4750165f0..c458e4d64 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -7,9 +7,9 @@ from pytest_mock import MockerFixture from keboola_mcp_server.clients.client import KeboolaClient -from keboola_mcp_server.config import Config, MetadataField +from keboola_mcp_server.config import Config, MetadataField, ServerRuntimeInfo from keboola_mcp_server.links import Link -from keboola_mcp_server.mcp import SCOPE_KEY, SessionScope, resolve_scope_secret +from keboola_mcp_server.mcp import OAUTH_SESSION_ID_KEY, SCOPE_KEY, ServerState, SessionScope, resolve_scope_secret from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, @@ -442,6 +442,36 @@ async def test_set_project_scope_subset_exchanges_and_stores( assert SessionScope.from_token(result.scope_token, resolve_scope_secret(Config())) == scope +@pytest.mark.asyncio +async def test_set_project_scope_persists_to_db_and_omits_scope_token_for_oauth_session( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + # An OAuth-authenticated session (OAUTH_SESSION_ID_KEY present) persists the scope on its + # oauth_sessions row instead of minting a scope_token -- the opaque OAuth access token already + # resolves back to that row on every subsequent call, so there's nothing left to resend. + _prep_client(mcp_context_client, mocker) + mcp_context_client.session.state[OAUTH_SESSION_ID_KEY] = 'session-1' + session_store = mocker.Mock() + session_store.update_scope = mocker.AsyncMock() + mcp_context_client.request_context.lifespan_context = ServerState( + config=Config(), runtime_info=ServerRuntimeInfo(transport='stdio'), session_store=session_store + ) + minted = SimpleNamespace(access_token='kbc_at_scoped', expires_at=1234.0, read_only=False) + mocker.patch('keboola_mcp_server.tools.project.exchange_scoped_token', new=mocker.AsyncMock(return_value=minted)) + + result = await set_project_scope(mcp_context_client, project_ids=[18, 83]) + + session_store.update_scope.assert_awaited_once() + call = session_store.update_scope.await_args + assert call.args == ('session-1',) + assert call.kwargs['project_ids'] == [18, 83] + assert call.kwargs['scoped_token'] == 'kbc_at_scoped' + assert call.kwargs['confirmed'] is True + # Nothing left for the caller to resend -- the server persisted the scope itself. + assert result.scope_token is None + assert 'no need to resend' in result.llm_instruction + + @pytest.mark.asyncio async def test_set_project_scope_all_introspects_then_exchanges( mcp_context_client: Context, mocker: MockerFixture diff --git a/uv.lock b/uv.lock index 3fe81c914..1e57ac8b3 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.77.0" +version = "1.78.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 0b28517987bdf0674515c66c56d12adf2e0a5d8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 09:44:52 +0200 Subject: [PATCH 63/89] feat(PSGO-261): tag merged fan-out items with their source project Closes the attribution gap flagged in the RFC's fan-out follow-up: multi-project read results merge every project's structured_content lists together, and until now nothing survived to say which project a given item came from once merged -- only the "=== project N ===" text envelope did, which structured-output clients don't parse. MultiProjectMiddleware._tag_items_with_project stamps _scope_project_id (not the RFC's originally-sketched source_project -- that name is already a real field on bucket/table models for Keboola's own linked-bucket provenance, and would have silently clobbered it) onto every dict item before _deep_merge concatenates the per-project lists. Only fires for genuine 2+ project fan-out; single-target calls already know their project unambiguously. Bump to 1.79.0 (new capability). --- feature_spec/pat_token_support/RFC.md | 26 +++++++++++++++++++ pyproject.toml | 2 +- src/keboola_mcp_server/mcp.py | 34 +++++++++++++++++++++++-- tests/test_mcp.py | 36 +++++++++++++++++++++++++++ uv.lock | 2 +- 5 files changed, 96 insertions(+), 4 deletions(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 98405047c..7bb632bbb 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -551,6 +551,32 @@ Concrete, with a 2-project read: Follow-up (not in this increment): make fan-out concurrent, catch per-project errors into a per-project `{project_id, ok|error}` envelope, and stamp `source_project` on merged rows. +## Resolved: structured_content attribution (PSGO-261, follow-up to the fan-out gap above) + +Error isolation shipped separately (`MultiProjectMiddleware.on_call_tool`'s per-project try/except, +collecting failures into retry-hint text notes rather than failing the whole call — see the code). +This closes the remaining half: attribution in `structured_content`. + +- **Field name is `_scope_project_id`, not `source_project` as originally sketched above.** + `source_project` is already a real field on bucket/table output models (`storage/tools.py:127,331`) + — Keboola's own cross-project *linked-bucket* provenance (which project a shared/linked bucket + originated from), a pre-existing and unrelated concept. Stamping that name here would have silently + overwritten real data on any linked bucket/table in a fanned-out result. `_scope_project_id` (leading + underscore, MCP-scope-specific name) avoids the collision; no output model in this codebase uses + that name today. +- **Mechanism:** `MultiProjectMiddleware._tag_items_with_project` stamps `_scope_project_id` onto every + dict item inside each project's structured payload, before `_deep_merge` concatenates the per-project + lists together — so the field survives the merge on every list item, not just the top level. + Non-dict list items (e.g. a plain list of ids) are left untouched — nothing to attribute. +- **Only applies to genuine fan-out (2+ targets).** A single-target call (scope of one, or narrowed to + one via `project_ids`) returns `call_next()` directly and never reaches `_merge` — it doesn't need + the tag, the whole session already knows which project it hit. +- **Schema safety:** no output model in this codebase sets `extra='forbid'` (`ConfigDict`), so no + generated JSON schema declares `additionalProperties: false` — adding this key doesn't violate any + existing tool's declared output schema. +- Text-content attribution (`=== project N ===`) is unchanged and still emitted alongside — this adds + the same information to `structured_content` for callers that only read that half of the result. + ## Tool gating: call-time, not list-time (why hide-then-reveal was reverted) We first tried **scope-first tool visibility**: while a programmatic session's scope was unconfirmed, diff --git a/pyproject.toml b/pyproject.toml index b60ed8ebc..fc17c0768 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.78.0" +version = "1.79.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index d19d0ae94..5e068c89c 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -1229,6 +1229,34 @@ def _truncate_lists(sc: Any, limit: int) -> Any: return sc[:limit] return sc + # Key stamped onto every dict item in a merged multi-project structured_content, so a client + # reading only structured_content (not the `=== project N ===` text envelope) can still tell + # which project an item came from once results are concatenated. Leading underscore + a name + # unlikely to collide with any real Keboola field (see PSGO-261 RFC addendum: merged-result + # project attribution). No output schema in this codebase sets extra='forbid'/additionalProperties: + # false, so an extra key here doesn't break schema validation for any existing tool. + _PROJECT_ATTRIBUTION_KEY = '_scope_project_id' + + @staticmethod + def _tag_items_with_project(sc: Any, project_id: int) -> Any: + """Stamps ``project_id`` onto every dict item in ``sc``'s top-level lists (non-dict items -- + e.g. a list of plain strings/ids -- are left alone; nothing to attribute). + """ + if not isinstance(sc, dict): + return sc + tagged = dict(sc) + for key, value in sc.items(): + if isinstance(value, list): + tagged[key] = [ + ( + {**item, MultiProjectMiddleware._PROJECT_ATTRIBUTION_KEY: project_id} + if isinstance(item, dict) + else item + ) + for item in value + ] + return tagged + @staticmethod def _merge(results: list[tuple[int, 'ToolResult']], errors: 'list[tuple[int, str]] | None' = None) -> 'ToolResult': # Deep-merge the per-project structured payloads into one schema-valid object (lists concatenated @@ -1247,7 +1275,7 @@ def _merge(results: list[tuple[int, 'ToolResult']], errors: 'list[tuple[int, str per_project_counts: list[tuple[int, int]] = [] total_items = 0 for project_id, result in results: - sc = result.structured_content + sc = MultiProjectMiddleware._tag_items_with_project(result.structured_content, project_id) item_count = MultiProjectMiddleware._largest_list_len(sc) per_project_counts.append((project_id, item_count)) total_items += item_count @@ -1256,7 +1284,9 @@ def _merge(results: list[tuple[int, 'ToolResult']], errors: 'list[tuple[int, str sc if merged_structured is None else MultiProjectMiddleware._deep_merge(merged_structured, sc) ) - # Small enough: full detail with per-project text envelopes (attribution the model can read). + # Small enough: full detail, with per-project text envelopes AND a `_scope_project_id` on every + # merged structured_content item -- attribution survives whichever half of the result a caller + # actually reads. if total_items <= MultiProjectMiddleware._FANOUT_MAX_ITEMS: content: list[Any] = list(error_notes) for project_id, result in results: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 2c951c3eb..01cd969e4 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1574,9 +1574,45 @@ def _items_result(n: int) -> ToolResult: def test_merge_small_keeps_full_detail(self) -> None: merged = MultiProjectMiddleware._merge([(11, self._items_result(2)), (22, self._items_result(3))]) # Under the cap: per-project text envelopes + fully merged lists; counters summed. + # Non-dict list items (plain ints here) are left alone -- nothing to attribute. assert merged.structured_content == {'buckets': [0, 1, 0, 1, 2], 'total': 5} assert [c.text for c in merged.content] == ['=== project 11 ===', '2 items', '=== project 22 ===', '3 items'] + @staticmethod + def _dict_items_result(project_id: int, n: int) -> ToolResult: + return ToolResult( + content=[mt.TextContent(type='text', text=f'{n} items')], + structured_content={'tables': [{'id': f'p{project_id}-t{i}'} for i in range(n)], 'total': n}, + ) + + def test_tag_items_with_project_stamps_dict_items_only(self) -> None: + tagged = MultiProjectMiddleware._tag_items_with_project( + {'tables': [{'id': 't1'}, {'id': 't2'}], 'ids': [1, 2], 'total': 2}, project_id=42 + ) + assert tagged == { + 'tables': [{'id': 't1', '_scope_project_id': 42}, {'id': 't2', '_scope_project_id': 42}], + 'ids': [1, 2], # non-dict items untouched + 'total': 2, + } + + def test_tag_items_with_project_passes_through_non_dict_and_none(self) -> None: + assert MultiProjectMiddleware._tag_items_with_project(None, project_id=42) is None + assert MultiProjectMiddleware._tag_items_with_project([1, 2, 3], project_id=42) == [1, 2, 3] + + def test_merge_small_stamps_project_id_on_dict_items_in_structured_content(self) -> None: + # Attribution must survive a client that reads only structured_content, not the text envelope. + merged = MultiProjectMiddleware._merge( + [(11, self._dict_items_result(11, 2)), (22, self._dict_items_result(22, 1))] + ) + assert merged.structured_content == { + 'tables': [ + {'id': 'p11-t0', '_scope_project_id': 11}, + {'id': 'p11-t1', '_scope_project_id': 11}, + {'id': 'p22-t0', '_scope_project_id': 22}, + ], + 'total': 3, + } + @pytest.mark.asyncio async def test_fan_out_partial_failure_returns_successes_with_retry_hint(self) -> None: scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) diff --git a/uv.lock b/uv.lock index 1e57ac8b3..5227280f7 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.78.0" +version = "1.79.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 10e3b3d21bbe8f75219c6c0423d36738ca27d6a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 09:55:28 +0200 Subject: [PATCH 64/89] feat(PSGO-261): partition oauth_sessions by month, 2-month retention Resolves the RFC's "session expiry / cleanup" open question: rows never got deleted. Converts oauth_sessions to PARTITION BY RANGE (created_at), one partition per month, dropped wholesale once older than the 2-month retention window (an instant DROP TABLE, no vacuum needed, unlike a DELETE sweep). - Migration 0002 recreates the table as partitioned, copying existing rows across (safe: no production OAuth sessions exist on this schema yet). - session_store/retention.py's ensure_partitions() creates the current + next month's partition ahead of time (a RANGE insert with no matching partition raises immediately) and drops ones past retention. Idempotent, safe to have missed a run. - New `keboola-mcp-server gc-sessions` CLI subcommand wires this as a schedule-triggered job (kbc-stacks monthly CronJob, follow-up), separate from the deploy-triggered `migrate` hook. - Trade-off documented in the RFC: PostgreSQL requires partitioned-table UNIQUE/PK indexes to include the partition key, so access_token_hash / refresh_token_hash / id are only enforced unique per-partition now, not table-wide -- a cross-month collision on a 256-bit random token is cryptographically negligible. Verified end-to-end against the real docker-compose Postgres (migration, create/lookup/revoke, idempotent + first-run partition creation, stale partition drop). Bump to 1.80.0 (new capability). --- feature_spec/oauth_session_persistence/RFC.md | 41 +++++++- pyproject.toml | 2 +- src/keboola_mcp_server/cli.py | 34 +++++++ .../0002_partition_oauth_sessions.sql | 81 ++++++++++++++++ .../session_store/retention.py | 85 +++++++++++++++++ tests/session_store/test_migrator.py | 21 ++++- tests/session_store/test_retention.py | 93 +++++++++++++++++++ tests/test_cli.py | 70 +++++++++++++- uv.lock | 2 +- 9 files changed, 421 insertions(+), 8 deletions(-) create mode 100644 src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql create mode 100644 src/keboola_mcp_server/session_store/retention.py create mode 100644 tests/session_store/test_retention.py diff --git a/feature_spec/oauth_session_persistence/RFC.md b/feature_spec/oauth_session_persistence/RFC.md index c634bfed9..0f2ccf6c3 100644 --- a/feature_spec/oauth_session_persistence/RFC.md +++ b/feature_spec/oauth_session_persistence/RFC.md @@ -57,7 +57,7 @@ unaffected) or the header/PAT-supplied-token flow (still fully stateless, unaffe | Column | Type | Notes | |---|---|---| -| `id` | `uuid`, PK | Internal row id | +| `id` | `uuid`, PK | Internal row id. As of the partitioning resolution below, PK is `(id, created_at)` — required by `PARTITION BY RANGE (created_at)` — with a plain non-unique index kept on `id` alone for lookup speed | | `access_token_hash` | `bytea`, unique, indexed | `sha256` of the opaque access token the client holds — store the hash, not the token, so a DB read alone can't leak a live bearer credential | | `refresh_token_hash` | `bytea`, unique, indexed, nullable | Same, for the opaque refresh token | | `client_id` | `text` | The OAuth client (`claude.ai`, etc.) — audit/introspection only | @@ -235,9 +235,42 @@ its embedded expiry). would make future rotation non-disruptive if we want to add it later — flagging now so the column format (prefix the ciphertext with a key-version byte) is decided before Phase 1's migration ships, not retrofitted after real rows exist. -3. **Session expiry / cleanup.** Rows never get deleted automatically today's plan — need a retention - policy (e.g. delete rows with `revoked_at` set or `last_used_at` older than N days) — a cron/cleanup - job, explicitly out of scope for v1 but should be tracked as an immediate follow-up, not forgotten. +3. **Session expiry / cleanup — RESOLVED: monthly `RANGE` partitioning on `created_at`, 2-month + retention.** Rather than a `DELETE ... WHERE` sweep (which bloats the table with dead tuples until + `VACUUM` reclaims them, and gets slower as the table grows), `oauth_sessions` becomes a + `PARTITION BY RANGE (created_at)` table with one partition per month + (`oauth_sessions_YYYY_MM`). Dropping a whole month is an instant `DROP TABLE`, no vacuum needed. + + - **Retention window: 2 months.** E.g. in July, June's partition is the oldest kept; it gets + dropped once August starts (so at most 2 full months of session history ever exist). A session + genuinely still in use well past that keeps refreshing (`last_used_at`) but its *row* still ages + out with its creation month — acceptable, since a session that old should reasonably force + re-login rather than live forever; this isn't meant to be a durable audit log. + - **Maintenance is a monthly job, not per-request logic** (`session_store/retention.py`, + `ensure_partitions()`): each run (a) creates the partition for the *current* and *next* month if + missing — created ahead of time, not on first use, because a `RANGE`-partitioned `INSERT` with no + matching partition raises immediately, it does not fall through to a partition created moments + later — and (b) drops any partition whose entire month is older than the retention cutoff. + Idempotent (`CREATE TABLE IF NOT EXISTS` / `DROP TABLE IF EXISTS`-equivalent checks), safe to + re-run, safe to have missed a run or several (it always computes from "now", not from a + last-run watermark). + - **Wired via a new CLI subcommand + a monthly kbc-stacks CronJob** (`keboola-mcp-server + gc-sessions`), the same shape as the existing `migrate` pre-install/pre-upgrade hook Job but + schedule-triggered instead of deploy-triggered — deploys don't happen monthly on a + reliable cadence, so partition upkeep can't piggyback on the migration Job. + - **Trade-off, explicit not accidental: uniqueness is now per-partition, not table-wide.** + PostgreSQL requires a partitioned table's `UNIQUE` (and `PRIMARY KEY`) indexes to include the + partition key. `access_token_hash`/`refresh_token_hash` — and `id` itself — can therefore only be + enforced unique *within* a given month's partition, not globally across the whole table. A + same-hash collision across two different months on a 256-bit random token (`generate_opaque_token`, + `session_store/repository.py:20`) is cryptographically negligible (same reasoning already applied + to `id`'s `gen_random_uuid()`), so this is an acceptable relaxation of a guarantee that was never + meaningfully load-bearing to begin with — not a real security gap. + - **Migration (`0002_partition_oauth_sessions.sql`) recreates the table rather than converting it + in place**, copying any existing rows across into whichever partition (or the `DEFAULT` catch-all) + their `created_at` lands in. Safe because no production OAuth sessions exist on this schema yet + (dev/testing stacks only, as of this writing) — if that's no longer true when this ships, a + data-preserving rewrite would be needed instead of a straight copy. 4. **Does Postgres downtime take down OAuth login entirely, or degrade gracefully?** With no DB, no OAuth session can be created or validated — this is a new hard dependency for the OAuth path (by design, per Scope: "no silent in-memory fallback for a production auth path"). Confirm this is diff --git a/pyproject.toml b/pyproject.toml index fc17c0768..ef5a07250 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.79.0" +version = "1.80.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 93cb6c720..60adcf4c7 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -124,6 +124,13 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: 'Intended to run as a one-shot job before the server deployment rolls out.', ) + subparsers.add_parser( + 'gc-sessions', + help='Ensures upcoming oauth_sessions partitions exist and drops ones past the retention ' + 'window, then exits. Intended to run monthly (e.g. a kbc-stacks CronJob), independent of ' + 'deployments.', + ) + return parser.parse_args(args) @@ -262,6 +269,29 @@ async def _run_migrate() -> None: print('✓ Schema already up to date -- no migrations applied.') +async def _run_gc_sessions() -> None: + """Ensures upcoming oauth_sessions partitions exist and drops ones past the retention window, + then exits. Reads the DSN from the same env vars the server itself uses, so this can share the + exact same envFrom secret as the deployment (see cli.py's `migrate` command). + """ + import asyncpg + + from keboola_mcp_server.session_store.retention import ensure_partitions + + config = Config().replace_by(os.environ) + if not config.postgres_dsn: + raise RuntimeError('A Postgres DSN is required to run gc-sessions: set MCP_DB_URL (or KBC_POSTGRES_DSN).') + + pool = await asyncpg.create_pool(config.postgres_dsn) + try: + result = await ensure_partitions(pool) + finally: + await pool.close() + + created, dropped = result['created'], result['dropped'] + print(f"✓ Partitions created: {', '.join(created) or 'none'}; dropped: {', '.join(dropped) or 'none'}") + + async def run_server(args: list[str] | None = None) -> None: """Runs the MCP server in async mode.""" parsed_args = parse_args(args) @@ -308,6 +338,10 @@ async def run_server(args: list[str] | None = None) -> None: await _run_migrate() return + if parsed_args.command == 'gc-sessions': + await _run_gc_sessions() + return + # Create config from the CLI arguments config = Config( storage_api_url=parsed_args.api_url, diff --git a/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql b/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql new file mode 100644 index 000000000..5b1de30ff --- /dev/null +++ b/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql @@ -0,0 +1,81 @@ +-- Partitions oauth_sessions by month (RANGE on created_at) for time-boundable retention (RFC +-- oauth_session_persistence, "Session expiry / cleanup" open question). Dropping a whole month is +-- an instant DROP TABLE, no VACUUM needed, unlike a DELETE ... WHERE sweep. See +-- session_store/retention.py for the ongoing monthly maintenance (creates upcoming partitions +-- ahead of time, drops ones older than the retention window) -- this migration only performs the +-- one-time structural conversion. +-- +-- Trade-off (explicit, not accidental): PostgreSQL requires a partitioned table's UNIQUE/PRIMARY +-- KEY indexes to include the partition key. access_token_hash/refresh_token_hash/id can therefore +-- only be enforced unique WITHIN a partition (a calendar month), not table-wide. A same-hash +-- collision across two different months on a 256-bit random token is cryptographically negligible +-- -- an acceptable relaxation, not a real gap. +-- +-- Recreates the table rather than converting it in place, copying any existing rows across (they +-- land in whichever partition -- or the DEFAULT catch-all -- their created_at falls into). Safe +-- because no production OAuth sessions exist on this schema yet (dev/testing stacks only). + +ALTER TABLE oauth_sessions RENAME TO oauth_sessions_pre_partition; + +-- Index/constraint names are global per-schema, not per-table -- renaming the table alone leaves +-- these attached to it under their old names, colliding with the new table's indexes below. +ALTER TABLE oauth_sessions_pre_partition RENAME CONSTRAINT oauth_sessions_pkey TO oauth_sessions_pre_partition_pkey; +ALTER INDEX oauth_sessions_access_token_hash_idx RENAME TO oauth_sessions_pre_partition_access_token_hash_idx; +ALTER INDEX oauth_sessions_refresh_token_hash_idx RENAME TO oauth_sessions_pre_partition_refresh_token_hash_idx; + +CREATE TABLE oauth_sessions ( + id UUID NOT NULL DEFAULT gen_random_uuid(), + access_token_hash BYTEA NOT NULL, + refresh_token_hash BYTEA, + client_id TEXT NOT NULL, + user_email TEXT, + kbc_access_token_enc BYTEA NOT NULL, + kbc_refresh_token_enc BYTEA NOT NULL, + kbc_access_expires_at TIMESTAMPTZ NOT NULL, + scope_project_ids INTEGER[], + scope_read_only BOOLEAN NOT NULL DEFAULT FALSE, + scope_confirmed BOOLEAN NOT NULL DEFAULT FALSE, + scope_scoped_token_enc BYTEA, + scope_scoped_expires_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + last_used_at TIMESTAMPTZ NOT NULL DEFAULT now(), + revoked_at TIMESTAMPTZ, + PRIMARY KEY (id, created_at) +) PARTITION BY RANGE (created_at); + +-- Plain (non-unique) index on id alone: revoke()/rotate_kbc_tokens()/rotate_opaque_tokens() all +-- look up by id, and the composite PK above doesn't help a query that only has id. +CREATE INDEX oauth_sessions_id_idx ON oauth_sessions (id); + +CREATE UNIQUE INDEX oauth_sessions_access_token_hash_idx ON oauth_sessions (access_token_hash, created_at); +CREATE UNIQUE INDEX oauth_sessions_refresh_token_hash_idx ON oauth_sessions (refresh_token_hash, created_at) + WHERE refresh_token_hash IS NOT NULL; + +-- Catch-all for rows outside any explicit month partition -- notably the rows copied over from +-- oauth_sessions_pre_partition below, and a safety net if partition maintenance ever lags. Never +-- touched by session_store/retention.py's cleanup (only oauth_sessions_YYYY_MM names are). +CREATE TABLE oauth_sessions_default PARTITION OF oauth_sessions DEFAULT; + +-- This month's and next month's partitions, so writes never fail for lack of one -- a +-- RANGE-partitioned INSERT with no matching partition raises immediately, it does not fall +-- through to a partition created moments later. session_store/retention.py takes over creating +-- further-ahead partitions (and dropping old ones) every month after this. +DO $$ +DECLARE + this_month DATE := date_trunc('month', now()); + next_month DATE := this_month + INTERVAL '1 month'; +BEGIN + EXECUTE format( + 'CREATE TABLE oauth_sessions_%s PARTITION OF oauth_sessions FOR VALUES FROM (%L) TO (%L)', + to_char(this_month, 'YYYY_MM'), this_month, next_month + ); + EXECUTE format( + 'CREATE TABLE oauth_sessions_%s PARTITION OF oauth_sessions FOR VALUES FROM (%L) TO (%L)', + to_char(next_month, 'YYYY_MM'), next_month, next_month + INTERVAL '1 month' + ); +END $$; + +INSERT INTO oauth_sessions SELECT * FROM oauth_sessions_pre_partition; + +DROP TABLE oauth_sessions_pre_partition; diff --git a/src/keboola_mcp_server/session_store/retention.py b/src/keboola_mcp_server/session_store/retention.py new file mode 100644 index 000000000..46b8efdd4 --- /dev/null +++ b/src/keboola_mcp_server/session_store/retention.py @@ -0,0 +1,85 @@ +"""Monthly partition maintenance for oauth_sessions (RFC oauth_session_persistence, "Session +expiry / cleanup"). Two responsibilities, both idempotent and safe to re-run or to have missed a +run (each call computes everything from "now", not from a last-run watermark): + + - Ensure a partition exists for the current month and the next, so writes never fail for lack of + one -- a RANGE-partitioned INSERT with no matching partition raises immediately, it does not + fall through to a partition created moments later. + - Drop partitions whose entire month is older than the retention window. + +Intended to run as a monthly job (`keboola-mcp-server gc-sessions`), separate from the deploy-time +`migrate` command -- deploys don't happen on a reliable monthly cadence, so this can't piggyback +on that hook. +""" + +import logging +import re +from datetime import date, datetime, timezone + +import asyncpg + +LOG = logging.getLogger(__name__) + +DEFAULT_RETENTION_MONTHS = 2 + +_PARTITION_NAME_RE = re.compile(r'^oauth_sessions_(\d{4})_(\d{2})$') + + +def _month_start(d: date) -> date: + return d.replace(day=1) + + +def _add_months(d: date, n: int) -> date: + month_index = d.month - 1 + n + year = d.year + month_index // 12 + month = month_index % 12 + 1 + return date(year, month, 1) + + +def _partition_name(month_start: date) -> str: + return f'oauth_sessions_{month_start:%Y_%m}' + + +async def ensure_partitions( + pool: asyncpg.Pool, *, retention_months: int = DEFAULT_RETENTION_MONTHS +) -> dict[str, list[str]]: + """Creates this month's + next month's partition if missing; drops partitions entirely older + than ``retention_months`` back from the current month. + + :return: ``{'created': [...], 'dropped': [...]}`` partition names, for the CLI to report. + """ + this_month = _month_start(datetime.now(timezone.utc).date()) + cutoff = _add_months(this_month, -retention_months) + + created: list[str] = [] + async with pool.acquire() as conn: + for offset in (0, 1): + start = _add_months(this_month, offset) + end = _add_months(this_month, offset + 1) + name = _partition_name(start) + exists = await conn.fetchval('SELECT to_regclass($1) IS NOT NULL', name) + if not exists: + # DDL bounds can't be bound query parameters -- start/end are computed dates, not + # user input, so direct formatting here carries no injection risk. + await conn.execute( + f'CREATE TABLE {name} PARTITION OF oauth_sessions ' + f"FOR VALUES FROM ('{start.isoformat()}') TO ('{end.isoformat()}')" + ) + LOG.info(f'Created oauth_sessions partition: {name} [{start}, {end})') + created.append(name) + + rows = await conn.fetch( + "SELECT tablename FROM pg_tables WHERE tablename ~ '^oauth_sessions_[0-9]{4}_[0-9]{2}$'" + ) + dropped: list[str] = [] + for row in rows: + name = row['tablename'] + match = _PARTITION_NAME_RE.match(name) + assert match is not None # guaranteed by the query's own regex filter + partition_month = date(int(match.group(1)), int(match.group(2)), 1) + if partition_month < cutoff: + await conn.execute(f'DROP TABLE IF EXISTS {name}') + LOG.info(f'Dropped expired oauth_sessions partition: {name} (older than {cutoff})') + dropped.append(name) + + return {'created': created, 'dropped': dropped} diff --git a/tests/session_store/test_migrator.py b/tests/session_store/test_migrator.py index d8018cd2d..61226a0bb 100644 --- a/tests/session_store/test_migrator.py +++ b/tests/session_store/test_migrator.py @@ -21,7 +21,7 @@ async def test_applies_migrations_once() -> None: pool = await asyncpg.create_pool(TEST_DSN) try: applied = await apply_migrations(pool) - assert applied == ['0001_oauth_sessions.sql'] + assert applied == ['0001_oauth_sessions.sql', '0002_partition_oauth_sessions.sql'] # Re-running is a no-op -- the table already exists, so re-applying the DDL would fail # if the tracking table didn't correctly skip it. @@ -42,3 +42,22 @@ async def test_creates_oauth_sessions_table() -> None: assert {'access_token_hash', 'kbc_access_token_enc', 'scope_project_ids', 'revoked_at'} <= names finally: await pool.close() + + +async def test_partitions_by_month_with_current_and_next_ready() -> None: + pool = await asyncpg.create_pool(TEST_DSN) + try: + await apply_migrations(pool) + is_partitioned = await pool.fetchval("SELECT relkind = 'p' FROM pg_class WHERE relname = 'oauth_sessions'") + assert is_partitioned is True + + tables = { + r['tablename'] + for r in await pool.fetch("SELECT tablename FROM pg_tables WHERE tablename LIKE 'oauth_sessions%'") + } + # This month's + next month's partition exist immediately, plus the DEFAULT catch-all -- + # writes never fail for lack of a partition even before the monthly gc-sessions job runs. + assert 'oauth_sessions_default' in tables + assert sum(1 for t in tables if t not in ('oauth_sessions', 'oauth_sessions_default')) == 2 + finally: + await pool.close() diff --git a/tests/session_store/test_retention.py b/tests/session_store/test_retention.py new file mode 100644 index 000000000..5bc229912 --- /dev/null +++ b/tests/session_store/test_retention.py @@ -0,0 +1,93 @@ +from datetime import date + +import asyncpg +import pytest +import pytest_asyncio + +from keboola_mcp_server.session_store.migrator import apply_migrations +from keboola_mcp_server.session_store.retention import _add_months, _month_start, ensure_partitions +from tests.session_store.conftest import TEST_DSN, requires_postgres + + +@pytest.mark.parametrize( + ('start', 'n', 'expected'), + [ + (date(2026, 7, 15), 0, date(2026, 7, 1)), + (date(2026, 7, 1), 1, date(2026, 8, 1)), + (date(2026, 12, 1), 1, date(2027, 1, 1)), + (date(2026, 7, 1), -2, date(2026, 5, 1)), + (date(2026, 1, 1), -2, date(2025, 11, 1)), + ], +) +def test_add_months(start: date, n: int, expected: date) -> None: + assert _add_months(_month_start(start), n) == expected + + +@pytest.mark.asyncio +@requires_postgres +class TestEnsurePartitions: + @pytest_asyncio.fixture(autouse=True) + async def _clean_slate(self): + pool = await asyncpg.create_pool(TEST_DSN) + try: + await pool.execute('DROP TABLE IF EXISTS oauth_sessions, schema_migrations CASCADE') + await apply_migrations(pool) + finally: + await pool.close() + + @staticmethod + async def _existing_partitions(pool: asyncpg.Pool) -> set[str]: + rows = await pool.fetch( + "SELECT tablename FROM pg_tables WHERE tablename ~ '^oauth_sessions_[0-9]{4}_[0-9]{2}$'" + ) + return {r['tablename'] for r in rows} + + async def test_is_idempotent(self) -> None: + pool = await asyncpg.create_pool(TEST_DSN) + try: + # Migration 0002 already created this month's + next month's partition. + before = await self._existing_partitions(pool) + result = await ensure_partitions(pool) + assert result == {'created': [], 'dropped': []} + assert await self._existing_partitions(pool) == before + finally: + await pool.close() + + async def test_drops_only_partitions_older_than_retention(self) -> None: + pool = await asyncpg.create_pool(TEST_DSN) + try: + this_month = _month_start(date.today()) + stale = _add_months(this_month, -3) + kept = _add_months(this_month, -1) + for month_start in (stale, kept): + name = f'oauth_sessions_{month_start:%Y_%m}' + end = _add_months(month_start, 1) + await pool.execute( + f"CREATE TABLE {name} PARTITION OF oauth_sessions FOR VALUES FROM ('{month_start}') TO ('{end}')" + ) + + result = await ensure_partitions(pool, retention_months=2) + + assert result['dropped'] == [f'oauth_sessions_{stale:%Y_%m}'] + remaining = await self._existing_partitions(pool) + assert f'oauth_sessions_{stale:%Y_%m}' not in remaining + assert f'oauth_sessions_{kept:%Y_%m}' in remaining + finally: + await pool.close() + + async def test_creates_missing_current_and_next_month(self) -> None: + pool = await asyncpg.create_pool(TEST_DSN) + try: + # Simulate a fresh table with no partitions ensured yet. + for name in await self._existing_partitions(pool): + await pool.execute(f'DROP TABLE {name}') + + result = await ensure_partitions(pool) + + this_month = _month_start(date.today()) + next_month = _add_months(this_month, 1) + expected = {f'oauth_sessions_{this_month:%Y_%m}', f'oauth_sessions_{next_month:%Y_%m}'} + assert set(result['created']) == expected + assert await self._existing_partitions(pool) == expected + finally: + await pool.close() diff --git a/tests/test_cli.py b/tests/test_cli.py index a56beec9c..dcef87833 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,7 +2,7 @@ import pytest -from keboola_mcp_server.cli import _run_migrate, parse_args +from keboola_mcp_server.cli import _run_gc_sessions, _run_migrate, parse_args def test_parse_args_migrate() -> None: @@ -10,6 +10,11 @@ def test_parse_args_migrate() -> None: assert args.command == 'migrate' +def test_parse_args_gc_sessions() -> None: + args = parse_args(['gc-sessions']) + assert args.command == 'gc-sessions' + + class TestRunMigrate: @pytest.mark.asyncio async def test_requires_postgres_dsn(self, monkeypatch) -> None: @@ -67,3 +72,66 @@ async def test_closes_pool_even_if_migration_fails(self, monkeypatch) -> None: await _run_migrate() pool.close.assert_awaited_once() + + +class TestRunGcSessions: + @pytest.mark.asyncio + async def test_requires_postgres_dsn(self, monkeypatch) -> None: + monkeypatch.delenv('MCP_DB_URL', raising=False) + monkeypatch.delenv('KBC_POSTGRES_DSN', raising=False) + monkeypatch.delenv('KBC_MCP_DB_URL', raising=False) + with pytest.raises(RuntimeError, match='Postgres DSN'): + await _run_gc_sessions() + + @pytest.mark.asyncio + async def test_reports_created_and_dropped_partitions_and_closes_pool(self, monkeypatch, capsys) -> None: + monkeypatch.setenv('MCP_DB_URL', 'postgresql://u:p@host/db') + pool = MagicMock() + pool.close = AsyncMock() + with ( + patch('asyncpg.create_pool', AsyncMock(return_value=pool)) as create_pool, + patch( + 'keboola_mcp_server.session_store.retention.ensure_partitions', + AsyncMock(return_value={'created': ['oauth_sessions_2026_09'], 'dropped': ['oauth_sessions_2026_06']}), + ), + ): + await _run_gc_sessions() + + create_pool.assert_awaited_once_with('postgresql://u:p@host/db') + pool.close.assert_awaited_once() + out = capsys.readouterr().out + assert 'oauth_sessions_2026_09' in out + assert 'oauth_sessions_2026_06' in out + + @pytest.mark.asyncio + async def test_reports_none_when_nothing_changed(self, monkeypatch, capsys) -> None: + monkeypatch.setenv('MCP_DB_URL', 'postgresql://u:p@host/db') + pool = MagicMock() + pool.close = AsyncMock() + with ( + patch('asyncpg.create_pool', AsyncMock(return_value=pool)), + patch( + 'keboola_mcp_server.session_store.retention.ensure_partitions', + AsyncMock(return_value={'created': [], 'dropped': []}), + ), + ): + await _run_gc_sessions() + + assert 'none' in capsys.readouterr().out + + @pytest.mark.asyncio + async def test_closes_pool_even_if_it_fails(self, monkeypatch) -> None: + monkeypatch.setenv('MCP_DB_URL', 'postgresql://u:p@host/db') + pool = MagicMock() + pool.close = AsyncMock() + with ( + patch('asyncpg.create_pool', AsyncMock(return_value=pool)), + patch( + 'keboola_mcp_server.session_store.retention.ensure_partitions', + AsyncMock(side_effect=RuntimeError('boom')), + ), + ): + with pytest.raises(RuntimeError, match='boom'): + await _run_gc_sessions() + + pool.close.assert_awaited_once() diff --git a/uv.lock b/uv.lock index 5227280f7..272fe32fa 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.79.0" +version = "1.80.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From b6fc78259c58824f67ad4cb92bfb074cfc1fd278 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 10:35:29 +0200 Subject: [PATCH 65/89] refactor(PSGO-261): split mcp.py under 1k lines; address code-quality review Thermo-nuclear + ponytail review of this PR found mcp.py had grown from 613 to 1384 lines. Extracts two self-contained pieces into their own modules: - scope.py: SessionScope, resolve_scope_secret, and the scope_token/ SCOPE_KEY/OAUTH_SESSION_ID_KEY constants. - multiproject.py: MultiProjectMiddleware (fan-out + merge engine) in full. - build_tracing_headers moves to config.py (depends only on ServerRuntimeInfo) so multiproject.py doesn't need to import it back from mcp.py, avoiding a circular import. - BOOTSTRAP_TOOLS moves to tools/constants.py, the one constant genuinely shared between ToolsFilteringMiddleware (stays in mcp.py) and MultiProjectMiddleware (moved out). mcp.py: 1384 -> 947 lines. Test suite split to match (test_multiproject.py). Import sites (server.py, tools/project.py, tests) updated to import from wherever each symbol actually lives now, rather than re-exporting through mcp.py -- the honest fix, and the only way to avoid the cycle. Other findings addressed: - clients/auth_bridge.py: collapsed _AuthBridgeClient/_AuthBridgeExchangeError (one subclass each) into their single subclasses. - tools/project.py: collapsed set_project_scope's two near-identical except blocks (only the re-raise condition differed). - session_store/retention.py: dropped a redundant SQL regex (duplicating the Python one) and a can't-happen assert. - session_store/{migrations/0002,retention.py,cli.py}: migration 0002 no longer hand-rolls partition creation in a SQL DO block -- the `migrate` CLI command now calls ensure_partitions() right after applying migrations (also what the monthly gc-sessions job calls), one Python-side mechanism for all partition creation instead of two. - mcp.py: SessionStateMiddleware._read_scope_from_request no longer guards against context.message being absent (a required MiddlewareContext field, can't happen in production) -- only whether it HAS .arguments varies. Considered and NOT adopted (documented, not silently dropped): - Squashing migrations 0001+0002: both have almost certainly already run against live dev-stack Postgres databases; squashing risks stranding a stack caught mid-migration with no down-migration support. - Dropping crypto.py's key-version prefix byte: the RFC explicitly reasons about deciding this encoding before real rows exist, not retrofitting later -- removing it now would undo that already-made call. - Reducing postgres_dsn's 3 env-var aliases: matches the alias/prefix mechanism every other Config field already uses, not a bespoke pattern. All 1792 tests pass; black/isort/flake8/TOOLS.md clean. Bump to 1.81.0. --- integtests/conftest.py | 21 +- pyproject.toml | 2 +- src/keboola_mcp_server/cli.py | 7 + src/keboola_mcp_server/clients/auth_bridge.py | 18 +- src/keboola_mcp_server/config.py | 18 + src/keboola_mcp_server/mcp.py | 480 ++--------------- src/keboola_mcp_server/multiproject.py | 376 ++++++++++++++ src/keboola_mcp_server/scope.py | 80 +++ src/keboola_mcp_server/server.py | 9 +- .../0002_partition_oauth_sessions.sql | 27 +- .../session_store/retention.py | 7 +- src/keboola_mcp_server/tools/constants.py | 5 + src/keboola_mcp_server/tools/project.py | 24 +- tests/clients/test_client.py | 4 +- tests/session_store/test_migrator.py | 11 +- tests/session_store/test_retention.py | 2 +- tests/test_cli.py | 30 +- tests/test_mcp.py | 474 +---------------- tests/test_multiproject.py | 488 ++++++++++++++++++ tests/tools/test_project.py | 3 +- uv.lock | 2 +- 21 files changed, 1086 insertions(+), 1002 deletions(-) create mode 100644 src/keboola_mcp_server/multiproject.py create mode 100644 src/keboola_mcp_server/scope.py create mode 100644 tests/test_multiproject.py diff --git a/integtests/conftest.py b/integtests/conftest.py index 84c47264f..f6d8c5235 100644 --- a/integtests/conftest.py +++ b/integtests/conftest.py @@ -19,6 +19,8 @@ from mcp.shared.context import RequestContext from mcp.types import ClientCapabilities, Implementation, InitializeRequestParams +import keboola_mcp_server.mcp +import keboola_mcp_server.multiproject from integtests.project_lock import ( DEFAULT_MAX_WAIT_MINUTES, DEFAULT_POLL_INTERVAL_SECONDS, @@ -28,8 +30,8 @@ verify_project_endpoint, ) from keboola_mcp_server.clients.client import KeboolaClient -from keboola_mcp_server.config import Config, ServerRuntimeInfo -from keboola_mcp_server.mcp import ServerState, SessionStateMiddleware +from keboola_mcp_server.config import Config, ServerRuntimeInfo, build_tracing_headers +from keboola_mcp_server.mcp import ServerState from keboola_mcp_server.server import create_server from keboola_mcp_server.workspace import WorkspaceManager @@ -123,17 +125,20 @@ def _init_with_integtest_client_info(self, *args: Any, **kwargs: Any) -> None: @pytest.fixture(scope='session', autouse=True) def _patch_session_middleware_user_agent() -> Generator[None, None, None]: # Force a distinct User-Agent for outbound Keboola API requests during integration tests. + # build_tracing_headers is imported by name into both mcp.py (SessionStateMiddleware) and + # multiproject.py (MultiProjectMiddleware), so both call sites need patching -- patching only + # keboola_mcp_server.config.build_tracing_headers wouldn't affect either already-imported name. monkeypatch = pytest.MonkeyPatch() - original_get_headers = SessionStateMiddleware._get_headers.__func__ - def _get_headers_with_integtest_ua( - cls: type[SessionStateMiddleware], runtime_info: ServerRuntimeInfo - ) -> dict[str, Any]: - headers = original_get_headers(cls, runtime_info) + def _build_tracing_headers_with_integtest_ua(runtime_info: ServerRuntimeInfo) -> dict[str, Any]: + headers = build_tracing_headers(runtime_info) headers['User-Agent'] = INTEGTEST_USER_AGENT return headers - monkeypatch.setattr(SessionStateMiddleware, '_get_headers', classmethod(_get_headers_with_integtest_ua)) + monkeypatch.setattr(keboola_mcp_server.mcp, 'build_tracing_headers', _build_tracing_headers_with_integtest_ua) + monkeypatch.setattr( + keboola_mcp_server.multiproject, 'build_tracing_headers', _build_tracing_headers_with_integtest_ua + ) try: yield finally: diff --git a/pyproject.toml b/pyproject.toml index ef5a07250..eb8b9b0e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.80.0" +version = "1.81.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 60adcf4c7..2ca9385a4 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -252,6 +252,7 @@ async def _run_migrate() -> None: import asyncpg from keboola_mcp_server.session_store.migrator import apply_migrations + from keboola_mcp_server.session_store.retention import ensure_partitions config = Config().replace_by(os.environ) if not config.postgres_dsn: @@ -260,6 +261,10 @@ async def _run_migrate() -> None: pool = await asyncpg.create_pool(config.postgres_dsn) try: applied = await apply_migrations(pool) + # Bootstraps this month's + next month's oauth_sessions partition right after the schema + # exists, so the app never hits a RANGE-partitioned INSERT with no matching partition on + # first use -- the same call the monthly gc-sessions job makes on an ongoing basis. + partitions = await ensure_partitions(pool) finally: await pool.close() @@ -267,6 +272,8 @@ async def _run_migrate() -> None: print(f"✓ Applied {len(applied)} migration(s): {', '.join(applied)}") else: print('✓ Schema already up to date -- no migrations applied.') + if partitions['created']: + print(f"✓ Ensured oauth_sessions partitions: {', '.join(partitions['created'])}") async def _run_gc_sessions() -> None: diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index ac5cd304c..f93c7c0be 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -48,8 +48,8 @@ def is_programmatic_token(token: str | None) -> bool: return bare.startswith(_ACCESS_TOKEN_PREFIX) or bare.startswith(_PAT_PREFIX) -class _AuthBridgeExchangeError(RuntimeError): - """Base for auth-bridge exchange failures. +class OAuthTokenExchangeError(RuntimeError): + """Raised when the auth-bridge fails to exchange a league OAuth token for a programmatic session. :ivar status_code: The client-facing HTTP status (resolver 400/401/403 pass through; 5xx/timeout/network map to 502). @@ -63,8 +63,9 @@ def __str__(self) -> str: return self.args[0] -class _AuthBridgeClient: - """Shared setup for auth-bridge clients: base URL, SA-token path, timeout, transport.""" +class OAuthSessionExchanger: + """Exchanges a league OAuth access token (``claudai projectless`` scope) for a whole-stack + Keboola programmatic session (PSGO-261 oauth_session_exchange RFC).""" def __init__( self, @@ -89,15 +90,6 @@ def _read_sa_jwt(self) -> str: # Read per call — the kubelet rotates the projected token in place. return read_service_account_jwt(self._kubernetes_token_path) - -class OAuthTokenExchangeError(_AuthBridgeExchangeError): - """Raised when the auth-bridge fails to exchange a league OAuth token for a programmatic session.""" - - -class OAuthSessionExchanger(_AuthBridgeClient): - """Exchanges a league OAuth access token (``claudai projectless`` scope) for a whole-stack - Keboola programmatic session (PSGO-261 oauth_session_exchange RFC).""" - async def exchange(self, *, oauth_access_token: str) -> dict: """ Exchanges ``oauth_access_token`` for a ``CliTokenResponse`` (same shape as a PKCE login). diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 5724442ec..de68ee269 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -263,6 +263,24 @@ class ServerRuntimeInfo: """The version of the FastMCP library.""" +def build_tracing_headers(runtime_info: ServerRuntimeInfo) -> dict[str, Any]: + """Additional headers for requests made to Connection/downstream services, identifying this + MCP server for tracing. Depends only on ServerRuntimeInfo, so it lives here rather than in + mcp.py -- shared by SessionStateMiddleware and MultiProjectMiddleware's per-project client + construction, which live in separate modules.""" + return { + 'User-Agent': ( + f'Keboola MCP Server/{runtime_info.server_version} app_env={runtime_info.app_env} ' + f'transport={runtime_info.transport}' + ), + 'MCP-Server-Transport': runtime_info.transport or 'NA', + 'MCP-Server-Versions': ( + f'keboola-mcp-server/{runtime_info.server_version} mcp/{runtime_info.mcp_library_version} ' + f'fastmcp/{runtime_info.fastmcp_library_version}' + ), + } + + class MetadataField: """ Predefined names of Keboola metadata fields. diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 5e068c89c..cddb9974d 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -8,9 +8,7 @@ import asyncio import dataclasses import logging -import secrets import textwrap -import time from collections.abc import Awaitable, Callable, Iterable from typing import Any, TypeVar from unittest.mock import MagicMock @@ -18,16 +16,13 @@ import toon_format from fastmcp import Context, FastMCP from fastmcp.exceptions import ToolError -from fastmcp.exceptions import ValidationError as FastMCPValidationError from fastmcp.server import middleware as fmw from fastmcp.server.dependencies import get_http_request from fastmcp.server.middleware import CallNext, MiddlewareContext from fastmcp.tools import Tool -from fastmcp.tools.tool import ToolResult from mcp import types as mt from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from pydantic import BaseModel -from pydantic import ValidationError as PydanticValidationError from pydantic_core import to_json from starlette.applications import Starlette from starlette.requests import Request @@ -37,99 +32,33 @@ from keboola_mcp_server.clients.auth_bridge import is_programmatic_token, strip_bearer from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient -from keboola_mcp_server.config import Config, ServerRuntimeInfo, deployed_sa_token_path, is_same_stack +from keboola_mcp_server.config import ( + Config, + ServerRuntimeInfo, + build_tracing_headers, + deployed_sa_token_path, + is_same_stack, +) from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt from keboola_mcp_server.oauth import ProxyAccessToken +from keboola_mcp_server.scope import ( + OAUTH_SESSION_ID_KEY, + SCOPE_KEY, + SCOPE_TOKEN_ARG, + SessionScope, + resolve_scope_secret, +) from keboola_mcp_server.session_store.repository import SessionStore -from keboola_mcp_server.tools.constants import MODIFY_FLOW_TOOL_NAME, SEMANTIC_TOOLS_TAG, UPDATE_FLOW_TOOL_NAME +from keboola_mcp_server.tools.constants import ( + BOOTSTRAP_TOOLS, + MODIFY_FLOW_TOOL_NAME, + SEMANTIC_TOOLS_TAG, + UPDATE_FLOW_TOOL_NAME, +) from keboola_mcp_server.workspace import WorkspaceManager LOG = logging.getLogger(__name__) CONVERSATION_ID = 'conversation_id' -SCOPE_KEY = 'project_scope' - -# The OAuth session's DB row id (see session_store.repository.OAuthSession), stashed on -# ctx.session.state so set_project_scope can persist a newly-confirmed scope back to Postgres -# instead of only returning a scope_token. Absent for non-OAuth (PAT/header-token) sessions, which -# have no session row to persist against -- those keep relying on scope_token. -OAUTH_SESSION_ID_KEY = 'oauth_session_id' - -# Tools that must not be fanned out across multiple projects, even when a multi-project scope is -# active and they are read-only: the scope/auth tools operate on the whole-stack token (not a single -# project), and get_project_info resolves through the active project's WorkspaceManager (workspace id -# / sql dialect), so it reports the active project only. -# query_data is intentionally NOT here: the fan-out swaps a per-project WorkspaceManager (see -# MultiProjectMiddleware._swap_project) so a query runs against the workspace of each targeted -# project — narrow to one with the project_ids filter, or run across all scoped projects. -_NO_FANOUT_TOOLS = {'get_accessible_projects', 'set_project_scope', 'get_project_info'} - -# Tools allowed before the user has confirmed a project scope. Everything else is blocked with a -# message telling the assistant to ask the user which projects to work on first (ask-first UX). -_BOOTSTRAP_TOOLS = {'get_accessible_projects', 'set_project_scope'} - -# Optional per-call argument injected on fan-out-eligible read tools to restrict a single call to a -# subset of the scoped projects (consumed and stripped by MultiProjectMiddleware.on_call_tool). -_PROJECT_FILTER_ARG = 'project_ids' - -# Per-call argument that carries the confirmed multi-project scope forward (consumed and stripped -# by SessionStateMiddleware.on_request). See SessionScope.to_token/from_token: under the server's -# default stateless-HTTP transport a fresh, empty session is built for every request (the mcp -# 2026-07-28 RC formalizes this across the spec, dropping Mcp-Session-Id/session pinning entirely), -# so nothing survives in ctx.session.state between one tool call and the next -- on one replica or -# many, even within a single process. A scope set via "set_project_scope" only persists if the -# caller resends the token it returned. -_SCOPE_TOKEN_ARG = 'scope_token' - -# Process-local fallback signing key for scope_token, used when no shared KBC_JWT_SECRET is -# configured (e.g. local stdio/login sessions). A per-process secret is enough there since a stdio -# process serves exactly one conversation end-to-end; deployed multi-replica setups already require -# a shared jwt_secret for the OAuth-provider JWTs (see oauth.py), which this reuses. -_FALLBACK_SCOPE_SECRET = secrets.token_hex(32) - - -def resolve_scope_secret(config: Config) -> str: - """The HMAC key used to sign/verify ``scope_token`` -- shared across replicas when - ``config.jwt_secret`` (``KBC_JWT_SECRET``) is configured, otherwise a process-local fallback.""" - return config.jwt_secret or _FALLBACK_SCOPE_SECRET - - -@dataclasses.dataclass(frozen=True) -class SessionScope: - """In-conversation multi-project scope (PSGO-261 increment 2). - - Persisted on the session across the per-request state rebuild. ``project_ids`` is the - user-selected set; ``scoped_token`` is the child access token minted by /v1/auth/pat/exchange - and narrowed to those projects (re-minted from the parent when near expiry). - """ - - project_ids: list[int] - read_only: bool = False - scoped_token: str | None = None - scoped_expires_at: float | None = None - confirmed: bool = False - """True once the user has explicitly chosen a scope via ``set_project_scope``. The default - auto-leased scope is unconfirmed, which gates data tools until the user decides.""" - - @property - def active_project_id(self) -> int | None: - return self.project_ids[0] if self.project_ids else None - - @property - def is_near_expiry(self) -> bool: - if self.scoped_expires_at is None: - return False - return time.time() >= (self.scoped_expires_at - 60) - - def to_token(self, secret: str) -> str: - """Signs this scope into the opaque ``scope_token`` a caller resends on later calls.""" - return encode_jwt(dataclasses.asdict(self), secret) - - @classmethod - def from_token(cls, token: str, secret: str) -> 'SessionScope': - """Inverse of ``to_token``. Raises on a missing/invalid/tampered token -- callers should - treat any exception as "no scope" rather than fail the request.""" - return cls(**decode_jwt(token, secret)) - R = TypeVar('R') T = TypeVar('T') @@ -397,10 +326,10 @@ async def on_list_tools( for tool in tools: params = dict(tool.parameters or {}) props = dict(params.get('properties') or {}) - if _SCOPE_TOKEN_ARG in props: + if SCOPE_TOKEN_ARG in props: patched.append(tool) continue - props[_SCOPE_TOKEN_ARG] = { + props[SCOPE_TOKEN_ARG] = { 'type': 'string', 'description': ( 'Opaque token returned by "set_project_scope" (also echoed by ' @@ -413,24 +342,6 @@ async def on_list_tools( patched.append(tool.model_copy(update={'parameters': params})) return patched - @classmethod - def _get_headers(cls, runtime_info: ServerRuntimeInfo) -> dict[str, Any]: - """ - :param runtime_info: Runtime information - :return: Additional headers for the requests used for tracing the MCP server - """ - return { - 'User-Agent': ( - f'Keboola MCP Server/{runtime_info.server_version} app_env={runtime_info.app_env} ' - f'transport={runtime_info.transport}' - ), - 'MCP-Server-Transport': runtime_info.transport or 'NA', - 'MCP-Server-Versions': ( - f'keboola-mcp-server/{runtime_info.server_version} mcp/{runtime_info.mcp_library_version} ' - f'fastmcp/{runtime_info.fastmcp_library_version}' - ), - } - @classmethod def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_storage_api_url: str | None) -> Config: """ @@ -488,10 +399,13 @@ def _read_scope_from_request(cls, context: fmw.MiddlewareContext[Any], config: C MultiProjectMiddleware then steers the caller back through get_accessible_projects / set_project_scope. """ - args = getattr(getattr(context, 'message', None), 'arguments', None) + # context.message always exists (a required MiddlewareContext field); only whether it HAS + # .arguments varies by request type (a ListToolsRequest has none, a CallToolRequestParams + # does), hence the single getattr here. + args = getattr(context.message, 'arguments', None) if not isinstance(args, dict): return None - token = args.pop(_SCOPE_TOKEN_ARG, None) + token = args.pop(SCOPE_TOKEN_ARG, None) if not token: return None try: @@ -719,7 +633,7 @@ async def create_session_state( storage_api_url=config.storage_api_url, storage_api_token=storage_token, bearer_token=bearer_token, - headers={**cls._get_headers(runtime_info), **extra_headers}, + headers={**build_tracing_headers(runtime_info), **extra_headers}, readonly=readonly, own_stack_storage_api_url=own_stack_storage_api_url, ).with_branch_id(config.branch_id) @@ -945,7 +859,7 @@ async def on_call_tool( # (X-KBC-ProjectId); calling it here pre-scope would 401 before the tool's own body (which # establishes that context, e.g. via introspect_token) ever runs. Mirrors the same exemption # in on_list_tools and MultiProjectMiddleware. - if tool.name in _BOOTSTRAP_TOOLS: + if tool.name in BOOTSTRAP_TOOLS: return await call_next(context) token_info = await self.get_token_info(context.fastmcp_context) @@ -971,342 +885,6 @@ async def on_call_tool( return await call_next(context) -class MultiProjectMiddleware(fmw.Middleware): - """Fans a read-only tool call out across every project in the active multi-project scope. - - Single-project (or no) scope is an unchanged passthrough. With >1 project selected, a read-only - tool runs once per project — the active ``KeboolaClient`` in session state is swapped to each - project's client and the per-project results are labelled with a per-project text envelope. Their - structured content is deep-merged (lists concatenated across projects, counters summed) into one - schema-valid object, degrading to count-first with a truncated sample past ``_FANOUT_MAX_ITEMS``. - Write tools never fan out: they target the active project only, so the agent can never write to - multiple projects without the user explicitly re-scoping (PSGO-261 decision D8). - """ - - async def on_call_tool( - self, - context: MiddlewareContext[mt.CallToolRequestParams], - call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult], - ) -> mt.CallToolResult: - ctx = context.fastmcp_context - state = ctx.session.state - scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None - name = context.message.name - - # Ask-first gate: until the user confirms a scope via set_project_scope, block data tools and - # tell the assistant to ask the user which projects to work on. Only applies when a scope has - # been auto-leased (local programmatic session); deployed/legacy sessions have no scope. - if isinstance(scope, SessionScope) and not scope.confirmed and name not in _BOOTSTRAP_TOOLS: - raise ToolError( - f'This session can access {len(scope.project_ids)} Keboola project(s), but no scope has ' - 'been confirmed yet. Call "get_accessible_projects", show the user their projects, and ask ' - 'whether to work across ALL of them or a subset. Then call "set_project_scope" ' - '(no arguments = all projects, or pass the chosen project ids, optionally read_only=true). ' - 'This confirmation is required once per session.' - ) - - # No auto-leased scope (deployed / legacy) or a bootstrap/scope tool: pass through untouched. - # Bootstrap tools own a real `project_ids` argument, so we must not strip it. - if not isinstance(scope, SessionScope) or name in _BOOTSTRAP_TOOLS: - return await call_next(context) - # Workspace-bound and write tools always target the active project (no fan-out, filter ignored). - if name in _NO_FANOUT_TOOLS: - return await call_next(context) - tool = await ctx.fastmcp.get_tool(name) - if not is_read_only_tool(tool): - return await call_next(context) - - # Read tool: consume the optional per-call project filter (advertised via on_list_tools) so the - # tool never receives it, then narrow this call's target projects to the requested subset. - requested = None - args = getattr(context.message, 'arguments', None) - if isinstance(args, dict): - requested = args.pop(_PROJECT_FILTER_ARG, None) - - targets = list(scope.project_ids) - if requested is not None: - # Omit the filter to run across the full scope; an explicit empty list is a caller mistake - # (it must not silently fall through to the whole scope). - if not requested: - raise ToolError( - f'"{_PROJECT_FILTER_ARG}" must be a non-empty list of project ids, ' - 'or omitted to run across the full scope.' - ) - outside = [p for p in requested if p not in scope.project_ids] - if outside: - raise ToolError( - f'Project(s) {outside} are outside the current scope {scope.project_ids}. ' - 'Call "set_project_scope" to change the scope first.' - ) - targets = [p for p in scope.project_ids if p in requested] - if not targets: - return await call_next(context) - - server_state = ServerState.from_context(ctx) - original_client = state.get(KeboolaClient.STATE_KEY) - original_workspace = state.get(WorkspaceManager.STATE_KEY) - is_real_client = isinstance(original_client, KeboolaClient) - # Default (auto-leased) scope carries no minted token; fall back to the active client's token. - base_token = scope.scoped_token or (original_client.token if is_real_client else '') - # The active client's own URL — the current request/session's, not the startup config's - # (which can differ or be unset for streamable-HTTP setups that supply it per request). - storage_api_url = original_client.storage_api_url if is_real_client else server_state.config.storage_api_url - - # A single target (scope of one, or narrowed to one via the filter) runs once against that - # project only — one call, that project's X-KBC-ProjectId, no per-project envelope. - if len(targets) == 1: - target = targets[0] - if target == scope.active_project_id: - return await call_next(context) - try: - await self._swap_project(state, server_state, storage_api_url, base_token, target, scope.read_only) - return await call_next(context) - finally: - state[KeboolaClient.STATE_KEY] = original_client - state[WorkspaceManager.STATE_KEY] = original_workspace - - results: list[tuple[int, ToolResult]] = [] - errors: list[tuple[int, str]] = [] - try: - for project_id in targets: - await self._swap_project(state, server_state, storage_api_url, base_token, project_id, scope.read_only) - # Isolate per-project failures: one project's error (e.g. Queue 401, a transient 5xx) - # must not discard the other projects' good results. Collect it and keep going, so the - # agent gets a partial response plus a retry hint. CancelledError is BaseException, so - # `except Exception` lets client cancellation propagate. - try: - results.append((project_id, await call_next(context))) - except (FastMCPValidationError, PydanticValidationError): - # Argument-level validation error: the same bad arguments fail identically in - # every project, so fanning out would emit N identical copies plus a confusing - # "failed for all N projects" aggregate. Abort and surface the single clean error. - raise - except Exception as e: - LOG.warning(f'Fan-out call failed for project {project_id}: {e}', exc_info=True) - errors.append((project_id, str(e))) - finally: - state[KeboolaClient.STATE_KEY] = original_client - state[WorkspaceManager.STATE_KEY] = original_workspace - - # Every project failed → nothing partial to return; surface a single aggregate error. - if not results and errors: - detail = '; '.join(f'project {pid}: {msg}' for pid, msg in errors) - raise ToolError(f'The tool failed for all {len(errors)} scoped project(s): {detail}') - - return self._merge(results, errors) - - async def on_list_tools( - self, - context: MiddlewareContext[mt.ListToolsRequest], - call_next: CallNext[mt.ListToolsRequest, list[Tool]], - ) -> list[Tool]: - # Advertise the optional per-call `project_ids` filter on fan-out-eligible read tools while a - # multi-project scope is active, so the assistant can target a subset (e.g. a single project) - # without changing the session scope. The value is consumed and stripped in on_call_tool. - tools = await call_next(context) - ctx = context.fastmcp_context - state = getattr(ctx.session, 'state', None) - scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None - - # NOTE: we intentionally do NOT hide data tools before a scope is confirmed. Hiding relied on - # the client re-fetching the tool list after notifications/tools/list_changed, which Claude Code - # (and others) don't do mid-session — that left the newly-unlocked tools invisible until a - # reconnect. Instead every tool stays listed and the call-time ask-first gate (on_call_tool) - # steers the user to set_project_scope first; once scoped, the already-listed tools just work. - if not (isinstance(scope, SessionScope) and scope.confirmed and len(scope.project_ids) > 1): - return tools - - patched: list[Tool] = [] - for tool in tools: - if tool.name in _BOOTSTRAP_TOOLS or tool.name in _NO_FANOUT_TOOLS or not is_read_only_tool(tool): - patched.append(tool) - continue - params = dict(tool.parameters or {}) - props = dict(params.get('properties') or {}) - if _PROJECT_FILTER_ARG in props: - patched.append(tool) - continue - props[_PROJECT_FILTER_ARG] = { - 'type': 'array', - 'items': {'type': 'integer'}, - 'description': ( - 'Optional. Restrict this call to these project ids (a subset of the confirmed ' - 'multi-project scope). Omit to run across all scoped projects.' - ), - } - params['properties'] = props - patched.append(tool.model_copy(update={'parameters': params})) - return patched - - @classmethod - async def _swap_project( - cls, - state: dict[str, Any], - server_state: ServerState, - storage_api_url: str, - base_token: str, - project_id: int, - read_only: bool, - ) -> None: - """Points the session state at `project_id` for the duration of one fanned-out tool call. - - Swaps in a per-project `KeboolaClient` AND a `WorkspaceManager` built on it, so - workspace-bound reads (query_data) run against *this* project's workspace rather than the - active project's. The workspace is provisioned lazily on first use per project. - Note: rebuilt per call; caching across calls would need a store that survives the - per-request state rebuild — add if provisioning latency shows up in practice. - """ - client = await cls.client_for_project(server_state, storage_api_url, base_token, project_id, read_only) - state[KeboolaClient.STATE_KEY] = client - state[WorkspaceManager.STATE_KEY] = await WorkspaceManager.create( - client, server_state.config.workspace_schema, kubernetes_token_path=deployed_sa_token_path() - ) - - @staticmethod - async def client_for_project( - server_state: ServerState, storage_api_url: str, token: str, project_id: int, read_only: bool - ) -> KeboolaClient: - # `storage_api_url` is the current request/session URL (e.g. the active `KeboolaClient`'s), - # not `server_state.config.storage_api_url` — that's the startup/lifespan config, which can - # differ (or be unset) for streamable-HTTP setups that supply the URL per request. - # Normalize any inbound `Bearer ` scheme; KeboolaClient adds it back for bearer tokens, - # so a pre-prefixed value would otherwise become `Authorization: Bearer Bearer …`. - token = strip_bearer(token) - return await KeboolaClient( - storage_api_url=storage_api_url, - storage_api_token=token, - bearer_token=token, - headers={ - **SessionStateMiddleware._get_headers(server_state.runtime_info), - 'X-KBC-ProjectId': str(project_id), - }, - readonly=read_only or None, - ).with_branch_id(None) - - @staticmethod - def _deep_merge(a: Any, b: Any) -> Any: - """Merges two per-project structured outputs so the result still validates the tool's schema. - - Lists are concatenated (the combined slice across projects), nested objects merged key by key, - and numeric counters summed; any other scalar keeps the first project's value. This keeps every - required field present with its declared type, so the merged object validates against the - single-project output schema. - """ - if isinstance(a, list) and isinstance(b, list): - return a + b - if isinstance(a, dict) and isinstance(b, dict): - merged = dict(a) - for key, value in b.items(): - merged[key] = MultiProjectMiddleware._deep_merge(a[key], value) if key in a else value - return merged - if isinstance(a, bool) or isinstance(b, bool): - return a - if isinstance(a, (int, float)) and isinstance(b, (int, float)): - return a + b # counters like search "total" - return a - - # Total list items across projects before a fanned-out result degrades to count-first: instead of - # dumping every project's full listing (which, on big projects, overflows the context window in a - # single tool result), return per-project counts + a truncated sample + guidance to narrow. Small - # multi-project results stay fully detailed. Class attribute so tests can lower it. - _FANOUT_MAX_ITEMS = 200 - - @staticmethod - def _largest_list_len(sc: Any) -> int: - """Item count of a structured payload = the length of its largest top-level list (buckets/tables/hits).""" - if isinstance(sc, dict): - return max((len(v) for v in sc.values() if isinstance(v, list)), default=0) - if isinstance(sc, list): - return len(sc) - return 0 - - @staticmethod - def _truncate_lists(sc: Any, limit: int) -> Any: - """Truncate every top-level list to `limit` (schema-safe: a shorter list still validates).""" - if isinstance(sc, dict): - return {k: (v[:limit] if isinstance(v, list) else v) for k, v in sc.items()} - if isinstance(sc, list): - return sc[:limit] - return sc - - # Key stamped onto every dict item in a merged multi-project structured_content, so a client - # reading only structured_content (not the `=== project N ===` text envelope) can still tell - # which project an item came from once results are concatenated. Leading underscore + a name - # unlikely to collide with any real Keboola field (see PSGO-261 RFC addendum: merged-result - # project attribution). No output schema in this codebase sets extra='forbid'/additionalProperties: - # false, so an extra key here doesn't break schema validation for any existing tool. - _PROJECT_ATTRIBUTION_KEY = '_scope_project_id' - - @staticmethod - def _tag_items_with_project(sc: Any, project_id: int) -> Any: - """Stamps ``project_id`` onto every dict item in ``sc``'s top-level lists (non-dict items -- - e.g. a list of plain strings/ids -- are left alone; nothing to attribute). - """ - if not isinstance(sc, dict): - return sc - tagged = dict(sc) - for key, value in sc.items(): - if isinstance(value, list): - tagged[key] = [ - ( - {**item, MultiProjectMiddleware._PROJECT_ATTRIBUTION_KEY: project_id} - if isinstance(item, dict) - else item - ) - for item in value - ] - return tagged - - @staticmethod - def _merge(results: list[tuple[int, 'ToolResult']], errors: 'list[tuple[int, str]] | None' = None) -> 'ToolResult': - # Deep-merge the per-project structured payloads into one schema-valid object (lists concatenated - # across projects). Counters (e.g. bucket_counts, search total) are summed by _deep_merge, so they - # keep reflecting the true totals even if the item lists get truncated below. - # Per-project failures (partial success) are surfaced as retry-hint notes in the text content, - # so the model can re-run just the failed project(s) via the project_ids filter. - error_notes = [ - mt.TextContent( - type='text', - text=f'project {pid} failed (retry with project_ids=[{pid}]): {msg}', - ) - for pid, msg in (errors or []) - ] - merged_structured: Any = None - per_project_counts: list[tuple[int, int]] = [] - total_items = 0 - for project_id, result in results: - sc = MultiProjectMiddleware._tag_items_with_project(result.structured_content, project_id) - item_count = MultiProjectMiddleware._largest_list_len(sc) - per_project_counts.append((project_id, item_count)) - total_items += item_count - if sc is not None: - merged_structured = ( - sc if merged_structured is None else MultiProjectMiddleware._deep_merge(merged_structured, sc) - ) - - # Small enough: full detail, with per-project text envelopes AND a `_scope_project_id` on every - # merged structured_content item -- attribution survives whichever half of the result a caller - # actually reads. - if total_items <= MultiProjectMiddleware._FANOUT_MAX_ITEMS: - content: list[Any] = list(error_notes) - for project_id, result in results: - content.append(mt.TextContent(type='text', text=f'=== project {project_id} ===')) - content.extend(result.content or []) - return ToolResult(content=content, structured_content=merged_structured) - - # Count-first: the combined listing is too large for one result. Return per-project counts, a - # truncated sample (first _FANOUT_MAX_ITEMS), and guidance — instead of every project's full dump. - summary = ', '.join(f'project {pid}: {n}' for pid, n in per_project_counts) - note = ( - f'Multi-project result is large — {total_items} items across {len(results)} project(s) ' - f'({summary}). Showing the first {MultiProjectMiddleware._FANOUT_MAX_ITEMS} in structured_content; ' - f'counters reflect the true totals. Narrow with project_ids=[...] on this tool, or use the ' - f'search tool to find specific items.' - ) - truncated = MultiProjectMiddleware._truncate_lists(merged_structured, MultiProjectMiddleware._FANOUT_MAX_ITEMS) - return ToolResult(content=error_notes + [mt.TextContent(type='text', text=note)], structured_content=truncated) - - def _to_python(data: Any, exclude_none: bool = True) -> Any | None: if isinstance(data, BaseModel): return data.model_dump(exclude_none=exclude_none, by_alias=False) diff --git a/src/keboola_mcp_server/multiproject.py b/src/keboola_mcp_server/multiproject.py new file mode 100644 index 000000000..164278f92 --- /dev/null +++ b/src/keboola_mcp_server/multiproject.py @@ -0,0 +1,376 @@ +"""Multi-project read fan-out (PSGO-261): ``MultiProjectMiddleware`` runs a read-only tool call +once per project in the active multi-project scope and merges the results. + +Split out of ``mcp.py`` to keep that module focused on the core middleware/server wiring. +""" + +import logging +from typing import Any + +from fastmcp.exceptions import ToolError +from fastmcp.exceptions import ValidationError as FastMCPValidationError +from fastmcp.server import middleware as fmw +from fastmcp.server.middleware import CallNext, MiddlewareContext +from fastmcp.tools import Tool +from fastmcp.tools.tool import ToolResult +from mcp import types as mt +from pydantic import ValidationError as PydanticValidationError + +from keboola_mcp_server.clients.auth_bridge import strip_bearer +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.config import build_tracing_headers, deployed_sa_token_path +from keboola_mcp_server.mcp import ServerState, is_read_only_tool +from keboola_mcp_server.scope import SCOPE_KEY, SessionScope +from keboola_mcp_server.tools.constants import BOOTSTRAP_TOOLS +from keboola_mcp_server.workspace import WorkspaceManager + +LOG = logging.getLogger(__name__) + +# Tools that must not be fanned out across multiple projects, even when a multi-project scope is +# active and they are read-only: the scope/auth tools operate on the whole-stack token (not a single +# project), and get_project_info resolves through the active project's WorkspaceManager (workspace id +# / sql dialect), so it reports the active project only. +# query_data is intentionally NOT here: the fan-out swaps a per-project WorkspaceManager (see +# MultiProjectMiddleware._swap_project) so a query runs against the workspace of each targeted +# project — narrow to one with the project_ids filter, or run across all scoped projects. +_NO_FANOUT_TOOLS = {'get_accessible_projects', 'set_project_scope', 'get_project_info'} + +# Optional per-call argument injected on fan-out-eligible read tools to restrict a single call to a +# subset of the scoped projects (consumed and stripped by MultiProjectMiddleware.on_call_tool). +_PROJECT_FILTER_ARG = 'project_ids' + + +class MultiProjectMiddleware(fmw.Middleware): + """Fans a read-only tool call out across every project in the active multi-project scope. + + Single-project (or no) scope is an unchanged passthrough. With >1 project selected, a read-only + tool runs once per project — the active ``KeboolaClient`` in session state is swapped to each + project's client and the per-project results are labelled with a per-project text envelope. Their + structured content is deep-merged (lists concatenated across projects, counters summed) into one + schema-valid object, degrading to count-first with a truncated sample past ``_FANOUT_MAX_ITEMS``. + Write tools never fan out: they target the active project only, so the agent can never write to + multiple projects without the user explicitly re-scoping (PSGO-261 decision D8). + """ + + async def on_call_tool( + self, + context: MiddlewareContext[mt.CallToolRequestParams], + call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult], + ) -> mt.CallToolResult: + ctx = context.fastmcp_context + state = ctx.session.state + scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None + name = context.message.name + + # Ask-first gate: until the user confirms a scope via set_project_scope, block data tools and + # tell the assistant to ask the user which projects to work on. Only applies when a scope has + # been auto-leased (local programmatic session); deployed/legacy sessions have no scope. + if isinstance(scope, SessionScope) and not scope.confirmed and name not in BOOTSTRAP_TOOLS: + raise ToolError( + f'This session can access {len(scope.project_ids)} Keboola project(s), but no scope has ' + 'been confirmed yet. Call "get_accessible_projects", show the user their projects, and ask ' + 'whether to work across ALL of them or a subset. Then call "set_project_scope" ' + '(no arguments = all projects, or pass the chosen project ids, optionally read_only=true). ' + 'This confirmation is required once per session.' + ) + + # No auto-leased scope (deployed / legacy) or a bootstrap/scope tool: pass through untouched. + # Bootstrap tools own a real `project_ids` argument, so we must not strip it. + if not isinstance(scope, SessionScope) or name in BOOTSTRAP_TOOLS: + return await call_next(context) + # Workspace-bound and write tools always target the active project (no fan-out, filter ignored). + if name in _NO_FANOUT_TOOLS: + return await call_next(context) + tool = await ctx.fastmcp.get_tool(name) + if not is_read_only_tool(tool): + return await call_next(context) + + # Read tool: consume the optional per-call project filter (advertised via on_list_tools) so the + # tool never receives it, then narrow this call's target projects to the requested subset. + requested = None + args = getattr(context.message, 'arguments', None) + if isinstance(args, dict): + requested = args.pop(_PROJECT_FILTER_ARG, None) + + targets = list(scope.project_ids) + if requested is not None: + # Omit the filter to run across the full scope; an explicit empty list is a caller mistake + # (it must not silently fall through to the whole scope). + if not requested: + raise ToolError( + f'"{_PROJECT_FILTER_ARG}" must be a non-empty list of project ids, ' + 'or omitted to run across the full scope.' + ) + outside = [p for p in requested if p not in scope.project_ids] + if outside: + raise ToolError( + f'Project(s) {outside} are outside the current scope {scope.project_ids}. ' + 'Call "set_project_scope" to change the scope first.' + ) + targets = [p for p in scope.project_ids if p in requested] + if not targets: + return await call_next(context) + + server_state = ServerState.from_context(ctx) + original_client = state.get(KeboolaClient.STATE_KEY) + original_workspace = state.get(WorkspaceManager.STATE_KEY) + is_real_client = isinstance(original_client, KeboolaClient) + # Default (auto-leased) scope carries no minted token; fall back to the active client's token. + base_token = scope.scoped_token or (original_client.token if is_real_client else '') + # The active client's own URL — the current request/session's, not the startup config's + # (which can differ or be unset for streamable-HTTP setups that supply it per request). + storage_api_url = original_client.storage_api_url if is_real_client else server_state.config.storage_api_url + + # A single target (scope of one, or narrowed to one via the filter) runs once against that + # project only — one call, that project's X-KBC-ProjectId, no per-project envelope. + if len(targets) == 1: + target = targets[0] + if target == scope.active_project_id: + return await call_next(context) + try: + await self._swap_project(state, server_state, storage_api_url, base_token, target, scope.read_only) + return await call_next(context) + finally: + state[KeboolaClient.STATE_KEY] = original_client + state[WorkspaceManager.STATE_KEY] = original_workspace + + results: list[tuple[int, ToolResult]] = [] + errors: list[tuple[int, str]] = [] + try: + for project_id in targets: + await self._swap_project(state, server_state, storage_api_url, base_token, project_id, scope.read_only) + # Isolate per-project failures: one project's error (e.g. Queue 401, a transient 5xx) + # must not discard the other projects' good results. Collect it and keep going, so the + # agent gets a partial response plus a retry hint. CancelledError is BaseException, so + # `except Exception` lets client cancellation propagate. + try: + results.append((project_id, await call_next(context))) + except (FastMCPValidationError, PydanticValidationError): + # Argument-level validation error: the same bad arguments fail identically in + # every project, so fanning out would emit N identical copies plus a confusing + # "failed for all N projects" aggregate. Abort and surface the single clean error. + raise + except Exception as e: + LOG.warning(f'Fan-out call failed for project {project_id}: {e}', exc_info=True) + errors.append((project_id, str(e))) + finally: + state[KeboolaClient.STATE_KEY] = original_client + state[WorkspaceManager.STATE_KEY] = original_workspace + + # Every project failed → nothing partial to return; surface a single aggregate error. + if not results and errors: + detail = '; '.join(f'project {pid}: {msg}' for pid, msg in errors) + raise ToolError(f'The tool failed for all {len(errors)} scoped project(s): {detail}') + + return self._merge(results, errors) + + async def on_list_tools( + self, + context: MiddlewareContext[mt.ListToolsRequest], + call_next: CallNext[mt.ListToolsRequest, list[Tool]], + ) -> list[Tool]: + # Advertise the optional per-call `project_ids` filter on fan-out-eligible read tools while a + # multi-project scope is active, so the assistant can target a subset (e.g. a single project) + # without changing the session scope. The value is consumed and stripped in on_call_tool. + tools = await call_next(context) + ctx = context.fastmcp_context + state = getattr(ctx.session, 'state', None) + scope = state.get(SCOPE_KEY) if isinstance(state, dict) else None + + # NOTE: we intentionally do NOT hide data tools before a scope is confirmed. Hiding relied on + # the client re-fetching the tool list after notifications/tools/list_changed, which Claude Code + # (and others) don't do mid-session — that left the newly-unlocked tools invisible until a + # reconnect. Instead every tool stays listed and the call-time ask-first gate (on_call_tool) + # steers the user to set_project_scope first; once scoped, the already-listed tools just work. + if not (isinstance(scope, SessionScope) and scope.confirmed and len(scope.project_ids) > 1): + return tools + + patched: list[Tool] = [] + for tool in tools: + if tool.name in BOOTSTRAP_TOOLS or tool.name in _NO_FANOUT_TOOLS or not is_read_only_tool(tool): + patched.append(tool) + continue + params = dict(tool.parameters or {}) + props = dict(params.get('properties') or {}) + if _PROJECT_FILTER_ARG in props: + patched.append(tool) + continue + props[_PROJECT_FILTER_ARG] = { + 'type': 'array', + 'items': {'type': 'integer'}, + 'description': ( + 'Optional. Restrict this call to these project ids (a subset of the confirmed ' + 'multi-project scope). Omit to run across all scoped projects.' + ), + } + params['properties'] = props + patched.append(tool.model_copy(update={'parameters': params})) + return patched + + @classmethod + async def _swap_project( + cls, + state: dict[str, Any], + server_state: ServerState, + storage_api_url: str, + base_token: str, + project_id: int, + read_only: bool, + ) -> None: + """Points the session state at `project_id` for the duration of one fanned-out tool call. + + Swaps in a per-project `KeboolaClient` AND a `WorkspaceManager` built on it, so + workspace-bound reads (query_data) run against *this* project's workspace rather than the + active project's. The workspace is provisioned lazily on first use per project. + Note: rebuilt per call; caching across calls would need a store that survives the + per-request state rebuild — add if provisioning latency shows up in practice. + """ + client = await cls.client_for_project(server_state, storage_api_url, base_token, project_id, read_only) + state[KeboolaClient.STATE_KEY] = client + state[WorkspaceManager.STATE_KEY] = await WorkspaceManager.create( + client, server_state.config.workspace_schema, kubernetes_token_path=deployed_sa_token_path() + ) + + @staticmethod + async def client_for_project( + server_state: ServerState, storage_api_url: str, token: str, project_id: int, read_only: bool + ) -> KeboolaClient: + # `storage_api_url` is the current request/session URL (e.g. the active `KeboolaClient`'s), + # not `server_state.config.storage_api_url` — that's the startup/lifespan config, which can + # differ (or be unset) for streamable-HTTP setups that supply the URL per request. + # Normalize any inbound `Bearer ` scheme; KeboolaClient adds it back for bearer tokens, + # so a pre-prefixed value would otherwise become `Authorization: Bearer Bearer …`. + token = strip_bearer(token) + return await KeboolaClient( + storage_api_url=storage_api_url, + storage_api_token=token, + bearer_token=token, + headers={ + **build_tracing_headers(server_state.runtime_info), + 'X-KBC-ProjectId': str(project_id), + }, + readonly=read_only or None, + ).with_branch_id(None) + + @staticmethod + def _deep_merge(a: Any, b: Any) -> Any: + """Merges two per-project structured outputs so the result still validates the tool's schema. + + Lists are concatenated (the combined slice across projects), nested objects merged key by key, + and numeric counters summed; any other scalar keeps the first project's value. This keeps every + required field present with its declared type, so the merged object validates against the + single-project output schema. + """ + if isinstance(a, list) and isinstance(b, list): + return a + b + if isinstance(a, dict) and isinstance(b, dict): + merged = dict(a) + for key, value in b.items(): + merged[key] = MultiProjectMiddleware._deep_merge(a[key], value) if key in a else value + return merged + if isinstance(a, bool) or isinstance(b, bool): + return a + if isinstance(a, (int, float)) and isinstance(b, (int, float)): + return a + b # counters like search "total" + return a + + # Total list items across projects before a fanned-out result degrades to count-first: instead of + # dumping every project's full listing (which, on big projects, overflows the context window in a + # single tool result), return per-project counts + a truncated sample + guidance to narrow. Small + # multi-project results stay fully detailed. Class attribute so tests can lower it. + _FANOUT_MAX_ITEMS = 200 + + @staticmethod + def _largest_list_len(sc: Any) -> int: + """Item count of a structured payload = the length of its largest top-level list (buckets/tables/hits).""" + if isinstance(sc, dict): + return max((len(v) for v in sc.values() if isinstance(v, list)), default=0) + if isinstance(sc, list): + return len(sc) + return 0 + + @staticmethod + def _truncate_lists(sc: Any, limit: int) -> Any: + """Truncate every top-level list to `limit` (schema-safe: a shorter list still validates).""" + if isinstance(sc, dict): + return {k: (v[:limit] if isinstance(v, list) else v) for k, v in sc.items()} + if isinstance(sc, list): + return sc[:limit] + return sc + + # Key stamped onto every dict item in a merged multi-project structured_content, so a client + # reading only structured_content (not the `=== project N ===` text envelope) can still tell + # which project an item came from once results are concatenated. Leading underscore + a name + # unlikely to collide with any real Keboola field (see PSGO-261 RFC addendum: merged-result + # project attribution). No output schema in this codebase sets extra='forbid'/additionalProperties: + # false, so an extra key here doesn't break schema validation for any existing tool. + _PROJECT_ATTRIBUTION_KEY = '_scope_project_id' + + @staticmethod + def _tag_items_with_project(sc: Any, project_id: int) -> Any: + """Stamps ``project_id`` onto every dict item in ``sc``'s top-level lists (non-dict items -- + e.g. a list of plain strings/ids -- are left alone; nothing to attribute). + """ + if not isinstance(sc, dict): + return sc + tagged = dict(sc) + for key, value in sc.items(): + if isinstance(value, list): + tagged[key] = [ + ( + {**item, MultiProjectMiddleware._PROJECT_ATTRIBUTION_KEY: project_id} + if isinstance(item, dict) + else item + ) + for item in value + ] + return tagged + + @staticmethod + def _merge(results: list[tuple[int, 'ToolResult']], errors: 'list[tuple[int, str]] | None' = None) -> 'ToolResult': + # Deep-merge the per-project structured payloads into one schema-valid object (lists concatenated + # across projects). Counters (e.g. bucket_counts, search total) are summed by _deep_merge, so they + # keep reflecting the true totals even if the item lists get truncated below. + # Per-project failures (partial success) are surfaced as retry-hint notes in the text content, + # so the model can re-run just the failed project(s) via the project_ids filter. + error_notes = [ + mt.TextContent( + type='text', + text=f'project {pid} failed (retry with project_ids=[{pid}]): {msg}', + ) + for pid, msg in (errors or []) + ] + merged_structured: Any = None + per_project_counts: list[tuple[int, int]] = [] + total_items = 0 + for project_id, result in results: + sc = MultiProjectMiddleware._tag_items_with_project(result.structured_content, project_id) + item_count = MultiProjectMiddleware._largest_list_len(sc) + per_project_counts.append((project_id, item_count)) + total_items += item_count + if sc is not None: + merged_structured = ( + sc if merged_structured is None else MultiProjectMiddleware._deep_merge(merged_structured, sc) + ) + + # Small enough: full detail, with per-project text envelopes AND a `_scope_project_id` on every + # merged structured_content item -- attribution survives whichever half of the result a caller + # actually reads. + if total_items <= MultiProjectMiddleware._FANOUT_MAX_ITEMS: + content: list[Any] = list(error_notes) + for project_id, result in results: + content.append(mt.TextContent(type='text', text=f'=== project {project_id} ===')) + content.extend(result.content or []) + return ToolResult(content=content, structured_content=merged_structured) + + # Count-first: the combined listing is too large for one result. Return per-project counts, a + # truncated sample (first _FANOUT_MAX_ITEMS), and guidance — instead of every project's full dump. + summary = ', '.join(f'project {pid}: {n}' for pid, n in per_project_counts) + note = ( + f'Multi-project result is large — {total_items} items across {len(results)} project(s) ' + f'({summary}). Showing the first {MultiProjectMiddleware._FANOUT_MAX_ITEMS} in structured_content; ' + f'counters reflect the true totals. Narrow with project_ids=[...] on this tool, or use the ' + f'search tool to find specific items.' + ) + truncated = MultiProjectMiddleware._truncate_lists(merged_structured, MultiProjectMiddleware._FANOUT_MAX_ITEMS) + return ToolResult(content=error_notes + [mt.TextContent(type='text', text=note)], structured_content=truncated) diff --git a/src/keboola_mcp_server/scope.py b/src/keboola_mcp_server/scope.py new file mode 100644 index 000000000..6096a059e --- /dev/null +++ b/src/keboola_mcp_server/scope.py @@ -0,0 +1,80 @@ +"""In-conversation multi-project scope (PSGO-261 increment 2): the ``SessionScope`` model, its +``scope_token`` JWT round-trip, and the associated session-state keys. + +Split out of ``mcp.py`` so that module can stay focused on the middleware/server wiring itself +(``mcp.py``'s ``SessionStateMiddleware``/``MultiProjectMiddleware`` both depend on this). +""" + +import dataclasses +import secrets +import time + +from keboola_mcp_server.config import Config +from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt + +SCOPE_KEY = 'project_scope' + +# The OAuth session's DB row id (see session_store.repository.OAuthSession), stashed on +# ctx.session.state so set_project_scope can persist a newly-confirmed scope back to Postgres +# instead of only returning a scope_token. Absent for non-OAuth (PAT/header-token) sessions, which +# have no session row to persist against -- those keep relying on scope_token. +OAUTH_SESSION_ID_KEY = 'oauth_session_id' + +# Per-call argument that carries the confirmed multi-project scope forward (consumed and stripped +# by SessionStateMiddleware.on_request). See SessionScope.to_token/from_token: under the server's +# default stateless-HTTP transport a fresh, empty session is built for every request (the mcp +# 2026-07-28 RC formalizes this across the spec, dropping Mcp-Session-Id/session pinning entirely), +# so nothing survives in ctx.session.state between one tool call and the next -- on one replica or +# many, even within a single process. A scope set via "set_project_scope" only persists if the +# caller resends the token it returned. +SCOPE_TOKEN_ARG = 'scope_token' + +# Process-local fallback signing key for scope_token, used when no shared KBC_JWT_SECRET is +# configured (e.g. local stdio/login sessions). A per-process secret is enough there since a stdio +# process serves exactly one conversation end-to-end; deployed multi-replica setups already require +# a shared jwt_secret for the OAuth-provider JWTs (see oauth.py), which this reuses. +_FALLBACK_SCOPE_SECRET = secrets.token_hex(32) + + +def resolve_scope_secret(config: Config) -> str: + """The HMAC key used to sign/verify ``scope_token`` -- shared across replicas when + ``config.jwt_secret`` (``KBC_JWT_SECRET``) is configured, otherwise a process-local fallback.""" + return config.jwt_secret or _FALLBACK_SCOPE_SECRET + + +@dataclasses.dataclass(frozen=True) +class SessionScope: + """In-conversation multi-project scope (PSGO-261 increment 2). + + Persisted on the session across the per-request state rebuild. ``project_ids`` is the + user-selected set; ``scoped_token`` is the child access token minted by /v1/auth/pat/exchange + and narrowed to those projects (re-minted from the parent when near expiry). + """ + + project_ids: list[int] + read_only: bool = False + scoped_token: str | None = None + scoped_expires_at: float | None = None + confirmed: bool = False + """True once the user has explicitly chosen a scope via ``set_project_scope``. The default + auto-leased scope is unconfirmed, which gates data tools until the user decides.""" + + @property + def active_project_id(self) -> int | None: + return self.project_ids[0] if self.project_ids else None + + @property + def is_near_expiry(self) -> bool: + if self.scoped_expires_at is None: + return False + return time.time() >= (self.scoped_expires_at - 60) + + def to_token(self, secret: str) -> str: + """Signs this scope into the opaque ``scope_token`` a caller resends on later calls.""" + return encode_jwt(dataclasses.asdict(self), secret) + + @classmethod + def from_token(cls, token: str, secret: str) -> 'SessionScope': + """Inverse of ``to_token``. Raises on a missing/invalid/tampered token -- callers should + treat any exception as "no scope" rather than fail the request.""" + return cls(**decode_jwt(token, secret)) diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index c3acf7f2d..00ee9ae4d 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -18,13 +18,8 @@ from keboola_mcp_server.authorization import ToolAuthorizationMiddleware from keboola_mcp_server.config import Config, ServerRuntimeInfo, Transport, get_env_storage_api_url from keboola_mcp_server.errors import ValidationErrorMiddleware -from keboola_mcp_server.mcp import ( - KeboolaMcpServer, - MultiProjectMiddleware, - ServerState, - SessionStateMiddleware, - ToolsFilteringMiddleware, -) +from keboola_mcp_server.mcp import KeboolaMcpServer, ServerState, SessionStateMiddleware, ToolsFilteringMiddleware +from keboola_mcp_server.multiproject import MultiProjectMiddleware from keboola_mcp_server.oauth import SimpleOAuthProvider from keboola_mcp_server.preview import preview_config_diff from keboola_mcp_server.prompts.add_prompts import add_keboola_prompts diff --git a/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql b/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql index 5b1de30ff..09bac860c 100644 --- a/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql +++ b/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql @@ -53,29 +53,14 @@ CREATE UNIQUE INDEX oauth_sessions_refresh_token_hash_idx ON oauth_sessions (ref WHERE refresh_token_hash IS NOT NULL; -- Catch-all for rows outside any explicit month partition -- notably the rows copied over from --- oauth_sessions_pre_partition below, and a safety net if partition maintenance ever lags. Never --- touched by session_store/retention.py's cleanup (only oauth_sessions_YYYY_MM names are). +-- oauth_sessions_pre_partition below (this migration creates no month partitions itself; the +-- `migrate` CLI command calls session_store.retention.ensure_partitions() right after applying +-- migrations, which is also what the monthly gc-sessions job calls -- one Python-side mechanism +-- for all partition creation instead of duplicating it here in SQL too). Also a safety net if +-- partition maintenance ever lags. Never touched by ensure_partitions()'s cleanup (only +-- oauth_sessions_YYYY_MM names are). CREATE TABLE oauth_sessions_default PARTITION OF oauth_sessions DEFAULT; --- This month's and next month's partitions, so writes never fail for lack of one -- a --- RANGE-partitioned INSERT with no matching partition raises immediately, it does not fall --- through to a partition created moments later. session_store/retention.py takes over creating --- further-ahead partitions (and dropping old ones) every month after this. -DO $$ -DECLARE - this_month DATE := date_trunc('month', now()); - next_month DATE := this_month + INTERVAL '1 month'; -BEGIN - EXECUTE format( - 'CREATE TABLE oauth_sessions_%s PARTITION OF oauth_sessions FOR VALUES FROM (%L) TO (%L)', - to_char(this_month, 'YYYY_MM'), this_month, next_month - ); - EXECUTE format( - 'CREATE TABLE oauth_sessions_%s PARTITION OF oauth_sessions FOR VALUES FROM (%L) TO (%L)', - to_char(next_month, 'YYYY_MM'), next_month, next_month + INTERVAL '1 month' - ); -END $$; - INSERT INTO oauth_sessions SELECT * FROM oauth_sessions_pre_partition; DROP TABLE oauth_sessions_pre_partition; diff --git a/src/keboola_mcp_server/session_store/retention.py b/src/keboola_mcp_server/session_store/retention.py index 46b8efdd4..a307f7c2c 100644 --- a/src/keboola_mcp_server/session_store/retention.py +++ b/src/keboola_mcp_server/session_store/retention.py @@ -68,14 +68,13 @@ async def ensure_partitions( LOG.info(f'Created oauth_sessions partition: {name} [{start}, {end})') created.append(name) - rows = await conn.fetch( - "SELECT tablename FROM pg_tables WHERE tablename ~ '^oauth_sessions_[0-9]{4}_[0-9]{2}$'" - ) + rows = await conn.fetch("SELECT tablename FROM pg_tables WHERE tablename LIKE 'oauth_sessions_%'") dropped: list[str] = [] for row in rows: name = row['tablename'] match = _PARTITION_NAME_RE.match(name) - assert match is not None # guaranteed by the query's own regex filter + if match is None: + continue # oauth_sessions_default / oauth_sessions_pre_partition -- not a month partition partition_month = date(int(match.group(1)), int(match.group(2)), 1) if partition_month < cutoff: await conn.execute(f'DROP TABLE IF EXISTS {name}') diff --git a/src/keboola_mcp_server/tools/constants.py b/src/keboola_mcp_server/tools/constants.py index 26ce74581..a904aa99b 100644 --- a/src/keboola_mcp_server/tools/constants.py +++ b/src/keboola_mcp_server/tools/constants.py @@ -2,6 +2,11 @@ UPDATE_FLOW_TOOL_NAME = 'update_flow' MODIFY_FLOW_TOOL_NAME = 'modify_flow' +# Tools allowed before the user has confirmed a project scope. Everything else is blocked with a +# message telling the assistant to ask the user which projects to work on first (ask-first UX). +# Shared by mcp.py's ToolsFilteringMiddleware and multiproject.py's MultiProjectMiddleware. +BOOTSTRAP_TOOLS = {'get_accessible_projects', 'set_project_scope'} + # Tag for tools supporting config diff preview feature CONFIG_DIFF_PREVIEW_TAG = 'config-diff-preview' diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 496c9638c..6ef85e073 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -16,16 +16,10 @@ from keboola_mcp_server.config import MetadataField, deployed_sa_token_path from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager -from keboola_mcp_server.mcp import ( - OAUTH_SESSION_ID_KEY, - SCOPE_KEY, - MultiProjectMiddleware, - ServerState, - SessionScope, - process_concurrently, - resolve_scope_secret, -) +from keboola_mcp_server.mcp import ServerState, process_concurrently +from keboola_mcp_server.multiproject import MultiProjectMiddleware from keboola_mcp_server.resources.prompts import get_project_system_prompt +from keboola_mcp_server.scope import OAUTH_SESSION_ID_KEY, SCOPE_KEY, SessionScope, resolve_scope_secret from keboola_mcp_server.workspace import WorkspaceManager LOG = logging.getLogger(__name__) @@ -557,17 +551,15 @@ async def set_project_scope( scoped_expires_at=minted.expires_at, confirmed=True, ) - except httpx.HTTPStatusError as e: - if e.response.status_code in (400, 401, 403): + except Exception as e: + if isinstance(e, httpx.HTTPStatusError) and e.response.status_code in (400, 401, 403): # Client error (bad project_ids, invalid/insufficient token): the input or auth is wrong, # not the exchange endpoint — surface it instead of silently downgrading to an unscoped # whole-stack token, which would mislead the caller about what was actually scoped. raise - LOG.warning('Scoped-token exchange failed; scoping with the whole-stack token instead.', exc_info=True) - scope = SessionScope(project_ids=ids, read_only=read_only, confirmed=True) - except Exception: - # Network/timeout/unavailable exchange endpoint: fall back so scoping still works, narrowed - # per request by X-KBC-ProjectId, without the extra token-scoping security narrowing. + # Any other failure (network/timeout/unavailable exchange endpoint, or a non-400/401/403 + # HTTP status): fall back so scoping still works, narrowed per request by X-KBC-ProjectId, + # without the extra token-scoping security narrowing. LOG.warning('Scoped-token exchange failed; scoping with the whole-stack token instead.', exc_info=True) scope = SessionScope(project_ids=ids, read_only=read_only, confirmed=True) ctx.session.state[SCOPE_KEY] = scope diff --git a/tests/clients/test_client.py b/tests/clients/test_client.py index 06aa255aa..9985d5c92 100644 --- a/tests/clients/test_client.py +++ b/tests/clients/test_client.py @@ -12,7 +12,7 @@ from keboola_mcp_server.clients.client import KeboolaClient, get_metadata_property from keboola_mcp_server.clients.storage import AsyncStorageClient from keboola_mcp_server.config import ServerRuntimeInfo -from keboola_mcp_server.mcp import SessionStateMiddleware +from keboola_mcp_server.mcp import build_tracing_headers @pytest.fixture @@ -345,7 +345,7 @@ def runtime_config(self) -> ServerRuntimeInfo: @pytest.fixture def keboola_client_with_headers(self, runtime_config: ServerRuntimeInfo) -> KeboolaClient: - headers = SessionStateMiddleware._get_headers(runtime_config) + headers = build_tracing_headers(runtime_config) return KeboolaClient( storage_api_url='https://connection.nowhere', storage_api_token='test-token', headers=headers ) diff --git a/tests/session_store/test_migrator.py b/tests/session_store/test_migrator.py index 61226a0bb..fb2d8f9a3 100644 --- a/tests/session_store/test_migrator.py +++ b/tests/session_store/test_migrator.py @@ -44,7 +44,11 @@ async def test_creates_oauth_sessions_table() -> None: await pool.close() -async def test_partitions_by_month_with_current_and_next_ready() -> None: +async def test_partitions_table_with_default_catch_all() -> None: + # Migration 0002 only creates the structure + a DEFAULT catch-all partition -- creating this + # month's/next month's partition is the `migrate` CLI's job (it calls + # session_store.retention.ensure_partitions() right after this), not the migration's. One + # Python-side mechanism for partition creation instead of duplicating it here in SQL too. pool = await asyncpg.create_pool(TEST_DSN) try: await apply_migrations(pool) @@ -55,9 +59,6 @@ async def test_partitions_by_month_with_current_and_next_ready() -> None: r['tablename'] for r in await pool.fetch("SELECT tablename FROM pg_tables WHERE tablename LIKE 'oauth_sessions%'") } - # This month's + next month's partition exist immediately, plus the DEFAULT catch-all -- - # writes never fail for lack of a partition even before the monthly gc-sessions job runs. - assert 'oauth_sessions_default' in tables - assert sum(1 for t in tables if t not in ('oauth_sessions', 'oauth_sessions_default')) == 2 + assert tables == {'oauth_sessions', 'oauth_sessions_default'} finally: await pool.close() diff --git a/tests/session_store/test_retention.py b/tests/session_store/test_retention.py index 5bc229912..07a777999 100644 --- a/tests/session_store/test_retention.py +++ b/tests/session_store/test_retention.py @@ -45,7 +45,7 @@ async def _existing_partitions(pool: asyncpg.Pool) -> set[str]: async def test_is_idempotent(self) -> None: pool = await asyncpg.create_pool(TEST_DSN) try: - # Migration 0002 already created this month's + next month's partition. + await ensure_partitions(pool) # first run creates this month's + next month's partition before = await self._existing_partitions(pool) result = await ensure_partitions(pool) assert result == {'created': [], 'dropped': []} diff --git a/tests/test_cli.py b/tests/test_cli.py index dcef87833..7796e9bbb 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -35,12 +35,18 @@ async def test_applies_migrations_and_closes_pool(self, monkeypatch, capsys) -> 'keboola_mcp_server.session_store.migrator.apply_migrations', AsyncMock(return_value=['0001_oauth_sessions.sql']), ), + patch( + 'keboola_mcp_server.session_store.retention.ensure_partitions', + AsyncMock(return_value={'created': ['oauth_sessions_2026_07'], 'dropped': []}), + ), ): await _run_migrate() create_pool.assert_awaited_once_with('postgresql://u:p@host/db') pool.close.assert_awaited_once() - assert '0001_oauth_sessions.sql' in capsys.readouterr().out + out = capsys.readouterr().out + assert '0001_oauth_sessions.sql' in out + assert 'oauth_sessions_2026_07' in out @pytest.mark.asyncio async def test_no_pending_migrations_still_closes_pool(self, monkeypatch, capsys) -> None: @@ -50,6 +56,10 @@ async def test_no_pending_migrations_still_closes_pool(self, monkeypatch, capsys with ( patch('asyncpg.create_pool', AsyncMock(return_value=pool)), patch('keboola_mcp_server.session_store.migrator.apply_migrations', AsyncMock(return_value=[])), + patch( + 'keboola_mcp_server.session_store.retention.ensure_partitions', + AsyncMock(return_value={'created': [], 'dropped': []}), + ), ): await _run_migrate() @@ -73,6 +83,24 @@ async def test_closes_pool_even_if_migration_fails(self, monkeypatch) -> None: pool.close.assert_awaited_once() + @pytest.mark.asyncio + async def test_closes_pool_even_if_partition_ensure_fails(self, monkeypatch) -> None: + monkeypatch.setenv('MCP_DB_URL', 'postgresql://u:p@host/db') + pool = MagicMock() + pool.close = AsyncMock() + with ( + patch('asyncpg.create_pool', AsyncMock(return_value=pool)), + patch('keboola_mcp_server.session_store.migrator.apply_migrations', AsyncMock(return_value=[])), + patch( + 'keboola_mcp_server.session_store.retention.ensure_partitions', + AsyncMock(side_effect=RuntimeError('boom')), + ), + ): + with pytest.raises(RuntimeError, match='boom'): + await _run_migrate() + + pool.close.assert_awaited_once() + class TestRunGcSessions: @pytest.mark.asyncio diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 01cd969e4..9fc3ad1f7 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -7,30 +7,24 @@ import pytest from fastmcp import Context from fastmcp.exceptions import ToolError -from fastmcp.tools.tool import ToolResult -from mcp import types as mt from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from pydantic import BaseModel, Field -from pydantic import ValidationError as PydanticValidationError from starlette.requests import Request from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config, ServerRuntimeInfo from keboola_mcp_server.mcp import ( - SCOPE_KEY, AggregateError, - MultiProjectMiddleware, ServerState, - SessionScope, SessionStateMiddleware, ToolsFilteringMiddleware, _exclude_none_serializer, _filter_toon_nulls, process_concurrently, - resolve_scope_secret, toon_serializer, unwrap_results, ) +from keboola_mcp_server.scope import SessionScope, resolve_scope_secret from keboola_mcp_server.workspace import WorkspaceManager @@ -736,7 +730,7 @@ async def test_on_request_branch_handling(self, method: str, expected_branch_id: ctx.session = session ctx.request_context.lifespan_context = server_state - context = SimpleNamespace(method=method, fastmcp_context=ctx) + context = SimpleNamespace(message=SimpleNamespace(), method=method, fastmcp_context=ctx) expected_result = object() async def call_next(_): @@ -1164,8 +1158,8 @@ def test_read_scope_from_request_returns_none_when_absent_or_invalid(self, argum assert SessionStateMiddleware._read_scope_from_request(context, config) is None def test_read_scope_from_request_ignores_non_call_tool_requests(self) -> None: - # tools/list (and other non-call requests) carry no `.arguments` at all. - context = SimpleNamespace(method='tools/list', fastmcp_context=None) + # tools/list (and other non-call requests) have a .message, just not one with .arguments. + context = SimpleNamespace(message=SimpleNamespace(), method='tools/list', fastmcp_context=None) assert SessionStateMiddleware._read_scope_from_request(context, Config()) is None def test_wrong_secret_falls_back_to_no_scope_via_read_scope_from_request(self) -> None: @@ -1248,463 +1242,3 @@ async def call_next(_): tools = await SessionStateMiddleware().on_list_tools(context, call_next) assert 'scope_token' in tools[0].parameters['properties'] - - -class TestMultiProjectMiddleware: - """Read tools fan out across the scoped projects; writes and single-project scope do not.""" - - @staticmethod - def _ctx(scope: SessionScope | None, tool_name: str, read_only: bool, arguments: dict | None = None): - state: dict = {KeboolaClient.STATE_KEY: 'orig-client'} - if scope is not None: - state[SCOPE_KEY] = scope - ctx = MagicMock(spec=Context) - ctx.session = SimpleNamespace(state=state) - ctx.request_context.lifespan_context = ServerState( - config=Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x'), - runtime_info=ServerRuntimeInfo(transport='stdio'), - ) - tool = MagicMock() - tool.name = tool_name - if read_only: - tool.annotations.readOnlyHint = True - else: - tool.annotations = None - ctx.fastmcp.get_tool = AsyncMock(return_value=tool) - message = SimpleNamespace(name=tool_name, arguments=arguments if arguments is not None else {}) - context = SimpleNamespace(message=message, fastmcp_context=ctx) - return context, state - - @staticmethod - def _result(text: str) -> ToolResult: - return ToolResult( - content=[mt.TextContent(type='text', text=text)], - structured_content={'rows': [text]}, - ) - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ('scope', 'tool_name', 'read_only'), - [ - (None, 'get_tables', True), - (SessionScope(project_ids=[11], confirmed=True), 'get_tables', True), - (SessionScope(project_ids=[11, 22], confirmed=True), 'update_config', False), # write: no fan-out - (SessionScope(project_ids=[11, 22], confirmed=True), 'get_project_info', True), # excluded tool - ], - ids=['no_scope', 'single_project', 'write_tool', 'excluded_tool'], - ) - async def test_passthrough_calls_once(self, scope, tool_name, read_only) -> None: - context, _ = self._ctx(scope, tool_name, read_only) - calls = [] - - async def call_next(_): - calls.append(1) - return self._result('single') - - result = await MultiProjectMiddleware().on_call_tool(context, call_next) - assert len(calls) == 1 - assert result.content[0].text == 'single' - - @pytest.mark.asyncio - async def test_unconfirmed_scope_blocks_data_tools(self) -> None: - # Default (auto-leased, unconfirmed) scope: data tools are gated with an ask-first message. - scope = SessionScope(project_ids=[11, 22], confirmed=False) - context, _ = self._ctx(scope, 'get_tables', read_only=True) - - async def call_next(_): - raise AssertionError('call_next must not run for a gated tool') - - with pytest.raises(ToolError, match='no scope has been confirmed'): - await MultiProjectMiddleware().on_call_tool(context, call_next) - - @pytest.mark.asyncio - async def test_unconfirmed_scope_allows_bootstrap_tools(self) -> None: - scope = SessionScope(project_ids=[11, 22], confirmed=False) - context, _ = self._ctx(scope, 'get_accessible_projects', read_only=True) - calls = [] - - async def call_next(_): - calls.append(1) - return self._result('projects') - - result = await MultiProjectMiddleware().on_call_tool(context, call_next) - assert len(calls) == 1 - assert result.content[0].text == 'projects' - - @pytest.mark.asyncio - async def test_read_tool_fans_out_per_project(self) -> None: - scope = SessionScope( - project_ids=[11, 22], scoped_token='kbc_at_s', scoped_expires_at=time.time() + 3600, confirmed=True - ) - context, state = self._ctx(scope, 'get_tables', read_only=True) - active_clients: list = [] - active_workspaces: list = [] - - async def call_next(_): - active_clients.append(state[KeboolaClient.STATE_KEY]) - active_workspaces.append(state[WorkspaceManager.STATE_KEY]) - return self._result('rows') - - with ( - patch.object( - MultiProjectMiddleware, - 'client_for_project', - AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), - ), - patch.object( - WorkspaceManager, - 'create', - AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), - ), - ): - result = await MultiProjectMiddleware().on_call_tool(context, call_next) - - # Ran once per project, each against that project's client AND workspace. - assert active_clients == ['client-11', 'client-22'] - assert active_workspaces == ['wsm-client-11', 'wsm-client-22'] - # Active client and workspace restored afterwards. - assert state[KeboolaClient.STATE_KEY] == 'orig-client' - assert state.get(WorkspaceManager.STATE_KEY) is None - # Per-project results are labelled in the text content. - texts = [c.text for c in result.content] - assert texts == ['=== project 11 ===', 'rows', '=== project 22 ===', 'rows'] - # Structured output is deep-merged (list fields concatenated) so it still validates the schema. - assert result.structured_content == {'rows': ['rows', 'rows']} - - @pytest.mark.asyncio - async def test_swap_project_uses_active_client_url_and_sa_token_path(self, monkeypatch) -> None: - # _swap_project must use the CURRENT request's Storage API URL (the active client's), not - # server_state.config's startup/lifespan URL, and must pass the deployed SA token path - # through to WorkspaceManager.create exactly like create_session_state does. - monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') - scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) - context, state = self._ctx(scope, 'get_tables', read_only=True) - # server_state.config carries a different (stale/absent) URL than the active request client. - state[KeboolaClient.STATE_KEY] = KeboolaClient( - storage_api_url='https://connection.request.keboola.com', storage_api_token='kbc_at_s' - ) - seen_calls: list = [] - - async def fake_client_for_project(_ss, storage_api_url, _token, pid, _ro): - seen_calls.append((storage_api_url, pid)) - return f'client-{pid}' - - async def call_next(_): - return self._result('rows') - - with ( - patch.object(MultiProjectMiddleware, 'client_for_project', AsyncMock(side_effect=fake_client_for_project)), - patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')) as ws_create, - ): - await MultiProjectMiddleware().on_call_tool(context, call_next) - - assert seen_calls == [ - ('https://connection.request.keboola.com', 11), - ('https://connection.request.keboola.com', 22), - ] - for call in ws_create.await_args_list: - assert call.kwargs.get('kubernetes_token_path') == '/var/run/secrets/token' - - @pytest.mark.asyncio - async def test_query_data_targets_single_project_workspace(self) -> None: - # query_data is no longer excluded: with the project_ids filter it runs once against that - # project's own workspace, so the user can query any scoped project without re-scoping. - scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) - context, state = self._ctx(scope, 'query_data', read_only=True, arguments={'project_ids': [22]}) - seen_workspaces: list = [] - - async def call_next(_): - seen_workspaces.append(state[WorkspaceManager.STATE_KEY]) - return self._result('csv') - - with ( - patch.object( - MultiProjectMiddleware, - 'client_for_project', - AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), - ), - patch.object( - WorkspaceManager, - 'create', - AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), - ), - ): - result = await MultiProjectMiddleware().on_call_tool(context, call_next) - - assert seen_workspaces == ['wsm-client-22'] # ran against project 22's workspace - assert result.content[0].text == 'csv' - - @pytest.mark.asyncio - async def test_project_filter_single_target_runs_once(self) -> None: - # project_ids filter narrows a multi-project scope to one project: one call, that project's - # client, and the filter is stripped from the arguments the tool receives. - scope = SessionScope(project_ids=[11, 22, 33], scoped_token='kbc_at_s', confirmed=True) - context, state = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [22]}) - seen_clients: list = [] - seen_workspaces: list = [] - - async def call_next(_): - seen_clients.append(state[KeboolaClient.STATE_KEY]) - seen_workspaces.append(state[WorkspaceManager.STATE_KEY]) - return self._result('t') - - with ( - patch.object( - MultiProjectMiddleware, - 'client_for_project', - AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), - ), - patch.object( - WorkspaceManager, - 'create', - AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), - ), - ): - result = await MultiProjectMiddleware().on_call_tool(context, call_next) - - assert seen_clients == ['client-22'] # ran once, against project 22 only - assert seen_workspaces == ['wsm-client-22'] # its own workspace - assert state[KeboolaClient.STATE_KEY] == 'orig-client' # restored - assert 'project_ids' not in context.message.arguments # stripped before the tool - assert result.content[0].text == 't' # raw single-project result, not an envelope - - @pytest.mark.asyncio - async def test_project_filter_subset_fans_out(self) -> None: - scope = SessionScope(project_ids=[11, 22, 33], scoped_token='kbc_at_s', confirmed=True) - context, state = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [11, 33]}) - seen_clients: list = [] - - async def call_next(_): - seen_clients.append(state[KeboolaClient.STATE_KEY]) - return self._result('t') - - with ( - patch.object( - MultiProjectMiddleware, - 'client_for_project', - AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), - ), - patch.object( - WorkspaceManager, - 'create', - AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), - ), - ): - await MultiProjectMiddleware().on_call_tool(context, call_next) - - assert seen_clients == ['client-11', 'client-33'] # only the requested subset, in scope order - - @pytest.mark.asyncio - async def test_project_filter_outside_scope_raises(self) -> None: - scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) - context, _ = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [99]}) - - async def call_next(_): - raise AssertionError('must not run for an out-of-scope filter') - - with pytest.raises(ToolError, match='outside the current scope'): - await MultiProjectMiddleware().on_call_tool(context, call_next) - - @pytest.mark.asyncio - async def test_on_list_tools_injects_project_filter(self) -> None: - scope = SessionScope(project_ids=[11, 22], confirmed=True) - # a read fan-out tool, an excluded tool, and a write tool - read_tool = _tool('get_tables', read_only=True) - read_tool.parameters = {'type': 'object', 'properties': {'bucket_ids': {'type': 'array'}}} - excluded = _tool('get_project_info', read_only=True) - excluded.parameters = {'type': 'object', 'properties': {}} - write_tool = _tool('update_config', read_only=False) - write_tool.parameters = {'type': 'object', 'properties': {}} - for t in (read_tool, excluded, write_tool): - t.model_copy = lambda update, _t=t: SimpleNamespace(name=_t.name, parameters=update['parameters']) - - context, _ = self._ctx(scope, 'x', read_only=True) - - async def call_next(_): - return [read_tool, excluded, write_tool] - - tools = await MultiProjectMiddleware().on_list_tools(context, call_next) - by_name = {t.name: t for t in tools} - assert 'project_ids' in by_name['get_tables'].parameters['properties'] - assert 'project_ids' not in by_name['get_project_info'].parameters['properties'] - assert 'project_ids' not in by_name['update_config'].parameters['properties'] - - @pytest.mark.asyncio - async def test_on_list_tools_unconfirmed_scope_lists_all_tools(self) -> None: - # Data tools are NOT hidden before scope is confirmed: hiding relied on the client re-fetching - # after tools/list_changed, which Claude Code doesn't do mid-session. All tools stay listed; - # the call-time ask-first gate steers to set_project_scope instead. - scope = SessionScope(project_ids=[11, 22], confirmed=False) - context, _ = self._ctx(scope, 'x', read_only=True) - - async def call_next(_): - return [ - _tool('get_accessible_projects', read_only=True), - _tool('set_project_scope', read_only=True), - _tool('get_tables', read_only=True), - _tool('update_config', read_only=False), - ] - - tools = await MultiProjectMiddleware().on_list_tools(context, call_next) - assert {t.name for t in tools} == { - 'get_accessible_projects', - 'set_project_scope', - 'get_tables', - 'update_config', - } - - @pytest.mark.asyncio - async def test_on_list_tools_no_scope_is_passthrough(self) -> None: - # Legacy Storage-token session (no SessionScope): every tool stays advertised, unchanged. - context, _ = self._ctx(None, 'x', read_only=True) - - async def call_next(_): - return [_tool('get_tables', read_only=True), _tool('update_config', read_only=False)] - - tools = await MultiProjectMiddleware().on_list_tools(context, call_next) - assert {t.name for t in tools} == {'get_tables', 'update_config'} - - @staticmethod - def _items_result(n: int) -> ToolResult: - return ToolResult( - content=[mt.TextContent(type='text', text=f'{n} items')], - structured_content={'buckets': list(range(n)), 'total': n}, - ) - - def test_merge_small_keeps_full_detail(self) -> None: - merged = MultiProjectMiddleware._merge([(11, self._items_result(2)), (22, self._items_result(3))]) - # Under the cap: per-project text envelopes + fully merged lists; counters summed. - # Non-dict list items (plain ints here) are left alone -- nothing to attribute. - assert merged.structured_content == {'buckets': [0, 1, 0, 1, 2], 'total': 5} - assert [c.text for c in merged.content] == ['=== project 11 ===', '2 items', '=== project 22 ===', '3 items'] - - @staticmethod - def _dict_items_result(project_id: int, n: int) -> ToolResult: - return ToolResult( - content=[mt.TextContent(type='text', text=f'{n} items')], - structured_content={'tables': [{'id': f'p{project_id}-t{i}'} for i in range(n)], 'total': n}, - ) - - def test_tag_items_with_project_stamps_dict_items_only(self) -> None: - tagged = MultiProjectMiddleware._tag_items_with_project( - {'tables': [{'id': 't1'}, {'id': 't2'}], 'ids': [1, 2], 'total': 2}, project_id=42 - ) - assert tagged == { - 'tables': [{'id': 't1', '_scope_project_id': 42}, {'id': 't2', '_scope_project_id': 42}], - 'ids': [1, 2], # non-dict items untouched - 'total': 2, - } - - def test_tag_items_with_project_passes_through_non_dict_and_none(self) -> None: - assert MultiProjectMiddleware._tag_items_with_project(None, project_id=42) is None - assert MultiProjectMiddleware._tag_items_with_project([1, 2, 3], project_id=42) == [1, 2, 3] - - def test_merge_small_stamps_project_id_on_dict_items_in_structured_content(self) -> None: - # Attribution must survive a client that reads only structured_content, not the text envelope. - merged = MultiProjectMiddleware._merge( - [(11, self._dict_items_result(11, 2)), (22, self._dict_items_result(22, 1))] - ) - assert merged.structured_content == { - 'tables': [ - {'id': 'p11-t0', '_scope_project_id': 11}, - {'id': 'p11-t1', '_scope_project_id': 11}, - {'id': 'p22-t0', '_scope_project_id': 22}, - ], - 'total': 3, - } - - @pytest.mark.asyncio - async def test_fan_out_partial_failure_returns_successes_with_retry_hint(self) -> None: - scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) - context, state = self._ctx(scope, 'get_tables', read_only=True) - - async def call_next(_): - if state[KeboolaClient.STATE_KEY] == 'client-22': - raise RuntimeError('boom-22') - return self._result('rows') - - with ( - patch.object( - MultiProjectMiddleware, - 'client_for_project', - AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), - ), - patch.object( - WorkspaceManager, - 'create', - AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), - ), - ): - result = await MultiProjectMiddleware().on_call_tool(context, call_next) - - # Project 11 succeeded; project 22's failure is a retry hint, not a total failure. - assert result.structured_content == {'rows': ['rows']} - texts = [c.text for c in result.content] - assert any('project 22 failed' in t and 'project_ids=[22]' in t for t in texts) - - @pytest.mark.asyncio - async def test_fan_out_all_failed_raises_aggregate(self) -> None: - scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) - context, _ = self._ctx(scope, 'get_tables', read_only=True) - - async def call_next(_): - raise RuntimeError('down') - - with ( - patch.object( - MultiProjectMiddleware, - 'client_for_project', - AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), - ), - patch.object( - WorkspaceManager, - 'create', - AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), - ), - ): - with pytest.raises(ToolError, match='failed for all 2 scoped'): - await MultiProjectMiddleware().on_call_tool(context, call_next) - - @pytest.mark.asyncio - async def test_fan_out_validation_error_raised_once_not_per_project(self) -> None: - # A bad argument (e.g. get_components with no component_ids) fails identically in every - # project, so it must surface as ONE clean validation error, not N copies + an aggregate. - scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) - context, state = self._ctx(scope, 'get_tables', read_only=True) - calls = [] - - async def call_next(_): - calls.append(state[KeboolaClient.STATE_KEY]) - raise PydanticValidationError.from_exception_data('get_tables', []) - - with ( - patch.object( - MultiProjectMiddleware, - 'client_for_project', - AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), - ), - patch.object( - WorkspaceManager, - 'create', - AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), - ), - ): - with pytest.raises(PydanticValidationError): - await MultiProjectMiddleware().on_call_tool(context, call_next) - - # Aborted after the first project; not retried across the rest. - assert calls == ['client-11'] - - def test_merge_large_degrades_to_count_first(self, monkeypatch) -> None: - # Lower the cap so a modest result trips the count-first path. - monkeypatch.setattr(MultiProjectMiddleware, '_FANOUT_MAX_ITEMS', 3) - merged = MultiProjectMiddleware._merge([(11, self._items_result(2)), (22, self._items_result(3))]) - # Single guidance note (no per-project full dump), lists truncated, counters preserved. - assert len(merged.content) == 1 - note = merged.content[0].text - assert 'project 11: 2' in note - assert 'project 22: 3' in note - assert 'search tool' in note - assert 'project_ids' in note - assert len(merged.structured_content['buckets']) == 3 # truncated to the cap - assert merged.structured_content['total'] == 5 # true total preserved diff --git a/tests/test_multiproject.py b/tests/test_multiproject.py new file mode 100644 index 000000000..a021ee569 --- /dev/null +++ b/tests/test_multiproject.py @@ -0,0 +1,488 @@ +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastmcp import Context +from fastmcp.exceptions import ToolError +from fastmcp.tools.tool import ToolResult +from mcp import types as mt +from pydantic import ValidationError as PydanticValidationError + +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.config import Config, ServerRuntimeInfo +from keboola_mcp_server.mcp import ServerState +from keboola_mcp_server.multiproject import MultiProjectMiddleware +from keboola_mcp_server.scope import SCOPE_KEY, SessionScope +from keboola_mcp_server.workspace import WorkspaceManager + + +def _tool(name: str, read_only: bool = False, tags: set[str] | None = None) -> MagicMock: + tool = MagicMock() + tool.name = name + tool.tags = tags or set() + if read_only: + tool.annotations.readOnlyHint = True + else: + tool.annotations = None + return tool + + +class TestMultiProjectMiddleware: + """Read tools fan out across the scoped projects; writes and single-project scope do not.""" + + @staticmethod + def _ctx(scope: SessionScope | None, tool_name: str, read_only: bool, arguments: dict | None = None): + state: dict = {KeboolaClient.STATE_KEY: 'orig-client'} + if scope is not None: + state[SCOPE_KEY] = scope + ctx = MagicMock(spec=Context) + ctx.session = SimpleNamespace(state=state) + ctx.request_context.lifespan_context = ServerState( + config=Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x'), + runtime_info=ServerRuntimeInfo(transport='stdio'), + ) + tool = MagicMock() + tool.name = tool_name + if read_only: + tool.annotations.readOnlyHint = True + else: + tool.annotations = None + ctx.fastmcp.get_tool = AsyncMock(return_value=tool) + message = SimpleNamespace(name=tool_name, arguments=arguments if arguments is not None else {}) + context = SimpleNamespace(message=message, fastmcp_context=ctx) + return context, state + + @staticmethod + def _result(text: str) -> ToolResult: + return ToolResult( + content=[mt.TextContent(type='text', text=text)], + structured_content={'rows': [text]}, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ('scope', 'tool_name', 'read_only'), + [ + (None, 'get_tables', True), + (SessionScope(project_ids=[11], confirmed=True), 'get_tables', True), + (SessionScope(project_ids=[11, 22], confirmed=True), 'update_config', False), # write: no fan-out + (SessionScope(project_ids=[11, 22], confirmed=True), 'get_project_info', True), # excluded tool + ], + ids=['no_scope', 'single_project', 'write_tool', 'excluded_tool'], + ) + async def test_passthrough_calls_once(self, scope, tool_name, read_only) -> None: + context, _ = self._ctx(scope, tool_name, read_only) + calls = [] + + async def call_next(_): + calls.append(1) + return self._result('single') + + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + assert len(calls) == 1 + assert result.content[0].text == 'single' + + @pytest.mark.asyncio + async def test_unconfirmed_scope_blocks_data_tools(self) -> None: + # Default (auto-leased, unconfirmed) scope: data tools are gated with an ask-first message. + scope = SessionScope(project_ids=[11, 22], confirmed=False) + context, _ = self._ctx(scope, 'get_tables', read_only=True) + + async def call_next(_): + raise AssertionError('call_next must not run for a gated tool') + + with pytest.raises(ToolError, match='no scope has been confirmed'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + @pytest.mark.asyncio + async def test_unconfirmed_scope_allows_bootstrap_tools(self) -> None: + scope = SessionScope(project_ids=[11, 22], confirmed=False) + context, _ = self._ctx(scope, 'get_accessible_projects', read_only=True) + calls = [] + + async def call_next(_): + calls.append(1) + return self._result('projects') + + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + assert len(calls) == 1 + assert result.content[0].text == 'projects' + + @pytest.mark.asyncio + async def test_read_tool_fans_out_per_project(self) -> None: + scope = SessionScope( + project_ids=[11, 22], scoped_token='kbc_at_s', scoped_expires_at=time.time() + 3600, confirmed=True + ) + context, state = self._ctx(scope, 'get_tables', read_only=True) + active_clients: list = [] + active_workspaces: list = [] + + async def call_next(_): + active_clients.append(state[KeboolaClient.STATE_KEY]) + active_workspaces.append(state[WorkspaceManager.STATE_KEY]) + return self._result('rows') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), + ), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + # Ran once per project, each against that project's client AND workspace. + assert active_clients == ['client-11', 'client-22'] + assert active_workspaces == ['wsm-client-11', 'wsm-client-22'] + # Active client and workspace restored afterwards. + assert state[KeboolaClient.STATE_KEY] == 'orig-client' + assert state.get(WorkspaceManager.STATE_KEY) is None + # Per-project results are labelled in the text content. + texts = [c.text for c in result.content] + assert texts == ['=== project 11 ===', 'rows', '=== project 22 ===', 'rows'] + # Structured output is deep-merged (list fields concatenated) so it still validates the schema. + assert result.structured_content == {'rows': ['rows', 'rows']} + + @pytest.mark.asyncio + async def test_swap_project_uses_active_client_url_and_sa_token_path(self, monkeypatch) -> None: + # _swap_project must use the CURRENT request's Storage API URL (the active client's), not + # server_state.config's startup/lifespan URL, and must pass the deployed SA token path + # through to WorkspaceManager.create exactly like create_session_state does. + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True) + # server_state.config carries a different (stale/absent) URL than the active request client. + state[KeboolaClient.STATE_KEY] = KeboolaClient( + storage_api_url='https://connection.request.keboola.com', storage_api_token='kbc_at_s' + ) + seen_calls: list = [] + + async def fake_client_for_project(_ss, storage_api_url, _token, pid, _ro): + seen_calls.append((storage_api_url, pid)) + return f'client-{pid}' + + async def call_next(_): + return self._result('rows') + + with ( + patch.object(MultiProjectMiddleware, 'client_for_project', AsyncMock(side_effect=fake_client_for_project)), + patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')) as ws_create, + ): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_calls == [ + ('https://connection.request.keboola.com', 11), + ('https://connection.request.keboola.com', 22), + ] + for call in ws_create.await_args_list: + assert call.kwargs.get('kubernetes_token_path') == '/var/run/secrets/token' + + @pytest.mark.asyncio + async def test_query_data_targets_single_project_workspace(self) -> None: + # query_data is no longer excluded: with the project_ids filter it runs once against that + # project's own workspace, so the user can query any scoped project without re-scoping. + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'query_data', read_only=True, arguments={'project_ids': [22]}) + seen_workspaces: list = [] + + async def call_next(_): + seen_workspaces.append(state[WorkspaceManager.STATE_KEY]) + return self._result('csv') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), + ), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_workspaces == ['wsm-client-22'] # ran against project 22's workspace + assert result.content[0].text == 'csv' + + @pytest.mark.asyncio + async def test_project_filter_single_target_runs_once(self) -> None: + # project_ids filter narrows a multi-project scope to one project: one call, that project's + # client, and the filter is stripped from the arguments the tool receives. + scope = SessionScope(project_ids=[11, 22, 33], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [22]}) + seen_clients: list = [] + seen_workspaces: list = [] + + async def call_next(_): + seen_clients.append(state[KeboolaClient.STATE_KEY]) + seen_workspaces.append(state[WorkspaceManager.STATE_KEY]) + return self._result('t') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), + ), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_clients == ['client-22'] # ran once, against project 22 only + assert seen_workspaces == ['wsm-client-22'] # its own workspace + assert state[KeboolaClient.STATE_KEY] == 'orig-client' # restored + assert 'project_ids' not in context.message.arguments # stripped before the tool + assert result.content[0].text == 't' # raw single-project result, not an envelope + + @pytest.mark.asyncio + async def test_project_filter_subset_fans_out(self) -> None: + scope = SessionScope(project_ids=[11, 22, 33], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [11, 33]}) + seen_clients: list = [] + + async def call_next(_): + seen_clients.append(state[KeboolaClient.STATE_KEY]) + return self._result('t') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), + ), + ): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert seen_clients == ['client-11', 'client-33'] # only the requested subset, in scope order + + @pytest.mark.asyncio + async def test_project_filter_outside_scope_raises(self) -> None: + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, _ = self._ctx(scope, 'get_tables', read_only=True, arguments={'project_ids': [99]}) + + async def call_next(_): + raise AssertionError('must not run for an out-of-scope filter') + + with pytest.raises(ToolError, match='outside the current scope'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + @pytest.mark.asyncio + async def test_on_list_tools_injects_project_filter(self) -> None: + scope = SessionScope(project_ids=[11, 22], confirmed=True) + # a read fan-out tool, an excluded tool, and a write tool + read_tool = _tool('get_tables', read_only=True) + read_tool.parameters = {'type': 'object', 'properties': {'bucket_ids': {'type': 'array'}}} + excluded = _tool('get_project_info', read_only=True) + excluded.parameters = {'type': 'object', 'properties': {}} + write_tool = _tool('update_config', read_only=False) + write_tool.parameters = {'type': 'object', 'properties': {}} + for t in (read_tool, excluded, write_tool): + t.model_copy = lambda update, _t=t: SimpleNamespace(name=_t.name, parameters=update['parameters']) + + context, _ = self._ctx(scope, 'x', read_only=True) + + async def call_next(_): + return [read_tool, excluded, write_tool] + + tools = await MultiProjectMiddleware().on_list_tools(context, call_next) + by_name = {t.name: t for t in tools} + assert 'project_ids' in by_name['get_tables'].parameters['properties'] + assert 'project_ids' not in by_name['get_project_info'].parameters['properties'] + assert 'project_ids' not in by_name['update_config'].parameters['properties'] + + @pytest.mark.asyncio + async def test_on_list_tools_unconfirmed_scope_lists_all_tools(self) -> None: + # Data tools are NOT hidden before scope is confirmed: hiding relied on the client re-fetching + # after tools/list_changed, which Claude Code doesn't do mid-session. All tools stay listed; + # the call-time ask-first gate steers to set_project_scope instead. + scope = SessionScope(project_ids=[11, 22], confirmed=False) + context, _ = self._ctx(scope, 'x', read_only=True) + + async def call_next(_): + return [ + _tool('get_accessible_projects', read_only=True), + _tool('set_project_scope', read_only=True), + _tool('get_tables', read_only=True), + _tool('update_config', read_only=False), + ] + + tools = await MultiProjectMiddleware().on_list_tools(context, call_next) + assert {t.name for t in tools} == { + 'get_accessible_projects', + 'set_project_scope', + 'get_tables', + 'update_config', + } + + @pytest.mark.asyncio + async def test_on_list_tools_no_scope_is_passthrough(self) -> None: + # Legacy Storage-token session (no SessionScope): every tool stays advertised, unchanged. + context, _ = self._ctx(None, 'x', read_only=True) + + async def call_next(_): + return [_tool('get_tables', read_only=True), _tool('update_config', read_only=False)] + + tools = await MultiProjectMiddleware().on_list_tools(context, call_next) + assert {t.name for t in tools} == {'get_tables', 'update_config'} + + @staticmethod + def _items_result(n: int) -> ToolResult: + return ToolResult( + content=[mt.TextContent(type='text', text=f'{n} items')], + structured_content={'buckets': list(range(n)), 'total': n}, + ) + + def test_merge_small_keeps_full_detail(self) -> None: + merged = MultiProjectMiddleware._merge([(11, self._items_result(2)), (22, self._items_result(3))]) + # Under the cap: per-project text envelopes + fully merged lists; counters summed. + # Non-dict list items (plain ints here) are left alone -- nothing to attribute. + assert merged.structured_content == {'buckets': [0, 1, 0, 1, 2], 'total': 5} + assert [c.text for c in merged.content] == ['=== project 11 ===', '2 items', '=== project 22 ===', '3 items'] + + @staticmethod + def _dict_items_result(project_id: int, n: int) -> ToolResult: + return ToolResult( + content=[mt.TextContent(type='text', text=f'{n} items')], + structured_content={'tables': [{'id': f'p{project_id}-t{i}'} for i in range(n)], 'total': n}, + ) + + def test_tag_items_with_project_stamps_dict_items_only(self) -> None: + tagged = MultiProjectMiddleware._tag_items_with_project( + {'tables': [{'id': 't1'}, {'id': 't2'}], 'ids': [1, 2], 'total': 2}, project_id=42 + ) + assert tagged == { + 'tables': [{'id': 't1', '_scope_project_id': 42}, {'id': 't2', '_scope_project_id': 42}], + 'ids': [1, 2], # non-dict items untouched + 'total': 2, + } + + def test_tag_items_with_project_passes_through_non_dict_and_none(self) -> None: + assert MultiProjectMiddleware._tag_items_with_project(None, project_id=42) is None + assert MultiProjectMiddleware._tag_items_with_project([1, 2, 3], project_id=42) == [1, 2, 3] + + def test_merge_small_stamps_project_id_on_dict_items_in_structured_content(self) -> None: + # Attribution must survive a client that reads only structured_content, not the text envelope. + merged = MultiProjectMiddleware._merge( + [(11, self._dict_items_result(11, 2)), (22, self._dict_items_result(22, 1))] + ) + assert merged.structured_content == { + 'tables': [ + {'id': 'p11-t0', '_scope_project_id': 11}, + {'id': 'p11-t1', '_scope_project_id': 11}, + {'id': 'p22-t0', '_scope_project_id': 22}, + ], + 'total': 3, + } + + @pytest.mark.asyncio + async def test_fan_out_partial_failure_returns_successes_with_retry_hint(self) -> None: + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True) + + async def call_next(_): + if state[KeboolaClient.STATE_KEY] == 'client-22': + raise RuntimeError('boom-22') + return self._result('rows') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), + ), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + # Project 11 succeeded; project 22's failure is a retry hint, not a total failure. + assert result.structured_content == {'rows': ['rows']} + texts = [c.text for c in result.content] + assert any('project 22 failed' in t and 'project_ids=[22]' in t for t in texts) + + @pytest.mark.asyncio + async def test_fan_out_all_failed_raises_aggregate(self) -> None: + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, _ = self._ctx(scope, 'get_tables', read_only=True) + + async def call_next(_): + raise RuntimeError('down') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), + ), + ): + with pytest.raises(ToolError, match='failed for all 2 scoped'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + @pytest.mark.asyncio + async def test_fan_out_validation_error_raised_once_not_per_project(self) -> None: + # A bad argument (e.g. get_components with no component_ids) fails identically in every + # project, so it must surface as ONE clean validation error, not N copies + an aggregate. + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_tables', read_only=True) + calls = [] + + async def call_next(_): + calls.append(state[KeboolaClient.STATE_KEY]) + raise PydanticValidationError.from_exception_data('get_tables', []) + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object( + WorkspaceManager, + 'create', + AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), + ), + ): + with pytest.raises(PydanticValidationError): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + # Aborted after the first project; not retried across the rest. + assert calls == ['client-11'] + + def test_merge_large_degrades_to_count_first(self, monkeypatch) -> None: + # Lower the cap so a modest result trips the count-first path. + monkeypatch.setattr(MultiProjectMiddleware, '_FANOUT_MAX_ITEMS', 3) + merged = MultiProjectMiddleware._merge([(11, self._items_result(2)), (22, self._items_result(3))]) + # Single guidance note (no per-project full dump), lists truncated, counters preserved. + assert len(merged.content) == 1 + note = merged.content[0].text + assert 'project 11: 2' in note + assert 'project 22: 3' in note + assert 'search tool' in note + assert 'project_ids' in note + assert len(merged.structured_content['buckets']) == 3 # truncated to the cap + assert merged.structured_content['total'] == 5 # true total preserved diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index c458e4d64..0ff3d4b2c 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -9,7 +9,8 @@ from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config, MetadataField, ServerRuntimeInfo from keboola_mcp_server.links import Link -from keboola_mcp_server.mcp import OAUTH_SESSION_ID_KEY, SCOPE_KEY, ServerState, SessionScope, resolve_scope_secret +from keboola_mcp_server.mcp import ServerState +from keboola_mcp_server.scope import OAUTH_SESSION_ID_KEY, SCOPE_KEY, SessionScope, resolve_scope_secret from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, diff --git a/uv.lock b/uv.lock index 272fe32fa..4af3a08a7 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.80.0" +version = "1.81.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 11d7d45b9ec0e8ea34fb649f9b7d5f65e8e50415 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 13:12:12 +0200 Subject: [PATCH 66/89] fix(PSGO-261): keep exactly retention_months of oauth_sessions partitions ensure_partitions() computed cutoff as this_month - retention_months, which keeps retention_months+1 months of data (e.g. retention_months=2 kept 3 months: current + 2 prior). Fix the off-by-one so retention_months counts the current month too, matching the "2-month retention" the RFC documents. --- pyproject.toml | 2 +- src/keboola_mcp_server/cli.py | 6 +++--- .../session_store/retention.py | 17 +++++++++------- tests/session_store/test_retention.py | 20 +++++++++++++++++++ uv.lock | 2 +- 5 files changed, 35 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index eb8b9b0e7..2ac3c64bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.81.0" +version = "1.82.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 2ca9385a4..e0767a3c2 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -127,8 +127,8 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: subparsers.add_parser( 'gc-sessions', help='Ensures upcoming oauth_sessions partitions exist and drops ones past the retention ' - 'window, then exits. Intended to run monthly (e.g. a kbc-stacks CronJob), independent of ' - 'deployments.', + 'window, then exits. Intended to run on a recurring schedule (e.g. a kbc-stacks CronJob), ' + 'independent of deployments.', ) return parser.parse_args(args) @@ -263,7 +263,7 @@ async def _run_migrate() -> None: applied = await apply_migrations(pool) # Bootstraps this month's + next month's oauth_sessions partition right after the schema # exists, so the app never hits a RANGE-partitioned INSERT with no matching partition on - # first use -- the same call the monthly gc-sessions job makes on an ongoing basis. + # first use -- the same call the recurring gc-sessions job makes on an ongoing basis. partitions = await ensure_partitions(pool) finally: await pool.close() diff --git a/src/keboola_mcp_server/session_store/retention.py b/src/keboola_mcp_server/session_store/retention.py index a307f7c2c..64b501aad 100644 --- a/src/keboola_mcp_server/session_store/retention.py +++ b/src/keboola_mcp_server/session_store/retention.py @@ -1,15 +1,16 @@ -"""Monthly partition maintenance for oauth_sessions (RFC oauth_session_persistence, "Session -expiry / cleanup"). Two responsibilities, both idempotent and safe to re-run or to have missed a -run (each call computes everything from "now", not from a last-run watermark): +"""Partition maintenance for oauth_sessions (RFC oauth_session_persistence, "Session expiry / +cleanup"). Two responsibilities, both idempotent and safe to re-run or to have missed a run (each +call computes everything from "now", not from a last-run watermark): - Ensure a partition exists for the current month and the next, so writes never fail for lack of one -- a RANGE-partitioned INSERT with no matching partition raises immediately, it does not fall through to a partition created moments later. - Drop partitions whose entire month is older than the retention window. -Intended to run as a monthly job (`keboola-mcp-server gc-sessions`), separate from the deploy-time -`migrate` command -- deploys don't happen on a reliable monthly cadence, so this can't piggyback -on that hook. +Intended to run as a recurring job (`keboola-mcp-server gc-sessions`), separate from the +deploy-time `migrate` command -- deploys don't happen on a reliable cadence, so this can't +piggyback on that hook. Safe to run more often than monthly: the exists-check on creation and the +month-boundary check on drops make repeated runs within the same month no-ops. """ import logging @@ -49,7 +50,9 @@ async def ensure_partitions( :return: ``{'created': [...], 'dropped': [...]}`` partition names, for the CLI to report. """ this_month = _month_start(datetime.now(timezone.utc).date()) - cutoff = _add_months(this_month, -retention_months) + # retention_months counts the current month, so keep (retention_months - 1) months before it -- + # e.g. retention_months=2 on an August run keeps July + August, drops June. + cutoff = _add_months(this_month, -(retention_months - 1)) created: list[str] = [] async with pool.acquire() as conn: diff --git a/tests/session_store/test_retention.py b/tests/session_store/test_retention.py index 07a777999..a49d9648a 100644 --- a/tests/session_store/test_retention.py +++ b/tests/session_store/test_retention.py @@ -75,6 +75,26 @@ async def test_drops_only_partitions_older_than_retention(self) -> None: finally: await pool.close() + async def test_drops_partition_exactly_retention_months_old(self) -> None: + # Regression test: retention_months=2 must keep exactly 2 months (this + previous), so a + # partition dated retention_months back (2 months old) is dropped, not kept. + pool = await asyncpg.create_pool(TEST_DSN) + try: + this_month = _month_start(date.today()) + boundary = _add_months(this_month, -2) + name = f'oauth_sessions_{boundary:%Y_%m}' + end = _add_months(boundary, 1) + await pool.execute( + f"CREATE TABLE {name} PARTITION OF oauth_sessions FOR VALUES FROM ('{boundary}') TO ('{end}')" + ) + + result = await ensure_partitions(pool, retention_months=2) + + assert name in result['dropped'] + assert name not in await self._existing_partitions(pool) + finally: + await pool.close() + async def test_creates_missing_current_and_next_month(self) -> None: pool = await asyncpg.create_pool(TEST_DSN) try: diff --git a/uv.lock b/uv.lock index 4af3a08a7..d544295c3 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.81.0" +version = "1.82.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From a3d1ba22d1336b225b2abcf4d59a81886c696a37 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 13:38:45 +0200 Subject: [PATCH 67/89] fix(PSGO-261): move overlapping default-partition rows before attaching a new oauth_sessions partition ensure_partitions() plain `CREATE TABLE ... PARTITION OF ... FOR VALUES` fails with asyncpg.exceptions.CheckViolationError whenever oauth_sessions_default already holds rows in the range being carved out -- exactly what happens on a stack with real pre-existing sessions, since migration 0002 copies the legacy table's rows into the default partition before any month partition exists. Build the partition as a standalone table, move matching default rows into it first, then ATTACH -- works whether or not default has conflicting rows. --- pyproject.toml | 2 +- .../session_store/retention.py | 25 ++++++++++-- tests/session_store/test_retention.py | 38 ++++++++++++++++++- uv.lock | 2 +- 4 files changed, 60 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2ac3c64bd..6351d79b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.82.0" +version = "1.83.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/session_store/retention.py b/src/keboola_mcp_server/session_store/retention.py index 64b501aad..71f10fa0e 100644 --- a/src/keboola_mcp_server/session_store/retention.py +++ b/src/keboola_mcp_server/session_store/retention.py @@ -64,10 +64,27 @@ async def ensure_partitions( if not exists: # DDL bounds can't be bound query parameters -- start/end are computed dates, not # user input, so direct formatting here carries no injection risk. - await conn.execute( - f'CREATE TABLE {name} PARTITION OF oauth_sessions ' - f"FOR VALUES FROM ('{start.isoformat()}') TO ('{end.isoformat()}')" - ) + # + # Can't just `CREATE TABLE {name} PARTITION OF oauth_sessions FOR VALUES FROM ... TO + # ...`: oauth_sessions_default may already hold rows in that range (e.g. the + # one-time backlog migration 0002 copies over, or a period where partition + # maintenance lagged), and Postgres refuses to attach a new range partition while + # the default partition has matching rows. So build the partition as a standalone + # table, move any matching default rows into it first, then attach it -- this works + # whether or not default has conflicting rows. + async with conn.transaction(): + await conn.execute(f'CREATE TABLE {name} (LIKE oauth_sessions INCLUDING ALL)') + await conn.execute( + f'WITH moved AS (' + f' DELETE FROM oauth_sessions_default ' + f" WHERE created_at >= '{start.isoformat()}' AND created_at < '{end.isoformat()}' " + f' RETURNING *' + f') INSERT INTO {name} SELECT * FROM moved' + ) + await conn.execute( + f'ALTER TABLE oauth_sessions ATTACH PARTITION {name} ' + f"FOR VALUES FROM ('{start.isoformat()}') TO ('{end.isoformat()}')" + ) LOG.info(f'Created oauth_sessions partition: {name} [{start}, {end})') created.append(name) diff --git a/tests/session_store/test_retention.py b/tests/session_store/test_retention.py index a49d9648a..e04bc0d3c 100644 --- a/tests/session_store/test_retention.py +++ b/tests/session_store/test_retention.py @@ -1,4 +1,4 @@ -from datetime import date +from datetime import date, datetime, timedelta, timezone import asyncpg import pytest @@ -95,6 +95,42 @@ async def test_drops_partition_exactly_retention_months_old(self) -> None: finally: await pool.close() + async def test_creates_partition_when_default_has_overlapping_rows(self) -> None: + # Regression test: a plain `CREATE TABLE ... PARTITION OF` fails with a CheckViolationError + # if oauth_sessions_default already holds rows in the new partition's range -- e.g. the + # one-time backlog copied over by migration 0002 on a stack with pre-existing sessions. + # ensure_partitions() must move those rows into the new partition instead of erroring. + pool = await asyncpg.create_pool(TEST_DSN) + try: + for name in await self._existing_partitions(pool): + await pool.execute(f'DROP TABLE {name}') + + this_month = _month_start(date.today()) + mid_month = datetime(this_month.year, this_month.month, this_month.day, tzinfo=timezone.utc) + timedelta( + days=1 + ) + await pool.execute( + 'INSERT INTO oauth_sessions_default ' + '(access_token_hash, client_id, kbc_access_token_enc, kbc_refresh_token_enc, ' + 'kbc_access_expires_at, created_at) ' + "VALUES ($1, 'client', $2, $3, now() + interval '1 hour', $4)", + b'token-hash', + b'enc-access', + b'enc-refresh', + mid_month, + ) + + result = await ensure_partitions(pool) + + this_month_partition = f'oauth_sessions_{this_month:%Y_%m}' + assert this_month_partition in result['created'] + row = await pool.fetchrow(f'SELECT client_id FROM {this_month_partition}') + assert row['client_id'] == 'client' + default_count = await pool.fetchval('SELECT count(*) FROM oauth_sessions_default') + assert default_count == 0 + finally: + await pool.close() + async def test_creates_missing_current_and_next_month(self) -> None: pool = await asyncpg.create_pool(TEST_DSN) try: diff --git a/uv.lock b/uv.lock index d544295c3..695953bdb 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.82.0" +version = "1.83.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 8a7b5c156b6a99389ab3a742ab29a10e33728c88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 14:08:14 +0200 Subject: [PATCH 68/89] fix(PSGO-261): log lazy Keboola-session token refresh in load_access_token The near-expiry refresh in load_access_token was silent on success (only the failure path logged a warning), so there was no way to observe when a session's Keboola credential got lazily extended. Log an info line with the session id only -- no token values. --- pyproject.toml | 2 +- src/keboola_mcp_server/oauth.py | 1 + tests/test_oauth.py | 12 ++++++++++-- uv.lock | 2 +- 4 files changed, 13 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6351d79b8..7111cde7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.83.0" +version = "1.84.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index e322ca6bb..3683d7e03 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -464,6 +464,7 @@ async def load_access_token(self, token: str) -> AccessToken | None: session = dataclasses.replace( session, kbc_access_token=token_set.access_token, kbc_refresh_token=token_set.refresh_token ) + LOG.info(f'[load_access_token] Lazily refreshed near-expiry Keboola session: session_id={session.id}') proxy_token = ProxyAccessToken( token=token, diff --git a/tests/test_oauth.py b/tests/test_oauth.py index 9d5ea924c..d9dedbea7 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -1,4 +1,5 @@ import dataclasses +import logging import secrets import time from collections.abc import Mapping @@ -497,7 +498,7 @@ async def _failing_refresh_tokens(storage_api_url: str, *, refresh_token: str, t @pytest.mark.asyncio async def test_load_access_token_refreshes_near_expiry_session_transparently( - self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ): from keboola_mcp_server import oauth as oauth_module from keboola_mcp_server.auth_login import TokenSet @@ -516,7 +517,8 @@ async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, tran monkeypatch.setattr(oauth_module, 'refresh_tokens', _fake_refresh_tokens) - loaded = await oauth_provider.load_access_token(access_token) + with caplog.at_level(logging.INFO): + loaded = await oauth_provider.load_access_token(access_token) assert loaded is not None assert loaded.kbc_access_token == 'kbc_at_fresh' @@ -525,6 +527,12 @@ async def _fake_refresh_tokens(storage_api_url: str, *, refresh_token: str, tran assert stored is not None assert stored.kbc_access_token == 'kbc_at_fresh' assert stored.kbc_refresh_token == 'kbc_rt_fresh' + # Observable in logs (session id only, no token values) -- previously silent on success. + refresh_logs = [r for r in caplog.records if 'Lazily refreshed near-expiry' in r.message] + assert len(refresh_logs) == 1 + assert session.id in refresh_logs[0].message + assert 'kbc_at_fresh' not in refresh_logs[0].message + assert 'kbc_rt_fresh' not in refresh_logs[0].message @pytest.mark.asyncio async def test_load_access_token_tolerates_refresh_failure( diff --git a/uv.lock b/uv.lock index 695953bdb..55c26aa9c 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.83.0" +version = "1.84.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 0e64b371570fd14a0f551032a02d657b80d5709f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 14:18:44 +0200 Subject: [PATCH 69/89] fix(PSGO-261): require KBC_SESSION_ENCRYPTION_KEY when OAuth is enabled resolve_encryption_key() silently falls back to a process-local key when KBC_SESSION_ENCRYPTION_KEY is unset -- fine for local dev/tests, but in production this would make persisted OAuth sessions permanently undecryptable after every restart. Refuse to start instead, mirroring the existing Postgres DSN guard right above it. Addresses a Copilot review comment on PR #605. --- pyproject.toml | 2 +- src/keboola_mcp_server/server.py | 8 ++++++++ tests/test_server.py | 21 +++++++++++++++++++-- uv.lock | 2 +- 4 files changed, 29 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7111cde7a..77e70d51e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.84.0" +version = "2.0.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index 00ee9ae4d..fed8c1c91 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -215,6 +215,14 @@ def create_server( 'OAuth is configured (oauth_client_id/oauth_client_secret) but no Postgres DSN is set. ' 'Set MCP_DB_URL (or KBC_POSTGRES_DSN) so OAuth sessions can be stored.' ) + # Without an explicit key, resolve_encryption_key() falls back to a process-local one -- + # fine for local dev/tests, but in production it would silently make persisted sessions + # undecryptable after every restart (same "refuse to start" reasoning as the DSN check above). + if not config.session_encryption_key: + raise RuntimeError( + 'OAuth is configured (oauth_client_id/oauth_client_secret) but no session encryption key is ' + 'set. Set KBC_SESSION_ENCRYPTION_KEY so persisted OAuth sessions survive a process restart.' + ) session_store = PostgresSessionStore( config.postgres_dsn, encryption_key=resolve_encryption_key(config.session_encryption_key) ) diff --git a/tests/test_server.py b/tests/test_server.py index 4d42aa480..0b0027190 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -1,4 +1,5 @@ import asyncio +import base64 import json import subprocess import tempfile @@ -623,6 +624,8 @@ class TestCreateServerOAuthSessionStore: """OAuth sessions live in Postgres (oauth_session_persistence RFC) -- create_server() must refuse to enable OAuth without a DSN rather than silently falling back to something unrevoked.""" + _TEST_ENCRYPTION_KEY = base64.b64encode(b'0' * 32).decode() + @staticmethod def _oauth_config(**overrides) -> Config: return Config( @@ -636,13 +639,27 @@ def _oauth_config(**overrides) -> Config: def test_raises_without_postgres_dsn(self) -> None: with pytest.raises(RuntimeError, match='MCP_DB_URL'): - create_server(self._oauth_config(), runtime_info=ServerRuntimeInfo(transport='streamable-http')) + create_server( + self._oauth_config(session_encryption_key=self._TEST_ENCRYPTION_KEY), + runtime_info=ServerRuntimeInfo(transport='streamable-http'), + ) + + def test_raises_without_session_encryption_key(self) -> None: + # A silent fallback to a process-local key would make persisted OAuth sessions + # undecryptable after every restart -- refuse to start instead, same as the DSN check. + with pytest.raises(RuntimeError, match='KBC_SESSION_ENCRYPTION_KEY'): + create_server( + self._oauth_config(postgres_dsn='postgresql://u:p@host/db'), + runtime_info=ServerRuntimeInfo(transport='streamable-http'), + ) def test_constructs_session_store_when_dsn_is_set(self) -> None: from keboola_mcp_server.session_store.repository import PostgresSessionStore server = create_server( - self._oauth_config(postgres_dsn='postgresql://u:p@host/db'), + self._oauth_config( + postgres_dsn='postgresql://u:p@host/db', session_encryption_key=self._TEST_ENCRYPTION_KEY + ), runtime_info=ServerRuntimeInfo(transport='streamable-http'), ) assert isinstance(server, FastMCP) diff --git a/uv.lock b/uv.lock index 55c26aa9c..5c65e9737 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.84.0" +version = "2.0.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 5958223ebc20fd521dd3eb939be00622eece2d32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 31 Jul 2026 15:32:46 +0200 Subject: [PATCH 70/89] fix(PSGO-261): enforce real per-partition uniqueness on oauth_sessions token hashes The parent's (access_token_hash, created_at) / (refresh_token_hash, created_at) unique indexes never actually reject a duplicate hash -- created_at differs per row, so the composite pair is always distinct. PostgresSessionStore.get_by_access_token() does an unbounded UPDATE ... RETURNING * on access_token_hash alone, so a duplicate would update multiple rows while fetchrow() silently returns an arbitrary one. Add a plain (non-composite) unique index directly on each partition table -- Postgres only requires the partition key in indexes defined on the partitioned parent, not on a partition's own table. ensure_partitions() adds it to every new month partition; a new migration (0003) adds it to oauth_sessions_default retroactively, since 0002 is already applied on real stacks and can't be edited. Addresses a Copilot review comment on PR #605. --- pyproject.toml | 2 +- .../0003_default_partition_unique_indexes.sql | 6 +++++ .../session_store/retention.py | 23 ++++++++++------- tests/session_store/test_migrator.py | 25 ++++++++++++++++++- tests/session_store/test_retention.py | 23 +++++++++++++++++ uv.lock | 2 +- 6 files changed, 69 insertions(+), 12 deletions(-) create mode 100644 src/keboola_mcp_server/session_store/migrations/0003_default_partition_unique_indexes.sql diff --git a/pyproject.toml b/pyproject.toml index 77e70d51e..c5cdbf708 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "2.0.0" +version = "1.75.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/session_store/migrations/0003_default_partition_unique_indexes.sql b/src/keboola_mcp_server/session_store/migrations/0003_default_partition_unique_indexes.sql new file mode 100644 index 000000000..d3fb5dc47 --- /dev/null +++ b/src/keboola_mcp_server/session_store/migrations/0003_default_partition_unique_indexes.sql @@ -0,0 +1,6 @@ +-- 0002's composite (access_token_hash, created_at) index doesn't enforce hash uniqueness -- +-- created_at differs per row. A plain index on the partition table itself does. +-- retention.ensure_partitions() adds the same pair on every new month partition. +CREATE UNIQUE INDEX oauth_sessions_default_access_token_hash_uidx ON oauth_sessions_default (access_token_hash); +CREATE UNIQUE INDEX oauth_sessions_default_refresh_token_hash_uidx ON oauth_sessions_default (refresh_token_hash) + WHERE refresh_token_hash IS NOT NULL; diff --git a/src/keboola_mcp_server/session_store/retention.py b/src/keboola_mcp_server/session_store/retention.py index 71f10fa0e..af5b766de 100644 --- a/src/keboola_mcp_server/session_store/retention.py +++ b/src/keboola_mcp_server/session_store/retention.py @@ -62,16 +62,11 @@ async def ensure_partitions( name = _partition_name(start) exists = await conn.fetchval('SELECT to_regclass($1) IS NOT NULL', name) if not exists: - # DDL bounds can't be bound query parameters -- start/end are computed dates, not - # user input, so direct formatting here carries no injection risk. + # DDL bounds can't be bound query parameters -- start/end are computed, not user + # input, so direct formatting is safe. # - # Can't just `CREATE TABLE {name} PARTITION OF oauth_sessions FOR VALUES FROM ... TO - # ...`: oauth_sessions_default may already hold rows in that range (e.g. the - # one-time backlog migration 0002 copies over, or a period where partition - # maintenance lagged), and Postgres refuses to attach a new range partition while - # the default partition has matching rows. So build the partition as a standalone - # table, move any matching default rows into it first, then attach it -- this works - # whether or not default has conflicting rows. + # Postgres refuses to attach a new partition while default holds matching rows + # (e.g. migration 0002's backlog copy), so move any such rows in first. async with conn.transaction(): await conn.execute(f'CREATE TABLE {name} (LIKE oauth_sessions INCLUDING ALL)') await conn.execute( @@ -85,6 +80,16 @@ async def ensure_partitions( f'ALTER TABLE oauth_sessions ATTACH PARTITION {name} ' f"FOR VALUES FROM ('{start.isoformat()}') TO ('{end.isoformat()}')" ) + # The copied (access_token_hash, created_at) index doesn't actually enforce hash + # uniqueness (created_at differs per row) -- a plain index on the partition + # table itself, without the partition key, does. + await conn.execute( + f'CREATE UNIQUE INDEX {name}_access_token_hash_uidx ON {name} (access_token_hash)' + ) + await conn.execute( + f'CREATE UNIQUE INDEX {name}_refresh_token_hash_uidx ON {name} (refresh_token_hash) ' + f'WHERE refresh_token_hash IS NOT NULL' + ) LOG.info(f'Created oauth_sessions partition: {name} [{start}, {end})') created.append(name) diff --git a/tests/session_store/test_migrator.py b/tests/session_store/test_migrator.py index fb2d8f9a3..3e6c6a65c 100644 --- a/tests/session_store/test_migrator.py +++ b/tests/session_store/test_migrator.py @@ -21,7 +21,11 @@ async def test_applies_migrations_once() -> None: pool = await asyncpg.create_pool(TEST_DSN) try: applied = await apply_migrations(pool) - assert applied == ['0001_oauth_sessions.sql', '0002_partition_oauth_sessions.sql'] + assert applied == [ + '0001_oauth_sessions.sql', + '0002_partition_oauth_sessions.sql', + '0003_default_partition_unique_indexes.sql', + ] # Re-running is a no-op -- the table already exists, so re-applying the DDL would fail # if the tracking table didn't correctly skip it. @@ -62,3 +66,22 @@ async def test_partitions_table_with_default_catch_all() -> None: assert tables == {'oauth_sessions', 'oauth_sessions_default'} finally: await pool.close() + + +async def test_default_partition_rejects_duplicate_access_token_hash() -> None: + # Regression test: the parent's (access_token_hash, created_at) index alone doesn't reject a + # duplicate hash (created_at differs per row) -- migration 0003's plain index on the + # partition table itself must. + pool = await asyncpg.create_pool(TEST_DSN) + try: + await apply_migrations(pool) + insert = ( + 'INSERT INTO oauth_sessions_default ' + '(access_token_hash, client_id, kbc_access_token_enc, kbc_refresh_token_enc, kbc_access_expires_at) ' + "VALUES ($1, 'client', $2, $3, now())" + ) + await pool.execute(insert, b'dup-hash', b'enc-access', b'enc-refresh') + with pytest.raises(asyncpg.UniqueViolationError): + await pool.execute(insert, b'dup-hash', b'enc-access', b'enc-refresh') + finally: + await pool.close() diff --git a/tests/session_store/test_retention.py b/tests/session_store/test_retention.py index e04bc0d3c..f17e83ace 100644 --- a/tests/session_store/test_retention.py +++ b/tests/session_store/test_retention.py @@ -131,6 +131,29 @@ async def test_creates_partition_when_default_has_overlapping_rows(self) -> None finally: await pool.close() + async def test_created_partition_rejects_duplicate_access_token_hash(self) -> None: + # Regression test: the parent's composite index doesn't reject a duplicate hash (created_at + # differs per row) -- the plain index ensure_partitions() adds on the partition itself must. + pool = await asyncpg.create_pool(TEST_DSN) + try: + for name in await self._existing_partitions(pool): + await pool.execute(f'DROP TABLE {name}') + + result = await ensure_partitions(pool) + this_month_partition = f'oauth_sessions_{_month_start(date.today()):%Y_%m}' + assert this_month_partition in result['created'] + + insert = ( + f'INSERT INTO {this_month_partition} ' + '(access_token_hash, client_id, kbc_access_token_enc, kbc_refresh_token_enc, kbc_access_expires_at) ' + "VALUES ($1, 'client', $2, $3, now())" + ) + await pool.execute(insert, b'dup-hash', b'enc-access', b'enc-refresh') + with pytest.raises(asyncpg.UniqueViolationError): + await pool.execute(insert, b'dup-hash', b'enc-access', b'enc-refresh') + finally: + await pool.close() + async def test_creates_missing_current_and_next_month(self) -> None: pool = await asyncpg.create_pool(TEST_DSN) try: diff --git a/uv.lock b/uv.lock index 5c65e9737..8423256a4 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "2.0.0" +version = "1.75.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 18b961306d32f95b337912b9b5454845b2f293c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 4 Aug 2026 15:00:20 +0200 Subject: [PATCH 71/89] fix(PSGO-261): require explicit project_id on write tools instead of active-project re-scope Writes previously targeted the first scoped project implicitly, forcing a set_project_scope call to change which project a write lands on -- and reordering the whole scope's read fan-out as a side effect. Every write/modify/delete tool now declares project_id explicitly; ambiguous calls (2+ scoped projects, no project_id) raise instead of silently defaulting. MultiProjectMiddleware resolves and swaps the target ahead of ToolsFilteringMiddleware so authorization runs against the targeted project, mirroring the existing read fan-out ordering. Co-Authored-By: Claude Sonnet 5 --- TOOLS.md | 233 +++++++++++++++++- feature_spec/pat_token_support/RFC.md | 36 ++- src/keboola_mcp_server/multiproject.py | 66 ++++- src/keboola_mcp_server/scope.py | 17 ++ .../tools/components/tools.py | 7 + src/keboola_mcp_server/tools/data_apps.py | 6 + src/keboola_mcp_server/tools/flow/tools.py | 5 + src/keboola_mcp_server/tools/jobs.py | 2 + src/keboola_mcp_server/tools/oauth.py | 2 + src/keboola_mcp_server/tools/project.py | 12 +- src/keboola_mcp_server/tools/storage/tools.py | 2 + tests/test_multiproject.py | 71 +++++- 12 files changed, 444 insertions(+), 15 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index d3948de18..39cceef16 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -189,6 +189,18 @@ EXAMPLES: ], "default": null, "description": "The list of processors that will run after the configured component row runs." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -353,6 +365,18 @@ EXAMPLES: ], "default": null, "description": "Variable definitions to attach to this configuration. Each entry specifies a name, type (\"string\" or \"vault\"), and an optional default value. On creation, both `None` (omitted) and `[]` (empty list) mean \"do not attach variables\" \u2014 no `keboola.variables` config is created. To remove variables from an existing configuration, use `update_config` with `variables=[]`." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -510,6 +534,18 @@ EXAMPLES: ], "default": null, "description": "Variable definitions to attach to this transformation. Each entry specifies a name, type (\"string\" or \"vault\"), and an optional default value. On creation, both `None` (omitted) and `[]` (empty list) mean \"do not attach variables\" \u2014 no `keboola.variables` config is created. To remove variables from an existing transformation, use `update_sql_transformation` with `variables=[]`." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -1077,6 +1113,18 @@ WORKFLOW: ], "default": null, "description": "Variable definitions for this configuration. Provide a non-empty list to create or replace all variable definitions. Provide an empty list ([]) to remove all variables. Omit (None) to leave existing variables unchanged." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -1352,6 +1400,18 @@ WORKFLOW: ], "default": null, "description": "Enable or disable the configuration row. Set to True to disable execution (config row won't run), False to enable execution (config row will run). Only provide if changing the status, leave as null to preserve current state." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2018,6 +2078,18 @@ Example 4 - Update storage mappings: ], "default": null, "description": "Variable definitions for this transformation. Provide a non-empty list to create or replace all variable definitions. Provide an empty list ([]) to remove all variables. Omit (None) to leave existing variables unchanged." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2087,6 +2159,18 @@ additional token without invalidating any tokens already held by other clients. "configuration_id": { "description": "Storage configuration ID of the python-js data app.", "type": "string" + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2137,6 +2221,18 @@ in the response) or to `get_data_apps` for further work. "configuration_id": { "description": "Storage configuration ID of the python-js draft data app to delete.", "type": "string" + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2216,6 +2312,18 @@ Streamlit apps have no managed git repo, so `mode` has no effect on the deployed ], "default": null, "description": "Deployment mode. Set to \"dev\" to deploy a python-js draft as a **dev version of the data app** \u2014 the runtime uses a development `setup.sh` (hot reload), and the data-app proxy enables an auto-auth path so an iframe preview can render without a manual login. Only meaningful on **draft** configs (python-js apps with `isDraft=true`). Leave None (default) for prod redeploys and for Streamlit apps." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2502,6 +2610,18 @@ slug must be at most 63 characters (the DNS-label max), and note the UI's own UR ], "default": null, "description": "Folder name to organize this data app in the Keboola UI. Pass an empty string to remove an existing folder assignment. Existing folder names are returned in the response change_summary when no folder is provided and there are 20 or more data apps in the project. If there are 20 or more data apps, you should assign one of the existing folders or create a new one that clearly reflects the data app purpose." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2603,6 +2723,18 @@ SQL & DATA TYPE RULES: ], "default": null, "description": "Folder name to organize this data app in the Keboola UI. Pass an empty string to remove an existing folder assignment. Existing folder names are returned in the response change_summary when no folder is provided and there are 20 or more data apps in the project. If there are 20 or more data apps, you should assign one of the existing folders or create a new one that clearly reflects the data app purpose." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2709,6 +2841,18 @@ WHEN TO USE: "default": "", "description": "Folder name to organize this flow in the Keboola UI. Pass an empty string to remove an existing folder assignment. Existing folder names are returned in the response change_summary when no folder is provided and there are 20 or more flows in the project. If there are 20 or more flows, you should assign one of the existing folders or create a new one that clearly reflects the flow purpose.", "type": "string" + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -2780,6 +2924,18 @@ WHEN TO USE: "default": "", "description": "Folder name to organize this flow in the Keboola UI. Pass an empty string to remove an existing folder assignment. Existing folder names are returned in the response change_summary when no folder is provided and there are 20 or more flows in the project. If there are 20 or more flows, you should assign one of the existing folders or create a new one that clearly reflects the flow purpose.", "type": "string" + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -3116,6 +3272,18 @@ adjusting dependencies, or enabling/disabling flow execution ], "default": null, "description": "Folder name to organize this flow in the Keboola UI. Pass an empty string to remove an existing folder assignment. Existing folder names are returned in the response change_summary when no folder is provided and there are 20 or more flows in the project. If there are 20 or more flows, you should assign one of the existing folders or create a new one that clearly reflects the flow purpose." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -3248,6 +3416,18 @@ or enabling/disabling flow execution ], "default": null, "description": "Folder name to organize this flow in the Keboola UI. Pass an empty string to remove an existing folder assignment. Existing folder names are returned in the response change_summary when no folder is provided and there are 20 or more flows in the project. If there are 20 or more flows, you should assign one of the existing folders or create a new one that clearly reflects the flow purpose." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -3499,6 +3679,18 @@ Starts a new job for a given component or transformation. ], "default": null, "description": "Optional list of configuration row IDs to run. If not provided, all rows are executed." + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -3541,6 +3733,18 @@ configuration is created e.g. keboola.ex-google-analytics-v4 and keboola.ex-gmai "config_id": { "description": "The configuration ID for the component.", "type": "string" + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -3626,8 +3830,9 @@ Scopes the current session to a set of Keboola projects. Mints a scoped access token (narrowed to `project_ids`, optionally read-only) that is used for the rest of the conversation. Read-only tools then run against every scoped project in a single call; -write operations target the active (first) project only. Call this when the user states which -projects to work on; it can be called again any time to re-scope. +write/modify/delete tools take a `project_id` argument naming which scoped project to target (required +once 2+ projects are scoped). Call this when the user states which projects to work on; it can be +called again any time to re-scope. The server does not remember this scope between calls: pass the returned `scope_token` as the `scope_token` argument on every subsequent tool call in this conversation to keep it in effect. @@ -3683,6 +3888,18 @@ Updates the description of the current Keboola project. "description": { "description": "The new project description text.", "type": "string" + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ @@ -4703,6 +4920,18 @@ Usage examples (payload uses a list of DescriptionUpdate objects): "$ref": "#/$defs/DescriptionUpdate" }, "type": "array" + }, + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." } }, "required": [ diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 7bb632bbb..ef9e1a29d 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -331,6 +331,8 @@ the session's subject token for all downstream exchange/forwarding. automatically: with >1 project in scope they require a single confirmed target project. Server instructions state: **the agent must never write to more than one project without explicit user guidance or confirmation.** Bulk multi-project writes are possible but only on that explicit signal. + _(How "confirmed target" is expressed evolved — see "Decisions (increment 5)" below: an explicit + `project_id` tool argument, not the interim active-project/re-scope indirection.)_ ## Scope changes (relative to the base RFC) @@ -415,7 +417,7 @@ the per-project envelope shape change the docs), integration tests on a dev stac - **Q1 (exactly-1-in-scope) — confirmed: a single-project scope returns the raw result, not a 1-element envelope.** Single-project UX is byte-for-byte unchanged. - **Q4 (write-target confirmation) — lean: explicit `project_id` arg on write tools, honored only when - scope > 1.** (Still open; not blocking.) + scope > 1.** (Resolved as leaned — see "Decisions (increment 5)" below.) # Extension: query fan-out, dialect-aware bootstrap, per-service token gaps (PSGO-261, increment 3) @@ -605,6 +607,8 @@ achievable on clients that re-fetch on `list_changed`; we don't rely on it. (absent in Claude Code mid-session), so all tools stay listed and data tools are gated at call time. - **No phantom active project before a scope is confirmed**; after `set_project_scope` the `active_project_id` is the write / `query_data`-default target and is surfaced intentionally. + _(Superseded for writes by "Decisions (increment 5)" below: writes now take an explicit + `project_id` argument instead of implicitly targeting `active_project_id`.)_ - **Fan-out stays**, with the relevance/latency critiques answered by the `project_ids` filter and a (follow-up) concurrent loop; per-project error isolation is now implemented (partial results). - Fixed a latent bug: `set_project_scope` referenced `minted.read_only` on the exchange-failure path @@ -669,3 +673,33 @@ all: `create_session_state` forwards any programmatic token (`kbc_at_*`/`kbc_pat `resolve-storage-token` auth-bridge exchange this section referenced has been removed (see `oauth_session_exchange/RFC.md` Decision §6). Full multi-project scope now works the same way on the deployed server as it does locally. + +## Decisions (increment 5) — explicit `project_id` on write tools (resolves Q4) + +**Q4 (write-target confirmation), previously "still open; not blocking," is now resolved as leaned: +explicit `project_id` argument on every write/modify/delete tool, required once 2+ projects are +scoped.** Superseded is the interim behavior described above (line ~604, "increment 4"): a write +targeting `active_project_id` (the first scoped project) with no per-call target, requiring +`set_project_scope` to change which project a write lands on. That indirection was reported as +confusing in practice — writing to a different scoped project needlessly demanded a re-scope, which +also reorders the scope for every subsequent read fan-out. + +- **Every write tool now declares `project_id: str | None = None`** (a real, schema-visible + parameter — not a middleware-injected one, unlike the read-side `project_ids` filter). The LLM + states its target explicitly in the conversation. +- **`MultiProjectMiddleware._dispatch_write`** (not the tool body) resolves and swaps the target, + for the same reason `_swap_project` already runs ahead of `ToolsFilteringMiddleware` for read + fan-out: role/feature/branch authorization must be evaluated against the *targeted* project's + client, not whatever was active before the call. +- **Ambiguity is now a hard error, not a silent default:** 2+ scoped projects and no `project_id` → + `ToolError` naming the scoped projects and asking for one. Exactly one scoped project still + defaults `project_id` to it (unchanged single-project UX). +- **Read tools are unaffected** — they keep the existing `project_ids`-filtered fan-out; listing + needs no single target. + +This also folds in the one still-useful idea from the earlier, superseded MPA RFC (PR #500, +AI-3027, closed as superseded by this RFC): its "`project_id` as an explicit tool argument, chosen +over a header/middleware-only approach" recommendation, including the ambiguity rule (`from_project` +raising when 2+ projects are active and no `project_id` is given). Everything else in PR #500 (token +taxonomy, append-only project registry, Kai integration flow, 24h idle refresh) is already covered +by this RFC and the as-built code under different names. diff --git a/src/keboola_mcp_server/multiproject.py b/src/keboola_mcp_server/multiproject.py index 164278f92..2897bb67e 100644 --- a/src/keboola_mcp_server/multiproject.py +++ b/src/keboola_mcp_server/multiproject.py @@ -20,7 +20,7 @@ from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import build_tracing_headers, deployed_sa_token_path from keboola_mcp_server.mcp import ServerState, is_read_only_tool -from keboola_mcp_server.scope import SCOPE_KEY, SessionScope +from keboola_mcp_server.scope import PROJECT_ID_ARG, SCOPE_KEY, SessionScope from keboola_mcp_server.tools.constants import BOOTSTRAP_TOOLS from keboola_mcp_server.workspace import WorkspaceManager @@ -48,8 +48,8 @@ class MultiProjectMiddleware(fmw.Middleware): project's client and the per-project results are labelled with a per-project text envelope. Their structured content is deep-merged (lists concatenated across projects, counters summed) into one schema-valid object, degrading to count-first with a truncated sample past ``_FANOUT_MAX_ITEMS``. - Write tools never fan out: they target the active project only, so the agent can never write to - multiple projects without the user explicitly re-scoping (PSGO-261 decision D8). + Write tools never fan out: a write always targets exactly one project, named by its own + ``project_id`` argument (required once 2+ projects are scoped) -- see ``_dispatch_write``. """ async def on_call_tool( @@ -78,12 +78,12 @@ async def on_call_tool( # Bootstrap tools own a real `project_ids` argument, so we must not strip it. if not isinstance(scope, SessionScope) or name in BOOTSTRAP_TOOLS: return await call_next(context) - # Workspace-bound and write tools always target the active project (no fan-out, filter ignored). + # Workspace-bound tools always target the active project (no fan-out, filter ignored). if name in _NO_FANOUT_TOOLS: return await call_next(context) tool = await ctx.fastmcp.get_tool(name) if not is_read_only_tool(tool): - return await call_next(context) + return await self._dispatch_write(context, call_next, ctx, state, scope) # Read tool: consume the optional per-call project filter (advertised via on_list_tools) so the # tool never receives it, then narrow this call's target projects to the requested subset. @@ -164,6 +164,62 @@ async def on_call_tool( return self._merge(results, errors) + @staticmethod + def _resolve_write_target(scope: SessionScope, project_id: Any) -> int | None: + """Picks the single project a write call targets, or raises if that's ambiguous/invalid. + + ``project_id`` is required once 2+ projects are scoped (no more implicit "first project" + default); with exactly one scoped project it's optional and defaults to that project. + """ + if project_id is None: + if len(scope.project_ids) >= 2: + raise ToolError( + f'{len(scope.project_ids)} projects are scoped ({scope.project_ids}). ' + 'Pass project_id= (one of the scoped projects) -- a write targets exactly one project.' + ) + return scope.active_project_id + try: + target = int(project_id) + except (TypeError, ValueError): + raise ToolError(f'project_id must be an integer project id, got: {project_id!r}') + if target not in scope.project_ids: + raise ToolError( + f'Project {target} is outside the current scope {scope.project_ids}. ' + 'Call "set_project_scope" to change the scope first.' + ) + return target + + async def _dispatch_write( + self, + context: MiddlewareContext[mt.CallToolRequestParams], + call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult], + ctx: Any, + state: dict[str, Any], + scope: SessionScope, + ) -> mt.CallToolResult: + """Targets a write/modify/delete tool call at the project named by its ``project_id`` + argument (peeked, not popped -- it's a real declared tool parameter, not middleware-only). + """ + args = getattr(context.message, 'arguments', None) + project_id = args.get(PROJECT_ID_ARG) if isinstance(args, dict) else None + target = self._resolve_write_target(scope, project_id) + + if target is None or target == scope.active_project_id: + return await call_next(context) + + server_state = ServerState.from_context(ctx) + original_client = state.get(KeboolaClient.STATE_KEY) + original_workspace = state.get(WorkspaceManager.STATE_KEY) + is_real_client = isinstance(original_client, KeboolaClient) + base_token = scope.scoped_token or (original_client.token if is_real_client else '') + storage_api_url = original_client.storage_api_url if is_real_client else server_state.config.storage_api_url + try: + await self._swap_project(state, server_state, storage_api_url, base_token, target, scope.read_only) + return await call_next(context) + finally: + state[KeboolaClient.STATE_KEY] = original_client + state[WorkspaceManager.STATE_KEY] = original_workspace + async def on_list_tools( self, context: MiddlewareContext[mt.ListToolsRequest], diff --git a/src/keboola_mcp_server/scope.py b/src/keboola_mcp_server/scope.py index 6096a059e..0a10d0b23 100644 --- a/src/keboola_mcp_server/scope.py +++ b/src/keboola_mcp_server/scope.py @@ -8,12 +8,29 @@ import dataclasses import secrets import time +from typing import Annotated, Optional + +from pydantic import Field from keboola_mcp_server.config import Config from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt SCOPE_KEY = 'project_scope' +# Declared on every write/modify/delete tool; consumed by MultiProjectMiddleware.on_call_tool to +# pick which scoped project the call targets (see multiproject.py's write branch). Optional only +# when the scope resolves the target unambiguously (a single scoped project). +PROJECT_ID_ARG = 'project_id' +ProjectIdArg = Annotated[ + Optional[str], + Field( + description=( + 'Target Keboola project id for this write. Required when the session is scoped to 2+ ' + 'projects; optional (defaults to the single scoped project) otherwise.' + ) + ), +] + # The OAuth session's DB row id (see session_store.repository.OAuthSession), stashed on # ctx.session.state so set_project_scope can persist a newly-confirmed scope back to Postgres # instead of only returning a scope_token. Absent for non-OAuth (PAT/header-token) sessions, which diff --git a/src/keboola_mcp_server/tools/components/tools.py b/src/keboola_mcp_server/tools/components/tools.py index cc5d2a998..84ac1e3cb 100644 --- a/src/keboola_mcp_server/tools/components/tools.py +++ b/src/keboola_mcp_server/tools/components/tools.py @@ -49,6 +49,7 @@ toon_serializer_compact, unwrap_results, ) +from keboola_mcp_server.scope import ProjectIdArg from keboola_mcp_server.tools.components.model import ( Component, ComponentSummary, @@ -436,6 +437,7 @@ async def create_sql_transformation( ), ), ] = None, + project_id: ProjectIdArg = None, ) -> ConfigToolOutput: """ Creates an SQL transformation using the specified name, SQL query following the current SQL dialect, a detailed @@ -663,6 +665,7 @@ async def update_sql_transformation( ), ), ] = None, + project_id: ProjectIdArg = None, ) -> ConfigToolOutput: """ Updates an existing SQL transformation configuration by modifying its SQL code, storage mappings, @@ -1117,6 +1120,7 @@ async def create_config( ), ), ] = None, + project_id: ProjectIdArg = None, ) -> ConfigToolOutput: """ Creates a root component configuration using the specified name, component ID, configuration JSON, and description. @@ -1264,6 +1268,7 @@ async def add_config_row( list[dict[str, Any]] | None, Field(description='The list of processors that will run after the configured component row runs.'), ] = None, + project_id: ProjectIdArg = None, ) -> ConfigToolOutput: """ Creates a component configuration row in the specified configuration_id, using the specified name, @@ -1463,6 +1468,7 @@ async def update_config( ), ), ] = None, + project_id: ProjectIdArg = None, ) -> ConfigToolOutput: """ Updates an existing root component configuration by modifying its parameters, storage mappings, name or description. @@ -1730,6 +1736,7 @@ async def update_config_row( ), ), ] = None, + project_id: ProjectIdArg = None, ) -> ConfigToolOutput: """ Updates an existing component configuration row by modifying its parameters, storage mappings, name, or description. diff --git a/src/keboola_mcp_server/tools/data_apps.py b/src/keboola_mcp_server/tools/data_apps.py index 114b4c9dc..10d94da06 100644 --- a/src/keboola_mcp_server/tools/data_apps.py +++ b/src/keboola_mcp_server/tools/data_apps.py @@ -26,6 +26,7 @@ from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager from keboola_mcp_server.mcp import process_concurrently, toon_serializer_compact +from keboola_mcp_server.scope import ProjectIdArg from keboola_mcp_server.tools.components.utils import ( apply_folder_metadata, folder_field_description, @@ -519,6 +520,7 @@ async def modify_streamlit_data_app( str | None, Field(description=folder_field_description('data app', 'data apps')), ] = None, + project_id: ProjectIdArg = None, ) -> ModifiedDataAppOutput: """Creates or updates a Streamlit data app. @@ -957,6 +959,7 @@ async def modify_python_js_data_app( str | None, Field(description=folder_field_description('data app', 'data apps')), ] = None, + project_id: ProjectIdArg = None, ) -> ModifiedPythonJsDataAppOutput: """Creates or updates a python-js data app. @@ -1298,6 +1301,7 @@ async def modify_python_js_data_app( async def create_python_js_data_app_git_credential( ctx: Context, configuration_id: Annotated[str, Field(description='Storage configuration ID of the python-js data app.')], + project_id: ProjectIdArg = None, ) -> CreatedGitCredentialOutput: """Mints a one-time HTTPS token on a python-js **prod** data app so the caller can clone, pull, and push to the app's managed git repo over HTTPS. @@ -1620,6 +1624,7 @@ async def deploy_data_app( ), ), ] = None, + project_id: ProjectIdArg = None, ) -> DeploymentDataAppOutput: """Deploys/redeploys a data app or stops a running data app in the Keboola environment asynchronously, given the action and the configuration ID. @@ -1703,6 +1708,7 @@ async def delete_python_js_data_app_draft( configuration_id: Annotated[ str, Field(description='Storage configuration ID of the python-js draft data app to delete.') ], + project_id: ProjectIdArg = None, ) -> DeletedDraftOutput: """Deletes a python-js DRAFT data app — both the data-app instance (DSAPI) and its Storage configuration. diff --git a/src/keboola_mcp_server/tools/flow/tools.py b/src/keboola_mcp_server/tools/flow/tools.py index 3f696107a..d54844b93 100644 --- a/src/keboola_mcp_server/tools/flow/tools.py +++ b/src/keboola_mcp_server/tools/flow/tools.py @@ -26,6 +26,7 @@ from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import ProjectLinksManager from keboola_mcp_server.mcp import process_concurrently, toon_serializer_compact, unwrap_results +from keboola_mcp_server.scope import ProjectIdArg from keboola_mcp_server.tools.components.utils import ( build_folder_hint, clear_configuration_folder_metadata, @@ -164,6 +165,7 @@ async def create_flow( str, Field(description=folder_field_description('flow', 'flows')), ] = '', + project_id: ProjectIdArg = None, ) -> FlowToolOutput: """ Creates a new legacy (non-conditional) flow using `keboola.orchestrator`. @@ -259,6 +261,7 @@ async def create_conditional_flow( str, Field(description=folder_field_description('flow', 'flows')), ] = '', + project_id: ProjectIdArg = None, ) -> FlowToolOutput: """ Creates a new conditional flow configuration using `keboola.flow`. @@ -377,6 +380,7 @@ async def update_flow( str | None, Field(description=folder_field_description('flow', 'flows')), ] = None, + project_id: ProjectIdArg = None, ) -> FlowToolOutput: """ Updates an existing flow configuration (either legacy `keboola.orchestrator` or conditional `keboola.flow`). @@ -463,6 +467,7 @@ async def modify_flow( str | None, Field(description=folder_field_description('flow', 'flows')), ] = None, + project_id: ProjectIdArg = None, ) -> FlowToolOutput: """ Updates an existing flow configuration (either legacy `keboola.orchestrator` or conditional `keboola.flow`) or diff --git a/src/keboola_mcp_server/tools/jobs.py b/src/keboola_mcp_server/tools/jobs.py index 3efee4ed6..c2f99cecb 100644 --- a/src/keboola_mcp_server/tools/jobs.py +++ b/src/keboola_mcp_server/tools/jobs.py @@ -12,6 +12,7 @@ from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager from keboola_mcp_server.mcp import KeboolaMcpServer, process_concurrently, toon_serializer_compact, unwrap_results +from keboola_mcp_server.scope import ProjectIdArg LOG = logging.getLogger(__name__) @@ -435,6 +436,7 @@ async def run_job( description='Optional list of configuration row IDs to run. If not provided, all rows are executed.', ), ] = None, + project_id: ProjectIdArg = None, ) -> JobDetail: """ Starts a new job for a given component or transformation. diff --git a/src/keboola_mcp_server/tools/oauth.py b/src/keboola_mcp_server/tools/oauth.py index 34bbdbebd..d9378b269 100644 --- a/src/keboola_mcp_server/tools/oauth.py +++ b/src/keboola_mcp_server/tools/oauth.py @@ -12,6 +12,7 @@ from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.mcp import KeboolaMcpServer +from keboola_mcp_server.scope import ProjectIdArg LOG = logging.getLogger(__name__) @@ -37,6 +38,7 @@ async def create_oauth_url( ], config_id: Annotated[str, Field(description='The configuration ID for the component.')], ctx: Context, + project_id: ProjectIdArg = None, ) -> Annotated[str, Field(description='The OAuth authorization URL.')]: """ Generates an OAuth authorization URL for a Keboola component configuration. diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 6ef85e073..b70ce9cb1 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -19,7 +19,7 @@ from keboola_mcp_server.mcp import ServerState, process_concurrently from keboola_mcp_server.multiproject import MultiProjectMiddleware from keboola_mcp_server.resources.prompts import get_project_system_prompt -from keboola_mcp_server.scope import OAUTH_SESSION_ID_KEY, SCOPE_KEY, SessionScope, resolve_scope_secret +from keboola_mcp_server.scope import OAUTH_SESSION_ID_KEY, SCOPE_KEY, ProjectIdArg, SessionScope, resolve_scope_secret from keboola_mcp_server.workspace import WorkspaceManager LOG = logging.getLogger(__name__) @@ -206,6 +206,7 @@ async def update_project_description( str, Field(description='The new project description text.'), ], + project_id: ProjectIdArg = None, ) -> None: """Updates the description of the current Keboola project.""" client = KeboolaClient.from_state(ctx.session.state) @@ -516,8 +517,9 @@ async def set_project_scope( Mints a scoped access token (narrowed to `project_ids`, optionally read-only) that is used for the rest of the conversation. Read-only tools then run against every scoped project in a single call; - write operations target the active (first) project only. Call this when the user states which - projects to work on; it can be called again any time to re-scope. + write/modify/delete tools take a `project_id` argument naming which scoped project to target (required + once 2+ projects are scoped). Call this when the user states which projects to work on; it can be + called again any time to re-scope. The server does not remember this scope between calls: pass the returned `scope_token` as the `scope_token` argument on every subsequent tool call in this conversation to keep it in effect. @@ -588,8 +590,8 @@ async def set_project_scope( llm_instruction=( ( f'Session scoped to {len(ids)} projects. Read-only tools return results per project. ' - 'Write operations are not fanned out — they target the first scoped project; to write ' - f'elsewhere, re-scope to that project first (confirm with the user). {resend_instruction}' + 'Write operations require a project_id argument naming which scoped project to target ' + f'-- no re-scope needed to switch targets. {resend_instruction}' ) if multi else f'Session scoped to project {ids[0]}. {resend_instruction}' diff --git a/src/keboola_mcp_server/tools/storage/tools.py b/src/keboola_mcp_server/tools/storage/tools.py index 855e97072..6aa4b4deb 100644 --- a/src/keboola_mcp_server/tools/storage/tools.py +++ b/src/keboola_mcp_server/tools/storage/tools.py @@ -23,6 +23,7 @@ toon_serializer_compact, unwrap_results, ) +from keboola_mcp_server.scope import ProjectIdArg from keboola_mcp_server.tools.components.utils import get_nested from keboola_mcp_server.tools.storage.usage import ( ComponentUsageReference, @@ -1026,6 +1027,7 @@ async def update_descriptions( 'Examples: "bucket_id", "bucket_id.table_id", "bucket_id.table_id.column_name"' ), ], + project_id: ProjectIdArg = None, ) -> UpdateDescriptionsOutput: """Updates the description for a Keboola storage item. diff --git a/tests/test_multiproject.py b/tests/test_multiproject.py index a021ee569..826b02f32 100644 --- a/tests/test_multiproject.py +++ b/tests/test_multiproject.py @@ -66,10 +66,11 @@ def _result(text: str) -> ToolResult: [ (None, 'get_tables', True), (SessionScope(project_ids=[11], confirmed=True), 'get_tables', True), - (SessionScope(project_ids=[11, 22], confirmed=True), 'update_config', False), # write: no fan-out + # write, single scoped project: project_id is optional, defaults to the active project. + (SessionScope(project_ids=[11], confirmed=True), 'update_config', False), (SessionScope(project_ids=[11, 22], confirmed=True), 'get_project_info', True), # excluded tool ], - ids=['no_scope', 'single_project', 'write_tool', 'excluded_tool'], + ids=['no_scope', 'single_project', 'write_tool_single_project', 'excluded_tool'], ) async def test_passthrough_calls_once(self, scope, tool_name, read_only) -> None: context, _ = self._ctx(scope, tool_name, read_only) @@ -109,6 +110,72 @@ async def call_next(_): assert len(calls) == 1 assert result.content[0].text == 'projects' + @pytest.mark.asyncio + async def test_write_tool_targets_named_project(self) -> None: + # 2+ scoped projects, project_id names a non-active one: the client (and workspace) are + # swapped to that project for the single call, then restored. + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'update_config', read_only=False, arguments={'project_id': '22'}) + active_clients: list = [] + + async def call_next(_): + active_clients.append(state[KeboolaClient.STATE_KEY]) + return self._result('updated') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert active_clients == ['client-22'] + assert state[KeboolaClient.STATE_KEY] == 'orig-client' # restored + assert result.content[0].text == 'updated' + + @pytest.mark.asyncio + async def test_write_tool_no_swap_for_active_project(self) -> None: + # project_id names the already-active (first) project: no client swap needed. + scope = SessionScope(project_ids=[11, 22], confirmed=True) + context, state = self._ctx(scope, 'update_config', read_only=False, arguments={'project_id': '11'}) + calls = [] + + async def call_next(_): + calls.append(1) + return self._result('updated') + + with patch.object(MultiProjectMiddleware, 'client_for_project', AsyncMock()) as client_for_project: + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + client_for_project.assert_not_called() + assert len(calls) == 1 + assert result.content[0].text == 'updated' + + @pytest.mark.asyncio + async def test_write_tool_ambiguous_without_project_id_raises(self) -> None: + scope = SessionScope(project_ids=[11, 22], confirmed=True) + context, _ = self._ctx(scope, 'update_config', read_only=False, arguments={}) + + async def call_next(_): + raise AssertionError('call_next must not run for an ambiguous write') + + with pytest.raises(ToolError, match='2 projects are scoped'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + + @pytest.mark.asyncio + async def test_write_tool_project_id_outside_scope_raises(self) -> None: + scope = SessionScope(project_ids=[11, 22], confirmed=True) + context, _ = self._ctx(scope, 'update_config', read_only=False, arguments={'project_id': '33'}) + + async def call_next(_): + raise AssertionError('call_next must not run for an out-of-scope project_id') + + with pytest.raises(ToolError, match='outside the current scope'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + @pytest.mark.asyncio async def test_read_tool_fans_out_per_project(self) -> None: scope = SessionScope( From 099d13ecb40bf8a817fbde5deda508b27c67afb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 5 Aug 2026 07:01:51 +0200 Subject: [PATCH 72/89] feat(PSGO-261): surface organization_id/organization_name on get_accessible_projects Piggybacks on the per-project token verify already made for the SQL dialect (no extra API call) -- that response's organization field is the same one get_project_info already reads for organization_id. Co-Authored-By: Claude Sonnet 5 --- TOOLS.md | 12 +++-- src/keboola_mcp_server/tools/project.py | 72 +++++++++++++++---------- tests/tools/test_project.py | 25 ++++----- 3 files changed, 65 insertions(+), 44 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index 39cceef16..3f9c95522 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -50,7 +50,8 @@ providing their configuration IDs. - [modify_streamlit_data_app](#modify_streamlit_data_app): Creates or updates a Streamlit data app. ### Project Tools -- [get_accessible_projects](#get_accessible_projects): Lists the Keboola projects the current login can access across the stack, each with its SQL dialect. +- [get_accessible_projects](#get_accessible_projects): Lists the Keboola projects the current login can access across the stack, each with its SQL +dialect and organization. - [get_project_info](#get_project_info): Retrieves structured information about the current project, including essential context and base instructions for working with it (e. @@ -3766,14 +3767,15 @@ configuration is created e.g. keboola.ex-google-analytics-v4 and keboola.ex-gmai **Description**: -Lists the Keboola projects the current login can access across the stack, each with its SQL dialect. +Lists the Keboola projects the current login can access across the stack, each with its SQL +dialect and organization. Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), present the projects, and ask whether they want to work across all of them or a subset. Then call `set_project_scope` with their choice. This tool compacts several API calls (token introspection -plus a per-project token verify for the SQL dialect) into one result, so the assistant does not -need a separate get_project_info call per project. Pass with_llm_instruction=true on the first -call to also receive the base working instructions grouped by dialect. +plus a per-project token verify for the SQL dialect and organization) into one result, so the +assistant does not need a separate get_project_info call per project. Pass with_llm_instruction=true +on the first call to also receive the base working instructions grouped by dialect. **Input JSON Schema**: diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index b70ce9cb1..b47e920fc 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -288,6 +288,14 @@ def _sql_dialect_from_token(token_data: JsonDict) -> str | None: return None +def _organization_from_token(token_data: JsonDict) -> tuple[str | None, str | None]: + """Derives (organization_id, organization_name) from the token's organization field, the same + field get_project_info reads for organization_id.""" + organization = cast(JsonDict, token_data.get('organization') or {}) + org_id = organization.get('id') + return (str(org_id) if org_id is not None else None, organization.get('name')) + + class AccessibleProject(BaseModel): id: int = Field(description='The project id.') name: str | None = Field(default=None, description='The project name.') @@ -296,6 +304,8 @@ class AccessibleProject(BaseModel): sql_dialect: str | None = Field( default=None, description='The SQL dialect of the project ("Snowflake" or "BigQuery").' ) + organization_id: str | None = Field(default=None, description='The ID of the organization this project belongs to.') + organization_name: str | None = Field(default=None, description='The name of the organization, if known.') class BaseInstructionGroup(BaseModel): @@ -375,19 +385,22 @@ async def _persist_oauth_scope(ctx: Context, scope: SessionScope) -> bool: return True -async def _project_sql_dialect( +async def _project_verify_info( server_state: ServerState, storage_api_url: str, subject_token: str, project_id: int -) -> tuple[int, str | None]: - """Fetches one project's SQL dialect by verifying the parent token narrowed with X-KBC-ProjectId. +) -> tuple[int, str | None, str | None, str | None]: + """Fetches one project's SQL dialect + organization (id, name) via a single token verify, the + parent token narrowed with X-KBC-ProjectId. - No workspace is provisioned — the dialect comes from the token's owner.defaultBackend, so this is - a single cheap Storage API call per project. + No workspace is provisioned — the dialect comes from the token's owner.defaultBackend and the + organization from the token's organization field, so this is one cheap Storage API call per + project (the same call get_project_info makes for a single project). """ per_client = await MultiProjectMiddleware.client_for_project( server_state, storage_api_url, subject_token, project_id, read_only=True ) token_data = await per_client.storage_client.verify_token() - return project_id, _sql_dialect_from_token(token_data) + org_id, org_name = _organization_from_token(token_data) + return project_id, _sql_dialect_from_token(token_data), org_id, org_name @tool_errors() @@ -404,14 +417,15 @@ async def get_accessible_projects( ] = False, ) -> AccessibleProjects: """ - Lists the Keboola projects the current login can access across the stack, each with its SQL dialect. + Lists the Keboola projects the current login can access across the stack, each with its SQL + dialect and organization. Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), present the projects, and ask whether they want to work across all of them or a subset. Then call `set_project_scope` with their choice. This tool compacts several API calls (token introspection - plus a per-project token verify for the SQL dialect) into one result, so the assistant does not - need a separate get_project_info call per project. Pass with_llm_instruction=true on the first - call to also receive the base working instructions grouped by dialect. + plus a per-project token verify for the SQL dialect and organization) into one result, so the + assistant does not need a separate get_project_info call per project. Pass with_llm_instruction=true + on the first call to also receive the base working instructions grouped by dialect. """ client = KeboolaClient.from_state(ctx.session.state) subject_token = await _parent_subject_token(client) @@ -420,33 +434,37 @@ async def get_accessible_projects( scope = ctx.session.state.get(SCOPE_KEY) scoped_ids = scope.project_ids if isinstance(scope, SessionScope) and scope.confirmed else None - # Enrich each project with its SQL dialect (concurrently). Best-effort: a project whose verify - # fails simply keeps sql_dialect=None rather than failing the whole listing. + # Enrich each project with its SQL dialect + organization (concurrently). Best-effort: a project + # whose verify fails simply keeps these fields None rather than failing the whole listing. server_state = ServerState.from_context(ctx) - dialects: dict[int, str | None] = {} + verify_info: dict[int, tuple[str | None, str | None, str | None]] = {} results = await process_concurrently( [p.id for p in introspection.projects], - lambda pid: _project_sql_dialect(server_state, client.storage_api_url, subject_token, pid), + lambda pid: _project_verify_info(server_state, client.storage_api_url, subject_token, pid), ) for result in results: if isinstance(result, asyncio.CancelledError): raise result # never swallow cancellation — let it propagate if isinstance(result, BaseException): - LOG.warning(f'Could not resolve SQL dialect for a project: {result}', exc_info=result) + LOG.warning(f'Could not resolve SQL dialect/organization for a project: {result}', exc_info=result) continue - pid, dialect = result - dialects[pid] = dialect - - projects = [ - AccessibleProject( - id=p.id, - name=p.name, - role=p.role, - in_scope=scoped_ids is not None and p.id in scoped_ids, - sql_dialect=dialects.get(p.id), + pid, dialect, org_id, org_name = result + verify_info[pid] = (dialect, org_id, org_name) + + projects = [] + for p in introspection.projects: + dialect, org_id, org_name = verify_info.get(p.id, (None, None, None)) + projects.append( + AccessibleProject( + id=p.id, + name=p.name, + role=p.role, + in_scope=scoped_ids is not None and p.id in scoped_ids, + sql_dialect=dialect, + organization_id=org_id, + organization_name=org_name, + ) ) - for p in introspection.projects - ] # Optionally attach the base working instructions, grouped by dialect so the (large) prompt is # sent once per distinct dialect rather than duplicated per project. diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 0ff3d4b2c..37f3283cc 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -303,14 +303,15 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock introspect = mocker.patch( 'keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection) ) - # Per-project SQL dialect is resolved via a token verify narrowed by X-KBC-ProjectId; mock that. + # Per-project SQL dialect + organization are resolved via a token verify narrowed by + # X-KBC-ProjectId; mock that. mocker.patch( 'keboola_mcp_server.tools.project.ServerState.from_context', return_value=SimpleNamespace(config=Config()) ) - dialects = {18: 'BigQuery', 83: 'Snowflake'} + verify_info = {18: ('BigQuery', 'org-1', 'Org One'), 83: ('Snowflake', 'org-2', 'Org Two')} mocker.patch( - 'keboola_mcp_server.tools.project._project_sql_dialect', - new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, dialects[pid])), + 'keboola_mcp_server.tools.project._project_verify_info', + new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, *verify_info[pid])), ) # No scope confirmed yet. @@ -318,9 +319,9 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock introspect.assert_awaited_once_with(STACK, subject_token='kbc_at_parent') assert result.user_email == 'm@k.com' - assert [(p.id, p.name, p.role, p.sql_dialect) for p in result.projects] == [ - (18, 'A', 'admin', 'BigQuery'), - (83, 'B', 'admin', 'Snowflake'), + assert [(p.id, p.name, p.role, p.sql_dialect, p.organization_id, p.organization_name) for p in result.projects] == [ + (18, 'A', 'admin', 'BigQuery', 'org-1', 'Org One'), + (83, 'B', 'admin', 'Snowflake', 'org-2', 'Org Two'), ] assert result.scoped_project_ids is None assert result.read_only is None @@ -358,8 +359,8 @@ async def test_get_accessible_projects_llm_instructions_grouped_by_dialect( mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) dialects = {18: 'BigQuery', 86: 'BigQuery', 95: 'Snowflake'} mocker.patch( - 'keboola_mcp_server.tools.project._project_sql_dialect', - new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, dialects[pid])), + 'keboola_mcp_server.tools.project._project_verify_info', + new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, dialects[pid], None, None)), ) result = await get_accessible_projects(mcp_context_client, with_llm_instruction=True) @@ -384,8 +385,8 @@ async def test_get_accessible_projects_unknown_dialect_omits_snowflake_guidance( mocker.patch('keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection)) mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) mocker.patch( - 'keboola_mcp_server.tools.project._project_sql_dialect', - new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, None)), + 'keboola_mcp_server.tools.project._project_verify_info', + new=mocker.AsyncMock(side_effect=lambda _ss, _url, _tok, pid: (pid, None, None, None)), ) result = await get_accessible_projects(mcp_context_client, with_llm_instruction=True) @@ -409,7 +410,7 @@ async def test_get_accessible_projects_logs_dialect_failure_with_traceback( mocker.patch('keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection)) mocker.patch('keboola_mcp_server.tools.project.ServerState.from_context', return_value=mocker.Mock()) mocker.patch( - 'keboola_mcp_server.tools.project._project_sql_dialect', + 'keboola_mcp_server.tools.project._project_verify_info', new=mocker.AsyncMock(side_effect=RuntimeError('verify failed')), ) log_warning = mocker.patch('keboola_mcp_server.tools.project.LOG.warning') From 9e1b96761849a0d980d34423982e7562b104ca6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 6 Aug 2026 16:06:50 +0200 Subject: [PATCH 73/89] fix(PSGO-261): give get_project_info an explicit project_id like write tools get_project_info was in _NO_FANOUT_TOOLS and always reported on whichever project happened to be active (first in scope), with no way to name a different one -- the same "implicit active project" wart the write-tool targeting fix (787eca44) already removed elsewhere. It now takes project_id (required once 2+ projects are scoped, defaults to the single scoped project otherwise), resolved via the same single-target dispatch write tools use (renamed _dispatch_write/_resolve_write_target to _dispatch_single_target/_resolve_single_target since both now share it). Co-Authored-By: Claude Sonnet 5 --- TOOLS.md | 18 +++++++- src/keboola_mcp_server/multiproject.py | 52 +++++++++++++--------- src/keboola_mcp_server/tools/project.py | 4 +- tests/test_multiproject.py | 57 ++++++++++++++++++++++++- tests/test_server.py | 2 +- 5 files changed, 109 insertions(+), 24 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index 3f9c95522..0421a6bec 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -3807,14 +3807,28 @@ including essential context and base instructions for working with it (e.g., transformations, components, workflows, and dependencies). Always call this tool at least once at the start of a conversation -to establish the project context before using other tools. +to establish the project context before using other tools. Reports on exactly one project; +pass `project_id` to pick which when the session is scoped to 2+ projects. **Input JSON Schema**: ```json { "additionalProperties": false, - "properties": {}, + "properties": { + "project_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Target Keboola project id for this write. Required when the session is scoped to 2+ projects; optional (defaults to the single scoped project) otherwise." + } + }, "type": "object" } ``` diff --git a/src/keboola_mcp_server/multiproject.py b/src/keboola_mcp_server/multiproject.py index 2897bb67e..bfd9dceb8 100644 --- a/src/keboola_mcp_server/multiproject.py +++ b/src/keboola_mcp_server/multiproject.py @@ -26,14 +26,15 @@ LOG = logging.getLogger(__name__) -# Tools that must not be fanned out across multiple projects, even when a multi-project scope is -# active and they are read-only: the scope/auth tools operate on the whole-stack token (not a single -# project), and get_project_info resolves through the active project's WorkspaceManager (workspace id -# / sql dialect), so it reports the active project only. -# query_data is intentionally NOT here: the fan-out swaps a per-project WorkspaceManager (see -# MultiProjectMiddleware._swap_project) so a query runs against the workspace of each targeted -# project — narrow to one with the project_ids filter, or run across all scoped projects. -_NO_FANOUT_TOOLS = {'get_accessible_projects', 'set_project_scope', 'get_project_info'} +# Scope/auth tools that operate on the whole-stack token, not a single project -- never fanned out, +# never given a project_id (they don't have one to target). +_NO_FANOUT_TOOLS = {'get_accessible_projects', 'set_project_scope'} + +# Read tools that report on exactly one project (not a list to fan out over) and take an explicit +# project_id argument to say which -- same single-target resolution/swap as a write tool, just +# without the write semantics. get_project_info resolves through the active project's +# WorkspaceManager (workspace id / sql dialect), so it can only ever report one project at a time. +_SINGLE_TARGET_READ_TOOLS = {'get_project_info'} # Optional per-call argument injected on fan-out-eligible read tools to restrict a single call to a # subset of the scoped projects (consumed and stripped by MultiProjectMiddleware.on_call_tool). @@ -49,7 +50,9 @@ class MultiProjectMiddleware(fmw.Middleware): structured content is deep-merged (lists concatenated across projects, counters summed) into one schema-valid object, degrading to count-first with a truncated sample past ``_FANOUT_MAX_ITEMS``. Write tools never fan out: a write always targets exactly one project, named by its own - ``project_id`` argument (required once 2+ projects are scoped) -- see ``_dispatch_write``. + ``project_id`` argument (required once 2+ projects are scoped) -- see + ``_dispatch_single_target``. ``get_project_info`` uses the same single-target resolution + (it reports on the active project's WorkspaceManager, so it can't fan out either). """ async def on_call_tool( @@ -78,12 +81,16 @@ async def on_call_tool( # Bootstrap tools own a real `project_ids` argument, so we must not strip it. if not isinstance(scope, SessionScope) or name in BOOTSTRAP_TOOLS: return await call_next(context) - # Workspace-bound tools always target the active project (no fan-out, filter ignored). + # Whole-stack scope/auth tools: always the active project's client, no project_id to target. if name in _NO_FANOUT_TOOLS: return await call_next(context) + # Single-project-at-a-time tools (get_project_info) and all write tools resolve their own + # explicit project_id the same way -- one target, swap the client, no fan-out. + if name in _SINGLE_TARGET_READ_TOOLS: + return await self._dispatch_single_target(context, call_next, ctx, state, scope) tool = await ctx.fastmcp.get_tool(name) if not is_read_only_tool(tool): - return await self._dispatch_write(context, call_next, ctx, state, scope) + return await self._dispatch_single_target(context, call_next, ctx, state, scope) # Read tool: consume the optional per-call project filter (advertised via on_list_tools) so the # tool never receives it, then narrow this call's target projects to the requested subset. @@ -165,8 +172,9 @@ async def on_call_tool( return self._merge(results, errors) @staticmethod - def _resolve_write_target(scope: SessionScope, project_id: Any) -> int | None: - """Picks the single project a write call targets, or raises if that's ambiguous/invalid. + def _resolve_single_target(scope: SessionScope, project_id: Any) -> int | None: + """Picks the single project a write call or a single-target read targets, or raises if + that's ambiguous/invalid. ``project_id`` is required once 2+ projects are scoped (no more implicit "first project" default); with exactly one scoped project it's optional and defaults to that project. @@ -175,7 +183,7 @@ def _resolve_write_target(scope: SessionScope, project_id: Any) -> int | None: if len(scope.project_ids) >= 2: raise ToolError( f'{len(scope.project_ids)} projects are scoped ({scope.project_ids}). ' - 'Pass project_id= (one of the scoped projects) -- a write targets exactly one project.' + 'Pass project_id= (one of the scoped projects) -- this tool targets exactly one project.' ) return scope.active_project_id try: @@ -189,7 +197,7 @@ def _resolve_write_target(scope: SessionScope, project_id: Any) -> int | None: ) return target - async def _dispatch_write( + async def _dispatch_single_target( self, context: MiddlewareContext[mt.CallToolRequestParams], call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult], @@ -197,12 +205,13 @@ async def _dispatch_write( state: dict[str, Any], scope: SessionScope, ) -> mt.CallToolResult: - """Targets a write/modify/delete tool call at the project named by its ``project_id`` - argument (peeked, not popped -- it's a real declared tool parameter, not middleware-only). + """Targets a write/modify/delete tool, or a single-target read tool (get_project_info), at + the project named by its ``project_id`` argument (peeked, not popped -- it's a real + declared tool parameter, not middleware-only). """ args = getattr(context.message, 'arguments', None) project_id = args.get(PROJECT_ID_ARG) if isinstance(args, dict) else None - target = self._resolve_write_target(scope, project_id) + target = self._resolve_single_target(scope, project_id) if target is None or target == scope.active_project_id: return await call_next(context) @@ -243,7 +252,12 @@ async def on_list_tools( patched: list[Tool] = [] for tool in tools: - if tool.name in BOOTSTRAP_TOOLS or tool.name in _NO_FANOUT_TOOLS or not is_read_only_tool(tool): + if ( + tool.name in BOOTSTRAP_TOOLS + or tool.name in _NO_FANOUT_TOOLS + or tool.name in _SINGLE_TARGET_READ_TOOLS + or not is_read_only_tool(tool) + ): patched.append(tool) continue params = dict(tool.parameters or {}) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index b47e920fc..d8647dfd8 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -220,6 +220,7 @@ async def update_project_description( @tool_errors() async def get_project_info( ctx: Context, + project_id: ProjectIdArg = None, ) -> ProjectInfo: """ Retrieves structured information about the current project, @@ -227,7 +228,8 @@ async def get_project_info( (e.g., transformations, components, workflows, and dependencies). Always call this tool at least once at the start of a conversation - to establish the project context before using other tools. + to establish the project context before using other tools. Reports on exactly one project; + pass `project_id` to pick which when the session is scoped to 2+ projects. """ client = KeboolaClient.from_state(ctx.session.state) links_manager = await ProjectLinksManager.from_client(client) diff --git a/tests/test_multiproject.py b/tests/test_multiproject.py index 826b02f32..faa153086 100644 --- a/tests/test_multiproject.py +++ b/tests/test_multiproject.py @@ -68,7 +68,7 @@ def _result(text: str) -> ToolResult: (SessionScope(project_ids=[11], confirmed=True), 'get_tables', True), # write, single scoped project: project_id is optional, defaults to the active project. (SessionScope(project_ids=[11], confirmed=True), 'update_config', False), - (SessionScope(project_ids=[11, 22], confirmed=True), 'get_project_info', True), # excluded tool + (SessionScope(project_ids=[11, 22], confirmed=True), 'get_accessible_projects', True), # excluded tool ], ids=['no_scope', 'single_project', 'write_tool_single_project', 'excluded_tool'], ) @@ -176,6 +176,61 @@ async def call_next(_): with pytest.raises(ToolError, match='outside the current scope'): await MultiProjectMiddleware().on_call_tool(context, call_next) + @pytest.mark.asyncio + async def test_get_project_info_targets_named_project(self) -> None: + # Same single-target resolution as a write tool: 2+ scoped projects, project_id names a + # non-active one -- swap to it for the call, then restore. + scope = SessionScope(project_ids=[11, 22], scoped_token='kbc_at_s', confirmed=True) + context, state = self._ctx(scope, 'get_project_info', read_only=True, arguments={'project_id': '22'}) + active_clients: list = [] + + async def call_next(_): + active_clients.append(state[KeboolaClient.STATE_KEY]) + return self._result('info') + + with ( + patch.object( + MultiProjectMiddleware, + 'client_for_project', + AsyncMock(side_effect=lambda _ss, _url, _token, pid, _ro: f'client-{pid}'), + ), + patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')), + ): + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + assert active_clients == ['client-22'] + assert state[KeboolaClient.STATE_KEY] == 'orig-client' # restored + assert result.content[0].text == 'info' + + @pytest.mark.asyncio + async def test_get_project_info_defaults_for_single_scoped_project(self) -> None: + # Single scoped project, no project_id given: defaults to it, no swap. + scope = SessionScope(project_ids=[11], confirmed=True) + context, _ = self._ctx(scope, 'get_project_info', read_only=True, arguments={}) + calls = [] + + async def call_next(_): + calls.append(1) + return self._result('info') + + with patch.object(MultiProjectMiddleware, 'client_for_project', AsyncMock()) as client_for_project: + result = await MultiProjectMiddleware().on_call_tool(context, call_next) + + client_for_project.assert_not_called() + assert len(calls) == 1 + assert result.content[0].text == 'info' + + @pytest.mark.asyncio + async def test_get_project_info_ambiguous_without_project_id_raises(self) -> None: + scope = SessionScope(project_ids=[11, 22], confirmed=True) + context, _ = self._ctx(scope, 'get_project_info', read_only=True, arguments={}) + + async def call_next(_): + raise AssertionError('call_next must not run for an ambiguous get_project_info') + + with pytest.raises(ToolError, match='2 projects are scoped'): + await MultiProjectMiddleware().on_call_tool(context, call_next) + @pytest.mark.asyncio async def test_read_tool_fans_out_per_project(self) -> None: scope = SessionScope( diff --git a/tests/test_server.py b/tests/test_server.py index 0b0027190..057230242 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -154,7 +154,7 @@ async def test_tools_input_schema(self): missing_default.append(f'{tool.name}.{prop_name}') missing_properties.sort() - assert missing_properties == ['get_project_info'] + assert missing_properties == [] missing_type.sort() assert not missing_type, f'These tool params have no "type" info: {missing_type}' missing_default.sort() From 12b6ff85ca7bcb9656432cfab4bd5b4fc957fe26 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 6 Aug 2026 16:52:15 +0200 Subject: [PATCH 74/89] fix(PSGO-261): scope_token clutter + OAuth scoped-token never refreshing Two related fixes: 1. Stop advertising/needing scope_token when the transport already guarantees ctx.session persists across requests -- stdio always, and streamable-http with --no-stateless-http (a flag whose own help text already promised this and was never implemented). on_request now reuses an already-confirmed scope straight from ctx.session.state instead of unconditionally rebuilding it every call; on_list_tools and set_project_scope/get_accessible_projects stop advertising/returning scope_token in that mode. 2. Real bug this surfaced: a deployed OAuth session's scoped_token (minted once by set_project_scope) was never refreshed -- unlike the local path, _resolve_local_tokens returned early for deployed sessions without ever checking scope.is_near_expiry. Once that token expired mid-conversation, every fanned-out Storage call started 401ing for the rest of the session with no indication why. Deployed sessions now get the same near-expiry re-mint, persisted back to the OAuth session's Postgres row so it happens once per expiry, not once per request. Extracted the scope-persistence write (project.py's _persist_oauth_scope and this fix's on_request path did the same SessionStore.update_scope call) into a shared scope.persist_scope() helper. Co-Authored-By: Claude Sonnet 5 --- TOOLS.md | 6 +- src/keboola_mcp_server/cli.py | 4 +- src/keboola_mcp_server/config.py | 15 ++ src/keboola_mcp_server/mcp.py | 83 +++++++++-- src/keboola_mcp_server/scope.py | 24 +++- src/keboola_mcp_server/tools/project.py | 44 +++--- tests/test_config.py | 17 ++- tests/test_mcp.py | 176 +++++++++++++++++++++++- tests/tools/test_project.py | 36 +++-- 9 files changed, 361 insertions(+), 44 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index 0421a6bec..2261d010b 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -3850,8 +3850,10 @@ write/modify/delete tools take a `project_id` argument naming which scoped proje once 2+ projects are scoped). Call this when the user states which projects to work on; it can be called again any time to re-scope. -The server does not remember this scope between calls: pass the returned `scope_token` as the -`scope_token` argument on every subsequent tool call in this conversation to keep it in effect. +On most transports the server does not remember this scope between calls: pass the returned +`scope_token` as the `scope_token` argument on every subsequent tool call in this conversation +to keep it in effect. Not needed for a local server or an OAuth-authenticated session, both of +which persist the confirmed scope server-side instead -- `scope_token` is null there. **Input JSON Schema**: diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index e0767a3c2..c80c1268b 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -397,7 +397,9 @@ async def run_server(args: list[str] | None = None) -> None: mcp_server: FastMCP | None = None if parsed_args.transport in ['http-compat', 'streamable-http']: - http_runtime_config = ServerRuntimeInfo('http-compat/streamable-http') + http_runtime_config = ServerRuntimeInfo( + 'http-compat/streamable-http', stateless_http=parsed_args.stateless_http + ) mcp_server, custom_routes = create_server( config, runtime_info=http_runtime_config, custom_routes_handling='return' ) diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index de68ee269..7747a9adb 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -261,6 +261,21 @@ class ServerRuntimeInfo: """The version of the MCP library.""" fastmcp_library_version: str = importlib.metadata.version('fastmcp') """The version of the FastMCP library.""" + stateless_http: bool = True + """Only meaningful for streamable-http: whether the transport was started with the default + stateless session mode (a fresh session per request -- required for scaled/deployed servers + where any replica may handle any request) or `--no-stateless-http` (session pinned by + Mcp-Session-Id, for a single local server). Ignored for stdio, which is inherently + single-session -- see `session_state_persists`.""" + + @property + def session_state_persists(self) -> bool: + """True when the same `ctx.session` object (and thus its `.state` dict) is reused across + requests within one conversation: always for stdio (one process, one session, for the + whole conversation), and for streamable-http only when started with + `--no-stateless-http`. False for the deployed default (`--stateless-http`), where FastMCP + hands every request a fresh session object regardless of what this server does.""" + return self.transport == 'stdio' or not self.stateless_http def build_tracing_headers(runtime_info: ServerRuntimeInfo) -> dict[str, Any]: diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index cddb9974d..4167e072f 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -46,6 +46,7 @@ SCOPE_KEY, SCOPE_TOKEN_ARG, SessionScope, + persist_scope, resolve_scope_secret, ) from keboola_mcp_server.session_store.repository import SessionStore @@ -275,10 +276,30 @@ async def on_request( # SessionStore.update_scope) is read back here instead of round-tripping it as an argument. if scope is None: scope = self._read_persisted_oauth_scope(http_rq) + # stdio and --no-stateless-http streamable-http reuse the same ctx.session object (and + # its .state dict) across every request in the conversation, unlike the stateless-http + # default where FastMCP hands out a fresh session per request. On those transports, a + # scope already confirmed by an earlier set_project_scope call is still sitting in + # ctx.session.state -- reuse it instead of falling back to scope_token/auto-lease, so + # the caller never needs to resend scope_token at all. + if scope is None and runtime_info.session_state_persists: + scope = self._read_persisted_local_scope(ctx) if scope is None and not config.project_id and not is_list: scope = await self._autolease_default_scope(config) if not is_list: + scoped_token_before = scope.scoped_token if scope is not None else None config, scope = await self._resolve_local_tokens(config, scope) + # _resolve_local_tokens re-mints a near-expiry scoped_token for OAuth/deployed + # sessions too (not just local ones) -- persist the refresh to the OAuth session row + # immediately, so it's not silently re-attempted (and re-written) on every single + # request for the rest of this token's lifetime, only once per actual expiry. + if ( + scope is not None + and scope.scoped_token != scoped_token_before + and (oauth_session_id := self._read_oauth_session_id(http_rq)) is not None + and server_state.session_store is not None + ): + await persist_scope(server_state.session_store, oauth_session_id, scope) # TODO: We could probably get rid of the 'state' attribute set on ctx.session and just # pass KeboolaClient and WorkspaceManager instances to a tool as extra parameters. @@ -320,8 +341,19 @@ async def on_list_tools( whether a scope is currently confirmed can't be known while building this response. Showing the parameter always costs nothing when unused and is what lets the caller learn about it before ever calling `set_project_scope`. + + Skipped entirely when this session's transport persists `ctx.session.state` across requests + (stdio, or streamable-http with `--no-stateless-http`) -- there, `on_request` reuses the + already-confirmed scope straight from that state, so `scope_token` is dead weight. """ tools = await call_next(context) + ctx = getattr(context, 'fastmcp_context', None) + if ( + ctx is not None + and isinstance(ctx, Context) + and ServerState.from_context(ctx).runtime_info.session_state_persists + ): + return tools patched: list[Tool] = [] for tool in tools: params = dict(tool.parameters or {}) @@ -450,6 +482,20 @@ def _read_oauth_session_id(cls, http_rq: Request | None) -> str | None: access_token = cls._oauth_access_token(http_rq) return access_token.session_id if access_token is not None else None + @staticmethod + def _read_persisted_local_scope(ctx: Context) -> 'SessionScope | None': + """The scope confirmed by an earlier ``set_project_scope`` call on this same, still-live + ``ctx.session`` -- only ever meaningful when the transport pins one session object across + requests (see ``ServerRuntimeInfo.session_state_persists``); callers must check that first. + """ + # Real session objects (e.g. MiddlewareServerSession) have no `.state` attribute at all + # until this middleware sets one on a prior request -- getattr, not direct access. + state = getattr(ctx.session, 'state', None) + if not isinstance(state, dict): + return None + scope = state.get(SCOPE_KEY) + return scope if isinstance(scope, SessionScope) else None + @classmethod def _is_local_programmatic(cls, config: Config) -> bool: """True for a local (non-deployed) session carrying a Keboola programmatic token.""" @@ -534,18 +580,39 @@ async def _resolve_local_tokens( has explicitly narrowed scope (a minted scoped token is present), that token is re-minted from the parent when it nears expiry. The default (auto-leased) multi-project scope carries no minted token and simply uses the parent token, narrowed per request by ``X-KBC-ProjectId``. - On the deployed server (``KBC_KUBERNETES_TOKEN_PATH`` set) the per-request resolver exchange - already handles token freshness/narrowing once ``project_id`` is known -- but it only runs - once ``project_id`` is known, and nothing else threads a confirmed scope's active project id - into ``config`` for a deployed session. Without that, ``create_session_state`` keeps building - the active client from the unscoped whole-stack token with no ``X-KBC-ProjectId``, so every - call after ``set_project_scope`` 401s even though scoping itself succeeded. So still apply - just the active project id here for deployed sessions; the token itself is left alone since - the resolver-exchange path (keyed off that project id) handles narrowing it correctly. + On the deployed server (``KBC_KUBERNETES_TOKEN_PATH`` set), ``config.storage_token`` is + already the freshly-refreshed OAuth ``kbc_access_token`` (refreshed by + ``SimpleOAuthProvider.load_access_token``'s lazy refresh before this ever runs) -- that part + needs no help here. Nothing else threads a confirmed scope's active project id into + ``config`` for a deployed session, though: without it, ``create_session_state`` keeps + building the active client from the unscoped whole-stack token with no ``X-KBC-ProjectId``, + so every call after ``set_project_scope`` 401s even though scoping itself succeeded -- apply + just the active project id here. The confirmed scope's own ``scoped_token`` (minted once by + ``set_project_scope``, used by ``MultiProjectMiddleware`` for every fanned-out project once + 2+ are scoped -- including the first) *does* need the same near-expiry re-mint the local + branch below does, or it silently starts 401ing mid-conversation once it expires, with no + refresh ever attempted for the rest of the session (`on_request` persists the refreshed + token back to the OAuth session row afterward, so this happens at most once per expiry, not + every request). """ if not cls._is_local_programmatic(config): if scope and scope.project_ids and not config.project_id: config = dataclasses.replace(config, project_id=str(scope.active_project_id)) + if scope is not None and scope.scoped_token is not None and scope.is_near_expiry: + try: + minted = await exchange_scoped_token( + config.storage_api_url, + subject_token=strip_bearer(config.storage_token), + project_ids=scope.project_ids, + read_only=scope.read_only, + ) + scope = dataclasses.replace( + scope, scoped_token=minted.access_token, scoped_expires_at=minted.expires_at + ) + except Exception as e: + # Don't break the session if re-minting fails -- the caller keeps using the + # (possibly already-expired) scoped_token, same failure mode as before this fix. + LOG.warning(f'Could not refresh the deployed session scoped token: {e}', exc_info=True) return config, scope # Strip any inbound `Bearer ` scheme; introspect/exchange helpers add the scheme themselves, diff --git a/src/keboola_mcp_server/scope.py b/src/keboola_mcp_server/scope.py index 0a10d0b23..8a7d5a6e7 100644 --- a/src/keboola_mcp_server/scope.py +++ b/src/keboola_mcp_server/scope.py @@ -8,13 +8,17 @@ import dataclasses import secrets import time -from typing import Annotated, Optional +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Annotated, Optional from pydantic import Field from keboola_mcp_server.config import Config from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt +if TYPE_CHECKING: + from keboola_mcp_server.session_store.repository import SessionStore + SCOPE_KEY = 'project_scope' # Declared on every write/modify/delete tool; consumed by MultiProjectMiddleware.on_call_tool to @@ -95,3 +99,21 @@ def from_token(cls, token: str, secret: str) -> 'SessionScope': """Inverse of ``to_token``. Raises on a missing/invalid/tampered token -- callers should treat any exception as "no scope" rather than fail the request.""" return cls(**decode_jwt(token, secret)) + + +async def persist_scope(session_store: 'SessionStore', session_id: str, scope: SessionScope) -> None: + """Writes ``scope`` onto the OAuth session row ``session_id`` -- shared by ``set_project_scope`` + (a fresh confirmation) and ``SessionStateMiddleware.on_request`` (a near-expiry re-mint), so both + persist a refreshed ``scoped_token`` the same way.""" + await session_store.update_scope( + session_id, + project_ids=scope.project_ids, + read_only=scope.read_only, + confirmed=scope.confirmed, + scoped_token=scope.scoped_token, + scoped_expires_at=( + datetime.fromtimestamp(scope.scoped_expires_at, tz=timezone.utc) + if scope.scoped_expires_at is not None + else None + ), + ) diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index d8647dfd8..774c96a1a 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -1,6 +1,5 @@ import asyncio import logging -from datetime import datetime, timezone from typing import Annotated, Optional, cast import httpx @@ -19,7 +18,14 @@ from keboola_mcp_server.mcp import ServerState, process_concurrently from keboola_mcp_server.multiproject import MultiProjectMiddleware from keboola_mcp_server.resources.prompts import get_project_system_prompt -from keboola_mcp_server.scope import OAUTH_SESSION_ID_KEY, SCOPE_KEY, ProjectIdArg, SessionScope, resolve_scope_secret +from keboola_mcp_server.scope import ( + OAUTH_SESSION_ID_KEY, + SCOPE_KEY, + ProjectIdArg, + SessionScope, + persist_scope, + resolve_scope_secret, +) from keboola_mcp_server.workspace import WorkspaceManager LOG = logging.getLogger(__name__) @@ -372,18 +378,7 @@ async def _persist_oauth_scope(ctx: Context, scope: SessionScope) -> bool: session_store = ServerState.from_context(ctx).session_store if session_store is None: return False - await session_store.update_scope( - session_id, - project_ids=scope.project_ids, - read_only=scope.read_only, - confirmed=scope.confirmed, - scoped_token=scope.scoped_token, - scoped_expires_at=( - datetime.fromtimestamp(scope.scoped_expires_at, tz=timezone.utc) - if scope.scoped_expires_at is not None - else None - ), - ) + await persist_scope(session_store, session_id, scope) return True @@ -492,13 +487,15 @@ async def get_accessible_projects( 'projects or a subset, then call "set_project_scope" with the chosen project ids. Never write ' 'to more than one project without explicit user confirmation.' ) - is_oauth_persisted = False + is_persisted = False else: - is_oauth_persisted = bool(ctx.session.state.get(OAUTH_SESSION_ID_KEY)) + is_persisted = ( + bool(ctx.session.state.get(OAUTH_SESSION_ID_KEY)) or server_state.runtime_info.session_state_persists + ) instruction = ( f'Session is currently scoped to {len(scoped_ids)} project(s). Call "set_project_scope" to ' 'change the scope.' - if is_oauth_persisted + if is_persisted else f'Session is currently scoped to {len(scoped_ids)} project(s). Resend "scope_token" on every ' 'subsequent tool call to keep it in effect; call "set_project_scope" to change the scope.' ) @@ -509,7 +506,7 @@ async def get_accessible_projects( read_only=scope.read_only if scoped_ids is not None else None, scope_token=( scope.to_token(resolve_scope_secret(server_state.config)) - if scoped_ids is not None and not is_oauth_persisted + if scoped_ids is not None and not is_persisted else None ), base_instructions=base_instructions, @@ -541,8 +538,10 @@ async def set_project_scope( once 2+ projects are scoped). Call this when the user states which projects to work on; it can be called again any time to re-scope. - The server does not remember this scope between calls: pass the returned `scope_token` as the - `scope_token` argument on every subsequent tool call in this conversation to keep it in effect. + On most transports the server does not remember this scope between calls: pass the returned + `scope_token` as the `scope_token` argument on every subsequent tool call in this conversation + to keep it in effect. Not needed for a local server or an OAuth-authenticated session, both of + which persist the confirmed scope server-side instead -- `scope_token` is null there. """ client = KeboolaClient.from_state(ctx.session.state) parent_token = await _parent_subject_token(client) @@ -595,8 +594,9 @@ async def set_project_scope( LOG.debug(f'Could not send tools/list_changed after scoping: {e}') multi = len(ids) > 1 - persisted = await _persist_oauth_scope(ctx, scope) - scope_token = None if persisted else scope.to_token(resolve_scope_secret(ServerState.from_context(ctx).config)) + server_state = ServerState.from_context(ctx) + persisted = await _persist_oauth_scope(ctx, scope) or server_state.runtime_info.session_state_persists + scope_token = None if persisted else scope.to_token(resolve_scope_secret(server_state.config)) resend_instruction = ( 'The server persists this scope server-side for the rest of the conversation -- no need to resend it.' if persisted diff --git a/tests/test_config.py b/tests/test_config.py index f78cec800..0ace9bcde 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,7 +3,7 @@ import pytest -from keboola_mcp_server.config import Config, get_env_storage_api_url, is_same_stack +from keboola_mcp_server.config import Config, ServerRuntimeInfo, get_env_storage_api_url, is_same_stack class TestConfig: @@ -198,3 +198,18 @@ def test_is_same_stack(self, url: str | None, other_url: str | None, expected: b assert is_same_stack(url, other_url) is expected # The comparison is symmetric. assert is_same_stack(other_url, url) is expected + + +class TestServerRuntimeInfoSessionStatePersists: + def test_stdio_always_persists_regardless_of_stateless_http(self) -> None: + # stdio is one process/one session for the whole conversation -- the flag is meaningless there. + assert ServerRuntimeInfo(transport='stdio', stateless_http=True).session_state_persists is True + assert ServerRuntimeInfo(transport='stdio', stateless_http=False).session_state_persists is True + + def test_streamable_http_follows_stateless_http_flag(self) -> None: + assert ServerRuntimeInfo(transport='streamable-http', stateless_http=True).session_state_persists is False + assert ServerRuntimeInfo(transport='streamable-http', stateless_http=False).session_state_persists is True + + def test_defaults_to_stateless(self) -> None: + # Matches the CLI's --stateless-http default (scaled/deployed-safe). + assert ServerRuntimeInfo(transport='streamable-http').session_state_persists is False diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 9fc3ad1f7..240126cf0 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -24,7 +24,7 @@ toon_serializer, unwrap_results, ) -from keboola_mcp_server.scope import SessionScope, resolve_scope_secret +from keboola_mcp_server.scope import SCOPE_KEY, SessionScope, resolve_scope_secret from keboola_mcp_server.workspace import WorkspaceManager @@ -865,6 +865,67 @@ def test_apply_request_config_injects_exchanged_session_token(self): assert out_config.storage_token == 'kbc_at_exchanged' assert is_programmatic_token(out_config.storage_token) + @pytest.mark.asyncio + async def test_on_request_persists_remint_of_expiring_oauth_scoped_token(self, monkeypatch) -> None: + # End-to-end regression for the bug this fixes: a deployed OAuth session's scoped_token + # expiring mid-conversation silently 401ed every fanned-out call thereafter, since nothing + # ever refreshed it. on_request must re-mint it (via _resolve_local_tokens) and persist the + # refresh to the OAuth session row so it's fixed for the rest of the session, not just once. + from datetime import datetime, timezone + + from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser + from starlette.requests import Request + + from keboola_mcp_server.oauth import ProxyAccessToken + + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + access_token = ProxyAccessToken( + token='mcp_proxy', + client_id='claude.ai', + scopes=['claudai', 'projectless'], + expires_at=int(time.time() + 3600), + kbc_access_token='kbc_at_fresh_oauth', + session_id='session-1', + scope_project_ids=[18, 83], + scope_confirmed=True, + scope_scoped_token='kbc_at_stale', + scope_scoped_expires_at=datetime.fromtimestamp(time.time() - 1, tz=timezone.utc), + ) + http_rq = Request({'type': 'http', 'headers': [], 'user': AuthenticatedUser(access_token)}) + + session_store = AsyncMock() + config = Config(storage_api_url='https://connection.test.keboola.com') + server_state = ServerState( + config=config, + runtime_info=ServerRuntimeInfo(transport='http-compat/streamable-http'), + session_store=session_store, + ) + session = SimpleNamespace(state={}) + ctx = MagicMock(spec=Context) + ctx.session = session + ctx.request_context.lifespan_context = server_state + context = SimpleNamespace(message=SimpleNamespace(arguments={}), method='tools/call', fastmcp_context=ctx) + + minted = SimpleNamespace(access_token='kbc_at_reminted', expires_at=time.time() + 3600) + + async def call_next(_): + return 'ok' + + with ( + patch('keboola_mcp_server.mcp.get_http_request_or_none', return_value=http_rq), + patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock(return_value=minted)), + patch.object(SessionStateMiddleware, 'create_session_state', AsyncMock(return_value={})), + ): + result = await SessionStateMiddleware().on_request(context, call_next) + + assert result == 'ok' + session_store.update_scope.assert_awaited_once() + call = session_store.update_scope.await_args + assert call.args == ('session-1',) + assert call.kwargs['project_ids'] == [18, 83] + assert call.kwargs['scoped_token'] == 'kbc_at_reminted' + assert call.kwargs['scoped_expires_at'] == datetime.fromtimestamp(minted.expires_at, tz=timezone.utc) + class TestProgrammaticTokenForwarding: """A programmatic token (kbc_at_/kbc_pat_) is always forwarded downstream as a Bearer (PSGO-261). @@ -1001,6 +1062,54 @@ async def test_deployed_no_scope_or_already_set_project_id_is_noop(self, monkeyp assert out_config is config # project_id already set -- not overwritten assert out_scope is scope + @pytest.mark.asyncio + async def test_deployed_near_expiry_scoped_token_is_reminted(self, monkeypatch) -> None: + # Regression: a deployed (OAuth) session's scoped_token was never refreshed once minted by + # set_project_scope, so it silently started 401ing every fanned-out call once it expired, + # for the rest of the conversation, with no error pointing at the real cause. + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_fresh_oauth') + scope = SessionScope( + project_ids=[18, 83], scoped_token='kbc_at_stale', scoped_expires_at=time.time() - 1, confirmed=True + ) + minted = SimpleNamespace(access_token='kbc_at_reminted', expires_at=time.time() + 3600) + with patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock(return_value=minted)) as exch: + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + exch.assert_awaited_once_with( + 'https://connection.keboola.com', subject_token='kbc_at_fresh_oauth', project_ids=[18, 83], read_only=False + ) + assert out_scope.scoped_token == 'kbc_at_reminted' + assert out_scope.scoped_expires_at == minted.expires_at + assert out_config.project_id == '18' + + @pytest.mark.asyncio + async def test_deployed_scoped_token_not_near_expiry_is_untouched(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_fresh_oauth') + scope = SessionScope( + project_ids=[18, 83], scoped_token='kbc_at_live', scoped_expires_at=time.time() + 3600, confirmed=True + ) + with patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock()) as exch: + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + exch.assert_not_awaited() + assert out_scope is scope + + @pytest.mark.asyncio + async def test_deployed_remint_failure_keeps_old_scope(self, monkeypatch) -> None: + # Same failure mode as before this fix (the caller keeps using the stale token and 401s + # downstream) rather than crashing the request outright. + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_fresh_oauth') + scope = SessionScope( + project_ids=[18], scoped_token='kbc_at_stale', scoped_expires_at=time.time() - 1, confirmed=True + ) + with patch( + 'keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock(side_effect=RuntimeError('exchange down')) + ): + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + assert out_scope is scope + assert out_scope.scoped_token == 'kbc_at_stale' + @pytest.mark.asyncio async def test_legacy_token_is_noop(self, monkeypatch) -> None: monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) @@ -1226,6 +1335,30 @@ def test_read_oauth_session_id_returns_none_for_non_oauth_request(self) -> None: assert SessionStateMiddleware._read_oauth_session_id(None) is None assert SessionStateMiddleware._read_oauth_session_id(SimpleNamespace(scope={})) is None + def test_read_persisted_local_scope_returns_confirmed_scope_from_session_state(self) -> None: + scope = SessionScope(project_ids=[18, 83], confirmed=True) + ctx = SimpleNamespace(session=SimpleNamespace(state={SCOPE_KEY: scope})) + assert SessionStateMiddleware._read_persisted_local_scope(ctx) is scope + + def test_read_persisted_local_scope_returns_none_when_absent_or_invalid(self) -> None: + assert ( + SessionStateMiddleware._read_persisted_local_scope(SimpleNamespace(session=SimpleNamespace(state={}))) + is None + ) + assert ( + SessionStateMiddleware._read_persisted_local_scope( + SimpleNamespace(session=SimpleNamespace(state={SCOPE_KEY: 'not-a-scope'})) + ) + is None + ) + assert ( + SessionStateMiddleware._read_persisted_local_scope(SimpleNamespace(session=SimpleNamespace(state=None))) + is None + ) + # Regression: a real (non-mocked) session object has no `.state` attribute at all until this + # middleware sets one on a prior request -- must not raise AttributeError on the very first request. + assert SessionStateMiddleware._read_persisted_local_scope(SimpleNamespace(session=object())) is None + @pytest.mark.asyncio async def test_on_list_tools_advertises_scope_token_unconditionally(self) -> None: # Unlike MultiProjectMiddleware's `project_ids` filter, this must show up even with no scope @@ -1242,3 +1375,44 @@ async def call_next(_): tools = await SessionStateMiddleware().on_list_tools(context, call_next) assert 'scope_token' in tools[0].parameters['properties'] + + @pytest.mark.asyncio + async def test_on_list_tools_skips_scope_token_when_session_state_persists(self) -> None: + # stdio (and --no-stateless-http streamable-http) keep ctx.session.state across requests -- + # on_request reuses an already-confirmed scope straight from it, so there's nothing for the + # caller to resend and advertising scope_token would just be clutter. + tool = _tool('get_tables', read_only=True) + tool.parameters = {'type': 'object', 'properties': {}} + + ctx = MagicMock(spec=Context) + ctx.request_context.lifespan_context = ServerState( + config=Config(), runtime_info=ServerRuntimeInfo(transport='stdio') + ) + context = SimpleNamespace(method='tools/list', fastmcp_context=ctx) + + async def call_next(_): + return [tool] + + tools = await SessionStateMiddleware().on_list_tools(context, call_next) + + assert 'scope_token' not in tools[0].parameters['properties'] + + @pytest.mark.asyncio + async def test_on_list_tools_advertises_scope_token_when_session_state_does_not_persist(self) -> None: + tool = _tool('get_tables', read_only=True) + tool.parameters = {'type': 'object', 'properties': {}} + tool.model_copy = lambda update, _t=tool: SimpleNamespace(name=_t.name, parameters=update['parameters']) + + ctx = MagicMock(spec=Context) + ctx.request_context.lifespan_context = ServerState( + config=Config(), + runtime_info=ServerRuntimeInfo(transport='http-compat/streamable-http', stateless_http=True), + ) + context = SimpleNamespace(method='tools/list', fastmcp_context=ctx) + + async def call_next(_): + return [tool] + + tools = await SessionStateMiddleware().on_list_tools(context, call_next) + + assert 'scope_token' in tools[0].parameters['properties'] diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 37f3283cc..6742678a0 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -306,7 +306,8 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock # Per-project SQL dialect + organization are resolved via a token verify narrowed by # X-KBC-ProjectId; mock that. mocker.patch( - 'keboola_mcp_server.tools.project.ServerState.from_context', return_value=SimpleNamespace(config=Config()) + 'keboola_mcp_server.tools.project.ServerState.from_context', + return_value=SimpleNamespace(config=Config(), runtime_info=ServerRuntimeInfo(transport='stdio')), ) verify_info = {18: ('BigQuery', 'org-1', 'Org One'), 83: ('Snowflake', 'org-2', 'Org Two')} mocker.patch( @@ -329,17 +330,14 @@ async def test_get_accessible_projects(mcp_context_client: Context, mocker: Mock assert result.scope_token is None assert all(not p.in_scope for p in result.projects) - # Once scoped, the current scope is surfaced on the projects and at the top level, and echoed - # back as a scope_token the caller must resend on later calls (the server does not remember it). + # Once scoped, the current scope is surfaced on the projects and at the top level. On this + # (stdio) transport ctx.session.state persists across requests, so no scope_token is needed. mcp_context_client.session.state[SCOPE_KEY] = SessionScope(project_ids=[83], read_only=True, confirmed=True) result = await get_accessible_projects(mcp_context_client) assert result.scoped_project_ids == [83] assert result.read_only is True assert [(p.id, p.in_scope) for p in result.projects] == [(18, False), (83, True)] - assert result.scope_token is not None - assert SessionScope.from_token(result.scope_token, resolve_scope_secret(Config())) == SessionScope( - project_ids=[83], read_only=True, confirmed=True - ) + assert result.scope_token is None @pytest.mark.asyncio @@ -439,9 +437,31 @@ async def test_set_project_scope_subset_exchanges_and_stores( scope = mcp_context_client.session.state[SCOPE_KEY] assert scope.scoped_token == 'kbc_at_scoped' assert scope.project_ids == [18, 83] - # The server does not remember this scope between calls; the caller must resend scope_token. + # mcp_context_client's runtime is transport='stdio', which persists ctx.session.state across + # requests (ServerRuntimeInfo.session_state_persists) -- no scope_token needed to keep it in effect. + assert result.scope_token is None + assert 'persists this scope server-side' in result.llm_instruction + + +@pytest.mark.asyncio +async def test_set_project_scope_returns_scope_token_when_session_does_not_persist( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + # Deployed default: stateless-http streamable-http, a fresh ctx.session per request -- nothing + # server-side survives between calls, so the caller must resend scope_token. + mcp_context_client.request_context.lifespan_context = ServerState( + Config(), ServerRuntimeInfo(transport='http-compat/streamable-http', stateless_http=True) + ) + _prep_client(mcp_context_client, mocker) + minted = SimpleNamespace(access_token='kbc_at_scoped', expires_at=time.time() + 3600, read_only=False) + mocker.patch('keboola_mcp_server.tools.project.exchange_scoped_token', new=mocker.AsyncMock(return_value=minted)) + + result = await set_project_scope(mcp_context_client, project_ids=[18, 83]) + + scope = mcp_context_client.session.state[SCOPE_KEY] assert result.scope_token is not None assert SessionScope.from_token(result.scope_token, resolve_scope_secret(Config())) == scope + assert 'does not remember this scope' in result.llm_instruction @pytest.mark.asyncio From 1c4fefd94e04009e506cd8c381d82c13e6f3fc34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Fri, 7 Aug 2026 12:32:20 +0200 Subject: [PATCH 75/89] feat(PSGO-261): Kai (header-token) session-scope persistence + post-rebase ruff fixes Adds the design from pat_token_support/RFC.md increment 6: server-side scope persistence for deployed, non-OAuth, programmatic-token sessions (Kai), keyed by sha256(conversation_id:user_id) since the raw kbc_at_/kbc_pat_ token isn't stable across Kai's own refresh. New kai_sessions table (migration 0004), PostgresKaiScopeStore, read-side fallback in SessionStateMiddleware.on_request with subset-check invalidation (drop the whole scope if a previously-scoped project is no longer reachable), and write-side persistence from set_project_scope. Also fixes lint issues ruff 0.16 (picked up from the just-rebased RFC branch) flagged in pre-existing code: mutable class-attribute default, a two-call startswith merged into one, a redundant exception repr in a log call, non-tz-aware date.today() in tests, and a few unused tuple-unpack variables. --- feature_spec/pat_token_support/RFC.md | 96 ++++++++++ src/keboola_mcp_server/auth_login.py | 4 +- src/keboola_mcp_server/clients/auth_bridge.py | 2 +- src/keboola_mcp_server/mcp.py | 44 +++++ src/keboola_mcp_server/oauth.py | 2 +- src/keboola_mcp_server/server.py | 10 +- .../session_store/kai_scope.py | 91 +++++++++ .../migrations/0004_kai_sessions.sql | 15 ++ src/keboola_mcp_server/tools/project.py | 35 +++- tests/session_store/conftest.py | 18 +- tests/session_store/test_kai_scope.py | 46 +++++ tests/session_store/test_migrator.py | 16 +- tests/session_store/test_retention.py | 12 +- tests/test_mcp.py | 179 +++++++++++++++++- tests/test_multiproject.py | 2 +- tests/tools/test_project.py | 30 +++ 16 files changed, 582 insertions(+), 20 deletions(-) create mode 100644 src/keboola_mcp_server/session_store/kai_scope.py create mode 100644 src/keboola_mcp_server/session_store/migrations/0004_kai_sessions.sql create mode 100644 tests/session_store/test_kai_scope.py diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index ef9e1a29d..850921460 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -703,3 +703,99 @@ over a header/middleware-only approach" recommendation, including the ambiguity raising when 2+ projects are active and no `project_id` is given). Everything else in PR #500 (token taxonomy, append-only project registry, Kai integration flow, 24h idle refresh) is already covered by this RFC and the as-built code under different names. + +--- + +# Extension: Kai (header-token) session-scope persistence (PSGO-261, increment 6) + +## Context + +Kai currently authorizes with a legacy, project-bound Storage token and will transition to a +stack-wide programmatic token (`kbc_at_`/`kbc_pat_`), refreshed by Kai's own regime rather than +this server's PKCE store. Once that happens, every request Kai sends carries an **unscoped** +whole-stack token, and `set_project_scope`/`get_accessible_projects` need the same server-side +scope persistence OAuth sessions already get (§"Transport note", increment 5) — pushing the +`scope_token` round-trip onto an LLM-driven client is unreliable (nothing guarantees it survives +compaction, a fresh turn, or simply gets echoed back correctly). + +OAuth's persistence trick doesn't transfer directly, though: `SimpleOAuthProvider` mints its own +opaque token at login, so `sha256(opaque_token)` (`session_store/repository.py`) is a stable +Postgres key for the life of the session even as the *real* Keboola credential is refreshed +underneath it. Kai's raw token has no such stability — confirmed against the actual refresh code +in `auth_login.py`: `refresh_tokens()` returns a brand-new access-token string on every rotation, +and `create_pat()`'s response carries no separate token-id to key on either. Hashing the raw +inbound token would therefore silently drop the persisted scope on every Kai-side refresh. + +## Required behavior + +- **Persistence key:** `sha256(f'{conversation_id}:{user_id}')`, where `conversation_id` is the + existing `X-Conversation-Id`-derived `Config.conversation_id` (already flowing on every request + for tracing, confirmed stable for the life of one Kai chat session) and `user_id` is + `Introspection.user_id` (`auth_login.py`) resolved from the *current* request's token. Binding + to `user_id` — not just `conversation_id` — closes the gap a low-entropy or client-chosen + `conversation_id` would otherwise leave open: a collision (or reuse) only matches an existing row + if it also resolves to the same underlying Keboola identity, so a mismatched identity is a cache + miss, not a leaked scope, with no separate post-lookup equality check to forget. +- **Stored row:** `project_ids`, `read_only`, `confirmed` only — no `scoped_token`/expiry fields, + since Kai refreshes its own Keboola credential independently of this table; nothing here needs + to track the parent token's freshness. +- **Read-time validation, not a superset/subset hash:** a hash can only express exact-match + equality, not "grew is fine, shrank is not" — so the monotonicity rule is enforced in code, at + read time, against introspection data already being fetched: if + `set(row.project_ids) - {p.id for p in introspection.projects}` is non-empty (some previously + scoped project is no longer reachable), the row is dropped and the scope is treated as + unconfirmed. Projects *added* to the token's reach never invalidate an existing scope, since the + subset relation still holds. +- **On invalidation, drop the whole scope** (not auto-narrow to the intersection) — force a full + `get_accessible_projects` → `set_project_scope` redo so an access change is surfaced to the user + rather than silently absorbed. +- Applies only to deployed, non-OAuth, programmatic-token sessions with a `conversation_id` + present (`deployed_sa_token_path()` set, `is_programmatic_token(config.storage_token)`, no + `AuthenticatedUser`/`ProxyAccessToken` on the request). OAuth sessions keep using + `oauth_sessions`; local PKCE sessions keep using `ctx.session.state` (`session_state_persists`); + neither is affected by this table. + +## Resolution strategy + +- New table `kai_sessions` (migration `0004_kai_sessions.sql`), unpartitioned initially — same + starting point `oauth_sessions` had before partitioning became necessary (increment/migration + `0002`); add partitioning here too if/when retention needs it. +- New `session_store/kai_scope.py`: `KaiScope` (data) + `KaiScopeStore` (Protocol) + + `PostgresKaiScopeStore` (impl), deliberately **not** folded into `SessionStore`/`OAuthSession` — + different key scheme (composite hash vs. opaque-token hash), no encrypted credential fields (no + secret is stored, just a project-id list and two flags), different invalidation semantics + (subset-check + drop vs. revoke). Keeping it a separate small store avoids overloading the + OAuth-shaped `SessionStore` protocol with a second, structurally different session concept. + Same lazy-pool-on-first-use pattern as `PostgresSessionStore`. +- `ServerState.kai_scope_store: KaiScopeStore | None`, constructed in `server.py` whenever + `config.postgres_dsn` is set — **independent of whether OAuth is configured**, since Kai's path + needs no `oauth_client_id`/`session_encryption_key` (no OAuth login, no encrypted fields here). +- `SessionStateMiddleware.on_request` (`mcp.py`): a new fallback, + `_read_persisted_kai_scope`, slotted after `_read_persisted_local_scope` and before + `_autolease_default_scope` — mirrors `_read_persisted_oauth_scope`'s position in the chain but + reads from `kai_scope_store` instead of the OAuth session row, gated on the "deployed, + non-OAuth, programmatic, has conversation_id" condition above. Skipped for `/list` like every + other network-touching step in this chain. +- `tools/project.py`'s `set_project_scope`: a new `_persist_kai_scope`, called alongside the + existing `_persist_oauth_scope` — whichever one applies persists server-side and suppresses + `scope_token` in the response (`persisted = await _persist_oauth_scope(...) or await + _persist_kai_scope(...) or session_state_persists`, unchanged shape, one more branch). + +## Decisions (increment 6) + +- **Server-side persistence over client-side round-tripping**, confirmed: pushing scope state + into Kai/the LLM's own context is fragile (no guarantee of faithful round-trip across turns or + compaction); persisting server-side, looked up automatically on every request, needs no + cooperation from the calling LLM beyond sending the `conversation_id` header it already sends. +- **Composite key (`conversation_id` + `user_id`) over either alone.** `conversation_id` alone is + client-supplied and not guaranteed high-entropy; `user_id` alone is not conversation-scoped + (would incorrectly share scope across unrelated chats from the same person). Together they give + a key that's both stable across Kai's token refreshes and safe against a `conversation_id` + collision or reuse. +- **Drop-whole-scope over auto-narrow on a reachability shrink** — an explicit user decision + (over the friendlier-but-quieter auto-narrow-to-intersection alternative): surfacing an access + change via a forced re-scope beats silently continuing with whatever subset still works. +- **A new store/table over extending `oauth_sessions`/`SessionStore`** — the two session kinds + differ enough (key scheme, no encrypted fields, no OAuth-specific lifecycle) that folding Kai + scope into the OAuth-shaped protocol would blur its single responsibility for no real code + reuse (the two stores would share almost no method bodies). diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index ac76446b8..c3a9d3951 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -23,7 +23,7 @@ from dataclasses import asdict, dataclass from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path -from typing import cast +from typing import ClassVar, cast from urllib.parse import urlparse import httpx @@ -462,7 +462,7 @@ def forget_tokens(storage_api_url: str | None = None) -> bool: class _CallbackHandler(BaseHTTPRequestHandler): - result: dict = {} + result: ClassVar[dict] = {} def do_GET(self) -> None: # noqa: N802 (BaseHTTPRequestHandler API) query = urllib.parse.parse_qs(urlparse(self.path).query) diff --git a/src/keboola_mcp_server/clients/auth_bridge.py b/src/keboola_mcp_server/clients/auth_bridge.py index f93c7c0be..3409d8930 100644 --- a/src/keboola_mcp_server/clients/auth_bridge.py +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -45,7 +45,7 @@ def is_programmatic_token(token: str | None) -> bool: if not token: return False bare = strip_bearer(token) - return bare.startswith(_ACCESS_TOKEN_PREFIX) or bare.startswith(_PAT_PREFIX) + return bare.startswith((_ACCESS_TOKEN_PREFIX, _PAT_PREFIX)) class OAuthTokenExchangeError(RuntimeError): diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 4167e072f..8b7095c7b 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -49,6 +49,7 @@ persist_scope, resolve_scope_secret, ) +from keboola_mcp_server.session_store.kai_scope import KaiScopeStore from keboola_mcp_server.session_store.repository import SessionStore from keboola_mcp_server.tools.constants import ( BOOTSTRAP_TOOLS, @@ -121,6 +122,7 @@ class ServerState: config: Config runtime_info: ServerRuntimeInfo session_store: SessionStore | None = None + kai_scope_store: KaiScopeStore | None = None @property def own_stack_storage_api_url(self) -> str | None: @@ -284,6 +286,19 @@ async def on_request( # the caller never needs to resend scope_token at all. if scope is None and runtime_info.session_state_persists: scope = self._read_persisted_local_scope(ctx) + # Deployed, non-OAuth, programmatic-token sessions (Kai) carry no MCP-minted + # identifier and no persistent ctx.session -- kai_session_scope RFC persists their + # confirmed scope server-side instead, keyed by (conversation_id, token user id). + if ( + scope is None + and not is_list + and config.conversation_id + and deployed_sa_token_path() + and self._oauth_access_token(http_rq) is None + and is_programmatic_token(config.storage_token) + and server_state.kai_scope_store is not None + ): + scope = await self._read_persisted_kai_scope(config, server_state.kai_scope_store) if scope is None and not config.project_id and not is_list: scope = await self._autolease_default_scope(config) if not is_list: @@ -496,6 +511,35 @@ def _read_persisted_local_scope(ctx: Context) -> 'SessionScope | None': scope = state.get(SCOPE_KEY) return scope if isinstance(scope, SessionScope) else None + @classmethod + async def _read_persisted_kai_scope(cls, config: Config, store: KaiScopeStore) -> 'SessionScope | None': + """The scope confirmed by an earlier `set_project_scope` call on this Kai conversation, + looked up by `sha256(conversation_id:user_id)` (kai_session_scope RFC) rather than by + token hash, since Kai refreshes its raw token independently and its value isn't stable + across that refresh. Drops (and forgets) the stored scope -- rather than auto-narrowing or + trusting stale access -- if a previously scoped project is no longer reachable by the + current token; callers see this as "no scope yet" and are steered back through + get_accessible_projects / set_project_scope. + """ + try: + introspection = await introspect_token( + config.storage_api_url, subject_token=strip_bearer(config.storage_token) + ) + except Exception as e: + LOG.warning(f'Could not introspect Kai token for persisted scope lookup: {e}', exc_info=True) + return None + if introspection.user_id is None: + return None + stored = await store.get(config.conversation_id, introspection.user_id) + if stored is None: + return None + current_project_ids = {p.id for p in introspection.projects} + if not set(stored.project_ids).issubset(current_project_ids): + LOG.info('Persisted Kai scope references a project no longer reachable; dropping it.') + await store.drop(config.conversation_id, introspection.user_id) + return None + return SessionScope(project_ids=stored.project_ids, read_only=stored.read_only, confirmed=stored.confirmed) + @classmethod def _is_local_programmatic(cls, config: Config) -> bool: """True for a local (non-deployed) session carrying a Keboola programmatic token.""" diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index 3683d7e03..8094699df 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -549,7 +549,7 @@ async def exchange_refresh_token( error='invalid_grant', error_description=f'Failed to refresh token: status={e.response.status_code}' ) from e except httpx.HTTPError as e: - LOG.exception(f'[exchange_refresh_token] Could not reach Connection to refresh session: {e}') + LOG.exception('[exchange_refresh_token] Could not reach Connection to refresh session') raise TokenError( error='invalid_grant', error_description=f'Failed to refresh token: could not reach Connection ({e}).' ) from e diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index fed8c1c91..de1c20740 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -24,6 +24,7 @@ from keboola_mcp_server.preview import preview_config_diff from keboola_mcp_server.prompts.add_prompts import add_keboola_prompts from keboola_mcp_server.session_store.crypto import resolve_encryption_key +from keboola_mcp_server.session_store.kai_scope import PostgresKaiScopeStore from keboola_mcp_server.session_store.repository import PostgresSessionStore from keboola_mcp_server.tools.components.tools import add_component_tools from keboola_mcp_server.tools.data_apps import add_data_app_tools @@ -244,9 +245,16 @@ def create_server( oauth_provider = None session_store = None + # Kai session-scope persistence (pat_token_support/RFC.md, increment 6) needs only a Postgres + # DSN -- unlike OAuth sessions it stores no credential material, so no encryption key or + # oauth_client_id/secret is required. Independent of whether OAuth is configured above. + kai_scope_store = PostgresKaiScopeStore(config.postgres_dsn) if config.postgres_dsn else None + # Initialize FastMCP server with system lifespan LOG.info(f'Creating server with config: {config}') - server_state = ServerState(config=config, runtime_info=runtime_info, session_store=session_store) + server_state = ServerState( + config=config, runtime_info=runtime_info, session_store=session_store, kai_scope_store=kai_scope_store + ) mcp = KeboolaMcpServer( name='Keboola MCP Server', instructions=( diff --git a/src/keboola_mcp_server/session_store/kai_scope.py b/src/keboola_mcp_server/session_store/kai_scope.py new file mode 100644 index 000000000..55dc68b33 --- /dev/null +++ b/src/keboola_mcp_server/session_store/kai_scope.py @@ -0,0 +1,91 @@ +"""Postgres-backed scope persistence for deployed header-token (Kai) sessions. + +See ``feature_spec/pat_token_support/RFC.md`` ("Kai (header-token) session-scope persistence", +increment 6) for the design. Unlike OAuth sessions (`session_store/repository.py`), Kai's raw +Keboola token is refreshed by Kai's own regime and is not stable across that refresh, so rows are +keyed by ``sha256(conversation_id:user_id)`` rather than a hash of the token itself. No credential +material is stored here, so unlike `OAuthSession` nothing needs encryption at rest. +""" + +import asyncio +import dataclasses +import hashlib +from typing import Protocol + +import asyncpg + + +def _hash_key(conversation_id: str, user_id: int) -> bytes: + return hashlib.sha256(f'{conversation_id}:{user_id}'.encode()).digest() + + +@dataclasses.dataclass(frozen=True) +class KaiScope: + project_ids: list[int] + read_only: bool + confirmed: bool + + +class KaiScopeStore(Protocol): + async def get(self, conversation_id: str, user_id: int) -> KaiScope | None: ... + + async def upsert( + self, conversation_id: str, user_id: int, *, project_ids: list[int], read_only: bool, confirmed: bool + ) -> None: ... + + async def drop(self, conversation_id: str, user_id: int) -> None: ... + + +class PostgresKaiScopeStore: + """Schema migrations are NOT applied here -- see `PostgresSessionStore`'s docstring for why + (same reasoning, same `migrate` CLI/Job applies both). The connection pool is created lazily, + on first use, for the same sync-construction reason `PostgresSessionStore` does. + """ + + def __init__(self, dsn: str) -> None: + self._dsn = dsn + self._pool: asyncpg.Pool | None = None + self._pool_lock = asyncio.Lock() + + async def _get_pool(self) -> asyncpg.Pool: + if self._pool is None: + async with self._pool_lock: + if self._pool is None: # re-check: another task may have won the lock race first + self._pool = await asyncpg.create_pool(self._dsn) + return self._pool + + async def close(self) -> None: + if self._pool is not None: + await self._pool.close() + + async def get(self, conversation_id: str, user_id: int) -> KaiScope | None: + pool = await self._get_pool() + row = await pool.fetchrow( + 'UPDATE kai_sessions SET last_used_at = now() WHERE session_key = $1 RETURNING *', + _hash_key(conversation_id, user_id), + ) + if row is None: + return None + return KaiScope(project_ids=list(row['project_ids']), read_only=row['read_only'], confirmed=row['confirmed']) + + async def upsert( + self, conversation_id: str, user_id: int, *, project_ids: list[int], read_only: bool, confirmed: bool + ) -> None: + pool = await self._get_pool() + await pool.execute( + """ + INSERT INTO kai_sessions (session_key, project_ids, read_only, confirmed) + VALUES ($1, $2, $3, $4) + ON CONFLICT (session_key) DO UPDATE + SET project_ids = EXCLUDED.project_ids, read_only = EXCLUDED.read_only, + confirmed = EXCLUDED.confirmed, updated_at = now(), last_used_at = now() + """, + _hash_key(conversation_id, user_id), + project_ids, + read_only, + confirmed, + ) + + async def drop(self, conversation_id: str, user_id: int) -> None: + pool = await self._get_pool() + await pool.execute('DELETE FROM kai_sessions WHERE session_key = $1', _hash_key(conversation_id, user_id)) diff --git a/src/keboola_mcp_server/session_store/migrations/0004_kai_sessions.sql b/src/keboola_mcp_server/session_store/migrations/0004_kai_sessions.sql new file mode 100644 index 000000000..b8e04cc56 --- /dev/null +++ b/src/keboola_mcp_server/session_store/migrations/0004_kai_sessions.sql @@ -0,0 +1,15 @@ +-- kai_session_scope RFC (pat_token_support/RFC.md, increment 6): persisted multi-project scope +-- for deployed header-token (Kai) sessions. Kai's raw kbc_at_/kbc_pat_ token is refreshed +-- independently of this server and is not stable across that refresh, so rows are keyed by +-- sha256(conversation_id:user_id) instead of a token hash. No credential material is stored here +-- (unlike oauth_sessions) -- just the confirmed scope -- so no encryption is needed. + +CREATE TABLE kai_sessions ( + session_key BYTEA PRIMARY KEY, + project_ids INTEGER[] NOT NULL, + read_only BOOLEAN NOT NULL DEFAULT FALSE, + confirmed BOOLEAN NOT NULL DEFAULT FALSE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + last_used_at TIMESTAMPTZ NOT NULL DEFAULT now() +); diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 774c96a1a..1bc1dd25f 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -15,7 +15,7 @@ from keboola_mcp_server.config import MetadataField, deployed_sa_token_path from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager -from keboola_mcp_server.mcp import ServerState, process_concurrently +from keboola_mcp_server.mcp import CONVERSATION_ID, ServerState, process_concurrently from keboola_mcp_server.multiproject import MultiProjectMiddleware from keboola_mcp_server.resources.prompts import get_project_system_prompt from keboola_mcp_server.scope import ( @@ -370,7 +370,7 @@ class ProjectScope(BaseModel): async def _persist_oauth_scope(ctx: Context, scope: SessionScope) -> bool: """Persists ``scope`` on the caller's OAuth session row, if this is an OAuth-authenticated session (see mcp.OAUTH_SESSION_ID_KEY). No-op (returns False) for PAT/header-token sessions, - which have no session row to persist against -- those keep relying on scope_token. + which either use `_persist_kai_scope` or keep relying on scope_token. """ session_id = ctx.session.state.get(OAUTH_SESSION_ID_KEY) if not session_id: @@ -382,6 +382,31 @@ async def _persist_oauth_scope(ctx: Context, scope: SessionScope) -> bool: return True +async def _persist_kai_scope(ctx: Context, scope: SessionScope, client: KeboolaClient, parent_token: str) -> bool: + """Persists ``scope`` for a deployed, non-OAuth, programmatic-token session (Kai) -- see + `feature_spec/pat_token_support/RFC.md` ("Kai (header-token) session-scope persistence"). + No-op (returns False) when this isn't such a session, or no conversation_id/store is available. + """ + conversation_id = ctx.session.state.get(CONVERSATION_ID) + if not conversation_id or not deployed_sa_token_path() or not is_programmatic_token(client.bearer_token): + return False + server_state = ServerState.from_context(ctx) + store = server_state.kai_scope_store + if store is None: + return False + introspection = await introspect_token(client.storage_api_url, subject_token=parent_token) + if introspection.user_id is None: + return False + await store.upsert( + conversation_id, + introspection.user_id, + project_ids=scope.project_ids, + read_only=scope.read_only, + confirmed=scope.confirmed, + ) + return True + + async def _project_verify_info( server_state: ServerState, storage_api_url: str, subject_token: str, project_id: int ) -> tuple[int, str | None, str | None, str | None]: @@ -595,7 +620,11 @@ async def set_project_scope( multi = len(ids) > 1 server_state = ServerState.from_context(ctx) - persisted = await _persist_oauth_scope(ctx, scope) or server_state.runtime_info.session_state_persists + persisted = ( + await _persist_oauth_scope(ctx, scope) + or await _persist_kai_scope(ctx, scope, client, parent_token) + or server_state.runtime_info.session_state_persists + ) scope_token = None if persisted else scope.to_token(resolve_scope_secret(server_state.config)) resend_instruction = ( 'The server persists this scope server-side for the rest of the conversation -- no need to resend it.' diff --git a/tests/session_store/conftest.py b/tests/session_store/conftest.py index 5aa70e940..9e5850539 100644 --- a/tests/session_store/conftest.py +++ b/tests/session_store/conftest.py @@ -5,6 +5,7 @@ import pytest_asyncio from keboola_mcp_server.session_store.crypto import KEY_SIZE +from keboola_mcp_server.session_store.kai_scope import PostgresKaiScopeStore from keboola_mcp_server.session_store.migrator import apply_migrations from keboola_mcp_server.session_store.repository import PostgresSessionStore @@ -34,7 +35,7 @@ async def store(): try: # Clean slate per test: drop, then re-apply migrations -- standing in for the migration # Job that would normally run once, ahead of the app, in a real deployment. - await pool.execute('DROP TABLE IF EXISTS oauth_sessions, schema_migrations CASCADE') + await pool.execute('DROP TABLE IF EXISTS oauth_sessions, kai_sessions, schema_migrations CASCADE') await apply_migrations(pool) finally: await pool.close() @@ -43,3 +44,18 @@ async def store(): yield s finally: await s.close() + + +@pytest_asyncio.fixture +async def kai_store(): + pool = await asyncpg.create_pool(TEST_DSN) + try: + await pool.execute('DROP TABLE IF EXISTS oauth_sessions, kai_sessions, schema_migrations CASCADE') + await apply_migrations(pool) + finally: + await pool.close() + s = PostgresKaiScopeStore(TEST_DSN) + try: + yield s + finally: + await s.close() diff --git a/tests/session_store/test_kai_scope.py b/tests/session_store/test_kai_scope.py new file mode 100644 index 000000000..62f6b584f --- /dev/null +++ b/tests/session_store/test_kai_scope.py @@ -0,0 +1,46 @@ +import pytest + +from tests.session_store.conftest import requires_postgres + +pytestmark = [pytest.mark.asyncio, requires_postgres] + + +async def test_upsert_and_get(kai_store) -> None: + await kai_store.upsert('conv-1', 42, project_ids=[18, 83], read_only=False, confirmed=True) + + scope = await kai_store.get('conv-1', 42) + + assert scope is not None + assert scope.project_ids == [18, 83] + assert scope.read_only is False + assert scope.confirmed is True + + +async def test_get_unknown_returns_none(kai_store) -> None: + assert await kai_store.get('does-not-exist', 1) is None + + +async def test_different_user_id_is_a_different_row(kai_store) -> None: + # Same conversation_id, different user -- must not collide (the whole point of the + # composite key, see pat_token_support/RFC.md increment 6). + await kai_store.upsert('conv-1', 42, project_ids=[18], read_only=False, confirmed=True) + + assert await kai_store.get('conv-1', 999) is None + + +async def test_upsert_overwrites_existing_row(kai_store) -> None: + await kai_store.upsert('conv-1', 42, project_ids=[18], read_only=False, confirmed=True) + await kai_store.upsert('conv-1', 42, project_ids=[18, 83], read_only=True, confirmed=True) + + scope = await kai_store.get('conv-1', 42) + + assert scope.project_ids == [18, 83] + assert scope.read_only is True + + +async def test_drop_removes_the_row(kai_store) -> None: + await kai_store.upsert('conv-1', 42, project_ids=[18], read_only=False, confirmed=True) + + await kai_store.drop('conv-1', 42) + + assert await kai_store.get('conv-1', 42) is None diff --git a/tests/session_store/test_migrator.py b/tests/session_store/test_migrator.py index 3e6c6a65c..8deb1fad8 100644 --- a/tests/session_store/test_migrator.py +++ b/tests/session_store/test_migrator.py @@ -12,7 +12,7 @@ async def _clean_slate(): pool = await asyncpg.create_pool(TEST_DSN) try: - await pool.execute('DROP TABLE IF EXISTS oauth_sessions, schema_migrations CASCADE') + await pool.execute('DROP TABLE IF EXISTS oauth_sessions, kai_sessions, schema_migrations CASCADE') finally: await pool.close() @@ -25,6 +25,7 @@ async def test_applies_migrations_once() -> None: '0001_oauth_sessions.sql', '0002_partition_oauth_sessions.sql', '0003_default_partition_unique_indexes.sql', + '0004_kai_sessions.sql', ] # Re-running is a no-op -- the table already exists, so re-applying the DDL would fail @@ -68,6 +69,19 @@ async def test_partitions_table_with_default_catch_all() -> None: await pool.close() +async def test_creates_kai_sessions_table() -> None: + pool = await asyncpg.create_pool(TEST_DSN) + try: + await apply_migrations(pool) + columns = await pool.fetch( + "SELECT column_name FROM information_schema.columns WHERE table_name = 'kai_sessions'" + ) + names = {r['column_name'] for r in columns} + assert {'session_key', 'project_ids', 'read_only', 'confirmed'} <= names + finally: + await pool.close() + + async def test_default_partition_rejects_duplicate_access_token_hash() -> None: # Regression test: the parent's (access_token_hash, created_at) index alone doesn't reject a # duplicate hash (created_at differs per row) -- migration 0003's plain index on the diff --git a/tests/session_store/test_retention.py b/tests/session_store/test_retention.py index f17e83ace..a6101af49 100644 --- a/tests/session_store/test_retention.py +++ b/tests/session_store/test_retention.py @@ -30,7 +30,7 @@ class TestEnsurePartitions: async def _clean_slate(self): pool = await asyncpg.create_pool(TEST_DSN) try: - await pool.execute('DROP TABLE IF EXISTS oauth_sessions, schema_migrations CASCADE') + await pool.execute('DROP TABLE IF EXISTS oauth_sessions, kai_sessions, schema_migrations CASCADE') await apply_migrations(pool) finally: await pool.close() @@ -56,7 +56,7 @@ async def test_is_idempotent(self) -> None: async def test_drops_only_partitions_older_than_retention(self) -> None: pool = await asyncpg.create_pool(TEST_DSN) try: - this_month = _month_start(date.today()) + this_month = _month_start(datetime.now(tz=timezone.utc).date()) stale = _add_months(this_month, -3) kept = _add_months(this_month, -1) for month_start in (stale, kept): @@ -80,7 +80,7 @@ async def test_drops_partition_exactly_retention_months_old(self) -> None: # partition dated retention_months back (2 months old) is dropped, not kept. pool = await asyncpg.create_pool(TEST_DSN) try: - this_month = _month_start(date.today()) + this_month = _month_start(datetime.now(tz=timezone.utc).date()) boundary = _add_months(this_month, -2) name = f'oauth_sessions_{boundary:%Y_%m}' end = _add_months(boundary, 1) @@ -105,7 +105,7 @@ async def test_creates_partition_when_default_has_overlapping_rows(self) -> None for name in await self._existing_partitions(pool): await pool.execute(f'DROP TABLE {name}') - this_month = _month_start(date.today()) + this_month = _month_start(datetime.now(tz=timezone.utc).date()) mid_month = datetime(this_month.year, this_month.month, this_month.day, tzinfo=timezone.utc) + timedelta( days=1 ) @@ -140,7 +140,7 @@ async def test_created_partition_rejects_duplicate_access_token_hash(self) -> No await pool.execute(f'DROP TABLE {name}') result = await ensure_partitions(pool) - this_month_partition = f'oauth_sessions_{_month_start(date.today()):%Y_%m}' + this_month_partition = f'oauth_sessions_{_month_start(datetime.now(tz=timezone.utc).date()):%Y_%m}' assert this_month_partition in result['created'] insert = ( @@ -163,7 +163,7 @@ async def test_creates_missing_current_and_next_month(self) -> None: result = await ensure_partitions(pool) - this_month = _month_start(date.today()) + this_month = _month_start(datetime.now(tz=timezone.utc).date()) next_month = _add_months(this_month, 1) expected = {f'oauth_sessions_{this_month:%Y_%m}', f'oauth_sessions_{next_month:%Y_%m}'} assert set(result['created']) == expected diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 240126cf0..701eb4ffc 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,5 @@ import asyncio +import dataclasses import time from datetime import datetime, timedelta, timezone from types import SimpleNamespace @@ -926,6 +927,178 @@ async def call_next(_): assert call.kwargs['scoped_token'] == 'kbc_at_reminted' assert call.kwargs['scoped_expires_at'] == datetime.fromtimestamp(minted.expires_at, tz=timezone.utc) + @pytest.mark.asyncio + async def test_on_request_applies_persisted_kai_scope(self, monkeypatch) -> None: + # A deployed, non-OAuth, programmatic-token session (Kai) with no scope_token argument and + # no session_state_persists must fall back to the kai_scope_store, not auto-lease default. + from starlette.requests import Request + + from keboola_mcp_server.auth_login import Introspection, ProjectAccess + from keboola_mcp_server.session_store.kai_scope import KaiScope + + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + http_rq = Request({'type': 'http', 'headers': [(b'x-conversation-id', b'conv-1')], 'user': None}) + + kai_scope_store = AsyncMock() + kai_scope_store.get.return_value = KaiScope(project_ids=[18], read_only=False, confirmed=True) + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_kai') + server_state = ServerState( + config=config, + runtime_info=ServerRuntimeInfo(transport='http-compat/streamable-http'), + kai_scope_store=kai_scope_store, + ) + session = SimpleNamespace(state={}) + ctx = MagicMock(spec=Context) + ctx.session = session + ctx.request_context.lifespan_context = server_state + context = SimpleNamespace(message=SimpleNamespace(arguments={}), method='tools/call', fastmcp_context=ctx) + + captured_scopes = [] + + async def fake_create_session_state(cfg, _runtime_info, readonly=None, *, own_stack_storage_api_url): + return {} + + async def call_next(_): + captured_scopes.append(ctx.session.state.get(SCOPE_KEY)) + return 'ok' + + with ( + patch('keboola_mcp_server.mcp.get_http_request_or_none', return_value=http_rq), + patch.object(SessionStateMiddleware, 'create_session_state', side_effect=fake_create_session_state), + patch( + 'keboola_mcp_server.mcp.introspect_token', + AsyncMock( + return_value=Introspection( + user_id=42, user_email=None, user_name=None, projects=[ProjectAccess(id=18)] + ) + ), + ), + ): + result = await SessionStateMiddleware().on_request(context, call_next) + + assert result == 'ok' + kai_scope_store.get.assert_awaited_once_with('conv-1', 42) + assert captured_scopes == [SessionScope(project_ids=[18], read_only=False, confirmed=True)] + + +class TestReadPersistedKaiScope: + """Kai session-scope persistence (pat_token_support/RFC.md, increment 6): + SessionStateMiddleware._read_persisted_kai_scope.""" + + @staticmethod + def _introspection(project_ids: list[int], user_id: int | None = 42): + from keboola_mcp_server.auth_login import Introspection, ProjectAccess + + return Introspection( + user_id=user_id, + user_email='kai@keboola.com', + user_name='Kai', + projects=[ProjectAccess(id=pid) for pid in project_ids], + ) + + @pytest.mark.asyncio + async def test_returns_stored_scope_when_still_reachable(self) -> None: + from keboola_mcp_server.session_store.kai_scope import KaiScope + + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_x') + config = dataclasses.replace(config, conversation_id='conv-1') + store = AsyncMock() + store.get.return_value = KaiScope(project_ids=[18], read_only=False, confirmed=True) + + with patch( + 'keboola_mcp_server.mcp.introspect_token', + AsyncMock(return_value=self._introspection([18, 83])), + ): + scope = await SessionStateMiddleware._read_persisted_kai_scope(config, store) + + assert scope is not None + assert scope.project_ids == [18] + assert scope.confirmed is True + store.get.assert_awaited_once_with('conv-1', 42) + store.drop.assert_not_awaited() + + @pytest.mark.asyncio + async def test_drops_scope_when_a_project_is_no_longer_reachable(self) -> None: + from keboola_mcp_server.session_store.kai_scope import KaiScope + + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_x') + config = dataclasses.replace(config, conversation_id='conv-1') + store = AsyncMock() + store.get.return_value = KaiScope(project_ids=[18, 83], read_only=False, confirmed=True) + + with patch( + 'keboola_mcp_server.mcp.introspect_token', + AsyncMock(return_value=self._introspection([18])), # 83 dropped out + ): + scope = await SessionStateMiddleware._read_persisted_kai_scope(config, store) + + assert scope is None + store.drop.assert_awaited_once_with('conv-1', 42) + + @pytest.mark.asyncio + async def test_added_projects_do_not_invalidate_the_stored_scope(self) -> None: + from keboola_mcp_server.session_store.kai_scope import KaiScope + + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_x') + config = dataclasses.replace(config, conversation_id='conv-1') + store = AsyncMock() + store.get.return_value = KaiScope(project_ids=[18], read_only=False, confirmed=True) + + with patch( + 'keboola_mcp_server.mcp.introspect_token', + AsyncMock(return_value=self._introspection([18, 999])), # gained access to 999 + ): + scope = await SessionStateMiddleware._read_persisted_kai_scope(config, store) + + assert scope is not None + assert scope.project_ids == [18] + store.drop.assert_not_awaited() + + @pytest.mark.asyncio + async def test_no_stored_row_returns_none(self) -> None: + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_x') + config = dataclasses.replace(config, conversation_id='conv-1') + store = AsyncMock() + store.get.return_value = None + + with patch( + 'keboola_mcp_server.mcp.introspect_token', + AsyncMock(return_value=self._introspection([18])), + ): + scope = await SessionStateMiddleware._read_persisted_kai_scope(config, store) + + assert scope is None + + @pytest.mark.asyncio + async def test_unresolvable_user_id_returns_none_without_lookup(self) -> None: + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_x') + config = dataclasses.replace(config, conversation_id='conv-1') + store = AsyncMock() + + with patch( + 'keboola_mcp_server.mcp.introspect_token', + AsyncMock(return_value=self._introspection([18], user_id=None)), + ): + scope = await SessionStateMiddleware._read_persisted_kai_scope(config, store) + + assert scope is None + store.get.assert_not_awaited() + + @pytest.mark.asyncio + async def test_introspection_failure_returns_none(self) -> None: + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_x') + config = dataclasses.replace(config, conversation_id='conv-1') + store = AsyncMock() + + with patch( + 'keboola_mcp_server.mcp.introspect_token', + AsyncMock(side_effect=RuntimeError('network down')), + ): + scope = await SessionStateMiddleware._read_persisted_kai_scope(config, store) + + assert scope is None + store.get.assert_not_awaited() + class TestProgrammaticTokenForwarding: """A programmatic token (kbc_at_/kbc_pat_) is always forwarded downstream as a Bearer (PSGO-261). @@ -1090,7 +1263,7 @@ async def test_deployed_scoped_token_not_near_expiry_is_untouched(self, monkeypa project_ids=[18, 83], scoped_token='kbc_at_live', scoped_expires_at=time.time() + 3600, confirmed=True ) with patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock()) as exch: - out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + _out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) exch.assert_not_awaited() assert out_scope is scope @@ -1106,7 +1279,7 @@ async def test_deployed_remint_failure_keeps_old_scope(self, monkeypatch) -> Non with patch( 'keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock(side_effect=RuntimeError('exchange down')) ): - out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + _out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) assert out_scope is scope assert out_scope.scoped_token == 'kbc_at_stale' @@ -1153,7 +1326,7 @@ async def test_fresh_scoped_token_is_not_reminted(self, monkeypatch) -> None: patch('keboola_mcp_server.mcp.get_access_token', AsyncMock(return_value='kbc_at_parent')), patch('keboola_mcp_server.mcp.exchange_scoped_token', AsyncMock()) as exch, ): - out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) + out_config, _out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) exch.assert_not_awaited() assert out_config.storage_token == 'kbc_at_live' diff --git a/tests/test_multiproject.py b/tests/test_multiproject.py index faa153086..eb522a3a4 100644 --- a/tests/test_multiproject.py +++ b/tests/test_multiproject.py @@ -140,7 +140,7 @@ async def call_next(_): async def test_write_tool_no_swap_for_active_project(self) -> None: # project_id names the already-active (first) project: no client swap needed. scope = SessionScope(project_ids=[11, 22], confirmed=True) - context, state = self._ctx(scope, 'update_config', read_only=False, arguments={'project_id': '11'}) + context, _state = self._ctx(scope, 'update_config', read_only=False, arguments={'project_id': '11'}) calls = [] async def call_next(_): diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 6742678a0..d69295640 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -494,6 +494,36 @@ async def test_set_project_scope_persists_to_db_and_omits_scope_token_for_oauth_ assert 'no need to resend' in result.llm_instruction +@pytest.mark.asyncio +async def test_set_project_scope_persists_to_kai_scope_store_and_omits_scope_token( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + # A deployed, non-OAuth, programmatic-token session (Kai) persists the confirmed scope to + # kai_scope_store, keyed by (conversation_id, introspected user id) -- pat_token_support/RFC.md + # "Kai (header-token) session-scope persistence". No scope_token is needed afterward. + _prep_client(mcp_context_client, mocker) + mocker.patch('keboola_mcp_server.tools.project.deployed_sa_token_path', return_value='/var/run/secrets/token') + kai_scope_store = mocker.Mock() + kai_scope_store.upsert = mocker.AsyncMock() + mcp_context_client.request_context.lifespan_context = ServerState( + config=Config(), + runtime_info=ServerRuntimeInfo(transport='http-compat/streamable-http'), + kai_scope_store=kai_scope_store, + ) + introspection = SimpleNamespace(user_id=42, user_email='kai@keboola.com', projects=[]) + mocker.patch('keboola_mcp_server.tools.project.introspect_token', new=mocker.AsyncMock(return_value=introspection)) + minted = SimpleNamespace(access_token='kbc_at_scoped', expires_at=1234.0, read_only=False) + mocker.patch('keboola_mcp_server.tools.project.exchange_scoped_token', new=mocker.AsyncMock(return_value=minted)) + + result = await set_project_scope(mcp_context_client, project_ids=[18, 83]) + + kai_scope_store.upsert.assert_awaited_once_with( + 'convo-1234', 42, project_ids=[18, 83], read_only=False, confirmed=True + ) + assert result.scope_token is None + assert 'no need to resend' in result.llm_instruction + + @pytest.mark.asyncio async def test_set_project_scope_all_introspects_then_exchanges( mcp_context_client: Context, mocker: MockerFixture From 56b035b56e6ec3ba2a55beb94561372ff9a5dcc5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Sun, 9 Aug 2026 10:54:38 +0200 Subject: [PATCH 76/89] docs(PSGO-261): RFC increment 7 -- security hardening response to PR review Documents all 9 findings from Tomas Fejfar's review of PR #604 (7 confirmed as-stated, 1 confirmed-but-broader, 1 refuted/already-mitigated) and the fix design for each, before any implementation, per this repo's RFC-first convention. --- feature_spec/pat_token_support/RFC.md | 180 ++++++++++++++++++++++++++ 1 file changed, 180 insertions(+) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 850921460..dea4b84d6 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -799,3 +799,183 @@ inbound token would therefore silently drop the persisted scope on every Kai-sid differ enough (key scheme, no encrypted fields, no OAuth-specific lifecycle) that folding Kai scope into the OAuth-shaped protocol would blur its single responsibility for no real code reuse (the two stores would share almost no method bodies). + +--- + +# Extension: Security hardening — response to review (PSGO-261, increment 7) + +## Context + +Tomas Fejfar's review of PR #604 (2026-08-07, "Agentic review") raised 9 concerns. Each was +independently re-verified against the as-built code (file:line evidence) and cross-checked with a +second, independent security-review pass before any fix was designed — this section documents +what was actually found, not just what was claimed, since two items turned out different from +the original framing (one narrower, one broader; see below). + +## Verified findings + +1. **Header injection into `Config` → forgeable `scope_token`, CONFIRMED.** + `SessionStateMiddleware.apply_request_config` calls `config.replace_by(http_rq.headers)` with + no allowlist; `Config._read_options` matches *any* dataclass field against an `X-{name}` + header, including `jwt_secret`. Since `resolve_scope_secret(config)` reads `config.jwt_secret` + from that same per-request config, an `X-Jwt-Secret` header lets a caller choose the HMAC key + that both signs and verifies their own `scope_token` — full `project_ids` forgery. +2. **`scope_token` embeds a live bearer token, signed but not encrypted, CONFIRMED — broader than + first framed.** `jwt_utils.py`'s `encode_jwt`/`decode_jwt` are JWS (signature only) over + gzip+JSON; the payload is base64+gunzip-recoverable by anyone, without the secret. + `SessionScope.scoped_token` — a real, live Keboola access token, not just non-secret metadata + like `project_ids` — is itself a dataclass field, so it's embedded verbatim in the + client-visible token returned by `set_project_scope`/`get_accessible_projects` and resent as a + tool-call argument on every subsequent call: it lands in LLM context, client transcripts, and + client-side logs. No `exp` enforcement; decode failures (tampered, expired, or malformed) all + collapse into the same "no scope" outcome, with no revocation path. +3. **`read_only=True` fails open, CONFIRMED — broader scope than reported.** Not just + single-project scopes as originally described: `MultiProjectMiddleware` skips `_swap_project` + (the only code path that ever passes `readonly=scope.read_only` into a `KeboolaClient`) + whenever a call targets `scope.active_project_id` — true for every single-project scope *and* + the first/active project of any multi-project scope. `SessionStateMiddleware.create_session_state` + never passes `readonly=` at all from `on_request`, regardless of scope. So the active + project's writes are never locally read-only-restricted — enforcement depends entirely on the + minted `scoped_token` being genuinely read-only server-side, which doesn't exist when the + `/v1/auth/pat/exchange` call fails. Only *non-active* projects in a 2+ project scope get real + local enforcement today (via `client_for_project(readonly=scope.read_only or None)`). +4. **`normalize_storage_api_url` is a prefix check, not a domain allowlist, CONFIRMED.** + `hostname.startswith('connection.')` lets `connection.attacker.tld` pass. `is_same_stack` is a + correct exact-host match, but it's only ever applied when the server has its own configured + stack (`own_stack_storage_api_url` set); a server with no stack of its own (local mode, by + design, since it must accept the caller's URL) has no equivalent check before a caller-supplied + `X-Storage-Api-Url` host receives the live bearer token. +5. **`resolve_encryption_key`'s silent process-local fallback — REFUTED, already mitigated.** + `server.py` already refuses to start (`raise RuntimeError`) if OAuth is configured + (`oauth_client_id`/`oauth_client_secret` both set) without `KBC_SESSION_ENCRYPTION_KEY`, and + `PostgresSessionStore` is never constructed via any other path. The cross-replica + silent-decrypt-failure scenario the review described can't actually happen today. Documented + here so it isn't re-flagged as a live gap. +6. **MFA codes as CLI arguments, CONFIRMED.** `login --totp`/`--recovery` are plain `argparse` + string options — visible in shell history and `ps`/`/proc//cmdline` for the process + lifetime. Recovery codes are single-use, high-value. +7. **Verbatim auth-endpoint error bodies, CONFIRMED (minor nuance).** `elevate_session`/ + `create_pat` both raise `RuntimeError` including the raw `response.text` (`create_pat` also + `{payload=}`, which is `{name, expiresIn, scope}` — the MFA code itself is not in either + logged payload). Still real: no redaction, contradicting this RFC's general redaction stance. +8. **"Ask-first" is prompt-text, not access control — CONFIRMED, but narrower than it first + appears.** Re-verified exactly where this matters: the ask-first gate + (`MultiProjectMiddleware.on_call_tool`) only ever fires because `_autolease_default_scope` + (gated on a *local* programmatic session) auto-leases an unconfirmed, all-projects + `SessionScope` by default. OAuth and Kai sessions never do this — they simply have **no** scope + at all (not an auto-leased one) until `set_project_scope` runs, so neither grants usable + all-project access before an explicit choice; only the local `login`/env-var-token path has + this gap, and it's closed structurally rather than by better wording — see §Required behavior + below. +9. **Cross-process credential race, CONFIRMED.** No `asyncio`/`fcntl`/lock import anywhere in + `auth_login.py`; `save_tokens` does an unlocked read-modify-write on + `~/.keboola/mcp/credentials.json` with a rotating refresh token. As-built, `_store_key()` is + `hostname` alone, so two different local MCP client processes for the *same stack* (e.g. Claude + Desktop and a terminal `login`) genuinely share one entry today — this is confirmed as a real + design gap, not just a hypothetical. + +## Required behavior + +- **Config field allowlist for header-derived values.** `Config` gains an explicit + `_HEADER_ELIGIBLE_FIELDS` set (the fields legitimately meant to vary per request: + `storage_api_url`, `storage_token`, `branch_id`, `workspace_schema`, `workspace_id`, + `bearer_token`, `conversation_id`, `project_id`) and a new `replace_by_headers()` method that + only resolves `X-{name}` headers for fields in that set. `apply_request_config` uses it instead + of the unrestricted `replace_by`. Deployment-level fields (`jwt_secret`, `postgres_dsn`, + `session_encryption_key`, `oauth_client_id`/`oauth_client_secret`, `oauth_server_url`, + `mcp_server_url`) become permanently unreachable from any request header. Env-var (`KBC_{name}`) + and CLI-derived resolution is untouched — that input is already operator-trusted. +- **Keboola-domain allowlist for `normalize_storage_api_url`.** Replace the bare `connection.` + prefix check with a regex requiring both the `connection.` label and a genuine + `*.keboola.(com|dev)` suffix, mirroring the pattern `oauth.py`'s `_ALLOWED_DOMAINS` already uses + for redirect URIs. Applies uniformly to deployed (already double-covered by `is_same_stack`) + and local (previously uncovered) servers alike. +- **`read_only` is enforced locally for the active project too**, not just relying on the remote + scoped token: `create_session_state` now receives `readonly=(True if scope and scope.read_only + else None)` from `on_request`, so the base session client is built read-only whenever the + confirmed scope requests it — success or failure of the token exchange. Workspace provisioning + (a server-side plumbing GET+POST pair, not a user-visible mutation) is explicitly exempted via a + new `KeboolaClient.writable_storage_client`, so `query_data` keeps working against a read-only + scope that has no workspace yet. The `MultiProjectMiddleware` active-project shortcuts (read + fan-out and the write-dispatch path) are guarded with a `KeboolaClient.readonly` check so they + only skip the per-project client swap when the base client already matches the scope's + `read_only` — defense in depth, zero added cost for the common case once the above makes that + the normal state. `set_project_scope`'s exchange-failure fallback keeps working (some stacks + lack the exchange endpoint) but its `llm_instruction` now says explicitly whether read-only is + server-enforced (a real `scoped_token` exists) or only locally enforced (fallback path). + `KeboolaClient.with_branch_id()` — which rebuilds a fresh client for any non-default-branch + call (routine on a dev branch, not just adversarial) — is fixed to forward `readonly` into the + new client; a fresh `security-scanner` pass on the implementation caught this dropping + `readonly` silently, which would have reopened this exact fail-open bug on every branch switch. +- **`scope_token`'s payload is encrypted, not just signed.** `SessionScope.to_token`/`from_token` + move from `jwt_utils`'s JWS to AES-GCM authenticated encryption via the already-existing + `session_store/crypto.py` helpers and `resolve_encryption_key` — the same key OAuth sessions + already encrypt with. `resolve_scope_secret`/`_FALLBACK_SCOPE_SECRET` are removed in favour of + `resolve_scope_key`. A new `scope_token` is therefore ciphertext, not a + base64+gzip-recoverable signed blob; the live `scoped_token` it may carry is no longer readable + without the key. No backward-compatible legacy-JWS decode path: this feature has not shipped to + production (main has none of PSGO-261 yet), so there are no live tokens to migrate — a clean + replacement, not a staged one. (A separate design considered and rejected: a new + Postgres-backed `scope_sessions` table mirroring `kai_scope.py`, giving every client an opaque + handle instead of any client-held credential. Rejected as unwarranted complexity — + `scope_token` is only actually issued in the narrow + remaining case where neither OAuth nor Kai's Postgres-backed persistence applies; OAuth and Kai + sessions already never hand the client a live credential at all.) +- **MFA codes: prompt, don't require a CLI argument.** `login --pat` still accepts + `--totp`/`--recovery` as opt-in overrides for scripted/CI use (documented in `--help` as + shell-history/`ps`-visible), but when neither is supplied, prompts via `getpass.getpass()` — + hidden input on a real TTY, and a graceful (though visible, with a stderr warning) read from + stdin when piped/non-interactive, so scripted input still works without extra plumbing. +- **Auth-endpoint errors are redacted.** `elevate_session`/`create_pat` raise a generic + `RuntimeError(f'... failed ({status}). See debug logs for details.')`; the raw `response.text`/ + request `payload` move to `LOG.debug(...)` only. +- **Local sessions are scoped at login time, never auto-leased to everything.** This replaces + "document ask-first as guidance" with a structural fix: `login` (and `login --pat`) now require + an explicit project choice — prompted interactively (same "show projects, pick all or a subset" + flow already used in-conversation by `get_accessible_projects`/`set_project_scope`) when run + from a TTY without `--project-ids`/`--all`, required explicitly otherwise. `lease_pat`, which + previously always requested every accessible project, takes the same explicit choice. The + confirmed `project_ids`/`read_only` are persisted alongside the access/refresh tokens in the + stored credential entry, and the local-session bootstrap in `mcp.py` reads them back as an + already-`confirmed=True` `SessionScope` — `_autolease_default_scope`'s implicit + all-projects-then-ask-first default is removed for any session with a persisted choice. Since + OAuth and Kai sessions already never auto-lease (finding #8), this closes the gap at its actual + source (a local session existing before any explicit choice) rather than trying to make an + LLM-facing instruction into an enforcement boundary. +- **Credentials are keyed per interface, not just per stack.** `login` gains a profile identifier + (`--profile ` / `KBC_LOGIN_PROFILE`, defaulting to `'default'` so single-interface setups + are unaffected) naming which calling interface (Claude Desktop, Cursor, a terminal session) this + login is for. `_store_key()` becomes `(hostname, profile)`, and the on-disk schema nests entries + accordingly — removing finding #9's race by construction, since independent interfaces no + longer share an entry at all. The narrower race that remains — two concurrent requests *within + one process* both seeing "near expiry" and both refreshing — is closed with a plain + `asyncio.Lock` per `(hostname, profile)` in `get_access_token` (in-process; no file locking + needed for this case). A non-blocking `fcntl.flock` on a sibling `.lock` file around the on-disk + read-modify-write is kept as cheap defense-in-depth insurance (polling `LOCK_EX | LOCK_NB`, + never a blocking flock — degrades with a warning rather than stalling the event loop/MCP + handshake; the `fcntl` import is guarded for non-POSIX platforms), covering accidental + profile-sharing or a `login` run racing an already-running server for the same profile. + +## Explicitly out of scope this increment + +- Real MCP-elicitation-based (`elicitation/create`) human-in-the-loop confirmation for scoping — + superseded by login-time scoping, which removes the need for any runtime confirmation gate on + the local path. Worth revisiting only if a future flow reintroduces an unconfirmed-by-default + state. +- Windows-native file locking for the credential-lock insurance layer — CI and the documented + supported platforms are POSIX-only today; the `fcntl` import degrades cleanly rather than + crashing where absent. + +## Decisions (increment 7) + +- **Fix the flow, don't just document the gap**, for both #8 (ask-first) and #9 (credential + race): in both cases a structural fix (scope at login time; key credentials per interface) was + available and preferred over accepting the gap as a documented limitation. +- **Eliminate shared state before adding a lock**: #9's primary fix is removing the sharing + (per-profile keying), not the `fcntl.flock` layer, which is retained only as insurance for + whatever narrow sharing remains (in-process concurrency, accidental profile reuse). +- **Encrypt the existing `scope_token` fallback rather than build new server-side infrastructure** + for #2/#4: since OAuth and Kai already keep credentials server-side, the client-held-token case + is narrow enough that AES-GCM-encrypting the existing JWS payload is proportionate; a new + Postgres table mirroring `kai_scope.py` was considered and rejected as unneeded complexity for + that narrow remaining surface. From 3471706c45bdf4be08cd344b2774501b1d4b6f57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Sun, 9 Aug 2026 10:54:58 +0200 Subject: [PATCH 77/89] fix(PSGO-261): allowlist header-derived Config fields, close domain-check bypass Config.replace_by (used to apply per-request HTTP headers) resolved any dataclass field, including jwt_secret -- an X-Jwt-Secret header let a caller choose the HMAC key that signs and verifies their own scope_token, forging arbitrary project_ids. Add Config._HEADER_ELIGIBLE_FIELDS and a dedicated replace_by_headers() that only ever sets fields legitimately meant to vary per request; jwt_secret/postgres_dsn/session_encryption_key/oauth_*/ mcp_server_url become permanently unreachable from a header. normalize_storage_api_url was a bare `hostname.startswith('connection.')` prefix check, so connection.attacker.tld passed. Replace with a regex requiring a genuine connection.*.keboola.(com|dev) suffix, mirroring the domain-allowlist pattern oauth.py already uses for redirect URIs. --- src/keboola_mcp_server/clients/base.py | 12 ++++- src/keboola_mcp_server/config.py | 45 +++++++++++++++- tests/clients/test_base.py | 47 +++++++++++++++++ tests/test_config.py | 72 ++++++++++++++++++++++++++ tests/test_errors.py | 4 +- tests/test_server.py | 18 ++++--- 6 files changed, 185 insertions(+), 13 deletions(-) create mode 100644 tests/clients/test_base.py diff --git a/src/keboola_mcp_server/clients/base.py b/src/keboola_mcp_server/clients/base.py index 8975b3cce..fc0ba5bc4 100644 --- a/src/keboola_mcp_server/clients/base.py +++ b/src/keboola_mcp_server/clients/base.py @@ -1,5 +1,6 @@ import json import logging +import re from http import HTTPStatus from pathlib import Path from typing import Any, Union, cast @@ -18,15 +19,22 @@ LOG = logging.getLogger(__name__) +# A genuine Keboola stack host: a `connection.` label, any number of region/cloud-provider +# subdomain labels, ending in `.keboola.com` or `.keboola.dev`. `hostname.startswith('connection.')` +# alone is not a domain allowlist -- `connection.attacker.tld` would satisfy it -- see the +# "Security hardening" RFC increment. Mirrors the domain-allowlist pattern `oauth.py`'s +# `_ALLOWED_DOMAINS` already uses for redirect URIs, scoped to this server's own kind of host. +_STORAGE_API_HOST_RE = re.compile(r'^connection\.(?:[a-z0-9-]+\.)*keboola\.(?:com|dev)$', re.IGNORECASE) + def normalize_storage_api_url(storage_api_url: str) -> str: """ Validates a Keboola Storage API URL and returns its canonical ``https://connection.`` base. - :raises ValueError: if the host is missing or is not a ``connection.*`` host. + :raises ValueError: if the host is missing or is not a genuine ``connection.*.keboola.(com|dev)`` host. """ parsed = urlparse(storage_api_url) - if not parsed.hostname or not parsed.hostname.startswith('connection.'): + if not parsed.hostname or not _STORAGE_API_HOST_RE.fullmatch(parsed.hostname): raise ValueError(f'Invalid Keboola Storage API URL: {storage_api_url}') return urlunparse(('https', parsed.hostname, '', '', '', '')) diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 7747a9adb..804e2f228 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -7,7 +7,7 @@ import uuid from collections.abc import Mapping from dataclasses import dataclass, field -from typing import Any, Literal +from typing import Any, ClassVar, Literal from urllib.parse import urlparse, urlunparse LOG = logging.getLogger(__name__) @@ -68,6 +68,25 @@ class Config: Only consulted when the inbound Storage token is a Keboola programmatic token; the legacy project-bound Storage token derives its project from the token itself.""" + # Fields a per-request HTTP header may legitimately set (see `replace_by_headers`). Everything + # else -- jwt_secret, postgres_dsn, session_encryption_key, oauth_client_id/secret, + # oauth_server_url, mcp_server_url -- is deployment-level configuration and must only ever come + # from the process environment or CLI args, never a caller-supplied header. Without this + # allowlist, a header literally named (in any of the exact/`KBC_`/`X-` spellings `_read_options` + # accepts) e.g. `Jwt-Secret` would let a caller choose the HMAC key that verifies their own + # `scope_token`, forging arbitrary `project_ids` -- see the "Security hardening" RFC increment. + _HEADER_ELIGIBLE_FIELDS: ClassVar[frozenset[str]] = frozenset( + { + 'storage_api_url', + 'storage_token', + 'branch_id', + 'workspace_schema', + 'bearer_token', + 'conversation_id', + 'project_id', + } + ) + def __post_init__(self) -> None: for f in dataclasses.fields(self): if 'url' not in f.name: @@ -97,10 +116,18 @@ def _normalize(name: str) -> str: return name.lower().replace('_', '').replace('-', '') @classmethod - def _read_options(cls, d: Mapping[str, str]) -> Mapping[str, Any]: + def _read_options(cls, d: Mapping[str, str], *, allowed_fields: frozenset[str] | None = None) -> Mapping[str, Any]: + """:param allowed_fields: When given, only these field names are ever set -- fields + outside it are skipped entirely, under every naming convention (`X-{name}` and `KBC_{name}` + headers included). Used by `replace_by_headers` to keep deployment-level fields + unreachable from a request; `None` (the default, for env/CLI-derived input) leaves every + field reachable, since that input is already operator-trusted. + """ data = {cls._normalize(k): v for k, v in d.items()} options: dict[str, Any] = {} for f in dataclasses.fields(cls): + if allowed_fields is not None and f.name not in allowed_fields: + continue field_names = [f.name] + f.metadata.get('aliases', []) for name in field_names: @@ -142,9 +169,23 @@ def replace_by(self, d: Mapping[str, str]) -> 'Config': Creates new `Config` instance from the existing one by replacing the values from the input mapping. The keys in the input mapping can either be the names of the fields in `Config` class or their uppercase variant prefixed with 'KBC_'. + + For a per-request HTTP request's headers (untrusted caller input), use + `replace_by_headers` instead -- this method leaves every field reachable, which is only + safe for operator-trusted input (the process environment, CLI args). """ return dataclasses.replace(self, **self._read_options(d)) + def replace_by_headers(self, headers: Mapping[str, str]) -> 'Config': + """Like `replace_by`, but only ever sets fields in `_HEADER_ELIGIBLE_FIELDS` -- every + other field (`jwt_secret`, `postgres_dsn`, `session_encryption_key`, `oauth_client_id`/ + `oauth_client_secret`, `oauth_server_url`, `mcp_server_url`) is deployment-level + configuration and must never be settable by a caller-supplied header, under any of the + exact/`KBC_`/`X-` name spellings `_read_options` accepts -- see the "Security hardening" + RFC increment. + """ + return dataclasses.replace(self, **self._read_options(headers, allowed_fields=self._HEADER_ELIGIBLE_FIELDS)) + def __repr__(self) -> str: params: list[str] = [] for f in dataclasses.fields(self): diff --git a/tests/clients/test_base.py b/tests/clients/test_base.py new file mode 100644 index 000000000..7d02b6a11 --- /dev/null +++ b/tests/clients/test_base.py @@ -0,0 +1,47 @@ +import pytest + +from keboola_mcp_server.clients.base import normalize_storage_api_url + + +class TestNormalizeStorageApiUrl: + """`normalize_storage_api_url` requires a genuine Keboola stack domain, not just a + `connection.` prefix -- see the "Security hardening" RFC increment (`connection.attacker.tld` + previously passed, letting a caller-supplied host receive the live bearer token).""" + + @pytest.mark.parametrize( + ('url', 'expected'), + [ + ('https://connection.keboola.com', 'https://connection.keboola.com'), + ('https://connection.eu-central-1.keboola.com', 'https://connection.eu-central-1.keboola.com'), + ( + 'https://connection.north-europe.azure.keboola.com', + 'https://connection.north-europe.azure.keboola.com', + ), + ( + 'https://connection.europe-west3.gcp.keboola.com', + 'https://connection.europe-west3.gcp.keboola.com', + ), + ('https://connection.canary-orion.keboola.dev', 'https://connection.canary-orion.keboola.dev'), + ('https://connection.keboola.com:443', 'https://connection.keboola.com'), + ('https://connection.keboola.com/v2/storage', 'https://connection.keboola.com'), + ], + ) + def test_accepts_genuine_keboola_stack_hosts(self, url: str, expected: str) -> None: + assert normalize_storage_api_url(url) == expected + + @pytest.mark.parametrize( + 'url', + [ + 'https://connection.attacker.tld', + 'https://connection.attacker.example', + 'https://connection.keboola.com.attacker.example', + 'https://connection.keboola.com.example.com', + 'https://connection.example.com', + 'https://sapi.keboola.com', # no 'connection.' label at all + 'https://keboola.com', + '', + ], + ) + def test_rejects_lookalike_or_foreign_hosts(self, url: str) -> None: + with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): + normalize_storage_api_url(url) diff --git a/tests/test_config.py b/tests/test_config.py index 0ace9bcde..b1a12beaa 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -131,6 +131,78 @@ def test_url_field(self, url: str, expected: str) -> None: assert config.mcp_server_url == expected +class TestReplaceByHeaders: + """Deployment-level fields must never be settable by a per-request header, under any of the + exact/`KBC_`/`X-` name spellings `_read_options` accepts -- see the "Security hardening" RFC + increment (a caller-controlled `Jwt-Secret` header would otherwise let them forge their own + `scope_token`).""" + + @pytest.mark.parametrize( + 'headers', + [ + {'Jwt-Secret': 'attacker-chosen'}, + {'X-Jwt-Secret': 'attacker-chosen'}, + {'KBC-Jwt-Secret': 'attacker-chosen'}, + {'X-Postgres-Dsn': 'postgresql://evil'}, + {'X-Session-Encryption-Key': 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA='}, + {'X-Oauth-Client-Id': 'evil'}, + {'X-Oauth-Client-Secret': 'evil'}, + {'X-Oauth-Server-Url': 'https://evil.example'}, + {'X-Mcp-Server-Url': 'https://evil.example'}, + ], + ids=[ + 'jwt_secret_bare', + 'jwt_secret_x_prefixed', + 'jwt_secret_kbc_prefixed', + 'postgres_dsn', + 'session_encryption_key', + 'oauth_client_id', + 'oauth_client_secret', + 'oauth_server_url', + 'mcp_server_url', + ], + ) + def test_deployment_level_fields_are_unreachable(self, headers: Mapping[str, str]) -> None: + config = Config(jwt_secret='real-secret', postgres_dsn='postgresql://real') + out = config.replace_by_headers(headers) + assert out == config # nothing changed -- every one of these headers was ignored + + def test_allowlisted_fields_still_work(self) -> None: + config = Config() + out = config.replace_by_headers( + { + 'X-Storage-Api-Url': 'https://connection.keboola.com', + 'X-Branch-Id': '123', + 'X-Conversation-Id': 'conv-1', + } + ) + assert out.storage_api_url == 'https://connection.keboola.com' + assert out.branch_id == '123' + assert out.conversation_id == 'conv-1' + + def test_replace_by_is_unrestricted_for_trusted_input(self) -> None: + # replace_by (env/CLI, operator-trusted) is deliberately NOT subject to the same + # allowlist -- only replace_by_headers (untrusted per-request input) is restricted. + config = Config() + out = config.replace_by({'jwt_secret': 'ops-configured'}) + assert out.jwt_secret == 'ops-configured' + + +class TestServerRuntimeInfoSessionStatePersists: + def test_stdio_always_persists_regardless_of_stateless_http(self) -> None: + # stdio is one process/one session for the whole conversation -- the flag is meaningless there. + assert ServerRuntimeInfo(transport='stdio', stateless_http=True).session_state_persists is True + assert ServerRuntimeInfo(transport='stdio', stateless_http=False).session_state_persists is True + + def test_streamable_http_follows_stateless_http_flag(self) -> None: + assert ServerRuntimeInfo(transport='streamable-http', stateless_http=True).session_state_persists is False + assert ServerRuntimeInfo(transport='streamable-http', stateless_http=False).session_state_persists is True + + def test_defaults_to_stateless(self) -> None: + # Matches the CLI's --stateless-http default (scaled/deployed-safe). + assert ServerRuntimeInfo(transport='streamable-http').session_state_persists is False + + class TestEnvStorageApiUrl: @pytest.mark.parametrize( ('env', 'expected'), diff --git a/tests/test_errors.py b/tests/test_errors.py index 5c633b321..2a7c1104a 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -471,10 +471,8 @@ async def foo(_ctx: Context): [ ('https://connection.keboola.com', True), ('https://connection.north-europe.azure.keboola.com', False), - ('https://connection.keboola.com.attacker.example', False), - ('https://connection.attacker.example', False), ], - ids=['own_stack', 'other_stack', 'lookalike_suffix', 'foreign_domain'], + ids=['own_stack', 'other_stack'], ) async def test_event_step_up_header_only_for_own_stack( tmp_path, monkeypatch, mocker, empty_context: Context, session_storage_api_url: str, expect_step_up_header: bool diff --git a/tests/test_server.py b/tests/test_server.py index 057230242..4dcfacadc 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -211,7 +211,9 @@ async def test_own_stack_from_cli_parameter_only(tmp_path, monkeypatch): [ ( # config params in Config class Config( - storage_token='SAPI_1234', storage_api_url='http://connection.sapi', workspace_schema='WORKSPACE_1234' + storage_token='SAPI_1234', + storage_api_url='http://connection.test.keboola.com', + workspace_schema='WORKSPACE_1234', ), {}, ), @@ -219,16 +221,20 @@ async def test_own_stack_from_cli_parameter_only(tmp_path, monkeypatch): Config(), { 'KBC_STORAGE_TOKEN': 'SAPI_1234', - 'KBC_STORAGE_API_URL': 'http://connection.sapi', + 'KBC_STORAGE_API_URL': 'http://connection.test.keboola.com', 'KBC_WORKSPACE_SCHEMA': 'WORKSPACE_1234', }, ), ( # config params mixed up in both the Config class and the OS environment - Config(storage_api_url='http://connection.sapi'), + Config(storage_api_url='http://connection.test.keboola.com'), {'KBC_STORAGE_TOKEN': 'SAPI_1234', 'KBC_WORKSPACE_SCHEMA': 'WORKSPACE_1234'}, ), ( # the OS environment overrides the initial Config class - Config(storage_token='foo-bar', storage_api_url='http://connection.sapi', workspace_schema='xyz_123'), + Config( + storage_token='foo-bar', + storage_api_url='http://connection.test.keboola.com', + workspace_schema='xyz_123', + ), {'KBC_STORAGE_TOKEN': 'SAPI_1234', 'KBC_WORKSPACE_SCHEMA': 'WORKSPACE_1234'}, ), # TODO: Also test values obtained from an HTTP request. @@ -300,7 +306,7 @@ async def test_with_session_state_admin_role_tools(mocker, admin_info, expected_ os_mock = mocker.patch('keboola_mcp_server.server.os') os_mock.environ = { 'KBC_STORAGE_TOKEN': 'SAPI_1234', - 'KBC_STORAGE_API_URL': 'http://connection.sapi', + 'KBC_STORAGE_API_URL': 'http://connection.test.keboola.com', 'KBC_WORKSPACE_SCHEMA': 'WORKSPACE_1234', } @@ -507,7 +513,7 @@ async def test_json_logging(): '--transport', 'streamable-http', '--api-url', - 'http://connection.nowhere', + 'http://connection.test.keboola.com', '--storage-token', 'foo', '--log-config', From b0ed1ae21651812c08c026a86a1c95961dcfa4b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Sun, 9 Aug 2026 10:55:09 +0200 Subject: [PATCH 78/89] fix(PSGO-261): stop read_only from failing open for the active project MultiProjectMiddleware's active-project shortcuts skipped the per-project client swap -- the only place readonly=scope.read_only reached a KeboolaClient -- whenever a call targeted scope.active_project_id, true for every single-project scope and the first project of any multi-project scope. create_session_state never passed readonly= at all, so the active project's writes were never locally restricted regardless of a confirmed read-only scope. Thread readonly=(True if scope and scope.read_only else None) into create_session_state. Add KeboolaClient.writable_storage_client (workspace provisioning is server-side plumbing, not a user-visible mutation, so it must keep working under a read-only scope) and KeboolaClient.readonly; guard the MultiProjectMiddleware shortcuts with _active_client_honors_scope so they only skip the swap once the base client already matches the scope. Also fixes a regression this change would otherwise have reopened: KeboolaClient.with_branch_id() rebuilds a fresh client for any non-default branch (routine dev-branch usage) but dropped readonly when doing so, so a read-only session became fully writable again on a plain branch switch. Forward readonly=self.readonly in both non-self-returning branches. --- src/keboola_mcp_server/clients/client.py | 38 +++++++++-- src/keboola_mcp_server/multiproject.py | 17 ++++- src/keboola_mcp_server/workspace.py | 12 ++-- tests/clients/test_client.py | 85 ++++++++++++++++++++---- tests/test_multiproject.py | 65 ++++++++++++++++++ tests/test_workspace.py | 13 ++-- 6 files changed, 200 insertions(+), 30 deletions(-) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index 337c8837f..fd4252e9a 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -100,6 +100,7 @@ async def with_branch_id(self, branch_id: str | None) -> 'KeboolaClient': bearer_token=self._bearer_token, branch_id=None, headers=self._headers, + readonly=self.readonly, own_stack_storage_api_url=self._own_stack_storage_api_url, ) else: @@ -124,6 +125,7 @@ async def with_branch_id(self, branch_id: str | None) -> 'KeboolaClient': bearer_token=self._bearer_token, branch_id=normalized_branch_id, headers=self._headers, + readonly=self.readonly, own_stack_storage_api_url=self._own_stack_storage_api_url, ) @@ -280,6 +282,29 @@ def headers(self) -> dict[str, Any] | None: def storage_client(self) -> 'AsyncStorageClient': return self._storage_client + @property + def readonly(self) -> bool | None: + return self._storage_client.raw_client.readonly + + @property + def writable_storage_client(self) -> 'AsyncStorageClient': + """A Storage client identical to `storage_client` but never read-only. + + Used for server-side plumbing (workspace/config provisioning ahead of `query_data`) that + must succeed even under a read-only confirmed scope: the read-only guarantee is about + which tools the caller can use to mutate the project's own data, not whether the server + may provision the read-only workspace it needs to serve reads at all -- see the + "Security hardening" RFC increment. + """ + return AsyncStorageClient.create( + root_url=self._storage_api_url, + token=self._bearer_or_sapi_token, + branch_id=self._branch_id, + headers=self._headers, + readonly=None, + encryption_client=self._encryption_client, + ) + def step_up_storage_client(self, kubernetes_token_path: str) -> 'AsyncStorageClient': """ Returns a Storage client that keeps this client's user token and additionally @@ -287,8 +312,10 @@ def step_up_storage_client(self, kubernetes_token_path: str) -> 'AsyncStorageCli step-up header. Connection waives the permissions the user's token lacks on the step-up-enabled actions (workspace / config / event creation) when the ServiceAccount is authorized for them — no privileged token is minted and the - user's token stays the audited principal. The read-only write guard of the user's - Storage client is preserved; the header only widens server-side permissions. + user's token stays the audited principal. Always writable regardless of this client's + own read-only setting (see `writable_storage_client`) -- provisioning is server-side + plumbing, not a user-visible mutation, and step-up exists precisely to let it proceed on + a token that otherwise couldn't. The ServiceAccount JWT is a credential of the MCP server deployment itself, so it is only ever sent to the Keboola stack that this server belongs to. The Storage API URL of a @@ -296,7 +323,8 @@ def step_up_storage_client(self, kubernetes_token_path: str) -> 'AsyncStorageCli this server's own stack — resolved once when the server starts and passed to this client as `own_stack_storage_api_url` — before the header is attached. When the two differ, or when the server has no stack of its own (a locally run server), the step-up is skipped and this - client's plain Storage client is returned, so the JWT is never sent anywhere else. + client's plain (but still writable) Storage client is returned, so the JWT is never sent + anywhere else. The token file is read on each call so kubelet rotation needs no restart. @@ -309,7 +337,7 @@ def step_up_storage_client(self, kubernetes_token_path: str) -> 'AsyncStorageCli f"it is not the Storage API URL of this server's own stack " f'({self._own_stack_storage_api_url or "not configured"}).' ) - return self._storage_client + return self.writable_storage_client jwt = read_service_account_jwt(kubernetes_token_path) @@ -323,7 +351,7 @@ def step_up_storage_client(self, kubernetes_token_path: str) -> 'AsyncStorageCli token=self._bearer_or_sapi_token, branch_id=self._branch_id, headers=headers, - readonly=self._storage_client.raw_client.readonly, + readonly=None, ) @property diff --git a/src/keboola_mcp_server/multiproject.py b/src/keboola_mcp_server/multiproject.py index bfd9dceb8..ec6479b10 100644 --- a/src/keboola_mcp_server/multiproject.py +++ b/src/keboola_mcp_server/multiproject.py @@ -41,6 +41,19 @@ _PROJECT_FILTER_ARG = 'project_ids' +def _active_client_honors_scope(state: dict[str, Any], scope: SessionScope) -> bool: + """True when the base session client already matches ``scope.read_only`` -- the active- + project shortcuts below skip the per-project client swap only in that case (defense in + depth: `SessionStateMiddleware.create_session_state` already builds the base client + read-only whenever the scope requests it, so this is normally true and the shortcut's cost + stays zero; see the "Security hardening" RFC increment). + """ + if not scope.read_only: + return True + client = state.get(KeboolaClient.STATE_KEY) + return isinstance(client, KeboolaClient) and client.readonly is True + + class MultiProjectMiddleware(fmw.Middleware): """Fans a read-only tool call out across every project in the active multi-project scope. @@ -132,7 +145,7 @@ async def on_call_tool( # project only — one call, that project's X-KBC-ProjectId, no per-project envelope. if len(targets) == 1: target = targets[0] - if target == scope.active_project_id: + if target == scope.active_project_id and _active_client_honors_scope(state, scope): return await call_next(context) try: await self._swap_project(state, server_state, storage_api_url, base_token, target, scope.read_only) @@ -213,7 +226,7 @@ async def _dispatch_single_target( project_id = args.get(PROJECT_ID_ARG) if isinstance(args, dict) else None target = self._resolve_single_target(scope, project_id) - if target is None or target == scope.active_project_id: + if target is None or (target == scope.active_project_id and _active_client_honors_scope(state, scope)): return await call_next(context) server_state = ServerState.from_context(ctx) diff --git a/src/keboola_mcp_server/workspace.py b/src/keboola_mcp_server/workspace.py index bc890273d..6be5daf76 100644 --- a/src/keboola_mcp_server/workspace.py +++ b/src/keboola_mcp_server/workspace.py @@ -630,10 +630,12 @@ async def _provisioning_storage_client(self) -> AsyncStorageClient: step-up header — Connection waives permissions the user's token lacks when the ServiceAccount is authorized for workspace provisioning. No privileged token is ever minted; the audit trail stays on the user's token. - Otherwise the user's own Storage client is used unchanged. The SA JWT is attached - only when this manager's client talks to the server's own stack; - `KeboolaClient.step_up_storage_client()` falls back to the user's own client - otherwise. + Otherwise the user's own client is used, but always writable + (`KeboolaClient.writable_storage_client`) even under a read-only confirmed scope -- + provisioning is server-side plumbing, not a user-visible mutation, so it must succeed + even when the session itself can't write. The SA JWT is attached only when this + manager's client talks to the server's own stack; `KeboolaClient.step_up_storage_client()` + falls back to the plain (still writable) client otherwise. The step-up client is cached for this manager's lifetime, so the token file is read once — when the client is first built — not on every provisioning attempt. @@ -642,7 +644,7 @@ async def _provisioning_storage_client(self) -> AsyncStorageClient: rotation is picked up without restarting the server. """ if not self._kubernetes_token_path: - return self._client.storage_client + return self._client.writable_storage_client if self._provisioning_client is None: self._provisioning_client = self._client.step_up_storage_client(self._kubernetes_token_path) LOG.debug('Workspace provisioning storage client created.') diff --git a/tests/clients/test_client.py b/tests/clients/test_client.py index 9985d5c92..388126f96 100644 --- a/tests/clients/test_client.py +++ b/tests/clients/test_client.py @@ -17,7 +17,7 @@ @pytest.fixture def keboola_client() -> KeboolaClient: - return KeboolaClient(storage_api_url='https://connection.nowhere', storage_api_token='test-token') + return KeboolaClient(storage_api_url='https://connection.test.keboola.com', storage_api_token='test-token') @pytest.fixture @@ -257,7 +257,7 @@ async def test_trigger_event( if value } mock_client.post.assert_called_once_with( - 'https://connection.nowhere/v2/storage/events', + 'https://connection.test.keboola.com/v2/storage/events', params=None, headers={ 'Content-Type': 'application/json', @@ -327,7 +327,7 @@ async def test_token_create( # Verify the API call was made with correct parameters mock_client.post.assert_called_once_with( - 'https://connection.nowhere/v2/storage/tokens', + 'https://connection.test.keboola.com/v2/storage/tokens', params=None, headers={ 'Content-Type': 'application/json', @@ -347,7 +347,7 @@ def runtime_config(self) -> ServerRuntimeInfo: def keboola_client_with_headers(self, runtime_config: ServerRuntimeInfo) -> KeboolaClient: headers = build_tracing_headers(runtime_config) return KeboolaClient( - storage_api_url='https://connection.nowhere', storage_api_token='test-token', headers=headers + storage_api_url='https://connection.test.keboola.com', storage_api_token='test-token', headers=headers ) @pytest.mark.asyncio @@ -363,7 +363,7 @@ async def test_keboola_client_passing_headers(self, keboola_client_with_headers: mcp_version = importlib.metadata.version('mcp') fastmcp_version = importlib.metadata.version('fastmcp') mock_client.get.assert_called_once_with( - 'https://connection.nowhere/v2/storage/tokens/verify', + 'https://connection.test.keboola.com/v2/storage/tokens/verify', params=None, headers={ 'Content-Type': 'application/json', @@ -765,7 +765,10 @@ def test_uses_bearer_for_programmatic_token(self, tmp_path): assert 'X-StorageAPI-Token' not in headers @pytest.mark.parametrize('readonly', [None, True, False]) - def test_propagates_readonly_guard(self, tmp_path, readonly): + def test_always_writable_regardless_of_client_readonly(self, tmp_path, readonly): + # Provisioning (what step-up exists for) is server-side plumbing, not a user-visible + # mutation -- it must succeed even under a read-only confirmed scope. See the "Security + # hardening" RFC increment. token_file = tmp_path / 'token' token_file.write_text('sa-jwt') client = KeboolaClient( @@ -777,7 +780,7 @@ def test_propagates_readonly_guard(self, tmp_path, readonly): stepped = client.step_up_storage_client(str(token_file)) - assert stepped.raw_client.readonly == client.storage_client.raw_client.readonly + assert stepped.raw_client.readonly is None @pytest.mark.asyncio @pytest.mark.parametrize('branch_id', [None, '123'], ids=['main_branch', 'dev_branch']) @@ -801,6 +804,27 @@ async def test_own_stack_survives_branch_switch(self, tmp_path, mocker, branch_i == 'Bearer sa-jwt' ) + @pytest.mark.asyncio + @pytest.mark.parametrize('branch_id', [None, '123'], ids=['main_branch', 'dev_branch']) + @pytest.mark.parametrize('readonly', [None, True, False]) + async def test_with_branch_id_preserves_readonly(self, mocker, branch_id, readonly): + # Regression: with_branch_id() constructs a brand-new KeboolaClient for a non-default + # branch_id, and previously dropped `readonly` in doing so -- silently making a read-only + # confirmed scope's client writable again on the routine (non-adversarial) act of switching + # to a dev branch. See the "Security hardening" RFC increment. + client = KeboolaClient( + storage_api_url='https://connection.keboola.com', + storage_api_token='user-token', + branch_id='999', + readonly=readonly, + ) + client.storage_client.dev_branch_detail = mocker.AsyncMock(return_value={'isDefault': False}) + + branched = await client.with_branch_id(branch_id) + + assert branched is not client + assert branched.readonly == readonly + def test_fails_loudly_on_empty_token_file(self, tmp_path): token_file = tmp_path / 'token' token_file.write_text(' \n') @@ -826,18 +850,14 @@ def test_fails_loudly_on_missing_token_file(self, tmp_path): @pytest.mark.parametrize( ('storage_api_url', 'own_stack_storage_api_url'), [ - # Another Keboola stack ... + # Another (genuine) Keboola stack ... ('https://connection.north-europe.azure.keboola.com', OWN_STACK_URL), - # ... and hosts that only look like this server's stack. All of them satisfy the - # 'connection.' prefix that the Storage API URL itself is required to have. - ('https://connection.keboola.com.attacker.example', OWN_STACK_URL), - ('https://connection.attacker.example', OWN_STACK_URL), # A genuinely different port is a different endpoint. (OWN_STACK_URL, f'{OWN_STACK_URL}:8443'), # A server with no stack of its own (locally run) has no stack to step up on. (OWN_STACK_URL, None), ], - ids=['other_stack', 'lookalike_suffix', 'foreign_domain', 'other_port', 'no_own_stack'], + ids=['other_stack', 'other_port', 'no_own_stack'], ) def test_no_step_up_header_for_foreign_stack(self, tmp_path, storage_api_url, own_stack_storage_api_url): """The ServiceAccount JWT belongs to this server's stack and must not travel anywhere else.""" @@ -854,4 +874,41 @@ def test_no_step_up_header_for_foreign_stack(self, tmp_path, storage_api_url, ow stepped = client.step_up_storage_client(str(token_file)) assert 'X-Kubernetes-Authorization' not in (stepped.raw_client.headers or {}) - assert stepped is client.storage_client + # Falls back to the plain client, but still writable -- see `writable_storage_client`. + assert stepped.raw_client.readonly is None + + @pytest.mark.parametrize( + 'lookalike_url', + [ + # Satisfies the old 'connection.' prefix check but not a genuine keboola.com/dev + # suffix -- normalize_storage_api_url (Security hardening RFC increment) now rejects + # these outright, so they can never even become a session's storage_api_url, let + # alone reach the step-up destination check. + 'https://connection.keboola.com.attacker.example', + 'https://connection.attacker.example', + ], + ) + def test_lookalike_domain_rejected_before_construction(self, lookalike_url) -> None: + with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): + KeboolaClient( + storage_api_url=lookalike_url, + storage_api_token='user-token', + own_stack_storage_api_url=self.OWN_STACK_URL, + ) + + def test_foreign_stack_fallback_is_writable_even_under_a_readonly_client(self, tmp_path): + # The exemption must hold even when the caller's own client is genuinely read-only -- + # provisioning is server-side plumbing, not a user-visible mutation. + token_file = tmp_path / 'token' + token_file.write_text('sa-jwt') + client = KeboolaClient( + storage_api_url='https://connection.other.keboola.com', + storage_api_token='user-token', + readonly=True, + own_stack_storage_api_url=self.OWN_STACK_URL, + ) + assert client.readonly is True # sanity: the client itself really is read-only + + stepped = client.step_up_storage_client(str(token_file)) + + assert stepped.raw_client.readonly is None diff --git a/tests/test_multiproject.py b/tests/test_multiproject.py index eb522a3a4..f54048b61 100644 --- a/tests/test_multiproject.py +++ b/tests/test_multiproject.py @@ -608,3 +608,68 @@ def test_merge_large_degrades_to_count_first(self, monkeypatch) -> None: assert 'project_ids' in note assert len(merged.structured_content['buckets']) == 3 # truncated to the cap assert merged.structured_content['total'] == 5 # true total preserved + + +class TestActiveProjectReadOnlyGuard: + """The active-project shortcut only skips the per-project client swap when the base client + already honors the scope's read_only -- defense in depth for the fail-open case where + SessionStateMiddleware couldn't build the base client read-only (Security hardening RFC + increment).""" + + @staticmethod + def _ctx_with_client(scope: SessionScope, tool_name: str, read_only_tool: bool, client_readonly) -> tuple: + client = MagicMock(spec=KeboolaClient) + client.readonly = client_readonly + client.token = 'kbc_at_x' + client.storage_api_url = 'https://connection.keboola.com' + state: dict = {KeboolaClient.STATE_KEY: client, SCOPE_KEY: scope} + ctx = MagicMock(spec=Context) + ctx.session = SimpleNamespace(state=state) + ctx.request_context.lifespan_context = ServerState( + config=Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x'), + runtime_info=ServerRuntimeInfo(transport='stdio'), + ) + tool = MagicMock() + tool.name = tool_name + tool.annotations.readOnlyHint = read_only_tool if read_only_tool else None + ctx.fastmcp.get_tool = AsyncMock(return_value=tool) + message = SimpleNamespace(name=tool_name, arguments={}) + context = SimpleNamespace(message=message, fastmcp_context=ctx) + return context, state, client + + @pytest.mark.asyncio + async def test_skips_swap_when_base_client_already_readonly(self, mocker) -> None: + scope = SessionScope(project_ids=[11], read_only=True, confirmed=True) + context, _, _ = self._ctx_with_client(scope, 'get_tables', read_only_tool=True, client_readonly=True) + swap = mocker.patch.object(MultiProjectMiddleware, '_swap_project', new=AsyncMock()) + + async def call_next(_): + return 'ok' + + await MultiProjectMiddleware().on_call_tool(context, call_next) + swap.assert_not_called() + + @pytest.mark.asyncio + async def test_swaps_when_base_client_is_not_readonly_despite_readonly_scope(self, mocker) -> None: + # The fail-open case: the base client couldn't be built read-only (e.g. an older session + # predating the fix), so the active-project shortcut must not trust it -- fall through to + # a real per-project swap, which enforces read_only itself. + scope = SessionScope(project_ids=[11], read_only=True, confirmed=True) + context, state, _ = self._ctx_with_client(scope, 'get_tables', read_only_tool=True, client_readonly=None) + + async def fake_swap(state_, server_state, storage_api_url, base_token, project_id, read_only): + new_client = MagicMock(spec=KeboolaClient) + new_client.readonly = read_only or None + state_[KeboolaClient.STATE_KEY] = new_client + + mocker.patch.object(MultiProjectMiddleware, '_swap_project', new=AsyncMock(side_effect=fake_swap)) + mocker.patch.object(WorkspaceManager, 'create', new=AsyncMock(return_value='wsm')) + + captured_readonly = [] + + async def call_next(_): + captured_readonly.append(state[KeboolaClient.STATE_KEY].readonly) + return 'ok' + + await MultiProjectMiddleware().on_call_tool(context, call_next) + assert captured_readonly == [True] diff --git a/tests/test_workspace.py b/tests/test_workspace.py index 91ec1a5a5..c580f5141 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -146,6 +146,7 @@ async def test_workspace_creation_cleans_up_config_on_failure(): mock_client.branch_id = None mock_storage_client = AsyncMock() mock_client.storage_client = mock_storage_client + mock_client.writable_storage_client = mock_storage_client mock_storage_client.verify_token.return_value = {'owner': {'defaultBackend': 'snowflake'}} mock_storage_client.configuration_create.return_value = {'id': 'test-config-123', 'name': 'test'} @@ -189,6 +190,7 @@ async def test_workspace_creation_stops_on_terminal_error_status(job_detail: dic mock_client.branch_id = None mock_storage_client = AsyncMock() mock_client.storage_client = mock_storage_client + mock_client.writable_storage_client = mock_storage_client mock_storage_client.verify_token.return_value = {'owner': {'defaultBackend': 'snowflake'}} mock_storage_client.configuration_create.return_value = {'id': 'cfg-1', 'name': 'test'} @@ -210,6 +212,7 @@ async def test_workspace_creation_warning_with_id_uses_workspace(mocker): mock_client.branch_id = None mock_storage_client = AsyncMock() mock_client.storage_client = mock_storage_client + mock_client.writable_storage_client = mock_storage_client mock_storage_client.verify_token.return_value = {'owner': {'defaultBackend': 'snowflake'}} mock_storage_client.configuration_create.return_value = {'id': 'cfg-1', 'name': 'test'} @@ -470,14 +473,16 @@ async def test_workspace_creation_uses_step_up_client(tmp_path): @pytest.mark.asyncio async def test_provisioning_client_falls_back_to_user_client(): - """Without a Kubernetes token path the provisioning client is the user's own Storage client.""" + """Without a Kubernetes token path the provisioning client is the user's own client, but + always writable (see `KeboolaClient.writable_storage_client`) -- provisioning is server-side + plumbing, not a user-visible mutation, so it must succeed even under a read-only scope.""" mock_client = Mock(spec=KeboolaClient) - mock_storage_client = AsyncMock() - mock_client.storage_client = mock_storage_client + mock_writable_client = AsyncMock() + mock_client.writable_storage_client = mock_writable_client manager = WorkspaceManager(mock_client) - assert await manager._provisioning_storage_client() is mock_storage_client + assert await manager._provisioning_storage_client() is mock_writable_client mock_client.step_up_storage_client.assert_not_called() From de372da39ccf0a9250646275898c45638613a43b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Sun, 9 Aug 2026 10:55:21 +0200 Subject: [PATCH 79/89] fix(PSGO-261): encrypt scope_token instead of just signing it, wire login-time scoping SessionScope.to_token/from_token used jwt_utils's JWS (signature only) -- the gzip+JSON payload, which may embed a live scoped_token bearer credential, was base64+gunzip-recoverable without the secret. Switch to AES-GCM authenticated encryption via the existing session_store/crypto.py helpers and resolve_encryption_key (the same key OAuth sessions already encrypt with). resolve_scope_secret/_FALLBACK_SCOPE_SECRET are replaced by resolve_scope_key. Also wires SessionStateMiddleware to read a login-time-persisted project scope (auth_login.TokenSet.project_ids/read_only) before falling back to the old auto-lease-to-everything default, and makes set_project_scope's llm_instruction explicit about whether read_only is server-enforced (a real scoped_token exists) or only locally enforced (exchange-failure fallback). --- src/keboola_mcp_server/mcp.py | 38 ++++++- src/keboola_mcp_server/scope.py | 55 ++++++---- src/keboola_mcp_server/tools/project.py | 20 +++- tests/test_mcp.py | 139 ++++++++++++++++++++---- tests/tools/test_project.py | 37 ++++++- 5 files changed, 239 insertions(+), 50 deletions(-) diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 8b7095c7b..a1c7f0e19 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -47,7 +47,7 @@ SCOPE_TOKEN_ARG, SessionScope, persist_scope, - resolve_scope_secret, + resolve_scope_key, ) from keboola_mcp_server.session_store.kai_scope import KaiScopeStore from keboola_mcp_server.session_store.repository import SessionStore @@ -299,6 +299,13 @@ async def on_request( and server_state.kai_scope_store is not None ): scope = await self._read_persisted_kai_scope(config, server_state.kai_scope_store) + # Local sessions are scoped at `login` time now (see the "Security hardening" RFC + # increment) -- a persisted choice, once one exists, is used as a confirmed scope with + # no ask-first gate needed. Only a credential predating this choice (or a token + # supplied directly, never run through `login`) falls through to the old + # auto-lease-then-ask-first default below. + if scope is None and not config.project_id and not is_list: + scope = await self._read_persisted_login_scope(config) if scope is None and not config.project_id and not is_list: scope = await self._autolease_default_scope(config) if not is_list: @@ -328,8 +335,13 @@ async def on_request( LOG.info(f'Skipping branch validation for {context.method} request.') config = dataclasses.replace(config, branch_id=None) + # A read-only confirmed scope is enforced locally too (not just by the remote scoped + # token, which may not exist -- see set_project_scope's exchange-failure fallback and + # the "Security hardening" RFC increment): the base session client itself is built + # read-only whenever the scope requests it, success or failure of the token exchange. + readonly = True if scope is not None and scope.read_only else None state = await self.create_session_state( - config, runtime_info, own_stack_storage_api_url=own_stack_storage_api_url + config, runtime_info, readonly=readonly, own_stack_storage_api_url=own_stack_storage_api_url ) if scope is not None: state[SCOPE_KEY] = scope @@ -410,7 +422,11 @@ def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_sto :return: The configuration to use for this request. """ LOG.debug(f'Injecting headers: http_rq={http_rq}, headers={http_rq.headers}') - config = config.replace_by(http_rq.headers) + # Only fields meant to vary per request are settable from a header -- see + # Config._HEADER_ELIGIBLE_FIELDS / the "Security hardening" RFC increment. In particular + # this keeps `jwt_secret` (which would otherwise let a caller forge their own scope_token) + # and the other deployment-level fields permanently unreachable from a request. + config = config.replace_by_headers(http_rq.headers) if own_stack_storage_api_url and not is_same_stack(config.storage_api_url, own_stack_storage_api_url): LOG.warning( @@ -456,7 +472,7 @@ def _read_scope_from_request(cls, context: fmw.MiddlewareContext[Any], config: C if not token: return None try: - return SessionScope.from_token(token, resolve_scope_secret(config)) + return SessionScope.from_token(token, resolve_scope_key(config)) except Exception: LOG.warning('Ignoring invalid or expired scope_token.', exc_info=True) return None @@ -585,6 +601,20 @@ async def _maybe_use_stored_session(cls, config: Config, *, refresh: bool = True access_token = tokens.access_token return dataclasses.replace(config, storage_token=access_token) + @classmethod + async def _read_persisted_login_scope(cls, config: Config) -> 'SessionScope | None': + """The project scope chosen at `login` time (see `auth_login.TokenSet.project_ids`), if + any -- "Security hardening" RFC increment. Returns None (falls back to the old + auto-lease-all default) when this isn't a local programmatic session, or the stored + credential predates this choice / was never run through `login`'s prompt. + """ + if not cls._is_local_programmatic(config): + return None + tokens = load_tokens(config.storage_api_url) + if tokens is None or tokens.project_ids is None: + return None + return SessionScope(project_ids=tokens.project_ids, read_only=tokens.read_only, confirmed=True) + @classmethod async def _autolease_default_scope(cls, config: Config) -> 'SessionScope | None': """ diff --git a/src/keboola_mcp_server/scope.py b/src/keboola_mcp_server/scope.py index 8a7d5a6e7..d2f6c3449 100644 --- a/src/keboola_mcp_server/scope.py +++ b/src/keboola_mcp_server/scope.py @@ -1,12 +1,14 @@ """In-conversation multi-project scope (PSGO-261 increment 2): the ``SessionScope`` model, its -``scope_token`` JWT round-trip, and the associated session-state keys. +``scope_token`` round-trip, and the associated session-state keys. Split out of ``mcp.py`` so that module can stay focused on the middleware/server wiring itself (``mcp.py``'s ``SessionStateMiddleware``/``MultiProjectMiddleware`` both depend on this). """ +import base64 import dataclasses -import secrets +import gzip +import json import time from datetime import datetime, timezone from typing import TYPE_CHECKING, Annotated, Optional @@ -14,7 +16,7 @@ from pydantic import Field from keboola_mcp_server.config import Config -from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt +from keboola_mcp_server.session_store.crypto import decrypt, encrypt, resolve_encryption_key if TYPE_CHECKING: from keboola_mcp_server.session_store.repository import SessionStore @@ -50,17 +52,16 @@ # caller resends the token it returned. SCOPE_TOKEN_ARG = 'scope_token' -# Process-local fallback signing key for scope_token, used when no shared KBC_JWT_SECRET is -# configured (e.g. local stdio/login sessions). A per-process secret is enough there since a stdio -# process serves exactly one conversation end-to-end; deployed multi-replica setups already require -# a shared jwt_secret for the OAuth-provider JWTs (see oauth.py), which this reuses. -_FALLBACK_SCOPE_SECRET = secrets.token_hex(32) - -def resolve_scope_secret(config: Config) -> str: - """The HMAC key used to sign/verify ``scope_token`` -- shared across replicas when - ``config.jwt_secret`` (``KBC_JWT_SECRET``) is configured, otherwise a process-local fallback.""" - return config.jwt_secret or _FALLBACK_SCOPE_SECRET +def resolve_scope_key(config: Config) -> bytes: + """The AES-256 key used to encrypt/decrypt ``scope_token`` -- the same + ``KBC_SESSION_ENCRYPTION_KEY`` OAuth sessions already encrypt their stored credentials with + (shared across replicas when configured, otherwise a process-local fallback -- see + ``session_store.crypto.resolve_encryption_key``). ``scope_token`` may carry a live + ``scoped_token`` bearer credential, so it needs the same at-rest protection OAuth sessions + get, not just a signature -- see the "Security hardening" RFC increment. + """ + return resolve_encryption_key(config.session_encryption_key) @dataclasses.dataclass(frozen=True) @@ -90,15 +91,29 @@ def is_near_expiry(self) -> bool: return False return time.time() >= (self.scoped_expires_at - 60) - def to_token(self, secret: str) -> str: - """Signs this scope into the opaque ``scope_token`` a caller resends on later calls.""" - return encode_jwt(dataclasses.asdict(self), secret) + def to_token(self, key: bytes) -> str: + """Encrypts this scope into the opaque ``scope_token`` a caller resends on later calls. + + AES-GCM (authenticated encryption), not a bare signature: this may carry a live + ``scoped_token`` bearer credential, which must not be recoverable by anyone without the + key -- unlike a JWS, whose payload is trivially base64+gunzip-recoverable regardless of + whether the signature itself can be forged. See the "Security hardening" RFC increment. + """ + plaintext = gzip.compress(json.dumps(dataclasses.asdict(self)).encode('utf-8')) + return base64.urlsafe_b64encode(encrypt(plaintext, key)).decode('ascii').rstrip('=') @classmethod - def from_token(cls, token: str, secret: str) -> 'SessionScope': - """Inverse of ``to_token``. Raises on a missing/invalid/tampered token -- callers should - treat any exception as "no scope" rather than fail the request.""" - return cls(**decode_jwt(token, secret)) + def from_token(cls, token: str, key: bytes) -> 'SessionScope': + """Inverse of ``to_token``. Raises on a missing/invalid/tampered/wrong-key token -- + callers should treat any exception as "no scope" rather than fail the request. + """ + padded = token + '=' * (-len(token) % 4) + plaintext = decrypt(base64.urlsafe_b64decode(padded), key) + data = json.loads(gzip.decompress(plaintext).decode('utf-8')) + # Ignore any unknown keys rather than raising -- forward-compat if a future field is added + # to SessionScope after this token was minted. + known_fields = {f.name for f in dataclasses.fields(cls)} + return cls(**{k: v for k, v in data.items() if k in known_fields}) async def persist_scope(session_store: 'SessionStore', session_id: str, scope: SessionScope) -> None: diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 1bc1dd25f..d818bbd92 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -24,7 +24,7 @@ ProjectIdArg, SessionScope, persist_scope, - resolve_scope_secret, + resolve_scope_key, ) from keboola_mcp_server.workspace import WorkspaceManager @@ -530,7 +530,7 @@ async def get_accessible_projects( scoped_project_ids=scoped_ids, read_only=scope.read_only if scoped_ids is not None else None, scope_token=( - scope.to_token(resolve_scope_secret(server_state.config)) + scope.to_token(resolve_scope_key(server_state.config)) if scoped_ids is not None and not is_persisted else None ), @@ -625,13 +625,23 @@ async def set_project_scope( or await _persist_kai_scope(ctx, scope, client, parent_token) or server_state.runtime_info.session_state_persists ) - scope_token = None if persisted else scope.to_token(resolve_scope_secret(server_state.config)) + scope_token = None if persisted else scope.to_token(resolve_scope_key(server_state.config)) resend_instruction = ( 'The server persists this scope server-side for the rest of the conversation -- no need to resend it.' if persisted else 'The server does not remember this scope between calls -- pass "scope_token" as an argument on ' 'every subsequent tool call in this conversation.' ) + # Read-only is always enforced locally (this server blocks write operations regardless of + # scoped_token) -- but only backed by Connection itself when a real scoped_token exists. The + # exchange-failure fallback above has none, so say so explicitly rather than implying the same + # server-side guarantee the success path gets. + read_only_note = ( + ' (enforced by this server only -- the scoped-token exchange was unavailable, so Connection ' + 'itself does not additionally restrict this token.)' + if scope.read_only and scope.scoped_token is None + else '' + ) return ProjectScope( project_ids=ids, read_only=scope.read_only, @@ -640,9 +650,9 @@ async def set_project_scope( ( f'Session scoped to {len(ids)} projects. Read-only tools return results per project. ' 'Write operations require a project_id argument naming which scoped project to target ' - f'-- no re-scope needed to switch targets. {resend_instruction}' + f'-- no re-scope needed to switch targets. {resend_instruction}{read_only_note}' ) if multi - else f'Session scoped to project {ids[0]}. {resend_instruction}' + else f'Session scoped to project {ids[0]}. {resend_instruction}{read_only_note}' ), ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 701eb4ffc..70fefad50 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,5 @@ import asyncio +import base64 import dataclasses import time from datetime import datetime, timedelta, timezone @@ -25,7 +26,7 @@ toon_serializer, unwrap_results, ) -from keboola_mcp_server.scope import SCOPE_KEY, SessionScope, resolve_scope_secret +from keboola_mcp_server.scope import SCOPE_KEY, SCOPE_TOKEN_ARG, SessionScope, resolve_scope_key from keboola_mcp_server.workspace import WorkspaceManager @@ -760,6 +761,49 @@ async def fake_create_session_state(cfg, _runtime_info, readonly=None, *, own_st # whether the Kubernetes step-up header may be sent. assert captured_own_stack_urls == ['https://connection.test.keboola.com'] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ('scope', 'expected_readonly'), + [ + (None, None), + (SessionScope(project_ids=[18], read_only=False, confirmed=True), None), + (SessionScope(project_ids=[18], read_only=True, confirmed=True), True), + ], + ids=['no_scope', 'writable_scope', 'readonly_scope'], + ) + async def test_on_request_threads_scope_read_only_into_session_state(self, scope, expected_readonly) -> None: + # Security hardening RFC increment: a read-only confirmed scope must be enforced on the + # base session client too, not just relied on via the (possibly-absent) scoped_token. + config = Config(storage_api_url='https://connection.test.keboola.com', storage_token='kbc_at_x') + server_state = ServerState(config=config, runtime_info=ServerRuntimeInfo(transport='stdio')) + session = SimpleNamespace(state={}) + ctx = MagicMock(spec=Context) + ctx.session = session + ctx.request_context.lifespan_context = server_state + + args = {} + if scope is not None: + args[SCOPE_TOKEN_ARG] = scope.to_token(resolve_scope_key(config)) + context = SimpleNamespace(message=SimpleNamespace(arguments=args), method='tools/call', fastmcp_context=ctx) + + captured_readonly = [] + + async def fake_create_session_state(cfg, _runtime_info, readonly=None, *, own_stack_storage_api_url): + captured_readonly.append(readonly) + return {} + + async def call_next(_): + return 'ok' + + middleware = SessionStateMiddleware() + with ( + patch.object(middleware, 'create_session_state', side_effect=fake_create_session_state), + patch('keboola_mcp_server.mcp.get_http_request_or_none', return_value=None), + ): + await middleware.on_request(context, call_next) + + assert captured_readonly == [expected_readonly] + @pytest.mark.parametrize( ('server_storage_api_url', 'headers', 'expected_storage_api_url'), [ @@ -1199,6 +1243,40 @@ async def test_no_stored_session_is_noop(self, monkeypatch) -> None: assert out.storage_token is None +class TestReadPersistedLoginScope: + """Local sessions are scoped at `login` time (Security hardening RFC increment) -- + SessionStateMiddleware._read_persisted_login_scope.""" + + @pytest.mark.asyncio + async def test_returns_confirmed_scope_from_stored_credential(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x') + stored = SimpleNamespace(project_ids=[18, 83], read_only=True) + with patch('keboola_mcp_server.mcp.load_tokens', return_value=stored): + scope = await SessionStateMiddleware._read_persisted_login_scope(config) + + assert scope == SessionScope(project_ids=[18, 83], read_only=True, confirmed=True) + + @pytest.mark.asyncio + async def test_none_when_credential_predates_the_scoping_choice(self, monkeypatch) -> None: + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x') + stored = SimpleNamespace(project_ids=None, read_only=False) + with patch('keboola_mcp_server.mcp.load_tokens', return_value=stored): + scope = await SessionStateMiddleware._read_persisted_login_scope(config) + + assert scope is None + + @pytest.mark.asyncio + async def test_none_when_not_local_programmatic(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x') + with patch('keboola_mcp_server.mcp.load_tokens', side_effect=AssertionError('must not be called')): + scope = await SessionStateMiddleware._read_persisted_login_scope(config) + + assert scope is None + + class TestResolveLocalTokens: """SessionStateMiddleware keeps local tokens fresh and re-mints the scoped token (PSGO-261).""" @@ -1386,30 +1464,51 @@ async def test_autolease_noop_when_deployed(self, monkeypatch) -> None: assert await SessionStateMiddleware._autolease_default_scope(config) is None +def _b64_key(fill: bytes) -> str: + """A base64-encoded 32-byte KBC_SESSION_ENCRYPTION_KEY built from a repeated fill byte -- + deterministic test keys, distinct fills give distinct keys.""" + return base64.b64encode(fill * 32).decode('ascii') + + class TestScopeToken: """The multi-project scope is carried by the caller as the `scope_token` tool argument, not read back from ctx.session.state -- which is rebuilt empty on every request under this server's default stateless-HTTP transport, so nothing survives there between one tool call and the next. + Encrypted (AES-GCM), not just signed -- it may carry a live `scoped_token` bearer credential. """ + KEY_A = base64.b64decode(_b64_key(b'\x01')) + KEY_B = base64.b64decode(_b64_key(b'\x02')) + def test_round_trip(self) -> None: scope = SessionScope( project_ids=[11, 22], read_only=True, scoped_token='kbc_at_s', scoped_expires_at=1234.0, confirmed=True ) - token = scope.to_token('secret') - assert SessionScope.from_token(token, 'secret') == scope - - def test_wrong_secret_rejected(self) -> None: - token = SessionScope(project_ids=[11], confirmed=True).to_token('secret-a') + token = scope.to_token(self.KEY_A) + assert SessionScope.from_token(token, self.KEY_A) == scope + + def test_token_does_not_contain_the_scoped_token_in_the_clear(self) -> None: + # The whole point of encrypting rather than just signing: a live bearer credential must + # not be recoverable from the client-visible blob without the key. + scope = SessionScope(project_ids=[11], scoped_token='kbc_at_super_secret_live_token', confirmed=True) + token = scope.to_token(self.KEY_A) + assert 'kbc_at_super_secret_live_token' not in token + # Also not recoverable via a bare base64-decode (no key at all) -- unlike the old JWS. + padded = token + '=' * (-len(token) % 4) + assert b'kbc_at_super_secret_live_token' not in base64.urlsafe_b64decode(padded) + + def test_wrong_key_rejected(self) -> None: + token = SessionScope(project_ids=[11], confirmed=True).to_token(self.KEY_A) with pytest.raises(Exception, match='.+'): - SessionScope.from_token(token, 'secret-b') + SessionScope.from_token(token, self.KEY_B) - def test_resolve_scope_secret_prefers_configured_jwt_secret(self) -> None: - assert resolve_scope_secret(Config(jwt_secret='shared-secret')) == 'shared-secret' + def test_resolve_scope_key_prefers_configured_session_encryption_key(self) -> None: + key = _b64_key(b'\x03') + assert resolve_scope_key(Config(session_encryption_key=key)) == base64.b64decode(key) - def test_resolve_scope_secret_fallback_is_stable_within_process(self) -> None: + def test_resolve_scope_key_fallback_is_stable_within_process(self) -> None: config = Config() - assert resolve_scope_secret(config) == resolve_scope_secret(config) + assert resolve_scope_key(config) == resolve_scope_key(config) @staticmethod def _call_tool_context(arguments: dict) -> SimpleNamespace: @@ -1417,9 +1516,10 @@ def _call_tool_context(arguments: dict) -> SimpleNamespace: return SimpleNamespace(message=message, method='tools/call') def test_read_scope_from_request_decodes_and_pops_token(self) -> None: - config = Config(jwt_secret='shared-secret') + key = _b64_key(b'\x04') + config = Config(session_encryption_key=key) scope = SessionScope(project_ids=[11, 22], confirmed=True) - arguments = {'scope_token': scope.to_token('shared-secret'), 'other_arg': 1} + arguments = {'scope_token': scope.to_token(base64.b64decode(key)), 'other_arg': 1} context = self._call_tool_context(arguments) result = SessionStateMiddleware._read_scope_from_request(context, config) @@ -1431,11 +1531,11 @@ def test_read_scope_from_request_decodes_and_pops_token(self) -> None: @pytest.mark.parametrize( 'arguments', - [{}, {'scope_token': None}, {'scope_token': ''}, {'scope_token': 'not-a-valid-jwt'}], + [{}, {'scope_token': None}, {'scope_token': ''}, {'scope_token': 'not-a-valid-token'}], ids=['missing', 'none', 'empty', 'malformed'], ) def test_read_scope_from_request_returns_none_when_absent_or_invalid(self, arguments: dict) -> None: - config = Config(jwt_secret='shared-secret') + config = Config(session_encryption_key=_b64_key(b'\x05')) context = self._call_tool_context(dict(arguments)) assert SessionStateMiddleware._read_scope_from_request(context, config) is None @@ -1444,12 +1544,13 @@ def test_read_scope_from_request_ignores_non_call_tool_requests(self) -> None: context = SimpleNamespace(message=SimpleNamespace(), method='tools/list', fastmcp_context=None) assert SessionStateMiddleware._read_scope_from_request(context, Config()) is None - def test_wrong_secret_falls_back_to_no_scope_via_read_scope_from_request(self) -> None: - # A token minted with a different secret (e.g. a replica whose fallback secret differs) must + def test_wrong_key_falls_back_to_no_scope_via_read_scope_from_request(self) -> None: + # A token minted with a different key (e.g. a replica whose fallback key differs) must # degrade to "no scope" rather than raise -- the ask-first gate then re-prompts the caller. - token = SessionScope(project_ids=[11], confirmed=True).to_token('secret-a') + token = SessionScope(project_ids=[11], confirmed=True).to_token(self.KEY_A) context = self._call_tool_context({'scope_token': token}) - assert SessionStateMiddleware._read_scope_from_request(context, Config(jwt_secret='secret-b')) is None + config = Config(session_encryption_key=base64.b64encode(self.KEY_B).decode()) + assert SessionStateMiddleware._read_scope_from_request(context, config) is None @staticmethod def _http_rq_with_oauth_user(**access_token_kwargs) -> SimpleNamespace: diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index d69295640..044f6dd22 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -10,7 +10,7 @@ from keboola_mcp_server.config import Config, MetadataField, ServerRuntimeInfo from keboola_mcp_server.links import Link from keboola_mcp_server.mcp import ServerState -from keboola_mcp_server.scope import OAUTH_SESSION_ID_KEY, SCOPE_KEY, SessionScope, resolve_scope_secret +from keboola_mcp_server.scope import OAUTH_SESSION_ID_KEY, SCOPE_KEY, SessionScope, resolve_scope_key from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, @@ -460,7 +460,7 @@ async def test_set_project_scope_returns_scope_token_when_session_does_not_persi scope = mcp_context_client.session.state[SCOPE_KEY] assert result.scope_token is not None - assert SessionScope.from_token(result.scope_token, resolve_scope_secret(Config())) == scope + assert SessionScope.from_token(result.scope_token, resolve_scope_key(Config())) == scope assert 'does not remember this scope' in result.llm_instruction @@ -600,6 +600,39 @@ async def test_set_project_scope_falls_back_on_network_error( assert scope.scoped_token is None +@pytest.mark.asyncio +async def test_set_project_scope_read_only_fallback_notes_local_only_enforcement( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + # Security hardening RFC increment: when the exchange fails, read_only has no server-side + # backing (no scoped_token) -- the caller must be told explicitly, not left assuming the same + # guarantee the success path gets. + _prep_client(mcp_context_client, mocker) + mocker.patch( + 'keboola_mcp_server.tools.project.exchange_scoped_token', + new=mocker.AsyncMock(side_effect=httpx.ConnectTimeout('timed out')), + ) + + result = await set_project_scope(mcp_context_client, project_ids=[18], read_only=True) + + assert result.read_only is True + assert 'enforced by this server only' in result.llm_instruction + + +@pytest.mark.asyncio +async def test_set_project_scope_read_only_success_omits_local_only_note( + mcp_context_client: Context, mocker: MockerFixture +) -> None: + _prep_client(mcp_context_client, mocker) + minted = SimpleNamespace(access_token='kbc_at_scoped', expires_at=time.time() + 3600, read_only=True) + mocker.patch('keboola_mcp_server.tools.project.exchange_scoped_token', new=mocker.AsyncMock(return_value=minted)) + + result = await set_project_scope(mcp_context_client, project_ids=[18], read_only=True) + + assert result.read_only is True + assert 'enforced by this server only' not in result.llm_instruction + + @pytest.mark.asyncio @pytest.mark.parametrize( 'bearer', From 7584d63419ca63e388633ebb158fbe02087db487 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Sun, 9 Aug 2026 10:55:37 +0200 Subject: [PATCH 80/89] fix(PSGO-261): scope login-time, key credentials per interface, prompt MFA, redact errors Four related local-login hardening fixes: - login (and login --pat) now require an explicit project-scope choice -- prompted interactively, or --project-ids/--all when not run from a terminal -- persisted alongside the tokens as TokenSet.project_ids/ read_only. Replaces leasing/auto-scoping to every accessible project by default with only a prompt-text "ask first" gate that nothing enforced. - Credentials are keyed by (hostname, profile) instead of hostname alone (--profile / KBC_LOGIN_PROFILE, default "default"), so two local MCP client interfaces logged into the same stack (e.g. Claude Desktop and a terminal) no longer share one entry and its rotating refresh token. An asyncio.Lock per (hostname, profile) serializes concurrent in-process refreshes; a non-blocking fcntl.flock insurance layer covers the on-disk read-modify-write for whatever narrow sharing remains. - MFA codes (--totp/--recovery) default to a getpass.getpass() prompt instead of requiring a CLI argument that sits in shell history/`ps` for the process lifetime; the flags remain as opt-in overrides for scripted use. - elevate_session/create_pat no longer raise the raw auth-endpoint response body in the exception message; full detail moves to LOG.debug only. --- pyproject.toml | 2 +- src/keboola_mcp_server/auth_login.py | 191 +++++++++++++++++++++------ src/keboola_mcp_server/cli.py | 179 ++++++++++++++++++++++--- tests/test_auth_login.py | 155 ++++++++++++++++++++++ tests/test_cli.py | 143 +++++++++++++++++++- uv.lock | 12 +- 6 files changed, 616 insertions(+), 66 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c5cdbf708..22548bf12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.75.0" +version = "1.76.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index c3a9d3951..7cda71245 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -10,7 +10,10 @@ an injected httpx transport. """ +import asyncio import base64 +import contextlib +import dataclasses import hashlib import json import logging @@ -30,6 +33,11 @@ from keboola_mcp_server.clients.base import normalize_storage_api_url +try: + import fcntl +except ImportError: # pragma: no cover - non-POSIX; the cross-process lock degrades to a no-op + fcntl = None # type: ignore[assignment] + LOG = logging.getLogger(__name__) DEFAULT_CLIENT_ID = 'keboola-cli-demo' @@ -46,6 +54,26 @@ _LOGIN_CALLBACK_TIMEOUT_SECONDS = 300 _PAT_DEFAULT_EXPIRES_SECONDS = 30 * 24 * 60 * 60 # ~1 month _CREDENTIALS_PATH = Path.home() / '.keboola' / 'mcp' / 'credentials.json' +# Names which local interface (Claude Desktop, Cursor, a terminal `login`) a stored session +# belongs to, so two interfaces logged in to the *same* stack never share one entry (and its +# rotating refresh token) -- see the "Security hardening" RFC increment. Each interface's MCP +# client config sets this to a distinct value; a single-interface setup needs nothing set. +_PROFILE_ENV_VAR = 'KBC_LOGIN_PROFILE' +_DEFAULT_PROFILE = 'default' +# Cross-process insurance only (see `_credentials_lock`); the actual fix for the credential race +# is per-profile keying above. Non-blocking poll, never a blocking flock -- this runs inside +# `get_access_token`, which must never stall the event loop / MCP handshake. +_LOCK_POLL_INTERVAL_SECONDS = 0.05 +_LOCK_TIMEOUT_SECONDS = 10.0 +# One asyncio.Lock per (hostname, profile), serializing concurrent refreshes *within this +# process* -- the flock below only ever guards the on-disk file, not in-memory races. +_refresh_locks: dict[str, asyncio.Lock] = {} + + +def _resolve_profile(profile: str | None) -> str: + return profile or os.environ.get(_PROFILE_ENV_VAR) or _DEFAULT_PROFILE + + # Short connect timeout so an unreachable stack (e.g. VPN off — internal `.dev` stacks resolve to a # private 10.x IP) fails in a few seconds with a clear ConnectTimeout instead of blocking the full # window. A longer read timeout still tolerates a slow-but-reachable Connection. @@ -67,12 +95,20 @@ def _b64url(data: bytes) -> str: @dataclass(frozen=True) class TokenSet: - """A leased session: the access token plus what's needed to refresh it.""" + """A leased session: the access token plus what's needed to refresh it. + + ``project_ids``/``read_only`` are the scope chosen at `login` time (None only for a + credential predating this choice, or one never run through `login`'s prompt/flags) -- see + the "Security hardening" RFC increment: a local session is scoped before it's ever usable, + rather than auto-leased to everything with an unenforceable ask-first gate. + """ access_token: str refresh_token: str expires_at: float # epoch seconds session_id: str | None = None + project_ids: list[int] | None = None + read_only: bool = False @property def is_near_expiry(self) -> bool: @@ -213,7 +249,11 @@ async def elevate_session( json=payload, ) if response.is_error: - raise RuntimeError(f'POST /{_SUDO_PATH} failed ({response.status_code}): {response.text}') + # The response body may echo back request details; never surface it directly to the + # caller (it can end up in a CLI transcript/bug report) -- full detail goes to debug + # logs only. + LOG.debug(f'POST /{_SUDO_PATH} failed ({response.status_code}): {response.text}') + raise RuntimeError(f'POST /{_SUDO_PATH} failed ({response.status_code}). See debug logs for details.') body = cast(dict, response.json()) if response.content else {} return cast(str, body.get('token') or body.get('accessToken') or subject_token) @@ -247,8 +287,10 @@ async def create_pat( json=payload, ) if response.is_error: - # Surface the validation body so a wrong/missing field is visible (the schema is assumed). - raise RuntimeError(f'POST /{_PAT_PATH} failed ({response.status_code}) with {payload=}: {response.text}') + # Full detail (request payload + response body) to debug logs only -- never surfaced + # directly, since it can end up in a CLI transcript/bug report. + LOG.debug(f'POST /{_PAT_PATH} failed ({response.status_code}) with {payload=}: {response.text}') + raise RuntimeError(f'POST /{_PAT_PATH} failed ({response.status_code}). See debug logs for details.') body = cast(dict, response.json()) pat = body.get('token') or body.get('pat') or body.get('accessToken') if not pat: @@ -260,18 +302,24 @@ async def lease_pat( storage_api_url: str, *, subject_token: str, + project_ids: list[int] | None = None, totp_code: str | None = None, recovery_code: str | None = None, name: str = 'keboola-mcp-server', expires_in: int = _PAT_DEFAULT_EXPIRES_SECONDS, transport: httpx.AsyncBaseTransport | None = None, ) -> str: - """Leases a PAT over ALL accessible projects: introspect → sudo (MFA) → create PAT. + """Leases a PAT: introspect (or use the caller's explicit ``project_ids``) → sudo (MFA) → + create PAT. ``subject_token`` is the whole-stack session access token (``kbc_at_*``) from the PKCE login. + ``project_ids=None`` means "every project the token can currently reach" -- callers making an + explicit choice (e.g. `login --pat`'s scoping prompt) should always pass it explicitly instead + of relying on this default, per the "Security hardening" RFC increment. """ - introspection = await introspect_token(storage_api_url, subject_token=subject_token, transport=transport) - project_ids = [p.id for p in introspection.projects] + if project_ids is None: + introspection = await introspect_token(storage_api_url, subject_token=subject_token, transport=transport) + project_ids = [p.id for p in introspection.projects] if not project_ids: raise RuntimeError('The session token can not reach any projects; cannot create a PAT.') elevated = await elevate_session( @@ -329,11 +377,12 @@ async def refresh_tokens( return parse_token_response(cast(dict, response.json())) -# --- credential storage (mode-600 file, keyed by stack host) --- +# --- credential storage (mode-600 file, keyed by stack host + interface profile) --- -def _store_key(storage_api_url: str) -> str: - return cast(str, urlparse(storage_api_url).hostname) +def _store_key(storage_api_url: str, profile: str | None = None) -> str: + hostname = cast(str, urlparse(storage_api_url).hostname) + return f'{hostname}::{_resolve_profile(profile)}' def _read_store() -> dict: @@ -362,49 +411,116 @@ def _write_store(store: dict) -> None: json.dump(store, f, indent=2, ensure_ascii=False) -def load_tokens(storage_api_url: str) -> TokenSet | None: - entry = _read_store().get(_store_key(storage_api_url)) +@contextlib.asynccontextmanager +async def _credentials_lock(): + """Cross-process insurance around the on-disk read-modify-write (defense in depth; the + primary fix for the credential race is per-profile keying, see `_store_key`). Non-blocking + poll of a sibling `.lock` file -- never a blocking `flock`, which would stall the event loop + and could hang the MCP initialize handshake. Degrades to a no-op (with a warning) on timeout + or on a non-POSIX platform where `fcntl` is unavailable. + """ + if fcntl is None: + yield + return + lock_path = _CREDENTIALS_PATH.parent / (_CREDENTIALS_PATH.name + '.lock') + lock_path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600) + try: + deadline = time.time() + _LOCK_TIMEOUT_SECONDS + locked = False + while True: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + locked = True + break + except BlockingIOError: + if time.time() >= deadline: + LOG.warning('Timed out waiting for the credentials file lock; proceeding without it.') + break + await asyncio.sleep(_LOCK_POLL_INTERVAL_SECONDS) + try: + yield + finally: + if locked: + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + +def load_tokens(storage_api_url: str, *, profile: str | None = None) -> TokenSet | None: + entry = _read_store().get(_store_key(storage_api_url, profile)) if not entry: return None return TokenSet(**entry) -def save_tokens(storage_api_url: str, tokens: TokenSet) -> None: +def save_tokens(storage_api_url: str, tokens: TokenSet, *, profile: str | None = None) -> None: store = _read_store() - store[_store_key(storage_api_url)] = asdict(tokens) + store[_store_key(storage_api_url, profile)] = asdict(tokens) _write_store(store) async def get_access_token( storage_api_url: str, *, + profile: str | None = None, transport: httpx.AsyncBaseTransport | None = None, ) -> str: """ - Returns a valid access token for the stack, refreshing (and persisting the rotated + Returns a valid access token for the stack+profile, refreshing (and persisting the rotated pair) when near expiry. Raises if there are no stored credentials (run ``login``). + + Double-checked locking: an `asyncio.Lock` per (stack, profile) serializes concurrent + refreshes from within this process; a cross-process `flock` (see `_credentials_lock`) is + layered on top as insurance. After acquiring both, the stored tokens are re-read -- another + caller may have already refreshed while this one was waiting, in which case no network call + is made at all. """ - tokens = load_tokens(storage_api_url) + tokens = load_tokens(storage_api_url, profile=profile) if not tokens: raise RuntimeError( f'No stored credentials for {storage_api_url}. Run "keboola-mcp-server login --api-url " first.' ) - if tokens.is_near_expiry: - try: - tokens = await refresh_tokens(storage_api_url, refresh_token=tokens.refresh_token, transport=transport) - except httpx.HTTPStatusError as e: - # Dead token (refresh rejected). Drop the stale credentials and force a re-login. - _forget(storage_api_url) - raise RuntimeError( - f'Session for {storage_api_url} has expired; run "keboola-mcp-server login --api-url " again.' - ) from e - save_tokens(storage_api_url, tokens) - return tokens.access_token + if not tokens.is_near_expiry: + return tokens.access_token + + key = _store_key(storage_api_url, profile) + lock = _refresh_locks.setdefault(key, asyncio.Lock()) + async with lock: + async with _credentials_lock(): + tokens = load_tokens(storage_api_url, profile=profile) + if not tokens: + raise RuntimeError( + f'No stored credentials for {storage_api_url}. ' + 'Run "keboola-mcp-server login --api-url " first.' + ) + if not tokens.is_near_expiry: + return tokens.access_token + try: + refreshed = await refresh_tokens( + storage_api_url, refresh_token=tokens.refresh_token, transport=transport + ) + except httpx.HTTPStatusError as e: + # Dead token (refresh rejected). Only forget it if it's still the same refresh + # token we just tried -- another caller may have already rotated it, in which + # case dropping the (now newer) entry would just force an unnecessary re-login. + current = load_tokens(storage_api_url, profile=profile) + if current is not None and current.refresh_token == tokens.refresh_token: + _forget(storage_api_url, profile=profile) + raise RuntimeError( + f'Session for {storage_api_url} has expired; run "keboola-mcp-server login --api-url " again.' + ) from e + # The refresh response carries no scope -- carry the previously-persisted choice + # forward so a rotation never silently drops it. + refreshed = dataclasses.replace(refreshed, project_ids=tokens.project_ids, read_only=tokens.read_only) + save_tokens(storage_api_url, refreshed, profile=profile) + return refreshed.access_token async def ensure_access_token( storage_api_url: str, *, + profile: str | None = None, allow_interactive: bool = True, open_browser=webbrowser.open, transport: httpx.AsyncBaseTransport | None = None, @@ -425,23 +541,24 @@ async def ensure_access_token( client-driven OAuth regardless. """ try: - return await get_access_token(storage_api_url, transport=transport) + return await get_access_token(storage_api_url, profile=profile, transport=transport) except RuntimeError as exc: if not allow_interactive: raise LOG.info(f'No usable stored session for {storage_api_url} ({exc}); starting browser login.') - await perform_login(storage_api_url, open_browser=open_browser) - return await get_access_token(storage_api_url, transport=transport) + await perform_login(storage_api_url, profile=profile, open_browser=open_browser) + return await get_access_token(storage_api_url, profile=profile, transport=transport) -def _forget(storage_api_url: str) -> None: +def _forget(storage_api_url: str, *, profile: str | None = None) -> None: store = _read_store() - if store.pop(_store_key(storage_api_url), None) is not None: + if store.pop(_store_key(storage_api_url, profile), None) is not None: _write_store(store) -def forget_tokens(storage_api_url: str | None = None) -> bool: - """Deletes the stored PKCE session — for one stack, or all when ``storage_api_url`` is None. +def forget_tokens(storage_api_url: str | None = None, *, profile: str | None = None) -> bool: + """Deletes the stored PKCE session — for one stack+profile, or every stack/profile when + ``storage_api_url`` is None. Returns True if anything was removed. Used by the ``logout`` command so the next ``login`` starts a fresh browser flow (e.g. to switch user/token) instead of refreshing the old session. @@ -452,7 +569,7 @@ def forget_tokens(storage_api_url: str | None = None) -> bool: if storage_api_url is None: _write_store({}) return True - if store.pop(_store_key(storage_api_url), None) is not None: + if store.pop(_store_key(storage_api_url, profile), None) is not None: _write_store(store) return True return False @@ -476,7 +593,7 @@ def log_message(self, *args) -> None: # silence the default stderr logging pass -async def perform_login(storage_api_url: str, *, open_browser=webbrowser.open) -> TokenSet: +async def perform_login(storage_api_url: str, *, profile: str | None = None, open_browser=webbrowser.open) -> TokenSet: """Runs the interactive PKCE browser login and persists the resulting tokens.""" verifier = _b64url(secrets.token_bytes(48)) # 64 url-safe chars challenge = _b64url(hashlib.sha256(verifier.encode('ascii')).digest()) @@ -522,5 +639,5 @@ async def perform_login(storage_api_url: str, *, open_browser=webbrowser.open) - tokens = await exchange_code( storage_api_url, code=code, state=state, code_verifier=verifier, redirect_uri=redirect_uri ) - save_tokens(storage_api_url, tokens) + save_tokens(storage_api_url, tokens, profile=profile) return tokens diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index c80c1268b..aa10b612c 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -4,6 +4,7 @@ import asyncio import contextlib import dataclasses +import getpass import json import logging.config import os @@ -78,10 +79,34 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: help='Keboola Storage API URL (e.g. https://connection..keboola.com). ' 'Falls back to KBC_STORAGE_API_URL.', ) + login_parser.add_argument( + '--profile', + metavar='NAME', + help='Which local interface this login is for (Claude Desktop, Cursor, a terminal, ...). ' + 'Each interface needing its own session should use a distinct profile so they never share ' + 'one stored credential/refresh token. Falls back to KBC_LOGIN_PROFILE, then "default".', + ) + login_parser.add_argument( + '--project-ids', + metavar='ID[,ID...]', + help='Scope this login to these project ids (comma-separated). Skips the interactive prompt. ' + 'Required (with this or --all) when not run from a terminal.', + ) + login_parser.add_argument( + '--all', + dest='all_projects', + action='store_true', + help='Scope this login to every currently-accessible project. Skips the interactive prompt.', + ) + login_parser.add_argument( + '--read-only', + action='store_true', + help='Scope this login read-only (no write operations in any scoped project).', + ) login_parser.add_argument( '--pat', action='store_true', - help='After the browser login, lease a Personal Access Token (kbc_pat_) over all accessible ' + help='After the browser login, lease a Personal Access Token (kbc_pat_) over the scoped ' 'projects and print it. Requires an MFA code (--totp or --recovery).', ) login_parser.add_argument( @@ -90,9 +115,17 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: help='Also print the session access token (kbc_at_) to stdout — e.g. to pass as a header to a ' 'locally-run streamable-HTTP server. Note: it expires in ~1 hour.', ) - login_parser.add_argument('--totp', metavar='CODE', help='TOTP MFA code for the sudo elevation (--pat).') login_parser.add_argument( - '--recovery', metavar='CODE', help='Recovery MFA code for the sudo elevation (--pat); alternative to --totp.' + '--totp', + metavar='CODE', + help='TOTP MFA code for the sudo elevation (--pat). Visible in shell history/`ps` for the ' + 'process lifetime — prefer leaving this unset and entering the code at the prompt instead.', + ) + login_parser.add_argument( + '--recovery', + metavar='CODE', + help='Recovery MFA code for the sudo elevation (--pat); alternative to --totp. Single-use and ' + 'high-value — same shell-history/`ps` caveat as --totp; prefer the interactive prompt.', ) login_parser.add_argument( '--pat-name', @@ -104,7 +137,8 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: '--force', action='store_true', help='Force a fresh browser login even if a valid stored session exists (e.g. to switch ' - 'user/token). Without it, login refreshes the existing session.', + 'user/token). Without it, login refreshes the existing session. Also re-prompts for project ' + 'scope even if one is already stored.', ) logout_parser = subparsers.add_parser( @@ -116,6 +150,12 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: metavar='URL', help='Stack to log out of (default: KBC_STORAGE_API_URL). Use --all to clear every stack.', ) + logout_parser.add_argument( + '--profile', + metavar='NAME', + help='Which local interface to log out (see `login --profile`). Falls back to ' + 'KBC_LOGIN_PROFILE, then "default". Ignored with --all, which clears every profile.', + ) logout_parser.add_argument('--all', action='store_true', help='Delete stored sessions for all stacks.') subparsers.add_parser( @@ -172,9 +212,55 @@ async def _http_exception_handler(request: Request, exc: HTTPException): } +def _parse_project_ids(raw: str) -> list[int]: + try: + return [int(x.strip()) for x in raw.split(',') if x.strip()] + except ValueError: + raise RuntimeError(f'Could not parse --project-ids value: {raw!r} (expected comma-separated integers).') + + +def _prompt_project_selection(projects: list) -> tuple[list[int], bool]: + """Interactively asks which projects to scope this login to. Never returns an implicit + "everything" without the user seeing the list and choosing it -- see the "Security + hardening" RFC increment: a local session must be scoped before it's ever usable. + """ + print('\nAccessible projects:', file=sys.stderr) + for p in projects: + print(f' {p.id}' + (f' - {p.name}' if p.name else ''), file=sys.stderr) + raw = input('\nScope this login to which projects? [a]ll or comma-separated ids (default: all): ').strip() + if not raw or raw.lower() in ('a', 'all'): + project_ids = [p.id for p in projects] + else: + project_ids = _parse_project_ids(raw) + valid_ids = {p.id for p in projects} + if outside := [pid for pid in project_ids if pid not in valid_ids]: + raise RuntimeError(f'Project(s) {outside} are not accessible with this token.') + read_only = input('Read-only (no writes in any scoped project)? [y/N]: ').strip().lower() in ('y', 'yes') + return project_ids, read_only + + +def _prompt_mfa_code() -> tuple[str | None, str | None]: + """Prompts for a TOTP or recovery code via hidden input, instead of requiring a CLI argument + that would sit in shell history/`ps` for the process lifetime -- see the "Security hardening" + RFC increment. `getpass.getpass` degrades gracefully (visible input, with a stderr warning) on + a non-interactive stdin, so piped/scripted input still works. + """ + totp = getpass.getpass('TOTP code (leave blank to use a recovery code instead): ').strip() + if totp: + return totp, None + recovery = getpass.getpass('Recovery code: ').strip() + if not recovery: + raise RuntimeError('Leasing a PAT (--pat) requires an MFA code (TOTP or recovery).') + return None, recovery + + async def _run_login( api_url: str | None, *, + profile: str | None = None, + project_ids_arg: str | None = None, + all_projects: bool = False, + read_only: bool = False, pat: bool = False, totp: str | None = None, recovery: str | None = None, @@ -182,51 +268,94 @@ async def _run_login( show_token: bool = False, force: bool = False, ) -> None: - """Establishes a stored session and, with ``pat=True``, leases a PAT. + """Establishes a stored session, scoped to an explicit set of projects, and with ``pat=True`` + leases a PAT over that same scope. Refresh-first: if a stored session exists and its refresh token is still valid, this refreshes (no browser) — so re-running `login` an hour later just leases a fresh access token. A browser PKCE login runs only when there is no stored session or the refresh token itself is dead. - With ``pat=True``, additionally leases a Personal Access Token over all accessible projects - (introspect → sudo with the MFA code → create PAT) and prints it. + Project scope is chosen once, here, and persisted alongside the tokens (see + `auth_login.TokenSet`) — a local session is never auto-leased to every project with only a + prompt-text "ask first" gate; see the "Security hardening" RFC increment. Already-scoped + sessions keep their existing choice on a plain re-run; pass `--project-ids`/`--all` or + `--force` to change it. + + With ``pat=True``, additionally leases a Personal Access Token over the same scope + (sudo with the MFA code → create PAT) and prints it. """ - from keboola_mcp_server.auth_login import ensure_access_token, forget_tokens, lease_pat, load_tokens, perform_login + from keboola_mcp_server.auth_login import ( + ensure_access_token, + forget_tokens, + introspect_token, + lease_pat, + load_tokens, + perform_login, + save_tokens, + ) storage_api_url = api_url or os.environ.get('KBC_STORAGE_API_URL') if not storage_api_url: raise RuntimeError('A Storage API URL is required for login: pass --api-url or set KBC_STORAGE_API_URL.') - - if pat and bool(totp) == bool(recovery): - raise RuntimeError('Leasing a PAT (--pat) requires exactly one MFA code: pass --totp or --recovery.') + if project_ids_arg and all_projects: + raise RuntimeError('Pass either --project-ids or --all, not both.') if force: # Drop any stored session and always run the browser flow (e.g. to switch user/token). - forget_tokens(storage_api_url) - access_token = (await perform_login(storage_api_url)).access_token + forget_tokens(storage_api_url, profile=profile) + access_token = (await perform_login(storage_api_url, profile=profile)).access_token else: # Refresh-first, browser only when dead (interactive: this is the terminal `login` command). - access_token = await ensure_access_token(storage_api_url, allow_interactive=True) - tokens = load_tokens(storage_api_url) - remaining = max(0, int(tokens.expires_at - time.time())) if tokens else 0 - print(f'\n✓ Session ready for {storage_api_url} (access token expires in ~{remaining}s).') + access_token = await ensure_access_token(storage_api_url, profile=profile, allow_interactive=True) + tokens = load_tokens(storage_api_url, profile=profile) + assert tokens is not None # ensure_access_token/perform_login above always persist one + + if tokens.project_ids is not None and not force and not project_ids_arg and not all_projects: + # Already scoped from an earlier login (and not asked to change it) -- keep it as-is. + project_ids, project_read_only = tokens.project_ids, tokens.read_only + elif project_ids_arg: + project_ids, project_read_only = _parse_project_ids(project_ids_arg), read_only + elif all_projects: + introspection = await introspect_token(storage_api_url, subject_token=access_token) + project_ids, project_read_only = [p.id for p in introspection.projects], read_only + elif sys.stdin.isatty(): + introspection = await introspect_token(storage_api_url, subject_token=access_token) + project_ids, project_read_only = _prompt_project_selection(introspection.projects) + else: + raise RuntimeError( + 'A project scope is required for login: pass --project-ids or --all ' + '(not run from a terminal, so the interactive prompt is unavailable).' + ) + tokens = dataclasses.replace(tokens, project_ids=project_ids, read_only=project_read_only) + save_tokens(storage_api_url, tokens, profile=profile) + + remaining = max(0, int(tokens.expires_at - time.time())) + print( + f'\n✓ Session ready for {storage_api_url} (access token expires in ~{remaining}s), ' + f'scoped to {len(project_ids)} project(s)' + (', read-only' if project_read_only else '') + '.' + ) if show_token: # Explicitly requested (e.g. to pass as a header to a local streamable-HTTP server). print(f'\nAccess token (kbc_at_, expires in ~{remaining}s):\n\n {access_token}\n') if pat: + if bool(totp) == bool(recovery): + if totp or recovery: + raise RuntimeError('Leasing a PAT (--pat) requires exactly one MFA code: pass --totp or --recovery.') + totp, recovery = _prompt_mfa_code() pat_token = await lease_pat( storage_api_url, subject_token=access_token, + project_ids=project_ids, totp_code=totp, recovery_code=recovery, name=pat_name, ) - print(f'\n✓ Personal Access Token (valid ~1 month, all accessible projects):\n\n {pat_token}\n') + print(f'\n✓ Personal Access Token (valid ~1 month, {len(project_ids)} project(s)):\n\n {pat_token}\n') -async def _run_logout(api_url: str | None, *, all_stacks: bool = False) -> None: +async def _run_logout(api_url: str | None, *, profile: str | None = None, all_stacks: bool = False) -> None: """Deletes the stored PKCE session so the next login starts fresh.""" from keboola_mcp_server.auth_login import forget_tokens @@ -239,7 +368,7 @@ async def _run_logout(api_url: str | None, *, all_stacks: bool = False) -> None: raise RuntimeError( 'A Storage API URL is required for logout: pass --api-url, set KBC_STORAGE_API_URL, or use --all.' ) - removed = forget_tokens(storage_api_url) + removed = forget_tokens(storage_api_url, profile=profile) print(f'✓ Logged out of {storage_api_url}.' if removed else f'No stored session for {storage_api_url}.') @@ -328,6 +457,10 @@ async def run_server(args: list[str] | None = None) -> None: if parsed_args.command == 'login': await _run_login( getattr(parsed_args, 'api_url', None), + profile=getattr(parsed_args, 'profile', None), + project_ids_arg=getattr(parsed_args, 'project_ids', None), + all_projects=getattr(parsed_args, 'all_projects', False), + read_only=getattr(parsed_args, 'read_only', False), pat=getattr(parsed_args, 'pat', False), totp=getattr(parsed_args, 'totp', None), recovery=getattr(parsed_args, 'recovery', None), @@ -338,7 +471,11 @@ async def run_server(args: list[str] | None = None) -> None: return if parsed_args.command == 'logout': - await _run_logout(getattr(parsed_args, 'api_url', None), all_stacks=getattr(parsed_args, 'all', False)) + await _run_logout( + getattr(parsed_args, 'api_url', None), + profile=getattr(parsed_args, 'profile', None), + all_stacks=getattr(parsed_args, 'all', False), + ) return if parsed_args.command == 'migrate': diff --git a/tests/test_auth_login.py b/tests/test_auth_login.py index 88f573a97..dd15ed09b 100644 --- a/tests/test_auth_login.py +++ b/tests/test_auth_login.py @@ -1,8 +1,10 @@ """Tests for the local browser PKCE login + credential store (PSGO-261, Part B).""" +import asyncio import base64 import hashlib import json +import logging import stat import time from pathlib import Path @@ -335,3 +337,156 @@ def handler(request: httpx.Request) -> httpx.Response: ) assert pat == 'kbc_pat_leased' assert [p.split('/')[-1] for p in seen] == ['introspect', 'sudo', 'pat'] + + +@pytest.mark.asyncio +async def test_lease_pat_uses_explicit_project_ids_without_introspecting() -> None: + # An explicit choice (e.g. from `login`'s scoping prompt) must be used as-is -- lease_pat + # must not silently widen it back to every accessible project. + seen: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + path = request.url.path + seen.append(path) + if path.endswith('/token/introspect'): + raise AssertionError('must not introspect when project_ids is given explicitly') + if path.endswith('/auth/sudo'): + return httpx.Response(200, json={'token': 'kbc_sudo_1'}) + if path.endswith('/auth/pat'): + assert json.loads(request.content)['scope']['projects'] == ['18'] + return httpx.Response(201, json={'token': 'kbc_pat_leased'}) + raise AssertionError(f'unexpected path {path}') + + pat = await lease_pat( + STACK, + subject_token='kbc_at_parent', + project_ids=[18], + recovery_code='rec-9', + transport=httpx.MockTransport(handler), + ) + assert pat == 'kbc_pat_leased' + assert seen == ['/v1/auth/sudo', '/v1/auth/pat'] + + +# --- error redaction (Security hardening RFC increment) --- + + +@pytest.mark.asyncio +async def test_elevate_session_error_is_redacted(caplog) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(400, text='sensitive-detail-should-not-surface') + + with ( + caplog.at_level(logging.DEBUG, logger='keboola_mcp_server.auth_login'), + pytest.raises(RuntimeError) as exc_info, + ): + await elevate_session(STACK, subject_token='kbc_at_x', totp_code='1', transport=httpx.MockTransport(handler)) + assert 'sensitive-detail-should-not-surface' not in str(exc_info.value) + assert 'sensitive-detail-should-not-surface' in caplog.text + + +@pytest.mark.asyncio +async def test_create_pat_error_is_redacted(caplog) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(400, text='sensitive-detail-should-not-surface') + + with ( + caplog.at_level(logging.DEBUG, logger='keboola_mcp_server.auth_login'), + pytest.raises(RuntimeError) as exc_info, + ): + await create_pat( + STACK, + subject_token='kbc_sudo_1', + project_ids=[18], + name='demo', + transport=httpx.MockTransport(handler), + ) + assert 'sensitive-detail-should-not-surface' not in str(exc_info.value) + assert 'sensitive-detail-should-not-surface' in caplog.text + + +# --- per-profile credential keying (Security hardening RFC increment) --- + + +def test_different_profiles_same_stack_dont_collide(creds_file: Path) -> None: + save_tokens(STACK, TokenSet('kbc_at_desktop', 'kbc_rt_d', expires_at=time.time() + 3600), profile='desktop') + save_tokens(STACK, TokenSet('kbc_at_terminal', 'kbc_rt_t', expires_at=time.time() + 3600), profile='terminal') + + assert load_tokens(STACK, profile='desktop').access_token == 'kbc_at_desktop' + assert load_tokens(STACK, profile='terminal').access_token == 'kbc_at_terminal' + # No profile given resolves to the 'default' profile, distinct from either named one. + assert load_tokens(STACK) is None + + +def test_profile_env_var_is_the_default_when_none_given(creds_file: Path, monkeypatch) -> None: + monkeypatch.setenv('KBC_LOGIN_PROFILE', 'desktop') + save_tokens(STACK, TokenSet('kbc_at_desktop', 'kbc_rt', expires_at=time.time() + 3600), profile='desktop') + + assert load_tokens(STACK).access_token == 'kbc_at_desktop' + + +def test_forget_one_profile_leaves_other_profiles_of_same_stack(creds_file: Path) -> None: + save_tokens(STACK, TokenSet('a', 'r', expires_at=time.time() + 3600), profile='desktop') + save_tokens(STACK, TokenSet('b', 'r', expires_at=time.time() + 3600), profile='terminal') + + assert forget_tokens(STACK, profile='desktop') is True + assert load_tokens(STACK, profile='desktop') is None + assert load_tokens(STACK, profile='terminal') is not None + + +@pytest.mark.asyncio +async def test_get_access_token_preserves_scope_across_refresh(creds_file: Path) -> None: + save_tokens( + STACK, + TokenSet('kbc_at_old', 'kbc_rt_old', expires_at=time.time() + 5, project_ids=[18, 83], read_only=True), + ) + await get_access_token(STACK, transport=_token_response()) + + tokens = load_tokens(STACK) + assert tokens.access_token == 'kbc_at_new' + assert tokens.project_ids == [18, 83] + assert tokens.read_only is True + + +@pytest.mark.asyncio +async def test_concurrent_get_access_token_refreshes_once(creds_file: Path) -> None: + # Two callers racing a near-expiry refresh for the SAME (stack, profile) must only hit the + # network once -- the second one, after acquiring the lock, sees the already-refreshed token. + save_tokens(STACK, TokenSet('kbc_at_old', 'kbc_rt_old', expires_at=time.time() + 5)) + call_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal call_count + call_count += 1 + return httpx.Response( + 200, + json={'accessToken': 'kbc_at_new', 'refreshToken': 'kbc_rt_new', 'expiresIn': 3600, 'sessionId': 's'}, + ) + + transport = httpx.MockTransport(handler) + results = await asyncio.gather( + get_access_token(STACK, transport=transport), + get_access_token(STACK, transport=transport), + ) + assert results == ['kbc_at_new', 'kbc_at_new'] + assert call_count == 1 + + +@pytest.mark.asyncio +async def test_get_access_token_dead_refresh_does_not_clobber_newer_entry(creds_file: Path, monkeypatch) -> None: + # If the stored refresh token changed (another caller already rotated it) between our read + # and our failed refresh attempt, don't drop the newer entry. + save_tokens(STACK, TokenSet('kbc_at_old', 'kbc_rt_dead', expires_at=time.time() + 5)) + + async def fake_refresh(*_a, **_k): + # Simulate another process/caller rotating the token concurrently, then our own + # (now-stale) refresh attempt failing against the auth server. + save_tokens(STACK, TokenSet('kbc_at_newer', 'kbc_rt_newer', expires_at=time.time() + 3600)) + raise httpx.HTTPStatusError('dead', request=httpx.Request('POST', STACK), response=httpx.Response(401)) + + monkeypatch.setattr(auth_login, 'refresh_tokens', fake_refresh) + with pytest.raises(RuntimeError, match='has expired'): + await get_access_token(STACK) + + # The newer entry (written by the "other caller") must survive, not be forgotten. + assert load_tokens(STACK).access_token == 'kbc_at_newer' diff --git a/tests/test_cli.py b/tests/test_cli.py index 7796e9bbb..948bbeaf4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,8 +1,20 @@ +import time from unittest.mock import AsyncMock, MagicMock, patch import pytest -from keboola_mcp_server.cli import _run_gc_sessions, _run_migrate, parse_args +from keboola_mcp_server import auth_login +from keboola_mcp_server.auth_login import TokenSet +from keboola_mcp_server.cli import _run_gc_sessions, _run_login, _run_logout, _run_migrate, parse_args + +STACK = 'https://connection.keboola.com' + + +@pytest.fixture +def creds_file(tmp_path, monkeypatch): + path = tmp_path / 'creds' / 'credentials.json' + monkeypatch.setattr(auth_login, '_CREDENTIALS_PATH', path) + return path def test_parse_args_migrate() -> None: @@ -163,3 +175,132 @@ async def test_closes_pool_even_if_it_fails(self, monkeypatch) -> None: await _run_gc_sessions() pool.close.assert_awaited_once() + + +def _introspection(project_ids: list[int]): + from keboola_mcp_server.auth_login import Introspection, ProjectAccess + + return Introspection( + user_id=1, user_email='m@k.com', user_name='M', projects=[ProjectAccess(id=p) for p in project_ids] + ) + + +def _seed_unscoped_session(access_token: str = 'kbc_at_x') -> None: + """Simulates what `ensure_access_token`/`perform_login` normally persist -- a session with + no project scope chosen yet -- so `_run_login` (mocked past the actual network calls) has + something to `load_tokens` back.""" + auth_login.save_tokens(STACK, TokenSet(access_token, 'kbc_rt', expires_at=time.time() + 3600)) + + +class TestRunLogin: + """`login` scopes a session at login time (Security hardening RFC increment) -- never leaves + a local session auto-leased to everything with only a prompt-text ask-first gate.""" + + @pytest.mark.asyncio + async def test_project_ids_flag_persists_scope_without_prompting(self, creds_file, monkeypatch) -> None: + _seed_unscoped_session() + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + with patch('builtins.input', side_effect=AssertionError('must not prompt when --project-ids is given')): + await _run_login(STACK, project_ids_arg='18,83') + + tokens = auth_login.load_tokens(STACK) + assert tokens.project_ids == [18, 83] + assert tokens.read_only is False + + @pytest.mark.asyncio + async def test_all_flag_introspects_and_scopes_to_everything(self, creds_file, monkeypatch) -> None: + _seed_unscoped_session() + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + monkeypatch.setattr(auth_login, 'introspect_token', AsyncMock(return_value=_introspection([18, 83, 95]))) + await _run_login(STACK, all_projects=True, read_only=True) + + tokens = auth_login.load_tokens(STACK) + assert tokens.project_ids == [18, 83, 95] + assert tokens.read_only is True + + @pytest.mark.asyncio + async def test_interactive_prompt_scopes_selection(self, creds_file, monkeypatch) -> None: + _seed_unscoped_session() + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + monkeypatch.setattr(auth_login, 'introspect_token', AsyncMock(return_value=_introspection([18, 83, 95]))) + monkeypatch.setattr('sys.stdin.isatty', lambda: True) + with patch('builtins.input', side_effect=['18,83', 'y']): + await _run_login(STACK) + + tokens = auth_login.load_tokens(STACK) + assert tokens.project_ids == [18, 83] + assert tokens.read_only is True + + @pytest.mark.asyncio + async def test_non_interactive_without_scope_raises(self, creds_file, monkeypatch) -> None: + _seed_unscoped_session() + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + monkeypatch.setattr('sys.stdin.isatty', lambda: False) + with pytest.raises(RuntimeError, match='--project-ids'): + await _run_login(STACK) + + @pytest.mark.asyncio + async def test_plain_rerun_keeps_existing_persisted_scope(self, creds_file, monkeypatch) -> None: + auth_login.save_tokens( + STACK, TokenSet('kbc_at_old', 'kbc_rt', expires_at=time.time() + 3600, project_ids=[18], read_only=True) + ) + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_old')) + with patch('builtins.input', side_effect=AssertionError('must not re-prompt on a plain re-run')): + await _run_login(STACK) + + tokens = auth_login.load_tokens(STACK) + assert tokens.project_ids == [18] + assert tokens.read_only is True + + @pytest.mark.asyncio + async def test_force_reruns_the_prompt_even_with_an_existing_scope(self, creds_file, monkeypatch) -> None: + auth_login.save_tokens( + STACK, TokenSet('kbc_at_old', 'kbc_rt', expires_at=time.time() + 3600, project_ids=[18], read_only=True) + ) + monkeypatch.setattr(auth_login, 'forget_tokens', MagicMock(return_value=True)) + monkeypatch.setattr( + auth_login, + 'perform_login', + AsyncMock(return_value=TokenSet('kbc_at_new', 'kbc_rt_new', expires_at=time.time() + 3600)), + ) + await _run_login(STACK, project_ids_arg='83', force=True) + + tokens = auth_login.load_tokens(STACK) + assert tokens.project_ids == [83] + + @pytest.mark.asyncio + async def test_pat_prompts_for_mfa_when_neither_given(self, creds_file, monkeypatch) -> None: + _seed_unscoped_session() + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + lease_pat = AsyncMock(return_value='kbc_pat_new') + monkeypatch.setattr(auth_login, 'lease_pat', lease_pat) + with patch('getpass.getpass', side_effect=['123456']): + await _run_login(STACK, project_ids_arg='18,83', pat=True) + + lease_pat.assert_awaited_once() + assert lease_pat.await_args.kwargs['totp_code'] == '123456' + assert lease_pat.await_args.kwargs['recovery_code'] is None + assert lease_pat.await_args.kwargs['project_ids'] == [18, 83] + + @pytest.mark.asyncio + async def test_pat_explicit_totp_skips_prompt(self, creds_file, monkeypatch) -> None: + _seed_unscoped_session() + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + lease_pat = AsyncMock(return_value='kbc_pat_new') + monkeypatch.setattr(auth_login, 'lease_pat', lease_pat) + with patch('getpass.getpass', side_effect=AssertionError('must not prompt when --totp is given')): + await _run_login(STACK, project_ids_arg='18,83', pat=True, totp='654321') + + assert lease_pat.await_args.kwargs['totp_code'] == '654321' + + +class TestRunLogout: + @pytest.mark.asyncio + async def test_forgets_only_the_given_profile(self, creds_file) -> None: + auth_login.save_tokens(STACK, TokenSet('a', 'r', expires_at=time.time() + 3600), profile='desktop') + auth_login.save_tokens(STACK, TokenSet('b', 'r', expires_at=time.time() + 3600), profile='terminal') + + await _run_logout(STACK, profile='desktop') + + assert auth_login.load_tokens(STACK, profile='desktop') is None + assert auth_login.load_tokens(STACK, profile='terminal') is not None diff --git a/uv.lock b/uv.lock index 8423256a4..4561e9f57 100644 --- a/uv.lock +++ b/uv.lock @@ -22,7 +22,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "caio", marker = "python_full_version < '3.11'" }, + { name = "caio" }, ] sdist = { url = "https://files.pythonhosted.org/packages/67/e2/d7cb819de8df6b5c1968a2756c3cb4122d4fa2b8fc768b53b7c9e5edb646/aiofile-3.9.0.tar.gz", hash = "sha256:e5ad718bb148b265b6df1b3752c4d1d83024b93da9bd599df74b9d9ffcf7919b", size = 17943, upload-time = "2024-10-08T10:39:35.846Z" } wheels = [ @@ -41,7 +41,7 @@ resolution-markers = [ "python_full_version >= '3.11' and python_full_version < '3.13'", ] dependencies = [ - { name = "caio", marker = "python_full_version >= '3.11'" }, + { name = "caio" }, ] sdist = { url = "https://files.pythonhosted.org/packages/48/41/2fea7e193e061ce54eacc3b7bc0e6a99e4fcff43c78cf0a76dd781ed8334/aiofile-3.11.1.tar.gz", hash = "sha256:1f91912c6643d2a4e49ca4ae3514f0bf3867ce948a36d99a6411b8f4755f4cf9", size = 19342, upload-time = "2026-05-16T08:18:33.538Z" } wheels = [ @@ -1058,7 +1058,7 @@ name = "importlib-metadata" version = "9.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "zipp", marker = "python_full_version < '3.13'" }, + { name = "zipp" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } wheels = [ @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.75.0" +version = "1.76.0" source = { editable = "." } dependencies = [ { name = "asyncpg" }, @@ -2328,8 +2328,8 @@ name = "secretstorage" version = "3.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cryptography", marker = "python_full_version < '3.13' or sys_platform != 'win32'" }, - { name = "jeepney", marker = "python_full_version < '3.13' or sys_platform != 'win32'" }, + { name = "cryptography" }, + { name = "jeepney" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } wheels = [ From f214738049d1ab9b2025e5bbdc5409ec59ef558c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 11 Aug 2026 15:29:33 +0200 Subject: [PATCH 81/89] feat(PSGO-261): skip scoping when a session can only reach one project Applies to local login and OAuth: a session that can only reach one project has no scoping decision to make, so requiring it anyway is pure friction. --- feature_spec/pat_token_support/RFC.md | 25 +++++++ src/keboola_mcp_server/cli.py | 7 ++ src/keboola_mcp_server/oauth.py | 51 +++++++++++++- tests/test_cli.py | 18 +++++ tests/test_oauth.py | 99 +++++++++++++++++++++++++-- 5 files changed, 191 insertions(+), 9 deletions(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index dea4b84d6..dcf3ef284 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -979,3 +979,28 @@ the original framing (one narrower, one broader; see below). is narrow enough that AES-GCM-encrypting the existing JWS payload is proportionate; a new Postgres table mirroring `kai_scope.py` was considered and rejected as unneeded complexity for that narrow remaining surface. + +## Extension: single-project sessions never need scoping (increment 8) + +Follow-up observation, not from the review above: for both the local `login` flow (increment 7, +item 7) and the OAuth flow, a session whose token can reach exactly one project has no real +scoping decision to make -- prompting for it (locally) or requiring an explicit +`set_project_scope` call (OAuth) is pure friction. This is distinct from the "N of M projects" +case, which stays genuinely ambiguous for this server's OAuth grant (`claudai projectless` scope, +always whole-stack) -- introspection's count there is just the user's real total org membership, +not evidence of a prior scoping choice, so it isn't auto-confirmed. + +**Fix:** +- `cli.py`'s `_prompt_project_selection` skips the "which projects" question when introspection + returns exactly one project -- still asks read-only, then persists the single-project scope the + same way an explicit choice would be. +- `oauth.py`'s `exchange_authorization_code` introspects the freshly-exchanged session token + immediately after creating it; if exactly one project is reachable, it mints a scoped token + (mirroring `set_project_scope`'s own exchange-with-fallback pattern) and persists + `scope_confirmed=True`/`scope_project_ids=[that id]` on the session row right away -- no + `set_project_scope` call ever needed for that session. Best-effort: any introspection/exchange + failure here just leaves the session unconfirmed, exactly as before this fix; login itself never + fails because of it. +- `lease_pat`/`login --pat` need no separate change -- they already take an explicit + `project_ids` argument (increment 7), which now flows from the auto-detected single project when + applicable. diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index aa10b612c..5784c48a8 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -223,7 +223,14 @@ def _prompt_project_selection(projects: list) -> tuple[list[int], bool]: """Interactively asks which projects to scope this login to. Never returns an implicit "everything" without the user seeing the list and choosing it -- see the "Security hardening" RFC increment: a local session must be scoped before it's ever usable. + + Skips the "which projects" question when there's only one accessible project -- there's no + real choice to make, so asking it would just be friction; still asks read-only. """ + if len(projects) == 1: + print(f'\nOnly one accessible project ({projects[0].id}); scoping to it automatically.', file=sys.stderr) + read_only = input('Read-only (no writes in this project)? [y/N]: ').strip().lower() in ('y', 'yes') + return [projects[0].id], read_only print('\nAccessible projects:', file=sys.stderr) for p in projects: print(f' {p.id}' + (f' - {p.name}' if p.name else ''), file=sys.stderr) diff --git a/src/keboola_mcp_server/oauth.py b/src/keboola_mcp_server/oauth.py index 8094699df..921d2e38c 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -26,7 +26,13 @@ from pydantic import AnyHttpUrl, AnyUrl from starlette.exceptions import HTTPException -from keboola_mcp_server.auth_login import TokenSet, parse_token_response, refresh_tokens +from keboola_mcp_server.auth_login import ( + TokenSet, + exchange_scoped_token, + introspect_token, + parse_token_response, + refresh_tokens, +) from keboola_mcp_server.clients.auth_bridge import OAuthSessionExchanger, OAuthTokenExchangeError from keboola_mcp_server.config import deployed_sa_token_path from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt @@ -419,15 +425,56 @@ async def exchange_authorization_code( # Exchange the league OAuth access token for a whole-stack Keboola programmatic session. # The league token is used exactly once, here, and then never referenced again. token_set = await self._exchange_oauth_for_session(authorization_code.oauth_access_token.token) - access_token, refresh_token, _session = await self._session_store.create( + access_token, refresh_token, session = await self._session_store.create( client_id=client.client_id, user_email=None, kbc_access_token=token_set.access_token, kbc_refresh_token=token_set.refresh_token, kbc_access_expires_at=datetime.fromtimestamp(token_set.expires_at, tz=timezone.utc), ) + await self._auto_confirm_single_project_scope(session.id, token_set.access_token) return self._oauth_token(access_token, refresh_token, authorization_code.scopes) + async def _auto_confirm_single_project_scope(self, session_id: str, subject_token: str) -> None: + """If this freshly-created session's token can reach exactly one project, there's no real + scoping choice for the user to make -- confirm it immediately so the session is usable + without ever calling ``set_project_scope`` (mirrors the local ``login``/``login --pat`` + flow, which does the same for the same reason -- see the "Security hardening" RFC + increment). Any project count other than 1 is left untouched: this server's OAuth grant is + always whole-stack (``claudai projectless`` scope), so introspection's count there is just + the user's real total org membership, not a scoping decision to defer to. + + Best-effort: introspection/exchange failures here just leave the session unconfirmed, same + as before this method existed -- an explicit ``set_project_scope`` call still works. + """ + try: + introspection = await introspect_token(self._storage_api_url, subject_token=subject_token) + except Exception as e: + LOG.warning(f'Could not introspect new OAuth session for single-project auto-scope: {e}', exc_info=True) + return + if len(introspection.projects) != 1: + return + project_id = introspection.projects[0].id + scoped_token: str | None = None + scoped_expires_at: datetime | None = None + try: + minted = await exchange_scoped_token( + self._storage_api_url, subject_token=subject_token, project_ids=[project_id], read_only=False + ) + scoped_token = minted.access_token + scoped_expires_at = datetime.fromtimestamp(minted.expires_at, tz=timezone.utc) + except Exception as e: + LOG.warning(f'Scoped-token exchange failed while auto-confirming single project: {e}', exc_info=True) + await self._session_store.update_scope( + session_id, + project_ids=[project_id], + read_only=False, + confirmed=True, + scoped_token=scoped_token, + scoped_expires_at=scoped_expires_at, + ) + LOG.info(f'Session {session_id} auto-confirmed to its only accessible project ({project_id}).') + async def load_access_token(self, token: str) -> AccessToken | None: """ Loads an access token by looking up the opaque, randomly-generated token in the Postgres diff --git a/tests/test_cli.py b/tests/test_cli.py index 948bbeaf4..024f99c57 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -231,6 +231,24 @@ async def test_interactive_prompt_scopes_selection(self, creds_file, monkeypatch assert tokens.project_ids == [18, 83] assert tokens.read_only is True + @pytest.mark.asyncio + async def test_interactive_prompt_skips_project_question_with_only_one_project( + self, creds_file, monkeypatch + ) -> None: + # No real choice to make with a single accessible project -- don't ask which project(s), + # just auto-scope to it; still ask read-only. + _seed_unscoped_session() + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + monkeypatch.setattr(auth_login, 'introspect_token', AsyncMock(return_value=_introspection([18]))) + monkeypatch.setattr('sys.stdin.isatty', lambda: True) + with patch('builtins.input', side_effect=['y']) as mocked_input: + await _run_login(STACK) + assert mocked_input.call_count == 1 # only the read-only question, not a project-choice one + + tokens = auth_login.load_tokens(STACK) + assert tokens.project_ids == [18] + assert tokens.read_only is True + @pytest.mark.asyncio async def test_non_interactive_without_scope_raises(self, creds_file, monkeypatch) -> None: _seed_unscoped_session() diff --git a/tests/test_oauth.py b/tests/test_oauth.py index d9dedbea7..83e765348 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -6,6 +6,7 @@ from datetime import datetime, timedelta, timezone from http import HTTPStatus from typing import Any +from unittest import mock from urllib.parse import parse_qs, urlparse import httpx @@ -14,6 +15,7 @@ from mcp.shared.auth import InvalidRedirectUriError, OAuthClientInformationFull from pydantic import AnyHttpUrl, AnyUrl +from keboola_mcp_server.auth_login import Introspection, ProjectAccess, ScopedToken from keboola_mcp_server.clients.auth_bridge import OAuthTokenExchangeError from keboola_mcp_server.oauth import ( ProxyRefreshToken, @@ -26,6 +28,10 @@ JWT_KEY = 'secret' +def _project(project_id: int) -> ProjectAccess: + return ProjectAccess(id=project_id, name=None, role=None) + + class FakeSessionStore: """In-memory `SessionStore` (no real Postgres) for exercising `SimpleOAuthProvider` in isolation.""" @@ -331,15 +337,10 @@ async def test_authorize_redirects_to_consent_with_claudai_projectless_scope( query = parse_qs(parsed.query) assert query['scope'] == ['claudai projectless'] - @pytest.mark.asyncio - async def test_exchange_authorization_code_exchanges_for_session( - self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch - ): + @staticmethod + def _stub_exchanger(monkeypatch: pytest.MonkeyPatch, captured: dict[str, Any]) -> None: from keboola_mcp_server import oauth as oauth_module - monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: '/tmp/sa-token') - captured: dict[str, Any] = {} - class _FakeExchanger: def __init__(self, **kwargs): captured['init_kwargs'] = kwargs @@ -350,6 +351,28 @@ async def exchange(self, *, oauth_access_token: str): monkeypatch.setattr(oauth_module, 'OAuthSessionExchanger', _FakeExchanger) + @pytest.mark.asyncio + async def test_exchange_authorization_code_exchanges_for_session( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + from keboola_mcp_server import oauth as oauth_module + + monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: '/tmp/sa-token') + captured: dict[str, Any] = {} + self._stub_exchanger(monkeypatch, captured) + # Two reachable projects: not the single-project auto-confirm case, so the session should + # stay unconfirmed exactly as before that feature existed. Also proves introspection failures + # here are non-fatal to login -- see test_exchange_authorization_code_introspection_failure_is_non_fatal. + monkeypatch.setattr( + oauth_module, + 'introspect_token', + mock.AsyncMock( + return_value=Introspection( + user_id=1, user_email=None, user_name=None, projects=[_project(1), _project(2)] + ) + ), + ) + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) @@ -370,6 +393,68 @@ async def exchange(self, *, oauth_access_token: str): # forced-relogin window tied to the (1h) Keboola access token's lifetime. assert loaded.expires_at is None assert loaded_refresh.expires_at is None + assert loaded.scope_confirmed is False + + @pytest.mark.asyncio + async def test_exchange_authorization_code_auto_confirms_single_project( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + # No real scoping choice to make with only one reachable project -- see the "Security + # hardening" RFC increment: mirrors the same auto-confirm the local `login` flow does. + from keboola_mcp_server import oauth as oauth_module + + monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: '/tmp/sa-token') + captured: dict[str, Any] = {} + self._stub_exchanger(monkeypatch, captured) + monkeypatch.setattr( + oauth_module, + 'introspect_token', + mock.AsyncMock( + return_value=Introspection(user_id=1, user_email=None, user_name=None, projects=[_project(42)]) + ), + ) + monkeypatch.setattr( + oauth_module, + 'exchange_scoped_token', + mock.AsyncMock( + return_value=ScopedToken( + access_token='kbc_at_scoped', expires_at=time.time() + 3600, project_ids=[42], read_only=False + ) + ), + ) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) + oauth_token = await oauth_provider.exchange_authorization_code(client, auth_code) + + loaded = await oauth_provider.load_access_token(oauth_token.access_token) + assert loaded is not None + assert loaded.scope_confirmed is True + assert loaded.scope_project_ids == [42] + assert loaded.scope_read_only is False + assert loaded.scope_scoped_token == 'kbc_at_scoped' + + @pytest.mark.asyncio + async def test_exchange_authorization_code_introspection_failure_is_non_fatal( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + # Login must still succeed even if the best-effort auto-confirm can't run at all. + from keboola_mcp_server import oauth as oauth_module + + monkeypatch.setattr(oauth_module, 'deployed_sa_token_path', lambda: '/tmp/sa-token') + captured: dict[str, Any] = {} + self._stub_exchanger(monkeypatch, captured) + monkeypatch.setattr( + oauth_module, 'introspect_token', mock.AsyncMock(side_effect=httpx.ConnectError('unreachable')) + ) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) + oauth_token = await oauth_provider.exchange_authorization_code(client, auth_code) + + loaded = await oauth_provider.load_access_token(oauth_token.access_token) + assert loaded is not None + assert loaded.scope_confirmed is False @pytest.mark.asyncio async def test_exchange_authorization_code_maps_exchange_error( From 28dbd7056806324674fac78257f5bbfbcbe332b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 11 Aug 2026 15:55:20 +0200 Subject: [PATCH 82/89] fix(PSGO-261): extend local-login token fallback to streamable-http run_server's local-PKCE-login fallback (use tokens leased by a prior `login`, refreshing as needed, instead of requiring --storage-token) was wired into the stdio branch only. streamable-http/http-compat went straight to uvicorn.Config with just the CLI args, so running the server that way always required an explicit --storage-token even with a valid stored login session. Extract the fallback into _local_login_fallback and run it once, before the transport branch, for both. No-op when a token is already configured (CLI, env, or an existing storage_token) or OAuth is configured (the deployed server case, which authenticates per-session instead of via a local token). Also moves the KBC_* environment override (config.replace_by(os.environ)) earlier so this fallback can see an env-configured OAuth client id before deciding whether to log in. --- src/keboola_mcp_server/cli.py | 54 +++++++++++++++++++++----------- tests/test_cli.py | 59 ++++++++++++++++++++++++++++++++++- 2 files changed, 93 insertions(+), 20 deletions(-) diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 5784c48a8..8c7683d4d 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -261,6 +261,22 @@ def _prompt_mfa_code() -> tuple[str | None, str | None]: return None, recovery +async def _local_login_fallback(config: Config, *, allow_interactive: bool) -> Config: + """Fills in ``config.storage_token`` from the local PKCE `login` credential store when nothing + else has configured a token or an OAuth client -- so a locally-run server (stdio or + streamable-http alike) doesn't need `--storage-token`/`KBC_STORAGE_TOKEN` passed explicitly + once `login` has been run. No-op (returns ``config`` unchanged) when a token is already set, + there's no Storage API URL to log in against, or OAuth is configured (the deployed server + case, which authenticates per-session instead). + """ + if config.storage_token or not config.storage_api_url or config.oauth_client_id or config.oauth_client_secret: + return config + from keboola_mcp_server.auth_login import ensure_access_token + + access_token = await ensure_access_token(config.storage_api_url, allow_interactive=allow_interactive) + return dataclasses.replace(config, storage_token=access_token) + + async def _run_login( api_url: str | None, *, @@ -493,33 +509,33 @@ async def run_server(args: list[str] | None = None) -> None: await _run_gc_sessions() return - # Create config from the CLI arguments + # Create config from the CLI arguments, then apply KBC_* environment overrides up front (not + # just inside create_server, which does this again but too late for the local-login fallback + # below to see an env-configured OAuth client id / storage token). config = Config( storage_api_url=parsed_args.api_url, storage_token=parsed_args.storage_token, workspace_schema=parsed_args.workspace_schema, - ) + ).replace_by(os.environ) + + # Local dev convenience, for stdio and streamable-http alike: with no token configured (CLI, + # env, or OAuth) and a Storage API URL known, use the tokens leased by a prior browser `login` + # (refreshing them as needed) instead of requiring --storage-token/KBC_STORAGE_TOKEN to be + # passed explicitly. No-op for a deployed/OAuth-configured server -- see + # `_local_login_fallback`. + # + # Only run the interactive browser login when a real terminal is attached. For stdio, an MCP + # client launches this process with stdin/stdout as pipes (no TTY) and stdout as the JSON-RPC + # channel -- an interactive login there would corrupt the protocol and block the initialize + # handshake. In that case (and for any non-interactive streamable-http launch, e.g. a + # container) require a prior `login` (or a configured token) and fail fast with guidance + # instead. + allow_interactive = sys.stdin.isatty() and sys.stderr.isatty() + config = await _local_login_fallback(config, allow_interactive=allow_interactive) try: # Create and run the server if parsed_args.transport == 'stdio': - # Local/stdio needs only the stack URL: with no token configured, use the tokens - # leased by a prior browser `login` (refreshing them as needed). When no session is - # stored or it can no longer be refreshed, log in interactively on the spot — so the - # server can be started without a separate `login` step. - config = config.replace_by(os.environ) - if not config.storage_token and config.storage_api_url: - from keboola_mcp_server.auth_login import ensure_access_token - - # Only run the interactive browser login when a real terminal is attached. When an - # MCP client launches this stdio server, stdin/stdout are pipes (no TTY) and stdout - # is the JSON-RPC channel — an interactive login there would corrupt the protocol - # and block the initialize handshake. In that case require a prior `login` (or a - # configured token) and fail fast with guidance instead. - allow_interactive = sys.stdin.isatty() and sys.stderr.isatty() - access_token = await ensure_access_token(config.storage_api_url, allow_interactive=allow_interactive) - config = dataclasses.replace(config, storage_token=access_token) - runtime_config = ServerRuntimeInfo(transport=parsed_args.transport) keboola_mcp_server: FastMCP = create_server(config, runtime_info=runtime_config) if config.oauth_client_id or config.oauth_client_secret: diff --git a/tests/test_cli.py b/tests/test_cli.py index 024f99c57..de1370525 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -5,7 +5,15 @@ from keboola_mcp_server import auth_login from keboola_mcp_server.auth_login import TokenSet -from keboola_mcp_server.cli import _run_gc_sessions, _run_login, _run_logout, _run_migrate, parse_args +from keboola_mcp_server.cli import ( + _local_login_fallback, + _run_gc_sessions, + _run_login, + _run_logout, + _run_migrate, + parse_args, +) +from keboola_mcp_server.config import Config STACK = 'https://connection.keboola.com' @@ -27,6 +35,55 @@ def test_parse_args_gc_sessions() -> None: assert args.command == 'gc-sessions' +class TestLocalLoginFallback: + """Both stdio and streamable-http go through this so a locally-run server picks up a prior + `login`'s stored credentials instead of requiring --storage-token/KBC_STORAGE_TOKEN.""" + + @pytest.mark.asyncio + async def test_fills_in_token_from_login_store(self, monkeypatch) -> None: + monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) + config = Config(storage_api_url=STACK) + + result = await _local_login_fallback(config, allow_interactive=False) + + assert result.storage_token == 'kbc_at_x' + auth_login.ensure_access_token.assert_awaited_once_with(STACK, allow_interactive=False) + + @pytest.mark.asyncio + async def test_noop_when_token_already_set(self, monkeypatch) -> None: + ensure = AsyncMock() + monkeypatch.setattr(auth_login, 'ensure_access_token', ensure) + config = Config(storage_api_url=STACK, storage_token='kbc_at_already_set') + + result = await _local_login_fallback(config, allow_interactive=False) + + assert result is config + ensure.assert_not_awaited() + + @pytest.mark.asyncio + async def test_noop_without_storage_api_url(self, monkeypatch) -> None: + ensure = AsyncMock() + monkeypatch.setattr(auth_login, 'ensure_access_token', ensure) + config = Config() + + result = await _local_login_fallback(config, allow_interactive=False) + + assert result is config + ensure.assert_not_awaited() + + @pytest.mark.asyncio + async def test_noop_when_oauth_configured(self, monkeypatch) -> None: + # Deployed server: authenticates per-session via OAuth, not a locally stored token. + ensure = AsyncMock() + monkeypatch.setattr(auth_login, 'ensure_access_token', ensure) + config = Config(storage_api_url=STACK, oauth_client_id='id', oauth_client_secret='secret') + + result = await _local_login_fallback(config, allow_interactive=False) + + assert result is config + ensure.assert_not_awaited() + + class TestRunMigrate: @pytest.mark.asyncio async def test_requires_postgres_dsn(self, monkeypatch) -> None: From 6e1b70e31c2538f3e0e638dc88adc32028c12073 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Tue, 11 Aug 2026 17:17:49 +0200 Subject: [PATCH 83/89] fix(PSGO-261): stop instructing the LLM to re-scope an already-scoped session MultiProjectMiddleware's ask-first gate only fires for an unconfirmed scope, and a login-time-persisted (or single-project-auto-confirmed) scope is already confirmed=True -- tools were never actually blocked for these sessions. The redundant get_accessible_projects/set_project_scope calls seen in practice came from the server's static instructions string, which unconditionally told the LLM to call both "at the very START of the conversation, before doing anything else" regardless of whether a scope already existed -- for an already-scoped session this just means a wasted per-accessible-project verify fan-out for no reason. Reword the instructions and get_accessible_projects' docstring to be reactive: try the data tool you actually need first; only run the scoping dance if a tool call fails asking you to confirm scope. No enforcement change -- the gate's behavior was already correct, only the LLM-facing guidance was stale. --- TOOLS.md | 15 ++++++----- src/keboola_mcp_server/server.py | 33 ++++++++++++++----------- src/keboola_mcp_server/tools/project.py | 15 ++++++----- 3 files changed, 36 insertions(+), 27 deletions(-) diff --git a/TOOLS.md b/TOOLS.md index 2261d010b..7aa01c03a 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -3770,12 +3770,15 @@ configuration is created e.g. keboola.ex-google-analytics-v4 and keboola.ex-gmai Lists the Keboola projects the current login can access across the stack, each with its SQL dialect and organization. -Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), -present the projects, and ask whether they want to work across all of them or a subset. Then call -`set_project_scope` with their choice. This tool compacts several API calls (token introspection -plus a per-project token verify for the SQL dialect and organization) into one result, so the -assistant does not need a separate get_project_info call per project. Pass with_llm_instruction=true -on the first call to also receive the base working instructions grouped by dialect. +Only call this when a data tool call has actually failed asking you to confirm a project scope -- +the session may already be pre-scoped (e.g. the user chose specific projects at `login` time), in +which case data tools already work and this call would just be extra, unnecessary API traffic. +When a scope genuinely is needed: present the projects, ask whether the user wants to work across +all of them or a subset, then call `set_project_scope` with their choice. This tool compacts +several API calls (token introspection plus a per-project token verify for the SQL dialect and +organization) into one result, so the assistant does not need a separate get_project_info call per +project. Pass with_llm_instruction=true on the first call to also receive the base working +instructions grouped by dialect. **Input JSON Schema**: diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py index de1c20740..2d9c9157d 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -259,21 +259,24 @@ def create_server( name='Keboola MCP Server', instructions=( 'This server supports multi-project mode for stack-wide Keboola programmatic tokens ' - '(kbc_at_/kbc_pat_). When the session uses such a token, data tools are BLOCKED until a ' - 'project scope is confirmed. So at the very START of the conversation, before doing anything ' - 'else: call "get_accessible_projects", show the user their projects, and ASK whether to work ' - 'across ALL of them or a subset. Do not decide for them. Then call "set_project_scope" with ' - 'their answer (no arguments = all projects, or the chosen project ids, optionally ' - 'read_only=true). Both tools return a "scope_token" -- the server does not remember the ' - 'scope between calls, so resend that value as the "scope_token" argument on every ' - 'subsequent tool call in this conversation. After that, read-only tools return results per ' - 'project. Never write to more than one project without explicit user confirmation — write ' - 'operations target the active (first-scoped) project only. If instead the session uses a ' - 'legacy project-scoped Storage API token, it is already bound to a single project: use the ' - 'tools directly — "get_accessible_projects" / "set_project_scope" do not apply (they will ' - 'report that no programmatic token is present). Note: outside the Storage API, some tools ' - 'may need per-project token support not yet available on every stack; surface such errors ' - 'plainly rather than retrying.' + '(kbc_at_/kbc_pat_). A session may already be pre-scoped -- e.g. the user chose specific ' + 'projects at `login` time, or the token can only reach one project -- in which case data ' + 'tools work immediately with no further action needed. Do NOT call "get_accessible_projects" ' + 'or "set_project_scope" preemptively "just in case": just call the data tool you actually ' + 'need. Only if a data tool call fails with an error asking you to confirm a project scope ' + '(this happens when the session truly has none yet): call "get_accessible_projects", show ' + 'the user their projects, and ASK whether to work across ALL of them or a subset. Do not ' + 'decide for them. Then call "set_project_scope" with their answer (no arguments = all ' + 'projects, or the chosen project ids, optionally read_only=true). Both tools return a ' + '"scope_token" -- the server does not remember the scope between calls, so resend that value ' + 'as the "scope_token" argument on every subsequent tool call in this conversation. After ' + 'that, read-only tools return results per project. Never write to more than one project ' + 'without explicit user confirmation — write operations target the active (first-scoped) ' + 'project only. If instead the session uses a legacy project-scoped Storage API token, it is ' + 'already bound to a single project: use the tools directly — "get_accessible_projects" / ' + '"set_project_scope" do not apply (they will report that no programmatic token is present). ' + 'Note: outside the Storage API, some tools may need per-project token support not yet ' + 'available on every stack; surface such errors plainly rather than retrying.' ), lifespan=create_keboola_lifespan(server_state), auth=oauth_provider, diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index d818bbd92..297e40882 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -442,12 +442,15 @@ async def get_accessible_projects( Lists the Keboola projects the current login can access across the stack, each with its SQL dialect and organization. - Call this early in a conversation when the user logs in with a stack-wide token (PKCE login), - present the projects, and ask whether they want to work across all of them or a subset. Then call - `set_project_scope` with their choice. This tool compacts several API calls (token introspection - plus a per-project token verify for the SQL dialect and organization) into one result, so the - assistant does not need a separate get_project_info call per project. Pass with_llm_instruction=true - on the first call to also receive the base working instructions grouped by dialect. + Only call this when a data tool call has actually failed asking you to confirm a project scope -- + the session may already be pre-scoped (e.g. the user chose specific projects at `login` time), in + which case data tools already work and this call would just be extra, unnecessary API traffic. + When a scope genuinely is needed: present the projects, ask whether the user wants to work across + all of them or a subset, then call `set_project_scope` with their choice. This tool compacts + several API calls (token introspection plus a per-project token verify for the SQL dialect and + organization) into one result, so the assistant does not need a separate get_project_info call per + project. Pass with_llm_instruction=true on the first call to also receive the base working + instructions grouped by dialect. """ client = KeboolaClient.from_state(ctx.session.state) subject_token = await _parent_subject_token(client) From 4c6264f8da71f8f30d954c8912ac887f5ff90178 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 12 Aug 2026 09:25:22 +0200 Subject: [PATCH 84/89] fix(PSGO-261): stop X-KBC-ProjectId from overriding a confirmed scope project_id is a header-eligible Config field. _resolve_local_tokens's deployed/OAuth branch only applied a confirmed scope's active project id when config.project_id wasn't already set, so a request carrying X-KBC-ProjectId kept that header's value even after set_project_scope confirmed a narrower scope. MultiProjectMiddleware's active-project fast paths compare only the logical target against scope.active_project_id, never what project the base client was actually built with, so the mismatched client went unnoticed -- letting any caller able to attach one header redirect every default-target call to a project outside the confirmed scope, using the full unscoped token. Drop the `not config.project_id` guard so a confirmed scope's active project always wins, matching the local-programmatic branch's existing (unaffected, unconditional) behavior. A tool wanting a different scoped project still has its own project_id argument, validated separately by MultiProjectMiddleware._dispatch_single_target. Found by a full-PR security audit (thermo-nuclear/ponytail/security-scanner across the whole diff, at the user's request), not the original review. --- feature_spec/pat_token_support/RFC.md | 25 +++++++++++++++++++++++++ pyproject.toml | 2 +- src/keboola_mcp_server/mcp.py | 9 ++++++++- tests/test_mcp.py | 16 ++++++++++++++-- uv.lock | 2 +- 5 files changed, 49 insertions(+), 5 deletions(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index dcf3ef284..f10012fcb 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -1004,3 +1004,28 @@ not evidence of a prior scoping choice, so it isn't auto-confirmed. - `lease_pat`/`login --pat` need no separate change -- they already take an explicit `project_ids` argument (increment 7), which now flows from the auto-detected single project when applicable. + +## Extension: `X-KBC-ProjectId` could override a confirmed scope (increment 12) + +Found by a full-PR security audit, not from the original review: `project_id` is a header-eligible +`Config` field (`_HEADER_ELIGIBLE_FIELDS`), and `_resolve_local_tokens`'s deployed/OAuth branch +only applied a confirmed scope's active project id when `config.project_id` wasn't already set +(`if scope and scope.project_ids and not config.project_id: ...`). A request carrying +`X-KBC-ProjectId` therefore kept that header's value even after `set_project_scope` confirmed a +different, narrower scope -- and `MultiProjectMiddleware`'s active-project fast paths only compare +the *logical* target against `scope.active_project_id`, never inspect what project the base +client was actually built with, so the mismatched client was used unnoticed. Net effect: any +caller able to attach one header could redirect every default-target call to a project outside +what the user confirmed, using the full unscoped token -- defeating the scoping guarantee for +OAuth and Kai/header-token sessions (the local-programmatic branch was never affected -- it +already unconditionally overwrote `project_id` from the scope, no guard). + +**Fix:** drop the `not config.project_id` guard -- once a confirmed multi-project scope exists, +`project_id` always comes from `scope.active_project_id`, matching the local branch's existing +(safe) behavior. A tool wanting a *different* scoped project still has its own `project_id` +argument, validated against `scope.project_ids` by `MultiProjectMiddleware._dispatch_single_target` +-- this only affects which project the un-swapped base client targets. Considered and rejected: an +additional check on `MultiProjectMiddleware`'s side comparing the base client's actual +`X-KBC-ProjectId` header against the scope (defense-in-depth) -- redundant once the root cause is +fixed at the source, and it broke existing mock-based tests for no real security benefit; the +mcp.py-side fix alone closes the gap. diff --git a/pyproject.toml b/pyproject.toml index 22548bf12..31be328df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.76.0" +version = "1.76.1" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index a1c7f0e19..53402d6a3 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -670,7 +670,14 @@ async def _resolve_local_tokens( every request). """ if not cls._is_local_programmatic(config): - if scope and scope.project_ids and not config.project_id: + if scope and scope.project_ids: + # Always the scope's own active project, never a caller-supplied X-KBC-ProjectId -- + # project_id is header-eligible (Config._HEADER_ELIGIBLE_FIELDS), and once a scope is + # confirmed the header must not be able to silently redirect the base client to a + # project outside (or merely different from) what the user confirmed. A tool wanting a + # *different* one of the scoped projects still has its own project_id argument + # (MultiProjectMiddleware._dispatch_single_target), validated against scope.project_ids + # there -- this is only about which project the un-swapped base client targets. config = dataclasses.replace(config, project_id=str(scope.active_project_id)) if scope is not None and scope.scoped_token is not None and scope.is_near_expiry: try: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 70fefad50..6b86dbc05 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1305,12 +1305,24 @@ async def test_deployed_with_confirmed_scope_applies_active_project_id(self, mon assert out_scope is scope # untouched @pytest.mark.asyncio - async def test_deployed_no_scope_or_already_set_project_id_is_noop(self, monkeypatch) -> None: + async def test_deployed_no_scope_is_noop(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x', project_id='7') + out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, None) + assert out_config is config + assert out_scope is None + + @pytest.mark.asyncio + async def test_deployed_confirmed_scope_overrides_a_header_supplied_project_id(self, monkeypatch) -> None: + # Regression: project_id is header-eligible (X-KBC-ProjectId), so a caller could set + # config.project_id before scope resolution runs. A confirmed scope's active project must + # always win -- otherwise a session scoped to project 18 could be silently redirected to + # whatever project a request header names, via the base (un-swapped) client. monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x', project_id='7') scope = SessionScope(project_ids=[18], confirmed=True) out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, scope) - assert out_config is config # project_id already set -- not overwritten + assert out_config.project_id == '18' assert out_scope is scope @pytest.mark.asyncio diff --git a/uv.lock b/uv.lock index 4561e9f57..a31dbfa04 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.76.0" +version = "1.76.1" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 6fa5844af0d7efa14a56146eaa860fe31f552d73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 12 Aug 2026 09:29:53 +0200 Subject: [PATCH 85/89] test(PSGO-261): drop a redundant no-scope test added in c0d8fe02 test_deployed_no_scope_is_noop duplicated the pre-existing test_deployed_is_noop -- with scope=None neither test enters the branch this fix touches, so the project_id='7' it set was never read. The actual regression test for the fix is test_deployed_confirmed_scope_overrides_a_header_supplied_project_id. Found by /simplify (4 parallel reuse/simplification/efficiency/altitude reviews on commit c0d8fe02); the other three lenses found nothing to fix. --- tests/test_mcp.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 6b86dbc05..c3036e6e7 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1304,14 +1304,6 @@ async def test_deployed_with_confirmed_scope_applies_active_project_id(self, mon assert out_config.storage_token == 'kbc_at_x' # untouched; resolver-exchange narrows it assert out_scope is scope # untouched - @pytest.mark.asyncio - async def test_deployed_no_scope_is_noop(self, monkeypatch) -> None: - monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') - config = Config(storage_api_url='https://connection.keboola.com', storage_token='kbc_at_x', project_id='7') - out_config, out_scope = await SessionStateMiddleware._resolve_local_tokens(config, None) - assert out_config is config - assert out_scope is None - @pytest.mark.asyncio async def test_deployed_confirmed_scope_overrides_a_header_supplied_project_id(self, monkeypatch) -> None: # Regression: project_id is header-eligible (X-KBC-ProjectId), so a caller could set From f50a90b5d7631765328df3c84a4ba61a44e34064 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 12 Aug 2026 10:05:21 +0200 Subject: [PATCH 86/89] fix(PSGO-261): stop the login fallback from crashing streamable-http with no token _local_login_fallback made any transport attempt ensure_access_token whenever no token/OAuth is configured, not just stdio. streamable-http/ http-compat legitimately run with no default token, relying entirely on a per-request header -- with allow_interactive=False (no TTY in CI) and no stored local credential, ensure_access_token raised RuntimeError, uncaught, killing the server subprocess before it could start listening. Caught by CI: integtests/test_mcp_server.py::test_remote_setup and test_http_multiple_clients deliberately start streamable-http with no token, and started failing with "No stored credentials... Run login first" once the fallback got extended to that transport. _local_login_fallback gains a `required` param: stdio passes True (unchanged -- no other token source exists there), streamable-http/ http-compat pass False (catch and log instead of raising; the server starts normally and expects a token per request). --- feature_spec/pat_token_support/RFC.md | 17 +++++++++++++ pyproject.toml | 2 +- src/keboola_mcp_server/cli.py | 23 ++++++++++++++--- tests/test_cli.py | 36 ++++++++++++++++++++++++--- uv.lock | 2 +- 5 files changed, 71 insertions(+), 9 deletions(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index f10012fcb..67b9e8a90 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -1029,3 +1029,20 @@ additional check on `MultiProjectMiddleware`'s side comparing the base client's `X-KBC-ProjectId` header against the scope (defense-in-depth) -- redundant once the root cause is fixed at the source, and it broke existing mock-based tests for no real security benefit; the mcp.py-side fix alone closes the gap. + +## Fix: `_local_login_fallback` broke streamable-http with no configured token (increment 13) + +CI regression, caught by the integration-test suite (`integtests/test_mcp_server.py::test_remote_setup`, +`test_http_multiple_clients`): the increment-that-extended-`_local_login_fallback`-to-streamable-http +(RFC increment referenced above) made *any* transport attempt `ensure_access_token` whenever no +token/OAuth is configured, not just `stdio`. But `streamable-http`/`http-compat` legitimately run +with no default token at all, relying entirely on a per-request header (`X-Storage-Token`) -- +exactly what these integration tests deliberately exercise. With `allow_interactive=False` (no TTY +in CI) and no stored local-login credential, `ensure_access_token` raised `RuntimeError: No stored +credentials...`, uncaught, killing the server subprocess before it could even start listening. + +**Fix:** `_local_login_fallback` gains a `required: bool` parameter. `stdio` passes `True` (no +other token source exists there, so a missing credential must still fail startup with the "run +login" guidance -- unchanged behavior). `streamable-http`/`http-compat` pass `False`: a missing +local credential there is caught and logged, not raised -- `config` is returned unchanged and the +server starts normally, expecting a token per request. diff --git a/pyproject.toml b/pyproject.toml index 31be328df..d4d296a59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.76.1" +version = "1.76.2" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 8c7683d4d..31a72518f 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -261,19 +261,34 @@ def _prompt_mfa_code() -> tuple[str | None, str | None]: return None, recovery -async def _local_login_fallback(config: Config, *, allow_interactive: bool) -> Config: +async def _local_login_fallback(config: Config, *, allow_interactive: bool, required: bool) -> Config: """Fills in ``config.storage_token`` from the local PKCE `login` credential store when nothing else has configured a token or an OAuth client -- so a locally-run server (stdio or streamable-http alike) doesn't need `--storage-token`/`KBC_STORAGE_TOKEN` passed explicitly once `login` has been run. No-op (returns ``config`` unchanged) when a token is already set, there's no Storage API URL to log in against, or OAuth is configured (the deployed server case, which authenticates per-session instead). + + :param required: stdio has no other way to get a token (no per-request headers), so a missing + credential there must fail server startup with the "run login" guidance -- ``True`` + propagates that. streamable-http/http-compat can still get a token per request via a + header, so a missing local credential there is a legitimate, unconfigured-on-purpose state, + not an error -- ``False`` logs and leaves ``config`` unchanged instead of crashing startup. """ if config.storage_token or not config.storage_api_url or config.oauth_client_id or config.oauth_client_secret: return config from keboola_mcp_server.auth_login import ensure_access_token - access_token = await ensure_access_token(config.storage_api_url, allow_interactive=allow_interactive) + try: + access_token = await ensure_access_token(config.storage_api_url, allow_interactive=allow_interactive) + except RuntimeError: + if required: + raise + LOG.info( + f'No local login session for {config.storage_api_url} and none required for this transport -- ' + 'starting without a default token; callers must supply one per request.' + ) + return config return dataclasses.replace(config, storage_token=access_token) @@ -531,7 +546,9 @@ async def run_server(args: list[str] | None = None) -> None: # container) require a prior `login` (or a configured token) and fail fast with guidance # instead. allow_interactive = sys.stdin.isatty() and sys.stderr.isatty() - config = await _local_login_fallback(config, allow_interactive=allow_interactive) + config = await _local_login_fallback( + config, allow_interactive=allow_interactive, required=parsed_args.transport == 'stdio' + ) try: # Create and run the server diff --git a/tests/test_cli.py b/tests/test_cli.py index de1370525..c65e5e02b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -44,7 +44,7 @@ async def test_fills_in_token_from_login_store(self, monkeypatch) -> None: monkeypatch.setattr(auth_login, 'ensure_access_token', AsyncMock(return_value='kbc_at_x')) config = Config(storage_api_url=STACK) - result = await _local_login_fallback(config, allow_interactive=False) + result = await _local_login_fallback(config, allow_interactive=False, required=True) assert result.storage_token == 'kbc_at_x' auth_login.ensure_access_token.assert_awaited_once_with(STACK, allow_interactive=False) @@ -55,7 +55,7 @@ async def test_noop_when_token_already_set(self, monkeypatch) -> None: monkeypatch.setattr(auth_login, 'ensure_access_token', ensure) config = Config(storage_api_url=STACK, storage_token='kbc_at_already_set') - result = await _local_login_fallback(config, allow_interactive=False) + result = await _local_login_fallback(config, allow_interactive=False, required=True) assert result is config ensure.assert_not_awaited() @@ -66,7 +66,7 @@ async def test_noop_without_storage_api_url(self, monkeypatch) -> None: monkeypatch.setattr(auth_login, 'ensure_access_token', ensure) config = Config() - result = await _local_login_fallback(config, allow_interactive=False) + result = await _local_login_fallback(config, allow_interactive=False, required=True) assert result is config ensure.assert_not_awaited() @@ -78,11 +78,39 @@ async def test_noop_when_oauth_configured(self, monkeypatch) -> None: monkeypatch.setattr(auth_login, 'ensure_access_token', ensure) config = Config(storage_api_url=STACK, oauth_client_id='id', oauth_client_secret='secret') - result = await _local_login_fallback(config, allow_interactive=False) + result = await _local_login_fallback(config, allow_interactive=False, required=True) assert result is config ensure.assert_not_awaited() + @pytest.mark.asyncio + async def test_required_raises_when_no_stored_session(self, monkeypatch) -> None: + # stdio has no other way to get a token (no per-request headers) -- a missing local + # credential there must fail server startup with the "run login" guidance. + monkeypatch.setattr( + auth_login, 'ensure_access_token', AsyncMock(side_effect=RuntimeError('no stored credentials')) + ) + config = Config(storage_api_url=STACK) + + with pytest.raises(RuntimeError, match='no stored credentials'): + await _local_login_fallback(config, allow_interactive=False, required=True) + + @pytest.mark.asyncio + async def test_not_required_starts_without_a_token_when_no_stored_session(self, monkeypatch) -> None: + # streamable-http/http-compat can still get a token per request via a header -- a missing + # local credential there is a legitimate, unconfigured-on-purpose state, not a startup error + # (regression: this used to crash the server subprocess before it could even start + # listening, e.g. in integtests that deliberately run streamable-http with no token at all). + monkeypatch.setattr( + auth_login, 'ensure_access_token', AsyncMock(side_effect=RuntimeError('no stored credentials')) + ) + config = Config(storage_api_url=STACK) + + result = await _local_login_fallback(config, allow_interactive=False, required=False) + + assert result is config + assert result.storage_token is None + class TestRunMigrate: @pytest.mark.asyncio diff --git a/uv.lock b/uv.lock index a31dbfa04..1d8bbc146 100644 --- a/uv.lock +++ b/uv.lock @@ -1240,7 +1240,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.76.1" +version = "1.76.2" source = { editable = "." } dependencies = [ { name = "asyncpg" }, From 13d2c8ede5c7f642bfa578ef0dced1561425cdfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Wed, 12 Aug 2026 10:14:15 +0200 Subject: [PATCH 87/89] docs(PSGO-261): restructure RFC around architecture, not a fix changelog The RFC had grown into a chronological log of 13 "increment" postscripts, several of them narrow post-review bug-fix writeups (single-project auto-scope, an X-KBC-ProjectId override, a streamable-http startup crash) mixed in with the actual design decisions. Restructured around what the system does today -- token paths, multi-project scope, scope persistence, security invariants, decisions, testing -- dropping resolved-question transcripts and reverted-approach narration that no longer help a reader understand the current design. 1046 -> 233 lines; no behavior change. Also imports feature_spec/mpa_support/mpa-plan.md from the closed, unmerged PR #451 (feature/mpa-support) into main, where it previously only existed on that PR's branch -- with an annotated comparison against this RFC's design (config shape, where "which project(s)" lives, OAuth support) so the relationship is checkable in-repo instead of only in a closed PR's diff. --- feature_spec/mpa_support/mpa-plan.md | 348 +++++++++++++++++++++++++++ 1 file changed, 348 insertions(+) create mode 100644 feature_spec/mpa_support/mpa-plan.md diff --git a/feature_spec/mpa_support/mpa-plan.md b/feature_spec/mpa_support/mpa-plan.md new file mode 100644 index 000000000..b49d96204 --- /dev/null +++ b/feature_spec/mpa_support/mpa-plan.md @@ -0,0 +1,348 @@ +# Multi-Project Architecture (MPA) - Implementation Plan + +> **Status: superseded, kept for historical reference.** This plan originated in PR #451 +> (`davidesner`, branch `feature/mpa-support`), which was closed without merging. It is imported +> into `main` here (it previously existed only on that closed PR's branch) so the design comparison +> below is checkable in-repo rather than only in a closed PR's diff. +> +> The multi-project problem this plan addresses was solved differently by PSGO-261 +> (`feature_spec/pat_token_support/RFC.md`, see its "Relationship to prior multi-project attempts" +> section): a single stack-wide programmatic token + server-driven discovery +> (`get_accessible_projects`) + transparent per-project fan-out, instead of this plan's static +> per-project tokens in `mcp.json` + a middleware-injected `project_id`/`branch_id` parameter + the +> **agent** calling once per project. Two differences worth calling out explicitly: +> +> - **Config shape**: this plan's `ProjectConfig`/`projects: tuple[...]` static list (one Storage +> token per project, configured ahead of time) has no equivalent in the shipped design — a single +> whole-stack `kbc_at_*`/`kbc_pat_*` token is introspected at runtime to discover reachable +> projects instead, so no per-project config entries ever need to be authored or kept in sync. +> - **Where "which project(s)" lives**: this plan puts `project_id`/`branch_id` on every tool call +> (middleware-injected, agent-driven, once per project for an N-project operation). The shipped +> design puts it on session **scope** instead (`set_project_scope`, confirmed once, reused for +> every subsequent call) for reads, and keeps only writes as an explicit per-call `project_id` +> argument — see PSGO-261 RFC "Decisions" #3/#4 for why each mechanism was chosen where it was. +> +> OAuth+MPA (this plan's stated "not supported yet" limitation) is not a gap in the shipped design: +> OAuth sessions get full multi-project scope support there. + +## Context + +Users of the Keboola MCP server typically have access to multiple projects within an organization. Currently, the server supports only a single project per session, requiring users to re-login to switch projects. This plan adds multi-project support so that all tools can operate across projects via a `project_id` parameter, with per-project branch management and a CLI init flow that creates Storage tokens from a Manage token. + +## Key Design Decisions + +### Middleware-Based Project + Branch Resolution + +Rather than modifying all 31 tool function signatures to add `project_id` and `branch_id`, we use a **middleware approach**: + +1. A `ProjectResolutionMiddleware` dynamically injects optional `project_id` and `branch_id` parameters into every tool's JSON schema during `on_list_tools` +2. On `on_call_tool`, the middleware extracts `project_id` and `branch_id` from arguments, resolves the correct `KeboolaClient` + `WorkspaceManager`, and places them in `ctx.session.state` under the same legacy keys +3. **Zero changes to existing tool functions** — they still call `KeboolaClient.from_state(ctx.session.state)` as before +4. Full backward compatibility: single-project mode works identically (no extra params injected) + +### Stateless Branch Handling + +The server is stateless (no persistent DB). Branch selection is per-tool-call: +- `branch_id` is an optional middleware-injected parameter on every tool call +- Default branch comes from the per-project config in `mcp.json` (or main if not set) +- The `create_branch` and `list_branches` tools let users discover/create branches +- No session state or "switch branch" concept needed + +### Param Visibility Rules (Backward Compatible) + +Parameters are only injected when they provide value — if a value is fixed in config, the param is hidden: + +| Config scenario | project_id visible | branch_id visible | +|---|---|---| +| Legacy (env vars, no config file) | No | No | +| 1 project, branch_id set in config | No | No | +| 1 project, no branch_id in config | No | Yes | +| 2+ projects, all with branch_id | Yes | No | +| 2+ projects, some/all without branch_id | Yes | Yes | + +This means: +- **Legacy mode is 100% unchanged** — no config file = no extra params +- **Single-project with fixed branch** = behaves exactly like legacy +- **Agent can only change what the config allows** — if branch_id is in config, it's locked + +### OAuth Compatibility + +OAuth mode currently provides a single project token. MPA is not supported with OAuth yet — this is documented as a known limitation. No code changes for OAuth+MPA in this iteration. + +--- + +## Phase 1: Foundation — Config & Multi-Project State + +### 1.1 Extend Config (`src/keboola_mcp_server/config.py`) + +Add `ProjectConfig` dataclass: +```python +@dataclass(frozen=True) +class ProjectConfig: + project_id: str + storage_api_url: str + storage_token: str + branch_id: Optional[str] = None + workspace_schema: Optional[str] = None + alias: Optional[str] = None + forbid_main_branch_writes: bool = False +``` + +Add to existing `Config`: +- `projects: tuple[ProjectConfig, ...] = ()` +- `default_project_id: Optional[str] = None` +- `forbid_main_branch_writes: bool = False` (global default) +- `is_mpa_mode` property: `return len(self.projects) > 0` +- `from_config_file(path: Path) -> Config` classmethod to load `mcp.json` + +### 1.2 Create ProjectRegistry (`src/keboola_mcp_server/project_registry.py` — new file) + +```python +@dataclass +class ProjectContext: + project_id: str + client: KeboolaClient + workspace_manager: WorkspaceManager + alias: str | None + forbid_main_branch_writes: bool + +class ProjectRegistry: + STATE_KEY = 'project_registry' + projects: dict[str, ProjectContext] # keyed by project_id + default_project_id: str | None + + def get_project(self, project_id: str | None) -> ProjectContext + def list_projects(self) -> list[ProjectContext] + def inject_into_state(self, state: dict, project_id: str | None) -> None + @classmethod + def from_state(cls, state) -> ProjectRegistry +``` + +### 1.3 Update SessionStateMiddleware (`src/keboola_mcp_server/mcp.py`) + +In `create_session_state`, when `config.is_mpa_mode`: +- Create `KeboolaClient` + `WorkspaceManager` for each project (concurrently via `asyncio.gather`) +- Build `ProjectRegistry`, store in state +- Also inject default project's client/workspace under legacy keys (for middleware that runs before project resolution, e.g. `ToolsFilteringMiddleware`) + +--- + +## Phase 2: Project Resolution Middleware + +### 2.1 Create ProjectResolutionMiddleware (`src/keboola_mcp_server/mcp.py`) + +`on_list_tools`: +- If 2+ projects in config: inject optional `project_id` parameter into every tool's JSON schema (description lists available project IDs/aliases) +- If any project has no `branch_id` fixed in config: inject optional `branch_id` parameter into every tool's JSON schema +- Skip injection for project-agnostic tools (`docs_query`) +- In legacy mode (no config file): no-op + +`on_call_tool`: +- Extract and pop `project_id` from `context.message.arguments` (if present) +- Extract and pop `branch_id` from `context.message.arguments` (if present) +- Resolve `ProjectContext` from registry (use default/only project if `project_id` not specified; error if ambiguous with 2+ projects) +- Determine effective branch: explicit `branch_id` arg > project's config `branch_id` > main (None) +- If effective branch differs from project's default, call `client.with_branch_id(branch_id)` to get a branch-specific client +- Inject resolved client + workspace_manager into `ctx.session.state` under legacy keys +- Check `forbid_main_branch_writes`: if tool is a write op, effective branch is main, and setting is True → raise `ToolError` +- In legacy mode (no config file): no-op (session state already set by SessionStateMiddleware as today) + +### 2.2 Register Middleware (`src/keboola_mcp_server/server.py`) + +Middleware chain order: +```python +middleware=[ + SessionStateMiddleware(), + ProjectResolutionMiddleware(), # NEW — after session, before auth + ToolAuthorizationMiddleware(), + ToolsFilteringMiddleware(), + ValidationErrorMiddleware(), +] +``` + +--- + +## Phase 3: Branch Tools & Write Protection + +### 3.1 Add `dev_branch_create` to Storage Client (`src/keboola_mcp_server/clients/storage.py`) + +```python +async def dev_branch_create(self, name: str, description: str = '') -> JsonDict: + return await self.post(endpoint='dev-branches', data={'name': name, 'description': description}) +``` + +### 3.2 Create Branch Tools (`src/keboola_mcp_server/tools/branches.py` — new file) + +**`list_branches`** (readOnlyHint=True): +- Calls `client.storage_client.branches_list()` +- Returns list of branches with id, name, isDefault, created, description + +**`create_branch`** (destructiveHint=False): +- Parameters: `name: str`, `description: str = ''` +- Calls `client.storage_client.dev_branch_create(name, description)` +- Returns created branch info + +Register via `add_branch_tools(mcp)` in `server.py`. + +### 3.3 Main Branch Write Protection + +In `ProjectResolutionMiddleware.on_call_tool`: +- After project resolution, if the tool is not read-only AND the client is on main branch (branch_id is None) AND `forbid_main_branch_writes` is True for this project (or globally): + - Raise `ToolError`: "Write operations on the main branch are forbidden. Create a development branch first using `create_branch`, then specify the branch when calling tools." + +--- + +## Phase 4: get_project_info Changes + +### 4.1 Update `get_project_info` (`src/keboola_mcp_server/tools/project.py`) + +In MPA mode when no specific `project_id` is given (or a special "all" mode), return a new `MultiProjectInfo` model: + +```python +class MultiProjectInfo(BaseModel): + projects: list[ProjectInfo] # per-project info (without llm_instruction) + llm_instruction: str # shared, returned once +``` + +Each project entry includes: `project_id`, `project_name`, `project_description`, `organization_id`, `sql_dialect`, `conditional_flows`, `links`, `user_role`, `toolset_restrictions`. + +In single-project mode, behavior is unchanged (returns `ProjectInfo` as today). + +--- + +## Phase 5: CLI Init Command + +### 5.1 Add ManageClient (`src/keboola_mcp_server/clients/manage.py` — new file) + +Async HTTP client for Manage API (based on reference CLI pattern): +- Auth header: `X-KBC-ManageApiToken: ` +- `verify_token()` → `GET /manage/tokens/verify` +- `get_project(project_id)` → `GET /manage/projects/{project_id}` +- `list_organization_projects(org_id)` → `GET /manage/organizations/{org_id}/projects` +- `create_project_token(project_id, description, ...)` → `POST /manage/projects/{project_id}/tokens` + +Token creation payload (matching reference CLI): +```json +{ + "description": "keboola-mcp-server", + "canManageBuckets": true, + "canReadAllFileUploads": true, + "canReadAllProjectEvents": true, + "canManageDevBranches": true, + "canManageTokens": true +} +``` + +### 5.2 Add `init` CLI Command (`src/keboola_mcp_server/cli.py`) + +New `init` subcommand added to argparse: +``` +python -m keboola_mcp_server init \ + --manage-token \ + --api-url https://connection.north-europe.azure.keboola.com \ + [--project-ids 12345,67890] \ + [--all] \ + --output mcp.json \ + [--forbid-main-branch-writes] +``` + +Flow: +1. Verify manage token → `GET /manage/tokens/verify` +2. Get org from token info, list all projects in org +3. Project selection (three modes): + - `--project-ids 12345,67890`: Use specific projects (non-interactive) + - `--all`: Add all projects in the organization + - Neither flag: Interactive prompt listing available projects for user selection +3. For each selected project, create Storage API token via manage API +4. Write `mcp.json` with format: +```json +{ + "version": 1, + "default_project_id": "12345", + "forbid_main_branch_writes": false, + "projects": [ + { + "project_id": "12345", + "alias": "my-project", + "storage_api_url": "https://connection.north-europe.azure.keboola.com", + "token": "" + } + ] +} +``` +5. **Manage token is NOT stored** in the config file + +### 5.3 Add `--config-file` to `run_server` (`src/keboola_mcp_server/cli.py`) + +``` +python -m keboola_mcp_server --transport stdio --config-file mcp.json +``` + +When `--config-file` is provided, load `Config.from_config_file(path)` instead of using CLI args for token/URL. OAuth and other server settings can still come from env vars. + +--- + +## Phase 6: ToolsFilteringMiddleware Updates + +### 6.1 Adapt for MPA (`src/keboola_mcp_server/mcp.py`) + +`ToolsFilteringMiddleware` currently calls `verify_token()` to get project features and token role. In MPA mode: +- `on_list_tools`: Use the default project's client (already injected by SessionStateMiddleware under legacy keys) +- `on_call_tool`: By this point, `ProjectResolutionMiddleware` has already injected the correct project's client, so `ToolsFilteringMiddleware` works without changes + +Consider caching `verify_token()` results in `ProjectContext` during session creation to avoid repeated API calls. + +--- + +## Files Summary + +### New Files +| File | Purpose | +|------|---------| +| `src/keboola_mcp_server/project_registry.py` | ProjectContext, ProjectRegistry | +| `src/keboola_mcp_server/clients/manage.py` | Async ManageClient for Manage API | +| `src/keboola_mcp_server/tools/branches.py` | list_branches, create_branch tools | + +### Modified Files +| File | Changes | +|------|---------| +| `src/keboola_mcp_server/config.py` | ProjectConfig dataclass, MPA fields, config file loading | +| `src/keboola_mcp_server/mcp.py` | ProjectResolutionMiddleware, SessionStateMiddleware MPA support, write protection | +| `src/keboola_mcp_server/server.py` | Register new middleware + branch tools | +| `src/keboola_mcp_server/cli.py` | `init` subcommand, `--config-file` flag | +| `src/keboola_mcp_server/clients/storage.py` | `dev_branch_create` method | +| `src/keboola_mcp_server/tools/project.py` | MultiProjectInfo response for MPA mode | + +### Test Files +| File | Tests | +|------|-------| +| `tests/test_config.py` | ProjectConfig, is_mpa_mode, from_config_file | +| `tests/test_project_registry.py` (new) | Registry creation, project resolution, defaults, errors | +| `tests/test_mcp.py` | MPA session state, ProjectResolutionMiddleware, write protection | +| `tests/tools/test_branches.py` (new) | list_branches, create_branch | +| `tests/tools/test_project.py` | Multi-project info response | + +--- + +## Verification Plan + +1. **Unit tests**: Run `tox` — all existing tests must pass (backward compatibility), plus new MPA tests +2. **Single-project mode**: Start server with existing env vars / CLI args → verify all tools work exactly as before (no `project_id` param visible) +3. **MPA mode**: Start server with `--config-file mcp.json` containing 2+ projects → verify: + - `get_project_info` returns all projects + - Tools accept `project_id` parameter + - Default project used when `project_id` omitted + - Error when `project_id` missing and no default +4. **Init command**: Run `init` with a manage token → verify `mcp.json` created with correct tokens, manage token not stored +5. **Branch tools**: Create branch, list branches, verify per-project branch management +6. **Write protection**: Enable `forbid_main_branch_writes`, attempt a write tool on main → verify rejection, create branch and retry → verify success + +--- + +## Important Considerations + +- **OAuth + MPA**: OAuth provides a single bearer token for one project. MPA in OAuth mode is not supported initially — only SAPI token mode. Document this clearly. +- **Workspace creation**: Each project needs its own workspace (async). Use `asyncio.gather` for concurrent creation during session init. +- **Token info caching**: Cache `verify_token()` results in `ProjectContext` to avoid redundant API calls per tool invocation. +- **Schema injection**: Modifying Tool JSON schema dynamically requires working with the raw dict returned by `tool.parameters`. Add `project_id` as an optional string property. +- **Cross-stack projects**: While the initial version supports only same-organization projects, `ProjectConfig.storage_api_url` is per-project, so cross-stack support is architecturally possible. \ No newline at end of file From bf10d41c6e628ed221bfb1e741034259dfdf017f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 13 Aug 2026 09:59:04 +0200 Subject: [PATCH 88/89] docs(PSGO-261): fix RFC claim about login-time scoped-token minting The RFC (just restructured) claimed a local login session persists a minted scoped_token for an explicit project subset. That feature was implemented and then explicitly reverted earlier -- TokenSet has no scoped_token/explicitly_scoped fields, and _read_persisted_login_scope builds its SessionScope with none. Corrected: login-time narrowing is local-guard-only (X-KBC-ProjectId per request), not Connection-enforced; only set_project_scope and the single-project auto-confirms (login, OAuth) actually mint a token. Found while verifying RFC/code conformity ahead of undrafting PR #605. --- feature_spec/pat_token_support/RFC.md | 28 ++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/feature_spec/pat_token_support/RFC.md b/feature_spec/pat_token_support/RFC.md index 67b9e8a90..ba401d8b5 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -207,7 +207,33 @@ OAuth is **not** removed (the MCP protocol needs it for HTTP transport). The OAu 4. **SA token path env var** — align with the workspace step-up var (`b971146f`) and the Go services' `*_KUBERNETES_TOKEN_PATH` convention; share one file-read helper. 5. **Refresh + dead token** — the server **always refreshes during usage** when it holds the token pair; when the token is dead (refresh fails), it clears stored credentials and **enforces re-login**. ---- +1. **`project_id`/scope is explicit session state, never silently derived from the token itself** — + a whole-stack PAT has no implicit project. `get_accessible_projects` + `set_project_scope` are + the only mechanism; there is no separate "select-project" tool or header-only path once a scope + exists. +2. **In-conversation narrowing (`set_project_scope`) and the single-project auto-confirms (local + `login`, OAuth) mint a real token via `pat/exchange` when they can**, preferred over + advisory-only narrowing so a bug elsewhere can't reach an out-of-scope project even by accident; + the fallback (unminted, per-request-header-narrowed) exists only for stacks lacking the exchange + endpoint and says so explicitly to the caller. `login`'s own project-selection prompt (a genuine + subset, not the single-project case) is narrower: it persists the choice but does not mint a + token for it, relying on this server's own per-request guard alone. +3. **Fan-out via active-project indirection, not a per-tool `projects[]` parameter** — existing + tool call sites are unchanged; a dispatch-layer middleware swaps the active client/workspace for + the duration of one call. Read results use a per-project envelope, never a semantic merge. +4. **Writes require an explicit `project_id` tool argument once 2+ projects are scoped** — chosen + over implicitly targeting the "active" project, which was reported as confusing (re-scoping to + change a write's target also reordered every subsequent read fan-out). +5. **Scope persistence is per-session-type, not a single shared mechanism** — OAuth and Kai + sessions persist server-side (Postgres) since the client already sends a stable identifier on + every request; local/stateless sessions round-trip an opaque, encrypted `scope_token` instead, + since there's no server-side store to key against. +6. **Security fixes address the flow, not just the symptom** — e.g. local sessions are scoped at + `login` time (removing the unconfirmed-by-default state entirely) rather than documented as a + limitation; credential races are closed by removing the shared state (per-interface keying) + rather than only adding a lock around it. +7. **`clientId` for PKCE** is the demo value `keboola-cli-demo`, configurable via + `KBC_PKCE_CLIENT_ID`; refresh tokens are treated as opaque strings (no prefix assumptions). # Extension: Multi-project scope via introspect + scoped exchange (PSGO-261, increment 2) From 4a54436bf74b2898be1ecee659c49b16910adbc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Va=C5=A1ko?= Date: Thu, 13 Aug 2026 11:16:27 +0200 Subject: [PATCH 89/89] =?UTF-8?q?fix(PSGO-261):=20post-rebase=20fixups=20?= =?UTF-8?q?=E2=80=94=20stale=20call=20sites,=20dead=20import,=20dup=20test?= =?UTF-8?q?,=20ruff=200.16=20lint?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rebase onto main (post-#604 merge) surfaced a few loose ends: two test call sites still missing own_stack_storage_api_url after the security-fix conflict resolution, an unused jwt_utils import left over from a mcp.py merge, a duplicated TestServerRuntimeInfoSessionStatePersists class from resolving two separate conflicts, and lint findings only visible under ruff 0.16 (CI's pinned version) that this session's 0.15 venv didn't catch. --- src/keboola_mcp_server/auth_login.py | 2 +- src/keboola_mcp_server/jwt_utils.py | 3 ++- src/keboola_mcp_server/mcp.py | 1 - src/keboola_mcp_server/scope.py | 4 ++-- src/keboola_mcp_server/tools/project.py | 4 ++-- tests/test_cli.py | 12 ++++++------ tests/test_config.py | 15 --------------- tests/test_mcp.py | 6 ++++-- tests/test_multiproject.py | 8 ++++---- 9 files changed, 21 insertions(+), 34 deletions(-) diff --git a/src/keboola_mcp_server/auth_login.py b/src/keboola_mcp_server/auth_login.py index 7cda71245..2bc1b4798 100644 --- a/src/keboola_mcp_server/auth_login.py +++ b/src/keboola_mcp_server/auth_login.py @@ -581,7 +581,7 @@ def forget_tokens(storage_api_url: str | None = None, *, profile: str | None = N class _CallbackHandler(BaseHTTPRequestHandler): result: ClassVar[dict] = {} - def do_GET(self) -> None: # noqa: N802 (BaseHTTPRequestHandler API) + def do_GET(self) -> None: query = urllib.parse.parse_qs(urlparse(self.path).query) type(self).result = {k: v[0] for k, v in query.items()} self.send_response(200) diff --git a/src/keboola_mcp_server/jwt_utils.py b/src/keboola_mcp_server/jwt_utils.py index 5562cc5ac..ef6f6a2de 100644 --- a/src/keboola_mcp_server/jwt_utils.py +++ b/src/keboola_mcp_server/jwt_utils.py @@ -9,7 +9,8 @@ import gzip import json -from typing import Any, Mapping +from collections.abc import Mapping +from typing import Any import jwt.api_jws diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index 53402d6a3..4c7d0ae00 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -39,7 +39,6 @@ deployed_sa_token_path, is_same_stack, ) -from keboola_mcp_server.jwt_utils import decode_jwt, encode_jwt from keboola_mcp_server.oauth import ProxyAccessToken from keboola_mcp_server.scope import ( OAUTH_SESSION_ID_KEY, diff --git a/src/keboola_mcp_server/scope.py b/src/keboola_mcp_server/scope.py index d2f6c3449..f9914e80a 100644 --- a/src/keboola_mcp_server/scope.py +++ b/src/keboola_mcp_server/scope.py @@ -11,7 +11,7 @@ import json import time from datetime import datetime, timezone -from typing import TYPE_CHECKING, Annotated, Optional +from typing import TYPE_CHECKING, Annotated from pydantic import Field @@ -28,7 +28,7 @@ # when the scope resolves the target unambiguously (a single scoped project). PROJECT_ID_ARG = 'project_id' ProjectIdArg = Annotated[ - Optional[str], + str | None, Field( description=( 'Target Keboola project id for this write. Required when the session is scoped to 2+ ' diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 297e40882..9f30b4251 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -1,6 +1,6 @@ import asyncio import logging -from typing import Annotated, Optional, cast +from typing import Annotated, cast import httpx from fastmcp import Context, FastMCP @@ -546,7 +546,7 @@ async def get_accessible_projects( async def set_project_scope( ctx: Context, project_ids: Annotated[ - Optional[list[int]], + list[int] | None, Field( description='The project ids to scope the session to. ' 'Omit or pass null to scope to ALL accessible projects.' diff --git a/tests/test_cli.py b/tests/test_cli.py index c65e5e02b..5919978a3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -174,9 +174,9 @@ async def test_closes_pool_even_if_migration_fails(self, monkeypatch) -> None: 'keboola_mcp_server.session_store.migrator.apply_migrations', AsyncMock(side_effect=RuntimeError('boom')), ), + pytest.raises(RuntimeError, match='boom'), ): - with pytest.raises(RuntimeError, match='boom'): - await _run_migrate() + await _run_migrate() pool.close.assert_awaited_once() @@ -192,9 +192,9 @@ async def test_closes_pool_even_if_partition_ensure_fails(self, monkeypatch) -> 'keboola_mcp_server.session_store.retention.ensure_partitions', AsyncMock(side_effect=RuntimeError('boom')), ), + pytest.raises(RuntimeError, match='boom'), ): - with pytest.raises(RuntimeError, match='boom'): - await _run_migrate() + await _run_migrate() pool.close.assert_awaited_once() @@ -255,9 +255,9 @@ async def test_closes_pool_even_if_it_fails(self, monkeypatch) -> None: 'keboola_mcp_server.session_store.retention.ensure_partitions', AsyncMock(side_effect=RuntimeError('boom')), ), + pytest.raises(RuntimeError, match='boom'), ): - with pytest.raises(RuntimeError, match='boom'): - await _run_gc_sessions() + await _run_gc_sessions() pool.close.assert_awaited_once() diff --git a/tests/test_config.py b/tests/test_config.py index b1a12beaa..32ad2a49a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -270,18 +270,3 @@ def test_is_same_stack(self, url: str | None, other_url: str | None, expected: b assert is_same_stack(url, other_url) is expected # The comparison is symmetric. assert is_same_stack(other_url, url) is expected - - -class TestServerRuntimeInfoSessionStatePersists: - def test_stdio_always_persists_regardless_of_stateless_http(self) -> None: - # stdio is one process/one session for the whole conversation -- the flag is meaningless there. - assert ServerRuntimeInfo(transport='stdio', stateless_http=True).session_state_persists is True - assert ServerRuntimeInfo(transport='stdio', stateless_http=False).session_state_persists is True - - def test_streamable_http_follows_stateless_http_flag(self) -> None: - assert ServerRuntimeInfo(transport='streamable-http', stateless_http=True).session_state_persists is False - assert ServerRuntimeInfo(transport='streamable-http', stateless_http=False).session_state_persists is True - - def test_defaults_to_stateless(self) -> None: - # Matches the CLI's --stateless-http default (scaled/deployed-safe). - assert ServerRuntimeInfo(transport='streamable-http').session_state_persists is False diff --git a/tests/test_mcp.py b/tests/test_mcp.py index c3036e6e7..7df7dd052 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -905,7 +905,7 @@ def test_apply_request_config_injects_exchanged_session_token(self): http_rq = Request({'type': 'http', 'headers': [], 'user': AuthenticatedUser(access_token)}) config = Config(storage_api_url='https://connection.test.keboola.com') - out_config = SessionStateMiddleware.apply_request_config(http_rq, config) + out_config = SessionStateMiddleware.apply_request_config(http_rq, config, own_stack_storage_api_url=None) assert out_config.storage_token == 'kbc_at_exchanged' assert is_programmatic_token(out_config.storage_token) @@ -1169,7 +1169,9 @@ async def test_forwards_bearer_regardless_of_deployment_or_project_id( runtime_info = ServerRuntimeInfo(transport='http') with patch.object(WorkspaceManager, 'create', AsyncMock(return_value='wsm')): - state = await SessionStateMiddleware.create_session_state(config, runtime_info) + state = await SessionStateMiddleware.create_session_state( + config, runtime_info, own_stack_storage_api_url=None + ) client = state[KeboolaClient.STATE_KEY] assert client.bearer_token == 'kbc_at_abc' diff --git a/tests/test_multiproject.py b/tests/test_multiproject.py index f54048b61..c75e25c3d 100644 --- a/tests/test_multiproject.py +++ b/tests/test_multiproject.py @@ -561,9 +561,9 @@ async def call_next(_): 'create', AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), + pytest.raises(ToolError, match='failed for all 2 scoped'), ): - with pytest.raises(ToolError, match='failed for all 2 scoped'): - await MultiProjectMiddleware().on_call_tool(context, call_next) + await MultiProjectMiddleware().on_call_tool(context, call_next) @pytest.mark.asyncio async def test_fan_out_validation_error_raised_once_not_per_project(self) -> None: @@ -588,9 +588,9 @@ async def call_next(_): 'create', AsyncMock(side_effect=lambda client, _schema, kubernetes_token_path=None: f'wsm-{client}'), ), + pytest.raises(PydanticValidationError), ): - with pytest.raises(PydanticValidationError): - await MultiProjectMiddleware().on_call_tool(context, call_next) + await MultiProjectMiddleware().on_call_tool(context, call_next) # Aborted after the first project; not retried across the rest. assert calls == ['client-11']