diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d341d7555..cc5db3c6a 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 }} + # 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 }} run: | uv run tox -e integtests diff --git a/README.md b/README.md index bee7d0788..ab9341a2a 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, 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. -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 diff --git a/TOOLS.md b/TOOLS.md index fb0861cf4..7aa01c03a 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -50,9 +50,12 @@ 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 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. +- [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 @@ -187,6 +190,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": [ @@ -351,6 +366,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": [ @@ -508,6 +535,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": [ @@ -1075,6 +1114,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": [ @@ -1350,6 +1401,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": [ @@ -2016,6 +2079,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": [ @@ -2085,6 +2160,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": [ @@ -2135,6 +2222,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": [ @@ -2214,6 +2313,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": [ @@ -2500,6 +2611,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": [ @@ -2601,6 +2724,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": [ @@ -2707,6 +2842,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": [ @@ -2778,6 +2925,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": [ @@ -3114,6 +3273,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": [ @@ -3246,6 +3417,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": [ @@ -3497,6 +3680,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": [ @@ -3539,6 +3734,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": [ @@ -3552,6 +3759,44 @@ 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, each with its SQL +dialect and organization. + +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**: +```json +{ + "additionalProperties": false, + "properties": { + "with_llm_instruction": { + "default": false, + "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" + } + }, + "type": "object" +} +``` + +--- ## get_project_info **Annotations**: `read-only` @@ -3565,14 +3810,81 @@ 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": { + "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" +} +``` + +--- + +## 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/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. + +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**: ```json { "additionalProperties": false, - "properties": {}, + "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" } ``` @@ -3597,6 +3909,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": [ @@ -4617,6 +4941,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/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/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 diff --git a/feature_spec/oauth_session_exchange/RFC.md b/feature_spec/oauth_session_exchange/RFC.md new file mode 100644 index 000000000..111f52a00 --- /dev/null +++ b/feature_spec/oauth_session_exchange/RFC.md @@ -0,0 +1,104 @@ +# 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` 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. + +## 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). **`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/feature_spec/oauth_session_persistence/RFC.md b/feature_spec/oauth_session_persistence/RFC.md new file mode 100644 index 000000000..0f2ccf6c3 --- /dev/null +++ b/feature_spec/oauth_session_persistence/RFC.md @@ -0,0 +1,278 @@ +# 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. 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 | +| `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 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 +(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. + +`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` +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 — 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 + 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 fc4834a93..a22385953 100644 --- a/feature_spec/pat_token_support/RFC.md +++ b/feature_spec/pat_token_support/RFC.md @@ -206,3 +206,906 @@ 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**. + +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) + +> 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 } } + # NB: project ids are sent as STRINGS — the exchange API 400s on integers (auth_login.py:166) + +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. 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` + 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. + _(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) + +**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.** (Resolved as leaned — see "Decisions (increment 5)" below.) + +# 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** | `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) + +- **`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 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) + +## 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. + +## 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, +`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) + +- **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. + _(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 + 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. + +## 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. + +## 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. + +--- + +# 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). + +--- + +# 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. + +## 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. + +## 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. + +## 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. + +## Extension: `scope_token` was replayable across callers (increment 14) + +Raised in review (Tomas Fejfar), distinct from increment 7's finding #2: that fix +(`SessionScope.to_token`/`from_token` moving from JWS to AES-GCM) addressed *confidentiality* -- +the embedded live `scoped_token` is no longer recoverable without the key. It did nothing about +*replay*: `scope_token` decryption never checked who was presenting it, only that the ciphertext +authenticated against the single deployment-wide key. Anyone who obtained a valid `scope_token` +string some other way -- a shared/exported conversation transcript, client-side logs, an +observability platform sitting on the MCP traffic -- could resend it verbatim as their own +`scope_token` argument from an unrelated (even freshly self-registered) session on the same +deployment, and `MultiProjectMiddleware` would use the embedded `scoped_token` as `base_token` for +every fanned-out call, read *and* write (`multiproject.py`'s `base_token = scope.scoped_token or +...`), fully impersonating whoever the scope was minted for. + +This only matters on the deployed server: `set_project_scope`/`get_accessible_projects` only ever +return `scope_token` to the client when the scope isn't persisted server-side (`persisted = +_persist_oauth_scope(...) or _persist_kai_scope(...) or session_state_persists`) -- OAuth and Kai +sessions persist it in Postgres instead, and `stdio`/`--no-stateless-http` local sessions keep it +in `ctx.session.state`. A local server has no cross-caller boundary to defend in the first place +(`login` already grants its one user the whole stack), so the fix below is a no-op there by +design, not by omission. + +**Fix:** `session_store/crypto.py`'s `encrypt`/`decrypt` gain an optional `aad` (AES-GCM +additional authenticated data) parameter -- authenticated but never transmitted, so both sides +must already agree on it out of band. `scope.py` adds `resolve_scope_binding_aad(storage_token)`, +returning `sha256(storage_token)` when `deployed_sa_token_path()` is set (and the caller has a +token), `None` otherwise (local). `SessionScope.to_token`/`from_token` take an `aad` param and +thread it into `encrypt`/`decrypt`. Both mint sites (`set_project_scope`, `get_accessible_projects` +in `tools/project.py`) bind to `client.token`; the read side +(`SessionStateMiddleware._read_scope_from_request`) binds to `config.storage_token` -- the same +underlying value, since the deployed branch of `_resolve_local_tokens` never overwrites +`config.storage_token` with a narrowed `scoped_token` (only the local-programmatic branch does +that, which is exactly the branch this fix doesn't apply to). A `scope_token` minted while serving +caller A's request now fails AES-GCM authentication -- and is treated as "no scope", same as any +other invalid token -- when replayed by caller B, whose own `storage_token` hashes to a different +`aad`. diff --git a/integtests/conftest.py b/integtests/conftest.py index 133d0406e..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 @@ -45,6 +47,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 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' # 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' @@ -118,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: @@ -205,6 +215,60 @@ 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 + + +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 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') 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_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', diff --git a/integtests/test_pat_multiproject.py b/integtests/test_pat_multiproject.py new file mode 100644 index 000000000..8380ba9a2 --- /dev/null +++ b/integtests/test_pat_multiproject.py @@ -0,0 +1,109 @@ +"""Integration tests that run the SAME questions under different authentication (PSGO-261). + +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'] + + +def _make_server(config: Config) -> FastMCP: + server = create_server(config, runtime_info=ServerRuntimeInfo(transport='stdio')) + assert isinstance(server, FastMCP) + return server + + +@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_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' + assert all(p.sql_dialect in ('Snowflake', 'BigQuery') for p in result.projects) + assert result.base_instructions, 'base_instructions expected when with_llm_instruction=true' + assert {g.sql_dialect for g in result.base_instructions} == {p.sql_dialect for p in result.projects} diff --git a/pyproject.toml b/pyproject.toml index 5e8c8429d..acee5cd9f 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.3" 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/auth_login.py b/src/keboola_mcp_server/auth_login.py new file mode 100644 index 000000000..2bc1b4798 --- /dev/null +++ b/src/keboola_mcp_server/auth_login.py @@ -0,0 +1,643 @@ +"""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 asyncio +import base64 +import contextlib +import dataclasses +import hashlib +import json +import logging +import os +import secrets +import sys +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 ClassVar, cast +from urllib.parse import urlparse + +import httpx + +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' +_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' +_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' +# 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. +_AUTH_TIMEOUT = httpx.Timeout(connect=5.0, read=30.0, write=10.0, pool=5.0) + + +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: + return normalize_storage_api_url(storage_api_url) + + +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. + + ``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: + 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')), + ) + + +@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=_AUTH_TIMEOUT, 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=_AUTH_TIMEOUT, 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 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=_AUTH_TIMEOUT, 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: + # 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) + + +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=_AUTH_TIMEOUT, 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: + # 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: + 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, + 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: 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. + """ + 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( + 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, + *, + 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=_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())) + + +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=_AUTH_TIMEOUT, 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 + interface profile) --- + + +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: + 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: + # 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) + + +@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, *, profile: str | None = None) -> None: + store = _read_store() + 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+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, 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 + + 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, +) -> 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 + (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. + """ + try: + 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, profile=profile, open_browser=open_browser) + return await get_access_token(storage_api_url, profile=profile, transport=transport) + + +def _forget(storage_api_url: str, *, profile: str | None = None) -> None: + store = _read_store() + 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, *, 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. + """ + 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, profile), None) is not None: + _write_store(store) + return True + return False + + +# --- interactive browser login (not unit-tested; exercises a real browser + loopback) --- + + +class _CallbackHandler(BaseHTTPRequestHandler): + result: ClassVar[dict] = {} + + 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) + 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, *, 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()) + 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)}' + # 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 = {} + # 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): + 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…', file=sys.stderr, 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, profile=profile) + return tokens diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 21f56be9c..31a72518f 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -3,11 +3,14 @@ import argparse import asyncio import contextlib +import dataclasses +import getpass import json import logging.config import os import pathlib import sys +import time import traceback import pydantic @@ -55,8 +58,119 @@ 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') + 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.', + ) + 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 the scoped ' + '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). 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', + metavar='STR', + 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. Also re-prompts for project ' + 'scope even if one is already stored.', + ) + + 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( + '--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( + '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.', + ) + + subparsers.add_parser( + 'gc-sessions', + help='Ensures upcoming oauth_sessions partitions exist and drops ones past the retention ' + 'window, then exits. Intended to run on a recurring schedule (e.g. a kbc-stacks CronJob), ' + 'independent of deployments.', + ) + return parser.parse_args(args) @@ -98,6 +212,260 @@ 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. + + 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) + 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 _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 + + 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) + + +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, + pat_name: str = 'keboola-mcp-server', + show_token: bool = False, + force: bool = False, +) -> None: + """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. + + 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, + 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 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, 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, 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, {len(project_ids)} project(s)):\n\n {pat_token}\n') + + +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 + + 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, profile=profile) + 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 + 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 migrations: set MCP_DB_URL (or KBC_POSTGRES_DSN).') + + 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 recurring gc-sessions job makes on an ongoing basis. + partitions = await ensure_partitions(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.') + if partitions['created']: + print(f"✓ Ensured oauth_sessions partitions: {', '.join(partitions['created'])}") + + +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) @@ -124,11 +492,62 @@ async def run_server(args: list[str] | None = None) -> None: stream=sys.stderr, ) - # Create config from the CLI arguments + 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), + 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), + profile=getattr(parsed_args, 'profile', None), + all_stacks=getattr(parsed_args, 'all', False), + ) + return + + if parsed_args.command == 'migrate': + await _run_migrate() + return + + if parsed_args.command == 'gc-sessions': + await _run_gc_sessions() + return + + # 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, required=parsed_args.transport == 'stdio' ) try: @@ -155,14 +574,16 @@ 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' ) 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') 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..3409d8930 --- /dev/null +++ b/src/keboola_mcp_server/clients/auth_bridge.py @@ -0,0 +1,140 @@ +"""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 +from http import HTTPStatus +from typing import cast + +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_' +_PAT_PREFIX = 'kbc_pat_' +_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( + {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, _PAT_PREFIX)) + + +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). + """ + + 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).""" + + 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 + + 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) + + 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). + """ + # 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. + sa_jwt = self._read_sa_jwt() + headers = { + 'Content-Type': 'application/json', + 'Accept': 'application/json', + 'X-Kubernetes-Authorization': f'Bearer {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/clients/base.py b/src/keboola_mcp_server/clients/base.py index 1d1c14341..fc0ba5bc4 100644 --- a/src/keboola_mcp_server/clients/base.py +++ b/src/keboola_mcp_server/clients/base.py @@ -1,7 +1,10 @@ import json import logging +import re 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 @@ -16,6 +19,39 @@ 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 genuine ``connection.*.keboola.(com|dev)`` host. + """ + parsed = urlparse(storage_api_url) + 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, '', '', '', '')) + + +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: """ diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index 7cdd4acae..fd4252e9a 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 @@ -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, ) @@ -164,12 +166,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}', '', '', '', '')) @@ -179,7 +177,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 @@ -193,10 +191,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 +217,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, @@ -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,20 +337,21 @@ 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 = 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}' 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, + readonly=None, ) @property diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 6168aeb55..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__) @@ -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.""" @@ -39,10 +50,42 @@ 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 = 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. + + 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 """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: 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. + 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): @@ -73,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: @@ -118,15 +169,29 @@ 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): 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): @@ -237,6 +302,39 @@ 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]: + """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: 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/jwt_utils.py b/src/keboola_mcp_server/jwt_utils.py new file mode 100644 index 000000000..ef6f6a2de --- /dev/null +++ b/src/keboola_mcp_server/jwt_utils.py @@ -0,0 +1,25 @@ +""" +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 collections.abc import Mapping +from typing import Any + +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 c419e0fab..cd3174211 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 from collections.abc import Awaitable, Callable, Iterable from typing import Any, TypeVar @@ -29,11 +28,35 @@ 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, load_tokens +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, 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.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.scope import ( + OAUTH_SESSION_ID_KEY, + SCOPE_KEY, + SCOPE_TOKEN_ARG, + SessionScope, + persist_scope, + resolve_scope_binding_aad, + resolve_scope_key, +) +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, + MODIFY_FLOW_TOOL_NAME, + SEMANTIC_TOOLS_TAG, + UPDATE_FLOW_TOOL_NAME, +) from keboola_mcp_server.workspace import WorkspaceManager LOG = logging.getLogger(__name__) @@ -98,6 +121,8 @@ async def project_has_semantic_models(client: KeboolaClient) -> bool: 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: @@ -225,6 +250,79 @@ 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: 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 + # 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 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) + # 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) + # 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) + # 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) + # 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: + 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. @@ -232,14 +330,23 @@ 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) + # 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 + if oauth_session_id := self._read_oauth_session_id(http_rq): + state[OAUTH_SESSION_ID_KEY] = oauth_session_id ctx.session.state = state try: @@ -249,23 +356,50 @@ async def on_request( # ctx.session.state = {} pass - @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 + 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`. + + 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. """ - 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}' - ), - } + 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 {}) + 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 apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_storage_api_url: str | None) -> Config: @@ -288,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( @@ -298,19 +436,302 @@ 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}') 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, + # 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) return config + @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. + """ + # 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) + if not token: + return None + try: + return SessionScope.from_token( + token, resolve_scope_key(config), aad=resolve_scope_binding_aad(config.storage_token) + ) + except Exception: + 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 + + @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 + 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.""" + return ( + not deployed_sa_token_path() + and bool(config.storage_token) + and bool(config.storage_api_url) + and is_programmatic_token(config.storage_token) + ) + + @classmethod + 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. 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 deployed_sa_token_path(): + 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 + 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 + 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': + """ + 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 = strip_bearer(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}', exc_info=True) + 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), ``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: + # 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: + 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, + # 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) + 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}', exc_info=True) + 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 create_session_state( cls, @@ -343,11 +764,26 @@ 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 + extra_headers: dict[str, Any] = {} + if is_programmatic_token(storage_token): + # 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, - storage_api_token=config.storage_token, - bearer_token=config.bearer_token, - headers=cls._get_headers(runtime_info), + storage_api_token=storage_token, + bearer_token=bearer_token, + 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) @@ -365,7 +801,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 ) @@ -447,6 +883,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 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) + if is_programmatic_token(client.token): + 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() @@ -555,6 +1003,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/src/keboola_mcp_server/multiproject.py b/src/keboola_mcp_server/multiproject.py new file mode 100644 index 000000000..ec6479b10 --- /dev/null +++ b/src/keboola_mcp_server/multiproject.py @@ -0,0 +1,459 @@ +"""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 PROJECT_ID_ARG, SCOPE_KEY, SessionScope +from keboola_mcp_server.tools.constants import BOOTSTRAP_TOOLS +from keboola_mcp_server.workspace import WorkspaceManager + +LOG = logging.getLogger(__name__) + +# 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). +_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. + + 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: a write always targets exactly one project, named by its own + ``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( + 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) + # 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_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. + 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 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) + 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) + + @staticmethod + 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. + """ + 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) -- this tool 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_single_target( + 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, 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_single_target(scope, 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) + 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], + 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 tool.name in _SINGLE_TARGET_READ_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/oauth.py b/src/keboola_mcp_server/oauth.py index a17871e87..921d2e38c 100644 --- a/src/keboola_mcp_server/oauth.py +++ b/src/keboola_mcp_server/oauth.py @@ -1,5 +1,4 @@ -import gzip -import json +import dataclasses import logging import math import os @@ -7,23 +6,42 @@ import secrets import time from collections.abc import Mapping -from http.client import HTTPException +from datetime import datetime, timezone from typing import Any, cast 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, 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, + 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 +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')) @@ -113,14 +131,31 @@ 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. 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 + 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): - 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 + session_id: str | None = None class SimpleOAuthProvider(OAuthProvider): @@ -134,6 +169,7 @@ def __init__( client_secret: str, server_url: str, scope: str, + session_store: SessionStore, jwt_secret: str | None = None, ) -> None: """ @@ -146,18 +182,23 @@ 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._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 +269,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,104 +422,135 @@ 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 - ) - - # 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, - ) - 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)}', + # 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( client_id=client.client_id, - scopes=authorization_code.scopes, - expires_at=authorization_code.oauth_refresh_token.expires_at, - delegate=authorization_code.oauth_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=expires_in, - scope=' '.join(access_token.scopes), + 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), ) - - _log_debug( - f'[exchange_authorization_code] access_token={access_token}, refresh_token={refresh_token},' - f'oauth_token={oauth_token}' + 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, ) - - return oauth_token + 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 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 + ) + LOG.info(f'[load_access_token] Lazily refreshed near-expiry Keboola session: session_id={session.id}') + + 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, + 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 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( @@ -486,8 +560,12 @@ 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. + + 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. @@ -496,98 +574,69 @@ 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 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 - # 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 - ) - - # 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=oauth_access_token.scopes, - expires_at=oauth_access_token.expires_at, - delegate=oauth_access_token, - sapi_token=sapi_token, - ) - 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=oauth_refresh_token.scopes, - expires_at=oauth_refresh_token.expires_at, - delegate=oauth_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 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('[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 + + await self._session_store.rotate_kbc_tokens( + refresh_token.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), ) - 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(access_token.expires_at - time.time())), - scope=' '.join(access_token.scopes), - ) - - _log_debug( - f'[exchange_refresh_token] access_token={access_token}, refresh_token={refresh_token}, ' - f'oauth_token={oauth_token}' + expires_in=None, + scope=' '.join(scopes), ) - 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]: """ @@ -622,39 +671,38 @@ 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``. + + 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. """ - 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}', - }, + 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 TokenError( + error='invalid_request', + error_description='OAuth login is misconfigured: no Kubernetes ServiceAccount token available.', ) - 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}') + 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 TokenError(error='invalid_grant', error_description=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: @@ -665,15 +713,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/scope.py b/src/keboola_mcp_server/scope.py new file mode 100644 index 000000000..bd2ea60e3 --- /dev/null +++ b/src/keboola_mcp_server/scope.py @@ -0,0 +1,159 @@ +"""In-conversation multi-project scope (PSGO-261 increment 2): the ``SessionScope`` model, its +``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 gzip +import hashlib +import json +import time +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Annotated + +from pydantic import Field + +from keboola_mcp_server.config import Config, deployed_sa_token_path +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 + +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[ + str | None, + 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 +# 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' + + +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) + + +def resolve_scope_binding_aad(storage_token: str | None) -> bytes | None: + """Binds a `scope_token` to the caller's own presented Storage token on the deployed server, + so a token minted while serving one caller's request fails authentication if replayed by a + different caller who obtained the ciphertext some other way (a leaked transcript, client-side + logs) but doesn't hold the matching bearer credential -- see the "Security hardening" RFC + increment. Encryption alone (`SessionScope.to_token`) only stops a passive eavesdropper from + reading the embedded live `scoped_token`; it does nothing to stop this replay, which needs no + key at all, only the opaque string itself. + + A no-op (returns None) on a local server: a `login` session already grants the whole stack to + its single local user, so there is no cross-caller boundary to enforce, and the caller's own + presented token there is not a stable value to bind to (it may itself get narrowed to a scoped + token between calls -- see `SessionStateMiddleware._resolve_local_tokens`). + """ + if not deployed_sa_token_path() or not storage_token: + return None + return hashlib.sha256(storage_token.encode('utf-8')).digest() + + +@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, key: bytes, *, aad: bytes | None = None) -> 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. + + ``aad``, typically ``resolve_scope_binding_aad(caller's storage token)``, additionally + binds the ciphertext to the caller it was minted for -- pass the *same* value to + ``from_token`` or decryption fails. Without it, encryption alone stops eavesdropping but + not replay by a different caller who obtains the opaque string some other way. + """ + plaintext = gzip.compress(json.dumps(dataclasses.asdict(self)).encode('utf-8')) + return base64.urlsafe_b64encode(encrypt(plaintext, key, aad)).decode('ascii').rstrip('=') + + @classmethod + def from_token(cls, token: str, key: bytes, *, aad: bytes | None = None) -> 'SessionScope': + """Inverse of ``to_token``. Raises on a missing/invalid/tampered/wrong-key/wrong-aad + 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, aad) + 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: + """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/server.py b/src/keboola_mcp_server/server.py index 286d8bea9..2d9c9157d 100644 --- a/src/keboola_mcp_server/server.py +++ b/src/keboola_mcp_server/server.py @@ -18,15 +18,14 @@ 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, - 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 +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 from keboola_mcp_server.tools.doc import add_doc_tools @@ -208,6 +207,27 @@ 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.' + ) + # 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) + ) + oauth_provider = SimpleOAuthProvider( storage_api_url=config.storage_api_url, client_id=config.oauth_client_id, @@ -219,21 +239,56 @@ 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 + 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) + 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=( + 'This server supports multi-project mode for stack-wide Keboola programmatic tokens ' + '(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, middleware=[ LoggingMiddleware(log_level=logging.DEBUG), SessionStateMiddleware(), ToolAuthorizationMiddleware(), + # 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/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..71a764e4c --- /dev/null +++ b/src/keboola_mcp_server/session_store/crypto.py @@ -0,0 +1,67 @@ +"""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, aad: bytes | None = None) -> bytes: + """``aad`` (additional authenticated data) is bound into the auth tag but never transmitted -- + the caller must supply the identical value again to `decrypt`. Use it to tie ciphertext to a + caller identity so a token minted for one caller fails authentication if replayed by another, + even with the right key (see `scope.py`'s `resolve_scope_binding_aad`).""" + 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, aad) + return bytes([_KEY_VERSION]) + nonce + ciphertext + + +def decrypt(blob: bytes, key: bytes, aad: bytes | None = None) -> 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, aad) + except InvalidTag as e: + raise DecryptionError('Ciphertext failed authentication (wrong key/aad 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/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/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/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..09bac860c --- /dev/null +++ b/src/keboola_mcp_server/session_store/migrations/0002_partition_oauth_sessions.sql @@ -0,0 +1,66 @@ +-- 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 (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; + +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/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/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/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..a7eed78b2 --- /dev/null +++ b/src/keboola_mcp_server/session_store/repository.py @@ -0,0 +1,250 @@ +"""``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 asyncio +import dataclasses +import hashlib +import secrets +from datetime import datetime +from typing import Protocol + +import asyncpg + +from keboola_mcp_server.session_store import crypto + +# 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 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, + *, + 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: + """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() + + 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() + + 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() + pool = await self._get_pool() + row = await 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: + 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), + ) + return self._to_session(row) if row is not None else None + + async def get_by_refresh_token(self, refresh_token: str) -> OAuthSession | None: + 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), + ) + 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: + 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, + 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 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, + *, + project_ids: list[int], + read_only: bool, + confirmed: bool, + scoped_token: str | None, + scoped_expires_at: datetime | None, + ) -> None: + pool = await self._get_pool() + await 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: + pool = await self._get_pool() + await pool.execute('UPDATE oauth_sessions SET revoked_at = now() WHERE id = $1', session_id) 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..af5b766de --- /dev/null +++ b/src/keboola_mcp_server/session_store/retention.py @@ -0,0 +1,109 @@ +"""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 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 +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()) + # 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: + 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, not user + # input, so direct formatting is safe. + # + # 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( + 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()}')" + ) + # 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) + + 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) + 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}') + LOG.info(f'Dropped expired oauth_sessions partition: {name} (older than {cutoff})') + dropped.append(name) + + return {'created': created, 'dropped': dropped} 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/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/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 f9dc56a58..f722a82fb 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -1,17 +1,32 @@ +import asyncio import logging from typing import Annotated, cast +import httpx 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.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 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 ( + OAUTH_SESSION_ID_KEY, + SCOPE_KEY, + ProjectIdArg, + SessionScope, + persist_scope, + resolve_scope_binding_aad, + resolve_scope_key, +) from keboola_mcp_server.workspace import WorkspaceManager LOG = logging.getLogger(__name__) @@ -40,9 +55,52 @@ 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. + + 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: + return strip_bearer(cast(str, 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. @@ -155,6 +213,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) @@ -168,6 +227,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, @@ -175,7 +235,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) @@ -224,3 +285,382 @@ async def get_project_info( LOG.info('Returning unified 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 + + +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.') + 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.') + 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): + """Base working instructions shared by all projects of one SQL dialect (dialect-specific system prompt).""" + + 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.') + instructions: str = Field(description='The base working instructions for projects of this dialect.') + + +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.', + ) + 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=( + '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 (distinct from base_instructions).', + ) + + +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 | None = Field( + default=None, + description=( + '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 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: + return False + session_store = ServerState.from_context(ctx).session_store + if session_store is None: + return False + await persist_scope(session_store, session_id, scope) + 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]: + """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 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() + org_id, org_name = _organization_from_token(token_data) + return project_id, _sql_dialect_from_token(token_data), org_id, org_name + + +@tool_errors() +async def get_accessible_projects( + ctx: Context, + with_llm_instruction: Annotated[ + bool, + Field( + 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.' + ) + ), + ] = False, +) -> AccessibleProjects: + """ + Lists the Keboola projects the current login can access across the stack, each with its SQL + dialect and organization. + + 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) + introspection = await introspect_token(client.storage_api_url, subject_token=subject_token) + + 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 + 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) + 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_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/organization for a project: {result}', exc_info=result) + continue + 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, + ) + ) + + # Optionally attach the base working instructions, grouped by dialect so the (large) prompt is + # sent once per distinct dialect rather than duplicated per project. + 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) + base_instructions = [ + BaseInstructionGroup( + project_ids=ids, + sql_dialect=dialect, + # 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() + ] + + 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.' + ) + is_persisted = False + else: + 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_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( + 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_key(server_state.config), aad=resolve_scope_binding_aad(client.token)) + if scoped_ids is not None and not is_persisted + else None + ), + base_instructions=base_instructions, + llm_instruction=instruction, + ) + + +@tool_errors() +async def set_project_scope( + ctx: Context, + project_ids: Annotated[ + list[int] | None, + 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/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. + + 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) + + # 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) + 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: + 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 + # 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 + + # 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 + server_state = ServerState.from_context(ctx) + 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_key(server_state.config), aad=resolve_scope_binding_aad(client.token)) + ) + 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, + scope_token=scope_token, + llm_instruction=( + ( + 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}{read_only_note}' + ) + if multi + else f'Session scoped to project {ids[0]}. {resend_instruction}{read_only_note}' + ), + ) 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/src/keboola_mcp_server/workspace.py b/src/keboola_mcp_server/workspace.py index 0ec4ac7d3..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.') @@ -786,6 +788,25 @@ 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 (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 + elif duration > timeout_sec: LOG.info(f'Workspace creation timed out after {duration:.2f} seconds.') return None diff --git a/tests/clients/test_auth_bridge.py b/tests/clients/test_auth_bridge.py new file mode 100644 index 000000000..037259346 --- /dev/null +++ b/tests/clients/test_auth_bridge.py @@ -0,0 +1,132 @@ +"""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 OAuthSessionExchanger, OAuthTokenExchangeError, 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 test_invalid_storage_api_url_rejected(sa_token_file: Path) -> None: + with pytest.raises(ValueError, match='Invalid Keboola Storage API URL'): + OAuthSessionExchanger(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' + # 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' + + +@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: + 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/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/clients/test_client.py b/tests/clients/test_client.py index 644eca11d..388126f96 100644 --- a/tests/clients/test_client.py +++ b/tests/clients/test_client.py @@ -12,12 +12,12 @@ 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 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', @@ -345,9 +345,9 @@ 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 + 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', @@ -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(): @@ -802,8 +747,28 @@ 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): + 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( @@ -815,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']) @@ -839,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') @@ -864,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.""" @@ -892,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/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..9e5850539 --- /dev/null +++ b/tests/session_store/conftest.py @@ -0,0 +1,61 @@ +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.kai_scope import PostgresKaiScopeStore +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') + + +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, 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, kai_sessions, schema_migrations CASCADE') + await apply_migrations(pool) + finally: + await pool.close() + s = PostgresSessionStore(TEST_DSN, encryption_key=bytes([1] * KEY_SIZE)) + try: + 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_crypto.py b/tests/session_store/test_crypto.py new file mode 100644 index 000000000..e4b96f35a --- /dev/null +++ b/tests/session_store/test_crypto.py @@ -0,0 +1,86 @@ +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_round_trip_with_aad() -> None: + ciphertext = encrypt(b'kbc_at_secret', _key(), aad=b'caller-a') + assert decrypt(ciphertext, _key(), aad=b'caller-a') == b'kbc_at_secret' + + +def test_wrong_aad_fails() -> None: + ciphertext = encrypt(b'kbc_at_secret', _key(), aad=b'caller-a') + with pytest.raises(DecryptionError): + decrypt(ciphertext, _key(), aad=b'caller-b') + + +def test_missing_aad_fails_when_encrypted_with_one() -> None: + ciphertext = encrypt(b'kbc_at_secret', _key(), aad=b'caller-a') + with pytest.raises(DecryptionError): + decrypt(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_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 new file mode 100644 index 000000000..8deb1fad8 --- /dev/null +++ b/tests/session_store/test_migrator.py @@ -0,0 +1,101 @@ +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, kai_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', + '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 + # 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() + + +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) + 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%'") + } + assert tables == {'oauth_sessions', 'oauth_sessions_default'} + finally: + 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 + # 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_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/session_store/test_retention.py b/tests/session_store/test_retention.py new file mode 100644 index 000000000..a6101af49 --- /dev/null +++ b/tests/session_store/test_retention.py @@ -0,0 +1,172 @@ +from datetime import date, datetime, timedelta, timezone + +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, kai_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: + 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': []} + 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(datetime.now(tz=timezone.utc).date()) + 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_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(datetime.now(tz=timezone.utc).date()) + 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_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(datetime.now(tz=timezone.utc).date()) + 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_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(datetime.now(tz=timezone.utc).date()):%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: + # 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(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 + assert await self._existing_partitions(pool) == expected + finally: + await pool.close() diff --git a/tests/test_auth_login.py b/tests/test_auth_login.py new file mode 100644 index 000000000..dd15ed09b --- /dev/null +++ b/tests/test_auth_login.py @@ -0,0 +1,492 @@ +"""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 + +import httpx +import pytest + +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, + forget_tokens, + get_access_token, + introspect_token, + lease_pat, + 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_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) + + 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 + + +@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('=') + 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') + + +# --- 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 + + +# --- 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'] + + +@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 new file mode 100644 index 000000000..5919978a3 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,409 @@ +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from keboola_mcp_server import auth_login +from keboola_mcp_server.auth_login import TokenSet +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' + + +@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: + args = parse_args(['migrate']) + assert args.command == 'migrate' + + +def test_parse_args_gc_sessions() -> None: + args = parse_args(['gc-sessions']) + 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, required=True) + + 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, required=True) + + 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, required=True) + + 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, 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 + 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']), + ), + 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() + 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: + 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(return_value={'created': [], 'dropped': []}), + ), + ): + 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')), + ), + pytest.raises(RuntimeError, match='boom'), + ): + await _run_migrate() + + 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')), + ), + pytest.raises(RuntimeError, match='boom'), + ): + 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')), + ), + pytest.raises(RuntimeError, match='boom'), + ): + 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_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() + 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/tests/test_config.py b/tests/test_config.py index 22bbf9e9e..32ad2a49a 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: @@ -34,6 +34,26 @@ 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'), + ), + ( + {'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: @@ -78,12 +98,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)' + "jwt_secret=None, postgres_dsn='****', session_encryption_key='****', " + 'bearer_token=None, conversation_id=None, project_id=None)' ) @pytest.mark.parametrize( @@ -110,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_mcp.py b/tests/test_mcp.py index 8525e5c0c..4bb119310 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,7 @@ import asyncio +import base64 +import dataclasses +import time from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -6,6 +9,7 @@ import pytest from fastmcp import Context from fastmcp.exceptions import ToolError +from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from pydantic import BaseModel, Field from starlette.requests import Request @@ -22,6 +26,14 @@ toon_serializer, unwrap_results, ) +from keboola_mcp_server.scope import ( + SCOPE_KEY, + SCOPE_TOKEN_ARG, + SessionScope, + resolve_scope_binding_aad, + resolve_scope_key, +) +from keboola_mcp_server.workspace import WorkspaceManager class SimpleModel(BaseModel): @@ -409,6 +421,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'), @@ -511,6 +541,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', @@ -686,7 +738,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(_): @@ -715,6 +767,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'), [ @@ -797,3 +892,865 @@ 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') + + 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', + 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, own_stack_storage_api_url=None) + + 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) + + @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). + + 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 + @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_at_abc', project_id=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, own_stack_storage_api_url=None + ) + + 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: + """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_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', 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') + 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 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).""" + + @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_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_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.project_id == '18' + 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) + 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_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) + 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 + + +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(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, self.KEY_B) + + 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_key_fallback_is_stable_within_process(self) -> None: + config = Config() + assert resolve_scope_key(config) == resolve_scope_key(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: + 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(base64.b64decode(key)), '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-token'}], + ids=['missing', 'none', 'empty', 'malformed'], + ) + def test_read_scope_from_request_returns_none_when_absent_or_invalid(self, arguments: dict) -> None: + 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 + + def test_read_scope_from_request_ignores_non_call_tool_requests(self) -> 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_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(self.KEY_A) + context = self._call_tool_context({'scope_token': token}) + config = Config(session_encryption_key=base64.b64encode(self.KEY_B).decode()) + assert SessionStateMiddleware._read_scope_from_request(context, config) is None + + def test_resolve_scope_binding_aad_is_none_locally(self, monkeypatch) -> None: + # A local `login` session already grants the whole stack to its one user -- no cross-caller + # boundary to bind, so this must stay a no-op there (see the docstring). + monkeypatch.delenv('KBC_KUBERNETES_TOKEN_PATH', raising=False) + assert resolve_scope_binding_aad('kbc_at_caller') is None + + def test_resolve_scope_binding_aad_is_none_without_a_token(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + assert resolve_scope_binding_aad(None) is None + + def test_resolve_scope_binding_aad_differs_per_caller_when_deployed(self, monkeypatch) -> None: + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + aad_a = resolve_scope_binding_aad('kbc_at_caller_a') + aad_b = resolve_scope_binding_aad('kbc_at_caller_b') + assert aad_a is not None and aad_b is not None + assert aad_a != aad_b + # Stable for the same caller token (mint and later read must agree). + assert resolve_scope_binding_aad('kbc_at_caller_a') == aad_a + + def test_round_trip_with_binding_aad(self) -> None: + scope = SessionScope(project_ids=[11], scoped_token='kbc_at_s', confirmed=True) + token = scope.to_token(self.KEY_A, aad=b'caller-a') + assert SessionScope.from_token(token, self.KEY_A, aad=b'caller-a') == scope + + def test_a_different_callers_token_is_rejected(self) -> None: + # The core replay fix: a scope_token minted while serving caller A's request must not be + # usable by caller B, even with the right key -- only decryptable alongside A's own + # storage token, which B (having only obtained the opaque string, not A's credential) + # cannot supply. + scope = SessionScope(project_ids=[11], scoped_token='kbc_at_s', confirmed=True) + token = scope.to_token(self.KEY_A, aad=b'caller-a-token') + with pytest.raises(Exception, match='.+'): + SessionScope.from_token(token, self.KEY_A, aad=b'caller-b-token') + + def test_replayed_scope_token_from_another_caller_falls_back_to_no_scope(self, monkeypatch) -> None: + # End-to-end: on a deployed server, a scope_token minted for caller A's storage token, + # replayed by caller B (a different storage token), must degrade to "no scope" via + # _read_scope_from_request -- not silently grant B caller A's scoped access. + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/token') + key = _b64_key(b'\x06') + config_a = Config(session_encryption_key=key, storage_token='kbc_at_caller_a') + scope = SessionScope(project_ids=[11], scoped_token='kbc_at_victim_scoped', confirmed=True) + token = scope.to_token(base64.b64decode(key), aad=resolve_scope_binding_aad(config_a.storage_token)) + + context = self._call_tool_context({'scope_token': token}) + config_b = Config(session_encryption_key=key, storage_token='kbc_at_caller_b') + assert SessionStateMiddleware._read_scope_from_request(context, config_b) is None + + # The legitimate caller (same storage token the scope was minted for) still works. + context = self._call_tool_context({'scope_token': token}) + assert SessionStateMiddleware._read_scope_from_request(context, config_a) == scope + + @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 + + 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 + # 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'] + + @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/test_multiproject.py b/tests/test_multiproject.py new file mode 100644 index 000000000..c75e25c3d --- /dev/null +++ b/tests/test_multiproject.py @@ -0,0 +1,675 @@ +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), + # 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_accessible_projects', True), # 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) + 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_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_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( + 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}'), + ), + 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}'), + ), + 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 + + +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_oauth.py b/tests/test_oauth.py index f80d305ef..83e765348 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -1,17 +1,116 @@ +import dataclasses +import logging +import secrets import time from collections.abc import Mapping +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 import pytest -from mcp.server.auth.provider import AccessToken, RefreshToken +from mcp.server.auth.provider import AccessToken, AuthorizationParams, RefreshToken, TokenError 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.auth_login import Introspection, ProjectAccess, ScopedToken +from keboola_mcp_server.clients.auth_bridge import OAuthTokenExchangeError +from keboola_mcp_server.oauth import ( + ProxyRefreshToken, + SimpleOAuthProvider, + _ExtendedAuthorizationCode, + _OAuthClientInformationFull, +) +from keboola_mcp_server.session_store.repository import OAuthSession 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.""" + + 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: @@ -24,6 +123,7 @@ def oauth_provider(self) -> SimpleOAuthProvider: server_url='https://oauth', scope='scope', jwt_secret=JWT_KEY, + session_store=FakeSessionStore(), ) @staticmethod @@ -217,3 +317,356 @@ 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'] + + @staticmethod + def _stub_exchanger(monkeypatch: pytest.MonkeyPatch, captured: dict[str, Any]) -> None: + from keboola_mcp_server import oauth as oauth_module + + 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) + + @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()) + + 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' + # 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' + # 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 + 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( + 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') + + 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()) + + # 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.error == 'invalid_grant' + + @pytest.mark.asyncio + async def test_exchange_authorization_code_missing_sa_token_path( + 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: None) + + client = _OAuthClientInformationFull(redirect_uris=[AnyHttpUrl('http://foo')], client_id='foo-client-id') + auth_code = _ExtendedAuthorizationCode.model_validate(self.authorization_code()) + + with pytest.raises(TokenError) as exc: + await oauth_provider.exchange_authorization_code(client, auth_code) + assert exc.value.error == 'invalid_request' + + @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') + _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=None, + kbc_refresh_token='kbc_rt_old', + session_id=session.id, + ) + + 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' + 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_token_error( + self, oauth_provider: SimpleOAuthProvider, monkeypatch: pytest.MonkeyPatch + ): + 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') + _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=None, + kbc_refresh_token='kbc_rt_old', + session_id=session.id, + ) + + # 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.error == 'invalid_grant' + + @pytest.mark.asyncio + async def test_load_access_token_refreshes_near_expiry_session_transparently( + 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 + + 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) + + 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' + # 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' + # 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( + 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 02b025a64..4dcfacadc 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 @@ -14,6 +15,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 +28,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 @@ -60,6 +62,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 +86,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 +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() @@ -207,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', ), {}, ), @@ -215,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. @@ -296,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', } @@ -503,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', @@ -597,3 +607,73 @@ 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' + + +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( + 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(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', session_encryption_key=self._TEST_ENCRYPTION_KEY + ), + 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/tests/test_workspace.py b/tests/test_workspace.py index 766fee096..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'} @@ -171,6 +172,62 @@ async def test_workspace_creation_cleans_up_config_on_failure(): ) +@pytest.mark.asyncio +@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 + 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'} + mock_storage_client.workspace_create_for_config.return_value = {'id': 999} + 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 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_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'} + 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'), @@ -416,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() diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index 3e0ea2497..d925b9581 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -1,19 +1,36 @@ +import time +from types import SimpleNamespace + +import httpx import pytest from mcp.server.fastmcp import Context 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, 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_binding_aad, + resolve_scope_key, +) from keboola_mcp_server.tools.project import ( ProjectInfo, _get_toolset_restrictions, + _parent_subject_token, _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 +263,440 @@ 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_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) + 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) + ) + # 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(), runtime_info=ServerRuntimeInfo(transport='stdio')), + ) + verify_info = {18: ('BigQuery', 'org-1', 'Org One'), 83: ('Snowflake', 'org-2', 'Org Two')} + mocker.patch( + '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. + 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, 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 + 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. 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 None + + +@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_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) + + 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.base_instructions} + assert groups == {'BigQuery': [18, 86], 'Snowflake': [95]} + 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_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) + + 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_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_verify_info', + 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 +) -> 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] + scope = mcp_context_client.session.state[SCOPE_KEY] + assert scope.scoped_token == 'kbc_at_scoped' + assert scope.project_ids == [18, 83] + # 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_key(Config())) == scope + assert 'does not remember this scope' in result.llm_instruction + + +@pytest.mark.asyncio +async def test_set_project_scope_binds_scope_token_to_caller_on_deployed_server( + mcp_context_client: Context, mocker: MockerFixture, monkeypatch +) -> None: + # The replay fix: on a deployed server, the returned scope_token must only decrypt alongside + # the same caller's own storage token (client.token) it was minted for -- see + # resolve_scope_binding_aad. A different caller's token must fail, even with the right key. + monkeypatch.setenv('KBC_KUBERNETES_TOKEN_PATH', '/var/run/secrets/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] + key = resolve_scope_key(Config()) + assert SessionScope.from_token(result.scope_token, key, aad=resolve_scope_binding_aad('test-token')) == scope + with pytest.raises(Exception, match='.+'): + SessionScope.from_token(result.scope_token, key, aad=resolve_scope_binding_aad('kbc_at_someone_else')) + + +@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_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 +) -> 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 +@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_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', + [ + 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 +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=[]) diff --git a/uv.lock b/uv.lock index 2c80fe7ec..f39577d30 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 = [ @@ -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" @@ -990,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 = [ @@ -1172,9 +1240,10 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.75.3" +version = "1.76.3" 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" }, @@ -2258,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 = [