mirror of
https://github.com/anomalyco/opencode.git
synced 2026-10-26 18:30:00 +00:00
257 lines
8.5 KiB
TypeScript
257 lines
8.5 KiB
TypeScript
import { afterEach, describe, expect, jest, test } from "bun:test"
|
|
import { Client, type JSONRPCMessage, type Transport } from "@modelcontextprotocol/client"
|
|
import type { ServerWebSocket } from "bun"
|
|
import { connectClientMcp } from "../src/mcp"
|
|
|
|
type Peer = {
|
|
readonly request: Request
|
|
readonly socket: ServerWebSocket<{ id: number; request: Request }>
|
|
readonly client: Client
|
|
readonly closed: Promise<void>
|
|
}
|
|
|
|
const servers: Array<{ stop(force: boolean): Promise<void> }> = []
|
|
afterEach(async () => {
|
|
jest.useRealTimers()
|
|
await Promise.all(servers.splice(0).map((server) => server.stop(true)))
|
|
})
|
|
|
|
describe("connectClientMcp", () => {
|
|
test("serves registered tools to the opencode server", async () => {
|
|
const server = serve()
|
|
const connection = connectClientMcp({
|
|
url: server.url + "?ticket=t1",
|
|
name: "tui",
|
|
directory: "/repo",
|
|
headers: { authorization: "Basic secret" },
|
|
tools: [
|
|
{
|
|
name: "whoami",
|
|
description: "Return the calling session",
|
|
handler: (_args, ctx) => ({ content: [{ type: "text", text: ctx.sessionID ?? "none" }] }),
|
|
},
|
|
{
|
|
name: "echo",
|
|
inputSchema: { properties: { text: { type: "string" } }, required: ["text"] },
|
|
handler: (args) => ({ content: [{ type: "text", text: String(args.text) }] }),
|
|
},
|
|
{
|
|
name: "fail",
|
|
handler: () => {
|
|
throw new Error("boom")
|
|
},
|
|
},
|
|
],
|
|
})
|
|
const peer = await server.next()
|
|
|
|
const url = new URL(peer.request.url)
|
|
expect(url.pathname).toBe("/api/experimental/client-mcp")
|
|
expect(url.searchParams.get("name")).toBe("tui")
|
|
expect(url.searchParams.get("ticket")).toBe("t1")
|
|
expect(url.searchParams.get("location[directory]")).toBe("/repo")
|
|
expect(peer.request.headers.get("authorization")).toBe("Basic secret")
|
|
|
|
await peer.client.ping()
|
|
expect((await peer.client.listTools()).tools).toEqual([
|
|
{ name: "whoami", description: "Return the calling session", inputSchema: { type: "object" } },
|
|
{
|
|
name: "echo",
|
|
inputSchema: { type: "object", properties: { text: { type: "string" } }, required: ["text"] },
|
|
},
|
|
{ name: "fail", inputSchema: { type: "object" } },
|
|
])
|
|
expect(
|
|
await peer.client.callTool({ name: "whoami", arguments: {}, _meta: { "ai.opencode/sessionID": "ses_1" } }),
|
|
).toMatchObject({ content: [{ type: "text", text: "ses_1" }] })
|
|
expect(await peer.client.callTool({ name: "whoami", arguments: {} })).toMatchObject({
|
|
content: [{ type: "text", text: "none" }],
|
|
})
|
|
|
|
expect(await peer.client.callTool({ name: "echo", arguments: { text: "hi" } })).toMatchObject({
|
|
content: [{ type: "text", text: "hi" }],
|
|
})
|
|
expect(await peer.client.callTool({ name: "fail", arguments: {} })).toMatchObject({
|
|
content: [{ type: "text", text: "boom" }],
|
|
isError: true,
|
|
})
|
|
await expect(peer.client.callTool({ name: "missing", arguments: {} })).rejects.toThrow("Unknown tool: missing")
|
|
|
|
await connection.close()
|
|
await peer.closed
|
|
})
|
|
|
|
test("reconnects after the server drops the socket", async () => {
|
|
const server = serve()
|
|
const connection = connectClientMcp({
|
|
url: server.url,
|
|
name: "tui",
|
|
tools: [{ name: "ping", handler: () => ({ content: [{ type: "text", text: "pong" }] }) }],
|
|
})
|
|
const first = await server.next()
|
|
expect((await first.client.listTools()).tools).toHaveLength(1)
|
|
|
|
first.socket.close(1011, "internal")
|
|
const second = await server.next()
|
|
expect(await second.client.callTool({ name: "ping", arguments: {} })).toMatchObject({
|
|
content: [{ type: "text", text: "pong" }],
|
|
})
|
|
|
|
await connection.close()
|
|
})
|
|
|
|
test("close cancels a pending reconnect", async () => {
|
|
const server = serve()
|
|
const socket = trackedSocket()
|
|
const connection = connectClientMcp({ url: server.url, name: "tui", tools: [], WebSocket: socket.WebSocket })
|
|
const peer = await server.next()
|
|
|
|
jest.useFakeTimers()
|
|
peer.socket.close(1011, "internal")
|
|
await socket.closed(1)
|
|
await connection.close()
|
|
jest.advanceTimersByTime(10_000)
|
|
expect(socket.opened()).toBe(1)
|
|
})
|
|
|
|
test("keeps backing off while the opencode server never initializes", async () => {
|
|
const server = Bun.serve({
|
|
port: 0,
|
|
fetch(request, server) {
|
|
if (server.upgrade(request)) return
|
|
return new Response("upgrade required", { status: 426 })
|
|
},
|
|
websocket: {
|
|
open(socket) {
|
|
socket.send("not json")
|
|
socket.close(1011, "internal")
|
|
},
|
|
message() {},
|
|
},
|
|
})
|
|
servers.push(server)
|
|
const socket = trackedSocket()
|
|
jest.useFakeTimers()
|
|
const connection = connectClientMcp({
|
|
url: server.url.href,
|
|
name: "tui",
|
|
tools: [],
|
|
WebSocket: socket.WebSocket,
|
|
})
|
|
|
|
await socket.closed(1)
|
|
jest.advanceTimersByTime(500)
|
|
expect(socket.opened()).toBe(2)
|
|
await socket.closed(2)
|
|
jest.advanceTimersByTime(999)
|
|
expect(socket.opened()).toBe(2)
|
|
jest.advanceTimersByTime(1)
|
|
expect(socket.opened()).toBe(3)
|
|
|
|
await connection.close()
|
|
})
|
|
})
|
|
|
|
/** Counts sockets synchronously at construction, so fake-timer reconnects are observable without real waits. */
|
|
function trackedSocket() {
|
|
const state = { opened: 0, closed: 0, waiting: [] as Array<{ count: number; resolve: () => void }> }
|
|
class Tracked extends WebSocket {
|
|
constructor(url: string | URL, init?: string | string[]) {
|
|
super(url, init)
|
|
state.opened++
|
|
this.addEventListener("close", () => {
|
|
state.closed++
|
|
state.waiting = state.waiting.filter((waiter) => {
|
|
if (state.closed < waiter.count) return true
|
|
// Let connectClientMcp's own close listener run before the test continues.
|
|
setImmediate(waiter.resolve)
|
|
return false
|
|
})
|
|
})
|
|
}
|
|
}
|
|
return {
|
|
WebSocket: Tracked,
|
|
opened: () => state.opened,
|
|
closed: (count: number) =>
|
|
new Promise<void>((resolve) => {
|
|
if (state.closed >= count) return resolve()
|
|
state.waiting.push({ count, resolve })
|
|
}),
|
|
}
|
|
}
|
|
|
|
function serve() {
|
|
const waiting: Array<(peer: Peer) => void> = []
|
|
const ready: Peer[] = []
|
|
const transports = new Map<number, ReturnType<typeof socketTransport>>()
|
|
const state = { connections: 0 }
|
|
const server = Bun.serve<{ id: number; request: Request }>({
|
|
port: 0,
|
|
fetch(request, server) {
|
|
if (server.upgrade(request, { data: { id: ++state.connections, request } })) return
|
|
return new Response("upgrade required", { status: 426 })
|
|
},
|
|
websocket: {
|
|
open(socket) {
|
|
const transport = socketTransport(socket)
|
|
transports.set(socket.data.id, transport)
|
|
const client = new Client({ name: "opencode", version: "1.0.0" })
|
|
const closed = new Promise<void>((resolve) => {
|
|
client.onclose = resolve
|
|
})
|
|
void client.connect(transport).then(() => {
|
|
const peer = { request: socket.data.request, socket, client, closed }
|
|
const resolve = waiting.shift()
|
|
if (resolve) return resolve(peer)
|
|
ready.push(peer)
|
|
})
|
|
},
|
|
// Deliver outside the uWS callback: on Bun 1.4.4, replying from the SDK's microtasks inside it can make the
|
|
// server reject the client's next frame ("Received new message inside a fragmented message").
|
|
message(socket, message) {
|
|
const text = String(message)
|
|
setImmediate(() => transports.get(socket.data.id)?.deliver(text))
|
|
},
|
|
close(socket) {
|
|
setImmediate(() => transports.get(socket.data.id)?.ended())
|
|
},
|
|
},
|
|
})
|
|
const result = {
|
|
url: server.url.href,
|
|
next: () =>
|
|
new Promise<Peer>((resolve) => {
|
|
const peer = ready.shift()
|
|
if (peer) return resolve(peer)
|
|
waiting.push(resolve)
|
|
}),
|
|
stop: (force: boolean) => server.stop(force),
|
|
}
|
|
servers.push(result)
|
|
return result
|
|
}
|
|
|
|
function socketTransport(socket: ServerWebSocket<{ id: number; request: Request }>) {
|
|
const state = { closed: false }
|
|
const transport: Transport & { deliver(text: string): void; ended(): void } = {
|
|
async start() {},
|
|
async send(message: JSONRPCMessage) {
|
|
if (state.closed) throw new Error("Not connected")
|
|
socket.send(JSON.stringify(message))
|
|
},
|
|
async close() {
|
|
socket.close(1000)
|
|
transport.ended()
|
|
},
|
|
deliver(text) {
|
|
transport.onmessage?.(JSON.parse(text))
|
|
},
|
|
ended() {
|
|
if (state.closed) return
|
|
state.closed = true
|
|
transport.onclose?.()
|
|
},
|
|
}
|
|
return transport
|
|
}
|