feat(mcp): enhance MCP management and status reporting in agent loop
This commit is contained in:
Ethan Wang
parent
7b6f60b668
commit
39d55b6874
7 changed files with 722 additions and 101 deletions
|
|
@ -1032,6 +1032,47 @@ func (al *AgentLoop) ReloadProviderAndConfig(
|
|||
// Ensure shared tools are re-registered on the new registry
|
||||
registerSharedTools(al, cfg, al.bus, registry, provider)
|
||||
|
||||
var (
|
||||
newMCPManager mcpController
|
||||
mcpSummary mcpRegistrationSummary
|
||||
)
|
||||
if cfg.Tools.IsToolEnabled("mcp") && countEnabledMCPServers(cfg.Tools.MCP.Servers) > 0 {
|
||||
newMCPManager = newMCPController()
|
||||
|
||||
workspacePath := cfg.WorkspacePath()
|
||||
if defaultAgent := registry.GetDefaultAgent(); defaultAgent != nil && defaultAgent.Workspace != "" {
|
||||
workspacePath = defaultAgent.Workspace
|
||||
}
|
||||
|
||||
if err := newMCPManager.LoadFromMCPConfig(ctx, cfg.Tools.MCP, workspacePath); err != nil {
|
||||
if closeErr := newMCPManager.Close(); closeErr != nil {
|
||||
logger.ErrorCF("agent", "Failed to close MCP manager",
|
||||
map[string]any{
|
||||
"error": closeErr.Error(),
|
||||
})
|
||||
}
|
||||
return fmt.Errorf("failed to initialize MCP during reload: %w", err)
|
||||
}
|
||||
|
||||
mcpSummary = registerMCPToolsOnRegistry(registry, cfg, newMCPManager, newMCPManager.GetServers())
|
||||
if err := registerMCPDiscoveryToolsOnRegistry(registry, cfg); err != nil {
|
||||
if closeErr := newMCPManager.Close(); closeErr != nil {
|
||||
logger.ErrorCF("agent", "Failed to close MCP manager",
|
||||
map[string]any{
|
||||
"error": closeErr.Error(),
|
||||
})
|
||||
}
|
||||
return fmt.Errorf("failed to restore MCP discovery tools during reload: %w", err)
|
||||
}
|
||||
}
|
||||
if al.mediaStore != nil {
|
||||
for _, agentID := range registry.ListAgentIDs() {
|
||||
if agent, ok := registry.GetAgent(agentID); ok {
|
||||
agent.Tools.SetMediaStore(al.mediaStore)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Atomically swap the config and registry under write lock
|
||||
// This ensures readers see a consistent pair
|
||||
al.mu.Lock()
|
||||
|
|
@ -1053,6 +1094,7 @@ func (al *AgentLoop) ReloadProviderAndConfig(
|
|||
|
||||
al.mu.Unlock()
|
||||
|
||||
oldMCPManager := al.mcp.replaceForReload(newMCPManager, newMCPManager != nil)
|
||||
al.hookRuntime.reset(al)
|
||||
configureHookManagerFromConfig(al.hooks, cfg)
|
||||
|
||||
|
|
@ -1073,10 +1115,21 @@ func (al *AgentLoop) ReloadProviderAndConfig(
|
|||
}
|
||||
}
|
||||
}
|
||||
if oldMCPManager != nil {
|
||||
if err := oldMCPManager.Close(); err != nil {
|
||||
logger.ErrorCF("agent", "Failed to close previous MCP manager",
|
||||
map[string]any{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
logger.InfoCF("agent", "Provider and config reloaded successfully",
|
||||
map[string]any{
|
||||
"model": cfg.Agents.Defaults.GetModelName(),
|
||||
"mcp_server_count": mcpSummary.serverCount,
|
||||
"mcp_unique_tools": mcpSummary.uniqueTools,
|
||||
"mcp_total_registrations": mcpSummary.totalRegistrations,
|
||||
})
|
||||
|
||||
return nil
|
||||
|
|
@ -3432,6 +3485,9 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
|
|||
Config: cfg,
|
||||
ListAgentIDs: registry.ListAgentIDs,
|
||||
ListDefinitions: al.cmdRegistry.Definitions,
|
||||
GetMCPStatus: func() string {
|
||||
return formatMCPStatus(cfg, al.mcp.statusSnapshot())
|
||||
},
|
||||
GetEnabledChannels: func() []string {
|
||||
if al.channelManager == nil {
|
||||
return nil
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ package agent
|
|||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
|
|
@ -20,20 +22,52 @@ import (
|
|||
type mcpRuntime struct {
|
||||
initOnce sync.Once
|
||||
mu sync.Mutex
|
||||
manager *mcp.Manager
|
||||
attempted bool
|
||||
manager mcpController
|
||||
caller tools.MCPManager
|
||||
servers map[string]*mcp.ServerConnection
|
||||
initErr error
|
||||
lastErr error
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) setManager(manager *mcp.Manager) {
|
||||
type mcpController interface {
|
||||
tools.MCPManager
|
||||
LoadFromMCPConfig(ctx context.Context, mcpCfg config.MCPConfig, workspacePath string) error
|
||||
GetServers() map[string]*mcp.ServerConnection
|
||||
Close() error
|
||||
}
|
||||
|
||||
var newMCPController = func() mcpController {
|
||||
return mcp.NewManager()
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) setManager(manager mcpController) {
|
||||
r.mu.Lock()
|
||||
r.attempted = true
|
||||
r.manager = manager
|
||||
r.caller = manager
|
||||
if manager != nil {
|
||||
r.servers = manager.GetServers()
|
||||
} else {
|
||||
r.servers = nil
|
||||
}
|
||||
r.initErr = nil
|
||||
r.lastErr = nil
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) setInitErr(err error) {
|
||||
r.mu.Lock()
|
||||
r.attempted = true
|
||||
r.initErr = err
|
||||
r.lastErr = err
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) setStatusErr(err error) {
|
||||
r.mu.Lock()
|
||||
r.attempted = true
|
||||
r.lastErr = err
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
|
|
@ -43,20 +77,301 @@ func (r *mcpRuntime) getInitErr() error {
|
|||
return r.initErr
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) takeManager() *mcp.Manager {
|
||||
func (r *mcpRuntime) takeManager() mcpController {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
manager := r.manager
|
||||
r.manager = nil
|
||||
r.caller = nil
|
||||
r.servers = nil
|
||||
return manager
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) replaceForReload(manager mcpController, attempted bool) mcpController {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
oldManager := r.manager
|
||||
r.initOnce = sync.Once{}
|
||||
r.attempted = attempted
|
||||
r.manager = manager
|
||||
r.initErr = nil
|
||||
r.lastErr = nil
|
||||
if manager != nil {
|
||||
r.caller = manager
|
||||
r.servers = manager.GetServers()
|
||||
} else {
|
||||
r.caller = nil
|
||||
r.servers = nil
|
||||
}
|
||||
if attempted {
|
||||
r.initOnce.Do(func() {})
|
||||
}
|
||||
return oldManager
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) hasManager() bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.manager != nil
|
||||
}
|
||||
|
||||
type mcpStatusSnapshot struct {
|
||||
attempted bool
|
||||
lastErr error
|
||||
servers map[string]*mcp.ServerConnection
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) statusSnapshot() mcpStatusSnapshot {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
servers := make(map[string]*mcp.ServerConnection, len(r.servers))
|
||||
for name, conn := range r.servers {
|
||||
servers[name] = conn
|
||||
}
|
||||
|
||||
return mcpStatusSnapshot{
|
||||
attempted: r.attempted,
|
||||
lastErr: r.lastErr,
|
||||
servers: servers,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *mcpRuntime) registrationSnapshot() (tools.MCPManager, map[string]*mcp.ServerConnection) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
if r.caller == nil || len(r.servers) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
servers := make(map[string]*mcp.ServerConnection, len(r.servers))
|
||||
for name, conn := range r.servers {
|
||||
servers[name] = conn
|
||||
}
|
||||
return r.caller, servers
|
||||
}
|
||||
|
||||
func formatMCPStatus(cfg *config.Config, snap mcpStatusSnapshot) string {
|
||||
if cfg == nil {
|
||||
return "MCP status unavailable: config not loaded."
|
||||
}
|
||||
|
||||
if !cfg.Tools.IsToolEnabled("mcp") || !cfg.Tools.MCP.Enabled {
|
||||
return "MCP is disabled."
|
||||
}
|
||||
|
||||
configured := cfg.Tools.MCP.Servers
|
||||
if len(configured) == 0 {
|
||||
return "MCP is enabled, but no servers are configured."
|
||||
}
|
||||
|
||||
lines := []string{
|
||||
fmt.Sprintf("MCP Enabled: yes"),
|
||||
fmt.Sprintf("Initialization Attempted: %s", yesNo(snap.attempted)),
|
||||
fmt.Sprintf("Connected Servers: %d/%d", len(snap.servers), countEnabledMCPServers(configured)),
|
||||
}
|
||||
|
||||
if snap.lastErr != nil {
|
||||
lines = append(lines, fmt.Sprintf("Last Init Error: %s", snap.lastErr.Error()))
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(configured))
|
||||
for name, serverCfg := range configured {
|
||||
if !serverCfg.Enabled {
|
||||
continue
|
||||
}
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
|
||||
if len(names) == 0 {
|
||||
lines = append(lines, "No enabled MCP servers.")
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
lines = append(lines, "", "Servers:")
|
||||
for _, name := range names {
|
||||
serverCfg := configured[name]
|
||||
conn, connected := snap.servers[name]
|
||||
toolCount := 0
|
||||
if conn != nil {
|
||||
toolCount = len(conn.Tools)
|
||||
}
|
||||
lines = append(lines, fmt.Sprintf(
|
||||
"- %s: %s, transport=%s, tools=%d%s",
|
||||
name,
|
||||
connectionStatusLabel(connected),
|
||||
mcpTransportLabel(serverCfg),
|
||||
toolCount,
|
||||
mcpEndpointSummary(serverCfg),
|
||||
))
|
||||
}
|
||||
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
func countEnabledMCPServers(servers map[string]config.MCPServerConfig) int {
|
||||
count := 0
|
||||
for _, serverCfg := range servers {
|
||||
if serverCfg.Enabled {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func yesNo(v bool) string {
|
||||
if v {
|
||||
return "yes"
|
||||
}
|
||||
return "no"
|
||||
}
|
||||
|
||||
func connectionStatusLabel(connected bool) string {
|
||||
if connected {
|
||||
return "connected"
|
||||
}
|
||||
return "not connected"
|
||||
}
|
||||
|
||||
func mcpTransportLabel(serverCfg config.MCPServerConfig) string {
|
||||
transportType := strings.TrimSpace(serverCfg.Type)
|
||||
if transportType != "" {
|
||||
return transportType
|
||||
}
|
||||
if strings.TrimSpace(serverCfg.URL) != "" {
|
||||
return "sse"
|
||||
}
|
||||
if strings.TrimSpace(serverCfg.Command) != "" {
|
||||
return "stdio"
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
func mcpEndpointSummary(serverCfg config.MCPServerConfig) string {
|
||||
if url := strings.TrimSpace(serverCfg.URL); url != "" {
|
||||
return fmt.Sprintf(", url=%s", url)
|
||||
}
|
||||
if cmd := strings.TrimSpace(serverCfg.Command); cmd != "" {
|
||||
return fmt.Sprintf(", command=%s", cmd)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type mcpRegistrationSummary struct {
|
||||
serverCount int
|
||||
uniqueTools int
|
||||
totalRegistrations int
|
||||
agentCount int
|
||||
}
|
||||
|
||||
func registerMCPToolsOnRegistry(
|
||||
registry *AgentRegistry,
|
||||
cfg *config.Config,
|
||||
caller tools.MCPManager,
|
||||
servers map[string]*mcp.ServerConnection,
|
||||
) mcpRegistrationSummary {
|
||||
if registry == nil || cfg == nil || caller == nil || len(servers) == 0 {
|
||||
return mcpRegistrationSummary{}
|
||||
}
|
||||
if !cfg.Tools.IsToolEnabled("mcp") || !cfg.Tools.MCP.Enabled {
|
||||
return mcpRegistrationSummary{}
|
||||
}
|
||||
|
||||
agentIDs := registry.ListAgentIDs()
|
||||
summary := mcpRegistrationSummary{
|
||||
serverCount: len(servers),
|
||||
agentCount: len(agentIDs),
|
||||
}
|
||||
|
||||
for serverName, conn := range servers {
|
||||
serverCfg, ok := cfg.Tools.MCP.Servers[serverName]
|
||||
if !ok || !serverCfg.Enabled {
|
||||
continue
|
||||
}
|
||||
|
||||
summary.uniqueTools += len(conn.Tools)
|
||||
registerAsHidden := serverIsDeferred(cfg.Tools.MCP.Discovery.Enabled, serverCfg)
|
||||
|
||||
for _, tool := range conn.Tools {
|
||||
for _, agentID := range agentIDs {
|
||||
agent, ok := registry.GetAgent(agentID)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
mcpTool := tools.NewMCPTool(caller, serverName, tool)
|
||||
mcpTool.SetWorkspace(agent.Workspace)
|
||||
mcpTool.SetMaxInlineTextRunes(cfg.Tools.MCP.GetMaxInlineTextChars())
|
||||
|
||||
if registerAsHidden {
|
||||
agent.Tools.RegisterHidden(mcpTool)
|
||||
} else {
|
||||
agent.Tools.Register(mcpTool)
|
||||
}
|
||||
|
||||
summary.totalRegistrations++
|
||||
logger.DebugCF("agent", "Registered MCP tool",
|
||||
map[string]any{
|
||||
"agent_id": agentID,
|
||||
"server": serverName,
|
||||
"tool": tool.Name,
|
||||
"name": mcpTool.Name(),
|
||||
"deferred": registerAsHidden,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return summary
|
||||
}
|
||||
|
||||
func registerMCPDiscoveryToolsOnRegistry(registry *AgentRegistry, cfg *config.Config) error {
|
||||
if registry == nil || cfg == nil || !cfg.Tools.MCP.Enabled || !cfg.Tools.MCP.Discovery.Enabled {
|
||||
return nil
|
||||
}
|
||||
|
||||
useBM25 := cfg.Tools.MCP.Discovery.UseBM25
|
||||
useRegex := cfg.Tools.MCP.Discovery.UseRegex
|
||||
if !useBM25 && !useRegex {
|
||||
return fmt.Errorf(
|
||||
"tool discovery is enabled but neither 'use_bm25' nor 'use_regex' is set to true in the configuration",
|
||||
)
|
||||
}
|
||||
|
||||
ttl := cfg.Tools.MCP.Discovery.TTL
|
||||
if ttl <= 0 {
|
||||
ttl = 5
|
||||
}
|
||||
|
||||
maxSearchResults := cfg.Tools.MCP.Discovery.MaxSearchResults
|
||||
if maxSearchResults <= 0 {
|
||||
maxSearchResults = 5
|
||||
}
|
||||
|
||||
logger.InfoCF("agent", "Initializing tool discovery", map[string]any{
|
||||
"bm25": useBM25, "regex": useRegex, "ttl": ttl, "max_results": maxSearchResults,
|
||||
})
|
||||
|
||||
for _, agentID := range registry.ListAgentIDs() {
|
||||
agent, ok := registry.GetAgent(agentID)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
if useRegex {
|
||||
agent.Tools.Register(tools.NewRegexSearchTool(agent.Tools, ttl, maxSearchResults))
|
||||
}
|
||||
if useBM25 {
|
||||
agent.Tools.Register(tools.NewBM25SearchTool(agent.Tools, ttl, maxSearchResults))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ensureMCPInitialized loads MCP servers/tools once so both Run() and direct
|
||||
// agent mode share the same initialization path.
|
||||
func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
||||
|
|
@ -81,7 +396,7 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
|||
}
|
||||
|
||||
al.mcp.initOnce.Do(func() {
|
||||
mcpManager := mcp.NewManager()
|
||||
mcpManager := newMCPController()
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
workspacePath := al.cfg.WorkspacePath()
|
||||
|
|
@ -90,6 +405,7 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
|||
}
|
||||
|
||||
if err := mcpManager.LoadFromMCPConfig(ctx, al.cfg.Tools.MCP, workspacePath); err != nil {
|
||||
al.mcp.setStatusErr(err)
|
||||
logger.WarnCF("agent", "Failed to load MCP servers, MCP tools will not be available",
|
||||
map[string]any{
|
||||
"error": err.Error(),
|
||||
|
|
@ -103,68 +419,18 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
|||
return
|
||||
}
|
||||
|
||||
// Register MCP tools for all agents
|
||||
servers := mcpManager.GetServers()
|
||||
uniqueTools := 0
|
||||
totalRegistrations := 0
|
||||
agentIDs := al.registry.ListAgentIDs()
|
||||
agentCount := len(agentIDs)
|
||||
|
||||
for serverName, conn := range servers {
|
||||
uniqueTools += len(conn.Tools)
|
||||
|
||||
// Determine whether this server's tools should be deferred (hidden).
|
||||
// Per-server "deferred" field takes precedence over the global Discovery.Enabled.
|
||||
serverCfg := al.cfg.Tools.MCP.Servers[serverName]
|
||||
registerAsHidden := serverIsDeferred(al.cfg.Tools.MCP.Discovery.Enabled, serverCfg)
|
||||
|
||||
for _, tool := range conn.Tools {
|
||||
for _, agentID := range agentIDs {
|
||||
agent, ok := al.registry.GetAgent(agentID)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
mcpTool := tools.NewMCPTool(mcpManager, serverName, tool)
|
||||
mcpTool.SetWorkspace(agent.Workspace)
|
||||
mcpTool.SetMaxInlineTextRunes(al.cfg.Tools.MCP.GetMaxInlineTextChars())
|
||||
|
||||
if registerAsHidden {
|
||||
agent.Tools.RegisterHidden(mcpTool)
|
||||
} else {
|
||||
agent.Tools.Register(mcpTool)
|
||||
}
|
||||
|
||||
totalRegistrations++
|
||||
logger.DebugCF("agent", "Registered MCP tool",
|
||||
map[string]any{
|
||||
"agent_id": agentID,
|
||||
"server": serverName,
|
||||
"tool": tool.Name,
|
||||
"name": mcpTool.Name(),
|
||||
"deferred": registerAsHidden,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
summary := registerMCPToolsOnRegistry(al.registry, al.cfg, mcpManager, servers)
|
||||
logger.InfoCF("agent", "MCP tools registered successfully",
|
||||
map[string]any{
|
||||
"server_count": len(servers),
|
||||
"unique_tools": uniqueTools,
|
||||
"total_registrations": totalRegistrations,
|
||||
"agent_count": agentCount,
|
||||
"server_count": summary.serverCount,
|
||||
"unique_tools": summary.uniqueTools,
|
||||
"total_registrations": summary.totalRegistrations,
|
||||
"agent_count": summary.agentCount,
|
||||
})
|
||||
|
||||
// Initializes Discovery Tools only if enabled by configuration
|
||||
if al.cfg.Tools.MCP.Enabled && al.cfg.Tools.MCP.Discovery.Enabled {
|
||||
useBM25 := al.cfg.Tools.MCP.Discovery.UseBM25
|
||||
useRegex := al.cfg.Tools.MCP.Discovery.UseRegex
|
||||
|
||||
// Fail fast: If discovery is enabled but no search method is turned on
|
||||
if !useBM25 && !useRegex {
|
||||
al.mcp.setInitErr(fmt.Errorf(
|
||||
"tool discovery is enabled but neither 'use_bm25' nor 'use_regex' is set to true in the configuration",
|
||||
))
|
||||
if err := registerMCPDiscoveryToolsOnRegistry(al.registry, al.cfg); err != nil {
|
||||
al.mcp.setInitErr(err)
|
||||
if closeErr := mcpManager.Close(); closeErr != nil {
|
||||
logger.ErrorCF("agent", "Failed to close MCP manager",
|
||||
map[string]any{
|
||||
|
|
@ -174,35 +440,6 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
|||
return
|
||||
}
|
||||
|
||||
ttl := al.cfg.Tools.MCP.Discovery.TTL
|
||||
if ttl <= 0 {
|
||||
ttl = 5 // Default value
|
||||
}
|
||||
|
||||
maxSearchResults := al.cfg.Tools.MCP.Discovery.MaxSearchResults
|
||||
if maxSearchResults <= 0 {
|
||||
maxSearchResults = 5 // Default value
|
||||
}
|
||||
|
||||
logger.InfoCF("agent", "Initializing tool discovery", map[string]any{
|
||||
"bm25": useBM25, "regex": useRegex, "ttl": ttl, "max_results": maxSearchResults,
|
||||
})
|
||||
|
||||
for _, agentID := range agentIDs {
|
||||
agent, ok := al.registry.GetAgent(agentID)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
if useRegex {
|
||||
agent.Tools.Register(tools.NewRegexSearchTool(agent.Tools, ttl, maxSearchResults))
|
||||
}
|
||||
if useBM25 {
|
||||
agent.Tools.Register(tools.NewBM25SearchTool(agent.Tools, ttl, maxSearchResults))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
al.mcp.setManager(mcpManager)
|
||||
})
|
||||
|
||||
|
|
|
|||
|
|
@ -14,9 +14,12 @@ import (
|
|||
"testing"
|
||||
"time"
|
||||
|
||||
sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/channels"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
mcppkg "github.com/sipeed/picoclaw/pkg/mcp"
|
||||
"github.com/sipeed/picoclaw/pkg/media"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/routing"
|
||||
|
|
@ -75,6 +78,66 @@ type recordingProvider struct {
|
|||
lastMessages []providers.Message
|
||||
}
|
||||
|
||||
type fakeMCPController struct {
|
||||
closed bool
|
||||
servers map[string]*mcppkg.ServerConnection
|
||||
}
|
||||
|
||||
func (m *fakeMCPController) LoadFromMCPConfig(
|
||||
_ context.Context,
|
||||
mcpCfg config.MCPConfig,
|
||||
_ string,
|
||||
) error {
|
||||
m.servers = make(map[string]*mcppkg.ServerConnection)
|
||||
for serverName, serverCfg := range mcpCfg.Servers {
|
||||
if !serverCfg.Enabled {
|
||||
continue
|
||||
}
|
||||
m.servers[serverName] = &mcppkg.ServerConnection{
|
||||
Name: serverName,
|
||||
Tools: []*sdkmcp.Tool{
|
||||
{
|
||||
Name: "ping",
|
||||
Description: "Remote ping",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *fakeMCPController) GetServers() map[string]*mcppkg.ServerConnection {
|
||||
servers := make(map[string]*mcppkg.ServerConnection, len(m.servers))
|
||||
for name, conn := range m.servers {
|
||||
servers[name] = conn
|
||||
}
|
||||
return servers
|
||||
}
|
||||
|
||||
func (m *fakeMCPController) CallTool(
|
||||
ctx context.Context,
|
||||
serverName, toolName string,
|
||||
arguments map[string]any,
|
||||
) (*sdkmcp.CallToolResult, error) {
|
||||
if m.closed {
|
||||
return nil, fmt.Errorf("manager is closed")
|
||||
}
|
||||
return &sdkmcp.CallToolResult{
|
||||
Content: []sdkmcp.Content{
|
||||
&sdkmcp.TextContent{Text: fmt.Sprintf("%s:%s", serverName, toolName)},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *fakeMCPController) Close() error {
|
||||
m.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *recordingProvider) Chat(
|
||||
ctx context.Context,
|
||||
messages []providers.Message,
|
||||
|
|
@ -2331,6 +2394,216 @@ func TestProcessDirectWithChannel_TriggersMCPInitialization(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestReloadProviderAndConfig_RebuildsMCPToolsFromNewConfig(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
oldFactory := newMCPController
|
||||
newMCPController = func() mcpController {
|
||||
return &fakeMCPController{}
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
newMCPController = oldFactory
|
||||
})
|
||||
|
||||
oldCfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
Tools: config.ToolsConfig{
|
||||
MCP: config.MCPConfig{
|
||||
ToolConfig: config.ToolConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"stale": {Enabled: true},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
newCfg := &config.Config{
|
||||
Agents: oldCfg.Agents,
|
||||
Tools: config.ToolsConfig{
|
||||
MCP: config.MCPConfig{
|
||||
ToolConfig: config.ToolConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"fresh": {Enabled: true},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
al := NewAgentLoop(oldCfg, msgBus, &mockProvider{})
|
||||
defer al.Close()
|
||||
|
||||
al.mcp.setManager(&fakeMCPController{
|
||||
servers: map[string]*mcppkg.ServerConnection{
|
||||
"stale": {
|
||||
Name: "stale",
|
||||
Tools: []*sdkmcp.Tool{
|
||||
{
|
||||
Name: "ping",
|
||||
Description: "Stale ping",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if err := al.ReloadProviderAndConfig(context.Background(), &mockProvider{}, newCfg); err != nil {
|
||||
t.Fatalf("ReloadProviderAndConfig() error = %v", err)
|
||||
}
|
||||
|
||||
agent := al.GetRegistry().GetDefaultAgent()
|
||||
if agent == nil {
|
||||
t.Fatal("expected default agent after reload")
|
||||
}
|
||||
|
||||
if _, ok := agent.Tools.Get("mcp_stale_ping"); ok {
|
||||
t.Fatal("expected stale MCP tool to be removed after reload")
|
||||
}
|
||||
|
||||
toolName := "mcp_fresh_ping"
|
||||
if _, ok := agent.Tools.Get(toolName); !ok {
|
||||
t.Fatalf("expected MCP tool %q to be registered from reloaded config", toolName)
|
||||
}
|
||||
|
||||
result := agent.Tools.Execute(context.Background(), toolName, map[string]any{})
|
||||
if result == nil || result.IsError {
|
||||
t.Fatalf("expected MCP tool %q to execute successfully after reload, got %#v", toolName, result)
|
||||
}
|
||||
if !strings.Contains(result.ContentForLLM(), "fresh:ping") {
|
||||
t.Fatalf("unexpected MCP tool result after reload: %q", result.ContentForLLM())
|
||||
}
|
||||
}
|
||||
|
||||
func TestReloadProviderAndConfig_RemovesMCPToolsWhenServersDeleted(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
Tools: config.ToolsConfig{
|
||||
MCP: config.MCPConfig{
|
||||
ToolConfig: config.ToolConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"remote": {Enabled: true},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
reloadedCfg := &config.Config{
|
||||
Agents: cfg.Agents,
|
||||
Tools: config.ToolsConfig{
|
||||
MCP: config.MCPConfig{
|
||||
ToolConfig: config.ToolConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
al := NewAgentLoop(cfg, msgBus, &mockProvider{})
|
||||
defer al.Close()
|
||||
|
||||
al.mcp.setManager(&fakeMCPController{
|
||||
servers: map[string]*mcppkg.ServerConnection{
|
||||
"remote": {
|
||||
Name: "remote",
|
||||
Tools: []*sdkmcp.Tool{
|
||||
{
|
||||
Name: "ping",
|
||||
Description: "Remote ping",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if err := al.ReloadProviderAndConfig(context.Background(), &mockProvider{}, reloadedCfg); err != nil {
|
||||
t.Fatalf("ReloadProviderAndConfig() error = %v", err)
|
||||
}
|
||||
|
||||
agent := al.GetRegistry().GetDefaultAgent()
|
||||
if agent == nil {
|
||||
t.Fatal("expected default agent after reload")
|
||||
}
|
||||
if _, ok := agent.Tools.Get("mcp_remote_ping"); ok {
|
||||
t.Fatal("expected MCP tool to be removed when no servers are configured after reload")
|
||||
}
|
||||
if al.mcp.hasManager() {
|
||||
t.Fatal("expected MCP manager to be cleared when reload removes all servers")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatMCPStatus_IncludesConfiguredServers(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Tools: config.ToolsConfig{
|
||||
MCP: config.MCPConfig{
|
||||
ToolConfig: config.ToolConfig{Enabled: true},
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"remote-http": {
|
||||
Enabled: true,
|
||||
Type: "http",
|
||||
URL: "http://127.0.0.1:8080/mcp",
|
||||
},
|
||||
"local-stdio": {
|
||||
Enabled: true,
|
||||
Command: "npx",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
status := formatMCPStatus(cfg, mcpStatusSnapshot{
|
||||
attempted: true,
|
||||
servers: map[string]*mcppkg.ServerConnection{
|
||||
"remote-http": {
|
||||
Name: "remote-http",
|
||||
Tools: []*sdkmcp.Tool{{Name: "ping"}, {Name: "echo"}},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if !strings.Contains(status, "Initialization Attempted: yes") {
|
||||
t.Fatalf("status missing initialization state:\n%s", status)
|
||||
}
|
||||
if !strings.Contains(status, "Connected Servers: 1/2") {
|
||||
t.Fatalf("status missing connected server count:\n%s", status)
|
||||
}
|
||||
if !strings.Contains(status, "remote-http: connected, transport=http, tools=2") {
|
||||
t.Fatalf("status missing connected http server details:\n%s", status)
|
||||
}
|
||||
if !strings.Contains(status, "local-stdio: not connected, transport=stdio, tools=0, command=npx") {
|
||||
t.Fatalf("status missing disconnected stdio server details:\n%s", status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTargetReasoningChannelID_AllChannels(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
||||
if err != nil {
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ func TestBuiltinHelpHandler_ReturnsFormattedMessage(t *testing.T) {
|
|||
if !strings.Contains(reply, "/show [model|channel|agents]") {
|
||||
t.Fatalf("/help reply missing /show usage, got %q", reply)
|
||||
}
|
||||
if !strings.Contains(reply, "/list [models|channels|agents|skills]") {
|
||||
if !strings.Contains(reply, "/list [models|channels|agents|skills|mcp]") {
|
||||
t.Fatalf("/help reply missing /list usage, got %q", reply)
|
||||
}
|
||||
if !strings.Contains(reply, "/use <skill> <message>") {
|
||||
|
|
@ -174,6 +174,31 @@ func TestBuiltinListSkills_UsesRuntimeSkillNames(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestBuiltinListMCP_UsesRuntimeStatus(t *testing.T) {
|
||||
rt := &Runtime{
|
||||
GetMCPStatus: func() string {
|
||||
return "MCP Enabled: yes\nConnected Servers: 1/1\n\nServers:\n- remote: connected, transport=http, tools=3, url=http://127.0.0.1:8080/mcp"
|
||||
},
|
||||
}
|
||||
defs := BuiltinDefinitions()
|
||||
ex := NewExecutor(NewRegistry(defs), rt)
|
||||
|
||||
var reply string
|
||||
res := ex.Execute(context.Background(), Request{
|
||||
Text: "/list mcp",
|
||||
Reply: func(text string) error {
|
||||
reply = text
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if res.Outcome != OutcomeHandled {
|
||||
t.Fatalf("/list mcp: outcome=%v, want=%v", res.Outcome, OutcomeHandled)
|
||||
}
|
||||
if !strings.Contains(reply, "Connected Servers: 1/1") || !strings.Contains(reply, "remote: connected") {
|
||||
t.Fatalf("/list mcp reply=%q, want MCP status summary", reply)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuiltinUseCommand_PassthroughsToAgentLogic(t *testing.T) {
|
||||
defs := BuiltinDefinitions()
|
||||
ex := NewExecutor(NewRegistry(defs), nil)
|
||||
|
|
|
|||
|
|
@ -64,6 +64,16 @@ func listCommand() Definition {
|
|||
))
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "mcp",
|
||||
Description: "MCP server status",
|
||||
Handler: func(_ context.Context, req Request, rt *Runtime) error {
|
||||
if rt == nil || rt.GetMCPStatus == nil {
|
||||
return req.Reply(unavailableMsg)
|
||||
}
|
||||
return req.Reply(rt.GetMCPStatus())
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ type Runtime struct {
|
|||
ListDefinitions func() []Definition
|
||||
ListSkillNames func() []string
|
||||
GetEnabledChannels func() []string
|
||||
GetMCPStatus func() string
|
||||
GetActiveTurn func() any // Returning any to avoid circular dependency with agent package
|
||||
SwitchModel func(value string) (oldModel string, err error)
|
||||
SwitchChannel func(value string) error
|
||||
|
|
|
|||
|
|
@ -64,6 +64,9 @@ func TestShowListHandlers_ListHandledOnAllChannels(t *testing.T) {
|
|||
ListSkillNames: func() []string {
|
||||
return []string{"shell"}
|
||||
},
|
||||
GetMCPStatus: func() string {
|
||||
return "MCP Enabled: yes\nConnected Servers: 0/1"
|
||||
},
|
||||
}
|
||||
ex := NewExecutor(NewRegistry(BuiltinDefinitions()), rt)
|
||||
|
||||
|
|
@ -101,4 +104,20 @@ func TestShowListHandlers_ListHandledOnAllChannels(t *testing.T) {
|
|||
if !strings.Contains(reply, "shell") {
|
||||
t.Fatalf("whatsapp /list skills reply=%q, expected installed skills content", reply)
|
||||
}
|
||||
|
||||
reply = ""
|
||||
res = ex.Execute(context.Background(), Request{
|
||||
Channel: "whatsapp",
|
||||
Text: "/list mcp",
|
||||
Reply: func(text string) error {
|
||||
reply = text
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if res.Outcome != OutcomeHandled {
|
||||
t.Fatalf("whatsapp /list mcp outcome=%v, want=%v", res.Outcome, OutcomeHandled)
|
||||
}
|
||||
if !strings.Contains(reply, "Connected Servers: 0/1") {
|
||||
t.Fatalf("whatsapp /list mcp reply=%q, expected mcp status content", reply)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue