[Proposal] Provide ServerSession in onConnect & global notification handler callbacks (#780)

Proposal: Provide ServerSession in onConnect & Global Notification Handler Callbacks (#780)

Motivation & Context

Addressing issue #780 in kotlin-sdk-server.

Currently, the server lifecycle and notification listener APIs in Server.kt exhibit significant ergonomic limitations:

1. No Session Parameter in onConnect:
   Server.onConnect is declared as:

   public fun onConnect(block: () -> Unit)
   

   When a new client session connects (e.g. via Streamable HTTP, WebSocket, or Stdio), the callback is triggered with zero arguments. To perform session-scoped operations (such as querying session.listRoots() or inspecting negotiated client capabilities), developers are forced to loop through server.sessions, making it difficult and error-prone to identify and initialize the specific session that just connected.
2. Notification Handlers Do Not Propagate to Future Sessions:
   Server.setNotificationHandler(...) only iterates through currently active sessions (sessions.forEach { ... }). Any session connecting after setNotificationHandler was invoked receives no handler.
3. Notification Callbacks Lack Session Origin Context:
   The handler signature (notification: T) -> Deferred<Unit> does not provide the originating ServerSession. In a multi-session server, it is impossible for the notification handler to know which client sent the notification or to respond back to that specific client session.

This proposal outlines a clean, backwards-compatible design to provide ServerSession in connection lifecycle callbacks and global notification handlers.

---

Proposed Architecture & Component Design

1. Enhanced onConnect & onClose Callbacks (Server.kt)

In kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt:

• Dual-Signature onConnect:
  Introduce a session-aware callback while preserving the parameterless overload for backward compatibility:

  private var _onConnectSession: ((ServerSession) -> Unit) = {}
  
  /**
   * Registers a callback to be invoked when a new [ServerSession] connects.
   */
  public fun onConnect(block: (session: ServerSession) -> Unit) {
      val old = _onConnectSession
      _onConnectSession = { session ->
          old(session)
          block(session)
      }
  }
  
  /**
   * Backward-compatible overload without parameters.
   */
  public fun onConnect(block: () -> Unit) {
      val old = _onConnect
      _onConnect = {
          old()
          block()
      }
  }
  

• Session Lifecycle Dispatch in createSession:

  session.connect(transport)
  sessionRegistry.addSession(session)
  notificationService.subscribeSession(session)
  
  // Attach registered global notification handlers (see below)
  attachGlobalNotificationHandlers(session)
  
  _onConnect()
  _onConnectSession(session)
  return session
  

• Symmetric onClose Support:
  Add public fun onClose(block: (session: ServerSession) -> Unit) to allow servers to cleanly release per-session caches, state handles, and background jobs upon disconnection.

2. Global Notification Handler Template Registry

In Server.kt:

• Introduce a template registry for handlers that should apply to both existing and future sessions:

  private val globalNotificationHandlers = mutableListOf<(ServerSession) -> Unit>()
  
  /**
   * Registers a notification handler for [method] that automatically attaches
   * to all currently active sessions and all future sessions upon connection.
   */
  public fun <T : Notification> setNotificationHandler(
      method: Method,
      handler: (session: ServerSession, notification: T) -> Deferred<Unit>
  ) {
      val attacher: (ServerSession) -> Unit = { session ->
          session.setNotificationHandler<T>(method) { notification ->
              handler(session, notification)
          }
      }
      globalNotificationHandlers.add(attacher)
      sessions.values.forEach(attacher)
  }
  

• Backward-Compatible Overload:
  Retain the existing (notification: T) -> Deferred<Unit> signature, delegating to the session-aware version:

  public fun <T : Notification> setNotificationHandler(
      method: Method,
      handler: (notification: T) -> Deferred<Unit>
  ) {
      setNotificationHandler<T>(method) { _, notification -> handler(notification) }
  }
  

3. Example Developer Experience

Developers can now configure per-session reactive flows cleanly and intuitively:

val server = McpServer(serverInfo, options) {
    onConnect { session ->
        logger.info { "Client connected with session ID: ${session.sessionId}" }
        // Inspect negotiated client capabilities
        if (session.clientCapabilities?.roots?.listChanged == true) {
            session.scope.launch {
                val roots = session.listRoots()
                logger.info { "Session roots initialized: ${roots.roots}" }
            }
        }
    }

    // Handles notifications with full session origin context for all connected and future clients
    setNotificationHandler<RootsListChangedNotification>(Method.Defined.NotificationsRootsListChanged) { session, _ ->
        session.scope.async {
            val updatedRoots = session.listRoots()
            handleRootsChanged(session, updatedRoots)
        }
    }

    onClose { session ->
        logger.info { "Session ${session.sessionId} disconnected; cleaning up resources." }
    }
}

4. Testing & Verification Plan

• Unit Tests (ServerSessionTest.kt / ServerLifecycleTest.kt):
  • Test onConnect(session): Assert that the newly connected session is passed directly to the lambda and matches server.sessions[sessionId].
  • Test setNotificationHandler with future connections: Register a handler before calling server.createSession(transport), emit a notification from client, and assert the handler is triggered with the correct session instance.
  • Test multiple handlers and chained registration order.
  • Assert backward compatibility: ensure existing code calling parameterless onConnect {} compiles and behaves identically.

---

Next Steps

Upon review and maintainer consensus:

1. Implement the overloaded onConnect(session: ServerSession -> Unit) and onClose(session: ServerSession -> Unit) in Server.kt.
2. Implement global notification handler template registration.
3. Add unit tests covering multi-session lifecycle and notification origin assertions.

AI assistance disclosure: AI was used to discover this opportunity and draft the change or text. The submission was checked against the prepared artifact and recorded verification evidence.

Start in kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt, especially createSession, onConnect, and setNotificationHandler. Then review the proposed coverage in ServerSessionTest.kt and ServerLifecycleTest.kt. Done means session-aware lifecycle callbacks, handlers reaching future sessions, notification origin context, and preserved parameterless behavior are covered by tests.