effect 4.0.0-beta.102 → 4.0.0-beta.103

package/src/unstable/ai/McpServer.ts

@@ -13,11 +13,14 @@
 import * as Arr from "../../Array.ts"
 import * as Cause from "../../Cause.ts"
 import * as Context from "../../Context.ts"
+import * as Data from "../../Data.ts"
 import * as Effect from "../../Effect.ts"
 import * as Exit from "../../Exit.ts"
 import * as Fiber from "../../Fiber.ts"
 import * as Layer from "../../Layer.ts"
+import * as LogLevel from "../../LogLevel.ts"
 import * as Option from "../../Option.ts"
+import * as Predicate from "../../Predicate.ts"
 import * as Queue from "../../Queue.ts"
 import * as RcMap from "../../RcMap.ts"
 import { CurrentLogLevel } from "../../References.ts"
@@ -40,29 +43,36 @@ import * as RpcMessage from "../rpc/RpcMessage.ts"
 import * as RpcSerialization from "../rpc/RpcSerialization.ts"
 import * as RpcServer from "../rpc/RpcServer.ts"
 import * as AiError from "./AiError.ts"
+import * as McpProtocolRegistry from "./internal/mcpProtocolRegistry.ts"
+import type * as McpProtocol from "./McpProtocol.ts"
 import {
   CallToolResult,
-  ClientNotificationRpcs,
+  CancelledNotification,
   ClientRpcs,
-  CompleteResult,
   Elicit,
   ElicitationDeclined,
   EnabledWhen,
   GetPromptResult,
   InternalError,
+  INVALID_REQUEST_ERROR_CODE,
   InvalidParams,
+  InvalidRequest,
   isParam,
   ListPromptsResult,
   ListResourcesResult,
   ListResourceTemplatesResult,
   ListToolsResult,
+  LoggingMessageNotification,
+  McpErrorBase,
   McpServerClient,
   McpServerClientMiddleware,
+  MethodNotFound,
+  ParseError,
   Prompt,
   Resource,
   ResourceTemplate,
+  ResourceUpdatedNotification,
   ServerNotificationRpcs,
-  ServerRequestRpcs,
   TextContent,
   Tool as McpTool
 } from "./McpSchema.ts"
@@ -70,8 +80,10 @@ import type {
   CallTool,
   ClientCapabilities,
   Complete,
+  CompleteResult,
   GetPrompt,
   Initialize,
+  LoggingLevel,
   Param,
   PromptArgument,
   PromptMessage,
@@ -81,6 +93,8 @@ import type {
 import * as Tool from "./Tool.ts"
 import type * as Toolkit from "./Toolkit.ts"
 
+type CompletionContext = typeof Complete.payloadSchema.Type["context"]
+
 /**
  * Service that stores and serves an MCP server's registered tools, resources,
  * prompts, completions, and outgoing notifications.
@@ -90,7 +104,7 @@ import type * as Toolkit from "./Toolkit.ts"
  * Handlers use this service to register capabilities and resolve incoming MCP
  * requests.
  *
- * @category server
+ * @category services
  * @since 4.0.0
  */
 export class McpServer extends Context.Service<McpServer, {
@@ -105,7 +119,7 @@ export class McpServer extends Context.Service<McpServer, {
   readonly addTool: (options: {
     readonly tool: McpTool
     readonly annotations: Context.Context<never>
-    readonly handle: (payload: any) => Effect.Effect<CallToolResult, never, McpServerClient>
+    readonly handle: (payload: any) => Effect.Effect<CallToolResult, InternalError | InvalidParams, McpServerClient>
   }) => Effect.Effect<void>
   readonly callTool: (
     requests: typeof CallTool.payloadSchema.Type
@@ -130,17 +144,27 @@ export class McpServer extends Context.Service<McpServer, {
       readonly template: ResourceTemplate
       readonly annotations: Context.Context<never>
       readonly routerPath: string
-      readonly completions: Record<string, (input: string) => Effect.Effect<CompleteResult, InternalError>>
+      readonly completions: Record<
+        string,
+        (
+          input: string,
+          context: CompletionContext
+        ) => Effect.Effect<CompleteResult, InternalError>
+      >
       readonly handle: (
         uri: string,
         params: Array<string>
-      ) => Effect.Effect<typeof ReadResourceResult.Type, InvalidParams | InternalError, McpServerClient>
+      ) => Effect.Effect<
+        typeof ReadResourceResult.Type,
+        InvalidParams | InternalError,
+        McpServerClient
+      >
     }
   ) => Effect.Effect<void>
 
   readonly findResource: (
     uri: string
-  ) => Effect.Effect<typeof ReadResourceResult.Type, InvalidParams | InternalError, McpServerClient>
+  ) => Effect.Effect<typeof ReadResourceResult.Type, McpErrorBase | InvalidParams | InternalError, McpServerClient>
 
   readonly prompts: ReadonlyArray<{
     readonly prompt: Prompt
@@ -151,7 +175,10 @@ export class McpServer extends Context.Service<McpServer, {
     readonly annotations: Context.Context<never>
     readonly completions: Record<
       string,
-      (input: string) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
+      (
+        input: string,
+        context: CompletionContext
+      ) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
     >
     readonly handle: (
       params: Record<string, string>
@@ -163,7 +190,7 @@ export class McpServer extends Context.Service<McpServer, {
 
   readonly completion: (
     complete: typeof Complete.payloadSchema.Type
-  ) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
+  ) => Effect.Effect<CompleteResult, InvalidParams | InternalError, McpServerClient>
 }>()("effect/ai/McpServer") {
   /**
    * Builds an MCP server service from registered tools, prompts, resources, and completions.
@@ -177,7 +204,11 @@ export class McpServer extends Context.Service<McpServer, {
         readonly handle: (
           uri: string,
           params: Array<string>
-        ) => Effect.Effect<typeof ReadResourceResult.Type, InternalError | InvalidParams, McpServerClient>
+        ) => Effect.Effect<
+          typeof ReadResourceResult.Type,
+          InternalError | InvalidParams,
+          McpServerClient
+        >
       } | {
         readonly _tag: "Resource"
         readonly effect: Effect.Effect<typeof ReadResourceResult.Type, InternalError, McpServerClient>
@@ -187,7 +218,10 @@ export class McpServer extends Context.Service<McpServer, {
       readonly tool: McpTool
       readonly annotations: Context.Context<never>
     }>()
-    const toolMap = new Map<string, (payload: any) => Effect.Effect<CallToolResult, InternalError, McpServerClient>>()
+    const toolMap = new Map<
+      string,
+      (payload: any) => Effect.Effect<CallToolResult, InternalError | InvalidParams, McpServerClient>
+    >()
     const resources: Array<{
       readonly resource: Resource
       readonly annotations: Context.Context<never>
@@ -206,7 +240,10 @@ export class McpServer extends Context.Service<McpServer, {
     >()
     const completionsMap = new Map<
       string,
-      (input: string) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
+      (
+        input: string,
+        context: CompletionContext
+      ) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
     >()
     const notificationsQueue = yield* Queue.make<RpcMessage.Request<any>>()
     const listChangedHandles = new Map<string, any>()
@@ -259,7 +296,9 @@ export class McpServer extends Context.Service<McpServer, {
           if (!handle) {
             return Effect.fail(new InvalidParams({ message: `Tool '${request.name}' not found` }))
           }
-          return handle(request.arguments)
+          return handle(request.arguments).pipe(
+            Effect.catchDefect(() => Effect.fail(new InternalError({ message: "Internal error" })))
+          )
         }),
       get resources() {
         return resources
@@ -286,7 +325,7 @@ export class McpServer extends Context.Service<McpServer, {
         Effect.suspend(() => {
           const match = matcher.find(uri)
           if (!match) {
-            return Effect.succeed({ contents: [] })
+            return Effect.fail(new McpErrorBase({ code: -32002, message: `Resource '${uri}' not found` }))
           } else if (match.handler._tag === "Resource") {
             return match.handler.effect
           }
@@ -321,7 +360,19 @@ export class McpServer extends Context.Service<McpServer, {
           ? `ref/resource/${ref.uri}/${complete.argument.name}`
           : `ref/prompt/${ref.name}/${complete.argument.name}`
         const handler = completionsMap.get(key)
-        return handler ? yield* handler(complete.argument.value) : CompleteResult.empty
+        if (!handler) {
+          return yield* new InvalidParams({ message: "Unknown completion reference or argument" })
+        }
+        const result = yield* handler(complete.argument.value, complete.context)
+        const values = Arr.take(result.completion.values, 100)
+        return {
+          completion: {
+            ...result.completion,
+            values,
+            hasMore: result.completion.hasMore === true ||
+              values.length < result.completion.values.length
+          }
+        }
       })
     })
   })
@@ -334,29 +385,63 @@ export class McpServer extends Context.Service<McpServer, {
   static readonly layer: Layer.Layer<McpServer | McpServerClient> = Layer.effect(McpServer)(McpServer.make) as any
 }
 
-const LATEST_PROTOCOL_VERSION = "2025-06-18"
-const SUPPORTED_PROTOCOL_VERSIONS = [
-  LATEST_PROTOCOL_VERSION,
-  "2025-03-26",
-  "2024-11-05",
-  "2024-10-07"
-]
 const MCP_SESSION_ID_HEADER = "mcp-session-id"
 const MCP_PROTOCOL_VERSION_HEADER = "mcp-protocol-version"
+const MCP_INVALID_BATCH_METHOD = "invalid/json-rpc-batch"
+const decodeCancelledNotification = Schema.decodeUnknownEffect(
+  Schema.toCodecJson(CancelledNotification.payloadSchema)
+)
+
+const requestKey = (requestId: string | number): string => `${typeof requestId}:${requestId}`
+
+type SessionLogLevel =
+  | {
+    readonly _tag: "Effect"
+    readonly level: LogLevel.LogLevel
+  }
+  | {
+    readonly _tag: "Mcp"
+    readonly level: LoggingLevel
+  }
+
+interface Session {
+  readonly initializePayload: typeof Initialize.payloadSchema.Type
+  readonly protocol: McpProtocol.ProtocolAdapter
+  readonly resourceSubscriptions: Set<string> | undefined
+  logLevel: SessionLogLevel
+}
+
+interface Sessions {
+  readonly bySessionId: Map<string, Session>
+  readonly byClientId: Map<number, Session>
+}
+
+class McpClientKey extends Data.Class<{
+  readonly clientId: number
+  readonly protocolVersion: string
+}> {}
 
 class McpProtocolState extends Context.Service<McpProtocolState, {
-  readonly clientSessions: Map<string, typeof Initialize.payloadSchema.Type>
-  readonly latestProtocolVersion: string
-  readonly supportedProtocolVersions: ReadonlySet<string>
+  readonly sessions: Sessions
+  readonly protocolRegistry: McpProtocolRegistry.ProtocolRegistry<McpProtocol.ProtocolAdapter>
 }>()("effect/ai/McpServer/McpProtocolState") {}
 
-const makeMcpProtocolState = (): McpProtocolState["Service"] => ({
-  clientSessions: new Map(),
-  latestProtocolVersion: LATEST_PROTOCOL_VERSION,
-  supportedProtocolVersions: new Set(SUPPORTED_PROTOCOL_VERSIONS)
+const makeMcpProtocolState = Effect.fnUntraced(function*(
+  protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
+) {
+  return McpProtocolState.of({
+    sessions: {
+      bySessionId: new Map(),
+      byClientId: new Map()
+    },
+    protocolRegistry: yield* McpProtocolRegistry.make(protocols)
+  })
 })
 
-const layerMcpProtocolState = Layer.sync(McpProtocolState, makeMcpProtocolState)
+const layerMcpProtocolState = (
+  protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
+): Layer.Layer<McpProtocolState, Cause.IllegalArgumentError> =>
+  Layer.effect(McpProtocolState)(makeMcpProtocolState(protocols))
 
 /**
  * Runs an MCP server over the current `RpcServer.Protocol`.
@@ -367,38 +452,54 @@ const layerMcpProtocolState = Layer.sync(McpProtocolState, makeMcpProtocolState)
  * tools, resources, and prompts, and forwards queued server notifications to
  * initialized clients.
  *
- * @category constructors
+ * @category running
  * @since 4.0.0
  */
 export const run: (options: {
   readonly name: string
   readonly version: string
+  readonly protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }) => Effect.Effect<
   never,
-  never,
+  Cause.IllegalArgumentError,
   McpServer | RpcServer.Protocol
 > = Effect.fnUntraced(function*(options: {
   readonly name: string
   readonly version: string
+  readonly protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }) {
-  const protocolState = Option.getOrElse(
-    yield* Effect.serviceOption(McpProtocolState),
-    makeMcpProtocolState
-  )
+  const protocolStateOption = yield* Effect.serviceOption(McpProtocolState)
+  const protocolState = Option.isSome(protocolStateOption)
+    ? protocolStateOption.value
+    : yield* makeMcpProtocolState(options.protocols)
+  return yield* runWithProtocolState(options, protocolState)
+})
+
+const runWithProtocolState = Effect.fnUntraced(function*(options: {
+  readonly name: string
+  readonly version: string
+  readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
+}, protocolState: McpProtocolState["Service"]) {
+  const serverScope = yield* Effect.scope
+  const protocolRegistry = protocolState.protocolRegistry
   const protocol = yield* RpcServer.Protocol
   const server = yield* McpServer
   const isHttp = Option.isSome(yield* Effect.serviceOption(HttpRouter.HttpRouter))
-  const clientSessions = protocolState.clientSessions
-  const handlers = yield* Layer.build(layerHandlers(options)).pipe(
-    Effect.provideService(McpProtocolState, protocolState)
-  )
+  const sessions = protocolState.sessions
+  const clientProtocols = new Map<number, McpProtocol.ProtocolAdapter>()
+  const cancelledRequests = new Map<number, Set<string>>()
+  const handlers = yield* Layer.build(layerHandlers(options, {
+    sessions,
+    protocolRegistry
+  }))
 
   const clients = yield* RcMap.make({
-    lookup: Effect.fnUntraced(function*(clientId: number) {
+    lookup: Effect.fnUntraced(function*(key: McpClientKey) {
+      const selectedProtocol = protocolRegistry.select(key.protocolVersion)
       let write!: (message: RpcMessage.FromServerEncoded) => Effect.Effect<void>
-      const client = yield* RpcClient.make(ServerRequestRpcs, {
+      const client = yield* RpcClient.make(selectedProtocol.serverRequestRpcs, {
         spanPrefix: "McpServer/Client"
       }).pipe(
         Effect.provideServiceEffect(
@@ -409,7 +510,7 @@ export const run: (options: {
             return {
               send(id, request, _transferables) {
                 cid = id
-                return protocol.send(clientId, {
+                return protocol.send(key.clientId, {
                   ...request,
                   headers: undefined,
                   traceId: undefined,
@@ -430,35 +531,62 @@ export const run: (options: {
     idleTimeToLive: 10000
   })
 
-  const clientMiddleware = McpServerClientMiddleware.of((effect, { client, headers, rpc }) => {
-    const initializePayload = getInitializedClient(clientSessions, client.id, headers)
-    const isInitialize = rpc._tag === "initialize"
-    if (!isInitialize && !initializePayload) {
+  const clientMiddleware = McpServerClientMiddleware.of((effect, { client, headers, payload, rpc }) => {
+    const session = getClientSession(sessions, client.id, headers)
+    const isInitialize = rpc._tag.endsWith("/initialize")
+    if (!isInitialize && !session) {
       const fiber = Fiber.getCurrent()!
       const httpRequest = Context.getOrUndefined(fiber.context, HttpServerRequest.HttpServerRequest)
       if (httpRequest) {
         appendPreResponseHandlerUnsafe(
           httpRequest,
-          () => Effect.succeed(HttpServerResponse.empty({ status: 404 }))
+          () =>
+            Effect.succeed(
+              HttpServerResponse.empty({
+                status: headers[MCP_SESSION_ID_HEADER] === undefined ? 400 : 404
+              })
+            )
         )
       }
       return Effect.die(new Error(`Mcp-Session-Id does not exist`))
     }
-    return Effect.provideService(
-      effect,
-      McpServerClient,
-      McpServerClient.of({
-        clientId: client.id,
-        initializePayload: initializePayload!,
-        getClient: RcMap.get(clients, client.id).pipe(
-          Effect.map(({ client }) => client)
-        )
-      })
+    const selectedProtocol = session?.protocol ?? protocolForInternalTag(protocolRegistry, rpc._tag)
+    return effect.pipe(
+      Effect.provideService(
+        McpServerClient,
+        McpServerClient.of({
+          clientId: client.id,
+          protocolVersion: selectedProtocol.protocolVersion,
+          initializePayload: session?.initializePayload ?? payload as typeof Initialize.payloadSchema.Type,
+          getClient: RcMap.get(
+            clients,
+            new McpClientKey({
+              clientId: client.id,
+              protocolVersion: selectedProtocol.protocolVersion
+            })
+          ).pipe(
+            Effect.map(({ client }) => client)
+          )
+        })
+      ),
+      Effect.provideService(CurrentLogLevel, effectLogLevel(session?.logLevel))
     )
   })
 
   const patchedProtocol = RpcServer.Protocol.of({
     ...protocol,
+    send: (clientId, response) => {
+      if (response._tag === "Exit") {
+        const requests = cancelledRequests.get(clientId)
+        if (requests?.delete(requestKey(response.requestId)) === true) {
+          if (requests.size === 0) {
+            cancelledRequests.delete(clientId)
+          }
+          return Effect.void
+        }
+      }
+      return protocol.send(clientId, response)
+    },
     run: (f) =>
       protocol.run((clientId, request_) => {
         const fiber = Fiber.getCurrent()!
@@ -483,56 +611,161 @@ export const run: (options: {
         }
         switch (request._tag) {
           case "Request": {
-            if (httpRequest !== undefined) {
-              const client = getInitializedClient(clientSessions, clientId, httpRequest.headers)
-              if (client) {
+            const headers = isHttp
+              ? Context.getUnsafe(
+                Fiber.getCurrent()!.context,
+                HttpServerRequest.HttpServerRequest
+              ).headers
+              : Headers.fromInput(request.headers)
+            const session = getClientSession(sessions, clientId, headers)
+            const selectedProtocol = session?.protocol ??
+              (request.tag === "initialize"
+                ? protocolRegistry.select(getOfferedProtocolVersion(request.payload))
+                : protocolRegistry.protocols[0])
+            clientProtocols.set(clientId, selectedProtocol)
+            if (request.tag === MCP_INVALID_BATCH_METHOD) {
+              return protocol.send(clientId, {
+                _tag: "Exit",
+                requestId: request.id,
+                exit: {
+                  _tag: "Failure",
+                  cause: [{
+                    _tag: "Fail",
+                    error: new InvalidRequest({ message: "JSON-RPC batches are not supported" })
+                  }]
+                }
+              })
+            }
+            if (isHttp) {
+              const fiber = Fiber.getCurrent()!
+              const httpRequest = Context.getUnsafe(fiber.context, HttpServerRequest.HttpServerRequest)
+              if (session) {
                 appendPreResponseHandlerUnsafe(httpRequest, (_, res) =>
                   Effect.succeed(
-                    HttpServerResponse.setHeader(res, MCP_PROTOCOL_VERSION_HEADER, client.protocolVersion)
+                    HttpServerResponse.setHeader(
+                      res,
+                      MCP_PROTOCOL_VERSION_HEADER,
+                      session.protocol.protocolVersion
+                    )
                   ))
               }
             }
-            const rpc = ClientNotificationRpcs.requests.get(request.tag)
-            if (rpc) {
-              const headers = Headers.fromInput(request.headers)
-              if (!getInitializedClient(clientSessions, clientId, headers)) {
-                if (httpRequest !== undefined) {
+            const routedRequest = protocolRegistry.routeClientRequest(selectedProtocol, request)
+            const rpc = protocolRegistry.clientRpcs.requests.get(routedRequest.tag)
+            if (rpc && selectedProtocol.clientNotificationRpcs.requests.has(request.tag)) {
+              if (!session) {
+                if (httpRequest) {
                   appendPreResponseHandlerUnsafe(
                     httpRequest,
-                    () => Effect.succeed(HttpServerResponse.empty({ status: 404 }))
+                    () =>
+                      Effect.succeed(
+                        HttpServerResponse.empty({
+                          status: headers[MCP_SESSION_ID_HEADER] === undefined ? 400 : 404
+                        })
+                      )
                   )
                 }
                 return Effect.void
               }
               if (request.tag === "notifications/cancelled") {
-                return f(clientId, {
-                  _tag: "Interrupt",
-                  requestId: String((request.payload as any).requestId)
-                })
+                return decodeCancelledNotification(request.payload).pipe(
+                  Effect.flatMap(({ requestId }) => {
+                    const requests = cancelledRequests.get(clientId) ?? new Set<string>()
+                    requests.add(requestKey(requestId))
+                    cancelledRequests.set(clientId, requests)
+                    return f(clientId, {
+                      _tag: "Interrupt",
+                      requestId: RpcMessage.RequestId(requestId)
+                    })
+                  }),
+                  Effect.catchCause(() => Effect.void)
+                )
               }
-              const handler = handlers.mapUnsafe.get(request.tag) as Rpc.Handler<string>
-              return handler
-                ? handler.handler(request.payload, {
-                  rpc,
-                  requestId: RpcMessage.RequestId(request.id),
-                  client: new Rpc.ServerClient(clientId),
-                  headers
-                }) as any as Effect.Effect<void>
-                : Effect.void
+              return selectedProtocol.payloadCodecs(rpc).decode(request.payload).pipe(
+                Effect.flatMap((payload) => {
+                  if (
+                    request.tag === "notifications/roots/list_changed" &&
+                    session.initializePayload.capabilities.roots?.listChanged === true
+                  ) {
+                    if (httpRequest !== undefined) {
+                      return Effect.void
+                    }
+                    return RcMap.get(
+                      clients,
+                      new McpClientKey({
+                        clientId,
+                        protocolVersion: selectedProtocol.protocolVersion
+                      })
+                    ).pipe(
+                      Effect.flatMap(({ client }) => client["roots/list"](undefined)),
+                      Effect.scoped,
+                      Effect.ignoreCause,
+                      Effect.forkIn(serverScope),
+                      Effect.asVoid
+                    )
+                  }
+                  const handler = handlers.mapUnsafe.get(rpc.key) as Rpc.Handler<string> | undefined
+                  return handler
+                    ? handler.handler(payload, {
+                      rpc,
+                      requestId: RpcMessage.RequestId(request.id),
+                      client: new Rpc.ServerClient(clientId),
+                      headers
+                    }) as any as Effect.Effect<void>
+                    : Effect.void
+                }),
+                Effect.catchCause(() => Effect.void)
+              )
             }
-            return f(clientId, request)
+            if (!rpc) {
+              if (request.isNotification) {
+                return Effect.void
+              }
+              return protocol.send(clientId, {
+                _tag: "Exit",
+                requestId: request.id,
+                exit: {
+                  _tag: "Failure",
+                  cause: [{ _tag: "Fail", error: new MethodNotFound({ message: `Method not found: ${request.tag}` }) }]
+                }
+              })
+            }
+            return selectedProtocol.payloadCodecs(rpc).decode(request.payload).pipe(
+              Effect.matchEffect({
+                onSuccess: () => f(clientId, routedRequest),
+                onFailure: () =>
+                  request.isNotification
+                    ? Effect.void
+                    : protocol.send(clientId, {
+                      _tag: "Exit",
+                      requestId: request.id,
+                      exit: {
+                        _tag: "Failure",
+                        cause: [{ _tag: "Fail", error: new InvalidParams({ message: "Invalid method parameters" }) }]
+                      }
+                    })
+              })
+            )
           }
           case "Ping":
           case "Ack":
           case "Interrupt":
+            return f(clientId, request)
           case "Eof":
+            cancelledRequests.delete(clientId)
             return f(clientId, request)
           case "Pong":
           case "Exit":
           case "Chunk":
           case "ClientProtocolError":
           case "Defect":
-            return RcMap.get(clients, clientId).pipe(
+            return RcMap.get(
+              clients,
+              new McpClientKey({
+                clientId,
+                protocolVersion: getProtocolForClient(clientProtocols, clientId, protocolRegistry).protocolVersion
+              })
+            ).pipe(
               Effect.flatMap(({ write }) => write(request)),
               Effect.scoped
             )
@@ -540,23 +773,52 @@ export const run: (options: {
       })
   })
 
-  const encodeNotification = Schema.encodeUnknownEffect(
-    Schema.Union(Array.from(ServerNotificationRpcs.requests.values(), (rpc) => rpc.payloadSchema))
-  )
   yield* Queue.take(server.notificationsQueue).pipe(
     Effect.flatMap(Effect.fnUntraced(function*(request) {
-      const encoded = yield* encodeNotification(request.payload)
-      const message: RpcMessage.RequestEncoded = {
-        _tag: "Request",
-        tag: request.tag,
-        payload: encoded
-      } as any
       const clientIds = yield* patchedProtocol.clientIds
+      for (const clientId of clientProtocols.keys()) {
+        if (!clientIds.has(clientId)) {
+          clientProtocols.delete(clientId)
+          sessions.byClientId.delete(clientId)
+        }
+      }
       for (const clientId of server.initializedClients.keys()) {
         if (!clientIds.has(clientId)) {
           server.initializedClients.delete(clientId)
           continue
         }
+        const selectedProtocol = clientProtocols.get(clientId)
+        if (!selectedProtocol) {
+          continue
+        }
+        const rpc = selectedProtocol.serverNotificationRpcs.requests.get(request.tag)
+        if (!rpc) {
+          continue
+        }
+        if (request.tag === "notifications/message") {
+          const { level } = yield* Schema.decodeUnknownEffect(
+            LoggingMessageNotification.payloadSchema
+          )(request.payload)
+          if (!isMcpLogLevelEnabled(level, sessions.byClientId.get(clientId)?.logLevel)) {
+            continue
+          }
+        }
+        if (request.tag === "notifications/resources/updated") {
+          const { uri } = yield* Schema.decodeUnknownEffect(
+            ResourceUpdatedNotification.payloadSchema
+          )(request.payload)
+          if (sessions.byClientId.get(clientId)?.resourceSubscriptions?.has(uri) !== true) {
+            continue
+          }
+        }
+        const encoded = yield* selectedProtocol.payloadCodecs(rpc).encode(request.payload)
+        // TODO: Extend RpcServer.Protocol's outbound message contract with server-originated
+        // notifications so MCP does not need to treat this notification as an RPC response.
+        const message: RpcMessage.RequestEncoded = {
+          _tag: "Request",
+          tag: request.tag,
+          payload: encoded
+        } as any
         yield* patchedProtocol.send(clientId, message as any)
       }
     })),
@@ -565,7 +827,7 @@ export const run: (options: {
     Effect.forkScoped
   )
 
-  return yield* RpcServer.make(ClientRpcs, {
+  return yield* RpcServer.make(protocolRegistry.clientRpcs, {
     spanPrefix: "McpServer",
     disableFatalDefects: true
   }).pipe(
@@ -607,10 +869,11 @@ export const run: (options: {
 export const layer = (options: {
   readonly name: string
   readonly version: string
+  readonly protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
-}): Layer.Layer<McpServer | McpServerClient, never, RpcServer.Protocol> =>
+}): Layer.Layer<McpServer | McpServerClient, Cause.IllegalArgumentError, RpcServer.Protocol> =>
   layerWithProtocolState(options).pipe(
-    Layer.provide(layerMcpProtocolState)
+    Layer.provide(layerMcpProtocolState(options.protocols))
   )
 
 const layerWithProtocolState = (options: {
@@ -618,63 +881,133 @@ const layerWithProtocolState = (options: {
   readonly version: string
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }): Layer.Layer<McpServer | McpServerClient, never, RpcServer.Protocol | McpProtocolState> =>
-  Layer.effectDiscard(Effect.forkScoped(run(options))).pipe(
+  Layer.effectDiscard(
+    Effect.gen(function*() {
+      const protocolState = yield* McpProtocolState
+      yield* Effect.forkScoped(runWithProtocolState(options, protocolState))
+    })
+  ).pipe(
     Layer.provideMerge(McpServer.layer)
   )
 
+const StdioInitializeRequest = Schema.Struct({
+  method: Schema.Literal("initialize"),
+  params: Schema.Struct({
+    protocolVersion: Schema.String
+  })
+})
+
+const StdioInvalidBatchExit = Schema.Struct({
+  _tag: Schema.Literal("Exit"),
+  requestId: Schema.Null,
+  exit: Schema.Struct({
+    cause: Schema.Unknown
+  })
+})
+
+const decodeStdioInitializeRequest = Schema.decodeUnknownOption(StdioInitializeRequest)
+const decodeStdioInvalidBatchExit = Schema.decodeUnknownOption(StdioInvalidBatchExit)
+
+const makeStdioSerialization = (
+  protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
+): RpcSerialization.RpcSerialization["Service"] =>
+  RpcSerialization.RpcSerialization.of({
+    contentType: "application/json-rpc",
+    includesFraming: true,
+    makeUnsafe: () => {
+      const framing = RpcSerialization.ndjson.makeUnsafe()
+      const jsonRpc = RpcSerialization.jsonRpc().makeUnsafe()
+      const protocolsByVersion = new Map<string, McpProtocol.ProtocolAdapter>(
+        protocols.map((protocol) => [protocol.protocolVersion, protocol])
+      )
+      let selectedProtocol = protocols[0]
+      return {
+        decode: (data) => {
+          const frames = framing.decode(data)
+          const messages: Array<unknown> = []
+          for (const frame of frames) {
+            const entries = Array.isArray(frame) ? frame : [frame]
+            const initialize = Arr.findFirst(entries, (entry) => decodeStdioInitializeRequest(entry))
+            selectedProtocol = Option.match(initialize, {
+              onNone: () => selectedProtocol,
+              onSome: ({ params }) => protocolsByVersion.get(params.protocolVersion) ?? protocols[0]
+            })
+            if (Array.isArray(frame) && !selectedProtocol.transport.acceptsJsonRpcBatches) {
+              messages.push({
+                _tag: "Request",
+                id: null,
+                tag: MCP_INVALID_BATCH_METHOD,
+                payload: null,
+                headers: []
+              })
+            } else {
+              messages.push(...jsonRpc.decode(JSON.stringify(frame)))
+            }
+          }
+          return messages
+        },
+        encode: (response) => {
+          const invalidBatchExit = decodeStdioInvalidBatchExit(response)
+          if (Option.isSome(invalidBatchExit)) {
+            return framing.encode({
+              jsonrpc: "2.0",
+              id: null,
+              error: {
+                _tag: "Cause",
+                code: INVALID_REQUEST_ERROR_CODE,
+                message: "JSON-RPC batches are not supported",
+                data: invalidBatchExit.value.exit.cause
+              }
+            })
+          }
+          const encoded = jsonRpc.encode(response)
+          return encoded === undefined ? undefined : `${encoded}\n`
+        }
+      }
+    }
+  })
+
 /**
  * Runs the McpServer, using stdio for input and output.
  *
- * **Example** (Running an MCP server over stdio)
+ * **Example** (Configuring an MCP server over stdio)
  *
- * ```ts
- * import { Effect, Layer, Logger, Schema } from "effect"
- * import { NodeRuntime, NodeStdio } from "@effect/platform-node"
- * import { McpSchema, McpServer } from "effect/unstable/ai"
+ * ```ts import.meta.vitest
+ * import { Effect, Layer, Schema } from "effect"
+ * import { McpProtocol, McpSchema, McpServer } from "effect/unstable/ai"
  *
  * const idParam = McpSchema.param("id", Schema.Number)
  *
- * // Define a resource template for a README file
  * const ReadmeTemplate = McpServer.resource`file://readme/${idParam}`({
  *   name: "README Template",
- *   // You can add auto-completion for the ID parameter
  *   completion: {
- *     id: (_) => Effect.succeed([1, 2, 3, 4, 5])
+ *     id: () => Effect.succeed([1, 2, 3])
  *   },
- *   content: Effect.fn(function*(_uri, id) {
- *     return `# MCP Server Demo - ID: ${id}`
- *   })
+ *   content: (_uri, id) => Effect.succeed(`# MCP Server Demo - ID: ${id}`)
  * })
  *
- * // Define a test prompt with parameters
  * const TestPrompt = McpServer.prompt({
  *   name: "Test Prompt",
- *   description: "A test prompt to demonstrate MCP server capabilities",
+ *   description: "Looks up flight booking details",
  *   parameters: {
  *     flightNumber: Schema.String
  *   },
  *   completion: {
- *     flightNumber: () => Effect.succeed(["FL123", "FL456", "FL789"])
+ *     flightNumber: () => Effect.succeed(["FL123", "FL456"])
  *   },
  *   content: ({ flightNumber }) =>
  *     Effect.succeed(`Get the booking details for flight number: ${flightNumber}`)
  * })
  *
- * // Merge all the resources and prompts into a single server layer
- * const ServerLayer = Layer.mergeAll(
- *   ReadmeTemplate,
- *   TestPrompt
- * ).pipe(
- *   // Provide the MCP server implementation
+ * const ServerLayer = Layer.mergeAll(ReadmeTemplate, TestPrompt).pipe(
  *   Layer.provide(McpServer.layerStdio({
  *     name: "Demo Server",
  *     version: "1.0.0",
- *   })),
- *   Layer.provide(NodeStdio.layer),
- *   Layer.provide(Layer.succeed(Logger.LogToStderr)(true))
+ *     protocols: [McpProtocol.v2025_06_18]
+ *   }))
  * )
  *
- * Layer.launch(ServerLayer).pipe(NodeRuntime.runMain)
+ * Layer.isLayer(ServerLayer) // => true
  * ```
  *
  * @category layers
@@ -683,11 +1016,14 @@ const layerWithProtocolState = (options: {
 export const layerStdio = (options: {
   readonly name: string
   readonly version: string
+  readonly protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
-}): Layer.Layer<McpServer | McpServerClient, never, Stdio> =>
+}): Layer.Layer<McpServer | McpServerClient, Cause.IllegalArgumentError, Stdio> =>
   layer(options).pipe(
     Layer.provide(RpcServer.layerProtocolStdio),
-    Layer.provide(RpcSerialization.layerNdJsonRpc())
+    Layer.provide(
+      Layer.succeed(RpcSerialization.RpcSerialization)(makeStdioSerialization(options.protocols))
+    )
   )
 
 /**
@@ -701,7 +1037,7 @@ export const layerStdio = (options: {
  *
  * POST serves JSON-RPC and accepted notification-only requests return `202`.
  * Unsupported protocol versions return `400`; methods without MCP handlers
- * return `405`.
+ * return `405`. Browser Origins are rejected unless listed in `allowedOrigins`.
  *
  * @see {@link layerStdio} for exposing the server over stdio
  * @see {@link layer} for the base MCP server layer without a transport protocol
@@ -713,18 +1049,28 @@ export const layerHttp = (options: {
   readonly name: string
   readonly version: string
   readonly path: HttpRouter.PathInput
+  readonly protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter>
+  readonly allowedOrigins?: ReadonlyArray<string> | undefined
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
-}): Layer.Layer<McpServer | McpServerClient, never, HttpRouter.HttpRouter> => {
-  const protocolState = layerMcpProtocolState
+}): Layer.Layer<McpServer | McpServerClient, Cause.IllegalArgumentError, HttpRouter.HttpRouter> => {
+  const protocolState = layerMcpProtocolState(options.protocols)
   const methodNotAllowedResponse = HttpServerResponse.empty({
     status: 405,
     headers: { allow: "POST" }
   })
+  const methodNotAllowed = (request: HttpServerRequest.HttpServerRequest) =>
+    Effect.succeed(
+      request.headers.origin !== undefined &&
+        !options.allowedOrigins?.includes(request.headers.origin)
+        ? HttpServerResponse.empty({ status: 403 })
+        : methodNotAllowedResponse
+    )
   const routes = Layer.mergeAll(
-    HttpRouter.add("GET", options.path, methodNotAllowedResponse),
-    HttpRouter.add("PUT", options.path, methodNotAllowedResponse),
-    HttpRouter.add("PATCH", options.path, methodNotAllowedResponse),
-    HttpRouter.add("DELETE", options.path, methodNotAllowedResponse)
+    HttpRouter.add("GET", options.path, methodNotAllowed),
+    HttpRouter.add("PUT", options.path, methodNotAllowed),
+    HttpRouter.add("PATCH", options.path, methodNotAllowed),
+    HttpRouter.add("DELETE", options.path, methodNotAllowed),
+    HttpRouter.add("OPTIONS", options.path, methodNotAllowed)
   )
   return Layer.merge(layerWithProtocolState(options), routes).pipe(
     Layer.provide(layerMcpProtocolHttp(options)),
@@ -735,6 +1081,7 @@ export const layerHttp = (options: {
 
 const layerMcpProtocolHttp = (options: {
   readonly path: HttpRouter.PathInput
+  readonly allowedOrigins?: ReadonlyArray<string> | undefined
 }): Layer.Layer<
   RpcServer.Protocol,
   never,
@@ -745,14 +1092,142 @@ const layerMcpProtocolHttp = (options: {
     const { httpEffect, protocol } = yield* RpcServer.makeProtocolWithHttpEffect
     const router = yield* HttpRouter.HttpRouter
     yield* router.add("POST", options.path, (request) => {
+      if (
+        request.headers.origin !== undefined &&
+        !options.allowedOrigins?.includes(request.headers.origin)
+      ) {
+        return Effect.succeed(HttpServerResponse.empty({ status: 403 }))
+      }
+      const contentType = request.headers["content-type"]?.split(";", 1)[0]?.trim().toLowerCase()
+      if (contentType !== "application/json") {
+        return Effect.succeed(HttpServerResponse.empty({ status: 415 }))
+      }
+      const accepted = new Set<string>()
+      for (const entry of request.headers.accept?.split(",") ?? []) {
+        const [mediaType, ...parameters] = entry.split(";").map((part) => part.trim().toLowerCase())
+        let quality = 1
+        for (const parameter of parameters) {
+          const [name, value] = parameter.split("=", 2).map((part) => part.trim())
+          if (name === "q") {
+            quality = value === undefined ? Number.NaN : Number(value)
+          }
+        }
+        if (mediaType !== undefined && quality > 0 && quality <= 1) {
+          accepted.add(mediaType)
+        }
+      }
+      if (!accepted.has("application/json") || !accepted.has("text/event-stream")) {
+        return Effect.succeed(HttpServerResponse.empty({ status: 406 }))
+      }
       const protocolVersion = request.headers[MCP_PROTOCOL_VERSION_HEADER]
       if (
         protocolVersion !== undefined &&
-        !state.supportedProtocolVersions.has(protocolVersion)
+        !state.protocolRegistry.protocols.some((protocol) => protocol.protocolVersion === protocolVersion)
       ) {
         return Effect.succeed(HttpServerResponse.empty({ status: 400 }))
       }
-      return httpEffect
+      return request.text.pipe(
+        Effect.matchEffect({
+          onFailure: () =>
+            Effect.succeed(
+              HttpServerResponse.jsonUnsafe({
+                jsonrpc: "2.0",
+                id: null,
+                error: new ParseError({
+                  message: "Parse error"
+                })
+              })
+            ),
+          onSuccess: (body) => {
+            return Effect.match(Schema.decodeUnknownEffect(Schema.UnknownFromJsonString)(body), {
+              onFailure: () => ({
+                _tag: "Error" as const,
+                id: null,
+                error: new ParseError({ message: "Parse error" })
+              }),
+              onSuccess: (input) => {
+                if (Array.isArray(input)) {
+                  const sessionId = request.headers[MCP_SESSION_ID_HEADER]
+                  const session = sessionId === undefined ? undefined : state.sessions.bySessionId.get(sessionId)
+                  let selectedProtocol = session?.protocol ?? state.protocolRegistry.protocols[0]
+                  for (const entry of input) {
+                    if (
+                      Predicate.hasProperty(entry, "method") &&
+                      entry.method === "initialize" &&
+                      Predicate.hasProperty(entry, "params") &&
+                      Predicate.hasProperty(entry.params, "protocolVersion") &&
+                      typeof entry.params.protocolVersion === "string"
+                    ) {
+                      selectedProtocol = state.protocolRegistry.select(entry.params.protocolVersion)
+                      break
+                    }
+                  }
+                  if (!selectedProtocol.transport.acceptsJsonRpcBatches) {
+                    return { _tag: "HttpError" as const, status: 400 }
+                  }
+                }
+                const hasId = Predicate.hasProperty(input, "id")
+                const id = hasId && (typeof input.id === "string" || typeof input.id === "number")
+                  ? input.id
+                  : null
+                const isJsonRpc = Predicate.hasProperty(input, "jsonrpc") && input.jsonrpc === "2.0"
+                const hasValidRequestId = hasId === false || typeof input.id === "string" ||
+                  typeof input.id === "number"
+                const isRequest = isJsonRpc && hasValidRequestId &&
+                  Predicate.hasProperty(input, "method") && typeof input.method === "string"
+                const hasValidResponseId = hasId &&
+                  (typeof input.id === "string" || typeof input.id === "number" || input.id === null)
+                const hasResult = Predicate.hasProperty(input, "result")
+                const hasError = Predicate.hasProperty(input, "error")
+                const isResponse = isJsonRpc && hasValidResponseId && hasResult !== hasError
+                const isInitialize = isRequest && input.method === "initialize"
+                const sessionId = request.headers[MCP_SESSION_ID_HEADER]
+                const session = sessionId === undefined ? undefined : state.sessions.bySessionId.get(sessionId)
+                if (isInitialize && sessionId !== undefined) {
+                  return {
+                    _tag: "HttpError" as const,
+                    status: session === undefined ? 404 : 400
+                  }
+                }
+                if (!isInitialize && isRequest && session === undefined) {
+                  return {
+                    _tag: "HttpError" as const,
+                    status: sessionId === undefined ? 400 : 404
+                  }
+                }
+                if (
+                  session !== undefined &&
+                  session.protocol.transport.requiresVersionHeader &&
+                  protocolVersion !== session.protocol.protocolVersion
+                ) {
+                  return { _tag: "HttpError" as const, status: 400 }
+                }
+                return isRequest || isResponse
+                  ? { _tag: "Success" as const }
+                  : {
+                    _tag: "Error" as const,
+                    id,
+                    error: new InvalidRequest({ message: "Invalid Request" })
+                  }
+              }
+            }).pipe(
+              Effect.flatMap((decoded) =>
+                decoded._tag === "HttpError"
+                  ? Effect.succeed(HttpServerResponse.empty({ status: decoded.status }))
+                  : decoded._tag === "Error"
+                  ? Effect.succeed(
+                    HttpServerResponse.jsonUnsafe({
+                      jsonrpc: "2.0",
+                      id: decoded.id,
+                      error: decoded.error
+                    })
+                  )
+                  : httpEffect
+              )
+            )
+          }
+        })
+      )
     })
     return protocol
   }))
@@ -768,7 +1243,7 @@ const toolErrorResult = (message: string): CallToolResult =>
 /**
  * Registers a `Toolkit` with the `McpServer`.
  *
- * @category tools
+ * @category handlers
  * @since 4.0.0
  */
 export const registerToolkit: <Tools extends Record<string, Tool.Any>>(
@@ -791,10 +1266,12 @@ export const registerToolkit: <Tools extends Record<string, Tool.Any>>(
     const annotations = tool.annotations
     const toolMeta = Context.getOrUndefined(annotations, Tool.Meta)
     const isDeclaredFailure = Schema.is(tool.failureSchema)
+    const outputSchema = Tool.getJsonSchemaFromSchema(tool.successSchema)
     const mcpTool = new McpTool({
       name: tool.name,
       description: Tool.getDescription(tool),
       inputSchema: Tool.getJsonSchema(tool),
+      ...(outputSchema.type === "object" ? { outputSchema } : {}),
       annotations: {
         ...(Context.getOption(tool.annotations, Tool.Title).pipe(
           Option.map((title) => ({ title })),
@@ -833,7 +1310,7 @@ export const registerToolkit: <Tools extends Record<string, Tool.Any>>(
             if (AiError.isAiError(error)) {
               const reason = (error as AiError.AiError).reason
               return reason._tag === "ToolParameterValidationError"
-                ? Effect.succeed(toolErrorResult(reason.message))
+                ? Effect.fail(new InvalidParams({ message: reason.message }))
                 : Effect.succeed(toolErrorResult(INTERNAL_TOOL_ERROR_MESSAGE))
             }
             if (isDeclaredFailure(error)) {
@@ -854,7 +1331,7 @@ export const registerToolkit: <Tools extends Record<string, Tool.Any>>(
 /**
  * Registers an `AiToolkit` with the `McpServer`.
  *
- * @category tools
+ * @category layers
  * @since 4.0.0
  */
 export const toolkit = <Tools extends Record<string, Tool.Any>>(
@@ -878,7 +1355,11 @@ export const toolkit = <Tools extends Record<string, Tool.Any>>(
 export type ValidateCompletions<Completions, Keys extends string> =
   & Completions
   & {
-    readonly [K in keyof Completions]: K extends Keys ? (input: string) => any : never
+    readonly [K in keyof Completions]: K extends Keys ? (
+        input: string,
+        context: CompletionContext
+      ) => any
+      : never
   }
 
 /**
@@ -897,7 +1378,10 @@ export type ResourceCompletions<Schemas extends ReadonlyArray<Schema.Constraint>
   readonly [
     K in Extract<keyof Schemas, `${number}`> as Schemas[K] extends Param<infer Id, infer _S> ? Id
       : `param${K}`
-  ]: (input: string) => Effect.Effect<Array<Schemas[K]["Type"]>, any, any>
+  ]: (
+    input: string,
+    context: CompletionContext
+  ) => Effect.Effect<Array<Schemas[K]["Type"]>, any, any>
 }
 
 /**
@@ -911,7 +1395,7 @@ export type ResourceCompletions<Schemas extends ReadonlyArray<Schema.Constraint>
  *
  * @see {@link resource} for the layer-based resource registration wrapper
  *
- * @category resources
+ * @category handlers
  * @since 4.0.0
  */
 export const registerResource: {
@@ -926,7 +1410,7 @@ export const registerResource: {
    *
    * @see {@link resource} for the layer-based resource registration wrapper
    *
-   * @category resources
+   * @category handlers
    * @since 4.0.0
    */
   <E, R>(
@@ -956,7 +1440,7 @@ export const registerResource: {
    *
    * @see {@link resource} for the layer-based resource registration wrapper
    *
-   * @category resources
+   * @category handlers
    * @since 4.0.0
    */
   <const Schemas extends ReadonlyArray<Schema.Constraint>>(segments: TemplateStringsArray, ...schemas: Schemas): <
@@ -1034,7 +1518,15 @@ export const registerResource: {
     readonly mimeType?: string | undefined
     readonly audience?: ReadonlyArray<"user" | "assistant"> | undefined
     readonly priority?: number | undefined
-    readonly completion?: Record<string, (input: string) => Effect.Effect<any>> | undefined
+    readonly completion?:
+      | Record<
+        string,
+        (
+          input: string,
+          context: CompletionContext
+        ) => Effect.Effect<any>
+      >
+      | undefined
     readonly content: (uri: string, ...params: Array<any>) => Effect.Effect<
       typeof ReadResourceResult.Type | string | Uint8Array,
       E,
@@ -1050,13 +1542,20 @@ export const registerResource: {
       uriTemplate: uriPath,
       annotations: options!
     })
-    const completions: Record<string, (input: string) => Effect.Effect<CompleteResult, InternalError>> = Object.create(
-      null
-    )
+    const completions: Record<
+      string,
+      (
+        input: string,
+        context: CompletionContext
+      ) => Effect.Effect<CompleteResult, InternalError>
+    > = Object.create(null)
     for (const [param, handle] of Object.entries(options.completion ?? {})) {
       const encodeArray = Schema.encodeUnknownEffect(Schema.Array(params[param]))
-      const handler = (input: string) =>
-        handle(input).pipe(
+      const handler = (
+        input: string,
+        context: CompletionContext
+      ) =>
+        handle(input, context).pipe(
           Effect.flatMap(encodeArray),
           Effect.map((values) => ({
             completion: {
@@ -1105,7 +1604,7 @@ export const registerResource: {
  *
  * @see {@link registerResource} for the Effect-level resource registration API
  *
- * @category resources
+ * @category layers
  * @since 4.0.0
  */
 export const resource: {
@@ -1118,7 +1617,7 @@ export const resource: {
    *
    * @see {@link registerResource} for the Effect-level resource registration API
    *
-   * @category resources
+   * @category layers
    * @since 4.0.0
    */
   <E, R>(
@@ -1145,7 +1644,7 @@ export const resource: {
    *
    * @see {@link registerResource} for the Effect-level resource registration API
    *
-   * @category resources
+   * @category layers
    * @since 4.0.0
    */
   <const Schemas extends ReadonlyArray<Schema.Constraint>>(segments: TemplateStringsArray, ...schemas: Schemas): <
@@ -1204,7 +1703,7 @@ export const resource: {
  *
  * @see {@link prompt} for the layer-based prompt registration wrapper
  *
- * @category prompts
+ * @category handlers
  * @since 4.0.0
  */
 export const registerPrompt = <
@@ -1212,7 +1711,10 @@ export const registerPrompt = <
   R,
   Params extends Schema.Struct.Fields = {},
   const Completions extends {
-    readonly [K in keyof Params]?: (input: string) => Effect.Effect<Array<Params[K]>, any, any>
+    readonly [K in keyof Params]?: (
+      input: string,
+      context: CompletionContext
+    ) => Effect.Effect<Array<Params[K]>, any, any>
   } = {}
 >(
   options: {
@@ -1241,18 +1743,30 @@ export const registerPrompt = <
   const decode = options.parameters
     ? Schema.decodeEffect(Schema.Struct(props))
     : () => Effect.succeed({} as Params)
-  const completion: Record<string, (input: string) => Effect.Effect<any>> = options.completion ?? {}
+  const completion: Record<
+    string,
+    (
+      input: string,
+      context: CompletionContext
+    ) => Effect.Effect<any>
+  > = options.completion ?? {}
   return Effect.gen(function*() {
     const registry = yield* McpServer
     const services = yield* Effect.context<Exclude<R | Schema.Struct.DecodingServices<Params>, McpServerClient>>()
     const completions: Record<
       string,
-      (input: string) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
+      (
+        input: string,
+        context: CompletionContext
+      ) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
     > = Object.create(null)
     for (const [param, handle] of Object.entries(completion)) {
       const encodeArray = Schema.encodeEffect(Schema.Array(props[param]))
-      const handler = (input: string) =>
-        handle(input).pipe(
+      const handler = (
+        input: string,
+        context: CompletionContext
+      ) =>
+        handle(input, context).pipe(
           Effect.flatMap(encodeArray),
           Effect.map((values) => ({
             completion: {
@@ -1276,20 +1790,23 @@ export const registerPrompt = <
       handle: (params) =>
         decode(params).pipe(
           Effect.mapError((error) => new InvalidParams({ message: error.message })),
-          Effect.flatMap((params) => options.content(params as any)),
-          Effect.map((messages) => {
-            messages = typeof messages === "string" ?
-              [{
-                role: "user",
-                content: TextContent.make({ text: messages })
-              }] :
-              messages
-            return new GetPromptResult({ messages, description: prompt.description })
-          }),
-          Effect.catchCause((cause) => {
-            const prettyError = Cause.prettyErrors(cause)[0]
-            return Effect.fail(new InternalError({ message: prettyError.message }))
-          }),
+          Effect.flatMap((params) =>
+            options.content(params as any).pipe(
+              Effect.map((messages) => {
+                messages = typeof messages === "string" ?
+                  [{
+                    role: "user",
+                    content: TextContent.make({ text: messages })
+                  }] :
+                  messages
+                return new GetPromptResult({ messages, description: prompt.description })
+              }),
+              Effect.catchCause((cause) => {
+                const prettyError = Cause.prettyErrors(cause)[0]
+                return Effect.fail(new InternalError({ message: prettyError.message }))
+              })
+            )
+          ),
           Effect.provideContext(services as Context.Context<unknown>)
         )
     })
@@ -1311,7 +1828,7 @@ export const registerPrompt = <
  *
  * @see {@link registerPrompt} for the Effect-level prompt registration API
  *
- * @category prompts
+ * @category layers
  * @since 4.0.0
  */
 export const prompt = <
@@ -1319,7 +1836,10 @@ export const prompt = <
   R,
   Params extends Schema.Struct.Fields = {},
   const Completions extends {
-    readonly [K in keyof Params]?: (input: string) => Effect.Effect<Array<Params[K]["Type"]>, any, any>
+    readonly [K in keyof Params]?: (
+      input: string,
+      context: CompletionContext
+    ) => Effect.Effect<Array<Params[K]["Type"]>, any, any>
   } = {}
 >(
   options: {
@@ -1346,7 +1866,7 @@ export const prompt = <
  * Accepted content is decoded with the supplied schema, declined requests fail
  * with `ElicitationDeclined`, and canceled requests interrupt the effect.
  *
- * @category elicitation
+ * @category accessors
  * @since 4.0.0
  */
 export const elicit: <S extends Schema.ConstraintEncoder<Record<string, unknown>, unknown>>(options: {
@@ -1383,7 +1903,7 @@ export const elicit: <S extends Schema.ConstraintEncoder<Record<string, unknown>
 /**
  * Accesses the current client's capabilities.
  *
- * @category capabilities
+ * @category accessors
  * @since 4.0.0
  */
 export const clientCapabilities: Effect.Effect<
@@ -1442,145 +1962,191 @@ const layerHandlers = (serverInfo: {
   readonly name: string
   readonly version: string
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
+}, options: {
+  readonly sessions: Sessions
+  readonly protocolRegistry: McpProtocolRegistry.ProtocolRegistry<McpProtocol.ProtocolAdapter>
 }) =>
-  ClientRpcs.toLayer(
+  Layer.effectContext(
     Effect.gen(function*() {
       const server = yield* McpServer
-      const protocolState = yield* McpProtocolState
-      const clientSessions = protocolState.clientSessions
-      let currentLogLevel = yield* CurrentLogLevel
-
-      return ClientRpcs.of({
-        // Requests
-        ping: () => Effect.succeed({}),
-        initialize(params, { client }) {
-          const requestedVersion = protocolState.supportedProtocolVersions.has(params.protocolVersion)
-            ? params.protocolVersion
-            : protocolState.latestProtocolVersion
-          if (requestedVersion !== params.protocolVersion) {
-            params = {
-              ...params,
-              protocolVersion: requestedVersion
+      const currentLogLevel = yield* CurrentLogLevel
+      const contextMap = new Map<string, unknown>()
+
+      for (const protocol of options.protocolRegistry.protocols) {
+        const selectedProtocol = protocol
+        const wireHandlers = ClientRpcs.of({
+          // Requests
+          ping: () => Effect.succeed({}),
+          initialize(params, { client }) {
+            const capabilities: Types.DeepMutable<typeof ServerCapabilities.Type> = {
+              completions: {},
+              logging: {}
             }
-          }
-          const capabilities: Types.DeepMutable<typeof ServerCapabilities.Type> = {
-            completions: {}
-          }
-          if (server.tools.length > 0) {
-            capabilities.tools = { listChanged: true }
-          }
-          if (server.resources.length > 0 || server.resourceTemplates.length > 0) {
-            capabilities.resources = {
-              listChanged: true,
-              subscribe: false
+            if (server.tools.length > 0) {
+              capabilities.tools = { listChanged: true }
             }
-          }
-          if (server.prompts.length > 0) {
-            capabilities.prompts = { listChanged: true }
-          }
-          if (serverInfo.extensions) {
-            capabilities.extensions = serverInfo.extensions as any
-          }
-          return Effect.withFiber((fiber) => {
-            const httpRequest = Context.getOrUndefined(fiber.context, HttpServerRequest.HttpServerRequest)
-            if (httpRequest) {
-              const sessionId = crypto.randomUUID()
-              clientSessions.set(sessionId, params)
-              appendPreResponseHandlerUnsafe(httpRequest, (_req, res) =>
-                Effect.succeed(HttpServerResponse.setHeaders(res, {
-                  [MCP_SESSION_ID_HEADER]: sessionId,
-                  [MCP_PROTOCOL_VERSION_HEADER]: requestedVersion
-                })))
-            } else {
-              clientSessions.set(String(client.id), params)
+            if (server.resources.length > 0 || server.resourceTemplates.length > 0) {
+              capabilities.resources = {
+                listChanged: true,
+                subscribe: true
+              }
             }
-            return Effect.succeed({
-              capabilities,
-              serverInfo,
-              protocolVersion: requestedVersion
-            })
-          })
-        },
-        "completion/complete": (r) =>
-          server.completion(r).pipe(
-            Effect.provideService(CurrentLogLevel, currentLogLevel)
-          ),
-        "logging/setLevel": ({ level }) =>
-          Effect.sync(() => {
-            switch (level) {
-              case "notice":
-              case "info":
-                currentLogLevel = "Info"
-                break
-              case "error":
-                currentLogLevel = "Error"
-                break
-              case "debug":
-                currentLogLevel = "Debug"
-                break
-              case "warning":
-                currentLogLevel = "Warn"
-                break
-              case "critical":
-              case "alert":
-              case "emergency":
-                currentLogLevel = "Fatal"
-                break
+            if (server.prompts.length > 0) {
+              capabilities.prompts = { listChanged: true }
             }
-          }),
-        "prompts/get": (r) =>
-          server.getPromptResult(r).pipe(
-            Effect.provideService(CurrentLogLevel, currentLogLevel)
-          ),
-        "prompts/list": (_, { client, headers }) =>
-          Effect.sync(() => {
-            const initialized = getInitializedClient(clientSessions, client.id, headers)
-            return new ListPromptsResult({ prompts: filterByClient(initialized, server.prompts, "prompt") })
-          }),
-        "resources/list": (_, { client, headers }) =>
-          Effect.sync(() => {
-            const initialized = getInitializedClient(clientSessions, client.id, headers)
-            return new ListResourcesResult({ resources: filterByClient(initialized, server.resources, "resource") })
-          }),
-        "resources/read": ({ uri }) =>
-          server.findResource(uri).pipe(
-            Effect.provideService(CurrentLogLevel, currentLogLevel)
-          ),
-        "resources/subscribe": () =>
-          InternalError.notImplemented,
-        "resources/unsubscribe": () =>
-          InternalError.notImplemented,
-        "resources/templates/list": (_, { client, headers }) =>
-          Effect.sync(() => {
-            const initialized = getInitializedClient(clientSessions, client.id, headers)
-            return new ListResourceTemplatesResult({
-              resourceTemplates: filterByClient(initialized, server.resourceTemplates, "template")
-            })
-          }),
-        "tools/call": (r) =>
-          server.callTool(r).pipe(
-            Effect.provideService(CurrentLogLevel, currentLogLevel)
-          ),
-        "tools/list": (_, { client, headers }) =>
-          Effect.sync(() => {
-            const initialized = getInitializedClient(clientSessions, client.id, headers)
-            return new ListToolsResult({
-              tools: filterByClient(initialized, server.tools, "tool")
+            if (serverInfo.extensions) {
+              capabilities.extensions = serverInfo.extensions as any
+            }
+            return Effect.withFiber((fiber) => {
+              const httpRequest = Context.getOrUndefined(fiber.context, HttpServerRequest.HttpServerRequest)
+              if (httpRequest !== undefined && capabilities.resources !== undefined) {
+                capabilities.resources.subscribe = false
+              }
+              const session: Session = {
+                initializePayload: params,
+                protocol: selectedProtocol,
+                resourceSubscriptions: capabilities.resources?.subscribe === true ? new Set() : undefined,
+                logLevel: { _tag: "Effect", level: currentLogLevel }
+              }
+              if (httpRequest) {
+                const sessionId = crypto.randomUUID()
+                options.sessions.bySessionId.set(sessionId, session)
+                appendPreResponseHandlerUnsafe(httpRequest, (_req, res) =>
+                  Effect.succeed(HttpServerResponse.setHeaders(res, {
+                    [MCP_SESSION_ID_HEADER]: sessionId,
+                    [MCP_PROTOCOL_VERSION_HEADER]: selectedProtocol.protocolVersion
+                  })))
+              } else {
+                options.sessions.byClientId.set(client.id, session)
+              }
+              return Effect.succeed({
+                capabilities,
+                serverInfo,
+                protocolVersion: selectedProtocol.protocolVersion
+              })
             })
-          }),
+          },
+          "completion/complete": (r) =>
+            server.completion(r),
+          "logging/setLevel": ({ level }, { client, headers }) =>
+            Effect.sync(() => {
+              const session = getClientSession(options.sessions, client.id, headers)
+              if (session) {
+                session.logLevel = { _tag: "Mcp", level }
+              }
+              return {}
+            }),
+          "prompts/get": (r) =>
+            server.getPromptResult(r),
+          "prompts/list": (_, { client, headers }) =>
+            Effect.sync(() => {
+              const initialized = getClientSession(options.sessions, client.id, headers)?.initializePayload
+              return new ListPromptsResult({ prompts: filterByClient(initialized, server.prompts, "prompt") })
+            }),
+          "resources/list": (_, { client, headers }) =>
+            Effect.sync(() => {
+              const initialized = getClientSession(options.sessions, client.id, headers)?.initializePayload
+              return new ListResourcesResult({ resources: filterByClient(initialized, server.resources, "resource") })
+            }),
+          "resources/read": ({ uri }) => server.findResource(uri),
+          "resources/subscribe": ({ uri }, { client, headers }) =>
+            Effect.gen(function*() {
+              const subscriptions = getClientSession(
+                options.sessions,
+                client.id,
+                headers
+              )?.resourceSubscriptions
+              if (subscriptions === undefined) {
+                return yield* new MethodNotFound({
+                  message: "Resource subscriptions are not supported"
+                })
+              }
+              subscriptions.add(uri)
+              return {}
+            }),
+          "resources/unsubscribe": ({ uri }, { client, headers }) =>
+            Effect.gen(function*() {
+              const subscriptions = getClientSession(
+                options.sessions,
+                client.id,
+                headers
+              )?.resourceSubscriptions
+              if (subscriptions === undefined) {
+                return yield* new MethodNotFound({
+                  message: "Resource subscriptions are not supported"
+                })
+              }
+              subscriptions.delete(uri)
+              return {}
+            }),
+          "resources/templates/list": (_, { client, headers }) =>
+            Effect.sync(() => {
+              const initialized = getClientSession(options.sessions, client.id, headers)?.initializePayload
+              return new ListResourceTemplatesResult({
+                resourceTemplates: filterByClient(initialized, server.resourceTemplates, "template")
+              })
+            }),
+          "tools/call": (r) => server.callTool(r),
+          "tools/list": (_, { client, headers }) =>
+            Effect.sync(() => {
+              const initialized = getClientSession(options.sessions, client.id, headers)?.initializePayload
+              return new ListToolsResult({
+                tools: filterByClient(initialized, server.tools, "tool")
+              })
+            }),
 
-        // Notifications
-        "notifications/cancelled": (_) => Effect.void,
-        "notifications/initialized": (_, { client }) =>
-          Effect.sync(() => {
-            server.initializedClients.add(client.id)
-          }),
-        "notifications/progress": (_) => Effect.void,
-        "notifications/roots/list_changed": (_) => Effect.void
-      })
+          // Notifications
+          "notifications/cancelled": (_) => Effect.void,
+          "notifications/initialized": (_, { client, headers }) =>
+            Effect.sync(() => {
+              server.initializedClients.add(client.id)
+              const session = getClientSession(options.sessions, client.id, headers)
+              if (session) {
+                options.sessions.byClientId.set(client.id, session)
+              }
+            }),
+          "notifications/progress": (_) => Effect.void,
+          "notifications/roots/list_changed": (_) => Effect.void
+        })
+        yield* addProtocolHandlers(
+          options.protocolRegistry,
+          selectedProtocol,
+          selectedProtocol.clientRpcs,
+          wireHandlers,
+          contextMap
+        )
+      }
+      return Context.makeUnsafe(contextMap)
     })
   )
 
+const addProtocolHandlers = Effect.fnUntraced(function*<
+  ClientRpcs extends Rpc.Any
+>(
+  registry: McpProtocolRegistry.ProtocolRegistry<McpProtocol.ProtocolAdapter>,
+  protocol: McpProtocol.ProtocolAdapter,
+  clientRpcs: RpcGroup.RpcGroup<ClientRpcs>,
+  handlers: RpcGroup.HandlersFrom<ClientRpcs>,
+  contextMap: Map<string, unknown>
+) {
+  const handlerContext = yield* clientRpcs.toHandlers(handlers)
+  for (const rpcDefinition of clientRpcs.requests.values()) {
+    const routed = registry.routeClientRequest(protocol, {
+      _tag: "Request",
+      id: 0,
+      tag: rpcDefinition._tag,
+      payload: undefined,
+      headers: []
+    })
+    const namespacedRpc = registry.clientRpcs.requests.get(routed.tag)
+    const handler = handlerContext.mapUnsafe.get(rpcDefinition.key)
+    if (namespacedRpc === undefined || handler === undefined) {
+      return yield* Effect.die(`MCP handler registration invariant failed for ${routed.tag}`)
+    }
+    contextMap.set(namespacedRpc.key, handler)
+  }
+})
+
 const resolveResourceContent = (
   uri: string,
   content: typeof ReadResourceResult.Type | string | Uint8Array
@@ -1627,14 +2193,74 @@ const filterByClient = <
   return out
 }
 
-const getInitializedClient = (
-  sessions: Map<string, typeof Initialize.payloadSchema.Type>,
+const getClientSession = (
+  sessions: Sessions,
   clientId: number,
   headers: Headers.Headers
 ) => {
   const sessionId = headers[MCP_SESSION_ID_HEADER]
   if (sessionId === undefined) {
-    return sessions.get(String(clientId))
+    return sessions.byClientId.get(clientId)
   }
-  return sessions.get(sessionId)
+  return sessions.bySessionId.get(sessionId)
 }
+
+const mcpLogLevels: Record<LoggingLevel, {
+  readonly effect: LogLevel.LogLevel
+  readonly order: number
+}> = {
+  debug: { effect: "Debug", order: 0 },
+  info: { effect: "Info", order: 1 },
+  notice: { effect: "Info", order: 2 },
+  warning: { effect: "Warn", order: 3 },
+  error: { effect: "Error", order: 4 },
+  critical: { effect: "Fatal", order: 5 },
+  alert: { effect: "Fatal", order: 6 },
+  emergency: { effect: "Fatal", order: 7 }
+}
+
+const effectLogLevel = (logLevel: SessionLogLevel | undefined): LogLevel.LogLevel =>
+  logLevel?._tag === "Mcp" ? mcpLogLevels[logLevel.level].effect : logLevel?.level ?? "Info"
+
+const isMcpLogLevelEnabled = (
+  level: LoggingLevel,
+  minimum: SessionLogLevel | undefined
+): boolean =>
+  minimum?._tag === "Mcp"
+    ? mcpLogLevels[level].order >= mcpLogLevels[minimum.level].order
+    : LogLevel.isGreaterThanOrEqualTo(mcpLogLevels[level].effect, minimum?.level ?? "Info")
+
+const getOfferedProtocolVersion = (payload: unknown): string =>
+  typeof payload === "object" &&
+    payload !== null &&
+    "protocolVersion" in payload &&
+    typeof payload.protocolVersion === "string"
+    ? payload.protocolVersion
+    : ""
+
+const protocolForInternalTag = (
+  registry: McpProtocolRegistry.ProtocolRegistry<McpProtocol.ProtocolAdapter>,
+  tag: string
+): McpProtocol.ProtocolAdapter => {
+  for (const protocol of registry.protocols) {
+    const routed = registry.routeClientRequest(protocol, {
+      _tag: "Request",
+      id: 0,
+      tag: "",
+      payload: undefined,
+      headers: []
+    })
+    if (tag.startsWith(routed.tag)) {
+      return protocol
+    }
+  }
+  return registry.protocols[0]
+}
+
+const getProtocolForClient = (
+  clientProtocols: Map<number, McpProtocol.ProtocolAdapter>,
+  clientId: number,
+  registry: McpProtocolRegistry.ProtocolRegistry<McpProtocol.ProtocolAdapter>
+): McpProtocol.ProtocolAdapter =>
+  clientProtocols.get(clientId) ??
+    registry.protocols[0]