Skip to content

refactor(oauth): consolidate authorization flow primitives #14157

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wenxi-onyx wants to merge 1 commit into from
Open

refactor(oauth): consolidate authorization flow primitives #14157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions backend/onyx/auth/oauth_token_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import time
from collections.abc import Mapping
from typing import Any
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
from uuid import UUID
Expand Down Expand Up @@ -52,6 +53,26 @@ def validate_oauth_endpoint_url(url: str, *, resolve_dns: bool = True) -> None:
"accounts.google.com": {"access_type": "offline", "prompt": "consent"},
}

_PROTOCOL_AUTHORIZATION_PARAMS = frozenset(
{
"client_id",
"code_challenge",
"code_challenge_method",
"redirect_uri",
"resource",
"response_type",
"scope",
"state",
}
)


def conflicting_authorization_params(
additional_params: Mapping[str, object] | None,
) -> frozenset[str]:
"""Return configured parameters whose values the OAuth flow must own."""
return _PROTOCOL_AUTHORIZATION_PARAMS.intersection(additional_params or {})


def ensure_offline_access_auth_params(authorization_url: str) -> str:
"""Merge missing offline-access params into a built authorize URL.
Expand Down
26 changes: 20 additions & 6 deletions backend/onyx/oauth/authorization_attempt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import hashlib
import json
import re
import secrets
from collections.abc import Callable
from datetime import datetime, timedelta, timezone
from typing import Generic, cast

from pydantic import ValidationError
from pydantic import JsonValue, ValidationError

from onyx.cache.interface import CacheBackend
from onyx.error_handling.error_codes import OnyxErrorCode
Expand All @@ -20,16 +22,28 @@
MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS = 10 * 60


def canonical_json_fingerprint(value: JsonValue) -> str:

@cubic-dev-ai cubic-dev-ai Bot Aug 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P2: mcp_oauth_connection_headers_fingerprint passes list[tuple[str, str]], which is not included in Pydantic's JsonValue type, so the strict ty check fails on this call. Either convert the header pairs to JSON arrays before calling this helper or include that explicitly typed tuple-list shape in the helper's accepted input type.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At backend/onyx/oauth/authorization_attempt.py, line 24:

<comment>`mcp_oauth_connection_headers_fingerprint` passes `list[tuple[str, str]]`, which is not included in Pydantic's `JsonValue` type, so the strict `ty check` fails on this call. Either convert the header pairs to JSON arrays before calling this helper or include that explicitly typed tuple-list shape in the helper's accepted input type.</comment>

<file context>
@@ -20,6 +21,18 @@
 MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS = 10 * 60
 
 
+def canonical_json_fingerprint(value: JsonValue) -> str:
+    """Fingerprint configuration captured by a pending authorization attempt."""
+    serialized = json.dumps(
</file context>
Suggested change
def canonical_json_fingerprint(value: JsonValue) -> str:
def canonical_json_fingerprint(
value: JsonValue | list[tuple[str, str]],
) -> str:

"""Fingerprint configuration captured by a pending authorization attempt."""
serialized = json.dumps(
value,
ensure_ascii=True,
allow_nan=False,
sort_keys=True,
separators=(",", ":"),
)
return hashlib.sha256(serialized.encode()).hexdigest()


class AuthorizationAttemptStore(Generic[PayloadT]):
"""Stores typed OAuth attempts in a tenant-scoped cache."""

def __init__(
self,
cache: CacheBackend,
*,
cache_backend_provider: Callable[[], CacheBackend],
namespace: str,
payload_type: type[PayloadT],
ttl_seconds: int,
ttl_seconds: int = MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS,
) -> None:
if not _NAMESPACE_PATTERN.fullmatch(namespace) or len(namespace) > 64:
raise ValueError("OAuth attempt namespace is invalid")
Expand All @@ -39,7 +53,7 @@ def __init__(
f"{MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS} seconds"
)

self._cache = cache
self._cache_backend_provider = cache_backend_provider
self._namespace = namespace
self._attempt_type = cast(
type[AuthorizationAttempt[PayloadT]],
Expand Down Expand Up @@ -67,7 +81,7 @@ def store(
expires_at=expires_at,
payload=payload,
)
stored = self._cache.set_if_absent(
stored = self._cache_backend_provider().set_if_absent(
self._key(attempt.owner_id, attempt.state),
attempt.model_dump_json(),
ex=self._ttl_seconds,
Expand All @@ -85,7 +99,7 @@ def consume(
owner_id: str,
state: str,
) -> AuthorizationAttempt[PayloadT]:
stored = self._cache.getdel(self._key(owner_id, state))
stored = self._cache_backend_provider().getdel(self._key(owner_id, state))
if stored is None:
raise _invalid_attempt_error()

Expand Down
23 changes: 21 additions & 2 deletions backend/onyx/oauth/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,29 @@
from typing import Generic, TypeVar
from typing import Annotated, Generic, TypeVar

from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field

from onyx.utils.url import sanitize_next_url

PayloadT = TypeVar("PayloadT", bound=BaseModel)


def _validate_safe_oauth_return_path(value: str) -> str:
if sanitize_next_url(value) != value or any(
not character.isprintable() for character in value
):
raise ValueError("OAuth return path must be a safe internal path")
return value


SafeOAuthReturnPath = Annotated[
str,
Field(max_length=2048),
AfterValidator(_validate_safe_oauth_return_path),
]
OAuthConfigurationFingerprint = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")]
PKCECodeVerifier = Annotated[str, Field(min_length=43, max_length=128)]


class AuthorizationAttempt(BaseModel, Generic[PayloadT]):
"""One pending authorization request for an authenticated user.

Expand Down
19 changes: 5 additions & 14 deletions backend/onyx/server/documents/standard_oauth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

from fastapi import APIRouter, Depends, Query, Request
from pydantic import BaseModel, ConfigDict, ValidationError, field_validator
from pydantic import BaseModel, ConfigDict, ValidationError
from sqlalchemy.orm import Session

from onyx.auth.permissions import require_permission
Expand All @@ -21,10 +21,10 @@
AuthorizationAttemptStore,
generate_authorization_state,
)
from onyx.oauth.models import PKCECodeVerifier, SafeOAuthReturnPath
from onyx.server.documents.models import CredentialBase
from onyx.utils.logger import setup_logger
from onyx.utils.subclasses import find_all_subclasses_in_package
from onyx.utils.url import sanitize_next_url

logger = setup_logger()

Expand Down Expand Up @@ -83,25 +83,16 @@ def _validate_additional_kwargs(
class _ConnectorOAuthAttemptPayload(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)

desired_return_path: str
desired_return_path: SafeOAuthReturnPath
additional_kwargs: dict[str, str]
code_verifier: str | None = None

@field_validator("desired_return_path")
@classmethod
def validate_return_path(cls, value: str) -> str:
if sanitize_next_url(value) != value or any(
not character.isprintable() for character in value
):
raise ValueError("OAuth return path must be a local application path")
return value
code_verifier: PKCECodeVerifier | None = None


def _authorization_attempt_store(
source: DocumentSource,
) -> AuthorizationAttemptStore[_ConnectorOAuthAttemptPayload]:
return AuthorizationAttemptStore(
get_cache_backend(),
cache_backend_provider=get_cache_backend,
namespace=f"{_OAUTH_ATTEMPT_NAMESPACE_PREFIX}-{source.value}",
payload_type=_ConnectorOAuthAttemptPayload,
ttl_seconds=_OAUTH_STATE_EXPIRATION_SECONDS,
Expand Down
29 changes: 12 additions & 17 deletions backend/onyx/server/features/mcp/credentials.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,19 +19,18 @@
in ``onyx.db.mcp``. It must not import ``onyx.server.features.mcp.oauth``.
"""

import hashlib
import json
import time
from collections.abc import Mapping
from typing import cast

from mcp.shared.auth import OAuthClientInformationFull
from pydantic import BaseModel, ConfigDict
from pydantic import BaseModel, ConfigDict, JsonValue
from sqlalchemy.orm import Session

from onyx.db.enums import MCPAuthenticationPerformer, MCPAuthenticationType
from onyx.db.mcp import get_user_connection_config
from onyx.db.models import MCPConnectionConfig, MCPServer, User
from onyx.oauth.authorization_attempt import canonical_json_fingerprint
from onyx.server.features.mcp.models import (
DENYLISTED_MCP_HEADERS,
MCPAuthTemplate,
Expand Down Expand Up @@ -71,31 +70,27 @@ def mcp_token_expired(config_data: MCPConnectionData) -> bool:


def mcp_oauth_connection_headers_fingerprint(headers: dict[str, str]) -> str:
routing_headers = sorted(
(key.lower(), value)
for key, value in headers.items()
if key.lower() != "authorization"
)
serialized_headers = json.dumps(
routing_headers, ensure_ascii=True, separators=(",", ":")
)
return hashlib.sha256(serialized_headers.encode()).hexdigest()
routing_headers: JsonValue = [
[key, value]
for key, value in sorted(
(key.lower(), value)
for key, value in headers.items()
if key.lower() != "authorization"
)
]
return canonical_json_fingerprint(routing_headers)


def mcp_oauth_client_information_fingerprint(
client_information: OAuthClientInformationFull,
) -> str:
serialized_client = json.dumps(
return canonical_json_fingerprint(
client_information.model_dump(
mode="json",
exclude_none=True,
by_alias=True,
),
ensure_ascii=True,
sort_keys=True,
separators=(",", ":"),
)
return hashlib.sha256(serialized_client.encode()).hexdigest()


def requires_user_authentication(
Expand Down
45 changes: 16 additions & 29 deletions backend/onyx/server/features/mcp/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from mcp.types import Tool as MCPLibTool
from pydantic import AnyUrl, BaseModel, Field, model_validator

from onyx.auth.oauth_token_manager import conflicting_authorization_params
from onyx.db.enums import (
EndpointPolicy,
MCPAuthenticationPerformer,
Expand All @@ -19,21 +20,16 @@
MCPServerStatus,
MCPTransport,
)
from onyx.oauth.models import (
OAuthConfigurationFingerprint,
PKCECodeVerifier,
SafeOAuthReturnPath,
)

# Matches `{placeholder_name}` inside header value templates.
_PLACEHOLDER_RE = re.compile(r"\{([^}]+)\}")
# RFC 9110 field-name syntax: a non-empty sequence of HTTP token characters.
_HTTP_FIELD_NAME_RE = re.compile(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+")
RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS = {
"client_id",
"code_challenge",
"code_challenge_method",
"redirect_uri",
"resource",
"response_type",
"scope",
"state",
}


def _build_auto_substitution_map(*, user_email: str) -> dict[str, str]:
Expand Down Expand Up @@ -387,8 +383,8 @@ def validate_auth_configuration(self) -> "MCPToolCreateRequest":
raise ValueError(
"oauth_token_endpoint is required for known-provider OAuth mode"
)
reserved_params = RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS.intersection(
self.oauth_additional_auth_params or {}
reserved_params = conflicting_authorization_params(
self.oauth_additional_auth_params
)
if reserved_params:
raise ValueError(
Expand Down Expand Up @@ -511,7 +507,9 @@ class MCPOAuthConnectResponse(BaseModel):

class MCPUserOAuthConnectRequest(BaseModel):
server_id: int = Field(..., description="ID of the MCP server")
return_path: str = Field(..., description="Path to redirect to after callback")
return_path: SafeOAuthReturnPath = Field(
..., description="Path to redirect to after callback"
)
include_resource_param: bool = Field(..., description="Include resource parameter")
force_reauthentication: bool = Field(
default=False,
Expand Down Expand Up @@ -539,17 +537,6 @@ class MCPUserOAuthConnectRequest(BaseModel):
),
)

@model_validator(mode="after")
def validate_return_path(self) -> "MCPUserOAuthConnectRequest":
if (
not self.return_path.startswith("/")
or self.return_path.startswith("//")
or "\\" in self.return_path
or any(not character.isprintable() for character in self.return_path)
):
raise ValueError("return_path must be a safe internal path")
return self


class MCPUserOAuthConnectResponse(BaseModel):
server_id: int
Expand All @@ -574,7 +561,7 @@ def validate_outcome(self) -> "MCPUserOAuthConnectResponse":
class MCPPendingOAuthAuthorization(BaseModel):
authorization_url: str
state: str
code_verifier: str
code_verifier: PKCECodeVerifier


class MCPOAuthServerSnapshot(BaseModel):
Expand All @@ -592,12 +579,12 @@ class MCPOAuthServerSnapshot(BaseModel):
class MCPOAuthFlowState(BaseModel):
server_id: int
connection_config_id: int
return_path: str
code_verifier: str
return_path: SafeOAuthReturnPath
code_verifier: PKCECodeVerifier
redirect_uri: AnyUrl
server_snapshot: MCPOAuthServerSnapshot
connection_headers_fingerprint: str
client_information_fingerprint: str
connection_headers_fingerprint: OAuthConfigurationFingerprint
client_information_fingerprint: OAuthConfigurationFingerprint
protected_resource_metadata: ProtectedResourceMetadata | None = None
oauth_metadata: OAuthMetadata | None = None
authorization_server_url: str | None = None
Expand Down
8 changes: 4 additions & 4 deletions backend/onyx/server/features/mcp/oauth_flow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from onyx.auth.oauth_token_manager import (
OAuthFlowParams,
build_oauth_authorization_url,
conflicting_authorization_params,
)
from onyx.cache.factory import get_cache_backend
from onyx.db.enums import MCPTransport
Expand All @@ -36,7 +37,6 @@
)
from onyx.server.features.mcp.models import (
DENYLISTED_MCP_HEADERS,
RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS,
MCPOAuthFlowState,
MCPOAuthServerSnapshot,
MCPPendingOAuthAuthorization,
Expand Down Expand Up @@ -85,8 +85,8 @@ def _known_provider_flow_params(
"Known-provider OAuth requires oauth_authorization_endpoint, "
"oauth_token_endpoint, and a non-empty client_id.",
)
reserved_params = RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS.intersection(
mcp_server.oauth_additional_auth_params or {}
reserved_params = conflicting_authorization_params(
mcp_server.oauth_additional_auth_params
)
if reserved_params:
raise OnyxError(
Expand Down Expand Up @@ -180,7 +180,7 @@ def validate_mcp_oauth_flow_configuration(

def mcp_oauth_attempt_store() -> AuthorizationAttemptStore[MCPOAuthFlowState]:
return AuthorizationAttemptStore(
get_cache_backend(),
cache_backend_provider=get_cache_backend,
namespace="mcp",
payload_type=MCPOAuthFlowState,
ttl_seconds=MCP_OAUTH_FLOW_TTL_SECONDS,
Expand Down
Loading
Loading