Add OAuth2 SSO endpoint for external application integration #13227
+1,261
−10
Conversation
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
- Add response_type=oauth_token parameter to /sso/key/generate - Return JSON token response for OAuth2 flows instead of HTML redirect - Support CORS headers for cross-origin requests from external apps - Add oauth_token_flow state handling in SSO callback - Maintain backward compatibility with existing CLI and UI flows - Add comprehensive unit tests for OAuth2 functionality Enables external applications like RooCode to authenticate users via SSO and receive RFC 6749 compliant OAuth2 token responses.
|
The latest updates on your projects. Learn more about Vercel for Git ↗︎
|
- Change return type from Union[RedirectResponse, Dict] to Union[RedirectResponse, JSONResponse] - Resolves MyPy type checking error for OAuth2 JSON response handling
This refactoring consolidates OAuth2-related functionality into reusable utility classes, eliminating code duplication and improving maintainability. **New OAuth2 Utility Classes:** - OAuth2TokenManager: Centralized token response creation per RFC 6749 - OAuth2StateManager: Secure state parameter generation with timestamp/nonce validation - OAuth2URLManager: Unified URL parsing and redirect handling - OAuth2CORSHandler: Standardized CORS headers for OAuth2 responses **Refactored Implementation:** - Replace manual URL parsing with OAuth2URLManager.modify_url_for_oauth_flow() - Enhanced state management with 5-minute expiry and secure nonce generation - Consolidated token responses using OAuth2TokenManager.create_token_response() - Unified CORS handling with OAuth2CORSHandler.get_oauth_cors_headers() **Benefits:** - Eliminated duplicate code across 4+ locations - Enhanced security with proper state validation - Improved maintainability with centralized utilities - Better type safety with proper annotations **Testing:** - All 16 OAuth2 tests pass - Updated tests to use new utility classes - MyPy type checking clean (1022 files) - Ruff linting clean **Files Modified:** - litellm/proxy/management_endpoints/sso_helper_utils.py: +149 lines (new utilities) - litellm/proxy/management_endpoints/ui_sso.py: Refactored to use utilities - tests/.../test_oauth2_flow.py: Updated to use new utility classes - enterprise/.../custom_sso_handler.py: Fixed return type compatibility
Update return type annotation to support OAuth2 JSON responses alongside traditional redirect responses for consistency with main OAuth2 implementation.
This enhancement enables VSCode extension compatibility by supporting OAuth2 redirect flows alongside the existing JSON response format. **New Features:** - redirect_uri parameter support for OAuth2 flows - Dual response mode: redirect vs JSON based on redirect_uri presence - VSCode URI scheme validation (vscode://, vscode-insiders://) - Secure redirect_uri validation with allowlist **OAuth2 Flow Enhancements:** - Enhanced OAuth2URLManager with callback redirect URL building - Secure state management stores redirect_uri with nonce/timestamp validation - OAuth2StateManager extract_redirect_uri() method for state parsing - Comprehensive redirect_uri validation prevents XSS attacks **VSCode Extension Compatibility:** - URL format: /sso/key/generate?response_type=oauth_token&redirect_uri=vscode://ext/callback - Callback format: vscode://ext/callback?access_token=sk-xxx&token_type=Bearer&expires_in=86400 - Maintains full backward compatibility with existing JSON flows **Security:** - Validates redirect_uri against allowed schemes (vscode, https, http) - Prevents malicious redirects (javascript:, data:, file: schemes blocked) - Enhanced state parameter includes redirect_uri with secure validation - 5-minute state expiry with nonce protection **Testing:** - Added 6 new test cases for redirect_uri functionality - Tests cover VSCode URI validation, callback URL building, state management - All 23 OAuth2 tests passing with full coverage - Integration tests for VSCode extension flow structure **Backward Compatibility:** - Without redirect_uri: Returns JSON response (original behavior) - With redirect_uri: Redirects to callback with access_token query param - All existing flows continue to work unchanged **Implementation Details:** - OAuth2URLManager.build_callback_redirect_url() for VSCode callbacks - OAuth2URLManager.validate_redirect_uri() with security checks - Enhanced serve_login_page() and sso_login_redirect() endpoints - Modified auth_callback() to support both redirect and JSON responses This makes LiteLLM OAuth2 consistent with OpenRouter, Glama, and Requesty for seamless VSCode extension integration.
This enhancement significantly expands OAuth2 redirect_uri support to include all popular IDEs, editors, and development tools, making LiteLLM compatible with a wide range of desktop applications and extensions. **New Supported URI Schemes:** - IDE/Editor support: cursor, vscodium, code-oss, fleet, zed, sublime, atom, brackets, nova - Development tools: github-desktop, sourcetree, tower, fork - Legacy editors: textmate, coderunner - Custom schemes via OAUTH_ALLOWED_REDIRECT_SCHEMES environment variable **Key Features:** - Comprehensive URI scheme validation for 15+ popular IDEs/editors - Environment variable OAUTH_ALLOWED_REDIRECT_SCHEMES for custom schemes - Secure validation prevents malicious redirects while supporting legitimate tools - Maintains backward compatibility with existing vscode:// and https:// schemes **Security:** - Validates IDE URIs require proper format: scheme://extension-id/callback-path - Blocks dangerous schemes (javascript:, data:, file:) while allowing dev tools - Custom schemes from environment variable are safely validated - All schemes maintain existing security requirements **Testing:** - Added comprehensive tests for all new URI schemes - Environment variable configuration testing - Cursor editor specific flow testing - Integration tests for popular IDE flows - All 26 OAuth2 tests passing with expanded coverage **Configuration:** Set custom schemes via environment variable: ```bash export OAUTH_ALLOWED_REDIRECT_SCHEMES="myide,customapp,special-editor" ``` **Supported Flows:** - Cursor: cursor://extension/callback?access_token=... - Fleet: fleet://extension/callback?access_token=... - Zed: zed://extension/callback?access_token=... - And 10+ more popular development tools This makes LiteLLM OAuth2 compatible with virtually any desktop IDE or development tool that supports custom URI schemes for OAuth callbacks.
- Update modify_url_for_oauth_flow to preserve redirect_uri parameter - Add comprehensive documentation for redirect_uri handling - Ensure VSCode and IDE extension compatibility during OAuth flow - Maintain backward compatibility for existing OAuth implementations
- Add authentication check for OAuth token flows via Authorization header - Generate new session keys directly for authenticated users without SSO redirect - Implement _check_existing_authentication() to validate existing tokens - Implement _handle_authenticated_oauth_flow() for seamless key generation - Add comprehensive tests for both authenticated and non-authenticated scenarios - Improve user experience for IDE extensions like RooCode, VSCode, Cursor - Maintain full backward compatibility with existing OAuth flows Benefits: - Faster token generation for already-logged-in users - No unnecessary SSO redirects when user has valid session - Better UX for IDE extensions and external applications - Secure token validation before generating new OAuth session keys
7 tasks
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.
Title
Add OAuth2 SSO endpoint for external application integration
Relevant issues
Implements OAuth2 SSO flow to allow external applications (like RooCode) to authenticate users and obtain API tokens via a simple GET request.
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unitType
🆕 New Feature
Changes
OAuth2 SSO Integration for External Applications
This PR adds OAuth2-compatible SSO endpoints that allow external applications to authenticate users and obtain API tokens through a streamlined flow.
Key Features:
https://your-proxy.com/sso/key/generate?response_type=oauth_tokenaccess_token,token_type, andexpires_infields per RFC 6749OAUTH_ALLOWED_ORIGINSenvironment variableImplementation Details:
/sso/key/generateendpoint to detect OAuth token requests viaresponse_type=oauth_tokenparametercreate_oauth_token_response()helper function for standardized token formattingSecurity Features:
secrets.token_urlsafe()no-store,no-cache)Testing:
Usage Example for RooCode:
Response Format:
{ "access_token": "sk-litellm-xxxxxxxxxx", "token_type": "Bearer", "expires_in": 86400 }This implementation enables seamless integration with external applications while maintaining LiteLLM's existing security and authentication standards.