feat(PSGO-261): PAT/AT exchange, local PKCE login, multi-project scope, and OAuth session exchange - #605
Open
Matovidlo wants to merge 89 commits into
Open
Conversation
Contributor
Author
|
@claude review |
Contributor
There was a problem hiding this comment.
Pull request overview
Implements Part A of programmatic token (PAT/AT) support by detecting kbc_at_* / kbc_pat_* tokens and exchanging them via Connection’s auth-bridge for a legacy Storage API token, keeping downstream behavior unchanged for legacy tokens.
Changes:
- Added
StorageTokenResolver+is_programmatic_token()to perform the auth-bridge token exchange using a projected Kubernetes SA JWT. - Extended
Configwithproject_id(envKBC_PROJECT_ID/ headerX-KBC-ProjectId) and wired middleware to exchange programmatic tokens before creatingKeboolaClient. - Added unit tests for resolver behavior/mapping and for middleware wiring; bumped project version to
1.73.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/clients/auth_bridge.py |
New auth-bridge exchange client, token detection, and error mapping logic. |
src/keboola_mcp_server/mcp.py |
Middleware now exchanges programmatic tokens (using KBC_KUBERNETES_TOKEN_PATH) before instantiating KeboolaClient. |
src/keboola_mcp_server/config.py |
Added project_id config field with aliasing intended to support X-KBC-ProjectId and KBC_PROJECT_ID. |
tests/clients/test_auth_bridge.py |
New unit tests for token detection, request shape, and error mapping/no-leak guarantees. |
tests/test_mcp.py |
New unit tests for programmatic-token exchange wiring in SessionStateMiddleware. |
tests/test_config.py |
Updated Config.__repr__ expectation to include project_id. |
pyproject.toml |
Version bump to 1.73.0. |
uv.lock |
Lockfile updated for version bump. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Matovidlo
changed the title
Matovidlo
changed the title
Matovidlo
force-pushed
the
martinvasko-psgo-261-pat-token-exchange-implementation
branch
from
July 21, 2026 05:59
5b8c65e to
7357f6c
Compare
Matovidlo
added a commit
that referenced
this pull request
Jul 21, 2026
…/altitude) From the /simplify + ponytail + thermonuclear review passes, the safe, behavior-preserving wins: - Extract shared clients/base.py helpers normalize_storage_api_url() and read_service_account_jwt(), replacing the 3x duplicated URL validation and the 2x duplicated SA-JWT read across client.py, auth_bridge.py, auth_login.py. - Centralize the KBC_KUBERNETES_TOKEN_PATH am-I-deployed check into config.deployed_sa_token_path(), used by mcp.py (5 sites) and errors.py. - Compute MultiProjectMiddleware._largest_list_len once per project in _merge. - Harden the auth-bridge resolver: map a non-JSON / non-dict 200 body to a 502 StorageTokenExchangeError instead of letting it bubble as a 500 (Copilot #605). - Add KBC_PROJECT_ID / X-KBC-ProjectId config parsing regression cases (Copilot). Larger structural findings deferred to a focused follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Matovidlo
added a commit
that referenced
this pull request
Jul 21, 2026
Replace the nonstandard 'ponytail:' marker with 'Note:' so the caching-follow-up remark reads clearly to maintainers (Copilot review, #605). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Matovidlo
added a commit
that referenced
this pull request
Jul 21, 2026
- get_accessible_projects: an unresolved SQL dialect (None) now yields the no-dialect prompt instead of defaulting to Snowflake, so a BigQuery/unknown project isn't given Snowflake-specific SQL guidance (+ regression test). - perform_login: bound the loopback callback wait with a 300s server timeout and raise a clear error, so a closed tab / blocked browser can't hang login (and stdio auto-login). - Log swallowed scoped-exchange and per-project fan-out failures with exc_info=True so the tracebacks reach Datadog (matching the codebase convention). Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Matovidlo
added a commit
that referenced
this pull request
Jul 21, 2026
…ons, store polish - Promote MultiProjectMiddleware._client_for_project to public client_for_project so the project tool no longer reaches into a protected middleware member (also a thermonuclear finding); update its one external caller and the tests. - Reword the server instructions so they don't mislead legacy Storage-API-token sessions: multi-project gating applies to programmatic (kbc_at_/kbc_pat_) tokens; a legacy token is already single-project and uses tools directly. - credentials.json: json.dump(ensure_ascii=False) per the project JSON guideline. - Fix ensure_access_token docstring: the loopback wait is now bounded by _LOGIN_CALLBACK_TIMEOUT_SECONDS. Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Matovidlo
added a commit
that referenced
this pull request
Jul 21, 2026
…st comment - get_accessible_projects: re-raise asyncio.CancelledError from the concurrent dialect lookups instead of treating it as a best-effort per-project failure, so a cancelled request stops promptly rather than continuing work. - Reword the /list on_request comment: it skips the extra auth round-trips (introspect/ refresh/scoped-exchange), but create_session_state below may still make ordinary Storage calls (WorkspaceManager.create) — the old 'zero network' claim was misleading. Copilot review, #605. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…s naming convention Keeps the field named postgres_dsn internally (so __repr__'s dsn-keyword redaction still applies) while accepting MCP_DB_URL / KBC_MCP_DB_URL / KBC_POSTGRES_DSN as the env var, consistent with how other Config fields (e.g. storage_token/storage_api_token) already alias multiple accepted names. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ore (Phase 2) OAuth access/refresh tokens are no longer self-contained JWTs -- they're opaque, randomly-generated references to a row in Postgres (oauth_sessions), carrying the real (encrypted) Keboola access/refresh token server-side only. This buys two things the JWT design structurally couldn't: revocation (revoke_token now actually deletes the session instead of being a no-op) and transparent server-managed refresh (load_access_token refreshes the underlying Keboola credential in place when it's near expiry, so the client's own opaque token never needs to change). exchange_refresh_token rotates both the underlying Keboola credential and the client-facing opaque token pair (OAuth 2.1 refresh-token-rotation). Neither opaque token carries a client-visible expiry -- validity is revocation-based (row deleted/soft-revoked), not TTL-based. server.py now refuses to enable OAuth without a Postgres DSN configured (no silent in-memory fallback for a production auth path) and constructs the PostgresSessionStore, whose connection pool is lazily created on first use so create_server() stays a plain sync function -- forcing its many call sites (including a dozen-plus synchronous tests) to become async would have been a much bigger, unrelated change. Also adds the `keboola-mcp-server migrate` CLI subcommand: a thin wrapper over the migration runner added in the prior commit, meant to run as a one-shot Job before the server deployment rolls out (schema migrations are deliberately NOT auto-applied by the app itself). Multi-project scope-on-DB for OAuth sessions (RFC Phase 3, replacing scope_token for this session type specifically) is not part of this commit -- scope_token keeps working unchanged for every session type, OAuth included. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
OAuth-authenticated sessions no longer need to resend scope_token: the opaque OAuth access token already round-trips through the Postgres session store on every call (load_access_token), so set_project_scope now persists the confirmed scope there too (SessionStore.update_scope) instead of only minting a scope_token. PAT/header-token sessions have no session row to persist against and keep relying on scope_token as before. Bump to 1.78.0 (new capability, not just a fix).
Closes the attribution gap flagged in the RFC's fan-out follow-up: multi-project read results merge every project's structured_content lists together, and until now nothing survived to say which project a given item came from once merged -- only the "=== project N ===" text envelope did, which structured-output clients don't parse. MultiProjectMiddleware._tag_items_with_project stamps _scope_project_id (not the RFC's originally-sketched source_project -- that name is already a real field on bucket/table models for Keboola's own linked-bucket provenance, and would have silently clobbered it) onto every dict item before _deep_merge concatenates the per-project lists. Only fires for genuine 2+ project fan-out; single-target calls already know their project unambiguously. Bump to 1.79.0 (new capability).
Resolves the RFC's "session expiry / cleanup" open question: rows never got deleted. Converts oauth_sessions to PARTITION BY RANGE (created_at), one partition per month, dropped wholesale once older than the 2-month retention window (an instant DROP TABLE, no vacuum needed, unlike a DELETE sweep). - Migration 0002 recreates the table as partitioned, copying existing rows across (safe: no production OAuth sessions exist on this schema yet). - session_store/retention.py's ensure_partitions() creates the current + next month's partition ahead of time (a RANGE insert with no matching partition raises immediately) and drops ones past retention. Idempotent, safe to have missed a run. - New `keboola-mcp-server gc-sessions` CLI subcommand wires this as a schedule-triggered job (kbc-stacks monthly CronJob, follow-up), separate from the deploy-triggered `migrate` hook. - Trade-off documented in the RFC: PostgreSQL requires partitioned-table UNIQUE/PK indexes to include the partition key, so access_token_hash / refresh_token_hash / id are only enforced unique per-partition now, not table-wide -- a cross-month collision on a 256-bit random token is cryptographically negligible. Verified end-to-end against the real docker-compose Postgres (migration, create/lookup/revoke, idempotent + first-run partition creation, stale partition drop). Bump to 1.80.0 (new capability).
… review
Thermo-nuclear + ponytail review of this PR found mcp.py had grown from 613
to 1384 lines. Extracts two self-contained pieces into their own modules:
- scope.py: SessionScope, resolve_scope_secret, and the scope_token/
SCOPE_KEY/OAUTH_SESSION_ID_KEY constants.
- multiproject.py: MultiProjectMiddleware (fan-out + merge engine) in full.
- build_tracing_headers moves to config.py (depends only on
ServerRuntimeInfo) so multiproject.py doesn't need to import it back from
mcp.py, avoiding a circular import.
- BOOTSTRAP_TOOLS moves to tools/constants.py, the one constant genuinely
shared between ToolsFilteringMiddleware (stays in mcp.py) and
MultiProjectMiddleware (moved out).
mcp.py: 1384 -> 947 lines. Test suite split to match (test_multiproject.py).
Import sites (server.py, tools/project.py, tests) updated to import from
wherever each symbol actually lives now, rather than re-exporting through
mcp.py -- the honest fix, and the only way to avoid the cycle.
Other findings addressed:
- clients/auth_bridge.py: collapsed _AuthBridgeClient/_AuthBridgeExchangeError
(one subclass each) into their single subclasses.
- tools/project.py: collapsed set_project_scope's two near-identical except
blocks (only the re-raise condition differed).
- session_store/retention.py: dropped a redundant SQL regex (duplicating the
Python one) and a can't-happen assert.
- session_store/{migrations/0002,retention.py,cli.py}: migration 0002 no
longer hand-rolls partition creation in a SQL DO block -- the `migrate` CLI
command now calls ensure_partitions() right after applying migrations
(also what the monthly gc-sessions job calls), one Python-side mechanism
for all partition creation instead of two.
- mcp.py: SessionStateMiddleware._read_scope_from_request no longer guards
against context.message being absent (a required MiddlewareContext field,
can't happen in production) -- only whether it HAS .arguments varies.
Considered and NOT adopted (documented, not silently dropped):
- Squashing migrations 0001+0002: both have almost certainly already run
against live dev-stack Postgres databases; squashing risks stranding a
stack caught mid-migration with no down-migration support.
- Dropping crypto.py's key-version prefix byte: the RFC explicitly reasons
about deciding this encoding before real rows exist, not retrofitting
later -- removing it now would undo that already-made call.
- Reducing postgres_dsn's 3 env-var aliases: matches the alias/prefix
mechanism every other Config field already uses, not a bespoke pattern.
All 1792 tests pass; black/isort/flake8/TOOLS.md clean.
Bump to 1.81.0.
…ions ensure_partitions() computed cutoff as this_month - retention_months, which keeps retention_months+1 months of data (e.g. retention_months=2 kept 3 months: current + 2 prior). Fix the off-by-one so retention_months counts the current month too, matching the "2-month retention" the RFC documents.
…ng a new oauth_sessions partition ensure_partitions() plain `CREATE TABLE ... PARTITION OF ... FOR VALUES` fails with asyncpg.exceptions.CheckViolationError whenever oauth_sessions_default already holds rows in the range being carved out -- exactly what happens on a stack with real pre-existing sessions, since migration 0002 copies the legacy table's rows into the default partition before any month partition exists. Build the partition as a standalone table, move matching default rows into it first, then ATTACH -- works whether or not default has conflicting rows.
…token The near-expiry refresh in load_access_token was silent on success (only the failure path logged a warning), so there was no way to observe when a session's Keboola credential got lazily extended. Log an info line with the session id only -- no token values.
resolve_encryption_key() silently falls back to a process-local key when KBC_SESSION_ENCRYPTION_KEY is unset -- fine for local dev/tests, but in production this would make persisted OAuth sessions permanently undecryptable after every restart. Refuse to start instead, mirroring the existing Postgres DSN guard right above it. Addresses a Copilot review comment on PR #605.
…s token hashes The parent's (access_token_hash, created_at) / (refresh_token_hash, created_at) unique indexes never actually reject a duplicate hash -- created_at differs per row, so the composite pair is always distinct. PostgresSessionStore.get_by_access_token() does an unbounded UPDATE ... RETURNING * on access_token_hash alone, so a duplicate would update multiple rows while fetchrow() silently returns an arbitrary one. Add a plain (non-composite) unique index directly on each partition table -- Postgres only requires the partition key in indexes defined on the partitioned parent, not on a partition's own table. ensure_partitions() adds it to every new month partition; a new migration (0003) adds it to oauth_sessions_default retroactively, since 0002 is already applied on real stacks and can't be edited. Addresses a Copilot review comment on PR #605.
…active-project re-scope Writes previously targeted the first scoped project implicitly, forcing a set_project_scope call to change which project a write lands on -- and reordering the whole scope's read fan-out as a side effect. Every write/modify/delete tool now declares project_id explicitly; ambiguous calls (2+ scoped projects, no project_id) raise instead of silently defaulting. MultiProjectMiddleware resolves and swaps the target ahead of ToolsFilteringMiddleware so authorization runs against the targeted project, mirroring the existing read fan-out ordering. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ssible_projects Piggybacks on the per-project token verify already made for the SQL dialect (no extra API call) -- that response's organization field is the same one get_project_info already reads for organization_id. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…e tools get_project_info was in _NO_FANOUT_TOOLS and always reported on whichever project happened to be active (first in scope), with no way to name a different one -- the same "implicit active project" wart the write-tool targeting fix (787eca4) already removed elsewhere. It now takes project_id (required once 2+ projects are scoped, defaults to the single scoped project otherwise), resolved via the same single-target dispatch write tools use (renamed _dispatch_write/_resolve_write_target to _dispatch_single_target/_resolve_single_target since both now share it). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Two related fixes: 1. Stop advertising/needing scope_token when the transport already guarantees ctx.session persists across requests -- stdio always, and streamable-http with --no-stateless-http (a flag whose own help text already promised this and was never implemented). on_request now reuses an already-confirmed scope straight from ctx.session.state instead of unconditionally rebuilding it every call; on_list_tools and set_project_scope/get_accessible_projects stop advertising/returning scope_token in that mode. 2. Real bug this surfaced: a deployed OAuth session's scoped_token (minted once by set_project_scope) was never refreshed -- unlike the local path, _resolve_local_tokens returned early for deployed sessions without ever checking scope.is_near_expiry. Once that token expired mid-conversation, every fanned-out Storage call started 401ing for the rest of the session with no indication why. Deployed sessions now get the same near-expiry re-mint, persisted back to the OAuth session's Postgres row so it happens once per expiry, not once per request. Extracted the scope-persistence write (project.py's _persist_oauth_scope and this fix's on_request path did the same SessionStore.update_scope call) into a shared scope.persist_scope() helper. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ebase ruff fixes Adds the design from pat_token_support/RFC.md increment 6: server-side scope persistence for deployed, non-OAuth, programmatic-token sessions (Kai), keyed by sha256(conversation_id:user_id) since the raw kbc_at_/kbc_pat_ token isn't stable across Kai's own refresh. New kai_sessions table (migration 0004), PostgresKaiScopeStore, read-side fallback in SessionStateMiddleware.on_request with subset-check invalidation (drop the whole scope if a previously-scoped project is no longer reachable), and write-side persistence from set_project_scope. Also fixes lint issues ruff 0.16 (picked up from the just-rebased RFC branch) flagged in pre-existing code: mutable class-attribute default, a two-call startswith merged into one, a redundant exception repr in a log call, non-tz-aware date.today() in tests, and a few unused tuple-unpack variables.
…review Documents all 9 findings from Tomas Fejfar's review of PR #604 (7 confirmed as-stated, 1 confirmed-but-broader, 1 refuted/already-mitigated) and the fix design for each, before any implementation, per this repo's RFC-first convention.
…heck bypass
Config.replace_by (used to apply per-request HTTP headers) resolved any
dataclass field, including jwt_secret -- an X-Jwt-Secret header let a caller
choose the HMAC key that signs and verifies their own scope_token, forging
arbitrary project_ids. Add Config._HEADER_ELIGIBLE_FIELDS and a dedicated
replace_by_headers() that only ever sets fields legitimately meant to vary
per request; jwt_secret/postgres_dsn/session_encryption_key/oauth_*/
mcp_server_url become permanently unreachable from a header.
normalize_storage_api_url was a bare `hostname.startswith('connection.')`
prefix check, so connection.attacker.tld passed. Replace with a regex
requiring a genuine connection.*.keboola.(com|dev) suffix, mirroring the
domain-allowlist pattern oauth.py already uses for redirect URIs.
MultiProjectMiddleware's active-project shortcuts skipped the per-project client swap -- the only place readonly=scope.read_only reached a KeboolaClient -- whenever a call targeted scope.active_project_id, true for every single-project scope and the first project of any multi-project scope. create_session_state never passed readonly= at all, so the active project's writes were never locally restricted regardless of a confirmed read-only scope. Thread readonly=(True if scope and scope.read_only else None) into create_session_state. Add KeboolaClient.writable_storage_client (workspace provisioning is server-side plumbing, not a user-visible mutation, so it must keep working under a read-only scope) and KeboolaClient.readonly; guard the MultiProjectMiddleware shortcuts with _active_client_honors_scope so they only skip the swap once the base client already matches the scope. Also fixes a regression this change would otherwise have reopened: KeboolaClient.with_branch_id() rebuilds a fresh client for any non-default branch (routine dev-branch usage) but dropped readonly when doing so, so a read-only session became fully writable again on a plain branch switch. Forward readonly=self.readonly in both non-self-returning branches.
…ogin-time scoping SessionScope.to_token/from_token used jwt_utils's JWS (signature only) -- the gzip+JSON payload, which may embed a live scoped_token bearer credential, was base64+gunzip-recoverable without the secret. Switch to AES-GCM authenticated encryption via the existing session_store/crypto.py helpers and resolve_encryption_key (the same key OAuth sessions already encrypt with). resolve_scope_secret/_FALLBACK_SCOPE_SECRET are replaced by resolve_scope_key. Also wires SessionStateMiddleware to read a login-time-persisted project scope (auth_login.TokenSet.project_ids/read_only) before falling back to the old auto-lease-to-everything default, and makes set_project_scope's llm_instruction explicit about whether read_only is server-enforced (a real scoped_token exists) or only locally enforced (exchange-failure fallback).
…t MFA, redact errors Four related local-login hardening fixes: - login (and login --pat) now require an explicit project-scope choice -- prompted interactively, or --project-ids/--all when not run from a terminal -- persisted alongside the tokens as TokenSet.project_ids/ read_only. Replaces leasing/auto-scoping to every accessible project by default with only a prompt-text "ask first" gate that nothing enforced. - Credentials are keyed by (hostname, profile) instead of hostname alone (--profile / KBC_LOGIN_PROFILE, default "default"), so two local MCP client interfaces logged into the same stack (e.g. Claude Desktop and a terminal) no longer share one entry and its rotating refresh token. An asyncio.Lock per (hostname, profile) serializes concurrent in-process refreshes; a non-blocking fcntl.flock insurance layer covers the on-disk read-modify-write for whatever narrow sharing remains. - MFA codes (--totp/--recovery) default to a getpass.getpass() prompt instead of requiring a CLI argument that sits in shell history/`ps` for the process lifetime; the flags remain as opt-in overrides for scripted use. - elevate_session/create_pat no longer raise the raw auth-endpoint response body in the exception message; full detail moves to LOG.debug only.
Applies to local login and OAuth: a session that can only reach one project has no scoping decision to make, so requiring it anyway is pure friction.
run_server's local-PKCE-login fallback (use tokens leased by a prior `login`, refreshing as needed, instead of requiring --storage-token) was wired into the stdio branch only. streamable-http/http-compat went straight to uvicorn.Config with just the CLI args, so running the server that way always required an explicit --storage-token even with a valid stored login session. Extract the fallback into _local_login_fallback and run it once, before the transport branch, for both. No-op when a token is already configured (CLI, env, or an existing storage_token) or OAuth is configured (the deployed server case, which authenticates per-session instead of via a local token). Also moves the KBC_* environment override (config.replace_by(os.environ)) earlier so this fallback can see an env-configured OAuth client id before deciding whether to log in.
… session MultiProjectMiddleware's ask-first gate only fires for an unconfirmed scope, and a login-time-persisted (or single-project-auto-confirmed) scope is already confirmed=True -- tools were never actually blocked for these sessions. The redundant get_accessible_projects/set_project_scope calls seen in practice came from the server's static instructions string, which unconditionally told the LLM to call both "at the very START of the conversation, before doing anything else" regardless of whether a scope already existed -- for an already-scoped session this just means a wasted per-accessible-project verify fan-out for no reason. Reword the instructions and get_accessible_projects' docstring to be reactive: try the data tool you actually need first; only run the scoping dance if a tool call fails asking you to confirm scope. No enforcement change -- the gate's behavior was already correct, only the LLM-facing guidance was stale.
project_id is a header-eligible Config field. _resolve_local_tokens's deployed/OAuth branch only applied a confirmed scope's active project id when config.project_id wasn't already set, so a request carrying X-KBC-ProjectId kept that header's value even after set_project_scope confirmed a narrower scope. MultiProjectMiddleware's active-project fast paths compare only the logical target against scope.active_project_id, never what project the base client was actually built with, so the mismatched client went unnoticed -- letting any caller able to attach one header redirect every default-target call to a project outside the confirmed scope, using the full unscoped token. Drop the `not config.project_id` guard so a confirmed scope's active project always wins, matching the local-programmatic branch's existing (unaffected, unconditional) behavior. A tool wanting a different scoped project still has its own project_id argument, validated separately by MultiProjectMiddleware._dispatch_single_target. Found by a full-PR security audit (thermo-nuclear/ponytail/security-scanner across the whole diff, at the user's request), not the original review.
test_deployed_no_scope_is_noop duplicated the pre-existing test_deployed_is_noop -- with scope=None neither test enters the branch this fix touches, so the project_id='7' it set was never read. The actual regression test for the fix is test_deployed_confirmed_scope_overrides_a_header_supplied_project_id. Found by /simplify (4 parallel reuse/simplification/efficiency/altitude reviews on commit c0d8fe0); the other three lenses found nothing to fix.
…with no token _local_login_fallback made any transport attempt ensure_access_token whenever no token/OAuth is configured, not just stdio. streamable-http/ http-compat legitimately run with no default token, relying entirely on a per-request header -- with allow_interactive=False (no TTY in CI) and no stored local credential, ensure_access_token raised RuntimeError, uncaught, killing the server subprocess before it could start listening. Caught by CI: integtests/test_mcp_server.py::test_remote_setup and test_http_multiple_clients deliberately start streamable-http with no token, and started failing with "No stored credentials... Run login first" once the fallback got extended to that transport. _local_login_fallback gains a `required` param: stdio passes True (unchanged -- no other token source exists there), streamable-http/ http-compat pass False (catch and log instead of raising; the server starts normally and expects a token per request).
The RFC had grown into a chronological log of 13 "increment" postscripts, several of them narrow post-review bug-fix writeups (single-project auto-scope, an X-KBC-ProjectId override, a streamable-http startup crash) mixed in with the actual design decisions. Restructured around what the system does today -- token paths, multi-project scope, scope persistence, security invariants, decisions, testing -- dropping resolved-question transcripts and reverted-approach narration that no longer help a reader understand the current design. 1046 -> 233 lines; no behavior change. Also imports feature_spec/mpa_support/mpa-plan.md from the closed, unmerged PR #451 (feature/mpa-support) into main, where it previously only existed on that PR's branch -- with an annotated comparison against this RFC's design (config shape, where "which project(s)" lives, OAuth support) so the relationship is checkable in-repo instead of only in a closed PR's diff.
The RFC (just restructured) claimed a local login session persists a minted scoped_token for an explicit project subset. That feature was implemented and then explicitly reverted earlier -- TokenSet has no scoped_token/explicitly_scoped fields, and _read_persisted_login_scope builds its SessionScope with none. Corrected: login-time narrowing is local-guard-only (X-KBC-ProjectId per request), not Connection-enforced; only set_project_scope and the single-project auto-confirms (login, OAuth) actually mint a token. Found while verifying RFC/code conformity ahead of undrafting PR #605.
…p test, ruff 0.16 lint Rebase onto main (post-#604 merge) surfaced a few loose ends: two test call sites still missing own_stack_storage_api_url after the security-fix conflict resolution, an unused jwt_utils import left over from a mcp.py merge, a duplicated TestServerRuntimeInfoSessionStatePersists class from resolving two separate conflicts, and lint findings only visible under ruff 0.16 (CI's pinned version) that this session's 0.15 venv didn't catch.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Linear: PSGO-261
Implements the PAT-support RFC: the deployed exchange (Part A), the local browser PKCE login (Part B), and multi-project scope (increments 2–3). Stacked on the RFC branch (PR #604). Also implements the OAuth session-exchange RFC (Part D): the remote/HTTP OAuth login flow now feeds into this same PAT/multi-project pipeline instead of minting a project-bound legacy Storage token.
Change Type
Summary
Part A — deployed exchange. Accept programmatic bearer tokens (
kbc_at_*,kbc_pat_*) and exchange them at Connection's auth-bridge resolver for a legacy Storage token. The server authenticates with its projected SA JWT (X-Kubernetes-Authorization, read per request); user token asX-Subject-Token;projectIdin the body. Resolver400/401/403pass through;5xx/timeout/network →502. LegacyX-StorageAPI-Tokentraffic is untouched. No token material logged.Part B — local PKCE login. The locally-run (stdio) server authenticates with only
KBC_STORAGE_API_URL.keboola-mcp-server loginruns a browser PKCE flow, leases a whole-stack session (access + refresh), and stores it to a mode-600 file. The token is refreshed during usage (per request), not just at startup (/v1/auth/token/refresh); a dead token forces re-login. Locally (no SA token) the programmatic token is forwarded downstream as a Bearer; the resolver exchange runs only on the deployed server.Part C — multi-project scope (MPA). Token introspection (
/v1/auth/token/introspect) enumerates accessible projects; scoped-token exchange (/v1/auth/pat/exchange, project ids sent as strings) narrows to a chosen subset.get_accessible_projects(lists projects; surfaces the current scope / active project; enriches each project with itssql_dialectfrom a per-project token verify; optionalwith_llm_instruction=truereturnsbase_instructionsgrouped by dialect, deduplicated) andset_project_scope(confirms scope; mints a scoped token, falls back to the whole-stack token if exchange is unavailable).set_project_scope(once per session).MultiProjectMiddleware— the activeKeboolaClientis swapped per project and structured outputs are deep-merged (list fields concatenated) so results still validate each tool's output schema; text content is labelled per project. Writes never fan out (active project only; decision D8).project_idsargument (advertised on fan-out-eligible read tools, consumed+stripped by the middleware) lets a single call target a subset — e.g.get_tables(project_ids=[18])hits only project 18 without changing scope.query_datafans out with a per-project workspace: the fan-out swaps both theKeboolaClientand a per-projectWorkspaceManager, so a query runs inside the targeted project's own workspace (its BigQuery dataset / Snowflake schema). Narrow to one project withproject_ids, or run across all scoped projects.get_project_inforeports the active project only.Files:
clients/auth_bridge.py,auth_login.py(introspect + scoped exchange),config.py(project_id),mcp.py(SessionScope, token refresh,MultiProjectMiddlewareincl. per-project workspace swap, ask-first gate, per-call filter),tools/project.py(scope tools + dialect enrichment),server.py(middleware + MPA/ask-first instructions),cli.py(login).Part D — OAuth session exchange.
/oauth/authorizeis being removed by Connection outright, replaced by/oauth/consentrequesting scopeclaudai projectless. After the standard code→token exchange, the resulting league OAuth access token is exchanged exactly once — via a newOAuthSessionExchanger(sibling ofStorageTokenResolverinclients/auth_bridge.py, same SA-JWT/X-Subject-Tokenmechanism) hittingmanage/internal/auth-bridge/exchange-oauth-token— for a whole-stackkbc_at_*programmatic session. That session then flows through the exact same PSGO-261 pipeline as a directly-supplied PAT:is_programmatic_token()/get_accessible_projects/set_project_scopeall apply to OAuth logins for the first time.ProxyAccessToken/ProxyRefreshTokendrop the obsoletesapi_token/delegatefields (AI Service/Jobs Queue now speak bearer, Part A) in favor of the exchanged session's own access/refresh tokens. Refresh is fully decoupled from the league OAuth session (confirmed against Connection'sTokenRefreshProcessor):exchange_refresh_token()callsrefresh_tokens()directly instead of re-negotiating with Connection's OAuth server on every refresh.Files:
oauth.py(consent endpoint + scope,_exchange_oauth_for_session,ProxyAccessToken/ProxyRefreshTokenshape change, simplified refresh),clients/auth_bridge.py(OAuthSessionExchanger,OAuthTokenExchangeError),mcp.py(apply_request_config).Known limitations / caveats
X-KBC-ProjectIdheader narrowing used by fan-out.metastoreanddata-sciencealready use the bearer/PAT token (+ project header), butjobs_queue,ai_service, andsync_actionsstill pass the raw storage token — soget_jobs/run_jobetc. 401 under a PAT/multi-project session until those services accept PATs or the server runs deployed (resolver → real SAPI token). Documented as future PAT work in the RFC.query_datainto a project without an existing MCP workspace needs a non-read-only scope.query_dataruns in exactly one project's workspace.KBC_PROJECT_ID).X-KBC-ManageApiTokenis required alongsideX-Kubernetes-Authorization(we send both, matching Connection's own E2E test fixture, but our existingresolve-storage-tokensibling call works today with only the latter). Seefeature_spec/oauth_session_exchange/RFC.mdDecisions §1/§3 for the exact open questions.Testing
Streamable-HTTPtransports)toxgreen (pytest, black, isort, flake8, check-tools-docs). Unit tests:tests/clients/test_auth_bridge.py,tests/test_auth_login.py(introspect + scoped exchange),tests/test_mcp.py(TestProgrammaticTokenExchange,TestResolveLocalTokens,TestMultiProjectMiddlewareincl. fan-out, deep-merge, per-project workspace swap, ask-first gate, per-call filter),tests/tools/test_project.py(scope tools + dialect enrichment + groupedbase_instructions).Integration tests (
integtests/test_pat_multiproject.py): run the same read question undersapi/pat_single/pat_mpaauth over the same pool projects (auth as a parameter; lock + cleanup use the SAPI token, the PAT only needs read access). Gated on theINTEGTEST_STORAGE_PATsecret (one PAT whose user is a member of all pool projects) — the PAT modes skip until it is wired, so CI stays green. Verified live against a dev stack: PKCE login, introspection (6 projects), scoped exchange, per-project fan-out incl.query_dataon BigQuery + Snowflake, and a data-app create/deploy.Part D unit tests:
tests/clients/test_auth_bridge.py(OAuthSessionExchangersuccess/error-mapping/incomplete-response),tests/test_oauth.py(/oauth/consent+ scope, exchange on authorization-code + refresh, network-error mapping, decoupled-refresh regression),tests/test_mcp.py(apply_request_configinjects the exchanged session token). Part D's live end-to-end verification (real OAuth client against a dev stack running the new Connection endpoint) is still pending — see Known limitations above.Checklist
INTEGTEST_STORAGE_PATsecret is set)uv.lockfeature_spec/oauth_session_exchange/RFC.md(Part D) in this branch