From 2b3c22684e37beb325fd8858ca79170939d78215 Mon Sep 17 00:00:00 2001 From: Evanfeenstra Date: Fri, 10 Jul 2026 09:22:18 -0600 Subject: [PATCH] feat: surface edge counts in graph_search and graph_get Both Jarvis knowledge-graph tools now return an "edges" map ({EDGE_TYPE: count}) per node so an agent can see how connected a node is and which relationship types it can traverse next with graph_neighbors. - graph_search: sends include_edge_counts=true so Jarvis attaches the edges map inline (one call, no per-node round trips). - graph_get: fetches the existing /connection-counts endpoint and collapses its (edge_type, target_type, count) rows into {EDGE_TYPE: count} via the new exported collapseConnectionCounts helper. Adds an optional namespace param so the count is scoped correctly. Best effort: never fails the call if the lookup errors. Adds unit tests for collapseConnectionCounts and the updated search URL. --- mcp/src/repo/__tests__/toolsJarvis.test.ts | 45 ++++++++++++++- mcp/src/repo/toolsJarvis.ts | 66 +++++++++++++++++++++- 2 files changed, 106 insertions(+), 5 deletions(-) diff --git a/mcp/src/repo/__tests__/toolsJarvis.test.ts b/mcp/src/repo/__tests__/toolsJarvis.test.ts index ded46c159..0d65fa00a 100644 --- a/mcp/src/repo/__tests__/toolsJarvis.test.ts +++ b/mcp/src/repo/__tests__/toolsJarvis.test.ts @@ -1,5 +1,5 @@ import { test, expect } from "../../testkit.js"; -import { buildOntologyPayload } from "../toolsJarvis.js"; +import { buildOntologyPayload, collapseConnectionCounts } from "../toolsJarvis.js"; // ── graph_search URL construction helpers ──────────────────────────────────── // Simulate the URL-building logic from graph_search in toolsJarvis.ts so we @@ -27,6 +27,7 @@ function buildJarvisSearchUrl( const params = new URLSearchParams({ q, limit: String(limit) }); if (type) params.set("type", type); if (domains) params.set("domains", domains); + params.set("include_edge_counts", "true"); appendNs(params, namespace); return `${baseUrl}/v2/nodes?${params.toString()}`; } @@ -54,9 +55,14 @@ test.describe("graph_search URL construction (toolsJarvis.ts)", () => { expect(url).toContain("namespace=MyNamespace"); }); - test("without namespace, URL is identical to today's baseline", () => { + test("without namespace, URL carries q, limit, and include_edge_counts", () => { const url = buildJarvisSearchUrl(BASE, { q: "bitcoin", limit: 10 }); - expect(url).toBe(`${BASE}/v2/nodes?q=bitcoin&limit=10`); + expect(url).toBe(`${BASE}/v2/nodes?q=bitcoin&limit=10&include_edge_counts=true`); + }); + + test("always requests inline edge counts", () => { + const url = buildJarvisSearchUrl(BASE, { q: "bitcoin" }); + expect(url).toContain("include_edge_counts=true"); }); test("with type, domains, and namespace, all params appear", () => { @@ -227,3 +233,36 @@ test.describe("buildOntologyPayload", () => { expect(person?.description).toBe("A person node"); }); }); + +test.describe("collapseConnectionCounts", () => { + test("sums counts across target types per edge_type", () => { + const edges = collapseConnectionCounts([ + { edge_type: "CONTAINS", target_type: "File", count: 3 }, + { edge_type: "CONTAINS", target_type: "Function", count: 2 }, + { edge_type: "PART_OF", target_type: "Repository", count: 1 }, + ]); + expect(edges).toEqual({ CONTAINS: 5, PART_OF: 1 }); + }); + + test("works when target_type is absent", () => { + const edges = collapseConnectionCounts([ + { edge_type: "CITES", count: 4 }, + { edge_type: "CITES", count: 1 }, + ]); + expect(edges).toEqual({ CITES: 5 }); + }); + + test("returns empty object for empty or missing input", () => { + expect(collapseConnectionCounts([])).toEqual({}); + expect(collapseConnectionCounts(undefined as any)).toEqual({}); + }); + + test("coerces non-numeric counts and skips rows without edge_type", () => { + const edges = collapseConnectionCounts([ + { edge_type: "", target_type: "File", count: 9 } as any, + { edge_type: "HAS", target_type: "Tag", count: "2" as any }, + { edge_type: "HAS", target_type: "Tag", count: undefined as any }, + ]); + expect(edges).toEqual({ HAS: 2 }); + }); +}); diff --git a/mcp/src/repo/toolsJarvis.ts b/mcp/src/repo/toolsJarvis.ts index 9ac88b9ff..89282b8b8 100644 --- a/mcp/src/repo/toolsJarvis.ts +++ b/mcp/src/repo/toolsJarvis.ts @@ -37,6 +37,23 @@ function toPythonListLiteral(arr: string[]): string { return `[${arr.map((s) => `"${s}"`).join(",")}]`; } +/** + * Collapse Jarvis `/connection-counts` rows ([{edge_type, target_type, count}]) + * into a compact `{EDGE_TYPE: totalCount}` map, summing across target types. + * This mirrors the inline `edges` map returned by graph_search so both tools + * present connectivity the same way. + */ +export function collapseConnectionCounts( + counts: Array<{ edge_type: string; target_type?: string; count: number }>, +): Record { + const out: Record = {}; + for (const c of counts ?? []) { + if (!c?.edge_type) continue; + out[c.edge_type] = (out[c.edge_type] ?? 0) + Number(c.count ?? 0); + } + return out; +} + /** * Jarvis nodes keep their human label under wildly different keys depending on * node type. Try a generous ordered list of candidates — short identifier-like @@ -235,6 +252,8 @@ export function registerJarvisTools( description: "Search the Jarvis knowledge graph for ontology nodes — people, topics, episodes, clips, organizations, workflows, and more. " + "Unlike stakgraph_search (code nodes only), this queries the full Jarvis ontology. " + + "Each result includes an `edges` map ({EDGE_TYPE: count}) showing how connected the node is and " + + "which relationship types you can traverse next with graph_neighbors. " + "Call get_ontology first to discover valid values for the `type` parameter.", inputSchema: z.object({ q: z.string().describe("The search query"), @@ -281,6 +300,9 @@ export function registerJarvisTools( const params = new URLSearchParams({ q, limit: String(limit) }); if (type) params.set("type", type); if (domains) params.set("domains", domains); + // Ask Jarvis to attach a per-node {EDGE_TYPE: count} map inline so the + // agent can gauge connectivity and see hop targets in one call. + params.set("include_edge_counts", "true"); appendNamespace(params, namespace); const url = `${jarvisUrl}/v2/nodes?${params.toString()}`; console.log( @@ -307,6 +329,9 @@ export function registerJarvisTools( n.properties?.summary ?? n.properties?.text ?? "", + // {EDGE_TYPE: count} map of this node's relationships — shows how + // connected it is and which edge types graph_neighbors can follow. + edges: (n.edges ?? {}) as Record, })) ); } catch (err: any) { @@ -319,11 +344,26 @@ export function registerJarvisTools( description: "Resolve a single node in the Jarvis knowledge graph to its full content by ref_id. " + "Use the ref_id from graph_search or graph_neighbors results. " + - "Returns the node's ref_id, node_type, derived name, and properties.", + "Returns the node's ref_id, node_type, derived name, properties, and an " + + "`edges` map ({EDGE_TYPE: count}) showing how connected the node is and " + + "which relationship types you can traverse next with graph_neighbors.", inputSchema: z.object({ ref_id: z.string().describe("The ref_id of the node to resolve."), + namespace: z + .string() + .optional() + .describe( + "Scope edge-count computation to a Jarvis namespace (data partition). " + + "Only affects the `edges` map. Not an access-control boundary." + ), }), - execute: async ({ ref_id }: { ref_id: string }) => { + execute: async ({ + ref_id, + namespace, + }: { + ref_id: string; + namespace?: string; + }) => { // limit=1 keeps Jarvis from materializing the node's whole neighborhood // (which can OOM Neo4j for hub nodes) — we only read the node itself. const url = `${jarvisUrl}/v2/nodes/${encodeURIComponent(ref_id)}?limit=1`; @@ -342,11 +382,33 @@ export function registerJarvisTools( : data; if (!raw || !raw.ref_id) return `node not found: ${ref_id}`; const properties = (raw.properties ?? {}) as Record; + + // Fetch edge-type connectivity from the dedicated aggregation endpoint + // (cheap: counts only, no neighbor materialization). Collapse the + // (edge_type, target_type) breakdown into a {EDGE_TYPE: count} map so + // graph_get and graph_search present connectivity identically. Best + // effort — never fail the whole call if this lookup errors. + let edges: Record = {}; + try { + const ccParams = new URLSearchParams(); + appendNamespace(ccParams, namespace); + const ccQuery = ccParams.toString(); + const ccUrl = `${jarvisUrl}/v2/nodes/${encodeURIComponent(ref_id)}/connection-counts${ccQuery ? `?${ccQuery}` : ""}`; + const ccResp = await jarvisFetch(ccUrl, jarvisHeaders); + if (ccResp.ok) { + const ccData = (await ccResp.json()) as any; + edges = collapseConnectionCounts(ccData?.counts ?? []); + } + } catch { + // ignore — edges stays {} + } + return JSON.stringify({ ref_id: raw.ref_id, node_type: raw.node_type, name: deriveNodeName(raw, properties), properties: raw.properties, + edges, }); } catch (err: any) { return `graph_get failed: ${err?.message ?? String(err)}`;