perf(agent): load MCP tools asynchronously to speed up startup

Load MCP server tools in the background so the UI appears immediately
instead of blocking until all servers connect. The first LLM call
automatically waits for tools to be ready before proceeding.

Key changes:
- NewAgent() starts MCP loading in a background goroutine and returns
  immediately with core/extension tools only
- GenerateWithLoop() calls ensureMCPTools() to lazily wait and rebuild
  the fantasy agent with full tool set before first LLM call
- Parallelize LoadTools() across all configured MCP servers
- Add WaitForMCPTools() and MCPToolsReady() for status checking
- Refactor SetModel/SetExtraTools to use shared rebuildFantasyAgent()
- Expose async MCP status methods in public SDK
This commit is contained in:
Ed Zynda
2026-04-07 13:33:54 +03:00
parent 88870be4d2
commit 7e54710d4a
3 changed files with 252 additions and 155 deletions
+180 -140
View File
@@ -88,6 +88,10 @@ type StepUsageHandler func(inputTokens, outputTokens, cacheReadTokens, cacheCrea
// Core tools (bash, read, write, edit, grep, find, ls) are registered as direct
// AgentTool implementations — no MCP layer, no serialization overhead.
// Additional tools from external MCP servers can be loaded alongside core tools.
//
// When MCP servers are configured, tool loading happens in the background so the
// agent (and UI) can start immediately. The first LLM call automatically waits
// for MCP tools to finish loading before proceeding.
type Agent struct {
toolManager *tools.MCPToolManager
fantasyAgent fantasy.Agent
@@ -101,6 +105,18 @@ type Agent struct {
coreTools []fantasy.AgentTool
extraTools []fantasy.AgentTool
toolWrapper func([]fantasy.AgentTool) []fantasy.AgentTool // stored for SetModel rebuild
// providerOptions and modelConfig are stored for rebuilding the fantasy
// agent when MCP tools arrive asynchronously or on SetModel.
providerOptions fantasy.ProviderOptions
skipMaxOutputTokens bool
modelConfig *models.ProviderConfig
// mcpReady is closed when background MCP tool loading completes (success
// or failure). nil when no MCP servers are configured.
mcpReady chan struct{}
// mcpErr holds any error from background MCP loading.
mcpErr error
}
// GenerateWithLoopResult contains the result and conversation history from an agent interaction.
@@ -119,7 +135,10 @@ type GenerateWithLoopResult struct {
// NewAgent creates a new Agent with core tools and optional MCP tool integration.
// Core tools (bash, read, write, edit, grep, find, ls) are always registered.
// External MCP tools are loaded from the config if any MCP servers are configured.
// If MCP servers are configured, their tools are loaded in the background —
// the agent returns immediately and is usable with core tools only. The first
// LLM call (GenerateWithLoop) automatically waits for MCP tools to finish
// loading and rebuilds the agent with the full tool set.
func NewAgent(ctx context.Context, agentConfig *AgentConfig) (*Agent, error) {
// Create the LLM provider
providerResult, err := models.CreateProvider(ctx, agentConfig.ModelConfig)
@@ -134,33 +153,9 @@ func NewAgent(ctx context.Context, agentConfig *AgentConfig) (*Agent, error) {
coreTools = core.AllTools()
}
// Build the combined tool list: core tools + any external MCP tools
// Build the initial tool list: core tools + extension tools (no MCP yet).
allTools := make([]fantasy.AgentTool, len(coreTools))
copy(allTools, coreTools)
// Load external MCP tools if configured
var toolManager *tools.MCPToolManager
if agentConfig.MCPConfig != nil && len(agentConfig.MCPConfig.MCPServers) > 0 {
toolManager = tools.NewMCPToolManager()
toolManager.SetModel(providerResult.Model)
if agentConfig.AuthHandler != nil {
toolManager.SetAuthHandler(agentConfig.AuthHandler)
}
if agentConfig.DebugLogger != nil {
toolManager.SetDebugLogger(agentConfig.DebugLogger)
}
if err := toolManager.LoadTools(ctx, agentConfig.MCPConfig); err != nil {
// MCP tool loading failures are non-fatal; core tools still work
fmt.Printf("Warning: Failed to load MCP tools: %v\n", err)
} else {
mcpTools := toolManager.GetTools()
allTools = append(allTools, mcpTools...)
}
}
// Append any extra tools provided by extensions.
if len(agentConfig.ExtraTools) > 0 {
allTools = append(allTools, agentConfig.ExtraTools...)
@@ -172,6 +167,140 @@ func NewAgent(ctx context.Context, agentConfig *AgentConfig) (*Agent, error) {
}
// Build agent options
agentOpts := buildAgentOptions(agentConfig, providerResult, allTools)
// Create the agent
fantasyAgent := fantasy.NewAgent(providerResult.Model, agentOpts...)
// Determine provider type from model string
providerType := "default"
if agentConfig.ModelConfig != nil && agentConfig.ModelConfig.ModelString != "" {
if p, _, err := models.ParseModelString(agentConfig.ModelConfig.ModelString); err == nil {
providerType = p
}
}
a := &Agent{
fantasyAgent: fantasyAgent,
model: providerResult.Model,
providerCloser: providerResult.Closer,
maxSteps: agentConfig.MaxSteps,
systemPrompt: agentConfig.SystemPrompt,
loadingMessage: providerResult.Message,
providerType: providerType,
streamingEnabled: agentConfig.StreamingEnabled,
coreTools: coreTools,
extraTools: agentConfig.ExtraTools,
toolWrapper: agentConfig.ToolWrapper,
providerOptions: providerResult.ProviderOptions,
skipMaxOutputTokens: providerResult.SkipMaxOutputTokens,
modelConfig: agentConfig.ModelConfig,
}
// Start MCP tool loading in the background if servers are configured.
// The mcpReady channel is closed when loading completes (success or failure).
if agentConfig.MCPConfig != nil && len(agentConfig.MCPConfig.MCPServers) > 0 {
toolManager := tools.NewMCPToolManager()
toolManager.SetModel(providerResult.Model)
if agentConfig.AuthHandler != nil {
toolManager.SetAuthHandler(agentConfig.AuthHandler)
}
if agentConfig.DebugLogger != nil {
toolManager.SetDebugLogger(agentConfig.DebugLogger)
}
a.toolManager = toolManager
a.mcpReady = make(chan struct{})
go func() {
defer close(a.mcpReady)
if err := toolManager.LoadTools(ctx, agentConfig.MCPConfig); err != nil {
a.mcpErr = err
fmt.Printf("Warning: Failed to load MCP tools: %v\n", err)
}
}()
}
return a, nil
}
// WaitForMCPTools blocks until background MCP tool loading completes.
// Returns nil if no MCP servers are configured or if loading succeeded.
// Returns the loading error if all servers failed. Safe to call multiple times.
func (a *Agent) WaitForMCPTools() error {
if a.mcpReady == nil {
return nil
}
<-a.mcpReady
return a.mcpErr
}
// MCPToolsReady returns true if MCP tool loading has completed (or was never
// started). This is a non-blocking check useful for UI status display.
func (a *Agent) MCPToolsReady() bool {
if a.mcpReady == nil {
return true
}
select {
case <-a.mcpReady:
return true
default:
return false
}
}
// ensureMCPTools waits for MCP tools to load and rebuilds the fantasy agent
// with the full tool set. Called lazily before the first LLM call.
// This is idempotent — subsequent calls after the first rebuild are no-ops.
func (a *Agent) ensureMCPTools() {
if a.mcpReady == nil {
return
}
<-a.mcpReady
// If there are MCP tools, rebuild the fantasy agent to include them.
if a.toolManager != nil && len(a.toolManager.GetTools()) > 0 {
a.rebuildFantasyAgent()
}
// Nil out the channel so future calls are instant no-ops and we
// don't rebuild again.
a.mcpReady = nil
}
// rebuildFantasyAgent reconstructs the fantasy agent with the current full
// tool set (core + MCP + extension tools). Used after MCP tools arrive
// asynchronously and by SetModel.
func (a *Agent) rebuildFantasyAgent() {
allTools := make([]fantasy.AgentTool, len(a.coreTools))
copy(allTools, a.coreTools)
if a.toolManager != nil {
allTools = append(allTools, a.toolManager.GetTools()...)
}
if len(a.extraTools) > 0 {
allTools = append(allTools, a.extraTools...)
}
if a.toolWrapper != nil {
allTools = a.toolWrapper(allTools)
}
providerResult := &models.ProviderResult{
Model: a.model,
ProviderOptions: a.providerOptions,
SkipMaxOutputTokens: a.skipMaxOutputTokens,
}
agentOpts := buildAgentOptions(&AgentConfig{
ModelConfig: a.modelConfig,
SystemPrompt: a.systemPrompt,
MaxSteps: a.maxSteps,
}, providerResult, allTools)
a.fantasyAgent = fantasy.NewAgent(a.model, agentOpts...)
}
// buildAgentOptions constructs the fantasy.AgentOption slice from config,
// provider result, and the combined tool list. Shared by NewAgent,
// rebuildFantasyAgent, and SetModel.
func buildAgentOptions(agentConfig *AgentConfig, providerResult *models.ProviderResult, allTools []fantasy.AgentTool) []fantasy.AgentOption {
var agentOpts []fantasy.AgentOption
if agentConfig.SystemPrompt != "" {
@@ -217,31 +346,7 @@ func NewAgent(ctx context.Context, agentConfig *AgentConfig) (*Agent, error) {
}
}
// Create the agent
fantasyAgent := fantasy.NewAgent(providerResult.Model, agentOpts...)
// Determine provider type from model string
providerType := "default"
if agentConfig.ModelConfig != nil && agentConfig.ModelConfig.ModelString != "" {
if p, _, err := models.ParseModelString(agentConfig.ModelConfig.ModelString); err == nil {
providerType = p
}
}
return &Agent{
toolManager: toolManager,
fantasyAgent: fantasyAgent,
model: providerResult.Model,
providerCloser: providerResult.Closer,
maxSteps: agentConfig.MaxSteps,
systemPrompt: agentConfig.SystemPrompt,
loadingMessage: providerResult.Message,
providerType: providerType,
streamingEnabled: agentConfig.StreamingEnabled,
coreTools: coreTools,
extraTools: agentConfig.ExtraTools,
toolWrapper: agentConfig.ToolWrapper,
}, nil
return agentOpts
}
// GenerateWithLoop processes messages with a custom loop that displays tool calls in real-time.
@@ -266,6 +371,11 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
onStepUsage StepUsageHandler,
) (*GenerateWithLoopResult, error) {
// Wait for background MCP tool loading to complete and rebuild the
// fantasy agent with the full tool set. This is a no-op when no MCP
// servers are configured or tools have already been integrated.
a.ensureMCPTools()
// Inject tool output handler into context for use by core tools (e.g., bash).
if onToolOutput != nil {
ctx = core.ContextWithToolOutputCallback(ctx, onToolOutput)
@@ -657,38 +767,9 @@ func (a *Agent) GetExtensionToolCount() int {
// SetExtraTools replaces the agent's extra tools (e.g. extension-registered
// tools) and rebuilds the internal agent with the updated tool list. The
// model, system prompt, and all other configuration are preserved.
func (a *Agent) SetExtraTools(tools []fantasy.AgentTool) {
a.extraTools = tools
// Rebuild tool list (same as NewAgent / SetModel).
allTools := make([]fantasy.AgentTool, len(a.coreTools))
copy(allTools, a.coreTools)
if a.toolManager != nil {
allTools = append(allTools, a.toolManager.GetTools()...)
}
if len(a.extraTools) > 0 {
allTools = append(allTools, a.extraTools...)
}
if a.toolWrapper != nil {
allTools = a.toolWrapper(allTools)
}
// Rebuild agent options with the existing model.
var agentOpts []fantasy.AgentOption
if a.systemPrompt != "" {
agentOpts = append(agentOpts, fantasy.WithSystemPrompt(a.systemPrompt))
}
if len(allTools) > 0 {
agentOpts = append(agentOpts, fantasy.WithTools(allTools...))
}
if a.maxSteps > 0 {
agentOpts = append(agentOpts, fantasy.WithStopConditions(
fantasy.StepCountIs(a.maxSteps),
))
}
// Swap the fantasy agent (model and provider are unchanged).
a.fantasyAgent = fantasy.NewAgent(a.model, agentOpts...)
func (a *Agent) SetExtraTools(extraTools []fantasy.AgentTool) {
a.extraTools = extraTools
a.rebuildFantasyAgent()
}
// GetLoadingMessage returns the loading message from provider creation.
@@ -708,66 +789,14 @@ func (a *Agent) GetLoadedServerNames() []string {
// system prompt, and configuration are preserved. The old provider is closed
// if it has a closer. Returns the previous model string for notification.
func (a *Agent) SetModel(ctx context.Context, config *models.ProviderConfig) error {
// Ensure MCP tools are loaded before rebuilding (SetModel may be called
// before the first LLM call).
a.ensureMCPTools()
providerResult, err := models.CreateProvider(ctx, config)
if err != nil {
return fmt.Errorf("failed to create model provider: %v", err)
}
// Rebuild tool list (same as NewAgent).
allTools := make([]fantasy.AgentTool, len(a.coreTools))
copy(allTools, a.coreTools)
if a.toolManager != nil {
allTools = append(allTools, a.toolManager.GetTools()...)
}
if len(a.extraTools) > 0 {
allTools = append(allTools, a.extraTools...)
}
if a.toolWrapper != nil {
allTools = a.toolWrapper(allTools)
}
// Rebuild agent options.
var agentOpts []fantasy.AgentOption
if a.systemPrompt != "" {
agentOpts = append(agentOpts, fantasy.WithSystemPrompt(a.systemPrompt))
}
if len(allTools) > 0 {
agentOpts = append(agentOpts, fantasy.WithTools(allTools...))
}
if a.maxSteps > 0 {
agentOpts = append(agentOpts, fantasy.WithStopConditions(
fantasy.StepCountIs(a.maxSteps),
))
}
// Pass provider-specific options (e.g. OpenAI Responses API reasoning settings).
if providerResult.ProviderOptions != nil {
agentOpts = append(agentOpts, fantasy.WithProviderOptions(providerResult.ProviderOptions))
}
// Pass generation parameters when available.
// Skip max_output_tokens for providers that don't support it (e.g., Codex OAuth)
if config.MaxTokens > 0 && !providerResult.SkipMaxOutputTokens {
agentOpts = append(agentOpts, fantasy.WithMaxOutputTokens(int64(config.MaxTokens)))
}
if config.Temperature != nil {
agentOpts = append(agentOpts, fantasy.WithTemperature(float64(*config.Temperature)))
}
if config.TopP != nil {
agentOpts = append(agentOpts, fantasy.WithTopP(float64(*config.TopP)))
}
if config.TopK != nil {
agentOpts = append(agentOpts, fantasy.WithTopK(int64(*config.TopK)))
}
if config.FrequencyPenalty != nil {
agentOpts = append(agentOpts, fantasy.WithFrequencyPenalty(float64(*config.FrequencyPenalty)))
}
if config.PresencePenalty != nil {
agentOpts = append(agentOpts, fantasy.WithPresencePenalty(float64(*config.PresencePenalty)))
}
newFantasyAgent := fantasy.NewAgent(providerResult.Model, agentOpts...)
// Close old provider.
if a.providerCloser != nil {
_ = a.providerCloser.Close()
@@ -779,9 +808,11 @@ func (a *Agent) SetModel(ctx context.Context, config *models.ProviderConfig) err
}
// Swap fields.
a.fantasyAgent = newFantasyAgent
a.model = providerResult.Model
a.providerCloser = providerResult.Closer
a.providerOptions = providerResult.ProviderOptions
a.skipMaxOutputTokens = providerResult.SkipMaxOutputTokens
a.modelConfig = config
// Update provider type.
if config.ModelString != "" {
@@ -790,6 +821,9 @@ func (a *Agent) SetModel(ctx context.Context, config *models.ProviderConfig) err
}
}
// Rebuild the fantasy agent with the new model and current tool set.
a.rebuildFantasyAgent()
return nil
}
@@ -799,7 +833,13 @@ func (a *Agent) GetModel() fantasy.LanguageModel {
}
// Close closes the agent and cleans up resources.
// If MCP tools are still loading in the background, Close waits for them
// to finish before closing connections to avoid resource leaks.
func (a *Agent) Close() error {
// Wait for background MCP loading to finish before closing connections.
if a.mcpReady != nil {
<-a.mcpReady
}
var toolErr error
if a.toolManager != nil {
toolErr = a.toolManager.Close()
+55 -15
View File
@@ -4,8 +4,10 @@ import (
"context"
"encoding/json"
"fmt"
"maps"
"slices"
"strings"
"sync"
"charm.land/fantasy"
"github.com/mark3labs/kit/internal/config"
@@ -21,6 +23,7 @@ type MCPToolManager struct {
connectionPool *MCPConnectionPool
tools []fantasy.AgentTool
toolMap map[string]*toolMapping // maps prefixed tool names to their server and original name
mu sync.Mutex // protects tools and toolMap during parallel loading
model fantasy.LanguageModel // LLM model for sampling
authHandler MCPAuthHandler // OAuth handler for remote servers (nil = no OAuth)
config *config.Config
@@ -78,35 +81,62 @@ func (m *MCPToolManager) SetDebugLogger(logger DebugLogger) {
// Tools from different servers are prefixed with the server name to avoid naming conflicts.
// Returns an error only if all configured servers fail to load; partial failures are logged as warnings.
// This method is thread-safe and idempotent.
func (m *MCPToolManager) LoadTools(ctx context.Context, config *config.Config) error {
func (m *MCPToolManager) LoadTools(ctx context.Context, cfg *config.Config) error {
// Initialize connection pool
m.config = config
m.debug = config.Debug
m.config = cfg
m.debug = cfg.Debug
if m.debugLogger == nil {
m.debugLogger = NewSimpleDebugLogger(config.Debug)
m.debugLogger = NewSimpleDebugLogger(cfg.Debug)
}
m.connectionPool = NewMCPConnectionPool(DefaultConnectionPoolConfig(), m.model, config.Debug, m.authHandler)
m.connectionPool = NewMCPConnectionPool(DefaultConnectionPoolConfig(), m.model, cfg.Debug, m.authHandler)
m.connectionPool.SetDebugLogger(m.debugLogger)
var loadErrors []string
// Load all servers in parallel. Each server connection (subprocess
// spawn, MCP initialize handshake, ListTools) is independent and
// typically dominated by process startup latency. Running them
// concurrently reduces total wall-clock time from O(n * avg) to
// O(max).
type serverResult struct {
name string
err error
}
for serverName, serverConfig := range config.MCPServers {
if err := m.loadServerTools(ctx, serverName, serverConfig); err != nil {
loadErrors = append(loadErrors, fmt.Sprintf("server %s: %v", serverName, err))
fmt.Printf("Warning: Failed to load MCP server '%s': %v\n", serverName, err)
continue
results := make(chan serverResult, len(cfg.MCPServers))
var wg sync.WaitGroup
for serverName, serverConfig := range cfg.MCPServers {
wg.Add(1)
go func(name string, sc config.MCPServerConfig) {
defer wg.Done()
err := m.loadServerTools(ctx, name, sc)
results <- serverResult{name: name, err: err}
}(serverName, serverConfig)
}
// Close results channel once all goroutines finish.
go func() {
wg.Wait()
close(results)
}()
var loadErrors []string
for r := range results {
if r.err != nil {
loadErrors = append(loadErrors, fmt.Sprintf("server %s: %v", r.name, r.err))
fmt.Printf("Warning: Failed to load MCP server '%s': %v\n", r.name, r.err)
}
}
// If all servers failed to load, return an error
if len(loadErrors) == len(config.MCPServers) && len(config.MCPServers) > 0 {
if len(loadErrors) == len(cfg.MCPServers) && len(cfg.MCPServers) > 0 {
return fmt.Errorf("all MCP servers failed to load: %s", strings.Join(loadErrors, "; "))
}
return nil
}
// loadServerTools loads tools from a single MCP server
// loadServerTools loads tools from a single MCP server.
// Thread-safe: may be called concurrently for different servers.
func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string, serverConfig config.MCPServerConfig) error {
// Add debug logging
m.debugLogConnectionInfo(serverName, serverConfig)
@@ -134,6 +164,10 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
}
}
// Build tools locally before acquiring the lock.
var localTools []fantasy.AgentTool
localMap := make(map[string]*toolMapping)
// Convert MCP tools to fantasy AgentTools with prefixed names
for _, mcpTool := range listResults.Tools {
// Filter tools based on allowedTools/excludedTools
@@ -193,7 +227,7 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
serverConfig: serverConfig,
manager: m,
}
m.toolMap[prefixedName] = mapping
localMap[prefixedName] = mapping
// Create fantasy AgentTool
fantasyTool := &mcpFantasyTool{
@@ -206,9 +240,15 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
mapping: mapping,
}
m.tools = append(m.tools, fantasyTool)
localTools = append(localTools, fantasyTool)
}
// Merge into the manager under the lock.
m.mu.Lock()
maps.Copy(m.toolMap, localMap)
m.tools = append(m.tools, localTools...)
m.mu.Unlock()
return nil
}
+17
View File
@@ -113,15 +113,32 @@ func (m *Kit) GetLoadingMessage() string {
}
// GetLoadedServerNames returns the names of successfully loaded MCP servers.
// If MCP servers are still loading in the background, this returns only the
// servers that have completed loading so far.
func (m *Kit) GetLoadedServerNames() []string {
return m.agent.GetLoadedServerNames()
}
// GetMCPToolCount returns the number of tools loaded from external MCP servers.
// If MCP servers are still loading in the background, this returns the count
// of tools loaded so far (may be 0).
func (m *Kit) GetMCPToolCount() int {
return m.agent.GetMCPToolCount()
}
// WaitForMCPTools blocks until background MCP tool loading completes.
// Returns nil if no MCP servers are configured or if loading succeeded.
// Returns the loading error if all servers failed. Safe to call multiple times.
func (m *Kit) WaitForMCPTools() error {
return m.agent.WaitForMCPTools()
}
// MCPToolsReady returns true if MCP tool loading has completed (or was never
// started). This is a non-blocking check useful for UI status display.
func (m *Kit) MCPToolsReady() bool {
return m.agent.MCPToolsReady()
}
// GetExtensionToolCount returns the number of tools registered by extensions.
func (m *Kit) GetExtensionToolCount() int {
return m.agent.GetExtensionToolCount()