wip

IMax153 · IMax153 · commit 6ffb878534b9 · 2025-06-02T13:40:16.000-04:00
diff --git a/packages/ai/ai/src/McpSchema.ts b/packages/ai/ai/src/McpSchema.ts
@@ -1403,7 +1403,6 @@ export type SuccessEncoded<Group extends RpcGroup.Any> = RpcGroup.Rpcs<
   > ? {
       readonly _tag: "Success"
       readonly id: string | number
-      readonly method: _Tag
       readonly result: _Success["Encoded"]
     }
   : never
@@ -1420,7 +1419,6 @@ export type FailureEncoded<Group extends RpcGroup.Any> = RpcGroup.Rpcs<
   > ? {
       readonly _tag: "Failure"
       readonly id: string | number
-      readonly method: _Tag
       readonly error: _Error["Encoded"]
     }
   : never
diff --git a/packages/ai/ai/src/McpServer.ts b/packages/ai/ai/src/McpServer.ts
@@ -2,23 +2,29 @@
  * @since 1.0.0
  */
 import * as Headers from "@effect/platform/Headers"
-import * as Rpc from "@effect/rpc/Rpc"
+import type * as Rpc from "@effect/rpc/Rpc"
 import type * as RpcGroup from "@effect/rpc/RpcGroup"
 import { RequestId } from "@effect/rpc/RpcMessage"
 import * as RpcServer from "@effect/rpc/RpcServer"
 import type { NonEmptyReadonlyArray } from "effect/Array"
 import * as Cause from "effect/Cause"
 import * as Context from "effect/Context"
 import * as Effect from "effect/Effect"
+import * as Either from "effect/Either"
+import * as Exit from "effect/Exit"
 import * as Layer from "effect/Layer"
 import type { ParseError } from "effect/ParseResult"
 import { TreeFormatter } from "effect/ParseResult"
 import * as Schema from "effect/Schema"
 import * as Scope from "effect/Scope"
-import type { ClientRequestRpcs, Implementation, ServerCapabilities } from "./McpSchema.js"
-import { ClientRpcs } from "./McpSchema.js"
+import type { ClientRequestRpcs, FromServerEncoded, Implementation, ServerCapabilities } from "./McpSchema.js"
+import { ClientRpcs, InvalidParams, MethodNotFound } from "./McpSchema.js"
 import * as McpTransport from "./McpTransport.js"
 
+// TODO:
+// - Finish error types
+// - Investigate encoded errors in transport.send
+
 export class McpServer extends Context.Tag("@effect/ai/McpServer")<
   McpServer,
   {}
@@ -57,44 +63,51 @@ export const make = Effect.gen(function*() {
   type Schemas = {
     readonly decode: (u: unknown) => Effect.Effect<Rpc.Payload<Rpcs>, ParseError>
     readonly encodeChunk: (u: ReadonlyArray<unknown>) => Effect.Effect<NonEmptyReadonlyArray<unknown>, ParseError>
-    readonly encodeExit: (u: unknown) => Effect.Effect<Schema.ExitEncoded<unknown, unknown, unknown>, ParseError>
+    readonly encodeSuccess: (u: unknown) => Effect.Effect<unknown, ParseError>
+    readonly encodeFailure: (u: unknown) => Effect.Effect<unknown, ParseError>
     readonly context: Context.Context<never>
   }
 
+  type Client = {
+    readonly id: number
+    readonly schemas: Map<RequestId, Schemas>
+  }
+  const clients = new Map<number, Client>()
+
   const schemasCache = new WeakMap<any, Schemas>()
-  const getSchemas = (rpc: Rpc.AnyWithProps) => {
-    let schemas = schemasCache.get(rpc)
-    if (!schemas) {
+  const getRequestSchemas = (rpc: Rpc.AnyWithProps) => {
+    let metadata = schemasCache.get(rpc)
+    if (!metadata) {
       const entry = context.unsafeMap.get(rpc.key) as Rpc.Handler<Rpcs["_tag"]>
-      schemas = {
+      metadata = {
         decode: Schema.decodeUnknown(rpc.payloadSchema as any),
-        encodeChunk: Schema.encodeUnknown(Schema.Any) as any,
-        encodeExit: Schema.encodeUnknown(Rpc.exitSchema(rpc as any)) as any,
+        encodeChunk: Schema.encodeUnknown(Schema.Array(Schema.Any)) as any,
+        encodeSuccess: Schema.encodeUnknown(rpc.successSchema) as any,
+        encodeFailure: Schema.encodeUnknown(rpc.errorSchema as any) as any,
         context: entry.context
       }
-      schemasCache.set(rpc, schemas)
+      schemasCache.set(rpc, metadata)
     }
-    return schemas
+    return metadata
   }
 
-  type Client = {
-    readonly id: number
-    readonly schemas: Map<RequestId, Schemas>
-  }
-  const clients = new Map<number, Client>()
-
   const handleEncode = <A, R>(
     client: Client,
     requestId: RequestId,
-    effect: Effect.Effect<A, ParseError, R>,
-    onSuccess: (value: A) => any
+    encode: Effect.Effect<A, ParseError, R>,
+    onSuccess: (value: A) => FromServerEncoded | ReadonlyArray<FromServerEncoded>
   ) =>
     Effect.catchAllCause(
-      Effect.flatMap(effect, (value) => transport.send(client.id, onSuccess(value))),
+      Effect.flatMap(encode, (value) => transport.send(client.id, onSuccess(value))),
       (cause) => {
         client.schemas.delete(requestId)
         const _defect = Cause.squash(Cause.map(cause, TreeFormatter.formatErrorSync))
-        // TODO
+        // TODO: send appropriate error message for failure to encode response
+        // return transport.send(client.id, {
+        //   _tag: "Failure",
+        //   id: Number(requestId),
+        //   error: new
+        // })
         return Effect.void
       }
     )
@@ -108,41 +121,49 @@ export const make = Effect.gen(function*() {
       if (!client) return Effect.void
       switch (response._tag) {
         case "Chunk": {
-          const schemas = client.schemas.get(response.requestId)
-          if (!schemas) return Effect.void
+          const metadata = client.schemas.get(response.requestId)
+          if (!metadata) return Effect.void
           return handleEncode(
             client,
             response.requestId,
-            Effect.provide(schemas.encodeChunk(response.values), schemas.context),
+            Effect.provide(metadata.encodeChunk(response.values), metadata.context),
             (results) =>
               results.map((result) => ({
                 _tag: "Success",
                 id: Number(response.requestId),
-                method: "",
-                result
+                result: result as any
               }))
           )
         }
         case "Exit": {
-          const schemas = client.schemas.get(response.requestId)
-          if (!schemas) return Effect.void
+          const metadata = client.schemas.get(response.requestId)
+          if (!metadata) return Effect.void
           client.schemas.delete(response.requestId)
           return handleEncode(
             client,
             response.requestId,
-            Effect.provide(schemas.encodeExit(response.exit), schemas.context),
-            (exit) => {
-              if (exit._tag === "Success") {
-                return {
-                  _tag: "Success",
-                  id: Number(response.requestId),
-                  method: "",
-                  result: exit.value
-                }
-              }
-              // TODO
-              return Effect.void
-            }
+            Exit.match(response.exit, {
+              onFailure: (cause): Effect.Effect<Either.Either<unknown, unknown>, ParseError> =>
+                metadata.encodeFailure(Cause.squash(cause)).pipe(
+                  Effect.map((error) => Either.left(error))
+                ),
+              onSuccess: (value): Effect.Effect<Either.Either<unknown, unknown>, ParseError> =>
+                metadata.encodeSuccess(value).pipe(
+                  Effect.map((value) => Either.right(value))
+                )
+            }),
+            Either.match({
+              onLeft: (error) => ({
+                _tag: "Failure",
+                id: Number(response.requestId),
+                error: error as any
+              }),
+              onRight: (result) => ({
+                _tag: "Success",
+                id: Number(response.requestId),
+                result: result as any
+              })
+            })
           )
         }
         case "ClientEnd": {
@@ -170,23 +191,38 @@ export const make = Effect.gen(function*() {
 
       const rpc = ClientRpcs.requests.get(request.method)
       if (!rpc) {
-        // TODO: send error
-        return Effect.void
+        // For unknown notifications, ignore the request
+        if (request._tag === "Notification") return
+        // For unknown requests, respond with an error
+        return yield* transport.send(clientId, {
+          _tag: "Failure",
+          id: request.id,
+          error: new MethodNotFound({
+            message: "Method not found",
+            data: { method: request.method }
+          })
+        })
       }
 
       switch (request._tag) {
         case "Request": {
           const requestId = RequestId(typeof request.id === "string" ? request.id : BigInt(request.id))
-          const schemas = getSchemas(rpc)
+          const metadata = getRequestSchemas(rpc)
           return yield* Effect.matchEffect(
-            Effect.provide(schemas.decode(request.payload), schemas.context),
+            Effect.provide(metadata.decode(request.payload), metadata.context),
             {
               onFailure: (_error) => {
-                // TODO: send error
-                return Effect.void
+                return transport.send(clientId, {
+                  _tag: "Failure",
+                  id: request.id,
+                  error: new InvalidParams({
+                    message: `Invalid parameters for method: "${request.method}"`,
+                    data: { params: request.payload }
+                  })
+                })
               },
               onSuccess: (payload) => {
-                client.schemas.set(requestId, schemas)
+                client.schemas.set(requestId, metadata)
                 return server.write(client.id, {
                   _tag: "Request",
                   id: requestId,
@@ -215,42 +251,63 @@ export const make = Effect.gen(function*() {
   return {} as const
 })
 
+const ClientRpcHandlers = ClientRpcs.toLayer(
+  Effect.gen(function*() {
+    const serverInfo = yield* McpServerImplementation
+    const { capabilities = {} } = yield* McpServerOptions
+    return {
+      // Requests
+      ping: () => Effect.succeed({}),
+      initialize: Effect.fnUntraced(function*(params) {
+        const requestedVersion = params.protocolVersion
+        return {
+          capabilities,
+          serverInfo,
+          protocolVersion: SUPPORTED_PROTOCOL_VERSIONS.includes(requestedVersion)
+            ? requestedVersion
+            : LATEST_PROTOCOL_VERSION
+        }
+      }),
+      "completion/complete": () => Effect.dieMessage("Not implemented"),
+      "logging/setLevel": () => Effect.dieMessage("Not implemented"),
+      "prompts/get": () => Effect.dieMessage("Not implemented"),
+      "prompts/list": () => Effect.dieMessage("Not implemented"),
+      "resources/list": () => Effect.dieMessage("Not implemented"),
+      "resources/read": () => Effect.dieMessage("Not implemented"),
+      "resources/subscribe": () => Effect.dieMessage("Not implemented"),
+      "resources/unsubscribe": () => Effect.dieMessage("Not implemented"),
+      "resources/templates/list": () => Effect.dieMessage("Not implemented"),
+      "tools/call": () => Effect.dieMessage("Not implemented"),
+      "tools/list": () => Effect.dieMessage("Not implemented"),
+
+      // Notifications
+      "notifications/cancelled": () => Effect.dieMessage("Not implemented"),
+      "notifications/initialized": () => Effect.dieMessage("Not implemented"),
+      "notifications/progress": () => Effect.dieMessage("Not implemented"),
+      "notifications/roots/list_changed": () => Effect.dieMessage("Not implemented")
+    }
+  })
+)
+
 export const layer = (
   serverInfo: Context.Tag.Service<McpServerImplementation>,
   options?: Context.Tag.Service<McpServerOptions>
 ) =>
   Layer.effect(McpServer, make).pipe(
-    Layer.provide(ClientRpcs.toLayer(
-      Effect.gen(function*() {
-        const serverInfo = yield* McpServerImplementation
-        const { capabilities = {} } = yield* McpServerOptions
-        return {
-          ping: () => Effect.succeed({}),
-          initialize: Effect.fnUntraced(function*(params) {
-            const requestedVersion = params.protocolVersion
-            return {
-              capabilities,
-              serverInfo,
-              protocolVersion: SUPPORTED_PROTOCOL_VERSIONS.includes(requestedVersion)
-                ? requestedVersion
-                : LATEST_PROTOCOL_VERSION
-            }
-          }),
-          "notifications/initialized": () => Effect.void
-          // TODO: remove
-        } as any
-      })
-    )),
+    Layer.provide(ClientRpcHandlers),
     Layer.provideMerge(Layer.succeed(McpServerImplementation, serverInfo)),
     Layer.provideMerge(Layer.succeed(McpServerOptions, options ?? {}))
   )
 
+// Usage
+
 const MainLayer = layer({
   name: "Demo Server",
   version: "1.0.0"
 }, {
   capabilities: {
-    logging: {}
+    logging: {},
+    prompts: {}
   }
 }).pipe(
   Layer.provide(McpTransport.layerTransportStdio())
diff --git a/packages/ai/ai/src/McpTransport.ts b/packages/ai/ai/src/McpTransport.ts
@@ -19,6 +19,8 @@ import type {
 } from "./McpSchema.js"
 
 /**
+ * Represents a transport which can send and receive messages.
+ *
  * @since 1.0.0
  * @category Transport
  */
@@ -111,7 +113,7 @@ export const makeTransportStdio = Effect.fnUntraced(function*(options?: {
         case "Notification": {
           encoded = parser.encode({
             jsonrpc: JSON_RPC_VERSION,
-            method: message._tag,
+            method: message.method,
             params: message.payload
           })
           break
