diff --git a/flake.lock b/flake.lock
index 4c237c8..15262b0 100644
--- a/flake.lock
+++ b/flake.lock
@@ -1063,15 +1063,15 @@
     "nvim-treesitter-context": {
       "flake": false,
       "locked": {
-        "lastModified": 1652175020,
-        "narHash": "sha256-dYw/Y6+Eni7gTSjvCCl7E6ho8xSTxDYAwH3xGEhf54Q=",
-        "owner": "lewis6991",
+        "lastModified": 1681297663,
+        "narHash": "sha256-INYJBQDHgq+2V469JLIq54LolGb6T7nlgmmVDQcxwkk=",
+        "owner": "nvim-treesitter",
         "repo": "nvim-treesitter-context",
-        "rev": "c931a3136a696c1827eda78ac678aea542115bd0",
+        "rev": "0d730df898f3dc27fd88f03cfa6d26d2405554b4",
         "type": "github"
       },
       "original": {
-        "owner": "lewis6991",
+        "owner": "nvim-treesitter",
         "repo": "nvim-treesitter-context",
         "type": "github"
       }
diff --git a/flake.nix b/flake.nix
index 65a80c3..7678e24 100644
--- a/flake.nix
+++ b/flake.nix
@@ -87,7 +87,7 @@
       flake = false;
     };
     nvim-treesitter-context = {
-      url = "github:lewis6991/nvim-treesitter-context";
+      url = "github:nvim-treesitter/nvim-treesitter-context";
       flake = false;
     };
     nvim-lightbulb = {
diff --git a/modules/treesitter/context.nix b/modules/treesitter/context.nix
index 30472ae..08039bb 100644
--- a/modules/treesitter/context.nix
+++ b/modules/treesitter/context.nix
@@ -1,29 +1,84 @@
 {
+  pkgs,
   config,
   lib,
   ...
 }:
 with lib;
 with builtins; let
-  cfg = config.vim.treesitter;
+  treesitter = config.vim.treesitter;
+  cfg = treesitter.context;
 in {
-  options.vim.treesitter.context.enable = mkOption {
-    type = types.bool;
-    default = false;
-    description = "enable function context [nvim-treesitter-context]";
+  options.vim.treesitter.context = {
+    enable = mkEnableOption "context of current buffer contents [nvim-treesitter-context] ";
+
+    maxLines = mkOption {
+      description = "How many lines the window should span. Values <=0 mean no limit.";
+      type = types.int;
+      default = 0;
+    };
+
+    minWindowHeight = mkOption {
+      description = "Minimum editor window height to enable context. Values <= 0 mean no limit.";
+      type = types.int;
+      default = 0;
+    };
+
+    lineNumbers = mkOption {
+      description = "";
+      type = types.bool;
+      default = true;
+    };
+
+    multilineThreshold = mkOption {
+      description = "Maximum number of lines to collapse for a single context line.";
+      type = types.int;
+      default = 20;
+    };
+
+    trimScope = mkOption {
+      description = nvim.nmd.asciiDoc "Which context lines to discard if <<opt-vim.treesitter.context.maxLines>> is exceeded.";
+      type = types.enum ["inner" "outer"];
+      default = "outer";
+    };
+
+    mode = mkOption {
+      description = "Line used to calculate context.";
+      type = types.enum ["cursor" "topline"];
+      default = "cursor";
+    };
+
+    separator = mkOption {
+      description = nvim.nmd.asciiDoc ''
+        Separator between context and content. Should be a single character string, like '-'.
+
+        When separator is set, the context will only show up when there are at least 2 lines above cursorline.
+      '';
+      type = with types; nullOr str;
+      default = null;
+    };
+
+    zindex = mkOption {
+      description = "The Z-index of the context window.";
+      type = types.int;
+      default = 20;
+    };
   };
 
-  config = mkIf (cfg.enable && cfg.context.enable) {
-    vim.startPlugins = [
-      "nvim-treesitter-context"
-    ];
+  config = mkIf (treesitter.enable && cfg.enable) {
+    vim.startPlugins = ["nvim-treesitter-context"];
 
     vim.luaConfigRC.treesitter-context = nvim.dag.entryAnywhere ''
-      -- Treesitter Context config
       require'treesitter-context'.setup {
         enable = true,
-        throttle = true,
-        max_lines = 0
+        max_lines = ${toString cfg.maxLines},
+        min_window_height = ${toString cfg.minWindowHeight},
+        line_numbers = ${boolToString cfg.lineNumbers},
+        multiline_threshold = ${toString cfg.multilineThreshold},
+        trim_scope = '${cfg.trimScope}',
+        mode = '${cfg.mode}',
+        separator = ${nvim.lua.nullString cfg.separator},
+        max_lines = ${toString cfg.zindex},
       }
     '';
   };