diff --git a/spec/static_file_handler_spec.cr b/spec/static_file_handler_spec.cr
index 309e10f..5f8c029 100644
--- a/spec/static_file_handler_spec.cr
+++ b/spec/static_file_handler_spec.cr
@@ -141,11 +141,11 @@ describe Kemal::StaticFileHandler do
   end
 
   it "should handle setting custom headers" do
-    headers = Proc(HTTP::Server::Response, String, File::Info, Void).new do |response, path, stat|
+    headers = Proc(HTTP::Server::Context, String, File::Info, Void).new do |env, path, stat|
       if path =~ /\.html$/
-        response.headers.add("Access-Control-Allow-Origin", "*")
+        env.response.headers.add("Access-Control-Allow-Origin", "*")
       end
-      response.headers.add("Content-Size", stat.size.to_s)
+      env.response.headers.add("Content-Size", stat.size.to_s)
     end
 
     static_headers(&headers)
diff --git a/src/kemal/config.cr b/src/kemal/config.cr
index df7330a..f728f57 100644
--- a/src/kemal/config.cr
+++ b/src/kemal/config.cr
@@ -23,7 +23,7 @@ module Kemal
     property app_name, host_binding, ssl, port, env, public_folder, logging, running
     property always_rescue, server : HTTP::Server?, extra_options, shutdown_message
     property serve_static : (Bool | Hash(String, Bool))
-    property static_headers : (HTTP::Server::Response, String, File::Info -> Void)?
+    property static_headers : (HTTP::Server::Context, String, File::Info -> Void)?
     property? powered_by_header : Bool = true
 
     def initialize
diff --git a/src/kemal/helpers/helpers.cr b/src/kemal/helpers/helpers.cr
index 1b733c3..d8ef72b 100644
--- a/src/kemal/helpers/helpers.cr
+++ b/src/kemal/helpers/helpers.cr
@@ -134,7 +134,7 @@ def send_file(env : HTTP::Server::Context, path : String, mime_type : String? =
   filestat = File.info(file_path)
   attachment(env, filename, disposition)
 
-  Kemal.config.static_headers.try(&.call(env.response, file_path, filestat))
+  Kemal.config.static_headers.try(&.call(env, file_path, filestat))
 
   File.open(file_path) do |file|
     if env.request.method == "GET" && env.request.headers.has_key?("Range")
@@ -250,13 +250,13 @@ end
 # Adds headers to `Kemal::StaticFileHandler`. This is especially useful for `CORS`.
 #
 # ```
-# static_headers do |response, filepath, filestat|
+# static_headers do |env, filepath, filestat|
 #   if filepath =~ /\.html$/
-#     response.headers.add("Access-Control-Allow-Origin", "*")
+#     env.response.headers.add("Access-Control-Allow-Origin", "*")
 #   end
-#   response.headers.add("Content-Size", filestat.size.to_s)
+#   env.response.headers.add("Content-Size", filestat.size.to_s)
 # end
 # ```
-def static_headers(&headers : HTTP::Server::Response, String, File::Info -> Void)
+def static_headers(&headers : HTTP::Server::Context, String, File::Info -> Void)
   Kemal.config.static_headers = headers
 end