effect 4.0.0-beta.6 → 4.0.0-beta.62

package/src/unstable/ai/McpServer.ts

@@ -3,23 +3,28 @@
  */
 import * as Arr from "../../Array.ts"
 import * as Cause from "../../Cause.ts"
+import * as Context from "../../Context.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 Option from "../../Option.ts"
 import * as Queue from "../../Queue.ts"
 import * as RcMap from "../../RcMap.ts"
+import { CurrentLogLevel } from "../../References.ts"
 import * as Schema from "../../Schema.ts"
 import * as AST from "../../SchemaAST.ts"
-import * as ServiceMap from "../../ServiceMap.ts"
 import * as Sink from "../../Sink.ts"
 import type { Stdio } from "../../Stdio.ts"
 import * as Stream from "../../Stream.ts"
 import type * as Types from "../../Types.ts"
 import * as FindMyWay from "../http/FindMyWay.ts"
 import * as Headers from "../http/Headers.ts"
-import type * as HttpRouter from "../http/HttpRouter.ts"
-import type * as Rpc from "../rpc/Rpc.ts"
+import { appendPreResponseHandlerUnsafe } from "../http/HttpEffect.ts"
+import * as HttpRouter from "../http/HttpRouter.ts"
+import * as HttpServerRequest from "../http/HttpServerRequest.ts"
+import * as HttpServerResponse from "../http/HttpServerResponse.ts"
+import * as Rpc from "../rpc/Rpc.ts"
 import * as RpcClient from "../rpc/RpcClient.ts"
 import type * as RpcGroup from "../rpc/RpcGroup.ts"
 import * as RpcMessage from "../rpc/RpcMessage.ts"
@@ -32,6 +37,7 @@ import {
   CompleteResult,
   Elicit,
   ElicitationDeclined,
+  EnabledWhen,
   GetPromptResult,
   InternalError,
   InvalidParams,
@@ -51,10 +57,11 @@ import {
   Tool as McpTool
 } from "./McpSchema.ts"
 import type {
-  Annotations,
   CallTool,
+  ClientCapabilities,
   Complete,
   GetPrompt,
+  Initialize,
   Param,
   PromptArgument,
   PromptMessage,
@@ -68,30 +75,42 @@ import type * as Toolkit from "./Toolkit.ts"
  * @since 4.0.0
  * @category server
  */
-export class McpServer extends ServiceMap.Service<McpServer, {
+export class McpServer extends Context.Service<McpServer, {
   readonly notifications: RpcClient.RpcClient<RpcGroup.Rpcs<typeof ServerNotificationRpcs>>
   readonly notificationsQueue: Queue.Dequeue<RpcMessage.Request<any>>
   readonly initializedClients: Set<number>
 
-  readonly tools: ReadonlyArray<McpTool>
+  readonly tools: ReadonlyArray<{
+    readonly tool: McpTool
+    readonly annotations: Context.Context<never>
+  }>
   readonly addTool: (options: {
     readonly tool: McpTool
+    readonly annotations: Context.Context<never>
     readonly handle: (payload: any) => Effect.Effect<CallToolResult, never, McpServerClient>
   }) => Effect.Effect<void>
   readonly callTool: (
     requests: typeof CallTool.payloadSchema.Type
   ) => Effect.Effect<CallToolResult, InternalError | InvalidParams, McpServerClient>
 
-  readonly resources: ReadonlyArray<Resource>
-  readonly addResource: (
-    resource: Resource,
-    handle: Effect.Effect<typeof ReadResourceResult.Type, InternalError, McpServerClient>
-  ) => Effect.Effect<void>
+  readonly resources: ReadonlyArray<{
+    readonly resource: Resource
+    readonly annotations: Context.Context<never>
+  }>
+  readonly addResource: (options: {
+    readonly resource: Resource
+    readonly annotations: Context.Context<never>
+    readonly handle: Effect.Effect<typeof ReadResourceResult.Type, InternalError, McpServerClient>
+  }) => Effect.Effect<void>
 
-  readonly resourceTemplates: ReadonlyArray<ResourceTemplate>
+  readonly resourceTemplates: ReadonlyArray<{
+    readonly template: ResourceTemplate
+    readonly annotations: Context.Context<never>
+  }>
   readonly addResourceTemplate: (
     options: {
       readonly template: ResourceTemplate
+      readonly annotations: Context.Context<never>
       readonly routerPath: string
       readonly completions: Record<string, (input: string) => Effect.Effect<CompleteResult, InternalError>>
       readonly handle: (
@@ -105,9 +124,13 @@ export class McpServer extends ServiceMap.Service<McpServer, {
     uri: string
   ) => Effect.Effect<typeof ReadResourceResult.Type, InvalidParams | InternalError, McpServerClient>
 
-  readonly prompts: ReadonlyArray<Prompt>
+  readonly prompts: ReadonlyArray<{
+    readonly prompt: Prompt
+    readonly annotations: Context.Context<never>
+  }>
   readonly addPrompt: (options: {
     readonly prompt: Prompt
+    readonly annotations: Context.Context<never>
     readonly completions: Record<
       string,
       (input: string) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
@@ -140,11 +163,23 @@ export class McpServer extends ServiceMap.Service<McpServer, {
         readonly effect: Effect.Effect<typeof ReadResourceResult.Type, InternalError, McpServerClient>
       }
     >()
-    const tools = Arr.empty<McpTool>()
+    const tools = Arr.empty<{
+      readonly tool: McpTool
+      readonly annotations: Context.Context<never>
+    }>()
     const toolMap = new Map<string, (payload: any) => Effect.Effect<CallToolResult, InternalError, McpServerClient>>()
-    const resources: Array<Resource> = []
-    const resourceTemplates: Array<ResourceTemplate> = []
-    const prompts: Array<Prompt> = []
+    const resources: Array<{
+      readonly resource: Resource
+      readonly annotations: Context.Context<never>
+    }> = []
+    const resourceTemplates: Array<{
+      readonly template: ResourceTemplate
+      readonly annotations: Context.Context<never>
+    }> = []
+    const prompts: Array<{
+      readonly prompt: Prompt
+      readonly annotations: Context.Context<never>
+    }> = []
     const promptMap = new Map<
       string,
       (params: Record<string, string>) => Effect.Effect<GetPromptResult, InternalError | InvalidParams, McpServerClient>
@@ -188,13 +223,13 @@ export class McpServer extends ServiceMap.Service<McpServer, {
     return McpServer.of({
       notifications: notifications.client,
       notificationsQueue,
-      initializedClients: new Set<number>(),
+      initializedClients: new Set(),
       get tools() {
         return tools
       },
       addTool: (options) =>
         Effect.suspend(() => {
-          tools.push(options.tool)
+          tools.push(options)
           toolMap.set(options.tool.name, options.handle)
           return notifications.client["notifications/tools/list_changed"]({})
         }),
@@ -212,15 +247,15 @@ export class McpServer extends ServiceMap.Service<McpServer, {
       get resourceTemplates() {
         return resourceTemplates
       },
-      addResource: (resource, effect) =>
+      addResource: (options) =>
         Effect.suspend(() => {
-          resources.push(resource)
-          matcher.add(resource.uri, { _tag: "Resource", effect })
+          resources.push(options)
+          matcher.add(options.resource.uri, { _tag: "Resource", effect: options.handle })
           return notifications.client["notifications/resources/list_changed"]({})
         }),
-      addResourceTemplate: ({ completions, handle, routerPath, template }) =>
+      addResourceTemplate: ({ annotations, completions, handle, routerPath, template }) =>
         Effect.suspend(() => {
-          resourceTemplates.push(template)
+          resourceTemplates.push({ template, annotations })
           matcher.add(routerPath, { _tag: "ResourceTemplate", handle })
           for (const [param, handle] of Object.entries(completions)) {
             completionsMap.set(`ref/resource/${template.uriTemplate}/${param}`, handle)
@@ -246,7 +281,7 @@ export class McpServer extends ServiceMap.Service<McpServer, {
       },
       addPrompt: (options) =>
         Effect.suspend(() => {
-          prompts.push(options.prompt)
+          prompts.push(options)
           promptMap.set(options.prompt.name, options.handle)
           for (const [param, handle] of Object.entries(options.completions)) {
             completionsMap.set(`ref/prompt/${options.prompt.name}/${param}`, handle)
@@ -284,14 +319,18 @@ const SUPPORTED_PROTOCOL_VERSIONS = [
   "2024-11-05",
   "2024-10-07"
 ]
+const mcpSessionIdHeader = "mcp-session-id"
+const mcpProtocolVersionHeader = "mcp-protocol-version"
 
 /**
  * @since 4.0.0
  * @category constructors
  */
-export const run: (
-  options: { readonly name: string; readonly version: string }
-) => Effect.Effect<
+export const run: (options: {
+  readonly name: string
+  readonly version: string
+  readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
+}) => Effect.Effect<
   never,
   never,
   McpServer | RpcServer.Protocol
@@ -300,8 +339,10 @@ export const run: (
   readonly version: string
 }) {
   const protocol = yield* RpcServer.Protocol
-  const handlers = yield* Layer.build(layerHandlers(options))
   const server = yield* McpServer
+  const isHttp = Option.isSome(yield* Effect.serviceOption(HttpRouter.HttpRouter))
+  const clientSessions = new Map<string, typeof Initialize.payloadSchema.Type>()
+  const handlers = yield* Layer.build(layerHandlers(options, { clientSessions }))
 
   const clients = yield* RcMap.make({
     lookup: Effect.fnUntraced(function*(clientId: number) {
@@ -312,9 +353,11 @@ export const run: (
         Effect.provideServiceEffect(
           RpcClient.Protocol,
           RpcClient.Protocol.make(Effect.fnUntraced(function*(writeResponse) {
-            write = writeResponse
+            let cid = 0
+            write = (message) => writeResponse(cid, message)
             return {
-              send(request, _transferables) {
+              send(id, request, _transferables) {
+                cid = id
                 return protocol.send(clientId, {
                   ...request,
                   headers: undefined,
@@ -336,18 +379,32 @@ export const run: (
     idleTimeToLive: 10000
   })
 
-  const clientMiddleware = McpServerClientMiddleware.of((effect, { clientId }) =>
-    Effect.provideService(
+  const clientMiddleware = McpServerClientMiddleware.of((effect, { client, headers, rpc }) => {
+    const initializePayload = getInitializedClient(clientSessions, client.id, headers)
+    const isInitialize = rpc._tag === "initialize"
+    if (!isInitialize && !initializePayload) {
+      const fiber = Fiber.getCurrent()!
+      const httpRequest = Context.getOrUndefined(fiber.context, HttpServerRequest.HttpServerRequest)
+      if (httpRequest) {
+        appendPreResponseHandlerUnsafe(
+          httpRequest,
+          () => Effect.succeed(HttpServerResponse.empty({ status: 404 }))
+        )
+      }
+      return Effect.die(new Error(`Mcp-Session-Id does not exist`))
+    }
+    return Effect.provideService(
       effect,
       McpServerClient,
       McpServerClient.of({
-        clientId,
-        getClient: RcMap.get(clients, clientId).pipe(
+        clientId: client.id,
+        initializePayload: initializePayload!,
+        getClient: RcMap.get(clients, client.id).pipe(
           Effect.map(({ client }) => client)
         )
       })
     )
-  )
+  })
 
   const patchedProtocol = RpcServer.Protocol.of({
     ...protocol,
@@ -358,6 +415,17 @@ export const run: (
           | RpcMessage.FromClientEncoded
         switch (request._tag) {
           case "Request": {
+            if (isHttp) {
+              const fiber = Fiber.getCurrent()!
+              const httpRequest = Context.getUnsafe(fiber.context, HttpServerRequest.HttpServerRequest)
+              const client = getInitializedClient(clientSessions, clientId, httpRequest.headers)
+              if (client) {
+                appendPreResponseHandlerUnsafe(httpRequest, (_, res) =>
+                  Effect.succeed(
+                    HttpServerResponse.setHeader(res, mcpProtocolVersionHeader, client.protocolVersion)
+                  ))
+              }
+            }
             const rpc = ClientNotificationRpcs.requests.get(request.tag)
             if (rpc) {
               if (request.tag === "notifications/cancelled") {
@@ -371,9 +439,9 @@ export const run: (
                 ? handler.handler(request.payload, {
                   rpc,
                   requestId: RpcMessage.RequestId(request.id),
-                  clientId,
+                  client: new Rpc.ServerClient(clientId),
                   headers: Headers.fromInput(request.headers)
-                }) as Effect.Effect<void>
+                }) as any as Effect.Effect<void>
                 : Effect.void
             }
             return f(clientId, request)
@@ -408,7 +476,7 @@ export const run: (
         payload: encoded
       } as any
       const clientIds = yield* patchedProtocol.clientIds
-      for (const clientId of server.initializedClients) {
+      for (const clientId of server.initializedClients.keys()) {
         if (!clientIds.has(clientId)) {
           server.initializedClients.delete(clientId)
           continue
@@ -438,6 +506,7 @@ export const run: (
 export const layer = (options: {
   readonly name: string
   readonly version: string
+  readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }): Layer.Layer<McpServer | McpServerClient, never, RpcServer.Protocol> =>
   Layer.effectDiscard(Effect.forkScoped(run(options))).pipe(
     Layer.provideMerge(McpServer.layer)
@@ -503,6 +572,7 @@ export const layer = (options: {
 export const layerStdio = (options: {
   readonly name: string
   readonly version: string
+  readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }): Layer.Layer<McpServer | McpServerClient, never, Stdio> =>
   layer(options).pipe(
     Layer.provide(RpcServer.layerProtocolStdio),
@@ -519,6 +589,7 @@ export const layerHttp = (options: {
   readonly name: string
   readonly version: string
   readonly path: HttpRouter.PathInput
+  readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }): Layer.Layer<McpServer | McpServerClient, never, HttpRouter.HttpRouter> =>
   layer(options).pipe(
     Layer.provide(RpcServer.layerProtocolHttp(options)),
@@ -546,31 +617,35 @@ export const registerToolkit: <Tools extends Record<string, Tool.Any>>(
     never,
     Exclude<Tool.HandlersFor<Tools>, McpServerClient>
   >)
-  const services = yield* Effect.services<never>()
+  const services = yield* Effect.context<never>()
   for (const tool of Object.values(built.tools)) {
+    const annotations = tool.annotations
+    const toolMeta = Context.getOrUndefined(annotations, Tool.Meta)
     const mcpTool = new McpTool({
       name: tool.name,
       description: Tool.getDescription(tool),
       inputSchema: Tool.getJsonSchema(tool),
       annotations: {
-        ...(ServiceMap.getOption(tool.annotations, Tool.Title).pipe(
+        ...(Context.getOption(tool.annotations, Tool.Title).pipe(
           Option.map((title) => ({ title })),
           Option.getOrUndefined
         )),
-        readOnlyHint: ServiceMap.get(tool.annotations, Tool.Readonly),
-        destructiveHint: ServiceMap.get(tool.annotations, Tool.Destructive),
-        idempotentHint: ServiceMap.get(tool.annotations, Tool.Idempotent),
-        openWorldHint: ServiceMap.get(tool.annotations, Tool.OpenWorld)
-      }
+        readOnlyHint: Context.get(tool.annotations, Tool.Readonly),
+        destructiveHint: Context.get(tool.annotations, Tool.Destructive),
+        idempotentHint: Context.get(tool.annotations, Tool.Idempotent),
+        openWorldHint: Context.get(tool.annotations, Tool.OpenWorld)
+      },
+      _meta: toolMeta
     })
     yield* registry.addTool({
       tool: mcpTool,
+      annotations,
       handle(payload) {
         return built.handle(tool.name as any, payload).pipe(
           Stream.unwrap,
           Stream.run(Sink.last()),
           Effect.flatMap(Effect.fromOption),
-          Effect.provideServices(services as ServiceMap.ServiceMap<any>),
+          Effect.provideContext(services as Context.Context<any>),
           Effect.matchCause({
             onFailure: (cause) =>
               new CallToolResult({
@@ -659,6 +734,7 @@ export const registerResource: {
         E,
         R
       >
+      readonly annotations?: Context.Context<never> | undefined
     }
   ): Effect.Effect<void, never, Exclude<R, McpServerClient> | McpServer>
   /**
@@ -683,6 +759,7 @@ export const registerResource: {
       E,
       R
     >
+    readonly annotations?: Context.Context<never> | undefined
   }) => Effect.Effect<
     void,
     never,
@@ -699,26 +776,34 @@ export const registerResource: {
   >
 } = function() {
   if (arguments.length === 1) {
-    const options = arguments[0] as Resource & typeof Annotations.Type & {
-      readonly content: Effect.Effect<typeof ReadResourceResult.Type | string | Uint8Array>
+    const options = arguments[0] as {
+      readonly uri: string
+      readonly name: string
+      readonly description?: string | undefined
+      readonly mimeType?: string | undefined
+      readonly audience?: ReadonlyArray<"user" | "assistant"> | undefined
+      readonly priority?: number | undefined
+      readonly content: Effect.Effect<typeof ReadResourceResult.Type | string | Uint8Array, any, any>
+      readonly annotations?: Context.Context<never> | undefined
     }
     return Effect.gen(function*() {
-      const services = yield* Effect.services<any>()
+      const services = yield* Effect.context<any>()
       const registry = yield* McpServer
-      yield* registry.addResource(
-        new Resource({
+      yield* registry.addResource({
+        resource: new Resource({
           ...options,
           annotations: options
         }),
-        options.content.pipe(
-          Effect.provideServices(services),
+        handle: options.content.pipe(
+          Effect.provideContext(services),
           Effect.map((content) => resolveResourceContent(options.uri, content)),
           Effect.catchCause((cause) => {
             const prettyError = Cause.prettyErrors(cause)[0]
             return Effect.fail(new InternalError({ message: prettyError.message }))
           })
-        )
-      )
+        ),
+        annotations: options.annotations ?? Context.empty()
+      })
     })
   }
   const {
@@ -739,8 +824,9 @@ export const registerResource: {
       E,
       R
     >
+    readonly annotations?: Context.Context<never> | undefined
   }) {
-    const services = yield* Effect.services<any>()
+    const services = yield* Effect.context<any>()
     const registry = yield* McpServer
     const decode = Schema.decodeUnknownEffect(schema)
     const template = new ResourceTemplate({
@@ -765,7 +851,7 @@ export const registerResource: {
             const prettyError = Cause.prettyErrors(cause)[0]
             return Effect.fail(new InternalError({ message: prettyError.message }))
           }),
-          Effect.provideServices(services)
+          Effect.provideContext(services)
         )
       completions[param] = handler
     }
@@ -773,6 +859,7 @@ export const registerResource: {
       template,
       routerPath,
       completions,
+      annotations: options.annotations ?? Context.empty(),
       handle: (uri, params) =>
         decode(params).pipe(
           Effect.mapError((error) => new InvalidParams({ message: error.message })),
@@ -785,7 +872,7 @@ export const registerResource: {
               })
             )
           ),
-          Effect.provideServices(services)
+          Effect.provideContext(services)
         )
     })
   })
@@ -885,6 +972,7 @@ export const registerPrompt = <
     readonly parameters?: Params | undefined
     readonly completion?: ValidateCompletions<Completions, Extract<keyof Params, string>> | undefined
     readonly content: (params: Params) => Effect.Effect<Array<typeof PromptMessage.Type> | string, E, R>
+    readonly annotations?: Context.Context<never> | undefined
   }
 ): Effect.Effect<void, never, Exclude<Schema.Struct.DecodingServices<Params> | R, McpServerClient> | McpServer> => {
   const args = Arr.empty<typeof PromptArgument.Type>()
@@ -907,7 +995,7 @@ export const registerPrompt = <
   const completion: Record<string, (input: string) => Effect.Effect<any>> = options.completion ?? {}
   return Effect.gen(function*() {
     const registry = yield* McpServer
-    const services = yield* Effect.services<Exclude<R | Schema.Struct.DecodingServices<Params>, McpServerClient>>()
+    const services = yield* Effect.context<Exclude<R | Schema.Struct.DecodingServices<Params>, McpServerClient>>()
     const completions: Record<
       string,
       (input: string) => Effect.Effect<CompleteResult, InternalError, McpServerClient>
@@ -935,6 +1023,7 @@ export const registerPrompt = <
     yield* registry.addPrompt({
       prompt,
       completions,
+      annotations: options.annotations ?? Context.empty(),
       handle: (params) =>
         decode(params).pipe(
           Effect.mapError((error) => new InvalidParams({ message: error.message })),
@@ -943,7 +1032,7 @@ export const registerPrompt = <
             messages = typeof messages === "string" ?
               [{
                 role: "user",
-                content: TextContent.makeUnsafe({ text: messages })
+                content: TextContent.make({ text: messages })
               }] :
               messages
             return new GetPromptResult({ messages, description: prompt.description })
@@ -952,7 +1041,7 @@ export const registerPrompt = <
             const prettyError = Cause.prettyErrors(cause)[0]
             return Effect.fail(new InternalError({ message: prettyError.message }))
           }),
-          Effect.provideServices(services as ServiceMap.ServiceMap<unknown>)
+          Effect.provideContext(services as Context.Context<unknown>)
         )
     })
   })
@@ -980,6 +1069,7 @@ export const prompt = <
     readonly content: (
       params: Schema.Struct.Type<Params>
     ) => Effect.Effect<Array<typeof PromptMessage.Type> | string, E, R>
+    readonly annotations?: Context.Context<never> | undefined
   }
 ): Layer.Layer<never, never, Exclude<Schema.Struct.DecodingServices<Params> | R, McpServerClient>> =>
   Layer.effectDiscard(registerPrompt(options)).pipe(
@@ -992,21 +1082,21 @@ export const prompt = <
  * @since 4.0.0
  * @category elicitation
  */
-export const elicit: <S extends Schema.Codec<any, Record<string, unknown>, any, any>>(options: {
+export const elicit: <S extends Schema.Encoder<Record<string, unknown>, unknown>>(options: {
   readonly message: string
   readonly schema: S
 }) => Effect.Effect<
   S["Type"],
   ElicitationDeclined,
   McpServerClient | S["DecodingServices"]
-> = Effect.fnUntraced(function*<S extends Schema.Codec<any, Record<string, unknown>, any, any>>(options: {
+> = Effect.fnUntraced(function*<S extends Schema.Encoder<Record<string, unknown>, unknown>>(options: {
   readonly message: string
   readonly schema: S
 }) {
   const { getClient } = yield* McpServerClient
   const client = yield* getClient
   const schema = options.schema
-  const request = Elicit.payloadSchema.makeUnsafe({
+  const request = Elicit.payloadSchema.make({
     message: options.message,
     requestedSchema: Tool.getJsonSchemaFromSchema(schema)
   })
@@ -1023,6 +1113,18 @@ export const elicit: <S extends Schema.Codec<any, Record<string, unknown>, any,
   }
 }, Effect.scoped)
 
+/**
+ * Access the current client's capabilities.
+ *
+ * @since 4.0.0
+ * @category capabilities
+ */
+export const clientCapabilities: Effect.Effect<
+  ClientCapabilities,
+  never,
+  McpServerClient
+> = McpServerClient.useSync((_) => _.initializePayload.capabilities)
+
 // -----------------------------------------------------------------------------
 // Internal
 // -----------------------------------------------------------------------------
@@ -1054,7 +1156,8 @@ const compileUriTemplate = (segments: TemplateStringsArray, ...schemas: Readonly
       const key = String(i)
       arr.push(toCodecStringTree)
       routerPath += `:${key}${segment.replace(":", "::")}`
-      const paramName = isParam(toCodecStringTree) ? toCodecStringTree.name : `param${key}`
+      const schema = schemas[i]
+      const paramName = isParam(schema) ? (schema as Param<string, Schema.Top>).name : `param${key}`
       params[paramName] = toCodecStringTree
       uriPath += `{${paramName}}${segment}`
     }
@@ -1071,16 +1174,28 @@ const compileUriTemplate = (segments: TemplateStringsArray, ...schemas: Readonly
 const layerHandlers = (serverInfo: {
   readonly name: string
   readonly version: string
+  readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
+}, options: {
+  readonly clientSessions: Map<string, typeof Initialize.payloadSchema.Type>
 }) =>
   ClientRpcs.toLayer(
     Effect.gen(function*() {
       const server = yield* McpServer
+      let currentLogLevel = yield* CurrentLogLevel
 
       return ClientRpcs.of({
         // Requests
         ping: () => Effect.succeed({}),
-        initialize(params, { clientId }) {
-          const requestedVersion = params.protocolVersion
+        initialize(params, { client }) {
+          const requestedVersion = SUPPORTED_PROTOCOL_VERSIONS.includes(params.protocolVersion)
+            ? params.protocolVersion
+            : LATEST_PROTOCOL_VERSION
+          if (requestedVersion !== params.protocolVersion) {
+            params = {
+              ...params,
+              protocolVersion: requestedVersion
+            }
+          }
           const capabilities: Types.DeepMutable<typeof ServerCapabilities.Type> = {
             completions: {}
           }
@@ -1096,27 +1211,96 @@ const layerHandlers = (serverInfo: {
           if (server.prompts.length > 0) {
             capabilities.prompts = { listChanged: true }
           }
-          server.initializedClients.add(clientId)
-          return Effect.succeed({
-            capabilities,
-            serverInfo,
-            protocolVersion: SUPPORTED_PROTOCOL_VERSIONS.includes(requestedVersion)
-              ? requestedVersion
-              : LATEST_PROTOCOL_VERSION
+          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()
+              options.clientSessions.set(sessionId, params)
+              appendPreResponseHandlerUnsafe(httpRequest, (_req, res) =>
+                Effect.succeed(HttpServerResponse.setHeaders(res, {
+                  [mcpSessionIdHeader]: sessionId,
+                  [mcpProtocolVersionHeader]: requestedVersion
+                })))
+            } else {
+              options.clientSessions.set(String(client.id), params)
+            }
+            return Effect.succeed({
+              capabilities,
+              serverInfo,
+              protocolVersion: requestedVersion
+            })
           })
         },
-        "completion/complete": server.completion,
-        "logging/setLevel": () => InternalError.notImplemented.asEffect(),
-        "prompts/get": server.getPromptResult,
-        "prompts/list": () => Effect.sync(() => new ListPromptsResult({ prompts: server.prompts })),
-        "resources/list": () => Effect.sync(() => new ListResourcesResult({ resources: server.resources })),
-        "resources/read": ({ uri }) => server.findResource(uri),
-        "resources/subscribe": () => InternalError.notImplemented.asEffect(),
-        "resources/unsubscribe": () => InternalError.notImplemented.asEffect(),
-        "resources/templates/list": () =>
-          Effect.sync(() => new ListResourceTemplatesResult({ resourceTemplates: server.resourceTemplates })),
-        "tools/call": server.callTool,
-        "tools/list": () => Effect.sync(() => new ListToolsResult({ tools: server.tools })),
+        "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
+            }
+          }),
+        "prompts/get": (r) =>
+          server.getPromptResult(r).pipe(
+            Effect.provideService(CurrentLogLevel, currentLogLevel)
+          ),
+        "prompts/list": (_, { client, headers }) =>
+          Effect.sync(() => {
+            const initialized = getInitializedClient(options.clientSessions, client.id, headers)
+            return new ListPromptsResult({ prompts: filterByClient(initialized, server.prompts, "prompt") })
+          }),
+        "resources/list": (_, { client, headers }) =>
+          Effect.sync(() => {
+            const initialized = getInitializedClient(options.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.asEffect(),
+        "resources/unsubscribe": () =>
+          InternalError.notImplemented.asEffect(),
+        "resources/templates/list": (_, { client, headers }) =>
+          Effect.sync(() => {
+            const initialized = getInitializedClient(options.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(options.clientSessions, client.id, headers)
+            return new ListToolsResult({
+              tools: filterByClient(initialized, server.tools, "tool")
+            })
+          }),
 
         // Notifications
         "notifications/cancelled": (_) => Effect.void,
@@ -1148,3 +1332,39 @@ const resolveResourceContent = (
   }
   return content
 }
+
+const filterByClient = <
+  A extends {
+    readonly annotations: Context.Context<never>
+  },
+  P extends keyof A
+>(
+  client: typeof Initialize.payloadSchema.Type | undefined,
+  items: ReadonlyArray<A>,
+  prop: P
+): Array<A[P]> => {
+  if (!client) {
+    return items.map((item) => item[prop])
+  }
+  const out = Arr.empty<A[P]>()
+  for (let i = 0; i < items.length; i++) {
+    const item = items[i]
+    const enabledWhen = Context.getOrUndefined(item.annotations, EnabledWhen)
+    if (!enabledWhen || enabledWhen(client)) {
+      out.push(item[prop])
+    }
+  }
+  return out
+}
+
+const getInitializedClient = (
+  sessions: Map<string, typeof Initialize.payloadSchema.Type>,
+  clientId: number,
+  headers: Headers.Headers
+) => {
+  const sessionId = headers[mcpSessionIdHeader]
+  if (sessionId === undefined) {
+    return sessions.get(String(clientId))
+  }
+  return sessions.get(sessionId)
+}