diff --git a/vscode-extension/package.json b/vscode-extension/package.json
index 2bc3b25..0b5d6bb 100644
--- a/vscode-extension/package.json
+++ b/vscode-extension/package.json
@@ -83,7 +83,7 @@
     "compile": "bun run compile:client && bun run compile:server",
     "compile:client": "bun build client/src/extension.ts --outdir client/dist --target node --format cjs --external vscode",
     "compile:server": "bun build server/src/server.ts --outdir server/dist --target node --format cjs",
-    "watch": "bun run compile:client --watch",
+    "watch": "bun run compile:client --watch & bun run compile:server --watch",
     "package": "bun run compile:client --minify && bun run compile:server --minify",
     "check-types": "tsc --noEmit",
     "build-and-install": "bun run package && bunx @vscode/vsce package --allow-missing-repository && code --install-extension shrimp-*.vsix"
diff --git a/vscode-extension/server/src/semanticTokens.ts b/vscode-extension/server/src/semanticTokens.ts
index 08ab812..f7f60b7 100644
--- a/vscode-extension/server/src/semanticTokens.ts
+++ b/vscode-extension/server/src/semanticTokens.ts
@@ -55,10 +55,13 @@ function walkTree(node: SyntaxNode, document: TextDocument, builder: SemanticTok
 }
 
 // Map Lezer node IDs to semantic token type indices and modifiers
-function getTokenType(nodeTypeId: number, parentTypeId?: number): { type: number; modifiers: number } | undefined {
+function getTokenType(
+  nodeTypeId: number,
+  parentTypeId?: number
+): { type: number; modifiers: number } | undefined {
   switch (nodeTypeId) {
     case Terms.Identifier:
-      // Check parent to determine if this identifier is a function call or variable
+      // Check parent to determine context
       if (parentTypeId === Terms.FunctionCall) {
         return {
           type: TOKEN_TYPES.indexOf(SemanticTokenTypes.function),
@@ -77,12 +80,36 @@ function getTokenType(nodeTypeId: number, parentTypeId?: number): { type: number
           modifiers: 0,
         }
       }
+      if (parentTypeId === Terms.Params) {
+        return {
+          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.parameter),
+          modifiers: 0,
+        }
+      }
+      if (parentTypeId === Terms.DotGet) {
+        return {
+          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
+          modifiers: 0,
+        }
+      }
       // Otherwise it's a regular variable
       return {
         type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
         modifiers: 0,
       }
 
+    case Terms.IdentifierBeforeDot:
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
+        modifiers: 0,
+      }
+
+    case Terms.NamedArgPrefix:
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
+        modifiers: 0,
+      }
+
     case Terms.AssignableIdentifier:
       return {
         type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
@@ -124,24 +151,12 @@ function getTokenType(nodeTypeId: number, parentTypeId?: number): { type: number
 
     case Terms.keyword:
     case Terms.Do:
+    case Terms.colon:
       return {
         type: TOKEN_TYPES.indexOf(SemanticTokenTypes.keyword),
         modifiers: 0,
       }
 
-    case Terms.Params:
-    case Terms.NamedParam:
-      return {
-        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.parameter),
-        modifiers: 0,
-      }
-
-    case Terms.DotGet:
-      return {
-        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
-        modifiers: 0,
-      }
-
     case Terms.Regex:
       return {
         type: TOKEN_TYPES.indexOf(SemanticTokenTypes.regexp),
@@ -154,12 +169,6 @@ function getTokenType(nodeTypeId: number, parentTypeId?: number): { type: number
         modifiers: 0,
       }
 
-    case Terms.NamedArg:
-      return {
-        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
-        modifiers: 0,
-      }
-
     default:
       return undefined
   }