diff --git a/Package.swift b/Package.swift
index 834bacc..6e64cb6 100644
--- a/Package.swift
+++ b/Package.swift
@@ -27,7 +27,7 @@ let package = Package(
     ],
     dependencies: [
         // TODO: change to upstream once the upstream is tagged
-        .package(url: "https://github.com/andrea-scuderi/swift-aws-lambda-runtime.git", branch: "main"),
+        .package(url: "https://github.com/swift-server/swift-aws-lambda-runtime.git", from: "2.0.0-beta.1"),
         .package(url: "https://github.com/swift-server/swift-aws-lambda-events.git", from: "0.5.0"),
         .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.0.0"),
         .package(url: "https://github.com/soto-project/soto.git", from: "7.0.0"),
diff --git a/Tests/BreezeLambdaAPITests/BreezeLambdaHandlerTests.swift b/Tests/BreezeLambdaAPITests/BreezeLambdaHandlerTests.swift
index a245144..07dfdb1 100644
--- a/Tests/BreezeLambdaAPITests/BreezeLambdaHandlerTests.swift
+++ b/Tests/BreezeLambdaAPITests/BreezeLambdaHandlerTests.swift
@@ -336,7 +336,7 @@ final actor MockLambdaResponseStreamWriter: LambdaResponseStreamWriter {
         self.buffer = buffer
     }
 
-    func write(_ buffer: ByteBuffer) async throws {
+    func write(_ buffer: ByteBuffer, hasCustomHeaders: Bool) async throws {
         fatalError("Unexpected call")
     }