1  import type { McpServerConfig } from '@/types'
  2  import { hmrSingleton } from '@/lib/shared-utils'
  3  
  4  /**
  5   * Gateway-style runtime for SwarmClaw's MCP integration. Mirrors the behavior
  6   * of `@swarmclawai/mcp-gateway`'s router (alwaysExpose filter + mcp_tool_search
  7   * meta-tool) but stays inside SwarmClaw — no external dep on mcp-core (yet).
  8   *
  9   * Once `@swarmclawai/mcp-core` is published to npm, the plan is:
 10   *   1. Add `@swarmclawai/mcp-core` to SwarmClaw's dependencies.
 11   *   2. Replace `SessionToolPromoter`, `searchDiscoveredTools`, and
 12   *      `shouldExposeMcpTool` with imports from mcp-core. The field names on
 13   *      `DownstreamTool` differ from `DiscoveredTool` (`prefixedName` vs
 14   *      `langChainName`) — add a thin adapter during migration.
 15   *   3. Keep the hmrSingleton `state` here; mcp-core provides primitives, but
 16   *      SwarmClaw still owns process-wide session-scoped storage.
 17   *
 18   * Contains two pieces of shared state, both HMR-safe:
 19   *   - `SessionToolPromoter` instances keyed by sessionId. The agent calls
 20   *     `mcp_tool_search` to promote a lazy tool by name; the next turn's tool
 21   *     bind picks up the promoted name via `isPromoted`.
 22   *   - A discovery cache keyed by MCP server id, so even lazy servers have
 23   *     their tool schemas known in-process for `mcp_tool_search` to match
 24   *     against without a cold connect.
 25   */
 26  
 27  export interface DiscoveredTool {
 28    name: string // bare tool name as the downstream reported it
 29    langChainName: string // the `mcp_<server>_<tool>` name SwarmClaw binds it under
 30    description?: string
 31    inputSchema?: unknown
 32    serverId: string
 33    serverName: string
 34  }
 35  
 36  export class SessionToolPromoter {
 37    private readonly exposed = new Set<string>()
 38    allow(langChainName: string): boolean { return this.exposed.has(langChainName) }
 39    promote(langChainName: string): void { this.exposed.add(langChainName) }
 40    promoteMany(names: readonly string[]): void { for (const n of names) this.exposed.add(n) }
 41    promoted(): string[] { return Array.from(this.exposed) }
 42    clear(): void { this.exposed.clear() }
 43  }
 44  
 45  interface RuntimeState {
 46    promoters: Map<string, SessionToolPromoter>
 47    // serverId -> discovered tools (updated opportunistically when we connect)
 48    discovered: Map<string, DiscoveredTool[]>
 49  }
 50  
 51  const state = hmrSingleton<RuntimeState>('mcpGatewayRuntime', () => ({
 52    promoters: new Map<string, SessionToolPromoter>(),
 53    discovered: new Map<string, DiscoveredTool[]>(),
 54  }))
 55  
 56  export function getPromoter(sessionId: string): SessionToolPromoter {
 57    let p = state.promoters.get(sessionId)
 58    if (!p) {
 59      p = new SessionToolPromoter()
 60      state.promoters.set(sessionId, p)
 61    }
 62    return p
 63  }
 64  
 65  export function clearPromoter(sessionId: string): void {
 66    state.promoters.delete(sessionId)
 67  }
 68  
 69  export function recordDiscoveredTools(serverId: string, tools: DiscoveredTool[]): void {
 70    state.discovered.set(serverId, tools)
 71  }
 72  
 73  export function allDiscoveredTools(): DiscoveredTool[] {
 74    const out: DiscoveredTool[] = []
 75    for (const arr of state.discovered.values()) {
 76      for (const t of arr) out.push(t)
 77    }
 78    return out
 79  }
 80  
 81  /**
 82   * Decide whether a given tool, from a given server, should be bound on this
 83   * turn. Order of precedence:
 84   *   1. Per-agent eager allowlist (`mcpEagerTools`) — agent-scoped override
 85   *   2. Server-level `alwaysExpose`
 86   *      - true (default)  → bind
 87   *      - false           → skip unless promoted
 88   *      - string[]        → bind only if tool name is on the list
 89   *   3. Session promoter — if the agent has called `mcp_tool_search` this
 90   *      session and promoted this tool, bind it regardless of (2).
 91   */
 92  export function shouldExposeMcpTool(opts: {
 93    server: McpServerConfig
 94    toolName: string
 95    langChainName: string
 96    agentEagerTools?: readonly string[] | null
 97    promoter?: SessionToolPromoter | null
 98  }): boolean {
 99    const { server, toolName, langChainName, agentEagerTools, promoter } = opts
100    if (agentEagerTools && agentEagerTools.includes(toolName)) return true
101    if (agentEagerTools && agentEagerTools.includes(langChainName)) return true
102    if (promoter?.allow(langChainName)) return true
103    const mode = server.alwaysExpose
104    if (mode === undefined || mode === true) return true
105    if (mode === false) return false
106    if (Array.isArray(mode)) return mode.includes(toolName)
107    return true
108  }
109  
110  export interface ToolSearchMatch {
111    name: string // langChainName
112    server: string
113    description?: string
114    score: number
115  }
116  
117  export function searchDiscoveredTools(query: string, limit = 8): ToolSearchMatch[] {
118    const q = query.trim().toLowerCase()
119    if (!q) return []
120    const terms = q.split(/\s+/).filter((t) => t.length >= 2)
121    const clamped = Math.max(1, Math.min(limit, 50))
122    const scored = allDiscoveredTools().map((t): ToolSearchMatch => {
123      const haystack = `${t.langChainName} ${t.description ?? ''}`.toLowerCase()
124      let score = 0
125      if (haystack.includes(q)) score += 0.6
126      let termHits = 0
127      for (const term of terms) if (haystack.includes(term)) termHits += 1
128      if (terms.length) score += 0.4 * (termHits / terms.length)
129      return {
130        name: t.langChainName,
131        server: t.serverName,
132        description: t.description,
133        score: Math.min(1, score),
134      }
135    })
136    scored.sort((a, b) => b.score - a.score)
137    return scored.filter((m) => m.score > 0).slice(0, clamped)
138  }