lzh/picoclaw
1
0
Fork 0
mirror of https://github.com/sipeed/picoclaw.git synced 2026-06-12 18:08:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(agent): filter discovery by spawn permissions

Browse Source
This commit is contained in:
afjcjsbx
2026-05-07 18:26:09 +02:00
parent 96fd887cad
commit b8f4257cee
7 changed files with 170 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/agent/context.go
+5 -4
View File
@@ -26,7 +26,7 @@ type ContextBuilder struct {
skillsLoader *skills.SkillsLoader
memory *MemoryStore
splitOnMarker bool
agentDiscovery func(workspace string) []AgentDescriptor
agentDiscovery func(agentID string) []AgentDescriptor
promptRegistry *PromptRegistry
// Cache for system prompt to avoid rebuilding on every call.
@@ -68,13 +68,14 @@ func (cb *ContextBuilder) WithSplitOnMarker(enabled bool) *ContextBuilder {
}
func (cb *ContextBuilder) WithAgentDiscovery(
discover func(workspace string) []AgentDescriptor,
agentID string,
discover func(agentID string) []AgentDescriptor,
) *ContextBuilder {
cb.agentDiscovery = discover
if discover != nil {
if err := cb.RegisterPromptContributor(agentDiscoveryPromptContributor{
workspace: cb.workspace,
discover: discover,
agentID: agentID,
discover: discover,
}); err != nil {
logger.WarnCF("agent", "Failed to register agent discovery prompt contributor", map[string]any{
"error": err.Error(),
pkg/agent/discovery.go
+37 -2
View File
@@ -60,6 +60,41 @@ func (r *AgentRegistry) ListAgents(workspace string) []AgentDescriptor {
return descriptors
}
// ListSpawnableAgents returns descriptors only for agents the current agent is
// allowed to spawn. Restricted peers are intentionally omitted from discovery.
func (r *AgentRegistry) ListSpawnableAgents(agentID string) []AgentDescriptor {
r.mu.RLock()
defer r.mu.RUnlock()
parentID := routing.NormalizeAgentID(agentID)
parent, ok := r.agents[parentID]
if !ok || parent == nil {
return nil
}
ids := make([]string, 0, len(r.agents))
for id := range r.agents {
if id == parentID {
continue
}
if !agentAllowsSubagent(parent, id) {
continue
}
ids = append(ids, id)
}
sort.Strings(ids)
descriptors := make([]AgentDescriptor, 0, len(ids))
for _, id := range ids {
agent := r.agents[id]
if agent == nil {
continue
}
descriptors = append(descriptors, r.buildAgentDescriptorLocked(agent))
}
return descriptors
}
// GetAgentDescriptor returns the structured discovery payload for one agent.
func (r *AgentRegistry) GetAgentDescriptor(agentID string) (*AgentDescriptor, bool) {
r.mu.RLock()
@@ -195,7 +230,7 @@ func cleanWorkspacePath(path string) string {
}
func formatAgentDiscoverySection(agents []AgentDescriptor) string {
if len(agents) <= 1 {
if len(agents) == 0 {
return ""
}
@@ -212,7 +247,7 @@ func formatAgentDiscoverySection(agents []AgentDescriptor) string {
var header strings.Builder
header.WriteString("# Agent Discovery\n\n")
header.WriteString("This registry is authoritative for the current PicoClaw instance.\n")
header.WriteString("This registry lists the peer agents this agent is permitted to spawn.\n")
header.WriteString(
"Choose a peer based on its description. Use only agent IDs listed here when calling spawn.\n\n",
)
pkg/agent/discovery_test.go
+111 -6
View File
@@ -66,6 +66,30 @@ Handle support tickets carefully.
}
}
func TestAgentRegistry_ListSpawnableAgentsRespectsPermissions(t *testing.T) {
cfg := testCfg([]config.AgentConfig{
{
ID: "parent",
Default: true,
Subagents: &config.SubagentsConfig{
AllowAgents: []string{"child2", "child1"},
},
},
{ID: "child1"},
{ID: "child2"},
{ID: "restricted"},
})
registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
descriptors := registry.ListSpawnableAgents("parent")
if len(descriptors) != 2 {
t.Fatalf("expected 2 spawnable descriptors, got %d: %+v", len(descriptors), descriptors)
}
if descriptors[0].ID != "child1" || descriptors[1].ID != "child2" {
t.Fatalf("expected sorted spawnable peers only, got %+v", descriptors)
}
}
func TestContextBuilder_BuildMessagesIncludesAgentDiscoverySection(t *testing.T) {
mainWorkspace := setupWorkspace(t, map[string]string{
"AGENT.md": `---
@@ -90,9 +114,29 @@ Investigate deeply.
})
defer cleanupWorkspace(t, researchWorkspace)
restrictedWorkspace := setupWorkspace(t, map[string]string{
"AGENT.md": `---
name: Restricted Agent
description: Restricted specialist
---
# Agent
Handle restricted work.
`,
})
defer cleanupWorkspace(t, restrictedWorkspace)
cfg := testCfg([]config.AgentConfig{
{ID: "main", Default: true, Workspace: mainWorkspace},
{
ID: "main",
Default: true,
Workspace: mainWorkspace,
Subagents: &config.SubagentsConfig{
AllowAgents: []string{"research"},
},
},
{ID: "research", Workspace: researchWorkspace},
{ID: "restricted", Workspace: restrictedWorkspace},
})
cfg.Tools.ReadFile.Enabled = true
cfg.Tools.WriteFile.Enabled = true
@@ -121,13 +165,16 @@ Investigate deeply.
if !strings.Contains(systemPrompt, "# Agent Discovery") {
t.Fatalf("expected discovery section in system prompt, got %q", systemPrompt)
}
if !strings.Contains(systemPrompt, `"id": "main"`) ||
!strings.Contains(systemPrompt, `"id": "research"`) {
t.Fatalf("expected self and peer descriptors in discovery section, got %q", systemPrompt)
if strings.Contains(systemPrompt, `"id": "main"`) {
t.Fatalf("did not expect self descriptor in discovery section, got %q", systemPrompt)
}
if !strings.Contains(systemPrompt, `"name": "main"`) ||
if !strings.Contains(systemPrompt, `"id": "research"`) ||
!strings.Contains(systemPrompt, `"description": "Research specialist"`) {
t.Fatalf("expected minimal identity fields in discovery section, got %q", systemPrompt)
t.Fatalf("expected allowed peer descriptor in discovery section, got %q", systemPrompt)
}
if strings.Contains(systemPrompt, `"id": "restricted"`) ||
strings.Contains(systemPrompt, `"description": "Restricted specialist"`) {
t.Fatalf("did not expect restricted peer descriptor in discovery section, got %q", systemPrompt)
}
for _, forbidden := range []string{`"current_agent_id"`, `"available_tools"`, `"model"`, `"channels"`, `"skills"`, `"mcpServers"`, `"tools"`} {
if strings.Contains(systemPrompt, forbidden) {
@@ -136,6 +183,64 @@ Investigate deeply.
}
}
func TestContextBuilder_BuildMessagesOmitsAgentDiscoveryWithoutSpawnPermissions(t *testing.T) {
mainWorkspace := setupWorkspace(t, map[string]string{
"AGENT.md": `---
description: Main agent
---
# Agent
Generalist.
`,
})
defer cleanupWorkspace(t, mainWorkspace)
researchWorkspace := setupWorkspace(t, map[string]string{
"AGENT.md": `---
description: Research specialist
---
# Agent
Investigate deeply.
`,
})
defer cleanupWorkspace(t, researchWorkspace)
cfg := testCfg([]config.AgentConfig{
{ID: "main", Default: true, Workspace: mainWorkspace},
{ID: "research", Workspace: researchWorkspace},
})
cfg.Tools.ReadFile.Enabled = true
registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
mainAgent, ok := registry.GetAgent("main")
if !ok || mainAgent == nil {
t.Fatal("expected main agent")
}
messages := mainAgent.ContextBuilder.BuildMessages(
nil,
"",
"handle locally",
nil,
"telegram",
"chat-1",
"",
"",
)
if len(messages) == 0 {
t.Fatal("expected messages")
}
systemPrompt := messages[0].Content
if strings.Contains(systemPrompt, "# Agent Discovery") {
t.Fatalf("did not expect discovery section without spawn permissions, got %q", systemPrompt)
}
if strings.Contains(systemPrompt, `"id": "research"`) {
t.Fatalf("did not expect unauthorized peer identity in system prompt, got %q", systemPrompt)
}
}
func TestContextBuilder_BuildMessagesOmitsAgentDiscoverySectionForSingleton(t *testing.T) {
mainWorkspace := setupWorkspace(t, map[string]string{
"AGENT.md": `---
pkg/agent/prompt_contributors.go
+3 -3
View File
@@ -94,8 +94,8 @@ func (c mcpServerPromptContributor) ContributePrompt(
}
type agentDiscoveryPromptContributor struct {
workspace string
discover func(workspace string) []AgentDescriptor
agentID string
discover func(agentID string) []AgentDescriptor
}
func (c agentDiscoveryPromptContributor) PromptSource() PromptSourceDescriptor {
@@ -115,7 +115,7 @@ func (c agentDiscoveryPromptContributor) ContributePrompt(
if c.discover == nil {
return nil, nil
}
content := formatAgentDiscoverySection(c.discover(c.workspace))
content := formatAgentDiscoverySection(c.discover(c.agentID))
if strings.TrimSpace(content) == "" {
return nil, nil
}
pkg/agent/registry.go
+6 -3
View File
@@ -57,7 +57,7 @@ func NewAgentRegistry(
for _, instance := range registry.agents {
if instance.ContextBuilder != nil {
instance.ContextBuilder.WithAgentDiscovery(registry.ListAgents)
instance.ContextBuilder.WithAgentDiscovery(instance.ID, registry.ListSpawnableAgents)
}
}
@@ -119,10 +119,13 @@ func (r *AgentRegistry) CanSpawnSubagent(parentAgentID, targetAgentID string) bo
if !ok {
return false
}
if parent.Subagents == nil || parent.Subagents.AllowAgents == nil {
return agentAllowsSubagent(parent, routing.NormalizeAgentID(targetAgentID))
}
func agentAllowsSubagent(parent *AgentInstance, targetNorm string) bool {
if parent == nil || parent.Subagents == nil || parent.Subagents.AllowAgents == nil {
return false
}
targetNorm := routing.NormalizeAgentID(targetAgentID)
for _, allowed := range parent.Subagents.AllowAgents {
if allowed == "*" {
return true