diff --git a/mcp/src/graph/graph.ts b/mcp/src/graph/graph.ts index 677b2e2f3..f2528084c 100644 --- a/mcp/src/graph/graph.ts +++ b/mcp/src/graph/graph.ts @@ -11,7 +11,7 @@ import { GraphResponse, code_node_types, } from "./types.js"; -import { nameFileOnly, toReturnNode, formatNode, clean_node } from "./utils.js"; +import { nameFileOnly, toReturnNode, formatNode, clean_node, deser_multi } from "./utils.js"; import { getTokenizer } from "../repo/utils.js"; import { generate_services_config } from "./service.js"; @@ -520,6 +520,181 @@ function pathToSnippets(path: ShortestPath) { return r; } +const IMPACT_SEED_TYPES: NodeType[] = [ + "Function", + "Endpoint", + "Class", + "Trait", + "Datamodel", + "Page", + "Var", + "UnitTest", + "IntegrationTest", + "E2etest", +]; + +export interface ImpactParams { + files: string[]; + name: string; + node_type: string; + ref_id: string; + depth: number; + tests: boolean; +} + +interface ImpactSeed { + node_type: string; + name: string; + file: string; + ref_id: string; +} + +interface AffectedNode { + node_type: string; + name: string; + file: string; + ref_id: string; + start: number; +} + +export async function get_impact(p: ImpactParams): Promise { + let seeds: ImpactSeed[] = []; + + if (p.ref_id) { + const node = await db.find_node("" as NodeType, "", "", p.ref_id); + if (node) { + const nodeType = node.labels.find((l) => l !== "Data_Bank") || ""; + seeds.push({ + node_type: nodeType, + name: node.properties.name, + file: node.properties.file, + ref_id: node.ref_id || node.properties.ref_id || "", + }); + } + } else if (p.name) { + const nodeType = (p.node_type || "Function") as NodeType; + const node = await db.find_node(nodeType, p.name, "", ""); + if (node) { + const nt = node.labels.find((l) => l !== "Data_Bank") || nodeType; + seeds.push({ + node_type: nt, + name: node.properties.name, + file: node.properties.file, + ref_id: node.ref_id || node.properties.ref_id || "", + }); + } + } else if (p.files.length > 0) { + const seedTypes = p.node_type + ? [p.node_type as NodeType] + : IMPACT_SEED_TYPES; + const nodes = await db.find_nodes_by_files(p.files, seedTypes); + for (const node of nodes) { + const nt = node.labels.find((l) => l !== "Data_Bank") || ""; + seeds.push({ + node_type: nt, + name: node.properties.name, + file: node.properties.file, + ref_id: node.ref_id || node.properties.ref_id || "", + }); + } + } + + if (seeds.length === 0) { + return "No seed nodes found for the given input."; + } + + const seen = new Set(); + const affected: AffectedNode[] = []; + + for (const seed of seeds) { + seen.add(seed.ref_id || `${seed.name}::${seed.file}`); + } + + for (const seed of seeds) { + const record = await get_subtree({ + node_type: seed.node_type, + name: seed.name, + file: "", + ref_id: seed.ref_id, + tests: p.tests, + depth: p.depth, + direction: "up", + trim: [], + }); + if (!record) continue; + + const allNodes: Neo4jNode[] = deser_multi(record, "allNodes"); + for (const node of allNodes) { + const key = node.ref_id || node.properties.ref_id || `${node.properties.name}::${node.properties.file}`; + if (seen.has(key)) continue; + seen.add(key); + + const nt = node.labels.find((l) => l !== "Data_Bank") || "Unknown"; + affected.push({ + node_type: nt, + name: node.properties.name, + file: node.properties.file, + ref_id: node.ref_id || node.properties.ref_id || "", + start: node.properties.start || 0, + }); + } + } + + affected.sort((a, b) => a.node_type.localeCompare(b.node_type) || a.name.localeCompare(b.name)); + + const byType: Record = {}; + for (const node of affected) { + if (!byType[node.node_type]) byType[node.node_type] = []; + byType[node.node_type].push(node); + } + + let out = ""; + + out += `Seeds (${seeds.length} node${seeds.length === 1 ? "" : "s"}):\n`; + for (const s of seeds) { + out += ` ${s.node_type}: ${s.name} [${s.file}]\n`; + } + out += "\n"; + + if (affected.length === 0) { + out += "No upstream dependents found.\n"; + return out; + } + + const summaryParts = Object.entries(byType) + .map(([type, nodes]) => `${nodes.length} ${type}${nodes.length === 1 ? "" : "s"}`) + .join(", "); + out += `Impact: ${affected.length} affected nodes (${summaryParts})\n\n`; + + const displayOrder = [ + "Endpoint", "UnitTest", "IntegrationTest", "E2etest", + "Function", "Class", "Page", "Trait", "Datamodel", + ]; + const printed = new Set(); + for (const type of displayOrder) { + if (byType[type]) { + out += formatImpactGroup(type, byType[type]); + printed.add(type); + } + } + for (const [type, nodes] of Object.entries(byType)) { + if (!printed.has(type)) { + out += formatImpactGroup(type, nodes); + } + } + + return out; +} + +function formatImpactGroup(type: string, nodes: AffectedNode[]): string { + let out = `${type}:\n`; + for (const n of nodes) { + out += ` ${n.name} — ${n.file}:${n.start}\n`; + } + out += "\n"; + return out; +} + export async function get_graph( node_type: NodeType, edge_type: EdgeType, diff --git a/mcp/src/graph/neo4j.ts b/mcp/src/graph/neo4j.ts index c4e7d4738..735dce7f9 100644 --- a/mcp/src/graph/neo4j.ts +++ b/mcp/src/graph/neo4j.ts @@ -266,6 +266,25 @@ class Db { } } + async find_nodes_by_files( + files: string[], + node_types: NodeType[], + limit: number = 500, + ): Promise { + if (files.length === 0 || node_types.length === 0) return []; + const session = this.resilientSession(); + try { + const r = await session.run(Q.FIND_NODES_BY_FILES_QUERY, { + files, + node_types, + limit, + }); + return r.records.map((record) => deser_node(record, "n")); + } finally { + await session.close(); + } + } + async get_repositories(): Promise { const session = this.resilientSession(); try { diff --git a/mcp/src/graph/queries.ts b/mcp/src/graph/queries.ts index f57724147..cb229ef19 100644 --- a/mcp/src/graph/queries.ts +++ b/mcp/src/graph/queries.ts @@ -756,6 +756,14 @@ RETURN startNode, COLLECT(DISTINCT file) AS files `; +export const FIND_NODES_BY_FILES_QUERY = ` +MATCH (n:${Data_Bank}) +WHERE any(f IN $files WHERE n.file CONTAINS f) + AND any(label IN labels(n) WHERE label IN $node_types) +RETURN n +LIMIT toInteger($limit) +`; + export const REPO_SUBGRAPH_QUERY = ` WITH $node_label AS nodeLabel, $node_name as nodeName, diff --git a/mcp/src/repo/tools.ts b/mcp/src/repo/tools.ts index d88dcb0da..9b38976e2 100644 --- a/mcp/src/repo/tools.ts +++ b/mcp/src/repo/tools.ts @@ -70,6 +70,7 @@ type ToolName = | "stakgraph_search" | "stakgraph_map" | "stakgraph_code" + | "stakgraph_impact" | "jarvis" | "logs_agent"; @@ -81,6 +82,7 @@ const TOOL_NAMES: Set = new Set([ "ask_clarifying_questions", "list_concepts", "learn_concept", "learn_concepts", "list_workflows", "learn_workflow", "read_workflow_json", "vector_search", "stakgraph_search", "stakgraph_map", "stakgraph_code", + "stakgraph_impact", ]); export type SkillsConfig = Partial>; @@ -179,6 +181,8 @@ Rules: stakgraph_map: "Trace relationships from a node in the code graph. Use direction 'up' for callers and 'down' for callees.", stakgraph_code: "Retrieve actual source code for a specific node. Use ref_id from search results, or name+node_type to identify the node. Defaults to depth 1 (just the node itself).", + stakgraph_impact: + "Find all code affected by changes to specific files or functions. Shows affected endpoints, tests, and callers — the blast radius of a change.", jarvis: '', // virtual toggle: gates registration of get_ontology + graph_search tools. logs_agent: "Query runtime logs (CloudWatch / Quickwit). Use when the user asks about errors, performance, or runtime behaviour. Pass a focused, specific question.", @@ -586,8 +590,16 @@ export async function get_tools( return result.content?.[0]?.text ?? ""; }, }); + allTools.stakgraph_impact = tool({ + description: stak.ImpactTool.description || defaultDescriptions.stakgraph_impact, + inputSchema: stak.ImpactSchema, + execute: async (args: z.infer) => { + const result = await stak.impact(args); + return result.content?.[0]?.text ?? ""; + }, + }); console.log( - "===> registered stakgraph graph tools: stakgraph_search, stakgraph_map, stakgraph_code", + "===> registered stakgraph graph tools: stakgraph_search, stakgraph_map, stakgraph_code, stakgraph_impact", ); } diff --git a/mcp/src/tools/server.ts b/mcp/src/tools/server.ts index 11fb69d3d..3fdc93b68 100644 --- a/mcp/src/tools/server.ts +++ b/mcp/src/tools/server.ts @@ -64,6 +64,10 @@ server case stakgraph.GetRulesFilesTool.name: { return await stakgraph.getRulesFiles(); } + case stakgraph.ImpactTool.name: { + const fa = stakgraph.ImpactSchema.parse(args); + return await stakgraph.impact(fa); + } default: if (name.startsWith("stagehand_")) { return await stagehand.call(name, args || {}, extra.sessionId); diff --git a/mcp/src/tools/stakgraph/impact.ts b/mcp/src/tools/stakgraph/impact.ts new file mode 100644 index 000000000..7475e40f6 --- /dev/null +++ b/mcp/src/tools/stakgraph/impact.ts @@ -0,0 +1,84 @@ +import { z } from "zod"; +import { Tool } from "../types.js"; +import { parseSchema } from "../utils.js"; +import * as G from "../../graph/graph.js"; + +export const ImpactSchema = z.object({ + files: z + .array(z.string()) + .optional() + .describe( + "File paths to analyze impact for. All code nodes in these files become seeds for upstream traversal." + ), + name: z + .string() + .optional() + .describe("Name of a specific node to analyze impact for."), + node_type: z + .string() + .optional() + .describe( + "Filter seed nodes to this type (e.g. Function, Endpoint, Class)." + ), + ref_id: z + .string() + .optional() + .describe("ref_id of a specific node to start from."), + depth: z + .number() + .optional() + .default(3) + .describe("How many hops upstream to traverse (default: 3)."), + tests: z + .boolean() + .optional() + .default(true) + .describe("Whether to include affected tests (default: true)."), +}); + +export const ImpactTool: Tool = { + name: "stakgraph_impact", + description: + "Find all code affected by changes to the given files or functions. " + + "Returns affected endpoints, tests, callers, and other upstream dependents. " + + "Use this to understand the blast radius of a change.", + inputSchema: parseSchema(ImpactSchema), +}; + +export async function impact(args: z.infer) { + console.log("=> Running impact tool with args:", args); + + const hasFiles = args.files && args.files.length > 0; + const hasName = args.name && args.name.length > 0; + const hasRefId = args.ref_id && args.ref_id.length > 0; + + if (!hasFiles && !hasName && !hasRefId) { + return { + isError: true, + content: [ + { + type: "text", + text: "Error: provide at least one of files, name, or ref_id", + }, + ], + }; + } + + const result = await G.get_impact({ + files: args.files || [], + name: args.name || "", + node_type: args.node_type || "", + ref_id: args.ref_id || "", + depth: args.depth!, + tests: args.tests ?? true, + }); + + return { + content: [ + { + type: "text", + text: result, + }, + ], + }; +} diff --git a/mcp/src/tools/stakgraph/index.ts b/mcp/src/tools/stakgraph/index.ts index 8aba4883b..c1dc010a8 100644 --- a/mcp/src/tools/stakgraph/index.ts +++ b/mcp/src/tools/stakgraph/index.ts @@ -3,12 +3,14 @@ import { GetMapTool } from "./get_map.js"; import { GetRulesFilesTool } from "./get_rules_files.js"; import { ShortestPathTool } from "./shortest_path.js"; import { SearchTool } from "./search.js"; +import { ImpactTool } from "./impact.js"; export * from "./search.js"; export * from "./get_map.js"; export * from "./get_code.js"; export * from "./get_rules_files.js"; export * from "./shortest_path.js"; +export * from "./impact.js"; export const ALL_TOOLS = [ SearchTool, @@ -16,4 +18,5 @@ export const ALL_TOOLS = [ GetCodeTool, ShortestPathTool, GetRulesFilesTool, + ImpactTool, ];