Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 17 additions & 9 deletions src/main/scala/com/tjclp/fastmcp/server/FastMcpServer.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
*/
Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
@@ -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())
}
)
}