diff --git a/src/main/scala/com/tjclp/fastmcp/server/FastMcpServer.scala b/src/main/scala/com/tjclp/fastmcp/server/FastMcpServer.scala index 14c6a77..d4d0e7d 100644 --- a/src/main/scala/com/tjclp/fastmcp/server/FastMcpServer.scala +++ b/src/main/scala/com/tjclp/fastmcp/server/FastMcpServer.scala @@ -15,8 +15,8 @@ import scala.concurrent.ExecutionContext.Implicits.global import scala.jdk.CollectionConverters.* import scala.util.Failure import scala.util.Success - import core.* +import io.modelcontextprotocol.server.transport.StdioServerTransportProvider import server.manager.* // Needed for runToFuture onComplete /** Main server class for FastMCP-Scala @@ -157,13 +157,21 @@ class FastMcpServer( /** Run the server with stdio transport */ def runStdio(): ZIO[Any, Throwable, Unit] = - ZIO.attemptBlocking { - val stdioTransportProvider = - new io.modelcontextprotocol.server.transport.StdioServerTransportProvider() - this.setupServer(stdioTransportProvider) - JSystem.err.println(s"FastMCPScala server '${this.name}' running with stdio transport.") - Thread.sleep(Long.MaxValue) - }.unit + ZIO.scoped { // ⬅ drops the `Scope` requirement + ZIO.acquireRelease( + for { + provider <- ZIO.attempt(new StdioServerTransportProvider()) + _ <- ZIO.attempt(setupServer(provider)) + _ <- ZIO.attempt( + JSystem.err.println( + s"[FastMCPScala] '$name' running on stdio – press Ctrl-C to stop." + ) + ) + } yield () + )(_ => ZIO.attempt(underlyingJavaServer.foreach(_.close())).orDie) *> ZIO.never.as( + () + ) // ⬅ turn `Nothing` into `Unit` + } /** Set up the Java MCP Server with the given transport provider */ @@ -287,7 +295,7 @@ class FastMcpServer( // --- Build Server --- // Build the McpAsyncServer underlyingJavaServer = Some(serverBuilder.build()) - JSystem.err.println(s"MCP Server '$name' configured.") + JSystem.err.println(s"[FastMCPScala] MCP Server '$name' configured.") /** Creates a Java BiFunction handler for a specific static resource URI. Returns a Mono that * completes with the result. diff --git a/src/test/scala/com/tjclp/fastmcp/server/FastMcpServerShutdownSpec.scala b/src/test/scala/com/tjclp/fastmcp/server/FastMcpServerShutdownSpec.scala new file mode 100644 index 0000000..35cf728 --- /dev/null +++ b/src/test/scala/com/tjclp/fastmcp/server/FastMcpServerShutdownSpec.scala @@ -0,0 +1,36 @@ +package com.tjclp.fastmcp.server + +import java.util.concurrent.atomic.AtomicBoolean +import zio.* +import zio.test.* + +object FastMcpServerShutdownSpec extends ZIOSpecDefault { + + private final class MockCloseableServer(flag: AtomicBoolean) extends AutoCloseable { + override def close(): Unit = flag.set(true) + } + + def spec = + suite("FastMcpServer graceful-shutdown")( + test("runStdio closes the underlying server when interrupted") { + for { + closedFlag <- ZIO.succeed(new AtomicBoolean(false)) + + server <- ZIO.succeed( + new FastMcpServer() { + override def runStdio() = + ZIO.scoped { + ZIO.acquireRelease( + ZIO.succeed(new MockCloseableServer(closedFlag)) + )(srv => ZIO.succeed(srv.close())) *> ZIO.never + } + } + ) + + fiber <- server.runStdio().fork + _ <- TestClock.adjust(10.millis) + _ <- fiber.interrupt + } yield assertTrue(closedFlag.get()) + } + ) +}