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

package mcp

import (
	"context"
	"fmt"
	"strings"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/zzet/gortex/internal/query"
)

// registerWalkGraphTool wires walk_graph — a token-budgeted, free-form
// graph traversal. Unlike the fixed-purpose traversal tools
// (get_call_chain, get_dependents) the caller picks the edge kinds and
// the direction, and the walk auto-stops once the estimated encoded
// size of the result reaches a token budget.
func (s *Server) registerWalkGraphTool() {
	edgeKindList := strings.Join(query.KnownEdgeKinds(), ", ")
	s.addTool(
		mcp.NewTool("walk_graph",
			mcp.WithDescription("Token-budgeted free-form graph traversal from a start symbol. Pick the edge kinds and direction; the walk expands breadth-first and stops automatically once the encoded result would exceed the token budget. Use it to explore a neighbourhood when the fixed-purpose tools (get_call_chain, get_dependents) don't match the relationship you want to follow."),
			mcp.WithString("id", mcp.Required(), mcp.Description("Start symbol node ID (e.g. pkg/server.go::HandleRequest).")),
			mcp.WithString("edge_kinds", mcp.Description("Comma-separated edge kinds to follow (default: calls). Valid kinds: "+edgeKindList+".")),
			mcp.WithString("direction", mcp.Description("Traversal direction: out (default — follow outgoing edges), in (incoming), or both (undirected).")),
			mcp.WithString("community", mcp.Description("Constrain the walk to a single detected community. Accepts a community ID or label. Neighbours in a different community are dropped; structural nodes with no community membership still pass. No-op until community detection has run.")),
			mcp.WithNumber("token_budget", mcp.Description("Approximate token ceiling for the encoded result (default 6000). The walk stops adding nodes once the estimate would exceed this.")),
			mcp.WithNumber("max_depth", mcp.Description("Hard cap on BFS depth, applied even when the token budget would allow deeper expansion (default 8).")),
			mcp.WithString("format", mcp.Description("Output format: json (default), gcx (GCX1 compact wire format), or toon.")),
			mcp.WithString("repo", mcp.Description("Filter results to a specific repository prefix")),
			mcp.WithString("project", mcp.Description("Filter results to repositories in a specific project")),
			mcp.WithString("ref", mcp.Description("Filter results to repositories with a specific reference tag")),
			mcp.WithString("workspace", mcp.Description("Workspace override. In workspace-bound sessions this must match the active workspace.")),
			mcp.WithString("scope", mcp.Description("Saved scope name. Ignored for git diff scopes; explicit repo/project/ref filters take precedence.")),
		),
		s.handleWalkGraph,
	)
}

// handleWalkGraph runs WalkBudgeted with the request's parameters and
// returns the resulting subgraph. The response carries budget_hit and
// stopped_at_depth so the caller knows whether the walk was truncated.
func (s *Server) handleWalkGraph(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	id, err := req.RequireString("id")
	if err != nil {
		return mcp.NewToolResultError("id is required"), nil
	}

	edgeKinds, kindErr := query.ParseEdgeKindsCSV(req.GetString("edge_kinds", "calls"))
	if kindErr != nil {
		return mcp.NewToolResultError(kindErr.Error()), nil
	}
	if len(edgeKinds) == 0 {
		// An explicit empty value would otherwise mean "every kind";
		// walk_graph's documented default is calls.
		edgeKinds, _ = query.ParseEdgeKindsCSV("calls")
	}

	direction := strings.ToLower(strings.TrimSpace(req.GetString("direction", "out")))
	switch direction {
	case "", "out", "in", "both":
	default:
		return mcp.NewToolResultError(fmt.Sprintf("direction must be out, in, or both (got %q)", direction)), nil
	}

	tokenBudget := req.GetInt("token_budget", 6000)
	if tokenBudget <= 0 {
		tokenBudget = 6000
	}
	maxDepth := req.GetInt("max_depth", 8)
	if maxDepth <= 0 {
		maxDepth = 8
	}

	eng := s.engineFor(ctx)
	seed := eng.GetSymbol(id)
	if seed == nil {
		return mcp.NewToolResultError(fmt.Sprintf("symbol not found: %s", id)), nil
	}

	resolved, errResult := s.resolveScope(ctx, req, IntentReach)
	if errResult != nil {
		return errResult, nil
	}
	if !resolvedScopeAllowsNode(resolved, seed) {
		return symbolNotFoundGuidance(id), nil
	}

	// Optional community constraint: accept either a community ID or
	// label and resolve it to an ID, then hand the walk the node->comm
	// map so it can drop cross-community neighbours. A nil community
	// result (analysis not yet run) makes the filter a no-op.
	commID, nodeToComm := s.resolveCommunityFilter(req.GetString("community", ""))

	sg := eng.WalkBudgeted(id, query.WalkOptions{
		EdgeKinds:   edgeKinds,
		Direction:   direction,
		TokenBudget: tokenBudget,
		MaxDepth:    maxDepth,
		WorkspaceID: resolved.WorkspaceID,
		ProjectID:   resolved.ProjectID,
		RepoAllow:   resolved.RepoAllow,
		CommunityID: commID,
		NodeToComm:  nodeToComm,
	})

	budgetHit, stoppedAt := sg.BudgetHit, sg.StoppedAtDepth
	sg = filterSubGraphByResolvedScope(sg, resolved)
	// Keep the traversal budget fields stable after the defensive
	// post-filter.
	sg.BudgetHit = budgetHit
	sg.StoppedAtDepth = stoppedAt
	enrichSubGraphEdges(sg)
	return s.returnScopedSubGraph(ctx, req, sg, resolved)
}

// resolveCommunityFilter turns a community ID-or-label request argument
// into the (communityID, nodeToComm) pair WalkBudgeted's community gate
// needs. It mirrors winnowSymbols' label resolution: a label is mapped
// to its community ID up front; an ID passes through. When community
// detection has not run yet (s.getCommunities() is nil) or the argument
// is empty, both results are zero-valued — the walk's gate then no-ops.
//
// An unrecognised non-empty label is returned verbatim as the community
// ID; since no node carries that membership the gate drops every
// community-tagged neighbour, which is the correct "no such community"
// behaviour (an empty walk beats silently ignoring the filter).
func (s *Server) resolveCommunityFilter(arg string) (string, map[string]string) {
	want := strings.TrimSpace(arg)
	if want == "" {
		return "", nil
	}
	cr := s.getCommunities()
	if cr == nil {
		return "", nil
	}
	for _, com := range cr.Communities {
		if com.Label != "" && com.Label == want {
			return com.ID, cr.NodeToComm
		}
	}
	return want, cr.NodeToComm
}