diff --git a/samples/kotlin-mcp-server/src/main/kotlin/Main.kt b/samples/kotlin-mcp-server/src/main/kotlin/Main.kt
index 9808357..a73b5b7 100644
--- a/samples/kotlin-mcp-server/src/main/kotlin/Main.kt
+++ b/samples/kotlin-mcp-server/src/main/kotlin/Main.kt
@@ -14,12 +14,11 @@ import io.modelcontextprotocol.kotlin.sdk.PromptMessage
 import io.modelcontextprotocol.kotlin.sdk.Role
 import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
 import io.modelcontextprotocol.kotlin.sdk.Tool
-import io.modelcontextprotocol.kotlin.sdk.server.MCP
-import io.modelcontextprotocol.kotlin.sdk.server.SSEServerTransport
 import io.modelcontextprotocol.kotlin.sdk.server.Server
 import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
+import io.modelcontextprotocol.kotlin.sdk.server.SseServerTransport
 import io.modelcontextprotocol.kotlin.sdk.server.StdioServerTransport
-import kotlinx.coroutines.CompletableDeferred
+import io.modelcontextprotocol.kotlin.sdk.server.mcp
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.runBlocking
 import kotlinx.io.asSink
@@ -48,8 +47,6 @@ fun main(args: Array<String>) {
 }
 
 fun configureServer(): Server {
-    val def = CompletableDeferred<Unit>()
-
     val server = Server(
         Implementation(
             name = "mcp-kotlin test server",
@@ -61,10 +58,7 @@ fun configureServer(): Server {
                 resources = ServerCapabilities.Resources(subscribe = true, listChanged = true),
                 tools = ServerCapabilities.Tools(listChanged = true),
             )
-        ),
-        onCloseCallback = {
-            def.complete(Unit)
-        }
+        )
     )
 
     server.addPrompt(
@@ -129,7 +123,7 @@ fun runMcpServerUsingStdio() {
     runBlocking {
         server.connect(transport)
         val done = Job()
-        server.onCloseCallback = {
+        server.onClose {
             done.complete()
         }
         done.join()
@@ -146,7 +140,7 @@ fun runSseMcpServerWithPlainConfiguration(port: Int): Unit = runBlocking {
         install(SSE)
         routing {
             sse("/sse") {
-                val transport = SSEServerTransport("/message", this)
+                val transport = SseServerTransport("/message", this)
                 val server = configureServer()
 
                 // For SSE, you can also add prompts/tools/resources if needed:
@@ -154,7 +148,7 @@ fun runSseMcpServerWithPlainConfiguration(port: Int): Unit = runBlocking {
 
                 servers[transport.sessionId] = server
 
-                server.onCloseCallback = {
+                server.onClose {
                     println("Server closed")
                     servers.remove(transport.sessionId)
                 }
@@ -164,7 +158,7 @@ fun runSseMcpServerWithPlainConfiguration(port: Int): Unit = runBlocking {
             post("/message") {
                 println("Received Message")
                 val sessionId: String = call.request.queryParameters["sessionId"]!!
-                val transport = servers[sessionId]?.transport as? SSEServerTransport
+                val transport = servers[sessionId]?.transport as? SseServerTransport
                 if (transport == null) {
                     call.respond(HttpStatusCode.NotFound, "Session not found")
                     return@post
@@ -189,8 +183,8 @@ fun runSseMcpServerUsingKtorPlugin(port: Int): Unit = runBlocking {
     println("Use inspector to connect to the http://localhost:$port/sse")
 
     embeddedServer(CIO, host = "0.0.0.0", port = port) {
-        MCP {
-            return@MCP configureServer()
+        mcp {
+            return@mcp configureServer()
         }
     }.start(wait = true)
 }
\ No newline at end of file
diff --git a/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt b/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt
index e3f56f3..7fadede 100644
--- a/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt
+++ b/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt
@@ -121,7 +121,7 @@ fun `run mcp server`() {
     runBlocking {
         server.connect(transport)
         val done = Job()
-        server.onCloseCallback = {
+        server.onClose {
             done.complete()
         }
         done.join()
diff --git a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt
index 8724cf0..e24419e 100644
--- a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt
+++ b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt
@@ -30,7 +30,6 @@ public class ServerOptions(
  *
  * @param serverInfo Information about this server implementation (name, version).
  * @param options Configuration options for the server.
- * @param onCloseCallback A callback invoked when the server connection closes.
  */
 public open class Server(
     private val serverInfo: Implementation,
@@ -127,7 +126,6 @@ public open class Server(
 
     /**
      * Called when the server connection is closing.
-     * Invokes [onCloseCallback] if set.
      */
     override fun onClose() {
         logger.info { "Server connection closing" }
diff --git a/src/jvmTest/kotlin/InMemoryTransport.kt b/src/jvmTest/kotlin/InMemoryTransport.kt
index 987d166..ed849e2 100644
--- a/src/jvmTest/kotlin/InMemoryTransport.kt
+++ b/src/jvmTest/kotlin/InMemoryTransport.kt
@@ -1,6 +1,5 @@
 import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
 import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
-import io.modelcontextprotocol.kotlin.sdk.shared.Transport
 
 /**
  * In-memory transport for creating clients and servers that talk to each other within the same process.
diff --git a/src/jvmTest/kotlin/client/InMemoryTransportTest.kt b/src/jvmTest/kotlin/client/InMemoryTransportTest.kt
index 90acd1e..bfe01f6 100644
--- a/src/jvmTest/kotlin/client/InMemoryTransportTest.kt
+++ b/src/jvmTest/kotlin/client/InMemoryTransportTest.kt
@@ -96,11 +96,9 @@ class InMemoryTransportTest {
             clientTransport.close()
 
             assertThrows<IllegalStateException> {
-                runBlocking {
-                    clientTransport.send(
-                        InitializedNotification().toJSON()
-                    )
-                }
+                clientTransport.send(
+                    InitializedNotification().toJSON()
+                )
             }
         }
     }