diff --git a/code/custom/4coder_custom_hooks.cpp b/code/custom/4coder_custom_hooks.cpp
index ffecf40d..b8f972e5 100644
--- a/code/custom/4coder_custom_hooks.cpp
+++ b/code/custom/4coder_custom_hooks.cpp
@@ -281,20 +281,16 @@ function void custom_render_buffer(
     f32 mark_thickness = (f32)def_get_config_u64(app, vars_save_string_lit("mark_thickness"));
 
     // NOTE(allen): Token colorizing
-    Managed_Scope buffer_scope = buffer_get_managed_scope(app, buffer);
-    Buffer_Tree_Sitter_Data* tree_data = scope_attachment(app, buffer_scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
-    TSTree* tree = tree_sitter_buffer_get_tree_copy(tree_data);
-
-    Token_Array token_array = get_token_array_from_buffer(app, buffer);
+    Token_Array token_array;
     paint_text_color_fcolor(app, text_layout_id, visible_range, fcolor_id(defcolor_text_default)); // will get overridden by lang-specific token coloring below
-    if (token_array.tokens != 0 && tree)
+    if (use_tree_sitter_token_coloring)
     {
-      if (use_tree_sitter_token_coloring)
-      {
-        draw_tree_sitter_node_colors(app, text_layout_id, buffer);
-      }
-      else
-      {
+      draw_tree_sitter_node_colors(app, text_layout_id, buffer);
+    }
+    else
+    {
+      token_array = get_token_array_from_buffer(app, buffer);
+      if (token_array.tokens != 0) {
         draw_cpp_token_colors(app, text_layout_id, &token_array);
       }
     }
@@ -359,13 +355,17 @@ function void custom_render_buffer(
     // NOTE(allen): Whitespace highlight
     b64 show_whitespace = false;
     view_get_setting(app, view_id, ViewSetting_ShowWhitespace, &show_whitespace);
-    if (show_whitespace){
-        if (token_array.tokens == 0){
-            draw_whitespace_highlight(app, buffer, text_layout_id, cursor_roundness);
-        }
-        else{
-            draw_whitespace_highlight(app, text_layout_id, &token_array, cursor_roundness);
-        }
+    if (show_whitespace)
+    {
+      if (token_array.tokens == 0) token_array = get_token_array_from_buffer(app, buffer);
+      if (token_array.tokens == 0)
+      {
+          draw_whitespace_highlight(app, buffer, text_layout_id, cursor_roundness);
+      }
+      else
+      {
+          draw_whitespace_highlight(app, text_layout_id, &token_array, cursor_roundness);
+      }
     }
 
     // NOTE(allen): Cursor
diff --git a/code/custom/4coder_tree_sitter.cpp b/code/custom/4coder_tree_sitter.cpp
index a4e01b9e..0728fc77 100644
--- a/code/custom/4coder_tree_sitter.cpp
+++ b/code/custom/4coder_tree_sitter.cpp
@@ -577,6 +577,7 @@ function void
 draw_tree_sitter_node_colors(Application_Links* app, Text_Layout_ID text_layout_id, Buffer_ID buffer_id)
 {
   Tree_Sitter_Language_Definition* lang = tree_sitter_language_for_buffer(app, buffer_id);
+  if (!lang) return;
   TSQuery* query = lang->queries.ptr[Tree_Sitter_Language_Query_Highlights];
 
   Range_i64 visible_range = text_layout_get_visible_range(app, text_layout_id);