Add server initialization hook (#19)

This commit is contained in:
Mattt 2025-03-18 09:47:54 -07:00 • committed by GitHub
commit 2510dda70e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -155,9 +155,15 @@ public actor Server {
}
/// Start the server
public func start(transport: any Transport) async throws {
/// - Parameters:
/// - transport: The transport to use for the server
/// - initializeHook: An optional hook that runs when the client sends an initialize request
public func start(
transport: any Transport,
initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)? = nil
) async throws {
self.connection = transport
registerDefaultHandlers()
registerDefaultHandlers(initializeHook: initializeHook)
try await transport.connect()
await logger?.info(
@ -364,7 +370,9 @@ public actor Server {
}
}
private func registerDefaultHandlers() {
private func registerDefaultHandlers(
initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)?
) {
// Initialize
withMethodHandler(Initialize.self) { [weak self] params in
guard let self = self else {
@ -381,26 +389,30 @@ public actor Server {
"Unsupported protocol version: \(params.protocolVersion)")
}
// Call initialization hook if registered
if let hook = initializeHook {
try await hook(params.clientInfo, params.capabilities)
}
// Set initial state
await self.setInitialState(
clientInfo: params.clientInfo,
clientCapabilities: params.capabilities,
protocolVersion: params.protocolVersion
)
let result = Initialize.Result(
protocolVersion: Version.latest,
capabilities: await self.capabilities,
serverInfo: self.serverInfo,
instructions: nil
)
// Send initialized notification after a short delay
Task {
try? await Task.sleep(nanoseconds: 100_000_000) // 100ms
try? await self.notify(InitializedNotification.message())
}
return result
return Initialize.Result(
protocolVersion: Version.latest,
capabilities: await self.capabilities,
serverInfo: self.serverInfo,
instructions: nil
)
}
// Ping