1import type * as McpProtocol from "effect/ai/McpProtocol"
2import type * as McpSchema from "effect/ai/McpSchema"
3import * as McpServer from "effect/ai/McpServer"
4import type * as Arr from "effect/Array"
5import type * as Cause from "effect/Cause"
6import * as Context from "effect/Context"
7import * as Deferred from "effect/Deferred"
8import * as Effect from "effect/Effect"
9import type * as Fiber from "effect/Fiber"
10import * as Layer from "effect/Layer"
11import * as Queue from "effect/Queue"
12import * as Sink from "effect/Sink"
13import * as Stdio from "effect/Stdio"
14import * as Stream from "effect/Stream"
15
16export interface JsonRpcMessage {
17 readonly jsonrpc: "2.0"
18 readonly id?: string | number | null | undefined
19 readonly method?: string | undefined
20 readonly params?: unknown
21 readonly result?: unknown
22 readonly error?: unknown
23}
24
25export interface McpStdioRequestHandle {
26 readonly id: string | number
27 readonly response: Effect.Effect<JsonRpcMessage>
28 readonly cancel: (reason?: string) => Effect.Effect<void>
29}
30
31export interface McpStdioHarness {
32 readonly server: McpServer.McpServer["Service"]
33 readonly serverFiber: Fiber.Fiber<never, Cause.IllegalArgumentError>
34 readonly close: Effect.Effect<void>
35 readonly sendRaw: (message: unknown) => Effect.Effect<void>
36 readonly sendChunk: (chunk: string | Uint8Array) => Effect.Effect<void>
37 readonly sendRequest: (
38 method: string,
39 params?: unknown,
40 id?: string | number
41 ) => Effect.Effect<JsonRpcMessage>
42 readonly startRequest: (
43 method: string,
44 params?: unknown,
45 id?: string | number
46 ) => Effect.Effect<McpStdioRequestHandle>
47 readonly sendNotification: (method: string, params?: unknown) => Effect.Effect<void>
48 readonly initialize: (
49 capabilities?: typeof McpSchema.ClientCapabilities.Type
50 ) => Effect.Effect<JsonRpcMessage>
51 readonly takeMessage: Effect.Effect<JsonRpcMessage>
52 readonly takeFrame: Effect.Effect<unknown>
53 readonly takeRawStdout: Effect.Effect<string>
54 readonly takeStderr: Effect.Effect<string>
55 readonly awaitOutboundMethod: (method: string) => Effect.Effect<JsonRpcMessage>
56 readonly flushListChanged: Effect.Effect<void>
57 readonly respond: (id: string | number, result: unknown) => Effect.Effect<void>
58}
59
60const isJsonRpcMessage = (value: unknown): value is JsonRpcMessage =>
61 typeof value === "object" && value !== null && (value as JsonRpcMessage).jsonrpc === "2.0"
62
63const isResponse = (
64 message: JsonRpcMessage
65): message is JsonRpcMessage & { readonly id: string | number } =>
66 (typeof message.id === "string" || typeof message.id === "number") &&
67 message.method === undefined
68
69const requestKey = (id: string | number) => `${typeof id}:${id}`
70
71export const makeMcpStdioHarness = Effect.fnUntraced(function*<A>(
72 protocol: McpProtocol.ProtocolAdapter,
73 protocols: Arr.NonEmptyReadonlyArray<McpProtocol.ProtocolAdapter> = [protocol],
74 registrations?: Layer.Layer<A, never, McpServer.McpServer>
75) {
76 const stdin = yield* Queue.unbounded<Uint8Array, Cause.Done>()
77 const stdout = yield* Queue.unbounded<string | Uint8Array>()
78 const stderr = yield* Queue.unbounded<string | Uint8Array>()
79 const frames = yield* Queue.unbounded<unknown>()
80 const messages = yield* Queue.unbounded<JsonRpcMessage>()
81 const rawStdout = yield* Queue.unbounded<string>()
82 const rawStderr = yield* Queue.unbounded<string>()
83 const responseQueues = new Map<string, Queue.Queue<JsonRpcMessage>>()
84 const retainedMessages: Array<JsonRpcMessage> = []
85 const encoder = new TextEncoder()
86 const stdoutDecoder = new TextDecoder()
87 const stderrDecoder = new TextDecoder()
88 let nextRequestId = 1
89
90 const stdioLayer = Stdio.layerTest({
91 stdin: Stream.fromQueue(stdin),
92 stdout: () => Sink.forEach((chunk) => Queue.offer(stdout, chunk)),
93 stderr: () => Sink.forEach((chunk) => Queue.offer(stderr, chunk))
94 })
95 const ready = yield* Deferred.make<McpServer.McpServer["Service"]>()
96 const serverFiber = yield* Effect.gen(function*() {
97 const serverLayer = McpServer.layerStdio({
98 name: "McpConformance",
99 version: "1.0.0",
100 protocols
101 }).pipe(Layer.provide(stdioLayer))
102 const context = yield* Layer.build(
103 registrations === undefined
104 ? serverLayer
105 : registrations.pipe(Layer.provideMerge(serverLayer))
106 )
107 yield* Deferred.succeed(ready, Context.get(context, McpServer.McpServer))
108 return yield* Effect.never
109 }).pipe(Effect.scoped, Effect.forkScoped)
110 const server = yield* Deferred.await(ready)
111
112 const routeFrame = Effect.fnUntraced(function*(frame: unknown) {
113 yield* Queue.offer(frames, frame)
114 if (!isJsonRpcMessage(frame)) {
115 return
116 }
117 if (isResponse(frame)) {
118 const responseQueue = responseQueues.get(requestKey(frame.id))
119 if (responseQueue !== undefined) {
120 yield* Queue.offer(responseQueue, frame)
121 return
122 }
123 }
124 yield* Queue.offer(messages, frame)
125 })
126
127 yield* Effect.gen(function*() {
128 let pending = ""
129 while (true) {
130 const chunk = yield* Queue.take(stdout)
131 const text = typeof chunk === "string"
132 ? chunk
133 : stdoutDecoder.decode(chunk, { stream: true })
134 yield* Queue.offer(rawStdout, text)
135 pending += text
136 let newline = pending.indexOf("\n")
137 while (newline !== -1) {
138 const line = pending.slice(0, newline)
139 pending = pending.slice(newline + 1)
140 if (line.length > 0) {
141 yield* routeFrame(JSON.parse(line))
142 }
143 newline = pending.indexOf("\n")
144 }
145 }
146 }).pipe(Effect.forkScoped)
147
148 yield* Effect.gen(function*() {
149 while (true) {
150 const chunk = yield* Queue.take(stderr)
151 yield* Queue.offer(
152 rawStderr,
153 typeof chunk === "string" ? chunk : stderrDecoder.decode(chunk, { stream: true })
154 )
155 }
156 }).pipe(Effect.forkScoped)
157
158 const sendChunk = (chunk: string | Uint8Array) =>
159 Queue.offer(stdin, typeof chunk === "string" ? encoder.encode(chunk) : chunk)
160 const sendRaw = (message: unknown) => sendChunk(`${JSON.stringify(message)}\n`)
161 const requestMetadata = {
162 "io.modelcontextprotocol/protocolVersion": protocol.protocolVersion,
163 "io.modelcontextprotocol/clientCapabilities": {},
164 "io.modelcontextprotocol/clientInfo": { name: "stdio-client", version: "1.0.0" }
165 }
166 const withRequestMetadata = (params: unknown) =>
167 protocol.runtime._tag === "Stateless"
168 ? {
169 ...(typeof params === "object" && params !== null ? params : {}),
170 _meta: {
171 ...requestMetadata,
172 ...(typeof params === "object" && params !== null && "_meta" in params &&
173 typeof params._meta === "object" && params._meta !== null
174 ? params._meta
175 : {})
176 }
177 }
178 : params
179 const sendNotification = (method: string, params?: unknown) =>
180 sendRaw({
181 jsonrpc: "2.0",
182 method,
183 ...(params === undefined && protocol.runtime._tag !== "Stateless"
184 ? {}
185 : { params: withRequestMetadata(params) })
186 })
187 const startRequest = Effect.fnUntraced(function*(
188 method: string,
189 params?: unknown,
190 id: string | number = nextRequestId++
191 ) {
192 const responseQueue = yield* Queue.unbounded<JsonRpcMessage>()
193 const key = requestKey(id)
194 responseQueues.set(key, responseQueue)
195 yield* sendRaw({
196 jsonrpc: "2.0",
197 id,
198 method,
199 ...(params === undefined && protocol.runtime._tag !== "Stateless"
200 ? {}
201 : { params: withRequestMetadata(params) })
202 })
203 return {
204 id,
205 response: Queue.take(responseQueue).pipe(
206 Effect.ensuring(Effect.sync(() => responseQueues.delete(key)))
207 ),
208 cancel: (reason?: string) =>
209 sendNotification("notifications/cancelled", {
210 requestId: id,
211 ...(reason === undefined ? {} : { reason })
212 })
213 }
214 })
215 const sendRequest = Effect.fnUntraced(function*(
216 method: string,
217 params?: unknown,
218 id?: string | number
219 ) {
220 const request = yield* startRequest(method, params, id)
221 return yield* request.response
222 })
223 const takeMessage = Effect.suspend(() => {
224 const retained = retainedMessages.shift()
225 return retained === undefined ? Queue.take(messages) : Effect.succeed(retained)
226 })
227 const awaitOutboundMethod = Effect.fnUntraced(function*(method: string) {
228 const retainedIndex = retainedMessages.findIndex((message) => message.method === method)
229 if (retainedIndex !== -1) {
230 return retainedMessages.splice(retainedIndex, 1)[0]!
231 }
232 while (true) {
233 const message = yield* Queue.take(messages)
234 if (message.method === method) {
235 return message
236 }
237 retainedMessages.push(message)
238 }
239 })
240
241 return {
242 server,
243 serverFiber,
244 close: Queue.end(stdin),
245 sendRaw,
246 sendChunk,
247 sendRequest,
248 startRequest,
249 sendNotification,
250 initialize: Effect.fnUntraced(function*(capabilities = {}) {
251 if (protocol.runtime._tag === "Stateless") {
252 return yield* sendRequest("server/discover", {
253 _meta: {
254 ...requestMetadata,
255 "io.modelcontextprotocol/clientCapabilities": capabilities
256 }
257 })
258 }
259 const response = yield* sendRequest("initialize", {
260 protocolVersion: protocol.protocolVersion,
261 capabilities,
262 clientInfo: { name: "stdio-client", version: "1.0.0" }
263 })
264 yield* sendNotification("notifications/initialized")
265 return response
266 }),
267 takeMessage,
268 takeFrame: Queue.take(frames),
269 takeRawStdout: Queue.take(rawStdout),
270 takeStderr: Queue.take(rawStderr),
271 awaitOutboundMethod,
272 flushListChanged: Effect.yieldNow,
273 respond: (id, result) => sendRaw({ jsonrpc: "2.0", id, result })
274 } satisfies McpStdioHarness
275})