diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd09ad1df..f451ca831 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,8 @@ All notable changes to this project are documented below.
 The format is based on [keep a changelog](http://keepachangelog.com) and this project uses [semantic versioning](http://semver.org).
 
 ## [Unreleased]
+### Fixed
+- Correctly wire Go runtime shutdown function context.
 
 ## [3.22.0] - 2024-06-09
 ### Added
diff --git a/go.mod b/go.mod
index 258e816c8..1eb2e80ae 100644
--- a/go.mod
+++ b/go.mod
@@ -13,7 +13,7 @@ require (
 	github.com/gorilla/mux v1.8.1
 	github.com/gorilla/websocket v1.5.1
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0
-	github.com/heroiclabs/nakama-common v1.32.0
+	github.com/heroiclabs/nakama-common v1.32.1-0.20240617141901-09eb4ffdf9bf
 	github.com/heroiclabs/sql-migrate v0.0.0-20240528102547-233afc8cf05a
 	github.com/jackc/pgerrcode v0.0.0-20240316143900-6e2875d9b438
 	github.com/jackc/pgx/v5 v5.6.0
diff --git a/go.sum b/go.sum
index 53f20b7c1..31c75fed3 100644
--- a/go.sum
+++ b/go.sum
@@ -130,6 +130,8 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0/go.mod h1:P+Lt/0by1T8bfcF3z737
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/heroiclabs/nakama-common v1.32.0 h1:aCWyYf9mQzifeVu3bXBiRRL9Z/dGBgwY/rgUWoYCnQM=
 github.com/heroiclabs/nakama-common v1.32.0/go.mod h1:lPG64MVCs0/tEkh311Cd6oHX9NLx2vAPx7WW7QCJHQ0=
+github.com/heroiclabs/nakama-common v1.32.1-0.20240617141901-09eb4ffdf9bf h1:PpB1MCfShscfwNG04LnloAnobajN/zlKeeQNaYJe7+I=
+github.com/heroiclabs/nakama-common v1.32.1-0.20240617141901-09eb4ffdf9bf/go.mod h1:lPG64MVCs0/tEkh311Cd6oHX9NLx2vAPx7WW7QCJHQ0=
 github.com/heroiclabs/sql-migrate v0.0.0-20240528102547-233afc8cf05a h1:tuL2ZPaeCbNw8rXmV9ywd00nXRv95V4/FmbIGKLQJAE=
 github.com/heroiclabs/sql-migrate v0.0.0-20240528102547-233afc8cf05a/go.mod h1:hzCTPoEi/oml2BllVydJcNP63S7b56e5DzeQeLGvw1U=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
diff --git a/server/api_test.go b/server/api_test.go
index a6d487f62..89c9fc5b7 100644
--- a/server/api_test.go
+++ b/server/api_test.go
@@ -55,7 +55,7 @@ var (
 	}
 	metrics       = NewLocalMetrics(logger, logger, nil, cfg)
 	storageIdx, _ = NewLocalStorageIndex(logger, nil, &StorageConfig{DisableIndexOnly: false}, metrics)
-	_             = cfg.Validate(logger)
+	_             = ValidateConfig(logger, cfg)
 )
 
 type DummyMessageRouter struct{}
diff --git a/server/runtime_go.go b/server/runtime_go.go
index 855f7f02f..1027c1aa5 100644
--- a/server/runtime_go.go
+++ b/server/runtime_go.go
@@ -2674,6 +2674,7 @@ func (ri *RuntimeGoInitializer) RegisterFleetManager(fleetManager runtime.FleetM
 
 func (ri *RuntimeGoInitializer) RegisterShutdown(fn func(ctx context.Context, logger runtime.Logger, db *sql.DB, nk runtime.NakamaModule)) error {
 	ri.shutdownFunction = func(ctx context.Context) {
+		ctx = NewRuntimeGoContext(ctx, ri.node, ri.version, ri.env, RuntimeExecutionModeShutdown, nil, nil, 0, "", "", nil, "", "", "", "")
 		fn(ctx, ri.logger.WithField("mode", RuntimeExecutionModeShutdown.String()), ri.db, ri.nk)
 	}
 
diff --git a/server/shutdown.go b/server/shutdown.go
index 05b45332b..f00291257 100644
--- a/server/shutdown.go
+++ b/server/shutdown.go
@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"github.com/heroiclabs/nakama-common/runtime"
 	"go.uber.org/zap"
 	"os"
 	"time"
@@ -14,18 +15,20 @@ func HandleShutdown(ctx context.Context, logger *zap.Logger, matchRegistry Match
 	runtimeShutdownFnDone := make(chan struct{}, 1)
 
 	if graceSeconds != 0 {
-		timer = time.NewTimer(time.Duration(graceSeconds) * time.Second)
-		timerCh = timer.C
-
+		graceDuration := time.Duration(graceSeconds) * time.Second
 		if shutdownFn != nil {
 			go func() {
-				shutdownFn(ctx)
+				shCtx, _ := context.WithTimeoutCause(context.WithoutCancel(ctx), graceDuration, runtime.ErrGracePeriodExpired)
+				shutdownFn(shCtx)
 				close(runtimeShutdownFnDone)
 			}()
 		} else {
 			close(runtimeShutdownFnDone)
 		}
 
+		timer = time.NewTimer(graceDuration)
+		timerCh = timer.C
+
 		logger.Info("Shutdown started - use CTRL^C to force stop server", zap.Int("grace_period_sec", graceSeconds))
 	} else {
 		// No grace period.
diff --git a/vendor/github.com/heroiclabs/nakama-common/runtime/runtime.go b/vendor/github.com/heroiclabs/nakama-common/runtime/runtime.go
index 411f77761..042a54927 100644
--- a/vendor/github.com/heroiclabs/nakama-common/runtime/runtime.go
+++ b/vendor/github.com/heroiclabs/nakama-common/runtime/runtime.go
@@ -197,6 +197,8 @@ var (
 	ErrPartyRemove                   = errors.New("party could not remove")
 	ErrPartyRemoveSelf               = errors.New("party cannot remove self")
 
+	ErrGracePeriodExpired = errors.New("grace period expired")
+
 	ErrGroupNameInUse         = errors.New("group name in use")
 	ErrGroupPermissionDenied  = errors.New("group permission denied")
 	ErrGroupNoUpdateOps       = errors.New("no group updates")
diff --git a/vendor/modules.txt b/vendor/modules.txt
index 3e1f2ad9c..c89d64b2b 100644
--- a/vendor/modules.txt
+++ b/vendor/modules.txt
@@ -136,7 +136,7 @@ github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-openapiv2/internal/genopena
 github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-openapiv2/options
 github.com/grpc-ecosystem/grpc-gateway/v2/runtime
 github.com/grpc-ecosystem/grpc-gateway/v2/utilities
-# github.com/heroiclabs/nakama-common v1.32.0
+# github.com/heroiclabs/nakama-common v1.32.1-0.20240617141901-09eb4ffdf9bf
 ## explicit; go 1.19
 github.com/heroiclabs/nakama-common/api
 github.com/heroiclabs/nakama-common/rtapi