diff --git a/cli/src/utils/auth.ts b/cli/src/utils/auth.ts index 7d81f48ef9..cbf7ae9daa 100644 --- a/cli/src/utils/auth.ts +++ b/cli/src/utils/auth.ts @@ -14,8 +14,8 @@ import type { CiEnv } from '@codebuff/common/types/contracts/env' // User schema const userSchema = z.object({ id: z.string().optional(), - name: z.string(), - email: z.string(), + name: z.string().nullish(), + email: z.string().nullish(), authToken: z.string(), fingerprintId: z.string().optional(), fingerprintHash: z.string().optional(), diff --git a/common/src/mcp/__tests__/client-pool.test.ts b/common/src/mcp/__tests__/client-pool.test.ts new file mode 100644 index 0000000000..c19c979df9 --- /dev/null +++ b/common/src/mcp/__tests__/client-pool.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, it } from 'bun:test' + +import { MCPClientPool } from '../client-pool' + +type Config = { id: string } +type FakeClient = { id: string } + +describe('MCPClientPool', () => { + it('deduplicates concurrent connections and reports ready status', async () => { + let connects = 0 + const pool = new MCPClientPool({ + keyOf: (config) => config.id, + connect: async (config) => { + connects++ + await Bun.sleep(5) + return { id: config.id } + }, + close: async () => {}, + }) + + const [first, second] = await Promise.all([ + pool.get({ id: 'docs' }), + pool.get({ id: 'docs' }), + ]) + + expect(connects).toBe(1) + expect(first.client).toBe(second.client) + expect(pool.statuses()).toEqual([ + expect.objectContaining({ id: 'docs', state: 'ready' }), + ]) + }) + + it('removes failed connections so the next request can retry', async () => { + let attempts = 0 + const pool = new MCPClientPool({ + keyOf: (config) => config.id, + connect: async (config) => { + attempts++ + if (attempts === 1) throw new Error('offline') + return { id: config.id } + }, + close: async () => {}, + }) + + await expect(pool.get({ id: 'retry' })).rejects.toThrow('offline') + expect((await pool.get({ id: 'retry' })).client.id).toBe('retry') + expect(attempts).toBe(2) + }) + + it('closes one or every live client', async () => { + const closed: string[] = [] + const pool = new MCPClientPool({ + keyOf: (config) => config.id, + connect: async (config) => ({ id: config.id }), + close: async (client) => { + closed.push(client.id) + }, + }) + + await pool.get({ id: 'one' }) + await pool.get({ id: 'two' }) + expect(await pool.close('one')).toBe(true) + await pool.closeAll() + + expect(closed.sort()).toEqual(['one', 'two']) + expect(pool.statuses()).toEqual([]) + }) + + it('times out a stalled connection and allows a later retry', async () => { + let shouldHang = true + const pool = new MCPClientPool( + { + keyOf: (config) => config.id, + connect: async (config) => { + if (shouldHang) await new Promise(() => {}) + return { id: config.id } + }, + close: async () => {}, + }, + { connectTimeoutMs: 10 }, + ) + + await expect(pool.get({ id: 'slow' })).rejects.toThrow( + 'MCP connection timed out', + ) + shouldHang = false + expect((await pool.get({ id: 'slow' })).client.id).toBe('slow') + }) +}) diff --git a/common/src/mcp/client-pool.ts b/common/src/mcp/client-pool.ts new file mode 100644 index 0000000000..da63481613 --- /dev/null +++ b/common/src/mcp/client-pool.ts @@ -0,0 +1,125 @@ +export type MCPClientPoolStatus = { + id: string + state: 'connecting' | 'ready' + connectedAt: number | null + lastUsedAt: number +} + +type MCPClientPoolAdapter = { + keyOf: (config: TConfig) => string + connect: (config: TConfig) => Promise + close: (client: TClient) => Promise +} + +type MCPClientPoolOptions = { + connectTimeoutMs?: number +} + +type PoolEntry = { + client: TClient | null + connecting: Promise + connectedAt: number | null + lastUsedAt: number +} + +export function withTimeout( + promise: Promise, + timeoutMs: number, + message: string, +): Promise { + if (timeoutMs <= 0) return promise + + let timeout: ReturnType | undefined + const rejection = new Promise((_, reject) => { + timeout = setTimeout(() => reject(new Error(message)), timeoutMs) + }) + return Promise.race([promise, rejection]).finally(() => { + if (timeout) clearTimeout(timeout) + }) +} + +/** + * Reuses MCP transports across turns and owns their complete lifecycle. + * Concurrent callers for the same config share one handshake. + */ +export class MCPClientPool { + private readonly entries = new Map>() + private readonly connectTimeoutMs: number + + constructor( + private readonly adapter: MCPClientPoolAdapter, + options: MCPClientPoolOptions = {}, + ) { + this.connectTimeoutMs = options.connectTimeoutMs ?? 30_000 + } + + async get(config: TConfig): Promise<{ id: string; client: TClient }> { + const id = this.adapter.keyOf(config) + const existing = this.entries.get(id) + if (existing) { + existing.lastUsedAt = Date.now() + return { id, client: existing.client ?? (await existing.connecting) } + } + + const now = Date.now() + const rawConnection = this.adapter.connect(config) + const connecting = withTimeout( + rawConnection, + this.connectTimeoutMs, + `MCP connection timed out after ${this.connectTimeoutMs}ms`, + ) + const entry: PoolEntry = { + client: null, + connecting, + connectedAt: null, + lastUsedAt: now, + } + this.entries.set(id, entry) + + try { + const client = await connecting + entry.client = client + entry.connectedAt = Date.now() + entry.lastUsedAt = entry.connectedAt + return { id, client } + } catch (error) { + if (this.entries.get(id) === entry) this.entries.delete(id) + // A timed-out transport can still finish later. Close it instead of + // leaking a child process or socket no caller can reach. + rawConnection.then(this.adapter.close).catch(() => {}) + throw error + } + } + + getReady(id: string): TClient | undefined { + const entry = this.entries.get(id) + if (!entry?.client) return undefined + entry.lastUsedAt = Date.now() + return entry.client + } + + statuses(): MCPClientPoolStatus[] { + return [...this.entries.entries()] + .map(([id, entry]) => ({ + id, + state: entry.client ? ('ready' as const) : ('connecting' as const), + connectedAt: entry.connectedAt, + lastUsedAt: entry.lastUsedAt, + })) + .sort((a, b) => a.id.localeCompare(b.id)) + } + + async close(id: string): Promise { + const entry = this.entries.get(id) + if (!entry) return false + this.entries.delete(id) + const client = entry.client ?? (await entry.connecting.catch(() => null)) + if (client) await this.adapter.close(client) + return true + } + + async closeAll(): Promise { + const ids = [...this.entries.keys()] + await Promise.allSettled(ids.map((id) => this.close(id))) + } +}