Tool OAuth

def _oauth_get_credentials(
    self, redirect_uri: str, system_credentials: Mapping[str, Any], request: Request
) -> ToolOAuthCredentials:
    """
    Exchange authorization code for access token and refresh token. This is called
	to create ONE credential set for one account connection.
    """
    # Extract authorization code from OAuth callback
    code = request.args.get("code")
    if not code:
        raise ToolProviderOAuthError("Authorization code not provided")
    
    # Check for authorization errors from OAuth provider
    error = request.args.get("error")
    if error:
        error_description = request.args.get("error_description", "")
        raise ToolProviderOAuthError(f"OAuth authorization failed: {error} - {error_description}")
    
    # Exchange authorization code for tokens using OAuth Client Setup credentials

	# Assemble Gmail-specific payload
    data = {
        "client_id": system_credentials["client_id"],        # From OAuth Client Setup
        "client_secret": system_credentials["client_secret"], # From OAuth Client Setup
        "code": code,                                        # From user's authorization
        "grant_type": "authorization_code",                  # Standard OAuth flow type
        "redirect_uri": redirect_uri,                        # Must exactly match authorization URL
    }
    
    headers = {"Content-Type": "application/x-www-form-urlencoded"}
    
    try:
        response = requests.post(
            self._TOKEN_URL,
            data=data,
            headers=headers,
            timeout=10
        )
        response.raise_for_status()
        
        token_data = response.json()
        
        # Handle OAuth provider errors in response
        if "error" in token_data:
            error_desc = token_data.get('error_description', token_data['error'])
            raise ToolProviderOAuthError(f"Token exchange failed: {error_desc}")
        
        access_token = token_data.get("access_token")
        if not access_token:
            raise ToolProviderOAuthError("No access token received from provider")
        
        # Build credentials dict matching your credentials_schema
        credentials = {
            "access_token": access_token,
            "token_type": token_data.get("token_type", "Bearer"),
        }
        
        # Include refresh token if provided (critical for long-term access)
        refresh_token = token_data.get("refresh_token")
        if refresh_token:
            credentials["refresh_token"] = refresh_token
        
        # Handle token expiration - some providers don't provide expires_in
        expires_in = token_data.get("expires_in", 3600)  # Default to 1 hour
        expires_at = int(time.time()) + expires_in
        
        return ToolOAuthCredentials(credentials=credentials, expires_at=expires_at)
        
    except requests.RequestException as e:
        raise ToolProviderOAuthError(f"Network error during token exchange: {str(e)}")
    except Exception as e:
        raise ToolProviderOAuthError(f"Failed to exchange authorization code: {str(e)}")

def _oauth_refresh_credentials(
    self, redirect_uri: str, system_credentials: Mapping[str, Any], credentials: Mapping[str, Any]
) -> ToolOAuthCredentials:
    """
    Refresh the credentials using the refresh token. 
	Dify calls this automatically when tokens expire.
    """
    refresh_token = credentials.get("refresh_token")
    if not refresh_token:
        raise ToolProviderOAuthError("No refresh token available")

    # Standard OAuth refresh token flow
    data = {
        "client_id": system_credentials["client_id"],       # From OAuth Client Setup
        "client_secret": system_credentials["client_secret"], # From OAuth Client Setup
        "refresh_token": refresh_token,                     # From previous authorization
        "grant_type": "refresh_token",                      # OAuth refresh flow
    }

    headers = {"Content-Type": "application/x-www-form-urlencoded"}

    try:
        response = requests.post(
            self._TOKEN_URL,
            data=data,
            headers=headers,
            timeout=10
        )
        response.raise_for_status()

        token_data = response.json()

        # Handle refresh errors
        if "error" in token_data:
            error_desc = token_data.get('error_description', token_data['error'])
            raise ToolProviderOAuthError(f"Token refresh failed: {error_desc}")

        access_token = token_data.get("access_token")
        if not access_token:
            raise ToolProviderOAuthError("No access token received from provider")

        # Build new credentials, preserving existing refresh token
        new_credentials = {
            "access_token": access_token,
            "token_type": token_data.get("token_type", "Bearer"),
            "refresh_token": refresh_token,  # Keep existing refresh token
        }

        # Handle token expiration
        expires_in = token_data.get("expires_in", 3600)

        # update refresh token if new one provided
        new_refresh_token = token_data.get("refresh_token")
        if new_refresh_token:
            new_credentials["refresh_token"] = new_refresh_token

        # Calculate new expiration timestamp for Dify's token management
        expires_at = int(time.time()) + expires_in

        return ToolOAuthCredentials(credentials=new_credentials, expires_at=expires_at)

    except requests.RequestException as e:
        raise ToolProviderOAuthError(f"Network error during token refresh: {str(e)}")
    except Exception as e:
        raise ToolProviderOAuthError(f"Failed to refresh credentials: {str(e)}")