diff --git a/code/custom/4coder_tree_sitter.cpp b/code/custom/4coder_tree_sitter.cpp
index 72231c95..8ebb08e8 100644
--- a/code/custom/4coder_tree_sitter.cpp
+++ b/code/custom/4coder_tree_sitter.cpp
@@ -41,7 +41,7 @@ tree_sitter_query_new(Application_Links* app, TSLanguage* language, String_Const
   return result;
 }
 
-function void
+function Tree_Sitter_Language_Definition*
 tree_sitter_register_language(String_Const_u8 ext, TSLanguage* language, Tree_Sitter_Language_Queries queries, Tree_Sitter_Language_Flags flags)
 {
   Tree_Sitter_Language_Definition* lang = 0;
@@ -67,6 +67,8 @@ tree_sitter_register_language(String_Const_u8 ext, TSLanguage* language, Tree_Si
     lang->queries = queries;
     lang->flags = flags;
   }
+
+  return lang;
 }
 
 function Tree_Sitter_Language_Definition*
@@ -998,9 +1000,58 @@ tree_sitter_code_index_update_tick(Application_Links* app)
 // Token Highlighting
 ////////////////////////////////////////////////////////////////////
 
+function Tree_Sitter_Query_Color_Map*
+tree_sitter_get_language_query_color_map(Tree_Sitter_Language_Definition* lang, u64 hash, u64 desired_value)
+{
+  Tree_Sitter_Query_Color_Map* result = 0;
+  u64 index = hash & ArrayCount(lang->colors);
+  u64 start = index;
+  do {
+    Tree_Sitter_Query_Color_Map* at = &lang->colors[index];
+    if (at->hash == desired_value) {
+      result = at;
+      break;
+    }
+    if (at->hash == 0) break;
+    index += 1;
+    if (index >= ArrayCount(lang->colors)) index = 0;
+  } while (index != start);
+  return result;
+}
+
+function Tree_Sitter_Query_Color_Map*
+tree_sitter_get_language_query_color_map(Tree_Sitter_Language_Definition* lang, String_Const_u8 capture_name)
+{
+  u64 hash = table_hash_u8(capture_name.str, capture_name.size);
+  return tree_sitter_get_language_query_color_map(lang, hash, hash);
+}
+
+function void
+tree_sitter_set_language_query_color_map(Application_Links* app, Tree_Sitter_Language_Definition* lang, String_Const_u8 capture_name, String_Const_u8 color_id)
+{
+  u64 hash = table_hash_u8(capture_name.str, capture_name.size);
+  Tree_Sitter_Query_Color_Map* map = tree_sitter_get_language_query_color_map(lang, hash, 0);
+  Assert(map != 0);
+  map->hash = hash;
+  map->capture_name = capture_name;
+  map->color_id = managed_id_get(app, SCu8("colors"), color_id);
+}
+
+function void
+tree_sitter_languages_init_default_colors(Application_Links* app, Tree_Sitter_Language_Definition* lang)
+{
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("function"), SCu8("defcolor_function"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("type"), SCu8("defcolor_type"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("constant"), SCu8("defcolor_int_constnat"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("keyword"), SCu8("defcolor_keyword"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("string"), SCu8("defcolor_str_constant"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("comment"), SCu8("defcolor_comment"));
+}
+
 function bool
 tree_sitter_highlight_node(
   Application_Links* app,
+  Tree_Sitter_Language_Definition* lang,
   TSQuery* query,
   TSNode top_node,
   TSQueryCursor* query_cursor,
@@ -1020,24 +1071,11 @@ tree_sitter_highlight_node(
     const char* tmp = ts_query_capture_name_for_id(query, capture.index, &length);
     String_Const_u8 capture_name = SCu8((char*)tmp, length);
 
-    // TODO(PS): let each language provide a custom override
-    String_Const_u8 color_name = SCu8("defcolor_text_default");
-    if (string_match(capture_name, SCu8("function"))) {
-      color_name = SCu8("defcolor_function");
-    } else if (string_match(capture_name, SCu8("type"))) {
-      color_name = SCu8("defcolor_type");
-    } else if (string_match(capture_name, SCu8("constant")) || string_match(capture_name, SCu8("number"))) {
-      color_name = SCu8("defcolor_int_constant");
-    } else if (string_match(capture_name, SCu8("keyword"))) {
-      color_name = SCu8("defcolor_keyword");
-    } else if (string_match(capture_name, SCu8("string"))) {
-      color_name = SCu8("defcolor_str_constant");
-    } else if (string_match(capture_name, SCu8("comment"))) {
-      color_name = SCu8("defcolor_comment");
-    }
+    Managed_ID color_id = managed_id_get(app, SCu8("colors"), SCu8("defcolor_text_default"));;
+    Tree_Sitter_Query_Color_Map* color_map = tree_sitter_get_language_query_color_map(lang, capture_name);
+    if (color_map) color_id = color_map->color_id;
 
     Range_i64 highlight_range = tree_sitter_node_to_range(node);
-    Managed_ID color_id = managed_id_get(app, SCu8("colors"), color_name);
     if (color_id != 0)
     {
       paint_text_color_fcolor(app, text_layout_id, highlight_range, fcolor_id(color_id));
@@ -1088,20 +1126,20 @@ draw_tree_sitter_node_colors(Application_Links* app, Text_Layout_ID text_layout_
           TSNode node = ts_tree_cursor_current_node(&tree_cursor);
           Range_i64 child_range = tree_sitter_node_to_range(node);
           if (child_range.start > visible_range.end) break;
-          tree_sitter_highlight_node(app, query, node, query_cursor, text_layout_id);
+          tree_sitter_highlight_node(app, lang, query, node, query_cursor, text_layout_id);
         } while(ts_tree_cursor_goto_next_sibling(&tree_cursor));
       }
       else
       {
         // Pathological case - just highligh the whole document. This is probably bad
-        tree_sitter_highlight_node(app, query, root, query_cursor, text_layout_id);
+        tree_sitter_highlight_node(app, lang, query, root, query_cursor, text_layout_id);
       }
     }
     else
     {
       bool should_ascend = false;
       do {
-        should_ascend = !tree_sitter_highlight_node(app, query, visible_container, query_cursor, text_layout_id);
+        should_ascend = !tree_sitter_highlight_node(app, lang, query, visible_container, query_cursor, text_layout_id);
         if (should_ascend) visible_container = ts_node_parent(visible_container);
       } while (should_ascend && !ts_node_is_null(visible_container));
     }
diff --git a/code/custom/4coder_tree_sitter.h b/code/custom/4coder_tree_sitter.h
index 700233bb..9a410b53 100644
--- a/code/custom/4coder_tree_sitter.h
+++ b/code/custom/4coder_tree_sitter.h
@@ -23,6 +23,13 @@ enum
   Tree_Sitter_Language_Can_Receive_Virtual_Indent = 0x1
 };
 
+struct Tree_Sitter_Query_Color_Map
+{
+  u64 hash;
+  String_Const_u8 capture_name;
+  Managed_ID color_id;
+};
+
 struct Tree_Sitter_Language_Definition
 {
   String_Const_u8 extension;
@@ -32,6 +39,8 @@ struct Tree_Sitter_Language_Definition
   Tree_Sitter_Language_Queries queries;
   Tree_Sitter_Language_Flags flags;
 
+  Tree_Sitter_Query_Color_Map colors[1000];
+
   Tree_Sitter_Language_Definition* next;
 };
 
@@ -114,8 +123,12 @@ b8 use_tree_sitter_token_coloring = true;
 ///////////////////////////////////////////////////////
 // API
 
+function Tree_Sitter_Query_Color_Map* tree_sitter_get_language_query_color_map(Tree_Sitter_Language_Definition* lang, String_Const_u8 capture_name);
+function void tree_sitter_set_language_query_color_map(Application_Links* app, Tree_Sitter_Language_Definition* lang, String_Const_u8 capture_name, String_Const_u8 color_id);
+function void tree_sitter_languages_init_default_colors(Application_Links* app, Tree_Sitter_Language_Definition* lang);
+
 function TSQuery* tree_sitter_query_new(Application_Links* app, TSLanguage* language, String_Const_u8 query_string);
-function void tree_sitter_register_language(String_Const_u8 ext, TSLanguage* language, Tree_Sitter_Language_Queries queries, Tree_Sitter_Language_Flags flags);
+function Tree_Sitter_Language_Definition* tree_sitter_register_language(String_Const_u8 ext, TSLanguage* language, Tree_Sitter_Language_Queries queries, Tree_Sitter_Language_Flags flags);
 function Tree_Sitter_Language_Definition* tree_sitter_language_for_buffer(Application_Links* app, Buffer_ID buffer_id);
 
 function Tree_Sitter_Code_Index_Update_State tree_sitter_code_index_update_full_file(Application_Links *app, Buffer_ID buffer_id);
diff --git a/code/custom/languages/tree_sitter_bash.h b/code/custom/languages/tree_sitter_bash.h
index f0414be1..deb7bcfb 100644
--- a/code/custom/languages/tree_sitter_bash.h
+++ b/code/custom/languages/tree_sitter_bash.h
@@ -74,7 +74,10 @@ tree_sitter_register_bash(Application_Links* app)
   TSLanguage* language = tree_sitter_bash();
   Tree_Sitter_Language_Queries queries = {};
   queries.ptr[Tree_Sitter_Language_Query_Highlights] = tree_sitter_query_new(app, language, TS_BASH_HIGHLIGHT_QUERY);
-  tree_sitter_register_language(SCu8("sh"), language, queries, 0);
+
+  Tree_Sitter_Language_Definition* lang;
+  lang = tree_sitter_register_language(SCu8("sh"), language, queries, 0);
+  tree_sitter_languages_init_default_colors(app, lang);
 }
 
 #endif //TREE_SITTER_BASH_H
diff --git a/code/custom/languages/tree_sitter_cpp.h b/code/custom/languages/tree_sitter_cpp.h
index 72353565..509f121d 100644
--- a/code/custom/languages/tree_sitter_cpp.h
+++ b/code/custom/languages/tree_sitter_cpp.h
@@ -211,11 +211,22 @@ tree_sitter_register_cpp(Application_Links* app)
   Tree_Sitter_Language_Flags flags = (
      Tree_Sitter_Language_Can_Receive_Virtual_Indent
   );
-  tree_sitter_register_language(SCu8("c"),   language, queries, flags);
-  tree_sitter_register_language(SCu8("cpp"), language, queries, flags);
-  tree_sitter_register_language(SCu8("h"),   language, queries, flags);
-  tree_sitter_register_language(SCu8("hpp"), language, queries, flags);
-  tree_sitter_register_language(SCu8("cc"),  language, queries, flags);
+
+  Tree_Sitter_Language_Definition* lang;
+  lang = tree_sitter_register_language(SCu8("c"),   language, queries, flags);
+  tree_sitter_languages_init_default_colors(app, lang);
+
+  lang = tree_sitter_register_language(SCu8("cpp"), language, queries, flags);
+  tree_sitter_languages_init_default_colors(app, lang);
+
+  lang = tree_sitter_register_language(SCu8("h"),   language, queries, flags);
+  tree_sitter_languages_init_default_colors(app, lang);
+
+  lang = tree_sitter_register_language(SCu8("hpp"), language, queries, flags);
+  tree_sitter_languages_init_default_colors(app, lang);
+
+  lang = tree_sitter_register_language(SCu8("cc"),  language, queries, flags);
+  tree_sitter_languages_init_default_colors(app, lang);
 }
 
 #endif //TREE_SITTER_CPP_H
diff --git a/code/custom/languages/tree_sitter_jai.h b/code/custom/languages/tree_sitter_jai.h
index 6de8af95..2950dd73 100644
--- a/code/custom/languages/tree_sitter_jai.h
+++ b/code/custom/languages/tree_sitter_jai.h
@@ -299,7 +299,16 @@ tree_sitter_register_jai (Application_Links* app)
   Tree_Sitter_Language_Flags flags = (
      Tree_Sitter_Language_Can_Receive_Virtual_Indent
   );
-  tree_sitter_register_language(SCu8("jai"), language, queries, flags);
+  Tree_Sitter_Language_Definition* lang;
+  lang = tree_sitter_register_language(SCu8("jai"), language, queries, flags);
+
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("function"), SCu8("defcolor_function"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("function.call"), SCu8("defcolor_function"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("type"), SCu8("defcolor_type"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("constant"), SCu8("defcolor_int_constnat"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("keyword"), SCu8("defcolor_keyword"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("string"), SCu8("defcolor_str_constant"));
+  tree_sitter_set_language_query_color_map(app, lang, SCu8("comment"), SCu8("defcolor_comment"));
 }
 
 #endif //TREE_SITTER_JAI_H
diff --git a/code/custom/languages/tree_sitter_ts.h b/code/custom/languages/tree_sitter_ts.h
index 392cc192..4744b575 100644
--- a/code/custom/languages/tree_sitter_ts.h
+++ b/code/custom/languages/tree_sitter_ts.h
@@ -305,7 +305,9 @@ tree_sitter_register_ts (Application_Links* app)
   Tree_Sitter_Language_Flags flags = (
      Tree_Sitter_Language_Can_Receive_Virtual_Indent
   );
-  tree_sitter_register_language(SCu8("ts"), language, queries, flags);
+  Tree_Sitter_Language_Definition* lang;
+  lang = tree_sitter_register_language(SCu8("ts"), language, queries, flags);
+  tree_sitter_languages_init_default_colors(app, lang);
 }
 
 #endif //TREE_SITTER_TS_H