Split model loading and ClassLoader prep

modules/lsp/src/main/scala/playground/lsp/BuildLoader.scala

@@ -91,7 +91,7 @@ object BuildLoader {
       ): F[DynamicSchemaIndex] = Sync[F]
         .interruptibleMany {
           ModelLoader
-            .loadUnsafe(
+            .load(
               specs =
                 loaded
                   .config
@@ -106,15 +106,9 @@ object BuildLoader {
                       .toFile()
                   )
                   .toSet,
-              dependencies =
-                loaded.config.mavenDependencies ++ loaded.config.maven.foldMap(_.dependencies),
-              repositories =
-                loaded
-                  .config
-                  .mavenRepositories ++ loaded.config.maven.foldMap(_.repositories).map(_.url),
+              classLoader = ModelLoader.makeClassLoaderUnsafe(loaded.config),
             )
         }
-        .map(_._2)
         .flatMap(ModelReader.buildSchemaIndex[F])
 
     }

modules/lsp/src/main/scala/playground/lsp/ModelLoader.scala

@@ -1,12 +1,13 @@
 package playground.lsp
 
-// fork of smithy4s's ModelLoader
-
+import cats.implicits._
 import coursier._
 import coursier.cache.FileCache
 import coursier.parse.DependencyParser
 import coursier.parse.RepositoryParser
 import coursier.util.Task
+import fs2.io.file.Path
+import playground.BuildConfig
 import playground.lsp.buildinfo.BuildInfo
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.loader.ModelAssembler
@@ -15,39 +16,44 @@ import software.amazon.smithy.model.loader.ModelManifestException
 
 import java.io.File
 import java.net.URLClassLoader
+import java.nio.file.Paths
 import scala.concurrent.duration._
 import scala.jdk.CollectionConverters._
 import scala.util.chaining._
 
 object ModelLoader {
 
-  def loadUnsafe(
-    specs: Set[File],
-    dependencies: List[String],
-    repositories: List[String],
-  ): (
-    ClassLoader,
-    Model,
-  ) = {
+  def makeClassLoaderUnsafe(
+    buildConfig: BuildConfig
+  ): URLClassLoader = {
+    val dependencies =
+      buildConfig.mavenDependencies ++
+        buildConfig.maven.foldMap(_.dependencies)
+
+    val repositories =
+      buildConfig.mavenRepositories ++
+        buildConfig.maven.foldMap(_.repositories).map(_.url)
 
     val dependencyJars = resolveDependencies(dependencies, repositories)
 
-    val modelAssembler: ModelAssembler = Model
-      .assembler()
-      .putProperty(ModelAssembler.DISABLE_JAR_CACHE, true)
-      .pipe(addModelsFromJars(dependencyJars))
-      .pipe(addPlaygroundModels(this.getClass().getClassLoader()))
-      .pipe(addFileImports(specs))
-
-    (
-      new URLClassLoader(
-        dependencyJars.map(_.toURI().toURL()).toArray,
-        getClass().getClassLoader(),
-      ),
-      modelAssembler.assemble().unwrap(),
+    new URLClassLoader(
+      dependencyJars.map(_.toURI().toURL()).toArray,
+      getClass().getClassLoader(),
     )
   }
 
+  def load(
+    specs: Set[File],
+    classLoader: URLClassLoader,
+  ): Model = Model
+    .assembler()
+    .putProperty(ModelAssembler.DISABLE_JAR_CACHE, true)
+    .pipe(addModelsFromJars(classLoader.getURLs().map(_.toURI()).map(Paths.get(_).toFile())))
+    .pipe(addPlaygroundModels(this.getClass().getClassLoader()))
+    .pipe(addFileImports(specs))
+    .assemble()
+    .unwrap()
+
   private def addModelsFromJars(
     jarFiles: Iterable[File]
   ): ModelAssembler => ModelAssembler =

modules/lsp/src/main/scala/playground/lsp/PluginResolver.scala

@@ -8,19 +8,8 @@ import playground.plugins.PlaygroundPlugin
 trait PluginResolver[F[_]] {
 
   def resolve(
-    artifacts: List[String],
-    repositories: List[String],
-  ): F[List[PlaygroundPlugin]]
-
-  def resolveFromConfig(
     config: BuildConfig
-  ): F[List[PlaygroundPlugin]] = resolve(
-    config.mavenDependencies ++ config.maven.foldMap(_.dependencies) ++
-      config
-        .smithyPlayground
-        .foldMap(_.extensions),
-    config.mavenRepositories ++ config.maven.foldMap(_.repositories).map(_.url),
-  )
+  ): F[List[PlaygroundPlugin]]
 
 }
 
@@ -34,18 +23,9 @@ object PluginResolver {
     new PluginResolver[F] {
 
       def resolve(
-        artifacts: List[String],
-        repositories: List[String],
+        config: BuildConfig
       ): F[List[PlaygroundPlugin]] = Sync[F]
-        .interruptibleMany(
-          ModelLoader
-            .loadUnsafe(
-              specs = Set.empty,
-              dependencies = artifacts,
-              repositories = repositories,
-            )
-            ._1
-        )
+        .interruptibleMany(ModelLoader.makeClassLoaderUnsafe(config))
         .map(PlaygroundPlugin.getAllPlugins(_))
 
     }

modules/lsp/src/main/scala/playground/lsp/ServerBuilder.scala

@@ -65,7 +65,7 @@ object ServerBuilder {
       ): F[LanguageServer[F]] =
         for {
           dsi <- BuildLoader[F].buildSchemaIndex(buildInfo)
-          plugins <- PluginResolver[F].resolveFromConfig(buildInfo.config)
+          plugins <- PluginResolver[F].resolve(buildInfo.config)
           rep <- CommandResultReporter.instance[F]
         } yield {
           val runners = OperationRunner

modules/lsp/src/test/scala/playground/lsp/ModelLoaderTests.scala

@@ -1,5 +1,6 @@
 package playground.lsp
 
+import playground.BuildConfig
 import software.amazon.smithy.model.shapes.ShapeId
 import weaver._
 
@@ -32,21 +33,26 @@ object ModelLoaderTests extends FunSuite {
   test("Loader with dependencies can see external shapes") {
     val shapeId = ShapeId.from("alloy#UUID")
     val result = ModelLoader
-      .loadUnsafe(
+      .load(
         specs = Set.empty,
-        dependencies = List("com.disneystreaming.alloy:alloy-core:0.1.13"),
-        repositories = Nil,
+        ModelLoader.makeClassLoaderUnsafe(
+          BuildConfig(
+            mavenDependencies = List("com.disneystreaming.alloy:alloy-core:0.1.13"),
+            mavenRepositories = Nil,
+          )
+        ),
       )
-      ._2
       .expectShape(shapeId)
 
     assert.same(result.getId(), shapeId)
   }
 
   private def loadModelEmpty(
-  ) =
-    ModelLoader
-      .loadUnsafe(specs = Set.empty, dependencies = Nil, repositories = Nil)
-      ._2
+  ) = ModelLoader
+    .load(
+      specs = Set.empty,
+      classLoader = ModelLoader
+        .makeClassLoaderUnsafe(BuildConfig()),
+    )
 
 }

modules/lsp/src/test/scala/playground/lsp/PluginResolverTests.scala

@@ -1,24 +1,27 @@
 package playground.lsp
 
 import cats.effect.IO
+import playground.BuildConfig
 import weaver._
 
 object PluginResolverTests extends SimpleIOSuite {
   test("Empty plugin resolver finds no plugins") {
     PluginResolver
       .instance[IO]
-      .resolve(artifacts = Nil, repositories = Nil)
+      .resolve(BuildConfig(mavenDependencies = Nil, mavenRepositories = Nil))
       .map(assert.same(_, Nil))
   }
 
   test("Plugin resolver with a sample plugin artifact finds it") {
     PluginResolver
       .instance[IO]
       .resolve(
-        artifacts = List(
-          "com.kubukoz.playground::plugin-sample:latest.integration"
-        ),
-        repositories = Nil,
+        BuildConfig(
+          mavenDependencies = List(
+            "com.kubukoz.playground::plugin-sample:latest.integration"
+          ),
+          mavenRepositories = Nil,
+        )
       )
       .map(_.map(_.getClass().getName()))
       .map(assert.same(_, List("playground.sample.SamplePlaygroundPlugin")))