diff --git a/__tests__/kernel-retry-materialize.test.ts b/__tests__/kernel-retry-materialize.test.ts index 5de86ffcb..33811ec1e 100644 --- a/__tests__/kernel-retry-materialize.test.ts +++ b/__tests__/kernel-retry-materialize.test.ts @@ -25,6 +25,7 @@ import * as os from 'node:os'; import { CodeGraph } from '../src'; import { initGrammars, loadGrammarsForLanguages } from '../src/extraction/grammars'; import { tryKernelExtractRaw } from '../src/extraction/kernel'; +import { applyFrameworkExtracts } from '../src/extraction/framework-extracts'; import type { ExtractionResult } from '../src/types'; const KERNEL_PATH = path.join( @@ -95,6 +96,75 @@ describe.skipIf(!kernelBuilt || process.env.CODEGRAPH_KERNEL === '0')('kernel bu expect(nodes.map((n) => n.name)).toContain('target_fn'); expect(nodes.map((n) => n.name)).toContain('Adapter'); }); + + it('framework extract() output on the buffer transport is stored after decode', async () => { + const source = + 'from django.urls import path\n' + + 'from . import views\n' + + '\n' + + 'urlpatterns = [\n' + + ' path("orders//", views.order_detail, name="order-detail"),\n' + + ']\n'; + const filePath = 'urls.py'; + fs.writeFileSync(path.join(dir, filePath), source); + + const raw = tryKernelExtractRaw(filePath, source, 'python'); + expect(raw).not.toBeNull(); + const transport: ExtractionResult = { + nodes: [], + edges: [], + unresolvedReferences: [], + errors: raw!.errors, + durationMs: 0, + kernelBuffers: raw!.buffers, + kernelCounts: raw!.counts, + }; + // What parse-worker now does instead of falling back to decode: + // framework hooks run on (filePath, source) and append to the transport. + applyFrameworkExtracts(transport, filePath, source, 'python', ['django']); + const fwNodeCount = transport.nodes.length; + expect(fwNodeCount).toBeGreaterThan(0); + expect(transport.unresolvedReferences.length).toBeGreaterThan(0); + + const stats = fs.statSync(path.join(dir, filePath)); + const orchestrator = (cg as unknown as { orchestrator: { storeExtractionResult(f: string, c: string, l: string, s: fs.Stats, r: ExtractionResult): Promise } }).orchestrator; + await orchestrator.storeExtractionResult(filePath, source, 'python', stats, transport); + + // Kernel nodes + framework nodes both land — the appended transport + // arrays are pushed onto the decoded tables, not dropped. + const nodes = cg.getNodesInFile(filePath); + expect(nodes.length).toBe(raw!.counts.nodes + fwNodeCount); + expect(nodes.some((n) => n.name === 'order_detail' || n.kind === 'route' || n.filePath === filePath)).toBe(true); + }); + + it('materializeKernelResult appends framework nodes/refs after the kernel tables', async () => { + const { materializeKernelResult } = await import('../src/extraction/kernel'); + const source = + 'from django.urls import path\n' + + 'urlpatterns = [path("a/", views.a)]\n'; + const filePath = 'urls.py'; + + const raw = tryKernelExtractRaw(filePath, source, 'python'); + expect(raw).not.toBeNull(); + const transport: ExtractionResult = { + nodes: [], + edges: [], + unresolvedReferences: [], + errors: raw!.errors, + durationMs: 0, + kernelBuffers: raw!.buffers, + kernelCounts: raw!.counts, + }; + applyFrameworkExtracts(transport, filePath, source, 'python', ['django']); + const fwNodes = transport.nodes.length; + const fwRefs = transport.unresolvedReferences.length; + + const decoded = materializeKernelResult(transport, filePath, 'python'); + expect(decoded.nodes.length).toBe(raw!.counts.nodes + fwNodes); + expect(decoded.unresolvedReferences.length).toBeGreaterThanOrEqual(fwRefs); + // Append order: framework entries come after the kernel's own. + expect(decoded.nodes.slice(raw!.counts.nodes)).toEqual(transport.nodes); + }); }); /** diff --git a/src/extraction/framework-extracts.ts b/src/extraction/framework-extracts.ts new file mode 100644 index 000000000..4b67ca826 --- /dev/null +++ b/src/extraction/framework-extracts.ts @@ -0,0 +1,39 @@ +import { getAllFrameworkResolvers, getApplicableFrameworks } from '../resolution/frameworks'; +import type { ExtractionResult, Language } from '../types'; + +/** + * Append every applicable framework's `extract()` output to `result` — + * nodes then refs, extractor output first, a failed hook degrading to a + * warning. Runs inside extractFromSource for the decoded path, and on the + * store side of the kernel-buffer path, where the merge must land + * post-decode so framework projects keep the flat-buffer fast path. + */ +export function applyFrameworkExtracts( + result: ExtractionResult, + filePath: string, + source: string, + language: Language, + frameworkNames: string[] | undefined +): void { + if (!frameworkNames || frameworkNames.length === 0) return; + const applicable = getApplicableFrameworks( + getAllFrameworkResolvers().filter((r) => frameworkNames.includes(r.name)), + language + ); + for (const fw of applicable) { + if (!fw.extract) continue; + try { + const fwResult = fw.extract(filePath, source); + result.nodes.push(...fwResult.nodes); + result.unresolvedReferences.push(...fwResult.references); + } catch (err) { + result.errors.push({ + message: `Framework extractor '${fw.name}' failed: ${ + err instanceof Error ? err.message : String(err) + }`, + filePath, + severity: 'warning', + }); + } + } +} diff --git a/src/extraction/index.ts b/src/extraction/index.ts index 8c3fa8a76..dc085e754 100644 --- a/src/extraction/index.ts +++ b/src/extraction/index.ts @@ -2434,9 +2434,12 @@ export class ExtractionOrchestrator { const bp = walBackpressure?.(); if (bp) await bp; - // Kernel deferred-decode results carry table sizes in kernelCounts - // (their object arrays are empty — decode happens at the store). - const nodeCount = result.kernelCounts?.nodes ?? result.nodes.length; + // Kernel deferred-decode results carry table sizes in kernelCounts — + // decode happens at the store; result.nodes then holds only the + // framework extract() additions, so count both. + const nodeCount = result.kernelCounts !== undefined + ? result.kernelCounts.nodes + result.nodes.length + : result.nodes.length; const edgeCount = result.kernelCounts?.edges ?? result.edges.length; // A file whose grammar failed to load was never parsed (#2335). @@ -2460,6 +2463,7 @@ export class ExtractionOrchestrator { buffers: result.kernelBuffers, file: this.buildFileRecord(filePath, content, language, stats, nodeCount, result.errors), ...(result.unresolvedReferences.length > 0 ? { extraRefs: result.unresolvedReferences } : {}), + ...(result.nodes.length > 0 ? { extraNodes: result.nodes } : {}), }); } else { storeWriter.send(this.buildFreshStoreBundle(filePath, content, language, stats, result)); diff --git a/src/extraction/kernel/index.ts b/src/extraction/kernel/index.ts index 8424f2676..682bbcf59 100644 --- a/src/extraction/kernel/index.ts +++ b/src/extraction/kernel/index.ts @@ -310,8 +310,14 @@ export function materializeKernelResult( ); if (b.literalSource !== undefined) captureLiterals(b.literalSource, decoded.nodes); decoded.durationMs = result.durationMs; - // References read beside the tables (a CommonJS `require`) ride on the transport. + // References read beside the tables (a CommonJS `require`), then framework + // extract() output, ride on the transport — same append order as + // extractFromSource. So do the framework nodes and any merge warnings — + // which are the transport's error TAIL: the kernel's own errors already + // decoded above, so re-appending them would double-count. + if (result.nodes.length > 0) decoded.nodes.push(...result.nodes); if (result.unresolvedReferences.length > 0) decoded.unresolvedReferences.push(...result.unresolvedReferences); + if (result.errors.length > decoded.errors.length) decoded.errors.push(...result.errors.slice(decoded.errors.length)); return decoded; } diff --git a/src/extraction/parse-worker.ts b/src/extraction/parse-worker.ts index cb7b31787..4e30c3592 100644 --- a/src/extraction/parse-worker.ts +++ b/src/extraction/parse-worker.ts @@ -18,7 +18,7 @@ import { extractFromSource } from './tree-sitter'; import { detectLanguage } from './grammars'; import { tryKernelExtractRaw } from './kernel'; import { commonJsRequireRefs } from './commonjs-requires'; -import { getAllFrameworkResolvers, getApplicableFrameworks } from '../resolution/frameworks'; +import { applyFrameworkExtracts } from './framework-extracts'; import { mayDefineCppMacroMethods } from './languages/cpp-macro-methods'; import type { Language, ExtractionResult } from '../types'; @@ -43,20 +43,11 @@ parentPort!.on('message', async (msg: { type: string; id?: number; filePath?: st // Kernel deferred-decode fast path: ship the file's tables as flat // buffers and decode at the STORE boundary, so the main thread never // materializes per-node objects (nor pays their structured-clone cost — - // buffer clone is a flat memcpy). Only when no applicable framework has - // an extract() hook: those merge extra nodes/refs into the DECODED - // result inside extractFromSource, so such files keep the decoded path. - // So does a C++ file that may define methods with its own macros: they - // are minted after the kernel, in extractFromSource. + // buffer clone is a flat memcpy). Only a C++ file that may define + // methods with its own macros keeps the decoded path — those are + // minted after the kernel, in extractFromSource. let result: ExtractionResult | undefined; - const frameworksNeedDecode = - frameworkNames && frameworkNames.length > 0 - ? getApplicableFrameworks( - getAllFrameworkResolvers().filter((r) => frameworkNames.includes(r.name)), - language - ).some((fw) => !!fw.extract) - : false; - if (!frameworksNeedDecode && !(language === 'cpp' && mayDefineCppMacroMethods(content!))) { + if (!(language === 'cpp' && mayDefineCppMacroMethods(content!))) { const raw = tryKernelExtractRaw(filePath!, content!, language); if (raw) { result = { @@ -69,6 +60,11 @@ parentPort!.on('message', async (msg: { type: string; id?: number; filePath?: st kernelBuffers: raw.buffers, kernelCounts: raw.counts, }; + // A framework extract() hook no longer forces the decoded path: + // the hook needs only (filePath, source), so it runs here and its + // nodes/refs append onto the transport as extras — the store side + // pushes them after decode, same order as extractFromSource. + applyFrameworkExtracts(result, filePath!, content!, language, frameworkNames); } } result ??= extractFromSource(filePath!, content!, language, frameworkNames); diff --git a/src/extraction/store-worker.ts b/src/extraction/store-worker.ts index 15970589d..6afdf03e3 100644 --- a/src/extraction/store-worker.ts +++ b/src/extraction/store-worker.ts @@ -76,6 +76,10 @@ function decodeKernelBundle(bundle: KernelStoreBundle): StoreBundle { if (bundle.buffers.literalSource !== undefined) { captureLiterals(bundle.buffers.literalSource, decoded.nodes); } + // Framework extract() output and CommonJS refs ride the transport; both + // append after the kernel's tables, the same order extractFromSource + // produces (kernel → CommonJS → framework refs; kernel → framework nodes). + if (bundle.extraNodes?.length) decoded.nodes.push(...bundle.extraNodes); if (bundle.extraRefs?.length) decoded.unresolvedReferences.push(...bundle.extraRefs); return finalizeStoreBundle(decoded, bundle.filePath, bundle.language, bundle.file); } diff --git a/src/extraction/store-writer.ts b/src/extraction/store-writer.ts index 5fa049ffe..bfe74569d 100644 --- a/src/extraction/store-writer.ts +++ b/src/extraction/store-writer.ts @@ -53,8 +53,11 @@ export interface KernelStoreBundle { language: Language; buffers: NonNullable; file: FileRecord; - /** References read beside the kernel's tables (a CommonJS `require`). */ + /** References read beside the kernel's tables (a CommonJS `require`), + * then framework extract() references, in transport order. */ extraRefs?: ExtractionResult['unresolvedReferences']; + /** Framework extract() output nodes, appended after the kernel's. */ + extraNodes?: ExtractionResult['nodes']; } /** diff --git a/src/extraction/tree-sitter.ts b/src/extraction/tree-sitter.ts index 62f16d661..9a003bd54 100644 --- a/src/extraction/tree-sitter.ts +++ b/src/extraction/tree-sitter.ts @@ -44,10 +44,7 @@ import { CfmlExtractor } from './cfml-extractor'; import { tryKernelExtract, takeDeferredPreParse, tryKernelBindings, attachBindingNodeIds } from './kernel'; import { captureLiterals } from './literal-capture'; import { commonJsRequireRefs } from './commonjs-requires'; -import { - getAllFrameworkResolvers, - getApplicableFrameworks, -} from '../resolution/frameworks'; +import { applyFrameworkExtracts } from './framework-extracts'; // Re-export for backward compatibility export { generateNodeId } from './tree-sitter-helpers'; @@ -9931,29 +9928,7 @@ export function extractFromSource( } // Framework-specific extraction (routes, middleware, etc.) - if (frameworkNames && frameworkNames.length > 0) { - const allResolvers = getAllFrameworkResolvers(); - const applicable = getApplicableFrameworks( - allResolvers.filter((r) => frameworkNames.includes(r.name)), - detectedLanguage - ); - for (const fw of applicable) { - if (!fw.extract) continue; - try { - const fwResult = fw.extract(filePath, source); - result.nodes.push(...fwResult.nodes); - result.unresolvedReferences.push(...fwResult.references); - } catch (err) { - result.errors.push({ - message: `Framework extractor '${fw.name}' failed: ${ - err instanceof Error ? err.message : String(err) - }`, - filePath, - severity: 'warning', - }); - } - } - } + applyFrameworkExtracts(result, filePath, source, detectedLanguage, frameworkNames); return result; }