diff --git a/src/PowerShellEditorServices/Services/TextDocument/Handlers/PsesSemanticTokensHandler.cs b/src/PowerShellEditorServices/Services/TextDocument/Handlers/PsesSemanticTokensHandler.cs
index e5e69e60f..494843bb5 100644
--- a/src/PowerShellEditorServices/Services/TextDocument/Handlers/PsesSemanticTokensHandler.cs
+++ b/src/PowerShellEditorServices/Services/TextDocument/Handlers/PsesSemanticTokensHandler.cs
@@ -108,6 +108,11 @@ private static SemanticTokenType MapSemanticTokenType(Token token)
                 return SemanticTokenType.Operator;
             }
 
+            if ((token.TokenFlags & TokenFlags.AttributeName) != 0)
+            {
+                return SemanticTokenType.Decorator;
+            }
+
             if ((token.TokenFlags & TokenFlags.TypeName) != 0)
             {
                 return SemanticTokenType.Type;
@@ -142,8 +147,8 @@ private static SemanticTokenType MapSemanticTokenType(Token token)
                 case TokenKind.Number:
                     return SemanticTokenType.Number;
 
-                case TokenKind.Generic:
-                    return SemanticTokenType.Function;
+                case TokenKind.Label:
+                    return SemanticTokenType.Label;
             }
 
             return null;
diff --git a/test/PowerShellEditorServices.Test/Language/SemanticTokenTest.cs b/test/PowerShellEditorServices.Test/Language/SemanticTokenTest.cs
index c9f3c01d4..196fa9ca6 100644
--- a/test/PowerShellEditorServices.Test/Language/SemanticTokenTest.cs
+++ b/test/PowerShellEditorServices.Test/Language/SemanticTokenTest.cs
@@ -20,7 +20,8 @@ public void TokenizesFunctionElements()
         {
             const string text = @"
 function Get-Sum {
-    param( [int]$a, [int]$b )
+    param( [parameter()] [int]$a, [int]$b )
+    :loopLabel while (0) {break loopLabel}
     return $a + $b
 }
 ";
@@ -38,10 +39,21 @@ function Get-Sum {
                     case "function":
                     case "param":
                     case "return":
+                    case "while":
+                    case "break":
                         Assert.Single(mappedTokens, sToken => SemanticTokenType.Keyword == sToken.Type);
                         break;
-                    case "Get-Sum":
-                        Assert.Single(mappedTokens, sToken => SemanticTokenType.Function == sToken.Type);
+                    case "parameter":
+                        Assert.Single(mappedTokens, sToken => SemanticTokenType.Decorator == sToken.Type);
+                        break;
+                    case "0":
+                        Assert.Single(mappedTokens, sToken => SemanticTokenType.Number == sToken.Type);
+                        break;
+                    case ":loopLabel":
+                        Assert.Single(mappedTokens, sToken => SemanticTokenType.Label == sToken.Type);
+                        break;
+                    case "loopLabel":
+                        Assert.Single(mappedTokens, sToken => SemanticTokenType.Property == sToken.Type);
                         break;
                     case "$a":
                     case "$b":
@@ -74,7 +86,6 @@ public void TokenizesStringExpansion()
             Token stringExpandableToken = scriptFile.ScriptTokens[1];
             mappedTokens = new List<SemanticToken>(PsesSemanticTokensHandler.ConvertToSemanticTokens(stringExpandableToken));
             Assert.Collection(mappedTokens,
-                sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type),
                 sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type),
                 sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type)
             );
@@ -103,7 +114,11 @@ function Get-A*A {
                         Assert.Single(mappedTokens, sToken => SemanticTokenType.Keyword == sToken.Type);
                         break;
                     case "Get-A*A":
-                        Assert.Single(mappedTokens, sToken => SemanticTokenType.Function == sToken.Type);
+                        if (t.TokenFlags.HasFlag(TokenFlags.CommandName))
+                        {
+                            Assert.Single(mappedTokens, sToken => SemanticTokenType.Function == sToken.Type);
+                        }
+
                         break;
                 }
             }