feat: stakgraph_impact feature

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<string> {
+  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<string>();
+  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<string, AffectedNode[]> = {};
+  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<string>();
+  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,

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<Neo4jNode[]> {
+    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<Neo4jNode[]> {
     const session = this.resilientSession();
     try {

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,

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<string> = new Set<string>([
   "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<Record<string, boolean>>;
@@ -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<typeof stak.ImpactSchema>) => {
+        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",
     );
   }
 

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);

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<typeof ImpactSchema>) {
+  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,
+      },
+    ],
+  };
+}

mcp/src/tools/stakgraph/index.ts

@@ -3,17 +3,20 @@ 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,
   GetMapTool,
   GetCodeTool,
   ShortestPathTool,
   GetRulesFilesTool,
+  ImpactTool,
 ];