Compare commits

...
Sign in to create a new pull request.

4 commits

Author SHA1 Message Date
Aiden Cline
2a25bdce74 refactor(mcp): use tagged connection results 2026-06-09 13:54:05 -05:00
Aiden Cline
deebc338c8 refactor(mcp): simplify connection result flow 2026-06-09 13:49:07 -05:00
Aiden Cline
91a696da0b refactor(mcp): tighten connection result types 2026-06-09 13:46:50 -05:00
Aiden Cline
3fd1d4da7a refactor(mcp): simplify service helpers 2026-06-09 13:45:22 -05:00

View file

@ -10,7 +10,6 @@ import {
CallToolResultSchema, CallToolResultSchema,
ListToolsResultSchema, ListToolsResultSchema,
ToolSchema, ToolSchema,
type Tool as MCPToolDef,
ToolListChangedNotificationSchema, ToolListChangedNotificationSchema,
} from "@modelcontextprotocol/sdk/types.js" } from "@modelcontextprotocol/sdk/types.js"
import { Config } from "@/config/config" import { Config } from "@/config/config"
@ -71,6 +70,7 @@ export class NotFoundError extends Schema.TaggedErrorClass<NotFoundError>()("MCP
}) {} }) {}
type MCPClient = Client type MCPClient = Client
type MCPToolDef = Awaited<ReturnType<MCPClient["listTools"]>>["tools"][number]
const StatusConnected = Schema.Struct({ status: Schema.Literal("connected") }).annotate({ const StatusConnected = Schema.Struct({ status: Schema.Literal("connected") }).annotate({
identifier: "MCPStatusConnected", identifier: "MCPStatusConnected",
@ -114,7 +114,7 @@ function isMcpConfigured(entry: McpEntry): entry is ConfigMCPV1.Info {
const sanitize = (s: string) => s.replace(/[^a-zA-Z0-9_-]/g, "_") const sanitize = (s: string) => s.replace(/[^a-zA-Z0-9_-]/g, "_")
const MAX_LIST_PAGES = 1_000 const MAX_LIST_PAGES = 1_000
function remoteURL(key: string, value: string) { function remoteURL(value: string) {
if (URL.canParse(value)) return new URL(value) if (URL.canParse(value)) return new URL(value)
} }
@ -124,6 +124,10 @@ function isOutputSchemaValidationError(error: Error) {
) )
} }
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value)
}
async function paginate<T, R extends { nextCursor?: string }>( async function paginate<T, R extends { nextCursor?: string }>(
list: (cursor?: string) => Promise<R>, list: (cursor?: string) => Promise<R>,
items: (result: R) => T[], items: (result: R) => T[],
@ -168,12 +172,12 @@ function convertMcpTool(mcpTool: MCPToolDef, client: MCPClient, timeout?: number
const inputSchema = mcpTool.inputSchema const inputSchema = mcpTool.inputSchema
// Spread first, then override type to ensure it's always "object" // Spread first, then override type to ensure it's always "object"
const schema: JSONSchema7 = { const schema = {
...(inputSchema as JSONSchema7), ...inputSchema,
type: "object", type: "object",
properties: (inputSchema.properties ?? {}) as JSONSchema7["properties"], properties: inputSchema.properties ?? {},
additionalProperties: false, additionalProperties: false,
} } satisfies JSONSchema7
return dynamicTool({ return dynamicTool({
description: mcpTool.description ?? "", description: mcpTool.description ?? "",
@ -182,7 +186,7 @@ function convertMcpTool(mcpTool: MCPToolDef, client: MCPClient, timeout?: number
return client.callTool( return client.callTool(
{ {
name: mcpTool.name, name: mcpTool.name,
arguments: (args || {}) as Record<string, unknown>, arguments: isRecord(args) ? args : {},
}, },
CallToolResultSchema, CallToolResultSchema,
{ {
@ -196,25 +200,15 @@ function convertMcpTool(mcpTool: MCPToolDef, client: MCPClient, timeout?: number
} }
function defs(client: MCPClient, timeout?: number) { function defs(client: MCPClient, timeout?: number) {
return listTools(client, timeout ?? DEFAULT_TIMEOUT).pipe( return listTools(client, timeout ?? DEFAULT_TIMEOUT).pipe(Effect.catch(() => Effect.void))
Effect.catch((err) => {
return Effect.succeed(undefined)
}),
)
} }
function fetchFromClient<T extends { name: string }>( function fetchFromClient<T extends { name: string }>(
clientName: string, clientName: string,
client: Client, client: Client,
listFn: (c: Client) => Promise<T[]>, listFn: (c: Client) => Promise<T[]>,
label: string,
) { ) {
return Effect.tryPromise({ return Effect.tryPromise(() => listFn(client)).pipe(
try: () => listFn(client),
catch: (e: any) => {
return e
},
}).pipe(
Effect.map((items) => { Effect.map((items) => {
const out: Record<string, T & { client: string }> = {} const out: Record<string, T & { client: string }> = {}
const sanitizedClient = sanitize(clientName) const sanitizedClient = sanitize(clientName)
@ -227,11 +221,11 @@ function fetchFromClient<T extends { name: string }>(
) )
} }
interface CreateResult { type UnavailableStatus = Exclude<Status, { status: "connected" }>
mcpClient?: MCPClient type ConnectResult = { type: "connected"; client: MCPClient } | { type: "unavailable"; status: UnavailableStatus }
status: Status type CreateResult =
defs?: MCPToolDef[] | { type: "connected"; client: MCPClient; defs: MCPToolDef[] }
} | { type: "unavailable"; status: UnavailableStatus }
interface AuthResult { interface AuthResult {
authorizationUrl: string authorizationUrl: string
@ -308,20 +302,20 @@ export const layer = Layer.effect(
(t, exit) => (Exit.isFailure(exit) ? Effect.tryPromise(() => t.close()).pipe(Effect.ignore) : Effect.void), (t, exit) => (Exit.isFailure(exit) ? Effect.tryPromise(() => t.close()).pipe(Effect.ignore) : Effect.void),
) )
const DISABLED_RESULT: CreateResult = { status: { status: "disabled" } } const DISABLED_RESULT = {
type: "unavailable",
status: { status: "disabled" },
} satisfies CreateResult
const connectRemote = Effect.fn("MCP.connectRemote")(function* ( const connectRemote = Effect.fn("MCP.connectRemote")(function* (key: string, mcp: ConfigMCPV1.Remote) {
key: string,
mcp: ConfigMCPV1.Info & { type: "remote" },
) {
const oauthDisabled = mcp.oauth === false const oauthDisabled = mcp.oauth === false
const oauthConfig = typeof mcp.oauth === "object" ? mcp.oauth : undefined const oauthConfig = typeof mcp.oauth === "object" ? mcp.oauth : undefined
const url = remoteURL(key, mcp.url) const url = remoteURL(mcp.url)
if (!url) { if (!url) {
return { return {
client: undefined as MCPClient | undefined, type: "unavailable",
status: { status: "failed" as const, error: `Invalid MCP URL for "${key}"` }, status: { status: "failed", error: `Invalid MCP URL for "${key}"` },
} } satisfies ConnectResult
} }
let authProvider: McpOAuthProvider | undefined let authProvider: McpOAuthProvider | undefined
@ -361,7 +355,7 @@ export const layer = Layer.effect(
] ]
const connectTimeout = mcp.timeout ?? DEFAULT_TIMEOUT const connectTimeout = mcp.timeout ?? DEFAULT_TIMEOUT
let lastStatus: Status | undefined let lastStatus: UnavailableStatus | undefined
for (const { name, transport } of transports) { for (const { name, transport } of transports) {
const result = yield* connectTransport(transport, connectTimeout).pipe( const result = yield* connectTransport(transport, connectTimeout).pipe(
@ -374,7 +368,7 @@ export const layer = Layer.effect(
if (isAuthError) { if (isAuthError) {
if (lastError.message.includes("registration") || lastError.message.includes("client_id")) { if (lastError.message.includes("registration") || lastError.message.includes("client_id")) {
lastStatus = { lastStatus = {
status: "needs_client_registration" as const, status: "needs_client_registration",
error: "Server does not support dynamic client registration. Please provide clientId in config.", error: "Server does not support dynamic client registration. Please provide clientId in config.",
} }
return events return events
@ -387,7 +381,7 @@ export const layer = Layer.effect(
.pipe(Effect.ignore, Effect.as(undefined)) .pipe(Effect.ignore, Effect.as(undefined))
} else { } else {
pendingOAuthTransports.set(key, transport) pendingOAuthTransports.set(key, transport)
lastStatus = { status: "needs_auth" as const } lastStatus = { status: "needs_auth" }
return events return events
.publish(TuiEvent.ToastShow, { .publish(TuiEvent.ToastShow, {
title: "MCP Authentication Required", title: "MCP Authentication Required",
@ -399,27 +393,22 @@ export const layer = Layer.effect(
} }
} }
lastStatus = { status: "failed" as const, error: lastError.message } lastStatus = { status: "failed", error: lastError.message }
return Effect.succeed(undefined) return Effect.void
}), }),
) )
if (result) { if (result) return { type: "connected", client: result.client } satisfies ConnectResult
return { client: result.client as MCPClient | undefined, status: { status: "connected" } as Status }
}
// If this was an auth error, stop trying other transports // If this was an auth error, stop trying other transports
if (lastStatus?.status === "needs_auth" || lastStatus?.status === "needs_client_registration") break if (lastStatus?.status === "needs_auth" || lastStatus?.status === "needs_client_registration") break
} }
return { return {
client: undefined as MCPClient | undefined, type: "unavailable",
status: (lastStatus ?? { status: "failed", error: "Unknown error" }) as Status, status: lastStatus ?? { status: "failed", error: "Unknown error" },
} } satisfies ConnectResult
}) })
const connectLocal = Effect.fn("MCP.connectLocal")(function* ( const connectLocal = Effect.fn("MCP.connectLocal")(function* (mcp: ConfigMCPV1.Local) {
key: string,
mcp: ConfigMCPV1.Info & { type: "local" },
) {
const [cmd, ...args] = mcp.command const [cmd, ...args] = mcp.command
const cwd = yield* InstanceState.directory const cwd = yield* InstanceState.directory
const transport = new StdioClientTransport({ const transport = new StdioClientTransport({
@ -436,13 +425,13 @@ export const layer = Layer.effect(
const connectTimeout = mcp.timeout ?? DEFAULT_TIMEOUT const connectTimeout = mcp.timeout ?? DEFAULT_TIMEOUT
return yield* connectTransport(transport, connectTimeout).pipe( return yield* connectTransport(transport, connectTimeout).pipe(
Effect.map((client): { client: MCPClient | undefined; status: Status } => ({ Effect.map((client) => ({ type: "connected", client }) satisfies ConnectResult),
client, Effect.catch((error) => {
status: { status: "connected" },
})),
Effect.catch((error): Effect.Effect<{ client: MCPClient | undefined; status: Status }> => {
const msg = error instanceof Error ? error.message : String(error) const msg = error instanceof Error ? error.message : String(error)
return Effect.succeed({ client: undefined, status: { status: "failed", error: msg } }) return Effect.succeed({
type: "unavailable",
status: { status: "failed", error: msg },
} satisfies ConnectResult)
}), }),
) )
}) })
@ -452,35 +441,33 @@ export const layer = Layer.effect(
return DISABLED_RESULT return DISABLED_RESULT
} }
const { client: mcpClient, status } = const result = mcp.type === "remote" ? yield* connectRemote(key, mcp) : yield* connectLocal(mcp)
mcp.type === "remote"
? yield* connectRemote(key, mcp as ConfigMCPV1.Info & { type: "remote" })
: yield* connectLocal(key, mcp as ConfigMCPV1.Info & { type: "local" })
if (!mcpClient) { if (result.type === "unavailable") {
if (status.status !== "connected" && status.status !== "disabled") { yield* Effect.logWarning("server unavailable", { key, type: mcp.type, status: result.status.status })
yield* Effect.logWarning("server unavailable", { key, type: mcp.type, status: status.status }) return result satisfies CreateResult
}
return { status } satisfies CreateResult
} }
const listed = mcpClient.getServerCapabilities()?.tools ? yield* defs(mcpClient, mcp.timeout) : [] const listed = result.client.getServerCapabilities()?.tools ? yield* defs(result.client, mcp.timeout) : []
if (!listed) { if (!listed) {
yield* Effect.tryPromise(() => mcpClient.close()).pipe(Effect.ignore) yield* Effect.tryPromise(() => result.client.close()).pipe(Effect.ignore)
return { status: { status: "failed", error: "Failed to get tools" } } satisfies CreateResult return {
type: "unavailable",
status: { status: "failed", error: "Failed to get tools" },
} satisfies CreateResult
} }
return { mcpClient, status, defs: listed } satisfies CreateResult return { type: "connected", client: result.client, defs: listed } satisfies CreateResult
}) })
const cfgSvc = yield* Config.Service const cfgSvc = yield* Config.Service
const descendants = Effect.fnUntraced( const descendants = Effect.fnUntraced(
function* (pid: number) { function* (pid: number) {
if (process.platform === "win32") return [] as number[] if (process.platform === "win32") return Array<number>()
const pids: number[] = [] const pids: number[] = []
const queue = [pid] const queue = [pid]
while (queue.length > 0) { for (let index = 0; index < queue.length; index++) {
const current = queue.shift()! const current = queue[index]
const handle = yield* spawner.spawn(ChildProcess.make("pgrep", ["-P", String(current)], { stdin: "ignore" })) const handle = yield* spawner.spawn(ChildProcess.make("pgrep", ["-P", String(current)], { stdin: "ignore" }))
const text = yield* Stream.mkString(Stream.decodeText(handle.stdout)) const text = yield* Stream.mkString(Stream.decodeText(handle.stdout))
yield* handle.exitCode yield* handle.exitCode
@ -495,7 +482,7 @@ export const layer = Layer.effect(
return pids return pids
}, },
Effect.scoped, Effect.scoped,
Effect.catch(() => Effect.succeed([] as number[])), Effect.catch(() => Effect.succeed(Array<number>())),
) )
function watch(s: State, name: string, client: MCPClient, bridge: EffectBridge.Shape, timeout?: number) { function watch(s: State, name: string, client: MCPClient, bridge: EffectBridge.Shape, timeout?: number) {
@ -540,12 +527,14 @@ export const layer = Layer.effect(
const result = yield* create(key, mcp).pipe(Effect.catch(() => Effect.void)) const result = yield* create(key, mcp).pipe(Effect.catch(() => Effect.void))
if (!result) return if (!result) return
s.status[key] = result.status if (result.type === "unavailable") {
if (result.mcpClient) { s.status[key] = result.status
s.clients[key] = result.mcpClient return
s.defs[key] = result.defs!
watch(s, key, result.mcpClient, bridge, mcp.timeout)
} }
s.status[key] = { status: "connected" }
s.clients[key] = result.client
s.defs[key] = result.defs
watch(s, key, result.client, bridge, mcp.timeout)
}), }),
{ concurrency: "unbounded" }, { concurrency: "unbounded" },
) )
@ -628,14 +617,14 @@ export const layer = Layer.effect(
const s = yield* InstanceState.get(state) const s = yield* InstanceState.get(state)
const result = yield* create(name, mcp) const result = yield* create(name, mcp)
s.status[name] = result.status if (result.type === "unavailable") {
if (!result.mcpClient) { s.status[name] = result.status
yield* closeClient(s, name) yield* closeClient(s, name)
delete s.clients[name] delete s.clients[name]
return result.status return result.status
} }
return yield* storeClient(s, name, result.mcpClient, result.defs!, mcp.timeout) return yield* storeClient(s, name, result.client, result.defs, mcp.timeout)
}) })
const add = Effect.fn("MCP.add")(function* (name: string, mcp: ConfigMCPV1.Info) { const add = Effect.fn("MCP.add")(function* (name: string, mcp: ConfigMCPV1.Info) {
@ -666,92 +655,59 @@ export const layer = Layer.effect(
const config = cfg.mcp ?? {} const config = cfg.mcp ?? {}
const defaultTimeout = cfg.experimental?.mcp_timeout const defaultTimeout = cfg.experimental?.mcp_timeout
const connectedClients = Object.entries(s.clients).filter( for (const [clientName, client] of Object.entries(s.clients)) {
([clientName]) => s.status[clientName]?.status === "connected", if (s.status[clientName]?.status !== "connected") continue
) const mcpConfig = config[clientName]
const entry = mcpConfig && isMcpConfigured(mcpConfig) ? mcpConfig : s.config[clientName]
const listed = s.defs[clientName]
if (!listed) continue
yield* Effect.forEach( const timeout = entry?.timeout ?? defaultTimeout
connectedClients, for (const mcpTool of listed) {
([clientName, client]) => result[sanitize(clientName) + "_" + sanitize(mcpTool.name)] = convertMcpTool(mcpTool, client, timeout)
Effect.gen(function* () { }
const mcpConfig = config[clientName] }
const entry = mcpConfig && isMcpConfigured(mcpConfig) ? mcpConfig : s.config[clientName]
const listed = s.defs[clientName]
if (!listed) {
return
}
const timeout = entry?.timeout ?? defaultTimeout
for (const mcpTool of listed) {
result[sanitize(clientName) + "_" + sanitize(mcpTool.name)] = convertMcpTool(mcpTool, client, timeout)
}
}),
{ concurrency: "unbounded" },
)
return result return result
}) })
function collectFromConnected<T extends { name: string }>( function collectFromConnected<T extends { name: string }>(s: State, listFn: (c: Client) => Promise<T[]>) {
s: State,
listFn: (c: Client) => Promise<T[]>,
label: string,
) {
return Effect.forEach( return Effect.forEach(
Object.entries(s.clients).filter(([name]) => s.status[name]?.status === "connected"), Object.entries(s.clients).filter(([name]) => s.status[name]?.status === "connected"),
([clientName, client]) => ([clientName, client]) =>
fetchFromClient(clientName, client, listFn, label).pipe(Effect.map((items) => Object.entries(items ?? {}))), fetchFromClient(clientName, client, listFn).pipe(Effect.map((items) => Object.entries(items ?? {}))),
{ concurrency: "unbounded" }, { concurrency: "unbounded" },
).pipe(Effect.map((results) => Object.fromEntries<T & { client: string }>(results.flat()))) ).pipe(Effect.map((results) => Object.fromEntries<T & { client: string }>(results.flat())))
} }
const prompts = Effect.fn("MCP.prompts")(function* () { const prompts = Effect.fn("MCP.prompts")(function* () {
const s = yield* InstanceState.get(state) const s = yield* InstanceState.get(state)
return yield* collectFromConnected( return yield* collectFromConnected(s, (c) =>
s, c.getServerCapabilities()?.prompts
(c) => ? paginate(
c.getServerCapabilities()?.prompts (cursor) => c.listPrompts(cursor === undefined ? undefined : { cursor }),
? paginate( (result) => result.prompts,
(cursor) => c.listPrompts(cursor === undefined ? undefined : { cursor }), )
(result) => result.prompts, : Promise.resolve([]),
)
: Promise.resolve([]),
"prompts",
) )
}) })
const resources = Effect.fn("MCP.resources")(function* () { const resources = Effect.fn("MCP.resources")(function* () {
const s = yield* InstanceState.get(state) const s = yield* InstanceState.get(state)
return yield* collectFromConnected( return yield* collectFromConnected(s, (c) =>
s, c.getServerCapabilities()?.resources
(c) => ? paginate(
c.getServerCapabilities()?.resources (cursor) => c.listResources(cursor === undefined ? undefined : { cursor }),
? paginate( (result) => result.resources,
(cursor) => c.listResources(cursor === undefined ? undefined : { cursor }), )
(result) => result.resources, : Promise.resolve([]),
)
: Promise.resolve([]),
"resources",
) )
}) })
const withClient = Effect.fnUntraced(function* <A>( const withClient = Effect.fnUntraced(function* <A>(clientName: string, fn: (client: MCPClient) => Promise<A>) {
clientName: string,
fn: (client: MCPClient) => Promise<A>,
label: string,
meta?: Record<string, unknown>,
) {
const s = yield* InstanceState.get(state) const s = yield* InstanceState.get(state)
const client = s.clients[clientName] const client = s.clients[clientName]
if (!client) { if (!client) return undefined
return undefined return yield* Effect.tryPromise(() => fn(client)).pipe(Effect.orElseSucceed(() => undefined))
}
return yield* Effect.tryPromise({
try: () => fn(client),
catch: (e: any) => {
return e
},
}).pipe(Effect.orElseSucceed(() => undefined))
}) })
const getPrompt = Effect.fn("MCP.getPrompt")(function* ( const getPrompt = Effect.fn("MCP.getPrompt")(function* (
@ -759,15 +715,11 @@ export const layer = Layer.effect(
name: string, name: string,
args?: Record<string, string>, args?: Record<string, string>,
) { ) {
return yield* withClient(clientName, (client) => client.getPrompt({ name, arguments: args }), "getPrompt", { return yield* withClient(clientName, (client) => client.getPrompt({ name, arguments: args }))
promptName: name,
})
}) })
const readResource = Effect.fn("MCP.readResource")(function* (clientName: string, resourceUri: string) { const readResource = Effect.fn("MCP.readResource")(function* (clientName: string, resourceUri: string) {
return yield* withClient(clientName, (client) => client.readResource({ uri: resourceUri }), "readResource", { return yield* withClient(clientName, (client) => client.readResource({ uri: resourceUri }))
resourceUri,
})
}) })
const getMcpConfig = Effect.fnUntraced(function* (mcpName: string) { const getMcpConfig = Effect.fnUntraced(function* (mcpName: string) {
@ -790,7 +742,7 @@ export const layer = Layer.effect(
const mcpConfig = yield* requireMcpConfig(mcpName) const mcpConfig = yield* requireMcpConfig(mcpName)
if (mcpConfig.type !== "remote") throw new Error(`MCP server ${mcpName} is not a remote server`) if (mcpConfig.type !== "remote") throw new Error(`MCP server ${mcpName} is not a remote server`)
if (mcpConfig.oauth === false) throw new Error(`MCP server ${mcpName} has OAuth explicitly disabled`) if (mcpConfig.oauth === false) throw new Error(`MCP server ${mcpName} has OAuth explicitly disabled`)
const url = remoteURL(mcpName, mcpConfig.url) const url = remoteURL(mcpConfig.url)
if (!url) throw new Error(`Invalid MCP URL for "${mcpName}"`) if (!url) throw new Error(`Invalid MCP URL for "${mcpName}"`)
// OAuth config is optional - if not provided, we'll use auto-discovery // OAuth config is optional - if not provided, we'll use auto-discovery
@ -862,7 +814,7 @@ export const layer = Layer.effect(
: undefined : undefined
if (!client || !listed) { if (!client || !listed) {
yield* Effect.tryPromise(() => client?.close() ?? Promise.resolve()).pipe(Effect.ignore) yield* Effect.tryPromise(() => client?.close() ?? Promise.resolve()).pipe(Effect.ignore)
return { status: "failed", error: "Failed to get tools" } as Status return { status: "failed", error: "Failed to get tools" } satisfies Status
} }
const s = yield* InstanceState.get(state) const s = yield* InstanceState.get(state)
@ -910,14 +862,14 @@ export const layer = Layer.effect(
if (!transport) throw new Error(`No pending OAuth flow for MCP server: ${mcpName}`) if (!transport) throw new Error(`No pending OAuth flow for MCP server: ${mcpName}`)
const result = yield* Effect.tryPromise({ const result = yield* Effect.tryPromise({
try: () => transport.finishAuth(authorizationCode).then(() => true as const), try: () => transport.finishAuth(authorizationCode).then((): true => true),
catch: (error) => { catch: (error) => {
return error return error
}, },
}).pipe(Effect.option) }).pipe(Effect.option)
if (Option.isNone(result)) { if (Option.isNone(result)) {
return { status: "failed", error: "OAuth completion failed" } as Status return { status: "failed", error: "OAuth completion failed" } satisfies Status
} }
yield* auth.clearCodeVerifier(mcpName) yield* auth.clearCodeVerifier(mcpName)
@ -946,9 +898,9 @@ export const layer = Layer.effect(
const getAuthStatus = Effect.fn("MCP.getAuthStatus")(function* (mcpName: string) { const getAuthStatus = Effect.fn("MCP.getAuthStatus")(function* (mcpName: string) {
const entry = yield* auth.get(mcpName) const entry = yield* auth.get(mcpName)
if (!entry?.tokens) return "not_authenticated" as AuthStatus if (!entry?.tokens) return "not_authenticated"
const expired = yield* auth.isTokenExpired(mcpName) const expired = yield* auth.isTokenExpired(mcpName)
return (expired ? "expired" : "authenticated") as AuthStatus return expired ? "expired" : "authenticated"
}) })
return Service.of({ return Service.of({