feat(AI-3669): support pinning MCP queries to an explicit workspace ID - #654
feat(AI-3669): support pinning MCP queries to an explicit workspace ID#654Matovidlo wants to merge 6 commits into
Conversation
|
@claude review |
There was a problem hiding this comment.
Pull request overview
This PR adds support for pinning workspace resolution to an explicit Keboola workspace ID (e.g., when running inside a Data App) by introducing workspace_id into server configuration and threading it into WorkspaceManager, where it takes precedence over workspace_schema.
Changes:
- Add
Config.workspace_idwith per-request override via HTTP header and CLI flag, and pass it into session/workspace initialization. - Update
WorkspaceManagerto resolve a workspace by ID (with precedence over schema) and raise if the ID is not found. - Add unit tests for config parsing and workspace resolution-by-id behavior; bump project version to
1.75.0.
Reviewed changes
Copilot reviewed 7 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/keboola_mcp_server/config.py |
Adds workspace_id to server config model and documents precedence over workspace_schema. |
src/keboola_mcp_server/cli.py |
Adds --workspace-id CLI option and wires it into Config construction. |
src/keboola_mcp_server/mcp.py |
Threads workspace_id from request/config into WorkspaceManager.create() during session setup. |
src/keboola_mcp_server/workspace.py |
Implements workspace lookup by explicit ID with precedence and a not-found error. |
tests/test_config.py |
Extends config parsing tests to cover the new workspace_id field. |
tests/test_workspace.py |
Adds unit tests for _get_workspace() resolving by ID and failing when not found. |
pyproject.toml |
Bumps project version to 1.75.0. |
uv.lock |
Updates locked project version metadata to 1.75.0. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
4a4aa06 to
476e56b
Compare
keboola-pr-reviewer-bot
left a comment
keboola-pr-reviewer-bot
left a comment
There was a problem hiding this comment.
Verdict: needs_human (risk 3/5) · profile keboola-mcp-server
Adds an opt-in, per-request-header-settable workspace_id that redirects which workspace SQL queries target — a load-bearing surface worth a human glance.
Concerns:
src/keboola_mcp_server/workspace.py: New X-Workspace-Id header redirects which workspace SQL queries execute against.src/keboola_mcp_server/config.py: Per-request header-settable workspace selector on the auth boundary.
Suggested reviewers: @keboola/ai-swimlane-kai-assistant
|
@MiroCillik can you check it please ? 🙏 |
|
@Matovidlo and please, so the Workspace id will be passed to KAI (and then to MCP?) on each request (question/prompt)? Or when is the MCP server instantiated? |
|
@MiroCillik the Kai is connected to MCP during conversation, it start with it and during it it should be working with workspace. For Kai in client invoked in data app, the lifecykle should be somewhere on start of the app and end with its decrease |
MiroCillik
left a comment
MiroCillik
left a comment
There was a problem hiding this comment.
Review: pinning MCP queries to an explicit workspace ID
Reviewed at high effort (multi-agent finder + independent adversarial verification of every candidate). The mechanism is a clean mirror of the existing workspace_schema / X-Workspace-Schema path and the plumbing is correct; conventions all check out (minor bump 1.75.2 → 1.76.0, uv.lock synced, no tool signature change so TOOLS.md is unaffected).
10 findings are left as inline comments. Almost all of them trace back to one design decision rather than sloppy code: the pin lands on the session-wide WorkspaceManager rather than on the query path, and it is selected by an untrusted per-request header with no ownership, read-only, or branch validation.
What I'd fix before merging
get_workspace_id()is not query-only — the pinned id also reachestools/data_apps.py, where it gets written into other data apps' persistedWORKSPACE_IDsecret. This one is a concrete data-corruption bug, not a hypothetical (see comment onmcp.py:373).if self._workspace_id:is a truthiness check — an emptyX-Workspace-Idheader silently un-pins the session back to the unrestricted project-wide workspace. A restriction that disappears silently is worse than one that fails loudly.- No
readonlyvalidation on the pinned workspace — every other resolution path enforces read-only storage access (_find_ws_in_branchrequiresinfo.readonly,_create_wsalways passesread_only_storage_access=True) andquery_datahas no statement-type guard, so a pinned writable workspace turns the agent's SQL path read-write. - A bare
WORKSPACE_IDenv var (noKBC_prefix) pins every session — and that is exactly the variable Keboola injects into Data App containers, which is the deployment scenario this feature targets.
The rest (branch-scoped workspace_detail lookups, non-404 statuses bypassing the intended ValueError, _WspInfo credentials in an INFO log, the untested create() pass-through, missing README docs) are individually smaller but cheap to address.
Design note — this is the transport for the restriction, not the restriction
Because the workspace is chosen by a header the client supplies, the restriction is only as trustworthy as whoever sets it. That is the same trust model X-Workspace-Schema already had, so it is not a regression — but for the AJDA-3052 use case ("don't give data-app users a backdoor to all org data") a client-supplied header is not an enforcement boundary. The real boundary has to be the token's own grants, which is what keboola/connection#7981 / #7985 are building.
Explicitly out of scope for this PR: limiting the token itself. Note that this also means the pin narrows only query_data — bucket/table/column metadata still flows through the caller's Storage API token (get_tables, get_bucket_detail, get_project_info) and is unaffected, so the agent can still enumerate the whole project's structure. Worth stating in the PR description so the initiative doesn't assume it's covered here.
Process
The PR body is missing the mandated Release Notes section (Justification / Plans for Customer Communication / Impact Analysis / Deployment Plan / Rollback Plan / Post-Release Support Plan).
| client, | ||
| config.workspace_schema, | ||
| kubernetes_token_path=kubernetes_token_path, | ||
| workspace_id=config.workspace_id, |
There was a problem hiding this comment.
The pin is session-wide, not query-scoped — it leaks into other data apps' secrets.
get_workspace_id() (workspace.py:912) just returns (await self._get_workspace()).id, and it is used well beyond query execution: tools/data_apps.py:555 and :711 feed it into _get_secrets(workspace_id=...), and :1099 into legacy_secrets = {SECRET_WORKSPACE_ID: ...} — i.e. it is written into a data app's persisted configuration.
Failure scenario: Kai runs inside Data App A with the session pinned via X-Workspace-Id: <A's workspace>. A user asks it to create or update Data App B → B's config is written with A's workspace id as SECRET_WORKSPACE_ID. Deployed app B then queries through app A's workspace (wrong data scope, possibly broader than intended) and breaks permanently the moment app A or its workspace is deleted. get_branch_id() (workspace.py:916) leaks the same way.
Suggestion: scope the pin to the query path instead of the whole WorkspaceManager — e.g. keep the default-resolution workspace for get_workspace_id()/get_branch_id() consumers and use the pinned one only for execute_query/get_table_info/dialect resolution. If a session-wide pin is intentional, then data_apps.py must stop deriving the app's workspace id from the session manager.
There was a problem hiding this comment.
Proposed fix
The pin belongs to the query path; data_apps.py needs the MCP-managed workspace, i.e. what it got before this PR. Split the resolution rather than patching call sites ad hoc:
async def _get_workspace(self) -> _Workspace:
"""The workspace queries run against — honours an explicit `workspace_id` pin."""
if self._workspace:
return self._workspace
if self._workspace_id is not None:
self._workspace = self._init_workspace(await self._resolve_pinned())
return self._workspace
return await self._get_managed_workspace()
async def _get_managed_workspace(self) -> _Workspace:
"""The MCP-managed (or `workspace_schema`-pinned) workspace, ignoring `workspace_id`.
Data apps persist this id into their own configuration (`SECRET_WORKSPACE_ID`), so it must
never be the caller-supplied pin of whichever session happened to create the app.
"""
if self._managed_workspace:
return self._managed_workspace
... # existing workspace_schema -> _find_ws_in_branch -> _create_ws bodyThen add an explicit accessor and switch the leaking call sites:
async def get_data_app_workspace_id(self) -> int:
return (await self._get_managed_workspace()).idtools/data_apps.py:555,:711,:1099→get_data_app_workspace_id()- same treatment for
get_branch_id()at those three sites
_get_managed_workspace keeps honouring workspace_schema, so existing --workspace-schema deployments are unchanged and the whole fix is additive. Worth adding a docstring line on get_workspace_id() stating it is pin-aware and must not feed anything persisted — that's what stops a fourth call site reintroducing this.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Implemented essentially as you suggested, keeping get_workspace_id()/get_branch_id() pin-aware for their existing callers (tools/project.py:201, tools/storage/tools.py:722, which report the workspace actually being queried) rather than repointing them -- so instead I added get_data_app_workspace_id()/get_data_app_branch_id() backed by a new _get_managed_workspace() (your _get_managed_workspace, same split), and switched all three tools/data_apps.py call sites (555/557, 711/712, 1099) to the new accessors. Docstrings on all four methods now state which is pin-aware.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread you left on this same line with the concrete diff (mcp.py:373, second review pass).
| if self._workspace: | ||
| return self._workspace | ||
|
|
||
| if self._workspace_id: |
There was a problem hiding this comment.
Truthiness check → an empty header silently un-pins the workspace.
Config._read_options stores any present header value verbatim, including '', and mcp.py:291 applies request headers on top of the server config. So a deployment with KBC_WORKSPACE_ID=<app workspace> that receives a request with an empty X-Workspace-Id: (an unset template value on the kai-client / Data App side — a very common failure mode) ends up with workspace_id=''.
if self._workspace_id: is then falsy, this branch is skipped entirely, and execution falls through to _find_ws_in_branch() / _create_ws() — the MCP-managed workspace with read access to every bucket in the project. No error, no warning: the in-app data-access restriction is silently gone for that request.
| if self._workspace_id: | |
| if self._workspace_id is not None: |
...plus rejecting a blank value explicitly, e.g. normalise '' → None in Config.__post_init__ (as it already does for branch_id) and raise if a pin was expected but is empty. Given the security intent, failing loudly is the right default here.
There was a problem hiding this comment.
Proposed fix
Fix it at the config layer rather than at this if, so the un-pin cannot happen at all and every field benefits. Collapse the elif chain in Config._read_options into a candidate loop:
for name in field_names:
value: str | None = _NO_VALUE_MARKER
for dict_name in (name, f'KBC_{name}', f'X-{name}'):
# An empty header/env value means "not provided": an unset header template
# (`X-Workspace-Id:`) must not override a server-side pin with ''.
if (key := cls._normalize(dict_name)) in data and data[key] != '':
value = data[key]
breakThen keep if self._workspace_id is not None: here as defense in depth. Add empty-value cases to the existing test_from_dict / test_replace_by parametrize lists.
There was a problem hiding this comment.
Alternative approach 💡
Fixed at the config layer, but scoped rather than applied to every field: a blanket empty-means-absent change breaks the existing branch_id behavior ({'branch-id': ''} -> Config() is a real, tested feature -- clearing a server-configured branch via an empty header). Added an opt-in empty_means_absent field metadata flag instead, set on workspace_id and workspace_schema only. Combined with the require_prefix flag from the config.py:30 thread, _read_options now loops candidates per-field with both opt-ins honored. Kept is not None in workspace.py as defense in depth. Added the empty-header case to test_config.py.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (workspace.py:830, second review pass).
| # use the workspace that was explicitly requested (e.g. a Data App's own workspace) | ||
| # this workspace must never be written to the default branch metadata | ||
| LOG.info(f'Looking up workspace by id: {self._workspace_id}') | ||
| if info := await self._find_ws_by_id(self._workspace_id): |
There was a problem hiding this comment.
The pinned workspace is not validated as read-only — the header can widen access from read-only to read-write.
Every other resolution path enforces the read-only invariant:
_find_ws_in_branch(line 705):if info.id and info.backend and info.schema and info.readonly_create_ws: always provisions withread_only_storage_access=True
_find_ws_by_id checks only wi.id and wi.backend and wi.schema — never wi.readonly. And tools/sql.py has no statement-type guard (no SELECT-only check anywhere), so read-only storage access is the only thing preventing query_data from mutating data.
Failure scenario: X-Workspace-Id: <id of a writable transformation/sandbox workspace in the project> yields a read-write workspace, and the agent's next query_data call can run INSERT / UPDATE / DROP TABLE against project data. For a feature whose purpose is narrowing access, that inversion is worth an explicit guard:
if info := await self._find_ws_by_id(self._workspace_id):
if not info.readonly:
raise ValueError(f'Workspace {self._workspace_id} has no read-only storage access.')There was a problem hiding this comment.
Proposed fix — but decide the policy first
Which fix is correct depends on a fact I can't get from the repo: whether a Data App's platform-provisioned workspace actually has readOnlyStorageAccess: true. Worth checking before coding:
curl -s -H "X-StorageApi-Token: $TOKEN" \
"$KBC_URL/v2/storage/branch/default/workspaces/<data-app-workspace-id>" | jq '.readOnlyStorageAccess'- If
true→ enforce it, matching_find_ws_in_branch:if not info.readonly: raise ValueError(f'Workspace {self._workspace_id} has no read-only storage access.')
- If
false→ do not enforce; it would break the feature outright. InsteadLOG.warningon a writable pin and open a follow-up for aSELECT-only guard intools/sql.py, which is the real missing control here — today nothing stopsquery_datarunning DDL/DML against any writable workspace, pin or no pin.
Either way it's worth a line in the PR description, since it's the difference between this header narrowing access and widening it.
There was a problem hiding this comment.
Question ❓
Went with the warn-not-raise fallback for now -- I have no way to check readOnlyStorageAccess on a real Data App workspace from here. Added a LOG.warning('... has no read-only storage access.') plus a comment noting: promote to a raise once confirmed read-only, and either way the missing SELECT-only guard in tools/sql.py is a real follow-up regardless of this policy. Can you (or whoever has access to a stack with a live Data App workspace) run the curl check and confirm? Happy to flip this to the hard raise in a follow-up commit once we know.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Addressed with a warning for now, pending confirmation against a real workspace -- see the detailed reply on the follow-up thread (workspace.py:834, second review pass).
| # use the workspace that was explicitly requested (e.g. a Data App's own workspace) | ||
| # this workspace must never be written to the default branch metadata | ||
| LOG.info(f'Looking up workspace by id: {self._workspace_id}') | ||
| if info := await self._find_ws_by_id(self._workspace_id): |
There was a problem hiding this comment.
Non-404 statuses bypass the intended ValueError and surface as a raw HTTPStatusError.
_find_ws_by_id (line 669) maps only 404 → None and re-raises everything else. The header value is unvalidated and interpolated straight into the Storage API path by workspace_detail, whose signature declares workspace_id: int.
So X-Workspace-Id: my-app-workspace (non-numeric, or a typo) gets a 400 from Storage API, which propagates out of _get_workspace(), and every data tool fails with an opaque:
Client error '400 Bad Request' for url 'https://connection.../v2/storage/branch/default/workspaces/my-app-workspace'
...rather than the actionable No Keboola workspace found: workspace_id=... this branch was written to produce. Same for a 403 on a workspace the token may not read. Worth validating the id is an int at the Config boundary, and/or treating 400/403 like 404 here.
There was a problem hiding this comment.
Proposed fix
Two layers. First, map the non-404 statuses where the lookup happens (pairs well with the branch-fallback refactor suggested on line 839 — both want the HTTP call split into a _fetch_ws helper):
except HTTPStatusError as e:
if e.response.status_code in (400, 403, 404):
LOG.warning(f'Workspace lookup failed: id={workspace_id}, status={e.response.status_code}')
return None
raiseSecond, validate the shape at the config boundary so a typo'd id fails with something readable instead of reaching the Storage API at all — in Config.__post_init__, alongside the existing branch_id normalisation:
if self.workspace_id is not None and not str(self.workspace_id).isdigit():
raise ValueError(f'Invalid workspace_id: {self.workspace_id!r}')There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Both layers added: 400/403 now map to None in a shared _fetch_ws(client, workspace_id) helper (also used by the branch-fallback fix on the :870/:839 thread), and Config.__post_init__ now raises Invalid workspace_id: ... when it isn't all-digits, alongside the existing branch_id normalization.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (workspace.py:834, second review pass).
| # this workspace must never be written to the default branch metadata | ||
| LOG.info(f'Looking up workspace by id: {self._workspace_id}') | ||
| if info := await self._find_ws_by_id(self._workspace_id): | ||
| LOG.info(f'Found workspace: {info}') |
There was a problem hiding this comment.
This logs the whole _WspInfo, whose credentials field is the workspace's backend credential blob.
_WspInfo is a plain frozen dataclass with the default __repr__, so {info} prints every field — including credentials, which for BigQuery is a serialized service-account JSON (as _init_workspace shows by doing json.loads(info.credentials)).
On a BigQuery-backend project, the first request carrying X-Workspace-Id writes that service-account JSON — private_key included — into the application logs at INFO level, from where it ships to Datadog. Anyone with log-read access then has long-lived credentials to that dataset and can read project data entirely outside Keboola, with no MCP audit trail.
Note this line is reachable by a caller-chosen id now, which is what changes the exposure profile. Suggest logging only the safe fields:
LOG.info(f'Found workspace: id={info.id}, schema={info.schema}, backend={info.backend}, readonly={info.readonly}')(The identical pre-existing LOG.info(f'Found workspace: {info}') in the schema path has the same problem and is worth fixing in the same pass.)
There was a problem hiding this comment.
Proposed fix
Fix the type rather than the two call sites — that also covers the pre-existing schema-path log and any future interpolation of _WspInfo:
@dataclass(frozen=True, repr=False)
class _WspInfo:
...
def __repr__(self) -> str:
return (
f'_WspInfo(id={self.id!r}, schema={self.schema!r}, backend={self.backend!r}, '
f'credentials={"****" if self.credentials else None}, readonly={self.readonly!r})'
)(repr=False is needed so the dataclass machinery doesn't generate one over it.)
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Implemented exactly as suggested: @dataclass(frozen=True, repr=False) + a custom __repr__ redacting credentials. Every existing LOG.info(f'... {info}') call site (id path, schema path, default-branch path) is covered by this single fix, no per-call-site changes needed. Added a unit test asserting the secret never appears in repr(info).
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (workspace.py:835, second review pass).
| LOG.info(f'Looking up workspace by id: {self._workspace_id}') | ||
| if info := await self._find_ws_by_id(self._workspace_id): | ||
| LOG.info(f'Found workspace: {info}') | ||
| self._workspace = self._init_workspace(info) |
There was a problem hiding this comment.
No ownership validation, and it silently overrides the operator's own pin.
The resolved workspace is accepted with no check that it is the caller's / the app's workspace — any workspace id visible to the project token works. Combined with the documented precedence over workspace_schema, this means a request header beats server-side deployment configuration: apply_request_config() deliberately protects storage_api_url from headers, but nothing protects the workspace pin.
Failure scenario: a server deployed with KBC_WORKSPACE_SCHEMA=<restricted app schema> (the current way a deployment fences off data access) receives X-Workspace-Id: <another workspace in the same project>. The header wins, _get_workspace() returns the other workspace, data the operator intentionally fenced off becomes readable through query_data, and the audit trail shows a normal tool call.
At minimum: LOG.warning when a header-supplied workspace_id overrides a server-configured workspace_schema/workspace_id, and consider treating a server-side pin as authoritative (same reasoning as own_stack_storage_api_url).
There was a problem hiding this comment.
Proposed fix
Mirror the own_stack_storage_api_url reasoning in apply_request_config: a deployment that pinned itself shouldn't be overridable by a header, while the shared-server case — no server-side pin, header supplies it, which is the actual AJDA-3052 flow — keeps working:
server_config = config
config = config.replace_by(http_rq.headers)
if server_config.workspace_id or server_config.workspace_schema:
if config.workspace_id != server_config.workspace_id:
LOG.warning(f'Ignoring requested workspace_id "{config.workspace_id}"; server is pinned.')
config = dataclasses.replace(
config,
workspace_id=server_config.workspace_id,
workspace_schema=server_config.workspace_schema,
)
elif config.workspace_id and config.workspace_schema:
LOG.warning(f'Both workspace_id and workspace_schema supplied; using workspace_id={config.workspace_id}.')There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Implemented your apply_request_config snippet, with one small extension: made it symmetric over both workspace_id and workspace_schema (checked/restored together) since the same silent-override risk applies to a schema-based server pin, not just an id-based one. LOG.warning fires whenever either is overridden; a server with no pin of its own keeps taking either from the request headers unchanged.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (workspace.py:836, second review pass).
| self._workspace = self._init_workspace(info) | ||
| return self._workspace | ||
| else: | ||
| raise ValueError(f'No Keboola workspace found: workspace_id={self._workspace_id}') |
There was a problem hiding this comment.
workspace_detail is branch-scoped, so a cross-branch pin 404s and this error message is misleading.
AsyncStorageClient.workspace_detail hits branch/{self._branch_id}/workspaces/{workspace_id} (clients/storage.py:1111), where _branch_id defaults to 'default' (storage.py:321). On a project with the storage-branches feature, create() keeps the dev-branch client (line 592), so a session bound to dev branch 123 looks the pinned id up under branch/123/workspaces/<id>.
A Data App workspace that lives in the default/production branch → 404 → _find_ws_by_id returns None → this ValueError fires, and every query_data, get_table_detail (storage/tools.py:722), get_project_info (project.py:201) and data-app tool call in that session fails. The message says the workspace does not exist when it does. The mirror case (dev-branch workspace, default-branch session) fails too.
Worth noting _find_ws_by_id's only previous caller — _create_ws (line 794) — was inherently branch-safe because it looked up a workspace it had just created with the same client, so this is a new exposure for the helper. Either resolve the pinned id against the prod client explicitly, or make the error say which branch was searched.
There was a problem hiding this comment.
Proposed fix
Split the HTTP call into a helper and fall back to the default branch:
async def _find_ws_by_id(self, workspace_id: str | int) -> _WspInfo | None:
info = await self._fetch_ws(self._client, workspace_id)
if info is None and self._client.branch_id is not None:
# `workspace_detail` is branch-scoped; a pinned workspace (e.g. a Data App's own)
# may live in the default branch while this session is bound to a dev branch.
info = await self._fetch_ws(await self._client.with_branch_id(None), workspace_id)
return info_fetch_ws(client, id) is the current try/except body, and is also where the 400/403 mapping from the other thread belongs. And include the branch in the error message here so it stops being misleading:
raise ValueError(f'No Keboola workspace found: workspace_id={self._workspace_id}, branch_id={self._client.branch_id}')There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Implemented via the shared _fetch_ws helper: _find_ws_by_id tries this manager's client first, then the production-branch client via with_branch_id(None) if that differs. Also added branch_id={self._client.branch_id} to the final ValueError message so a genuine not-found no longer reads as if only one branch was ever checked.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (workspace.py:870, second review pass).
| """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 = None |
There was a problem hiding this comment.
A bare WORKSPACE_ID env var — no KBC_ prefix — will pin every session on the server.
_read_options probes three spellings per field, and the first is the bare field name: _normalize('workspace_id') == 'workspaceid'. Since server.py:192 runs config.replace_by(os.environ), a plain WORKSPACE_ID in the process environment is picked up — and it also silently overrides the new --workspace-id CLI flag.
That matters here specifically because WORKSPACE_ID is exactly the variable Keboola injects into Data App containers (see resources/data_app/sapi_query_data_code.py:11 and SECRET_WORKSPACE_ID in tools/data_apps.py). Run the MCP server or the in-platform agent in any such container — the AJDA-3052 scenario — and every session, for every user and branch, is pinned to that one workspace: users whose token cannot see it get ValueError: No Keboola workspace found: workspace_id=... on query_data / get_table_details / get_project_info; users who can see it silently query the wrong workspace.
workspace_schema has the same latent collision but WORKSPACE_SCHEMA isn't an injected variable, so this is the first field where it actually bites. Options: require the KBC_ prefix for this field, or rename it (e.g. pinned_workspace_id) so the bare spelling can't collide.
There was a problem hiding this comment.
Proposed fix
Add a per-field opt-out rather than changing the prefix rules globally — bare env/header names are relied on elsewhere, including the integtests (headers = {'storage_token': ..., 'workspace_schema': ...}):
workspace_id: str | None = field(default=None, metadata={'require_prefix': True})candidates = (f'KBC_{name}', f'X-{name}') if f.metadata.get('require_prefix') else (name, f'KBC_{name}', f'X-{name}')Heads-up: this breaks the two new test_config.py cases that pass a bare workspace_id key — switch them to KBC_WORKSPACE_ID and keep the X-Workspace-Id case.
The cleaner long-term fix is to require the KBC_ prefix for env vars generally (server.py:192 currently hands the whole environment to replace_by), but that risks breaking deployments relying on bare names, so the targeted flag is the safer call for this PR.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Implemented as suggested: require_prefix metadata on workspace_id only (kept workspace_schema and everything else on the existing bare/KBC_/X- candidate order, since integtests rely on the bare workspace_schema header key). Updated the two test_config.py cases that used a bare workspace_id key to KBC_WORKSPACE_ID, kept the X-Workspace-Id case, and added a numeric-shape validation test.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (config.py:30, second review pass).
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_get_workspace_resolves_by_id_when_set(): |
There was a problem hiding this comment.
The workspace_id pass-through in create() has zero coverage — I mutation-tested it.
Both new tests construct WorkspaceManager(...) directly, so the two changed cls(...) calls in create() (workspace.py:592 for storage-branches projects, :595 for the prod-client path) are never exercised. Deleting workspace_id=workspace_id from the storage-branches branch — exactly the omission that duplicated two-call shape invites — leaves the entire suite green while every dev-branch session falls back to the unrestricted default MCP workspace. The restriction stops applying and no test notices.
The existing parametrized test_workspace_manager_create_is_branch_aware (line 202) already asserts manager._workspace_schema == workspace_schema and is the natural place to add a workspace_id axis — which also satisfies CONTRIBUTING.md § Testing Requirements ("extend existing parametrized tests rather than adding new test functions for related scenarios").
There was a problem hiding this comment.
Proposed fix
Extend test_workspace_manager_create_is_branch_aware (line 202) with a workspace_id axis and assert manager._workspace_id == workspace_id. That is specifically what kills the mutation — two cases are enough:
dev_branch_with_sb_explicit_id(coversworkspace.py:592)default_branch_without_sb_explicit_id(coversworkspace.py:595)
Two more worth adding while here:
- an empty-header case in
test_config.py, for the''-means-unset change - one
_get_workspacetest asserting the data-app accessor still returns the managed workspace while the session is pinned — the regression test for themcp.py:373thread
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Added both named cases (dev_branch_with_sb_explicit_id, default_branch_without_sb_explicit_id) to test_workspace_manager_create_is_branch_aware, an empty-header case in test_config.py, and a regression test (test_get_data_app_workspace_id_ignores_the_pin) asserting get_data_app_workspace_id() returns the managed workspace while _get_workspace() still returns the pinned one in the same session.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (tests/test_workspace.py:261, second review pass).
| ) | ||
| 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.') |
There was a problem hiding this comment.
The new option is documented nowhere outside the source.
README.md mentions KBC_WORKSPACE_SCHEMA in ten places — the required-credentials list (line 134), a dedicated section (line 151), every client config sample, and the troubleshooting table (**Workspace Issues** | Confirm KBC_WORKSPACE_SCHEMA is correct) — while the new, functionally overlapping KBC_WORKSPACE_ID / --workspace-id / X-Workspace-Id appears in none of them.
The PR names kai-client, kai-agent and the kbc-ui embed as follow-up consumers; they can't discover the option or learn that it takes precedence over the schema. The troubleshooting row also now gives wrong advice for sessions pinned by ID.
There was a problem hiding this comment.
Proposed fix
Add a KBC_WORKSPACE_ID subsection next to KBC_WORKSPACE_SCHEMA (README.md:151) covering:
- what it is and that it takes precedence over
KBC_WORKSPACE_SCHEMA - that it is the option Data App / kai-agent callers supply, as
X-Workspace-Idper request - that the CLI equivalent is
--workspace-id
And update the troubleshooting row (line 441) to name both:
| **Workspace Issues** | Confirm `KBC_WORKSPACE_SCHEMA` / `KBC_WORKSPACE_ID` is correct |
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Added the KBC_WORKSPACE_ID subsection next to KBC_WORKSPACE_SCHEMA covering precedence, the X-Workspace-Id per-request form, and the --workspace-id CLI flag, plus the troubleshooting row naming both.
There was a problem hiding this comment.
Fixed ✅ in commit 71ea3bd
Resolved -- see the detailed reply on the follow-up thread (cli.py:56, second review pass).
Suggested sequencing for the review fixesI've replied to each review thread with a concrete proposed fix. Grouping them into three commits keeps the diff reviewable: 1.
2.
3.
One thing to settle before codingThe read-only question (thread on Not blocking this PRPer offline discussion: limiting the token itself is out of scope here, and so is restricting the metadata tools ( |
3fbc164 to
71ea3bd
Compare
cjayyy
left a comment
cjayyy
left a comment
There was a problem hiding this comment.
Code review of the workspace_id pin. Six findings below, most-severe first. Tests for the touched files pass locally (351 passed); the trust-boundary and read-only-enforcement questions already covered in the PR description are deliberately out of scope here.
The first three are the ones I'd want resolved before merge. The last one is arguably intentional-but-undocumented rather than a defect.
0dc5996 to
737b743
Compare
|
@cjayyy thanks for the review — all six addressed in 798c7b2 / e711a02:
Added/extended regression tests for all six in |
cjayyy
left a comment
cjayyy
left a comment
There was a problem hiding this comment.
Re-reviewed against e711a02a. Five of the six are genuinely fixed — I verified each behaviorally rather than just reading the diffs:
| # | Finding | Status |
|---|---|---|
| 1 | Malformed X-Workspace-Id → unhandled error |
✅ Fixed |
| 2 | 400/403 swallowed on creation path | ✅ Fixed |
| 3 | sql_dialect / workspace_id split-brain |
✅ Fixed |
| 4 | Pinned session auto-provisions a workspace | ✅ Fixed |
| 5 | Prod-branch fallback crosses branch isolation | |
| 6 | empty_means_absent on workspace_schema |
✅ Fixed |
Nice touches worth calling out: the strict flag threading is the right shape (creation path re-raises, caller-supplied pin stays lenient); the replace_by degradation correctly preserves a server-side pin when a junk header arrives (Config(workspace_id='111').replace_by({'X-Workspace-Id':'abc'}).workspace_id == '111') while the CLI/env constructor stays strict; and get_data_app_sql_dialect() mirrors the existing accessors cleanly.
One new issue from the #5 fix — details inline. Full suite: 1737 passed (the one test_json_logging failure is environmental, no bare python on PATH, and fails on main too).
Adds a `workspace_id` Config field (X-Workspace-Id header / KBC_WORKSPACE_ID env / --workspace-id CLI) so a caller can target a specific existing workspace instead of the default per-branch one, e.g. a Data App's own workspace when Kai queries run inside it. Takes precedence over workspace_schema when both are set.
…den its resolution - split _get_workspace() (pin-aware) from a new _get_managed_workspace() (workspace_schema + default MCP workspace, ignoring workspace_id); add get_data_app_workspace_id()/get_data_app_branch_id() backed by the latter and switch tools/data_apps.py's three call sites to them, so a session pinned to one Data App's workspace can no longer leak that workspace into a *different* app's own persisted WORKSPACE_ID secret - make a server-configured workspace_id/workspace_schema authoritative over a request header in apply_request_config(), warning when a header would have overridden it (mirrors the existing storage_api_url protection); a server with no pin of its own still takes it from the request - treat an empty workspace_id/workspace_schema header/env value as "not provided" (Config._read_options, opt-in via an `empty_means_absent` field metadata flag so branch_id's existing empty-clears-the-value behavior is unaffected), plus keep `is not None` in workspace.py as defense in depth - require the KBC_/X- prefix for workspace_id specifically (`require_prefix` metadata), since the bare WORKSPACE_ID spelling collides with the variable Keboola injects into Data App containers - validate workspace_id is numeric in Config.__post_init__ - resolve a pinned id against the production-branch client too (via a shared _fetch_ws helper), since a Data App's workspace is not tied to any particular branch and workspace_detail is branch-scoped; include the branch id in the not-found error message - treat 400/403 like 404 when resolving a pinned id - warn (not yet enforce -- unconfirmed whether a Data App's platform-provisioned workspace is actually read-only) when a pinned workspace has no read-only storage access - redact _WspInfo.__repr__ so logging it can never leak backend credentials
ruff (SIM102) flags the nested if as combinable; also fixes the CI lint failure that the previous commit's local ruff run (older pinned version) didn't catch.
… in sync, stop auto-provisioning for pinned sessions - _fetch_ws/_find_ws_by_id gain a strict mode: the post-creation lookup in _create_ws now re-raises a 400/403 instead of masking it as "not found" and leaving an orphaned workspace+config behind. The caller-supplied pin lookup keeps treating 400/403/404 alike. - the prod-branch fallback in _find_ws_by_id now rebinds this manager's client to the client the workspace was actually resolved on, so queries don't run against a production workspace through a dev-branch-bound client. - add get_data_app_sql_dialect(), mirroring the other managed-workspace accessors, and use it in modify_streamlit_data_app(_internal) instead of the pin-aware get_sql_dialect() -- the dialect baked into a deployed app's source code must match the workspace whose id/branch are persisted into that same app, not whichever workspace the session happened to be pinned to. - _get_managed_workspace() now raises instead of provisioning a brand new MCP-managed workspace when a session is pinned via workspace_id and none already exists -- that workspace would be billed to this session's token and provisioning can outright fail on a read-only token.
…ers, scope empty-header opt-out to workspace_id replace_by() previously let __post_init__'s workspace_id validation raise ValueError straight out of request-header handling: a junk X-Workspace-Id header became an unhandled server error (a 500 via preview.py, which maps ValueError to 400 only around the call this bypassed), and did so before the server-pin guard in mcp.py even ran, so a bad header from any caller could disrupt a server that never opted into honoring that header at all. Now a malformed workspace_id degrades to "not provided" instead. Also scope the empty_means_absent flag to workspace_id only. It was applied to the pre-existing workspace_schema field too, which silently changed behavior for existing callers: an empty X-Workspace-Schema header used to clear a server default (the multi-user opt-out the README describes) and would otherwise keep re-pinning to it instead.
…e prod-branch fallback The #5 fix rebound self._client to the prod client that resolved a pinned workspace, but self._client is shared by every other lookup this manager makes -- including _find_ws_in_branch for the *managed* workspace, which feeds get_data_app_workspace_id()/get_data_app_branch_id()/ get_data_app_sql_dialect() and gets persisted into a Data App's own config. A dev-branch session pinned to a workspace that only resolves via the prod fallback would then silently do its managed-workspace bookkeeping against production too -- reopening the #3 cross-branch bleed through a different door. _find_ws_by_id now returns the resolving client alongside the info instead of mutating self._client; _init_workspace takes an optional client override so only the workspace actually being constructed uses it, leaving the manager's own client (and everything else that reads it) untouched.
e711a02 to
1eb0bec
Compare
cjayyy
left a comment
cjayyy
left a comment
There was a problem hiding this comment.
Re-verified against 1eb0becc. All six findings are resolved. 1eb0becc takes the (info, client) approach — _find_ws_by_id/_create_ws return the resolving client and _init_workspace gained a client override, so only the workspace being constructed uses it and self._client is left alone.
| # | Finding | Status |
|---|---|---|
| 1 | Malformed X-Workspace-Id → unhandled error |
✅ |
| 2 | 400/403 swallowed on creation path | ✅ |
| 3 | sql_dialect / workspace_id split-brain |
✅ |
| 4 | Pinned session auto-provisions a workspace | ✅ |
| 5 | Prod-branch fallback crosses branch isolation | ✅ |
| 6 | empty_means_absent on workspace_schema |
✅ |
The #5 fix now satisfies both properties that were previously mutually exclusive. My original round-2 reproduction passes, and printing the values it asserts on:
pinned workspace id = 123
pinned ws bound client = prod <- #5: pin uses the client that resolved it
manager client branch_id = 456 <- #3: manager stays on its own branch
managed lookup branch_id = 456 <- managed bookkeeping no longer leaks to prod
I also mutation-tested test_pin_resolution_via_prod_fallback_does_not_leak_into_managed_lookup — reintroducing the single self._client = prod_client line on top of the new API fails it and test_find_ws_by_id_falls_back_to_production_branch, so the regression is genuinely pinned down rather than passing vacuously.
Spot-checked that the rebase didn't undo the earlier fixes: the replace_by degradation still preserves a server-side pin against a junk header ('111'), the constructor stays strict, X-Workspace-Schema: '' clears again while X-Workspace-Id: '' keeps its pin, and strict=True / get_data_app_sql_dialect() / the no-provision guard are all still in place. Also audited the tuple-return call sites — all three consumers unpack correctly, and the two _WspInfo-returning paths (workspace_schema, _find_ws_in_branch) correctly omit the client override.
Suite: 1738 passed. The single test_json_logging failure is environmental (no bare python on PATH) and fails on main too.
Two notes, neither blocking and neither a code request:
- My verification uses mocked clients throughout, so it establishes control flow, not behavior against a real stack. The
Streamable-HTTP/ canary testing checkboxes in the description are still unticked — worth exercising the pinned path end-to-end before this reaches production. - The two open design questions you documented (the header trust boundary, and warn-vs-reject for a writable pinned workspace) are unchanged by these fixes and still tracked for follow-up.
Nice work on the turnaround.
|
@MiroCillik one more round please just to be sure I did not miss anything |
Description
Linear: AI-3669
Change Type
Summary
Kai currently always queries the default per-branch, MCP-managed workspace — even when invoked
inside a Data App, which has its own dedicated workspace with its own (potentially more
restricted) permissions. This is the first step towards fixing that: it adds a
workspace_idoption to the server
Config, settable via CLI (--workspace-id), env (KBC_WORKSPACE_ID), or— since HTTP headers are read per-request — the
X-Workspace-Idheader. When set,WorkspaceManagernow resolves that specific workspace via the existing_find_ws_by_idlookupinstead of the default per-branch workspace, and takes precedence over the existing
workspace_schemaoption when both are set.This mirrors the
workspace_schema/X-Workspace-Schemamechanism that already existed, butkeyed by workspace ID instead of schema — Data App runtimes only know their own workspace as a
WORKSPACE_ID(seetools/data_apps.py), not a schema name, so the ID-based lookup is what acaller (e.g. kai-agent, when Kai is embedded in a running Data App) will actually be able to
supply.
Follow-up work (separate repos, not part of this PR):
kai-clientneeds to accept and forward aworkspace id as a header,
kai-agentneeds to thread that header through to both of its MCPclient paths (backend-side calls and the sandbox's own MCP client), and the Data App /
kbc-ui embed needs to actually supply the app's
WORKSPACE_IDinto the chat request.Review round 2 (@MiroCillik) — what changed
tools/data_apps.pybookkeeping call sites(
get_data_app_workspace_id()/get_data_app_branch_id()), which previously leaked a pinnedsession's workspace into a different data app's own persisted
WORKSPACE_IDsecret.workspace_id/workspace_schemais now authoritative over a requestheader (warns instead of silently being overridden); an empty header/env value no longer
un-pins a server-side default.
workspace_idnow requires theKBC_/X-prefix (a bareWORKSPACE_IDenv var collideswith the variable Keboola injects into Data App containers), is validated as numeric, and its
lookup now falls back to the production-branch client and maps 400/403 like 404.
_WspInfo's repr no longer printscredentials.Design note (not fully resolved in this PR)
— it narrows access only as much as whoever sets the header is trusted to. This is not a
regression (
workspace_schemahas always had the same trust model), but it means this PR doesnot by itself prevent a Data App user from reaching org data outside their app's workspace
(the AJDA-3052 "no backdoor" goal) —
get_tables/get_bucket_detail/get_project_infostillread through the caller's own Storage API token, unaffected by this pin. Limiting the token
itself is tracked separately in
keboola/connection#7981/#7985.gets a
LOG.warning, not a hard rejection — whether a Data App's platform-provisionedworkspace actually has
readOnlyStorageAccess: trueis unconfirmed against a real stack.Enforcing blindly could break the feature outright if it turns out to be
false. Onceconfirmed, this should either be promoted to a hard
raise(if read-only) or backed by aSELECT-only guard intools/sql.py(if not) —query_datahas no statement-type guard ofits own today regardless of this pin.
Testing
Streamable-HTTPtransports)Optional testing
canary-orionMCP (Streamable-HTTP)canary-orioncanary-orionUnit tests cover the
workspace_id/workspace_schemaconfig parsing (incl. the KBC_-prefixrequirement and empty-value handling),
WorkspaceManager's pin-vs-managed workspace split(precedence, not-found, non-readonly-warning, branch fallback, 400/403 handling), the
apply_request_configauthoritative-pin behavior, and the CLI--workspace-idflag. Fullpytestsuite passes (1732 passed).ruff format/ruff checkclean;tox -e check-tools-docsunaffected (no tool signature changes).
Checklist