The MCP (Model Context Protocol) subsystem enables tool-augmented LLM generation. It manages tool discovery, session pooling, and parallel tool execution for column generators that use external tools.
Source: packages/data-designer-engine/src/data_designer/engine/mcp/
MCP integration allows column generators to augment LLM completions with tool calls. When a column config specifies a tool_alias, the model facade routes tool calls through the MCP subsystem, which handles session management, tool schema discovery, and parallel execution.
The subsystem has three layers:
MCPIOService— low-level I/O: session pooling, tool listing, tool execution on a background async loopMCPFacade— scoped to a single tool config: schema formatting, completion response processing, tool call executionMCPRegistry— maps tool aliases to configs, lazy facade construction, health checks
The io.py module exposes MCP I/O through one shared MCPIOService instance (_MCP_IO_SERVICE) created at import; atexit registers shutdown. Async state (loop, sessions, caches) lives on that instance.
- Background async loop — runs on a daemon thread; sync callers use
asyncio.run_coroutine_threadsafeto bridge - Session pool —
_sessionskeyed by provider cache key (JSON of provider config);_get_or_create_sessionwith in-flight deduplication prevents redundant connections - Tool listing — cached per session; coalescing for concurrent list requests via
_inflight_toolsprevents duplicate discovery calls - Tool execution — parallel tool calls within a single completion response
Module-level functions (list_tools, call_tools, clear_session_pool) delegate to _MCP_IO_SERVICE.
Scoped to one ToolConfig. Provides the interface that ModelFacade uses:
get_tool_schemas()— returns tool schemas in OpenAI function-calling formatprocess_completion_response— extracts tool calls from a completion, executes them in parallel viaMCPIOService, returnsChatMessagelist with resultsrefuse_completion_response— handles tool-call turn limits (prevents infinite tool loops)
Tool result messages may contain either text or ordered multimodal content blocks. MCP image results, and generic
base64 payloads with image/* MIME metadata or image data URI prefixes, are preserved as canonical image_url data
URI blocks and translated by provider adapters at the API boundary. Models need VLM-capable provider support to
interpret those image results semantically.
Maps tool_alias → ToolConfig. Lazy MCPFacade construction mirrors ModelRegistry. Provides health checks for configured tools.
- Column config declares
tool_aliasreferencing a configured MCP tool - Generator's
ModelFacadeincludes tool schemas in the completion request - LLM returns a completion with tool calls
ModelFacadedelegates toMCPFacade.process_completion_responseMCPFacadeextracts tool calls, executes them in parallel viaMCPIOService- Tool results are formatted as
ChatMessages and fed back to the LLM for another completion round - Process repeats until the LLM produces a final response or the turn limit is reached
- Single background async loop avoids creating event loops per request. All MCP I/O funnels through one loop on a daemon thread, with sync callers bridging via
run_coroutine_threadsafe. - Session pooling with in-flight deduplication prevents redundant connections when multiple generators discover tools from the same provider concurrently.
- Tool schema coalescing — concurrent
list_toolscalls for the same session share a single in-flight request, reducing startup latency when many columns use the same tool. - Turn limits on tool loops prevent runaway tool-call chains.
refuse_completion_responsegracefully terminates when the limit is reached.
- System Architecture — where MCP fits in the stack
- Models — how
ModelFacadeintegrates MCP tool loops - Engine Layer —
ResourceProviderprovidesMCPRegistryto generators - Config Layer —
ToolConfigdefinition