diff --git a/runtime/v2/manager.go b/runtime/v2/manager.go
index dbd68260bd05..a773f52a1dfc 100644
--- a/runtime/v2/manager.go
+++ b/runtime/v2/manager.go
@@ -667,19 +667,19 @@ func registerMethod(
 
 	return string(requestName), stfRouter.RegisterHandler(string(requestName), func(
 		ctx context.Context,
-		msg appmodulev2.Message,
-	) (resp appmodulev2.Message, err error) {
+		msg transaction.Msg,
+	) (resp transaction.Msg, err error) {
 		res, err := md.Handler(ss, ctx, noopDecoder, messagePassingInterceptor(msg))
 		if err != nil {
 			return nil, err
 		}
-		return res.(appmodulev2.Message), nil
+		return res.(transaction.Msg), nil
 	})
 }
 
 func noopDecoder(_ interface{}) error { return nil }
 
-func messagePassingInterceptor(msg appmodulev2.Message) grpc.UnaryServerInterceptor {
+func messagePassingInterceptor(msg transaction.Msg) grpc.UnaryServerInterceptor {
 	return func(
 		ctx context.Context,
 		req interface{},
diff --git a/server/v2/stf/stf_router.go b/server/v2/stf/stf_router.go
index e4eaa00be2d0..41d5b805c612 100644
--- a/server/v2/stf/stf_router.go
+++ b/server/v2/stf/stf_router.go
@@ -160,7 +160,7 @@ func (r coreRouterImpl) CanInvoke(_ context.Context, typeURL string) error {
 	return nil
 }
 
-func (r coreRouterImpl) InvokeTyped(ctx context.Context, req, resp gogoproto.Message) error {
+func (r coreRouterImpl) InvokeTyped(ctx context.Context, req, resp transaction.Msg) error {
 	handlerResp, err := r.InvokeUntyped(ctx, req)
 	if err != nil {
 		return err
@@ -168,7 +168,7 @@ func (r coreRouterImpl) InvokeTyped(ctx context.Context, req, resp gogoproto.Mes
 	return merge(handlerResp, resp)
 }
 
-func (r coreRouterImpl) InvokeUntyped(ctx context.Context, req gogoproto.Message) (res gogoproto.Message, err error) {
+func (r coreRouterImpl) InvokeUntyped(ctx context.Context, req transaction.Msg) (res transaction.Msg, err error) {
 	typeName := msgTypeURL(req)
 	handler, exists := r.handlers[typeName]
 	if !exists {