zzet--gortex/internal/mcp/tools_extract_candidates_test.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/zzet/gortex/internal/graph"
)

func newExtractTestServer(t *testing.T) *Server {
	t.Helper()
	g := graph.New()

	// orchestrator: 50 lines, 5 callers, 10 distinct callees — prime candidate.
	g.AddNode(&graph.Node{
		ID: "p/big.go::Orchestrator", Name: "Orchestrator", Kind: graph.KindFunction,
		FilePath: "p/big.go", StartLine: 1, EndLine: 50,
	})
	// 5 callers of Orchestrator.
	for i := 1; i <= 5; i++ {
		callerID := fmt.Sprintf("p/caller%d.go::Call%d", i, i)
		g.AddNode(&graph.Node{ID: callerID, Name: fmt.Sprintf("Call%d", i), Kind: graph.KindFunction, FilePath: fmt.Sprintf("p/caller%d.go", i)})
		g.AddEdge(&graph.Edge{From: callerID, To: "p/big.go::Orchestrator", Kind: graph.EdgeCalls})
	}
	// 10 distinct callees from Orchestrator.
	for i := 1; i <= 10; i++ {
		calleeID := fmt.Sprintf("p/util%d.go::U%d", i, i)
		g.AddNode(&graph.Node{ID: calleeID, Name: fmt.Sprintf("U%d", i), Kind: graph.KindFunction, FilePath: fmt.Sprintf("p/util%d.go", i)})
		g.AddEdge(&graph.Edge{From: "p/big.go::Orchestrator", To: calleeID, Kind: graph.EdgeCalls})
	}

	// short: 10 lines, 5 callers, 2 callees — fails min_lines.
	g.AddNode(&graph.Node{
		ID: "p/small.go::Short", Name: "Short", Kind: graph.KindFunction,
		FilePath: "p/small.go", StartLine: 1, EndLine: 10,
	})
	for i := 1; i <= 5; i++ {
		callerID := fmt.Sprintf("p/sc%d.go::SC%d", i, i)
		g.AddNode(&graph.Node{ID: callerID, Name: fmt.Sprintf("SC%d", i), Kind: graph.KindFunction, FilePath: fmt.Sprintf("p/sc%d.go", i)})
		g.AddEdge(&graph.Edge{From: callerID, To: "p/small.go::Short", Kind: graph.EdgeCalls})
	}

	// solo: 50 lines, 1 caller, 10 callees — fails min_callers.
	g.AddNode(&graph.Node{
		ID: "p/solo.go::Solo", Name: "Solo", Kind: graph.KindFunction,
		FilePath: "p/solo.go", StartLine: 1, EndLine: 50,
	})
	g.AddNode(&graph.Node{ID: "p/solocaller.go::SoloCaller", Name: "SoloCaller", Kind: graph.KindFunction, FilePath: "p/solocaller.go"})
	g.AddEdge(&graph.Edge{From: "p/solocaller.go::SoloCaller", To: "p/solo.go::Solo", Kind: graph.EdgeCalls})

	// flat: 50 lines, 5 callers, 1 callee — fails min_fan_out.
	g.AddNode(&graph.Node{
		ID: "p/flat.go::Flat", Name: "Flat", Kind: graph.KindFunction,
		FilePath: "p/flat.go", StartLine: 1, EndLine: 50,
	})
	for i := 1; i <= 5; i++ {
		callerID := fmt.Sprintf("p/fc%d.go::FC%d", i, i)
		g.AddNode(&graph.Node{ID: callerID, Name: fmt.Sprintf("FC%d", i), Kind: graph.KindFunction, FilePath: fmt.Sprintf("p/fc%d.go", i)})
		g.AddEdge(&graph.Edge{From: callerID, To: "p/flat.go::Flat", Kind: graph.EdgeCalls})
	}
	g.AddNode(&graph.Node{ID: "p/flatcallee.go::FCU", Name: "FCU", Kind: graph.KindFunction, FilePath: "p/flatcallee.go"})
	g.AddEdge(&graph.Edge{From: "p/flat.go::Flat", To: "p/flatcallee.go::FCU", Kind: graph.EdgeCalls})

	return &Server{
		graph:      g,
		session:    newSessionState(),
		tokenStats: &tokenStats{},
		symHistory: &symbolHistory{entries: make(map[string][]SymbolModification)},
		sessions:   newSessionMap(),
		toolScopes: newScopeRegistry(),
	}
}

func callExtractHandler(t *testing.T, s *Server, args map[string]any) map[string]any {
	t.Helper()
	req := mcp.CallToolRequest{}
	req.Params.Arguments = args
	res, err := s.handleGetExtractionCandidates(context.Background(), req)
	require.NoError(t, err)
	require.NotNil(t, res)
	require.False(t, res.IsError)
	tc, ok := res.Content[0].(mcp.TextContent)
	require.True(t, ok)
	var m map[string]any
	require.NoError(t, json.Unmarshal([]byte(tc.Text), &m))
	return m
}

func TestExtractCandidates_OrchestratorSurfaces(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{})

	cands, _ := out["candidates"].([]any)
	require.NotEmpty(t, cands)
	first := cands[0].(map[string]any)
	assert.Equal(t, "p/big.go::Orchestrator", first["symbol_id"],
		"the 50-line / 5-caller / 10-callee function should rank highest")
	assert.EqualValues(t, 50, first["line_count"].(float64))
	assert.EqualValues(t, 5, first["caller_count"].(float64))
	assert.EqualValues(t, 10, first["fan_out"].(float64))
}

func TestExtractCandidates_FailsMinLines(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{})

	cands, _ := out["candidates"].([]any)
	for _, c := range cands {
		assert.NotEqual(t, "p/small.go::Short", c.(map[string]any)["symbol_id"],
			"Short fails min_lines=20")
	}
}

func TestExtractCandidates_FailsMinCallers(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{})

	cands, _ := out["candidates"].([]any)
	for _, c := range cands {
		assert.NotEqual(t, "p/solo.go::Solo", c.(map[string]any)["symbol_id"],
			"Solo fails min_callers=2")
	}
}

func TestExtractCandidates_FailsMinFanOut(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{})

	cands, _ := out["candidates"].([]any)
	for _, c := range cands {
		assert.NotEqual(t, "p/flat.go::Flat", c.(map[string]any)["symbol_id"],
			"Flat fails min_fan_out=5")
	}
}

func TestExtractCandidates_RelaxedThresholdsSurface(t *testing.T) {
	s := newExtractTestServer(t)
	// Drop all thresholds to 0 — Solo and Flat should now appear.
	out := callExtractHandler(t, s, map[string]any{
		"min_lines":   1,
		"min_callers": 0,
		"min_fan_out": 0,
	})
	cands, _ := out["candidates"].([]any)
	ids := map[string]bool{}
	for _, c := range cands {
		ids[c.(map[string]any)["symbol_id"].(string)] = true
	}
	assert.True(t, ids["p/solo.go::Solo"])
	assert.True(t, ids["p/flat.go::Flat"])
}

func TestExtractCandidates_RationaleMentionsSignals(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{})

	cands, _ := out["candidates"].([]any)
	require.NotEmpty(t, cands)
	rationale := cands[0].(map[string]any)["rationale"].(string)
	assert.Contains(t, rationale, "lines", "rationale names line count signal")
}

func TestExtractCandidates_PathPrefixScope(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{"path_prefix": "p/big"})

	cands, _ := out["candidates"].([]any)
	require.Len(t, cands, 1)
	assert.Equal(t, "p/big.go::Orchestrator", cands[0].(map[string]any)["symbol_id"])
}

func TestExtractCandidates_LimitTruncates(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{
		"min_lines":   1,
		"min_callers": 0,
		"min_fan_out": 0,
		"limit":       1,
	})
	cands, _ := out["candidates"].([]any)
	assert.Len(t, cands, 1)
	assert.Equal(t, true, out["truncated"])
}

func TestExtractCandidates_ThresholdsEchoed(t *testing.T) {
	s := newExtractTestServer(t)
	out := callExtractHandler(t, s, map[string]any{
		"min_lines":   30,
		"min_callers": 3,
		"min_fan_out": 7,
	})
	th := out["thresholds"].(map[string]any)
	assert.EqualValues(t, 30, th["min_lines"].(float64))
	assert.EqualValues(t, 3, th["min_callers"].(float64))
	assert.EqualValues(t, 7, th["min_fan_out"].(float64))
}

func TestExtractCandidates_RationaleSignalNaming(t *testing.T) {
	cases := map[string]struct {
		line, caller, fanOut int
		mustContain          []string
	}{
		"very-long":     {60, 2, 5, []string{"very long"}},
		"widely-called": {30, 15, 5, []string{"widely called"}},
		"orchestration": {30, 5, 20, []string{"orchestration shape"}},
		"baseline":      {21, 2, 5, []string{"long", "multi-caller", "complex body"}},
	}
	for name, c := range cases {
		got := buildExtractRationale(c.line, c.caller, c.fanOut)
		for _, sub := range c.mustContain {
			assert.Contains(t, got, sub, "%s: rationale=%q", name, got)
		}
	}
}