diff --git a/internal/cmd/commands/dev/dev.go b/internal/cmd/commands/dev/dev.go
index 9f865a8b44..b52b3d2a14 100644
--- a/internal/cmd/commands/dev/dev.go
+++ b/internal/cmd/commands/dev/dev.go
@@ -48,7 +48,7 @@ var (
 	_ cli.CommandAutocomplete = (*Command)(nil)
 )
 
-var extraSelfTerminationConditionFuncs []func(context.Context, chan struct{})
+var extraSelfTerminationConditionFuncs []func(*Command, chan struct{})
 
 type Command struct {
 	*base.Server
@@ -924,7 +924,7 @@ func (c *Command) Run(args []string) int {
 	}
 
 	for _, f := range extraSelfTerminationConditionFuncs {
-		f(c.Context, c.ServerSideShutdownCh)
+		f(c, c.ServerSideShutdownCh)
 	}
 
 	for !errorEncountered.Load() && !shutdownCompleted.Load() {
diff --git a/internal/cmd/commands/server/server.go b/internal/cmd/commands/server/server.go
index 25eef140f4..f1f6947d4a 100644
--- a/internal/cmd/commands/server/server.go
+++ b/internal/cmd/commands/server/server.go
@@ -41,7 +41,7 @@ var (
 	_ cli.CommandAutocomplete = (*Command)(nil)
 )
 
-var extraSelfTerminationConditionFuncs []func(context.Context, chan struct{})
+var extraSelfTerminationConditionFuncs []func(*Command, chan struct{})
 
 type Command struct {
 	*base.Server
@@ -700,7 +700,7 @@ func (c *Command) WaitForInterrupt() int {
 	}
 
 	for _, f := range extraSelfTerminationConditionFuncs {
-		f(c.Context, c.ServerSideShutdownCh)
+		f(c, c.ServerSideShutdownCh)
 	}
 
 	for !shutdownCompleted.Load() {