From ed0af0635e4b898c431b1102cf9d3a04d47fd02e Mon Sep 17 00:00:00 2001 From: carlos-alm <127798846+carlos-alm@users.noreply.github.com> Date: Sun, 1 Mar 2026 23:47:39 -0700 Subject: [PATCH 1/2] feat: add batch querying for multi-agent dispatch Add `codegraph batch [targets...]` CLI command, `batchData()` programmatic API, and `batch_query` MCP tool. Runs the same query against multiple targets in one call, returning all results in a single JSON payload with per-target error isolation. Supports 10 commands: fn-impact, context, explain, where, query, fn, impact, deps, flow, complexity. Accepts targets via positional args, --from-file (JSON array or newline-delimited), or --stdin. Impact: 3 functions changed, 3 affected --- src/batch.js | 90 +++++++++++++ src/cli.js | 52 ++++++++ src/index.js | 2 + src/mcp.js | 56 ++++++++ tests/integration/batch.test.js | 227 ++++++++++++++++++++++++++++++++ tests/unit/mcp.test.js | 1 + 6 files changed, 428 insertions(+) create mode 100644 src/batch.js create mode 100644 tests/integration/batch.test.js diff --git a/src/batch.js b/src/batch.js new file mode 100644 index 000000000..29bf88422 --- /dev/null +++ b/src/batch.js @@ -0,0 +1,90 @@ +/** + * Batch query orchestration — run the same query command against multiple targets + * and return all results in a single JSON payload. + * + * Designed for multi-agent swarms that need to dispatch 20+ queries in one call. + */ + +import { complexityData } from './complexity.js'; +import { flowData } from './flow.js'; +import { + contextData, + explainData, + fileDepsData, + fnDepsData, + fnImpactData, + impactAnalysisData, + queryNameData, + whereData, +} from './queries.js'; + +/** + * Map of supported batch commands → their data function + first-arg semantics. + * `sig` describes how the target string is passed to the data function: + * - 'name' → dataFn(target, dbPath, opts) + * - 'target' → dataFn(target, dbPath, opts) + * - 'file' → dataFn(target, dbPath, opts) + * - 'dbOnly' → dataFn(dbPath, { target, ...opts }) (target goes into opts) + */ +export const BATCH_COMMANDS = { + 'fn-impact': { fn: fnImpactData, sig: 'name' }, + context: { fn: contextData, sig: 'name' }, + explain: { fn: explainData, sig: 'target' }, + where: { fn: whereData, sig: 'target' }, + query: { fn: queryNameData, sig: 'name' }, + fn: { fn: fnDepsData, sig: 'name' }, + impact: { fn: impactAnalysisData, sig: 'file' }, + deps: { fn: fileDepsData, sig: 'file' }, + flow: { fn: flowData, sig: 'name' }, + complexity: { fn: complexityData, sig: 'dbOnly' }, +}; + +/** + * Run a query command against multiple targets, returning all results. + * + * @param {string} command - One of the keys in BATCH_COMMANDS + * @param {string[]} targets - List of target names/paths + * @param {string} [customDbPath] - Path to graph.db + * @param {object} [opts] - Shared options passed to every invocation + * @returns {{ command: string, total: number, succeeded: number, failed: number, results: object[] }} + */ +export function batchData(command, targets, customDbPath, opts = {}) { + const entry = BATCH_COMMANDS[command]; + if (!entry) { + throw new Error( + `Unknown batch command "${command}". Valid commands: ${Object.keys(BATCH_COMMANDS).join(', ')}`, + ); + } + + const results = []; + let succeeded = 0; + let failed = 0; + + for (const target of targets) { + try { + let data; + if (entry.sig === 'dbOnly') { + // complexityData(dbPath, { target, ...opts }) + data = entry.fn(customDbPath, { ...opts, target }); + } else { + // All other: dataFn(target, dbPath, opts) + data = entry.fn(target, customDbPath, opts); + } + results.push({ target, ok: true, data }); + succeeded++; + } catch (err) { + results.push({ target, ok: false, error: err.message }); + failed++; + } + } + + return { command, total: targets.length, succeeded, failed, results }; +} + +/** + * CLI wrapper — calls batchData and prints JSON to stdout. + */ +export function batch(command, targets, customDbPath, opts = {}) { + const data = batchData(command, targets, customDbPath, opts); + console.log(JSON.stringify(data, null, 2)); +} diff --git a/src/cli.js b/src/cli.js index 9ff8a5980..ca3d0d003 100644 --- a/src/cli.js +++ b/src/cli.js @@ -4,6 +4,7 @@ import fs from 'node:fs'; import path from 'node:path'; import { Command } from 'commander'; import { audit } from './audit.js'; +import { BATCH_COMMANDS, batch } from './batch.js'; import { buildGraph } from './builder.js'; import { loadConfig } from './config.js'; import { findCycles, formatCycles } from './cycles.js'; @@ -1150,4 +1151,55 @@ program } }); +program + .command('batch [targets...]') + .description( + `Run a query against multiple targets in one call. Output is always JSON.\nValid commands: ${Object.keys(BATCH_COMMANDS).join(', ')}`, + ) + .option('-d, --db ', 'Path to graph.db') + .option('--from-file ', 'Read targets from file (JSON array or newline-delimited)') + .option('--stdin', 'Read targets from stdin (JSON array)') + .option('--depth ', 'Traversal depth passed to underlying command') + .option('-f, --file ', 'Scope to file (partial match)') + .option('-k, --kind ', 'Filter by symbol kind') + .option('-T, --no-tests', 'Exclude test/spec files from results') + .option('--include-tests', 'Include test/spec files (overrides excludeTests config)') + .action(async (command, positionalTargets, opts) => { + if (opts.kind && !ALL_SYMBOL_KINDS.includes(opts.kind)) { + console.error(`Invalid kind "${opts.kind}". Valid: ${ALL_SYMBOL_KINDS.join(', ')}`); + process.exit(1); + } + + let targets; + if (opts.fromFile) { + const raw = fs.readFileSync(opts.fromFile, 'utf-8').trim(); + if (raw.startsWith('[')) { + targets = JSON.parse(raw); + } else { + targets = raw.split(/\r?\n/).filter(Boolean); + } + } else if (opts.stdin) { + const chunks = []; + for await (const chunk of process.stdin) chunks.push(chunk); + const raw = Buffer.concat(chunks).toString('utf-8').trim(); + targets = raw.startsWith('[') ? JSON.parse(raw) : raw.split(/\r?\n/).filter(Boolean); + } else { + targets = positionalTargets; + } + + if (!targets || targets.length === 0) { + console.error('No targets provided. Pass targets as arguments, --from-file, or --stdin.'); + process.exit(1); + } + + const batchOpts = { + depth: opts.depth ? parseInt(opts.depth, 10) : undefined, + file: opts.file, + kind: opts.kind, + noTests: resolveNoTests(opts), + }; + + batch(command, targets, opts.db, batchOpts); + }); + program.parse(); diff --git a/src/index.js b/src/index.js index 2f0fdcd62..f064c9944 100644 --- a/src/index.js +++ b/src/index.js @@ -7,6 +7,8 @@ // Audit (composite report) export { audit, auditData } from './audit.js'; +// Batch querying +export { BATCH_COMMANDS, batch, batchData } from './batch.js'; // Branch comparison export { branchCompareData, branchCompareMermaid } from './branch-compare.js'; // Graph building diff --git a/src/mcp.js b/src/mcp.js index bf8cbafea..30fcf6e7b 100644 --- a/src/mcp.js +++ b/src/mcp.js @@ -553,6 +553,52 @@ const BASE_TOOLS = [ required: ['target'], }, }, + { + name: 'batch_query', + description: + 'Run a query command against multiple targets in one call. Returns all results in a single JSON payload — ideal for multi-agent dispatch.', + inputSchema: { + type: 'object', + properties: { + command: { + type: 'string', + enum: [ + 'fn-impact', + 'context', + 'explain', + 'where', + 'query', + 'fn', + 'impact', + 'deps', + 'flow', + 'complexity', + ], + description: 'The query command to run for each target', + }, + targets: { + type: 'array', + items: { type: 'string' }, + description: 'List of target names (symbol names or file paths depending on command)', + }, + depth: { + type: 'number', + description: 'Traversal depth (for fn-impact, context, fn, flow)', + }, + file: { + type: 'string', + description: 'Scope to file (partial match)', + }, + kind: { + type: 'string', + enum: ALL_SYMBOL_KINDS, + description: 'Filter symbol kind', + }, + no_tests: { type: 'boolean', description: 'Exclude test files', default: false }, + }, + required: ['command', 'targets'], + }, + }, { name: 'branch_compare', description: @@ -1035,6 +1081,16 @@ export async function startMCPServer(customDbPath, options = {}) { }); break; } + case 'batch_query': { + const { batchData } = await import('./batch.js'); + result = batchData(args.command, args.targets, dbPath, { + depth: args.depth, + file: args.file, + kind: args.kind, + noTests: args.no_tests, + }); + break; + } case 'branch_compare': { const { branchCompareData, branchCompareMermaid } = await import('./branch-compare.js'); const bcData = await branchCompareData(args.base, args.target, { diff --git a/tests/integration/batch.test.js b/tests/integration/batch.test.js new file mode 100644 index 000000000..6e3020361 --- /dev/null +++ b/tests/integration/batch.test.js @@ -0,0 +1,227 @@ +/** + * Integration tests for src/batch.js + * + * Uses a hand-crafted in-memory DB (same pattern as queries.test.js / audit.test.js). + * + * Test graph (5 function nodes, 4 edges): + * authenticate → validateToken + * handleRoute → authenticate + * handleRoute → formatResponse + * formatResponse (leaf) + * validateToken (leaf) + */ + +import { execFileSync } from 'node:child_process'; +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import Database from 'better-sqlite3'; +import { afterAll, beforeAll, describe, expect, test } from 'vitest'; +import { BATCH_COMMANDS, batchData } from '../../src/batch.js'; +import { initSchema } from '../../src/db.js'; + +// ─── Helpers ─────────────────────────────────────────────────────────── + +function insertNode(db, name, kind, file, line) { + return db + .prepare('INSERT INTO nodes (name, kind, file, line) VALUES (?, ?, ?, ?)') + .run(name, kind, file, line).lastInsertRowid; +} + +function insertEdge(db, sourceId, targetId, kind = 'calls', confidence = 1.0) { + db.prepare( + 'INSERT INTO edges (source_id, target_id, kind, confidence, dynamic) VALUES (?, ?, ?, ?, 0)', + ).run(sourceId, targetId, kind, confidence); +} + +// ─── Fixture DB ──────────────────────────────────────────────────────── + +let tmpDir, dbPath; + +beforeAll(() => { + tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'codegraph-batch-')); + fs.mkdirSync(path.join(tmpDir, '.codegraph')); + dbPath = path.join(tmpDir, '.codegraph', 'graph.db'); + + const db = new Database(dbPath); + db.pragma('journal_mode = WAL'); + initSchema(db); + + // File nodes + const fileAuth = insertNode(db, 'src/auth.js', 'file', 'src/auth.js', 0); + const fileRoutes = insertNode(db, 'src/routes.js', 'file', 'src/routes.js', 0); + insertNode(db, 'src/utils.js', 'file', 'src/utils.js', 0); + + // Function nodes + const fnAuth = insertNode(db, 'authenticate', 'function', 'src/auth.js', 5); + const fnValidate = insertNode(db, 'validateToken', 'function', 'src/auth.js', 20); + const fnRoute = insertNode(db, 'handleRoute', 'function', 'src/routes.js', 10); + const fnFormat = insertNode(db, 'formatResponse', 'function', 'src/utils.js', 1); + + // File-level imports: routes → auth + insertEdge(db, fileRoutes, fileAuth, 'imports'); + + // Call edges + insertEdge(db, fnAuth, fnValidate); + insertEdge(db, fnRoute, fnAuth); + insertEdge(db, fnRoute, fnFormat); + + db.close(); +}); + +afterAll(() => { + fs.rmSync(tmpDir, { recursive: true, force: true }); +}); + +// ─── batchData: success cases ──────────────────────────────────────── + +describe('batchData — success', () => { + test('query: multiple targets both succeed', () => { + const data = batchData('query', ['authenticate', 'handleRoute'], dbPath); + expect(data.command).toBe('query'); + expect(data.total).toBe(2); + expect(data.succeeded).toBe(2); + expect(data.failed).toBe(0); + expect(data.results).toHaveLength(2); + expect(data.results[0].ok).toBe(true); + expect(data.results[0].target).toBe('authenticate'); + expect(data.results[0].data).toBeDefined(); + expect(data.results[1].ok).toBe(true); + expect(data.results[1].target).toBe('handleRoute'); + }); + + test('where: single target works', () => { + const data = batchData('where', ['authenticate'], dbPath); + expect(data.total).toBe(1); + expect(data.succeeded).toBe(1); + expect(data.results[0].ok).toBe(true); + expect(data.results[0].data.target).toBe('authenticate'); + }); + + test('explain: file targets', () => { + const data = batchData('explain', ['src/auth.js', 'src/utils.js'], dbPath); + expect(data.total).toBe(2); + expect(data.succeeded).toBe(2); + for (const r of data.results) { + expect(r.ok).toBe(true); + expect(r.data).toBeDefined(); + } + }); + + test('fn-impact: returns impact data', () => { + const data = batchData('fn-impact', ['authenticate'], dbPath); + expect(data.succeeded).toBe(1); + expect(data.results[0].data.name).toBe('authenticate'); + }); + + test('fn: returns dependency chain', () => { + const data = batchData('fn', ['handleRoute'], dbPath); + expect(data.succeeded).toBe(1); + expect(data.results[0].ok).toBe(true); + }); + + test('context: with depth option', () => { + const data = batchData('context', ['authenticate'], dbPath, { depth: 1 }); + expect(data.succeeded).toBe(1); + expect(data.results[0].ok).toBe(true); + }); +}); + +// ─── batchData: partial failure ────────────────────────────────────── + +describe('batchData — partial failure', () => { + test('non-existent target returns ok:true with empty results (no throw)', () => { + // fnImpactData returns { name, results: [] } for non-existent symbols — it doesn't throw + const data = batchData('fn-impact', ['authenticate', 'nonExistentSymbol'], dbPath); + expect(data.total).toBe(2); + expect(data.succeeded).toBe(2); + expect(data.failed).toBe(0); + const found = data.results.find((r) => r.target === 'authenticate'); + expect(found.ok).toBe(true); + expect(found.data.results.length).toBeGreaterThanOrEqual(1); + const notFound = data.results.find((r) => r.target === 'nonExistentSymbol'); + expect(notFound.ok).toBe(true); + expect(notFound.data.results).toEqual([]); + }); + + test('errors are captured per-target when data function throws', () => { + // Use a non-existent DB path to force an error per target + const badDb = path.join(tmpDir, '.codegraph', 'nonexistent.db'); + const data = batchData('query', ['anything'], badDb); + expect(data.total).toBe(1); + expect(data.failed).toBe(1); + expect(data.results[0].ok).toBe(false); + expect(data.results[0].error).toBeDefined(); + }); +}); + +// ─── batchData: edge cases ─────────────────────────────────────────── + +describe('batchData — edge cases', () => { + test('empty targets returns empty results', () => { + const data = batchData('query', [], dbPath); + expect(data.total).toBe(0); + expect(data.succeeded).toBe(0); + expect(data.failed).toBe(0); + expect(data.results).toEqual([]); + }); + + test('unknown command throws', () => { + expect(() => batchData('invalid-cmd', ['add'], dbPath)).toThrow(/Unknown batch command/); + }); + + test('BATCH_COMMANDS has all expected keys', () => { + const expected = [ + 'fn-impact', + 'context', + 'explain', + 'where', + 'query', + 'fn', + 'impact', + 'deps', + 'flow', + 'complexity', + ]; + for (const cmd of expected) { + expect(BATCH_COMMANDS).toHaveProperty(cmd); + } + }); + + test('shared opts are forwarded (noTests)', () => { + const data = batchData('query', ['authenticate'], dbPath, { noTests: true }); + expect(data.succeeded).toBe(1); + expect(data.results[0].ok).toBe(true); + }); +}); + +// ─── complexity (dbOnly sig) ───────────────────────────────────────── + +describe('batchData — complexity (dbOnly signature)', () => { + test('complexity command uses target as opts.target', () => { + const data = batchData('complexity', ['authenticate'], dbPath); + expect(data.total).toBe(1); + // complexityData returns functions array — it won't error for unknown targets + expect(data.results[0].ok).toBe(true); + expect(data.results[0].data).toHaveProperty('functions'); + }); +}); + +// ─── CLI smoke test ────────────────────────────────────────────────── + +describe('batch CLI', () => { + test('outputs valid JSON', () => { + const cliPath = path.resolve( + path.dirname(new URL(import.meta.url).pathname.replace(/^\/([A-Z]:)/i, '$1')), + '../../src/cli.js', + ); + const out = execFileSync('node', [cliPath, 'batch', 'query', 'authenticate', '--db', dbPath], { + encoding: 'utf-8', + timeout: 30_000, + }); + const parsed = JSON.parse(out); + expect(parsed.command).toBe('query'); + expect(parsed.total).toBe(1); + expect(parsed.results).toHaveLength(1); + }); +}); diff --git a/tests/unit/mcp.test.js b/tests/unit/mcp.test.js index 3135837cf..0e35001bd 100644 --- a/tests/unit/mcp.test.js +++ b/tests/unit/mcp.test.js @@ -35,6 +35,7 @@ const ALL_TOOL_NAMES = [ 'communities', 'code_owners', 'audit', + 'batch_query', 'branch_compare', 'list_repos', ]; From 2355b4364a041146dc5460fcd41e51cf0f29f36a Mon Sep 17 00:00:00 2001 From: carlos-alm <127798846+carlos-alm@users.noreply.github.com> Date: Mon, 2 Mar 2026 00:06:18 -0700 Subject: [PATCH 2/2] fix: correct spread order in batch.js comments to match implementation Impact: 1 functions changed, 1 affected --- src/batch.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/batch.js b/src/batch.js index 29bf88422..4e6778d72 100644 --- a/src/batch.js +++ b/src/batch.js @@ -24,7 +24,7 @@ import { * - 'name' → dataFn(target, dbPath, opts) * - 'target' → dataFn(target, dbPath, opts) * - 'file' → dataFn(target, dbPath, opts) - * - 'dbOnly' → dataFn(dbPath, { target, ...opts }) (target goes into opts) + * - 'dbOnly' → dataFn(dbPath, { ...opts, target }) (target goes into opts) */ export const BATCH_COMMANDS = { 'fn-impact': { fn: fnImpactData, sig: 'name' }, @@ -64,7 +64,7 @@ export function batchData(command, targets, customDbPath, opts = {}) { try { let data; if (entry.sig === 'dbOnly') { - // complexityData(dbPath, { target, ...opts }) + // complexityData(dbPath, { ...opts, target }) data = entry.fn(customDbPath, { ...opts, target }); } else { // All other: dataFn(target, dbPath, opts)