diff --git a/lua/lazyvim/plugins/extras/ui/treesitter-context.lua b/lua/lazyvim/plugins/extras/ui/treesitter-context.lua
index a1f8d5f3..019977ed 100644
--- a/lua/lazyvim/plugins/extras/ui/treesitter-context.lua
+++ b/lua/lazyvim/plugins/extras/ui/treesitter-context.lua
@@ -1,21 +1,22 @@
 -- Show context of the current function
 return {
   "nvim-treesitter/nvim-treesitter-context",
-  event = "LazyFile",
-  opts = { mode = "cursor", max_lines = 3 },
-  keys = {
-    {
-      "<leader>ut",
-      function()
-        local tsc = require("treesitter-context")
-        tsc.toggle()
-        if LazyVim.inject.get_upvalue(tsc.toggle, "enabled") then
-          LazyVim.info("Enabled Treesitter Context", { title = "Option" })
+  event = "VeryLazy",
+  opts = function()
+    local tsc = require("treesitter-context")
+
+    LazyVim.toggle.map("<leader>ut", {
+      name = "Treesitter Context",
+      get = tsc.enabled,
+      set = function(state)
+        if state then
+          tsc.enable()
         else
-          LazyVim.warn("Disabled Treesitter Context", { title = "Option" })
+          tsc.disable()
         end
       end,
-      desc = "Toggle Treesitter Context",
-    },
-  },
+    })
+
+    return { mode = "cursor", max_lines = 3 }
+  end,
 }