1import type * as McpProtocol from "effect/ai/McpProtocol"
2import type * as McpSchema from "effect/ai/McpSchema"
3import * as Cause from "effect/Cause"
4import * as Effect from "effect/Effect"
5import * as Exit from "effect/Exit"
6import * as Option from "effect/Option"
7import * as Queue from "effect/Queue"
8import * as Ref from "effect/Ref"
9import * as RpcClient from "effect/rpc/RpcClient"
10import type * as RpcClientError from "effect/rpc/RpcClientError"
11import type * as RpcGroup from "effect/rpc/RpcGroup"
12import type * as RpcSerialization from "effect/rpc/RpcSerialization"
13import * as Schema from "effect/Schema"
14
15export type ReverseMethod = "roots/list" | "sampling/createMessage" | "elicitation/create"
16
17export interface RecordedRequest {
18 readonly id: string | number
19 readonly method: ReverseMethod
20 readonly payload: unknown
21}
22
23export type Handler = (
24 request: RecordedRequest
25) => Effect.Effect<unknown, typeof McpSchema.McpError.Type>
26
27export interface McpTestPeerOptions {
28 readonly capabilities?: typeof McpSchema.ClientCapabilities.Type | undefined
29 readonly clientInfo?: typeof McpSchema.Implementation.Type | undefined
30 readonly handlers?: Partial<Record<ReverseMethod, Handler>> | undefined
31}
32
33export interface McpTestPeer {
34 readonly wireClient: RpcClient.RpcClient<
35 RpcGroup.Rpcs<typeof McpSchema.ServerRequestRpcs>,
36 RpcClientError.RpcClientError
37 >
38 readonly reverseClient: McpSchema.McpReverseClient
39 readonly requests: Effect.Effect<ReadonlyArray<RecordedRequest>>
40 readonly takeRequest: Effect.Effect<RecordedRequest>
41}
42
43const isReverseMethod = (method: string): method is ReverseMethod =>
44 ["roots/list", "sampling/createMessage", "elicitation/create"].includes(method)
45
46export const makeMcpTestPeer = Effect.fn("McpTestPeer.make")(function*(
47 protocol: McpProtocol.ProtocolAdapter,
48 options: McpTestPeerOptions = {}
49) {
50 const requests = yield* Ref.make<ReadonlyArray<RecordedRequest>>([])
51 const inbox = yield* Queue.unbounded<RecordedRequest>()
52 const handlers = options.handlers ?? {}
53
54 const rpcProtocol = yield* RpcClient.Protocol.make((writeResponse) =>
55 Effect.succeed({
56 send: (clientId, message) => {
57 if (message._tag !== "Request" || !isReverseMethod(message.tag)) {
58 return Effect.void
59 }
60 const request: RecordedRequest = {
61 id: message.id,
62 method: message.tag,
63 payload: message.payload
64 }
65 return Effect.gen(function*() {
66 yield* Ref.update(requests, (current) => [...current, request])
67 yield* Queue.offer(inbox, request)
68
69 const handler = handlers[request.method]
70 const result = yield* Effect.exit(
71 handler === undefined
72 ? Effect.fail({
73 code: -32601,
74 message: `No test peer handler for ${request.method}`
75 })
76 : handler(request)
77 )
78
79 if (Exit.isSuccess(result)) {
80 return yield* writeResponse(clientId, {
81 _tag: "Exit",
82 requestId: request.id,
83 exit: {
84 _tag: "Success",
85 value: result.value
86 }
87 })
88 }
89
90 const failure = Cause.findErrorOption(result.cause)
91 if (Option.isSome(failure)) {
92 return yield* writeResponse(clientId, {
93 _tag: "Exit",
94 requestId: request.id,
95 exit: {
96 _tag: "Failure",
97 cause: [{
98 _tag: "Fail",
99 error: failure.value
100 }]
101 }
102 })
103 }
104
105 return yield* Effect.die(Cause.squash(result.cause))
106 })
107 },
108 supportsAck: true,
109 supportsTransferables: false,
110 codecFor: Schema.toCodecJson as RpcSerialization.CodecFor
111 })
112 )
113
114 const wireClient = yield* RpcClient.make(
115 protocol.serverRequestRpcs as unknown as typeof McpSchema.ServerRequestRpcs
116 ).pipe(
117 Effect.provideService(RpcClient.Protocol, rpcProtocol)
118 )
119 const reverseClient = yield* protocol.makeReverseClient({
120 protocolVersion: protocol.protocolVersion,
121 clientCapabilities: options.capabilities ?? {},
122 clientInfo: options.clientInfo ?? { name: "McpTestPeer", version: "1.0.0" }
123 }).pipe(Effect.provideService(RpcClient.Protocol, rpcProtocol))
124
125 return {
126 wireClient,
127 reverseClient,
128 requests: Ref.get(requests),
129 takeRequest: Queue.take(inbox)
130 } satisfies McpTestPeer
131})