-
Notifications
You must be signed in to change notification settings - Fork 2.5k
Add PKCE integration for Hugging Face #7881
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.
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
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -55,6 +55,7 @@ import { formatLanguage } from "../../shared/language" | |
| import { WebviewMessage } from "../../shared/WebviewMessage" | ||
| import { EMBEDDING_MODEL_PROFILES } from "../../shared/embeddingModels" | ||
| import { ProfileValidator } from "../../shared/ProfileValidator" | ||
| import { HUGGING_FACE_OAUTH_CLIENT_ID } from "../../shared/oauth-constants" | ||
|
|
||
| import { Terminal } from "../../integrations/terminal/Terminal" | ||
| import { downloadTask } from "../../integrations/misc/export-markdown" | ||
|
|
@@ -1424,6 +1425,68 @@ export class ClineProvider | |
| await this.upsertProviderProfile(currentApiConfigName, newConfiguration) | ||
| } | ||
|
|
||
| // HuggingFace | ||
|
|
||
| async handleHuggingFaceCallback(code: string, returnedState?: string) { | ||
| let { apiConfiguration, currentApiConfigName = "default" } = await this.getState() | ||
|
|
||
| try { | ||
| // Retrieve stored PKCE verifier and state from extension state | ||
| const pkceSecret = await this.context.secrets.get("huggingFacePkce") | ||
| const pkceData = pkceSecret ? JSON.parse(pkceSecret) : undefined | ||
|
|
||
| if (!pkceData || !pkceData.verifier || !pkceData.state) { | ||
| throw new Error("PKCE verifier or state not found in extension state.") | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Error messages don't follow the i18n pattern used elsewhere in the codebase. Should use the translation function for consistency. |
||
| } | ||
|
|
||
| // Optional state validation (if state was provided in the callback) | ||
| if (returnedState && pkceData.state && returnedState !== pkceData.state) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The state validation is incomplete. It only checks for mismatch but doesn't verify the state format or prevent replay attacks. Consider adding timestamp validation to ensure the state is fresh (e.g., not older than 10 minutes). |
||
| // Clear stored data before throwing to avoid reuse | ||
| await this.context.secrets.delete("huggingFacePkce") | ||
| throw new Error("OAuth state mismatch.") | ||
| } | ||
|
|
||
| const verifier: string = pkceData.verifier | ||
|
|
||
| // Clear PKCE data to prevent reuse | ||
| await this.context.secrets.delete("huggingFacePkce") | ||
|
|
||
| const redirectUri = `${vscode.env.uriScheme}://${Package.publisher}.${Package.name}/huggingface` | ||
|
|
||
| const params = new URLSearchParams() | ||
| params.append("grant_type", "authorization_code") | ||
| params.append("code", code) | ||
| params.append("client_id", HUGGING_FACE_OAUTH_CLIENT_ID) | ||
| params.append("code_verifier", verifier) | ||
| params.append("redirect_uri", redirectUri) | ||
|
|
||
| const response = await axios.post("https://huggingface.co/oauth/token", params, { | ||
| headers: { "Content-Type": "application/x-www-form-urlencoded" }, | ||
| }) | ||
|
|
||
| const accessToken: string | undefined = response.data?.access_token | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Missing validation of the access token format. Consider validating the token structure and checking that it's a non-empty string before using it. |
||
|
|
||
| if (!accessToken) { | ||
| throw new Error("Invalid response from Hugging Face token endpoint") | ||
| } | ||
|
|
||
| const newConfiguration: ProviderSettings = { | ||
| ...apiConfiguration, | ||
| apiProvider: "huggingface", | ||
| huggingFaceApiKey: accessToken, | ||
| huggingFaceModelId: apiConfiguration?.huggingFaceModelId || "meta-llama/Llama-3.3-70B-Instruct", | ||
| huggingFaceInferenceProvider: apiConfiguration?.huggingFaceInferenceProvider || "auto", | ||
| } | ||
|
|
||
| await this.upsertProviderProfile(currentApiConfigName, newConfiguration) | ||
| } catch (error) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Error handling doesn't clean up stored PKCE data on all error paths. The secrets should be deleted in a finally block to ensure cleanup happens regardless of success or failure. |
||
| this.log( | ||
| `Error exchanging code for Hugging Face access token: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`, | ||
| ) | ||
| throw error | ||
| } | ||
| } | ||
|
|
||
| // Requesty | ||
|
|
||
| async handleRequestyCallback(code: string) { | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -978,6 +978,19 @@ export const webviewMessageHandler = async ( | |
| vscode.env.openExternal(vscode.Uri.parse(message.url)) | ||
| } | ||
| break | ||
| case "storeHuggingFacePkce": { | ||
| // Store PKCE verifier/state as a secret to avoid typing constraints on ContextProxy keys | ||
| const verifier = message.values?.verifier | ||
| const state = message.values?.state | ||
| if (typeof verifier === "string" && typeof state === "string" && verifier.length > 0 && state.length > 0) { | ||
| try { | ||
| await provider.context.secrets.store("huggingFacePkce", JSON.stringify({ verifier, state })) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Using context.secrets for temporary PKCE data isn't ideal as secrets are meant for long-term credentials. Consider using extension global state or in-memory storage for temporary OAuth flow data. |
||
| } catch (error) { | ||
| console.error("Failed to store Hugging Face PKCE data:", error) | ||
| } | ||
| } | ||
| break | ||
| } | ||
| case "checkpointDiff": | ||
| const result = checkoutDiffPayloadSchema.safeParse(message.payload) | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,10 @@ | ||
| /** | ||
| * OAuth Client IDs and Constants | ||
| * | ||
| * These are public OAuth client identifiers used in OAuth flows. | ||
| * They are safe to be exposed as they identify the application to OAuth providers. | ||
| * Unlike client secrets, these are designed to be public in OAuth 2.0 public clients. | ||
| */ | ||
|
|
||
| // Hugging Face OAuth client ID for PKCE flow | ||
| export const HUGGING_FACE_OAUTH_CLIENT_ID = "aba045f7-aceb-4e53-9247-5c85d7c2b7cb" | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Security concern: While OAuth client IDs are public, hardcoding them makes rotation difficult and prevents environment-specific configurations. Consider using environment variables or configuration files to allow overriding for different environments (dev/staging/prod). |
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,16 +1,17 @@ | ||
| import { useCallback, useState, useEffect, useMemo } from "react" | ||
| import { useEvent } from "react-use" | ||
| import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react" | ||
| import { VSCodeTextField, VSCodeButton } from "@vscode/webview-ui-toolkit/react" | ||
| import pkceChallenge from "pkce-challenge" | ||
|
|
||
| import type { ProviderSettings } from "@roo-code/types" | ||
|
|
||
| import { ExtensionMessage } from "@roo/ExtensionMessage" | ||
| import { vscode } from "@src/utils/vscode" | ||
| import { useAppTranslation } from "@src/i18n/TranslationContext" | ||
| import { VSCodeButtonLink } from "@src/components/common/VSCodeButtonLink" | ||
| import { SearchableSelect, type SearchableSelectOption } from "@src/components/ui" | ||
| import { cn } from "@src/lib/utils" | ||
| import { formatPrice } from "@/utils/formatPrice" | ||
| import { getHuggingFaceAuthUrl } from "@src/oauth/urls" | ||
|
|
||
| import { inputEventTransform } from "../transforms" | ||
|
|
||
|
|
@@ -39,9 +40,10 @@ type HuggingFaceProps = { | |
| value: ProviderSettings[keyof ProviderSettings], | ||
| isUserAction?: boolean, | ||
| ) => void | ||
| uriScheme?: string | ||
| } | ||
|
|
||
| export const HuggingFace = ({ apiConfiguration, setApiConfigurationField }: HuggingFaceProps) => { | ||
| export const HuggingFace = ({ apiConfiguration, setApiConfigurationField, uriScheme }: HuggingFaceProps) => { | ||
| const { t } = useAppTranslation() | ||
| const [models, setModels] = useState<HuggingFaceModel[]>([]) | ||
| const [loading, setLoading] = useState(false) | ||
|
|
@@ -109,6 +111,28 @@ export const HuggingFace = ({ apiConfiguration, setApiConfigurationField }: Hugg | |
| setApiConfigurationField, | ||
| ]) | ||
|
|
||
| // Start OAuth with PKCE | ||
| const handleStartOauth = useCallback(async () => { | ||
| try { | ||
| // Generate PKCE challenge using the library | ||
| const pkce = await pkceChallenge() | ||
| const state = crypto.randomUUID() // Use built-in UUID for state | ||
|
|
||
| // Store verifier/state in extension (secrets) | ||
| vscode.postMessage({ | ||
| type: "storeHuggingFacePkce", | ||
| values: { verifier: pkce.code_verifier, state }, | ||
| }) | ||
|
|
||
| const authUrl = getHuggingFaceAuthUrl(uriScheme, pkce.code_challenge, state) | ||
|
|
||
| // Open externally via extension | ||
| vscode.postMessage({ type: "openExternal", url: authUrl }) | ||
| } catch (e) { | ||
| console.error("Failed to start Hugging Face OAuth:", e) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Missing error handling for PKCE challenge generation. The catch block only logs to console but doesn't notify the user. Consider showing a user-friendly error message. |
||
| } | ||
| }, [uriScheme]) | ||
|
|
||
| const handleModelSelect = (modelId: string) => { | ||
| setApiConfigurationField("huggingFaceModelId", modelId) | ||
| // Reset provider selection when model changes | ||
|
|
@@ -180,9 +204,9 @@ export const HuggingFace = ({ apiConfiguration, setApiConfigurationField }: Hugg | |
| </div> | ||
|
|
||
| {!apiConfiguration?.huggingFaceApiKey && ( | ||
| <VSCodeButtonLink href="https://huggingface.co/settings/tokens" appearance="secondary"> | ||
| <VSCodeButton appearance="primary" onClick={handleStartOauth} style={{ width: "100%" }}> | ||
| {t("settings:providers.getHuggingFaceApiKey")} | ||
| </VSCodeButtonLink> | ||
| </VSCodeButton> | ||
| )} | ||
|
|
||
| <div className="flex flex-col gap-2"> | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Consider extracting common OAuth handling logic. The pattern here is duplicated across Glama, OpenRouter, and Requesty providers. A shared OAuth handler would reduce code duplication and make maintenance easier.