Skip to content

Add OAuth2 support#217

Draft
raulcabello wants to merge 11 commits into
mainfrom
agent-oauth2
Draft

Add OAuth2 support#217
raulcabello wants to merge 11 commits into
mainfrom
agent-oauth2

Conversation

@raulcabello

@raulcabello raulcabello commented Jun 11, 2026

Copy link
Copy Markdown
Collaborator

This PR introduces OAuth2 support using the Authorization Code Flow with Proof Key for Code Exchange (PKCE) for external MCPs configured as AIAgentConfigs.

Key Changes

1. Dynamic Discovery Endpoints

Exposed two new API endpoints to handle dynamic client registration and metadata discovery. These endpoints are designed to be consumed directly by the UI:

  • Metadata Discovery Endpoint (/oauth/metadata): Resolves OAuth configuration metadata for a targeted Model Context Protocol (MCP) server url:
curl -G "http://localhost:8000/oauth/metadata" \
  --data-urlencode "mcp_url=https://mcp.example.com/v1/mcp/authv2"
  • Dynamic Registration Endpoint (/oauth/dynamic-registration): Registers client configurations against a specified external metadata endpoint.
curl -G "http://localhost:8000/oauth/dynamic-registration" \
  --data-urlencode "metadata_endpoint=https://url.com"

2. Configuration & Secret Specification

The AIAgentConfig Custom Resource Definition (CRD) now supports referencing an authenticationSecret. The target Kubernetes secret must contain the following data keys: clientID, clientSecret, scope, metadata_endpoint

3. Authentication & Protocol Details

Flow Type: Utilizes the OAuth 2.0 Authorization Code Flow with PKCE (RFC 7636) to prevent authorization code interception attacks.

Token Storage: Access and refresh tokens are saved securely on the client side inside HTTP only cookies. The token expiration is currently hardcoded to one week.

Proactive Session Refresh: To minimize user interruption, if a refresh token is already present in the request context, the agent will always attempt a silent refresh before initiating a brand new login flow.

4. Initialization Strategy (Single vs. Multi-Agent)

The entry point for triggering the OAuth2 handshake depends heavily on the configured deployment topology:

Single Agent: The OAuth2 authorization flow is eagerly initiated at the very beginning of the session.

Multiple Agents: The OAuth2 flow is lazily initialized, triggering only when a specific subagent that requires OAuth2 is actively invoked by the user.

5. Multi-Tenant State Isolation & OAuthTokenStore

Temporary runtime state is held inside the OAuthTokenStore during the initial authorization handshake and retrieved on the redirect callback to validate the PKCE code_verifier.

Security & Isolation: The OAuthTokenStore uses the user's Rancher Session (R_SESS) token as the lookup key. This guarantees strict multi-tenant data boundary isolation; users can only access their own respective authorization records.

Housekeeping: Includes an automated cleanup mechanism to purge stale or unused entries from the token store.

6. Uses the Starlette Oauth Client https://docs.authlib.org/en/v0.14/client/starlette.html

How to Test / Verification Steps

Apply an AIAgentConfig CRD that sets up the authenticationSecret containing valid provider criteria (clientID, clientSecret, scope, metadata_endpoint).

Hit the dynamic discovery endpoints using the UI (or manually via the curl commands provided above) to confirm configurations resolve smoothly.

Test the single-agent setup to verify the flow is requested immediately on launch.

Test the multi-agent setup to verify the flow triggers only when interacting with a subagent requiring authentication.

Inspect the browser cookies to verify that access/refresh tokens are set properly with a 7-day expiration boundary.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR introduces OAuth 2.0 (PKCE) as a new per-agent MCP authentication mode, adds supporting FastAPI endpoints for callback/refresh/metadata/registration, and wires OAuth handling into the websocket-driven agent invocation flow (including deferred tool loading + post-auth tool reload).

Changes:

  • Add OAuth2 service layer (client/credentials/discovery/cookies/store/handler) and OAuth HTTP routes for callback + token refresh + discovery/registration helpers.
  • Extend agent configuration and CRDs to support authenticationType: OAUTH2, and integrate OAuth into agent tool loading / supervisor routing.
  • Add comprehensive unit tests for the new OAuth2 modules and update existing agent/websocket tests for new message structures and async MCP client creation.

Reviewed changes

Copilot reviewed 34 out of 36 changed files in this pull request and generated 11 comments.

Show a summary per file
File Description
uv.lock Locks dependency changes (incl. authlib) and wheel list updates.
pyproject.toml Adds authlib==1.7.0 dependency.
README.md Documents MCP authentication types, including OAuth2.
crd-generation/api/v1alpha1/aiagentconfig_types.go Allows OAUTH2 enum value for authenticationType.
chart/agent/templates/crds/ai.cattle.io_aiagentconfigs.yaml Adds OAUTH2 to CRD schema enum; improves caBundleRef schema/docs.
app/services/oauth2/store.py In-process TTL store for OAuth state and tokens.
app/services/oauth2/models.py OAuth exceptions + credentials dataclass.
app/services/oauth2/handler.py Websocket-side OAuth flow initiation, refresh attempt, and cookie injection.
app/services/oauth2/discovery.py Implements MCP OAuth discovery helpers (WWW-Authenticate + well-known endpoints).
app/services/oauth2/credentials.py Reads OAuth credentials/metadata from Kubernetes secrets.
app/services/oauth2/cookies.py Standardizes OAuth cookie naming.
app/services/oauth2/client.py Singleton Authlib OAuth client registry per agent.
app/services/oauth2/init.py Re-exports oauth2 module API.
app/services/agent/_constants.py Centralizes agent constants + new NeedsOauth2 exception.
app/services/agent/loader.py Adds AuthenticationType.OAUTH2.
app/services/agent/factory.py Makes MCP client creation async; adds OAuth2 header injection; defers tool loading on 401; adds tool reload helper.
app/services/agent/supervisor.py Tracks needs_oauth2 on child agents and raises to trigger OAuth flow.
app/services/agent/middleware/human_validation.py Adjusts imports; adds ExceptionGroup handling for tool-call failures.
app/services/agent/middleware/cancel_human_validation.py Adjusts imports to new constants location.
app/services/agent/middleware/_constants.py Removes middleware-local constant (moved to agent constants).
app/services/agent/middleware/init.py Re-exports constant from new location.
app/routers/websocket.py Catches OAuth requirements and reloads tools after authentication.
app/routers/oauth2.py Adds OAuth callback/refresh/metadata/registration endpoints.
app/routers/testui.html Adds client-side handling for OAuth popup + refresh messages.
app/main.py Registers oauth2 router and expands noisy log filtering.
app/controllers/ai_agent_config.py Awaits async MCP validation; registers OAuth clients from secrets; treats 401 for OAuth2 as “needs auth”.
tests/unit/services/oauth2/init.py Test package marker for oauth2 unit tests.
tests/unit/services/oauth2/test_store.py Unit tests for token/state store TTL + lifecycle behavior.
tests/unit/services/oauth2/test_handler.py Unit tests for websocket OAuth handler logic.
tests/unit/services/oauth2/test_discovery.py Unit tests for discovery logic.
tests/unit/services/oauth2/test_credentials.py Unit tests for Kubernetes-secret credential reading.
tests/unit/services/oauth2/test_cookies.py Unit tests for cookie naming utilities.
tests/unit/services/oauth2/test_client.py Unit tests for OAuth client manager + TLS verify logic.
tests/unit/services/agent/test_supervisor_agent.py Updates message assertions for message object usage.
tests/unit/services/agent/test_factory.py Updates tests for async create_mcp_client.
tests/unit/routers/test_websocket.py Updates supervisor child-agent mapping structure in tests.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread app/services/agent/factory.py
Comment thread app/services/agent/supervisor.py
Comment thread app/routers/websocket.py
Comment thread app/routers/websocket.py
Comment thread app/routers/websocket.py
Comment thread tests/unit/services/oauth2/test_handler.py Outdated
Comment thread README.md Outdated
Comment thread README.md Outdated
Comment thread app/routers/oauth2.py
Comment thread app/routers/oauth2.py
raulcabello and others added 2 commits June 11, 2026 15:50
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
await websocket.send_text(f'<token-refresh>{refresh_data}</token-refresh>')
# Wait for the refresh token response
#TODO check response!
response = await websocket.receive_text()

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: agree with UI what the response should be

Comment thread app/routers/oauth2.py
# Verify state to prevent CSRF attacks and retrieve stored data
session_token = request.cookies.get("R_SESS", "")
oauth_data = oauth_store.pop_state(state, session_token)
if not oauth_data:

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: this was done for the test UI. This will change

"""
token_refreshed = await _initiate_oauth_flow(agent_name, websocket)
if not token_refreshed:
await websocket.receive_text()

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: receive confirmation from UI!

@raulcabello raulcabello changed the title Agent oauth2 Add OAuth2 support Jun 12, 2026
@torchiaf torchiaf self-requested a review June 17, 2026 07:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants