effect 4.0.0-beta.100 → 4.0.0-beta.102

package/src/unstable/ai/McpServer.ts

@@ -39,6 +39,7 @@ import type * as RpcGroup from "../rpc/RpcGroup.ts"
 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 {
   CallToolResult,
   ClientNotificationRpcs,
@@ -340,8 +341,22 @@ const SUPPORTED_PROTOCOL_VERSIONS = [
   "2024-11-05",
   "2024-10-07"
 ]
-const mcpSessionIdHeader = "mcp-session-id"
-const mcpProtocolVersionHeader = "mcp-protocol-version"
+const MCP_SESSION_ID_HEADER = "mcp-session-id"
+const MCP_PROTOCOL_VERSION_HEADER = "mcp-protocol-version"
+
+class McpProtocolState extends Context.Service<McpProtocolState, {
+  readonly clientSessions: Map<string, typeof Initialize.payloadSchema.Type>
+  readonly latestProtocolVersion: string
+  readonly supportedProtocolVersions: ReadonlySet<string>
+}>()("effect/ai/McpServer/McpProtocolState") {}
+
+const makeMcpProtocolState = (): McpProtocolState["Service"] => ({
+  clientSessions: new Map(),
+  latestProtocolVersion: LATEST_PROTOCOL_VERSION,
+  supportedProtocolVersions: new Set(SUPPORTED_PROTOCOL_VERSIONS)
+})
+
+const layerMcpProtocolState = Layer.sync(McpProtocolState, makeMcpProtocolState)
 
 /**
  * Runs an MCP server over the current `RpcServer.Protocol`.
@@ -366,12 +381,19 @@ export const run: (options: {
 > = Effect.fnUntraced(function*(options: {
   readonly name: string
   readonly version: string
+  readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }) {
+  const protocolState = Option.getOrElse(
+    yield* Effect.serviceOption(McpProtocolState),
+    makeMcpProtocolState
+  )
   const protocol = yield* RpcServer.Protocol
   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 clientSessions = protocolState.clientSessions
+  const handlers = yield* Layer.build(layerHandlers(options)).pipe(
+    Effect.provideService(McpProtocolState, protocolState)
+  )
 
   const clients = yield* RcMap.make({
     lookup: Effect.fnUntraced(function*(clientId: number) {
@@ -439,24 +461,49 @@ export const run: (options: {
     ...protocol,
     run: (f) =>
       protocol.run((clientId, request_) => {
-        const request = request_ as any as
+        const fiber = Fiber.getCurrent()!
+        const request = request_ as unknown as
           | RpcMessage.FromServerEncoded
           | RpcMessage.FromClientEncoded
+        const httpRequest = isHttp
+          ? Context.getOrUndefined(fiber.context, HttpServerRequest.HttpServerRequest)
+          : undefined
+        if (httpRequest !== undefined && request._tag !== "Eof") {
+          appendPreResponseHandlerUnsafe(httpRequest, (_, response) =>
+            Effect.succeed(
+              response.status === 200 &&
+                response.body._tag === "Uint8Array" &&
+                response.body.contentLength === 0
+                ? HttpServerResponse.empty({
+                  headers: Headers.remove(response.headers, "content-type"),
+                  status: 202
+                })
+                : response
+            ))
+        }
         switch (request._tag) {
           case "Request": {
-            if (isHttp) {
-              const fiber = Fiber.getCurrent()!
-              const httpRequest = Context.getUnsafe(fiber.context, HttpServerRequest.HttpServerRequest)
+            if (httpRequest !== undefined) {
               const client = getInitializedClient(clientSessions, clientId, httpRequest.headers)
               if (client) {
                 appendPreResponseHandlerUnsafe(httpRequest, (_, res) =>
                   Effect.succeed(
-                    HttpServerResponse.setHeader(res, mcpProtocolVersionHeader, client.protocolVersion)
+                    HttpServerResponse.setHeader(res, MCP_PROTOCOL_VERSION_HEADER, client.protocolVersion)
                   ))
               }
             }
             const rpc = ClientNotificationRpcs.requests.get(request.tag)
             if (rpc) {
+              const headers = Headers.fromInput(request.headers)
+              if (!getInitializedClient(clientSessions, clientId, headers)) {
+                if (httpRequest !== undefined) {
+                  appendPreResponseHandlerUnsafe(
+                    httpRequest,
+                    () => Effect.succeed(HttpServerResponse.empty({ status: 404 }))
+                  )
+                }
+                return Effect.void
+              }
               if (request.tag === "notifications/cancelled") {
                 return f(clientId, {
                   _tag: "Interrupt",
@@ -469,7 +516,7 @@ export const run: (options: {
                   rpc,
                   requestId: RpcMessage.RequestId(request.id),
                   client: new Rpc.ServerClient(clientId),
-                  headers: Headers.fromInput(request.headers)
+                  headers
                 }) as any as Effect.Effect<void>
                 : Effect.void
             }
@@ -562,6 +609,15 @@ export const layer = (options: {
   readonly version: string
   readonly extensions?: Record<`${string}/${string}`, unknown> | undefined
 }): Layer.Layer<McpServer | McpServerClient, never, RpcServer.Protocol> =>
+  layerWithProtocolState(options).pipe(
+    Layer.provide(layerMcpProtocolState)
+  )
+
+const layerWithProtocolState = (options: {
+  readonly name: string
+  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.provideMerge(McpServer.layer)
   )
@@ -635,8 +691,7 @@ export const layerStdio = (options: {
   )
 
 /**
- * Registers an HTTP POST JSON-RPC route at `options.path` on the current
- * `HttpRouter`.
+ * Registers a Streamable HTTP MCP endpoint at `options.path`.
  *
  * **When to use**
  *
@@ -644,8 +699,9 @@ export const layerStdio = (options: {
  *
  * **Details**
  *
- * This layer composes `layer(options)`, `RpcServer.layerProtocolHttp(options)`,
- * and `RpcSerialization.layerJsonRpc()`.
+ * POST serves JSON-RPC and accepted notification-only requests return `202`.
+ * Unsupported protocol versions return `400`; methods without MCP handlers
+ * return `405`.
  *
  * @see {@link layerStdio} for exposing the server over stdio
  * @see {@link layer} for the base MCP server layer without a transport protocol
@@ -658,11 +714,56 @@ export const layerHttp = (options: {
   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)),
+}): Layer.Layer<McpServer | McpServerClient, never, HttpRouter.HttpRouter> => {
+  const protocolState = layerMcpProtocolState
+  const methodNotAllowedResponse = HttpServerResponse.empty({
+    status: 405,
+    headers: { allow: "POST" }
+  })
+  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)
+  )
+  return Layer.merge(layerWithProtocolState(options), routes).pipe(
+    Layer.provide(layerMcpProtocolHttp(options)),
+    Layer.provide(protocolState),
     Layer.provide(RpcSerialization.layerJsonRpc())
   )
+}
+
+const layerMcpProtocolHttp = (options: {
+  readonly path: HttpRouter.PathInput
+}): Layer.Layer<
+  RpcServer.Protocol,
+  never,
+  McpProtocolState | RpcSerialization.RpcSerialization | HttpRouter.HttpRouter
+> =>
+  Layer.effect(RpcServer.Protocol)(Effect.gen(function*() {
+    const state = yield* McpProtocolState
+    const { httpEffect, protocol } = yield* RpcServer.makeProtocolWithHttpEffect
+    const router = yield* HttpRouter.HttpRouter
+    yield* router.add("POST", options.path, (request) => {
+      const protocolVersion = request.headers[MCP_PROTOCOL_VERSION_HEADER]
+      if (
+        protocolVersion !== undefined &&
+        !state.supportedProtocolVersions.has(protocolVersion)
+      ) {
+        return Effect.succeed(HttpServerResponse.empty({ status: 400 }))
+      }
+      return httpEffect
+    })
+    return protocol
+  }))
+
+const INTERNAL_TOOL_ERROR_MESSAGE = "Tool execution failed due to an internal server error."
+
+const toolErrorResult = (message: string): CallToolResult =>
+  new CallToolResult({
+    isError: true,
+    content: [{ type: "text", text: message }]
+  })
 
 /**
  * Registers a `Toolkit` with the `McpServer`.
@@ -689,6 +790,7 @@ export const registerToolkit: <Tools extends Record<string, Tool.Any>>(
   for (const tool of Object.values(built.tools)) {
     const annotations = tool.annotations
     const toolMeta = Context.getOrUndefined(annotations, Tool.Meta)
+    const isDeclaredFailure = Schema.is(tool.failureSchema)
     const mcpTool = new McpTool({
       name: tool.name,
       description: Tool.getDescription(tool),
@@ -709,32 +811,41 @@ export const registerToolkit: <Tools extends Record<string, Tool.Any>>(
       tool: mcpTool,
       annotations,
       handle(payload) {
-        return built.handle(tool.name as any, payload).pipe(
+        return built.handle(tool.name as keyof Tools, payload).pipe(
           Stream.unwrap,
           Stream.run(Sink.last()),
           Effect.flatMap(Effect.fromOption),
-          Effect.provideContext(services as Context.Context<any>),
-          Effect.matchCause({
-            onFailure: (cause) =>
-              new CallToolResult({
-                isError: true,
-                content: [{
-                  type: "text",
-                  text: Cause.pretty(cause)
-                }]
-              }),
-            onSuccess: (result: any) =>
-              new CallToolResult({
-                isError: false,
-                structuredContent: typeof result.encodedResult === "object" ? result.encodedResult : undefined,
-                content: [{
-                  type: "text",
-                  text: JSON.stringify(result.encodedResult)
-                }]
-              })
+          Effect.provideContext(
+            services as Context.Context<Tool.HandlerServices<Tools[keyof Tools]>>
+          ),
+          Effect.map((result) =>
+            new CallToolResult({
+              isError: false,
+              structuredContent: typeof result.encodedResult === "object" ? result.encodedResult : undefined,
+              content: [{
+                type: "text",
+                text: JSON.stringify(result.encodedResult)
+              }]
+            })
+          ),
+          Effect.tapCause(Effect.logError),
+          Effect.catch((error) => {
+            if (AiError.isAiError(error)) {
+              const reason = (error as AiError.AiError).reason
+              return reason._tag === "ToolParameterValidationError"
+                ? Effect.succeed(toolErrorResult(reason.message))
+                : Effect.succeed(toolErrorResult(INTERNAL_TOOL_ERROR_MESSAGE))
+            }
+            if (isDeclaredFailure(error)) {
+              const message = error instanceof Error
+                ? error.message
+                : INTERNAL_TOOL_ERROR_MESSAGE
+              return Effect.succeed(toolErrorResult(message))
+            }
+            return Effect.succeed(toolErrorResult(INTERNAL_TOOL_ERROR_MESSAGE))
           }),
-          Effect.tapCause(Effect.log)
-        ) as any
+          Effect.catchDefect(() => Effect.succeed(toolErrorResult(INTERNAL_TOOL_ERROR_MESSAGE)))
+        )
       }
     })
   }
@@ -939,7 +1050,9 @@ export const registerResource: {
       uriTemplate: uriPath,
       annotations: options!
     })
-    const completions: Record<string, (input: string) => Effect.Effect<CompleteResult, InternalError>> = {}
+    const completions: Record<string, (input: string) => 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) =>
@@ -1135,7 +1248,7 @@ export const registerPrompt = <
     const completions: Record<
       string,
       (input: string) => 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) =>
@@ -1300,7 +1413,7 @@ const makeUriMatcher = <A>() => {
 const compileUriTemplate = (segments: TemplateStringsArray, ...schemas: ReadonlyArray<Schema.Constraint>) => {
   let routerPath = segments[0].replace(":", "::")
   let uriPath = segments[0]
-  const params: Record<string, Schema.Top> = {}
+  const params: Record<string, Schema.Top> = Object.create(null)
   let pathSchema = Schema.Tuple([]) as Schema.Top
   if (schemas.length > 0) {
     const arr: Array<Schema.Top> = []
@@ -1329,21 +1442,21 @@ 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
+      const protocolState = yield* McpProtocolState
+      const clientSessions = protocolState.clientSessions
       let currentLogLevel = yield* CurrentLogLevel
 
       return ClientRpcs.of({
         // Requests
         ping: () => Effect.succeed({}),
         initialize(params, { client }) {
-          const requestedVersion = SUPPORTED_PROTOCOL_VERSIONS.includes(params.protocolVersion)
+          const requestedVersion = protocolState.supportedProtocolVersions.has(params.protocolVersion)
             ? params.protocolVersion
-            : LATEST_PROTOCOL_VERSION
+            : protocolState.latestProtocolVersion
           if (requestedVersion !== params.protocolVersion) {
             params = {
               ...params,
@@ -1372,14 +1485,14 @@ const layerHandlers = (serverInfo: {
             const httpRequest = Context.getOrUndefined(fiber.context, HttpServerRequest.HttpServerRequest)
             if (httpRequest) {
               const sessionId = crypto.randomUUID()
-              options.clientSessions.set(sessionId, params)
+              clientSessions.set(sessionId, params)
               appendPreResponseHandlerUnsafe(httpRequest, (_req, res) =>
                 Effect.succeed(HttpServerResponse.setHeaders(res, {
-                  [mcpSessionIdHeader]: sessionId,
-                  [mcpProtocolVersionHeader]: requestedVersion
+                  [MCP_SESSION_ID_HEADER]: sessionId,
+                  [MCP_PROTOCOL_VERSION_HEADER]: requestedVersion
                 })))
             } else {
-              options.clientSessions.set(String(client.id), params)
+              clientSessions.set(String(client.id), params)
             }
             return Effect.succeed({
               capabilities,
@@ -1421,12 +1534,12 @@ const layerHandlers = (serverInfo: {
           ),
         "prompts/list": (_, { client, headers }) =>
           Effect.sync(() => {
-            const initialized = getInitializedClient(options.clientSessions, client.id, headers)
+            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(options.clientSessions, client.id, headers)
+            const initialized = getInitializedClient(clientSessions, client.id, headers)
             return new ListResourcesResult({ resources: filterByClient(initialized, server.resources, "resource") })
           }),
         "resources/read": ({ uri }) =>
@@ -1439,7 +1552,7 @@ const layerHandlers = (serverInfo: {
           InternalError.notImplemented,
         "resources/templates/list": (_, { client, headers }) =>
           Effect.sync(() => {
-            const initialized = getInitializedClient(options.clientSessions, client.id, headers)
+            const initialized = getInitializedClient(clientSessions, client.id, headers)
             return new ListResourceTemplatesResult({
               resourceTemplates: filterByClient(initialized, server.resourceTemplates, "template")
             })
@@ -1450,7 +1563,7 @@ const layerHandlers = (serverInfo: {
           ),
         "tools/list": (_, { client, headers }) =>
           Effect.sync(() => {
-            const initialized = getInitializedClient(options.clientSessions, client.id, headers)
+            const initialized = getInitializedClient(clientSessions, client.id, headers)
             return new ListToolsResult({
               tools: filterByClient(initialized, server.tools, "tool")
             })
@@ -1458,7 +1571,10 @@ const layerHandlers = (serverInfo: {
 
         // Notifications
         "notifications/cancelled": (_) => Effect.void,
-        "notifications/initialized": (_) => Effect.void,
+        "notifications/initialized": (_, { client }) =>
+          Effect.sync(() => {
+            server.initializedClients.add(client.id)
+          }),
         "notifications/progress": (_) => Effect.void,
         "notifications/roots/list_changed": (_) => Effect.void
       })
@@ -1516,7 +1632,7 @@ const getInitializedClient = (
   clientId: number,
   headers: Headers.Headers
 ) => {
-  const sessionId = headers[mcpSessionIdHeader]
+  const sessionId = headers[MCP_SESSION_ID_HEADER]
   if (sessionId === undefined) {
     return sessions.get(String(clientId))
   }