agent_source_test.go

package agent

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"testing"

	"github.com/panyam/mcpkit/core"
)

func childRunner(t *testing.T, turns ...StubTurn) *Runner {
	t.Helper()
	r, err := NewRunner(RunnerConfig{Provider: NewStubProvider(turns...)})
	if err != nil {
		t.Fatal(err)
	}
	return r
}

func TestAgentSource_DelegatesAndIsolates(t *testing.T) {
	ctx := context.Background()
	childStub := NewStubProvider(StubTurn{Text: "the child answer"})
	child, _ := NewRunner(RunnerConfig{Provider: childStub})
	as, err := NewAgentSource(AgentSourceConfig{Name: "researcher", Description: "does research", Runner: child})
	if err != nil {
		t.Fatal(err)
	}

	// the tool is exposed under its name
	defs, _ := as.Tools(ctx)
	if len(defs) != 1 || defs[0].Name != "researcher" {
		t.Fatalf("Tools = %+v, want one 'researcher'", defs)
	}

	res, err := as.Call(ctx, "researcher", map[string]any{"task": "find the answer"})
	if err != nil {
		t.Fatal(err)
	}
	if res.IsError || res.Content[0].Text != "the child answer" {
		t.Fatalf("call = %+v, want the child's final text", res)
	}

	// isolation: the child saw ONLY the task, not any parent history
	got := childStub.Requests()[0].Messages
	if len(got) != 1 || got[0].Role != RoleUser || got[0].Text != "find the answer" {
		t.Fatalf("child history = %+v, want just the task (isolated)", got)
	}
}

// TestAgentSource_Supervision drives a coordinator Runner whose tools are a
// MultiSource of two AgentSources — the "supervision falls out for free"
// claim, no new code.
func TestAgentSource_Supervision(t *testing.T) {
	ctx := context.Background()
	calc, _ := NewAgentSource(AgentSourceConfig{
		Name: "calc", Description: "arithmetic", Runner: childRunner(t, StubTurn{Text: "42"})})
	poet, _ := NewAgentSource(AgentSourceConfig{
		Name: "poet", Description: "verse", Runner: childRunner(t, StubTurn{Text: "a haiku"})})

	tools := NewMultiSource()
	if err := tools.Add("calc-agent", calc); err != nil {
		t.Fatal(err)
	}
	if err := tools.Add("poet-agent", poet); err != nil {
		t.Fatal(err)
	}

	coordStub := NewStubProvider(
		StubTurn{ToolCalls: []ToolCall{{ID: "c1", Name: "calc", Args: core.NewRawJSON(json.RawMessage(`{"task":"2+2"}`))}}},
		StubTurn{Text: "the sub-agent said 42"},
	)
	coord, _ := NewRunner(RunnerConfig{Provider: coordStub, Tools: tools})

	result, err := coord.Run(ctx, []Message{{Role: RoleUser, Text: "add two and two"}}, func(Event) {})
	if err != nil {
		t.Fatal(err)
	}
	if result.Text != "the sub-agent said 42" {
		t.Fatalf("coordinator final = %q", result.Text)
	}
	// the child's answer came back as the calc tool result
	fed := coordStub.Requests()[1].Messages
	var sawChild bool
	for _, m := range fed {
		if m.Role == RoleTool && strings.Contains(m.Text, "42") {
			sawChild = true
		}
	}
	if !sawChild {
		t.Fatal("child answer did not return to the coordinator as a tool result")
	}
}

func TestAgentSource_DepthCap(t *testing.T) {
	as, _ := NewAgentSource(AgentSourceConfig{
		Name: "deep", Description: "d", Runner: childRunner(t, StubTurn{Text: "x"}), MaxDepth: 2})

	// at depth 2 (== MaxDepth) the call is refused before running
	res, err := as.Call(withAgentDepth(context.Background(), 2), "deep", map[string]any{"task": "go"})
	if err != nil {
		t.Fatal(err)
	}
	if !res.IsError || !strings.Contains(res.Content[0].Text, "max depth") {
		t.Fatalf("depth cap = %+v, want an IsError max-depth refusal", res)
	}
}

func TestAgentSource_CallBudget(t *testing.T) {
	as, _ := NewAgentSource(AgentSourceConfig{
		Name: "b", Description: "d", Runner: childRunner(t, StubTurn{Text: "1"}, StubTurn{Text: "2"})})
	ctx := WithAgentCallBudget(context.Background(), 1)

	if res, _ := as.Call(ctx, "b", map[string]any{"task": "go"}); res.IsError {
		t.Fatalf("first call within budget should succeed: %+v", res)
	}
	res, _ := as.Call(ctx, "b", map[string]any{"task": "go"})
	if !res.IsError || !strings.Contains(res.Content[0].Text, "budget") {
		t.Fatalf("second call over budget = %+v, want an IsError budget refusal", res)
	}
}

func TestAgentSource_EventNesting(t *testing.T) {
	var mu sync.Mutex
	var got []SubAgentEvent
	sink := func(e SubAgentEvent) { mu.Lock(); got = append(got, e); mu.Unlock() }

	as, _ := NewAgentSource(AgentSourceConfig{
		Name: "worker", Description: "w", Runner: childRunner(t, StubTurn{Text: "done"}), OnEvent: sink})

	if _, err := as.Call(context.Background(), "worker", map[string]any{"task": "go"}); err != nil {
		t.Fatal(err)
	}
	if len(got) == 0 {
		t.Fatal("no sub-agent events forwarded")
	}
	// every forwarded event carries scope + depth on the envelope (Event stays flat)
	for _, ev := range got {
		if ev.Scope != "worker" || ev.Depth != 1 {
			t.Fatalf("envelope scope/depth = %q/%d, want worker/1", ev.Scope, ev.Depth)
		}
	}
	if got[0].Event.Kind != EventTurnBegin {
		t.Fatalf("first forwarded event = %v, want the child's turn-begin", got[0].Event.Kind)
	}
}

// TestAgentSource_NestedEventScope proves scopes compose down a tree: a
// sub-agent that spawns another sub-agent yields depth-2 events tagged with
// the joined path, all reaching one shared sink.
func TestAgentSource_NestedEventScope(t *testing.T) {
	var mu sync.Mutex
	var got []SubAgentEvent
	sink := func(e SubAgentEvent) { mu.Lock(); got = append(got, e); mu.Unlock() }

	inner, _ := NewAgentSource(AgentSourceConfig{
		Name: "inner", Description: "i", Runner: childRunner(t, StubTurn{Text: "deep"}), OnEvent: sink})
	innerTools := NewMultiSource()
	if err := innerTools.Add("inner-src", inner); err != nil {
		t.Fatal(err)
	}
	outerChild, _ := NewRunner(RunnerConfig{Tools: innerTools, Provider: NewStubProvider(
		StubTurn{ToolCalls: []ToolCall{{ID: "c1", Name: "inner", Args: core.NewRawJSON(json.RawMessage(`{"task":"dig"}`))}}},
		StubTurn{Text: "outer done"},
	)})
	outer, _ := NewAgentSource(AgentSourceConfig{
		Name: "outer", Description: "o", Runner: outerChild, OnEvent: sink})

	if _, err := outer.Call(context.Background(), "outer", map[string]any{"task": "start"}); err != nil {
		t.Fatal(err)
	}

	var sawOuter, sawInner bool
	for _, ev := range got {
		if ev.Scope == "outer" && ev.Depth == 1 {
			sawOuter = true
		}
		if ev.Scope == "outer/inner" && ev.Depth == 2 {
			sawInner = true
		}
	}
	if !sawOuter || !sawInner {
		t.Fatalf("nested scopes missing: sawOuter=%v sawInner=%v", sawOuter, sawInner)
	}
}

func TestAgentSource_MissingTaskAndValidation(t *testing.T) {
	as, _ := NewAgentSource(AgentSourceConfig{Name: "x", Description: "d", Runner: childRunner(t, StubTurn{Text: "ok"})})
	if res, _ := as.Call(context.Background(), "x", map[string]any{}); !res.IsError {
		t.Fatal("missing task should be an IsError result")
	}
	// unknown name is a dispatch error, not a result
	if _, err := as.Call(context.Background(), "nope", map[string]any{"task": "y"}); err == nil {
		t.Fatal("unknown tool name should be a dispatch error")
	}

	if _, err := NewAgentSource(AgentSourceConfig{Name: "x"}); err == nil {
		t.Fatal("missing Runner should error")
	}
	if _, err := NewAgentSource(AgentSourceConfig{Runner: childRunner(t, StubTurn{Text: "z"})}); err == nil {
		t.Fatal("missing Name should error")
	}
}

// TestAgentSource_TypedInput covers structured input (1033): with an InputSchema
// set, the tool advertises that schema and the child is seeded with the raw
// typed arguments JSON as its user turn, not the {task} string.
func TestAgentSource_TypedInput(t *testing.T) {
	ctx := context.Background()
	childStub := NewStubProvider(StubTurn{Text: "done"})
	child, _ := NewRunner(RunnerConfig{Provider: childStub})
	schema := json.RawMessage(`{"type":"object","properties":{"topic":{"type":"string"},"depth":{"type":"integer"}},"required":["topic"]}`)
	as, err := NewAgentSource(AgentSourceConfig{
		Name: "researcher", Description: "typed research", Runner: child, InputSchema: core.NewRawJSON(schema),
	})
	if err != nil {
		t.Fatal(err)
	}

	// the tool advertises the caller-provided schema, not {task}
	defs, _ := as.Tools(ctx)
	sj, _ := json.Marshal(defs[0].InputSchema)
	if !strings.Contains(string(sj), "topic") || strings.Contains(string(sj), "task") {
		t.Fatalf("tool schema should be the typed input, got %s", sj)
	}

	if _, err := as.Call(ctx, "researcher", map[string]any{"topic": "caching", "depth": 3}); err != nil {
		t.Fatal(err)
	}
	// the child's user turn is the raw typed args (a typed subtask), not a plucked task string
	seed := childStub.Requests()[0].Messages[0].Text
	if !strings.Contains(seed, "caching") || !strings.Contains(seed, "depth") {
		t.Fatalf("child seed should carry the typed args, got %q", seed)
	}
}

// TestAgentSource_ZeroInputSchemaIsAbsent pins the boundary the core.RawJSON
// change moved (#35). Presence used to be `InputSchema == nil`; on a struct
// type that test cannot be written, so it is now Len() == 0. An unset field
// must still mean the default {task} schema rather than an empty typed one,
// which would advertise a schema no caller can satisfy.
func TestAgentSource_ZeroInputSchemaIsAbsent(t *testing.T) {
	ctx := context.Background()
	child, _ := NewRunner(RunnerConfig{Provider: NewStubProvider(StubTurn{Text: "done"})})
	as, err := NewAgentSource(AgentSourceConfig{
		Name: "plain", Description: "untyped", Runner: child,
	})
	if err != nil {
		t.Fatal(err)
	}
	defs, _ := as.Tools(ctx)
	sj, _ := json.Marshal(defs[0].InputSchema)
	if !strings.Contains(string(sj), "task") {
		t.Fatalf("zero InputSchema should advertise the default {task} schema, got %s", sj)
	}
}

// TestAgentSource_StructuredOutput covers structured output (1033): a child whose
// Runner has a ResponseSchema returns its coerced JSON from Call, not the text.
func TestAgentSource_StructuredOutput(t *testing.T) {
	ctx := context.Background()
	schema := core.NewRawJSON(json.RawMessage(`{"type":"object","properties":{"score":{"type":"integer"}},"required":["score"]}`))
	childStub := NewStubProvider(
		StubTurn{Text: "the score is 7"}, // terminal text
		StubTurn{Text: `{"score":7}`},    // finalizing coercion
	)
	child, err := NewRunner(RunnerConfig{Provider: childStub, ResponseSchema: schema})
	if err != nil {
		t.Fatal(err)
	}
	as, _ := NewAgentSource(AgentSourceConfig{Name: "scorer", Description: "scores", Runner: child})

	res, err := as.Call(ctx, "scorer", map[string]any{"task": "score this"})
	if err != nil {
		t.Fatal(err)
	}
	if got := res.Content[0].Text; got != `{"score":7}` {
		t.Fatalf("structured output should be the coerced JSON, got %q", got)
	}
}