Merge pull request #1637 from krissetto/fix-subagent-logic

Fix subagent logic

Author: dgageot

pkg/runtime/runtime.go

@@ -1779,6 +1779,22 @@ func (r *LocalRuntime) handleTaskTransfer(ctx context.Context, sess *session.Ses
 
 	a := r.CurrentAgent()
 
+	// Validate that the target agent is in the current agent's sub-agents list
+	subAgents := a.SubAgents()
+	if !slices.ContainsFunc(subAgents, func(sa *agent.Agent) bool { return sa.Name() == params.Agent }) {
+		var subAgentNames []string
+		for _, sa := range subAgents {
+			subAgentNames = append(subAgentNames, sa.Name())
+		}
+		var errorMsg string
+		if len(subAgentNames) > 0 {
+			errorMsg = fmt.Sprintf("Agent %s cannot transfer task to %s: target agent not in sub-agents list. Available sub-agent IDs are: %s", a.Name(), params.Agent, strings.Join(subAgentNames, ", "))
+		} else {
+			errorMsg = fmt.Sprintf("Agent %s cannot transfer task to %s: target agent not in sub-agents list. This agent has no sub-agents configured.", a.Name(), params.Agent)
+		}
+		return tools.ResultError(errorMsg), nil
+	}
+
 	// Span for task transfer (optional)
 	ctx, span := r.startSpan(ctx, "runtime.task_transfer", trace.WithAttributes(
 		attribute.String("from.agent", a.Name()),

pkg/runtime/runtime_test.go

@@ -1364,3 +1364,75 @@ func TestToolRejectionWithoutReason(t *testing.T) {
 	require.Equal(t, "The user rejected the tool call.", toolResponse.Response)
 	require.NotContains(t, toolResponse.Response, "Reason:")
 }
+
+func TestTransferTaskRejectsNonSubAgent(t *testing.T) {
+	// root has librarian as sub-agent but NOT planner.
+	// planner exists in the team. transfer_task to planner should be rejected.
+	prov := &mockProvider{id: "test/mock-model", stream: &mockStream{}}
+
+	librarian := agent.New("librarian", "Library agent", agent.WithModel(prov))
+	root := agent.New("root", "Root agent", agent.WithModel(prov))
+	planner := agent.New("planner", "Planner agent", agent.WithModel(prov))
+
+	agent.WithSubAgents(librarian)(root)
+
+	tm := team.New(team.WithAgents(root, planner, librarian))
+
+	rt, err := NewLocalRuntime(tm, WithSessionCompaction(false), WithModelStore(mockModelStore{}))
+	require.NoError(t, err)
+
+	sess := session.New(session.WithUserMessage("Test"))
+	evts := make(chan Event, 128)
+
+	toolCall := tools.ToolCall{
+		ID:   "call_1",
+		Type: "function",
+		Function: tools.FunctionCall{
+			Name:      "transfer_task",
+			Arguments: `{"agent":"planner","task":"do something","expected_output":""}`,
+		},
+	}
+
+	result, err := rt.handleTaskTransfer(t.Context(), sess, toolCall, evts)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+	assert.True(t, result.IsError, "transfer to non-sub-agent should return an error result")
+	assert.Contains(t, result.Output, "cannot transfer task to planner")
+	assert.Contains(t, result.Output, "librarian")
+	assert.Equal(t, "root", rt.currentAgent, "current agent should remain root")
+}
+
+func TestTransferTaskAllowsSubAgent(t *testing.T) {
+	// Verify that transfer_task to a valid sub-agent is NOT rejected by the validation.
+	// We can't fully run the child session without a real model, so we just confirm
+	// it gets past validation (it will fail later due to mock stream being empty,
+	// which is fine — we only care that it's not blocked by the sub-agent check).
+	prov := &mockProvider{id: "test/mock-model", stream: newStreamBuilder().AddContent("done").AddStopWithUsage(10, 5).Build()}
+
+	librarian := agent.New("librarian", "Library agent", agent.WithModel(prov))
+	root := agent.New("root", "Root agent", agent.WithModel(prov))
+
+	agent.WithSubAgents(librarian)(root)
+
+	tm := team.New(team.WithAgents(root, librarian))
+
+	rt, err := NewLocalRuntime(tm, WithSessionCompaction(false), WithModelStore(mockModelStore{}))
+	require.NoError(t, err)
+
+	sess := session.New(session.WithUserMessage("Test"), session.WithToolsApproved(true))
+	evts := make(chan Event, 128)
+
+	toolCall := tools.ToolCall{
+		ID:   "call_1",
+		Type: "function",
+		Function: tools.FunctionCall{
+			Name:      "transfer_task",
+			Arguments: `{"agent":"librarian","task":"find a book","expected_output":"book title"}`,
+		},
+	}
+
+	result, err := rt.handleTaskTransfer(t.Context(), sess, toolCall, evts)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+	assert.False(t, result.IsError, "transfer to valid sub-agent should succeed")
+}

pkg/session/session.go

@@ -449,7 +449,7 @@ func buildInvariantSystemMessages(a *agent.Agent) []chat.Message {
 	var messages []chat.Message
 
 	if a.HasSubAgents() {
-		subAgents := append(a.SubAgents(), a.Parents()...)
+		subAgents := a.SubAgents()
 
 		var text strings.Builder
 		var validAgentIDs []string

pkg/session/session_test.go

@@ -1,9 +1,11 @@
 package session
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/docker/cagent/pkg/agent"
 	"github.com/docker/cagent/pkg/chat"
@@ -394,3 +396,42 @@ func TestGetLastUserMessages(t *testing.T) {
 		assert.Equal(t, "Third", msgs[1])
 	})
 }
+
+func TestTransferTaskPromptExcludesParents(t *testing.T) {
+	t.Parallel()
+
+	// Build hierarchy: planner -> root -> librarian
+	// root's sub-agents: [librarian]
+	// root's parents: [planner] (set by planner listing root as a sub-agent)
+	librarian := agent.New("librarian", "", agent.WithDescription("Library agent"))
+	root := agent.New("root", "You are the root agent",
+		agent.WithDescription("Root agent"),
+	)
+	planner := agent.New("planner", "",
+		agent.WithDescription("Planner agent"),
+	)
+	// Connect: root -> librarian (root has librarian as sub-agent)
+	agent.WithSubAgents(librarian)(root)
+	// Connect: planner -> root (planner has root as sub-agent, making root's parent = planner)
+	agent.WithSubAgents(root)(planner)
+
+	// Verify parent relationship was established
+	require.Len(t, root.Parents(), 1)
+	assert.Equal(t, "planner", root.Parents()[0].Name())
+
+	s := New()
+	messages := s.GetMessages(root)
+
+	// Find the system message about sub-agents
+	var subAgentMsg string
+	for _, msg := range messages {
+		if msg.Role == chat.MessageRoleSystem && strings.Contains(msg.Content, "transfer_task") {
+			subAgentMsg = msg.Content
+			break
+		}
+	}
+
+	require.NotEmpty(t, subAgentMsg, "should have a sub-agent system message")
+	assert.Contains(t, subAgentMsg, "librarian", "should list librarian as a valid sub-agent")
+	assert.NotContains(t, subAgentMsg, "planner", "should NOT list parent agent planner as a valid transfer target")
+}