From 697d7268312e2d8add181ed461a0c97f8d6379e5 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 18 Mar 2026 10:03:47 +0000 Subject: [PATCH] feat(agent): allow assigning MCP servers to specific personas - Added `MCPServers` and `SystemPrompt` to `AgentConfig`. - Configured default out-of-the-box personas (coding, google, communications, financial) with specific MCP assignments. - Added CLI commands `picoclaw agent create`, `assign-mcp`, and `remove-mcp`. - Updated agent MCP tool registration to filter tools per persona based on `MCPServers`. Agents with an empty list will not have any MCP tools registered. Co-authored-by: hobbyistlabs-coder <267281733+hobbyistlabs-coder@users.noreply.github.com> --- cmd/picoclaw/internal/agent/command.go | 4 + cmd/picoclaw/internal/agent/create.go | 129 +++++++++++++++++++++++++ cmd/picoclaw/internal/agent/mcp.go | 113 ++++++++++++++++++++++ pkg/agent/instance.go | 4 + pkg/agent/loop_mcp.go | 14 +++ pkg/config/agent.go | 16 +-- pkg/config/config_test.go | 4 +- pkg/config/defaults.go | 24 +++++ 8 files changed, 299 insertions(+), 9 deletions(-) create mode 100644 cmd/picoclaw/internal/agent/create.go create mode 100644 cmd/picoclaw/internal/agent/mcp.go diff --git a/cmd/picoclaw/internal/agent/command.go b/cmd/picoclaw/internal/agent/command.go index 47262fc85..6e8dc858f 100644 --- a/cmd/picoclaw/internal/agent/command.go +++ b/cmd/picoclaw/internal/agent/command.go @@ -26,5 +26,9 @@ func NewAgentCommand() *cobra.Command { cmd.Flags().StringVarP(&sessionKey, "session", "s", "cli:default", "Session key") cmd.Flags().StringVarP(&model, "model", "", "", "Model to use") + cmd.AddCommand(NewCreateCommand()) + cmd.AddCommand(NewAssignMCPCommand()) + cmd.AddCommand(NewRemoveMCPCommand()) + return cmd } diff --git a/cmd/picoclaw/internal/agent/create.go b/cmd/picoclaw/internal/agent/create.go new file mode 100644 index 000000000..4d72a6a1d --- /dev/null +++ b/cmd/picoclaw/internal/agent/create.go @@ -0,0 +1,129 @@ +package agent + +import ( + "bufio" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/google/uuid" + "github.com/spf13/cobra" + + "jane/cmd/picoclaw/internal" + "jane/pkg/config" +) + +func NewCreateCommand() *cobra.Command { + var ( + name string + workspace string + sysPrompt string + model string + interactive bool + ) + + cmd := &cobra.Command{ + Use: "create", + Short: "Create a new agent persona", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + if !interactive && name == "" { + return fmt.Errorf("name is required when not in interactive mode") + } + return createAgentCmd(name, workspace, sysPrompt, model, interactive) + }, + } + + cmd.Flags().StringVarP(&name, "name", "n", "", "Agent name") + cmd.Flags().StringVarP(&workspace, "workspace", "w", "", "Workspace path") + cmd.Flags().StringVarP(&sysPrompt, "system-prompt", "p", "", "System prompt / instructions") + cmd.Flags().StringVarP(&model, "model", "m", "", "Model configuration (primary)") + cmd.Flags().BoolVarP(&interactive, "interactive", "i", false, "Interactive mode") + + return cmd +} + +func createAgentCmd(name, workspace, sysPrompt, model string, interactive bool) error { + cfg, err := internal.LoadConfig() + if err != nil { + return fmt.Errorf("error loading config: %w", err) + } + + if interactive { + fmt.Printf("%s Creating new agent persona...\n\n", internal.Logo) + + reader := bufio.NewReader(os.Stdin) + + if name == "" { + fmt.Print("Agent Name: ") + nameInput, _ := reader.ReadString('\n') + nameInput = strings.TrimSpace(nameInput) + if nameInput != "" { + name = nameInput + } + } + + if workspace == "" { + fmt.Printf("Workspace path (default: ~/.picoclaw/workspace/%s): ", strings.ToLower(strings.ReplaceAll(name, " ", "_"))) + workspaceInput, _ := reader.ReadString('\n') + workspaceInput = strings.TrimSpace(workspaceInput) + if workspaceInput != "" { + workspace = workspaceInput + } + } + + if sysPrompt == "" { + fmt.Print("System Prompt (optional): ") + sysPromptInput, _ := reader.ReadString('\n') + sysPromptInput = strings.TrimSpace(sysPromptInput) + if sysPromptInput != "" { + sysPrompt = sysPromptInput + } + } + } + + if name == "" { + return fmt.Errorf("agent name is required") + } + + id := strings.ToLower(strings.ReplaceAll(name, " ", "-")) + id = strings.ReplaceAll(id, "_", "-") + + if id == "" { + id = uuid.New().String()[:8] + } + + if workspace == "" { + var homePath string + if picoclawHome := os.Getenv("PICOCLAW_HOME"); picoclawHome != "" { + homePath = picoclawHome + } else { + userHome, _ := os.UserHomeDir() + homePath = filepath.Join(userHome, ".picoclaw") + } + workspace = filepath.Join(homePath, "workspace", strings.ToLower(strings.ReplaceAll(name, " ", "_"))) + } + + var modelCfg *config.AgentModelConfig + if model != "" { + modelCfg = &config.AgentModelConfig{Primary: model} + } + + newAgent := config.AgentConfig{ + ID: id, + Name: name, + Workspace: workspace, + SystemPrompt: sysPrompt, + Model: modelCfg, + } + + cfg.Agents.List = append(cfg.Agents.List, newAgent) + + if err := config.SaveConfig(internal.GetConfigPath(), cfg); err != nil { + return fmt.Errorf("failed to save config: %w", err) + } + + fmt.Printf("\n✅ Successfully created agent persona '%s' with ID '%s'\n", name, id) + return nil +} diff --git a/cmd/picoclaw/internal/agent/mcp.go b/cmd/picoclaw/internal/agent/mcp.go new file mode 100644 index 000000000..ecacead40 --- /dev/null +++ b/cmd/picoclaw/internal/agent/mcp.go @@ -0,0 +1,113 @@ +package agent + +import ( + "fmt" + + "github.com/spf13/cobra" + + "jane/cmd/picoclaw/internal" + "jane/pkg/config" +) + +func NewAssignMCPCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "assign-mcp ", + Short: "Assign an MCP server to an agent persona", + Args: cobra.ExactArgs(2), + RunE: func(cmd *cobra.Command, args []string) error { + return assignMCPCmd(args[0], args[1]) + }, + } + return cmd +} + +func assignMCPCmd(agentID, mcpName string) error { + cfg, err := internal.LoadConfig() + if err != nil { + return fmt.Errorf("error loading config: %w", err) + } + + found := false + for i, a := range cfg.Agents.List { + if a.ID == agentID { + found = true + + // Check if already assigned + for _, m := range a.MCPServers { + if m == mcpName { + fmt.Printf("⚠️ MCP server '%s' is already assigned to agent '%s'\n", mcpName, agentID) + return nil + } + } + + cfg.Agents.List[i].MCPServers = append(cfg.Agents.List[i].MCPServers, mcpName) + break + } + } + + if !found { + return fmt.Errorf("agent with ID '%s' not found", agentID) + } + + if err := config.SaveConfig(internal.GetConfigPath(), cfg); err != nil { + return fmt.Errorf("failed to save config: %w", err) + } + + fmt.Printf("✅ Successfully assigned MCP server '%s' to agent '%s'\n", mcpName, agentID) + return nil +} + +func NewRemoveMCPCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "remove-mcp ", + Short: "Remove an MCP server from an agent persona", + Args: cobra.ExactArgs(2), + RunE: func(cmd *cobra.Command, args []string) error { + return removeMCPCmd(args[0], args[1]) + }, + } + return cmd +} + +func removeMCPCmd(agentID, mcpName string) error { + cfg, err := internal.LoadConfig() + if err != nil { + return fmt.Errorf("error loading config: %w", err) + } + + found := false + for i, a := range cfg.Agents.List { + if a.ID == agentID { + found = true + + newMCPs := make([]string, 0, len(a.MCPServers)) + removed := false + for _, m := range a.MCPServers { + if m == mcpName { + removed = true + continue + } + newMCPs = append(newMCPs, m) + } + + if !removed { + fmt.Printf("⚠️ MCP server '%s' is not assigned to agent '%s'\n", mcpName, agentID) + return nil + } + + cfg.Agents.List[i].MCPServers = newMCPs + break + } + } + + if !found { + return fmt.Errorf("agent with ID '%s' not found", agentID) + } + + if err := config.SaveConfig(internal.GetConfigPath(), cfg); err != nil { + return fmt.Errorf("failed to save config: %w", err) + } + + fmt.Printf("✅ Successfully removed MCP server '%s' from agent '%s'\n", mcpName, agentID) + return nil +} diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 0bfc1a241..f706efff2 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -39,6 +39,7 @@ type AgentInstance struct { Tools *tools.ToolRegistry Subagents *config.SubagentsConfig SkillsFilter []string + MCPServers []string Candidates []providers.FallbackCandidate // Router is non-nil when model routing is configured and the light model @@ -122,12 +123,14 @@ func NewAgentInstance( agentName := "" var subagents *config.SubagentsConfig var skillsFilter []string + var mcpServers []string if agentCfg != nil { agentID = routing.NormalizeAgentID(agentCfg.ID) agentName = agentCfg.Name subagents = agentCfg.Subagents skillsFilter = agentCfg.Skills + mcpServers = agentCfg.MCPServers } maxIter := defaults.MaxToolIterations @@ -246,6 +249,7 @@ func NewAgentInstance( Tools: toolsRegistry, Subagents: subagents, SkillsFilter: skillsFilter, + MCPServers: mcpServers, Candidates: candidates, Router: router, LightCandidates: lightCandidates, diff --git a/pkg/agent/loop_mcp.go b/pkg/agent/loop_mcp.go index 6f3158af5..755a23fe8 100644 --- a/pkg/agent/loop_mcp.go +++ b/pkg/agent/loop_mcp.go @@ -102,6 +102,20 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error { continue } + // Check if this agent explicitly restricts MCP servers + if agent != nil { + allowed := false + for _, allowedServer := range agent.MCPServers { + if allowedServer == serverName { + allowed = true + break + } + } + if !allowed { + continue // Skip registering this tool for this agent + } + } + mcpTool := tools.NewMCPTool(mcpManager, serverName, tool) if al.cfg.Tools.MCP.Discovery.Enabled { diff --git a/pkg/config/agent.go b/pkg/config/agent.go index 0176ce68a..483c3a9d9 100644 --- a/pkg/config/agent.go +++ b/pkg/config/agent.go @@ -47,13 +47,15 @@ func (m AgentModelConfig) MarshalJSON() ([]byte, error) { } type AgentConfig struct { - ID string `json:"id"` - Default bool `json:"default,omitempty"` - Name string `json:"name,omitempty"` - Workspace string `json:"workspace,omitempty"` - Model *AgentModelConfig `json:"model,omitempty"` - Skills []string `json:"skills,omitempty"` - Subagents *SubagentsConfig `json:"subagents,omitempty"` + ID string `json:"id"` + Default bool `json:"default,omitempty"` + Name string `json:"name,omitempty"` + SystemPrompt string `json:"system_prompt,omitempty"` + Workspace string `json:"workspace,omitempty"` + Model *AgentModelConfig `json:"model,omitempty"` + Skills []string `json:"skills,omitempty"` + MCPServers []string `json:"mcp_servers,omitempty"` + Subagents *SubagentsConfig `json:"subagents,omitempty"` } type SubagentsConfig struct { diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index fbe038a57..a5d8ef4e8 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -181,8 +181,8 @@ func TestConfig_BackwardCompat_NoAgentsList(t *testing.T) { t.Fatalf("unmarshal: %v", err) } - if len(cfg.Agents.List) != 1 { - t.Errorf("agents.list should have default clinician agent for backward compat, got %d", len(cfg.Agents.List)) + if len(cfg.Agents.List) != 5 { + t.Errorf("agents.list should have default agents for backward compat, got %d", len(cfg.Agents.List)) } if len(cfg.Bindings) != 0 { t.Errorf("bindings should be empty, got %d", len(cfg.Bindings)) diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index bd5b524d7..bc1e632fd 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -42,6 +42,30 @@ func DefaultConfig() *Config { Name: "Medical Persona", Workspace: filepath.Join(homePath, "Obsidian_Vault", "Patients"), }, + { + ID: "coding", + Name: "Coding Persona", + Workspace: filepath.Join(homePath, "workspace", "code"), + MCPServers: []string{"github", "bash"}, + }, + { + ID: "google", + Name: "Google Persona", + Workspace: filepath.Join(homePath, "workspace", "google"), + MCPServers: []string{"gmail", "calendar"}, + }, + { + ID: "communications", + Name: "Communications Persona", + Workspace: filepath.Join(homePath, "workspace", "communications"), + MCPServers: []string{"slack", "discord"}, + }, + { + ID: "financial", + Name: "Financial Persona", + Workspace: filepath.Join(homePath, "workspace", "finance"), + MCPServers: []string{"alpaca"}, + }, }, }, Bindings: []AgentBinding{},