diff --git a/runtime/queries/proto/highlights.scm b/runtime/queries/proto/highlights.scm
index 22df9002f..531209915 100644
--- a/runtime/queries/proto/highlights.scm
+++ b/runtime/queries/proto/highlights.scm
@@ -1,13 +1,27 @@
-(full_ident
-  (identifier) @variable)
+(package
+  (full_ident
+    (identifier) @module))
 
-(full_ident
-  (identifier)
-  (identifier) @variable.member)
+(extend
+  (full_ident
+    (identifier) @type))
+
+(constant
+  (full_ident
+    (identifier) @constant))
 
 (field
   (identifier) @property)
 
+(map_field
+  (identifier) @property)
+
+(oneof
+  (identifier) @type)
+
+(oneof_field
+  (identifier) @property)
+
 (field_option
   (identifier) @property)
 
@@ -17,24 +31,37 @@
 (block_lit
   (identifier) @property)
 
+; Extension option names, e.g. option (foo.bar) = ...
+(option
+  (full_ident
+    (identifier) @variable))
+
+(option
+  (full_ident
+    (identifier)
+    (identifier) @variable.member))
+
 [
-  "extend"
-  "extensions"
-  "oneof"
   "option"
-  "reserved"
   "syntax"
   "edition"
+] @keyword.directive
+
+[
+  "reserved"
   "to"
   "max"
 ] @keyword
 
 [
   "enum"
+  "extend"
+  "extensions"
   "group"
-  "service"
   "message"
   "map"
+  "oneof"
+  "service"
 ] @keyword.type
 
 "rpc" @keyword.function
@@ -60,22 +87,29 @@
 [
   (key_type)
   (type)
+] @type.builtin
+
+[
   (message_name)
   (enum_name)
   (service_name)
-  (rpc_name)
   (message_or_enum_type)
 ] @type
 
+(rpc_name) @function.method
+
 (enum_field
   (identifier) @constant)
 
 (string) @string
 
+(import
+  path: (string) @string.special.path)
+
 [
   "\"proto3\""
   "\"proto2\""
-] @string.special
+] @string.special.symbol
 
 (escape_sequence) @string.escape
 
@@ -90,9 +124,6 @@
 
 (comment) @comment @spell
 
-((comment) @comment.documentation
-  (#lua-match? @comment.documentation "^/[*][*][^*].*[*]/$"))
-
 [
   "("
   ")"
diff --git a/tests/query/highlights/proto/test.proto b/tests/query/highlights/proto/test.proto
index fdc610782..7f8c3abdf 100644
--- a/tests/query/highlights/proto/test.proto
+++ b/tests/query/highlights/proto/test.proto
@@ -1,11 +1,11 @@
 syntax = "proto2";
-//       ^^^^^^^^ @string.special
+//       ^^^^^^^^ @string.special.symbol
 package sls.asfd.asfd;
-//      ^^^ @variable
+//      ^^^ @module
 //         ^ @punctuation.delimiter
-//          ^^^^ @variable.member
+//          ^^^^ @module
 //              ^ @punctuation.delimiter
-//               ^^^^ @variable.member
+//               ^^^^ @module
 //                   ^ @punctuation.delimiter
 
 import "github.com/gogo/protobuf/gogoproto/gogo.proto";
@@ -19,7 +19,11 @@ import "github.com/gogo/protobuf/gogoproto/gogo.proto";
 //<- @keyword.import
 
 option (gogoproto.sizer_all) = true;
+//      ^^^^^^^^^ @variable
+//                ^^^^^^^^^ @variable.member
 option (gogoproto.marshaler_all) = true;
+//      ^^^^^^^^^ @variable
+//                ^^^^^^^^^^^^^ @variable.member
 option (gogoproto.unmarshaler_all.foo) = true;
 //      ^^^^^^^^^ @variable
 //                ^^^^^^^^^^^^^^^ @variable.member
@@ -30,12 +34,12 @@ message LogContent
 //<- @keyword.type
 {
     required string Key = 1;
-//           ^^^^^^ @type
+//           ^^^^^^ @type.builtin
 //                  ^^^ @property
 //                      ^ @operator
 //                        ^ @number
     required string Value = 2;
-//           ^^^^^^ @type
+//           ^^^^^^ @type.builtin
 //                  ^^^^^ @property
     required CustomType Value = 2;
 //  ^^^^^^^^ @keyword.modifier
@@ -46,7 +50,7 @@ message LogContent
     optional JsonFormat json_format = 6 [
 //  ^^^^^^^^ @keyword.modifier
       retention = RETENTION_RUNTIME,
-//                ^^^^^^^^^^^^^^^^^ @variable
+//                ^^^^^^^^^^^^^^^^^ @constant
 //                                 ^ @punctuation.delimiter
       targets = TARGET_TYPE_MESSAGE,
 //    ^^^^^^^ @property
@@ -60,3 +64,4 @@ message LogContent
     ];
 }
 
+