feat(AI-3669): support pinning MCP queries to an explicit workspace ID

README.md

@@ -159,6 +159,12 @@ This identifies your workspace in Keboola and is used for SQL queries. However,
 
 **Note**: KBC_WORKSPACE_SCHEMA is called Dataset Name in BigQuery workspaces, you simply click connect and copy the Dataset Name
 
+### KBC_WORKSPACE_ID
+
+Pins queries to one specific, already-existing workspace by its ID instead of the schema-based lookup above, and takes precedence over `KBC_WORKSPACE_SCHEMA` when both are set. This is the option a Data App / kai-agent caller supplies, as the `X-Workspace-Id` header, so that Kai embedded in that app queries only through its own workspace.
+
+Set via the `KBC_WORKSPACE_ID` environment variable, the `--workspace-id` CLI flag, or (per-request, for multi-user deployments) the `X-Workspace-Id` header.
+
 ### KBC_STORAGE_API_URL (Keboola Region)
 
 Your Keboola Region API URL depends on your deployment region. You can determine your region by looking at the URL in your browser when logged into your Keboola project:
@@ -438,7 +444,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 |
+| **Workspace Issues** | Confirm `KBC_WORKSPACE_SCHEMA` / `KBC_WORKSPACE_ID` is correct |
 | **Connection Timeout** | Check network connectivity |
 
 ## Development

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "keboola-mcp-server"
-version = "1.75.4"
+version = "1.76.0"
 description = "MCP server for interacting with Keboola Connection"
 readme = "README.md"
 requires-python = ">=3.10"

src/keboola_mcp_server/cli.py

@@ -53,6 +53,7 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace:
     )
     parser.add_argument('--storage-token', metavar='STR', help='Keboola Storage API token.')
     parser.add_argument('--workspace-schema', metavar='STR', help='Keboola Storage API workspace schema.')
+    parser.add_argument('--workspace-id', metavar='STR', help='Keboola Storage API workspace ID.')
     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('--log-config', type=pathlib.Path, metavar='PATH', help='Logging config file.')
@@ -129,6 +130,7 @@ async def run_server(args: list[str] | None = None) -> None:
         storage_api_url=parsed_args.api_url,
         storage_token=parsed_args.storage_token,
         workspace_schema=parsed_args.workspace_schema,
+        workspace_id=parsed_args.workspace_id,
     )
 
     try:

src/keboola_mcp_server/config.py

@@ -27,6 +27,18 @@ class Config:
     """The branch ID to access the storage API using the MCP tools."""
     workspace_schema: str | None = None
     """Workspace schema to access the buckets, tables and execute sql queries."""
+    workspace_id: str | None = field(default=None, metadata={'empty_means_absent': True, 'require_prefix': True})
+    """Workspace ID to access the buckets, tables and execute sql queries (e.g. a Data App's own
+    workspace, supplied per-request via the 'X-Workspace-Id' header). Takes precedence over
+    `workspace_schema` when both are set.
+
+    `require_prefix` is set because the bare `WORKSPACE_ID` env var is what Keboola injects into
+    Data App containers -- without it, that variable would pin every session on such a server.
+    `empty_means_absent` is set so an unset header template (`X-Workspace-Id:`) forwarded as an
+    empty string is not mistaken for an explicit pin to override a server-side default with. This
+    intentionally does NOT apply to `workspace_schema` (an empty `X-Workspace-Schema` header must
+    keep clearing it back to the MCP-managed workspace, same as `branch_id` below) nor to most
+    other fields -- it is opt-in per field precisely to avoid that kind of regression."""
     oauth_client_id: str | None = None
     """OAuth client ID registered in the Keboola OAuth Server."""
     oauth_client_secret: str | None = None
@@ -67,6 +79,9 @@ def __post_init__(self) -> None:
         if self.branch_id is not None and self.branch_id.lower() in ['', 'none', 'null', 'default', 'production']:
             object.__setattr__(self, 'branch_id', None)
 
+        if self.workspace_id is not None and not self.workspace_id.isdigit():
+            raise ValueError(f'Invalid workspace_id: {self.workspace_id!r}')
+
     @staticmethod
     def _normalize(name: str) -> str:
         """Removes dashes and underscores from the input string and turns it into lowercase."""
@@ -79,19 +94,27 @@ def _read_options(cls, d: Mapping[str, str]) -> Mapping[str, Any]:
         for f in dataclasses.fields(cls):
             field_names = [f.name] + f.metadata.get('aliases', [])
 
+            require_prefix = f.metadata.get('require_prefix', False)
+            empty_means_absent = f.metadata.get('empty_means_absent', False)
+
             for name in field_names:
                 value: str | None = _NO_VALUE_MARKER
-
-                if (dict_name := cls._normalize(name)) in data:
-                    value = data[dict_name]
-
-                elif (dict_name := cls._normalize(f'KBC_{name}')) in data:
-                    # environment variables start with KBC_
-                    value = data[dict_name]
-
-                elif (dict_name := cls._normalize(f'X-{name}')) in data:
-                    # HTTP headers start with X-
-                    value = data[dict_name]
+                # `require_prefix` skips the bare field name -- only the KBC_/X- prefixed forms
+                # count as "provided". Needed for fields whose bare name collides with a
+                # variable set by something other than this server's own config (see
+                # `workspace_id`'s docstring).
+                candidates = (f'KBC_{name}', f'X-{name}') if require_prefix else (name, f'KBC_{name}', f'X-{name}')
+
+                for candidate in candidates:
+                    if (dict_name := cls._normalize(candidate)) in data:
+                        candidate_value = data[dict_name]
+                        # An empty value means "not provided" for opted-in fields only -- an
+                        # unset header template (e.g. `X-Workspace-Id:`) must not be mistaken
+                        # for an explicit request to override a server-side default with ''.
+                        if empty_means_absent and candidate_value == '':
+                            continue
+                        value = candidate_value
+                        break
 
                 if value is not _NO_VALUE_MARKER:
                     if f.type == (bool | None):
@@ -118,8 +141,18 @@ 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_'.
+
+        A malformed `workspace_id` (the only field `__post_init__` validates) degrades to "not
+        provided" rather than raising: `d` is untrusted per-request input (an HTTP header), so a
+        junk value from a client should drop the pin, not turn into an unhandled server error.
         """
-        return dataclasses.replace(self, **self._read_options(d))
+        options = self._read_options(d)
+        try:
+            return dataclasses.replace(self, **options)
+        except ValueError as e:
+            LOG.warning(f'Ignoring invalid request header value(s): {e}')
+            options.pop('workspace_id', None)
+            return dataclasses.replace(self, **options)
 
     def __repr__(self) -> str:
         params: list[str] = []

src/keboola_mcp_server/mcp.py

@@ -288,8 +288,29 @@ 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}')
+        server_config = config
         config = config.replace_by(http_rq.headers)
 
+        # A workspace pin (`workspace_id`/`workspace_schema`) configured on the server fences
+        # off which project data a request can reach -- mirror the `own_stack_storage_api_url`
+        # reasoning above: a deployment that pinned itself must not be overridable by a header.
+        # A server with no pin of its own (the shared multi-tenant case, e.g. AJDA-3052's Data
+        # App flow) keeps taking the pin from the request, which is the only source it has.
+        if (server_config.workspace_id or server_config.workspace_schema) and (
+            config.workspace_id != server_config.workspace_id
+            or config.workspace_schema != server_config.workspace_schema
+        ):
+            LOG.warning(
+                f'Ignoring the requested workspace pin (workspace_id={config.workspace_id!r}, '
+                f'workspace_schema={config.workspace_schema!r}); this server is pinned to '
+                f'workspace_id={server_config.workspace_id!r}, workspace_schema={server_config.workspace_schema!r}.'
+            )
+            config = dataclasses.replace(
+                config,
+                workspace_id=server_config.workspace_id,
+                workspace_schema=server_config.workspace_schema,
+            )
+
         if own_stack_storage_api_url and not is_same_stack(config.storage_api_url, own_stack_storage_api_url):
             LOG.warning(
                 f'Ignoring the requested Storage API URL "{config.storage_api_url}"; '
@@ -367,7 +388,10 @@ async def create_session_state(
             # therefore attaches the JWT only when the target is this server's own stack.
             kubernetes_token_path = os.environ.get('KBC_KUBERNETES_TOKEN_PATH')
             workspace_manager = await WorkspaceManager.create(
-                client, config.workspace_schema, kubernetes_token_path=kubernetes_token_path
+                client,
+                config.workspace_schema,
+                kubernetes_token_path=kubernetes_token_path,
+                workspace_id=config.workspace_id,
             )
             state[WorkspaceManager.STATE_KEY] = workspace_manager
             LOG.info('Successfully initialized Storage API Workspace manager.')

src/keboola_mcp_server/tools/data_apps.py

@@ -552,9 +552,9 @@ async def modify_streamlit_data_app(
     links_manager = await ProjectLinksManager.from_client(client)
 
     project_id = await client.storage_client.project_id()
-    workspace_id = await workspace_manager.get_workspace_id()
-    sql_dialect = await workspace_manager.get_sql_dialect()
-    branch_id = await workspace_manager.get_branch_id()
+    workspace_id = await workspace_manager.get_data_app_workspace_id()
+    sql_dialect = await workspace_manager.get_data_app_sql_dialect()
+    branch_id = await workspace_manager.get_data_app_branch_id()
 
     secrets = _get_secrets(
         workspace_id=str(workspace_id),
@@ -708,8 +708,8 @@ async def modify_streamlit_data_app_internal(
     folder: str | None = None,
 ) -> tuple[DataApp, JsonDict, dict | None]:
     secrets = _get_secrets(
-        workspace_id=str(await workspace_manager.get_workspace_id()),
-        branch_id=str(await workspace_manager.get_branch_id()),
+        workspace_id=str(await workspace_manager.get_data_app_workspace_id()),
+        branch_id=str(await workspace_manager.get_data_app_branch_id()),
     )
     data_app = await _fetch_data_app(client, configuration_id=configuration_id, data_app_id=None)
     existing_config = data_app.configuration
@@ -720,7 +720,7 @@ async def modify_streamlit_data_app_internal(
         packages,
         authentication_type,
         secrets,
-        await workspace_manager.get_sql_dialect(),
+        await workspace_manager.get_data_app_sql_dialect(),
     )
     updated_config = cast(
         JsonDict,
@@ -1096,7 +1096,7 @@ async def modify_python_js_data_app(
     legacy_secrets: dict[str, Any] | None = None
     if not has_storage_workspace:
         workspace_manager = WorkspaceManager.from_state(ctx.session.state)
-        legacy_secrets = {SECRET_WORKSPACE_ID: str(await workspace_manager.get_workspace_id())}
+        legacy_secrets = {SECRET_WORKSPACE_ID: str(await workspace_manager.get_data_app_workspace_id())}
 
     if configuration_id:
         # Update existing python-js data app