From 0868f9fd48284fcf65c5c47d0c61fed28f00d241 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Wed, 22 Apr 2026 15:19:24 +0100 Subject: [PATCH] feat: enrich wrapped functions and support jsx --- .../src/enrichment/file-enricher.test.ts | 140 ++++++- .../agent/src/enrichment/file-enricher.ts | 101 ++++- packages/enricher/src/ast-helpers.ts | 23 ++ packages/enricher/src/call-detector.ts | 211 ++++++++++ .../enricher/src/comment-formatter.test.ts | 119 ++++++ packages/enricher/src/comment-formatter.ts | 79 +++- packages/enricher/src/detector.ts | 23 +- packages/enricher/src/enriched-result.ts | 8 +- packages/enricher/src/enricher.test.ts | 53 +++ packages/enricher/src/enricher.ts | 107 ++++- packages/enricher/src/import-resolver.test.ts | 158 ++++++++ packages/enricher/src/import-resolver.ts | 217 ++++++++++ packages/enricher/src/index.ts | 4 + packages/enricher/src/languages.ts | 56 +++ packages/enricher/src/parse-result.ts | 20 +- packages/enricher/src/parser-manager.ts | 6 +- packages/enricher/src/posthog-api.ts | 12 +- packages/enricher/src/types.ts | 61 ++- .../enricher/src/wrapper-detector.test.ts | 264 ++++++++++++ packages/enricher/src/wrapper-detector.ts | 375 ++++++++++++++++++ .../enricher/src/wrapper-integration.test.ts | 364 +++++++++++++++++ 21 files changed, 2351 insertions(+), 50 deletions(-) create mode 100644 packages/enricher/src/comment-formatter.test.ts create mode 100644 packages/enricher/src/import-resolver.test.ts create mode 100644 packages/enricher/src/import-resolver.ts create mode 100644 packages/enricher/src/wrapper-detector.test.ts create mode 100644 packages/enricher/src/wrapper-detector.ts create mode 100644 packages/enricher/src/wrapper-integration.test.ts diff --git a/packages/agent/src/enrichment/file-enricher.test.ts b/packages/agent/src/enrichment/file-enricher.test.ts index 751b70487..d53d04eea 100644 --- a/packages/agent/src/enrichment/file-enricher.test.ts +++ b/packages/agent/src/enrichment/file-enricher.test.ts @@ -8,11 +8,15 @@ function makeDeps(overrides: { parseRejects?: Error; isSupported?: boolean; getApiKey?: () => string | Promise; + findImportsInSource?: () => Promise; + getWrappersForFile?: () => Promise; }): { deps: FileEnrichmentDeps; parseSpy: ReturnType; enrichFromApiSpy: ReturnType; getApiKeySpy: ReturnType; + findImportsSpy: ReturnType; + getWrappersSpy: ReturnType; } { const enrichFromApiSpy = vi.fn(async () => ({ toInlineComments: () => @@ -29,11 +33,19 @@ function makeDeps(overrides: { }); const getApiKeySpy = vi.fn(overrides.getApiKey ?? (() => "phx_test")); + const findImportsSpy = vi.fn( + overrides.findImportsInSource ?? (async () => []), + ); + const getWrappersSpy = vi.fn( + overrides.getWrappersForFile ?? (async () => []), + ); const deps: FileEnrichmentDeps = { enricher: { isSupported: vi.fn(() => overrides.isSupported ?? true), parse: parseSpy, + findImportsInSource: findImportsSpy, + getWrappersForFile: getWrappersSpy, } as unknown as FileEnrichmentDeps["enricher"], apiConfig: { apiUrl: "https://test.posthog.com", @@ -42,7 +54,14 @@ function makeDeps(overrides: { }, }; - return { deps, parseSpy, enrichFromApiSpy, getApiKeySpy }; + return { + deps, + parseSpy, + enrichFromApiSpy, + getApiKeySpy, + findImportsSpy, + getWrappersSpy, + }; } describe("enrichFileForAgent", () => { @@ -97,8 +116,8 @@ describe("enrichFileForAgent", () => { expect(enrichFromApiSpy).not.toHaveBeenCalled(); }); - test("returns null and skips parse when content has no posthog reference", async () => { - const { deps, parseSpy } = makeDeps({}); + test("returns null and skips parse when content has no posthog reference AND no relative imports", async () => { + const { deps, parseSpy, findImportsSpy } = makeDeps({}); const result = await enrichFileForAgent( deps, "/tmp/code.ts", @@ -106,6 +125,121 @@ describe("enrichFileForAgent", () => { ); expect(result).toBeNull(); expect(parseSpy).not.toHaveBeenCalled(); + expect(findImportsSpy).not.toHaveBeenCalled(); + }); + + test("relative import with no resolvable wrapper → skips parse", async () => { + const { deps, parseSpy, findImportsSpy, getWrappersSpy } = makeDeps({ + findImportsInSource: async () => [ + { + localName: "foo", + importedName: "foo", + resolvedAbsPath: "/tmp/foo.ts", + }, + ], + getWrappersForFile: async () => [], + }); + const result = await enrichFileForAgent( + deps, + "/tmp/app.ts", + 'import { foo } from "./foo";\nfoo("x");', + ); + expect(result).toBeNull(); + expect(findImportsSpy).toHaveBeenCalled(); + expect(getWrappersSpy).toHaveBeenCalledWith("/tmp/foo.ts"); + expect(parseSpy).not.toHaveBeenCalled(); + }); + + test("relative import hitting a named wrapper triggers parse with context", async () => { + const wrapper = { + name: "track", + methodKind: "capture", + posthogMethod: "capture", + classification: { kind: "pass-through", paramIndex: 0 }, + isNamedExport: true, + isDefaultExport: false, + }; + const { deps, parseSpy, findImportsSpy, getWrappersSpy } = makeDeps({ + findImportsInSource: async () => [ + { + localName: "track", + importedName: "track", + resolvedAbsPath: "/tmp/telemetry.ts", + }, + ], + getWrappersForFile: async () => [wrapper], + }); + const result = await enrichFileForAgent( + deps, + "/tmp/app.ts", + 'import { track } from "./telemetry";\ntrack("x");', + ); + expect(result).toBe("enriched content"); + expect(findImportsSpy).toHaveBeenCalled(); + expect(getWrappersSpy).toHaveBeenCalledWith("/tmp/telemetry.ts"); + expect(parseSpy).toHaveBeenCalledWith( + expect.any(String), + expect.any(String), + expect.objectContaining({ + wrappersByLocalName: expect.any(Map), + }), + ); + const ctxArg = parseSpy.mock.calls[0][2] as { + wrappersByLocalName: Map; + }; + expect(ctxArg.wrappersByLocalName.get("track")).toEqual(wrapper); + }); + + test("file with posthog literal and no relative imports skips import resolution", async () => { + const { deps, findImportsSpy, parseSpy } = makeDeps({}); + await enrichFileForAgent(deps, "/tmp/code.ts", "posthog.capture('x');"); + expect(findImportsSpy).not.toHaveBeenCalled(); + expect(parseSpy).toHaveBeenCalled(); + }); + + test("file with only bare-package imports does not trigger import resolution", async () => { + const { deps, findImportsSpy } = makeDeps({}); + const content = [ + 'import React from "react";', + 'import { useState } from "react";', + 'import posthog from "posthog-js";', + "posthog.capture('x');", + ].join("\n"); + await enrichFileForAgent(deps, "/tmp/page.tsx", content); + expect(findImportsSpy).not.toHaveBeenCalled(); + }); + + test("file with direct posthog AND wrapper imports gets both enriched", async () => { + const wrapper = { + name: "track", + methodKind: "capture", + posthogMethod: "capture", + classification: { kind: "pass-through", paramIndex: 0 }, + isNamedExport: true, + isDefaultExport: false, + }; + const { deps, findImportsSpy, parseSpy } = makeDeps({ + findImportsInSource: async () => [ + { + localName: "track", + importedName: "track", + resolvedAbsPath: "/tmp/utils.ts", + }, + ], + getWrappersForFile: async () => [wrapper], + }); + const content = [ + 'import posthog from "posthog-js";', + 'import { track } from "./utils";', + "posthog.capture('direct');", + 'track("wrapper_call");', + ].join("\n"); + await enrichFileForAgent(deps, "/tmp/page.tsx", content); + expect(findImportsSpy).toHaveBeenCalled(); + const ctx = parseSpy.mock.calls[0][2] as { + wrappersByLocalName: Map; + }; + expect(ctx.wrappersByLocalName.get("track")).toEqual(wrapper); }); test("returns null when getApiKey yields empty string", async () => { diff --git a/packages/agent/src/enrichment/file-enricher.ts b/packages/agent/src/enrichment/file-enricher.ts index 5b9eae26d..4bce2658e 100644 --- a/packages/agent/src/enrichment/file-enricher.ts +++ b/packages/agent/src/enrichment/file-enricher.ts @@ -1,5 +1,11 @@ import * as path from "node:path"; -import { EXT_TO_LANG_ID, PostHogEnricher } from "@posthog/enricher"; +import { + EXT_TO_LANG_ID, + type ImportEdge, + type LocalWrapper, + type ParseContext, + PostHogEnricher, +} from "@posthog/enricher"; import type { PostHogAPIConfig } from "../types"; import type { Logger } from "../utils/logger"; @@ -27,6 +33,10 @@ export function createEnrichment( } const MAX_ENRICHMENT_BYTES = 1_000_000; +const MAX_RELATIVE_IMPORTS = 64; +const RELATIVE_IMPORT_REGEX = + /(?:^|\n)\s*(?:import\b[^\n]*['"]\.{1,2}\/|from\s+\.)/; +const POSTHOG_LITERAL_REGEX = /posthog/i; export async function enrichFileForAgent( deps: FileEnrichmentDeps, @@ -35,15 +45,29 @@ export async function enrichFileForAgent( ): Promise { if (!content || content.length > MAX_ENRICHMENT_BYTES) return null; - // Skip the tree-sitter parse for files with no PostHog references. - if (!/posthog/i.test(content)) return null; - const ext = path.extname(filePath).toLowerCase(); const langId = EXT_TO_LANG_ID[ext]; if (!langId || !deps.enricher.isSupported(langId)) return null; + const hasPostHogLiteral = POSTHOG_LITERAL_REGEX.test(content); + const hasRelativeImport = RELATIVE_IMPORT_REGEX.test(content); + let parseContext: ParseContext | undefined; + + // Build wrapper context whenever the file has relative imports — direct PostHog + // usage and wrapper usage can coexist in the same file, so we don't skip this + // just because `posthog` already appears literally. + if (hasRelativeImport) { + const absPath = path.resolve(filePath); + const ctx = await buildWrapperContext(deps, content, langId, absPath); + if (ctx) parseContext = ctx; + } + + // Bail only when nothing at all could be enriched: no direct posthog literal + // AND no resolvable wrappers. + if (!hasPostHogLiteral && !parseContext) return null; + try { - const parsed = await deps.enricher.parse(content, langId); + const parsed = await deps.enricher.parse(content, langId, parseContext); if (parsed.calls.length === 0 && parsed.initCalls.length === 0) { return null; } @@ -69,6 +93,7 @@ export async function enrichFileForAgent( deps.logger?.debug("File enriched", { filePath, calls: parsed.calls.length, + viaWrappers: parsed.calls.filter((c) => c.viaWrapper).length, }); return annotated; } catch (err) { @@ -80,3 +105,69 @@ export async function enrichFileForAgent( return null; } } + +async function buildWrapperContext( + deps: FileEnrichmentDeps, + content: string, + langId: string, + absPath: string, +): Promise { + let edges: ImportEdge[]; + try { + edges = await deps.enricher.findImportsInSource(content, langId, absPath); + } catch (err) { + deps.logger?.debug("Import resolution failed", { + absPath, + err: err instanceof Error ? err.message : String(err), + }); + return null; + } + + if (!edges.length) return null; + const bounded = edges.slice(0, MAX_RELATIVE_IMPORTS); + + const wrappersByLocalName = new Map(); + const namespaceWrappers = new Map>(); + + const resolutions = await Promise.all( + bounded.map(async (edge) => { + if (!edge.resolvedAbsPath) return null; + const wrappers = await deps.enricher.getWrappersForFile( + edge.resolvedAbsPath, + ); + if (!wrappers.length) return null; + return { edge, wrappers }; + }), + ); + + for (const entry of resolutions) { + if (!entry) continue; + const { edge, wrappers } = entry; + + if (edge.isNamespace) { + const nsMap = new Map(); + for (const w of wrappers) { + if (w.isNamedExport || w.isDefaultExport) { + nsMap.set(w.name, w); + } + } + if (nsMap.size) namespaceWrappers.set(edge.localName, nsMap); + continue; + } + + if (edge.isDefault) { + const target = wrappers.find((w) => w.isDefaultExport); + if (target) wrappersByLocalName.set(edge.localName, target); + continue; + } + + const target = wrappers.find( + (w) => w.name === edge.importedName && w.isNamedExport, + ); + if (target) wrappersByLocalName.set(edge.localName, target); + } + + if (!wrappersByLocalName.size && !namespaceWrappers.size) return null; + + return { wrappersByLocalName, namespaceWrappers }; +} diff --git a/packages/enricher/src/ast-helpers.ts b/packages/enricher/src/ast-helpers.ts index a82b45627..e2b26ee2a 100644 --- a/packages/enricher/src/ast-helpers.ts +++ b/packages/enricher/src/ast-helpers.ts @@ -158,3 +158,26 @@ export function walkNodes( }; visit(root); } + +const JSX_NODE_TYPES = new Set([ + "jsx_element", + "jsx_fragment", + "jsx_self_closing_element", + "jsx_opening_element", + "jsx_closing_element", + "jsx_attribute", +]); + +/** + * Returns true when `node` lives anywhere inside a JSX element — i.e. appending + * a trailing `// …` comment to the call's line would land inside JSX content + * rather than in a JavaScript statement context. + */ +export function isInsideJsx(node: Parser.SyntaxNode): boolean { + let cur: Parser.SyntaxNode | null = node.parent; + while (cur) { + if (JSX_NODE_TYPES.has(cur.type)) return true; + cur = cur.parent; + } + return false; +} diff --git a/packages/enricher/src/call-detector.ts b/packages/enricher/src/call-detector.ts index 9e783910a..0da96119d 100644 --- a/packages/enricher/src/call-detector.ts +++ b/packages/enricher/src/call-detector.ts @@ -9,11 +9,14 @@ import { extractClientName, extractParams, getCapture, + isInsideJsx, } from "./ast-helpers.js"; import type { ParserManager } from "./parser-manager.js"; import type { FlagAssignment, FunctionInfo, + LocalWrapper, + ParseContext, PostHogCall, PostHogInitCall, } from "./types.js"; @@ -25,6 +28,7 @@ export async function findPostHogCalls( pm: ParserManager, source: string, languageId: string, + context?: ParseContext, ): Promise { const ready = await pm.ensureReady(languageId); if (!ready) { @@ -60,6 +64,7 @@ export async function findPostHogCalls( const clientNode = getCapture(match.captures, "client"); const methodNode = getCapture(match.captures, "method"); const keyNode = getCapture(match.captures, "key"); + const callNode = getCapture(match.captures, "call"); if (!clientNode || !methodNode || !keyNode) { continue; @@ -102,6 +107,7 @@ export async function findPostHogCalls( line: keyNode.startPosition.row, keyStartCol: keyNode.startPosition.column, keyEndCol: keyNode.endPosition.column, + inJsx: callNode ? isInsideJsx(callNode) : undefined, }); } } @@ -167,6 +173,7 @@ export async function findPostHogCalls( const methodNode = getCapture(match.captures, "method"); const propNameNode = getCapture(match.captures, "prop_name"); const keyNode = getCapture(match.captures, "key"); + const callNode = getCapture(match.captures, "call"); if (!clientNode || !methodNode || !propNameNode || !keyNode) { continue; @@ -194,6 +201,7 @@ export async function findPostHogCalls( line: keyNode.startPosition.row, keyStartCol: keyNode.startPosition.column, keyEndCol: keyNode.endPosition.column, + inJsx: callNode ? isInsideJsx(callNode) : undefined, }); } } @@ -310,6 +318,7 @@ export async function findPostHogCalls( for (const match of matches) { const funcNode = getCapture(match.captures, "func_name"); const keyNode = getCapture(match.captures, "key"); + const callNode = getCapture(match.captures, "call"); if (!funcNode || !keyNode) { continue; } @@ -322,6 +331,7 @@ export async function findPostHogCalls( line: keyNode.startPosition.row, keyStartCol: keyNode.startPosition.column, keyEndCol: keyNode.endPosition.column, + inJsx: callNode ? isInsideJsx(callNode) : undefined, }); } } @@ -340,6 +350,7 @@ export async function findPostHogCalls( for (const match of matches) { const funcNode = getCapture(match.captures, "func_name"); const keyNode = getCapture(match.captures, "key"); + const callNode = getCapture(match.captures, "call"); if (!funcNode || !keyNode) { continue; } @@ -351,6 +362,7 @@ export async function findPostHogCalls( line: keyNode.startPosition.row, keyStartCol: keyNode.startPosition.column, keyEndCol: keyNode.endPosition.column, + inJsx: callNode ? isInsideJsx(callNode) : undefined, }); } } @@ -367,6 +379,7 @@ export async function findPostHogCalls( const clientNode = getCapture(match.captures, "client"); const methodNode = getCapture(match.captures, "method"); const argNode = getCapture(match.captures, "arg_id"); + const callNode = getCapture(match.captures, "call"); if (!clientNode || !methodNode || !argNode) { continue; } @@ -401,6 +414,7 @@ export async function findPostHogCalls( line, keyStartCol: argNode.startPosition.column, keyEndCol: argNode.endPosition.column, + inJsx: callNode ? isInsideJsx(callNode) : undefined, }); } } @@ -415,6 +429,7 @@ export async function findPostHogCalls( const clientNode = getCapture(match.captures, "client"); const methodNode = getCapture(match.captures, "method"); const firstArgNode = getCapture(match.captures, "first_arg"); + const callNode = getCapture(match.captures, "call"); if (!clientNode || !methodNode || !firstArgNode) { continue; } @@ -443,14 +458,208 @@ export async function findPostHogCalls( keyStartCol: firstArgNode.startPosition.column, keyEndCol: firstArgNode.endPosition.column, dynamic: true, + inJsx: callNode ? isInsideJsx(callNode) : undefined, }); matchedLines.add(line); } } + if (context?.wrappersByLocalName?.size) { + synthesizeBareWrapperCalls( + pm, + lang, + tree, + languageId, + context.wrappersByLocalName, + constantMap, + calls, + matchedLines, + ); + } + + if (context?.namespaceWrappers?.size) { + synthesizeNamespaceWrapperCalls( + pm, + lang, + tree, + languageId, + context.namespaceWrappers, + constantMap, + calls, + matchedLines, + ); + } + return calls; } +const WRAPPER_BARE_CALL_QUERIES: Record = { + javascript: `(call_expression function: (identifier) @func_name arguments: (arguments) @args) @call`, + javascriptreact: `(call_expression function: (identifier) @func_name arguments: (arguments) @args) @call`, + typescript: `(call_expression function: (identifier) @func_name arguments: (arguments) @args) @call`, + typescriptreact: `(call_expression function: (identifier) @func_name arguments: (arguments) @args) @call`, + python: `(call function: (identifier) @func_name arguments: (argument_list) @args) @call`, +}; + +const WRAPPER_NAMESPACE_CALL_QUERIES: Record = { + javascript: `(call_expression function: (member_expression object: (identifier) @ns property: (property_identifier) @method) arguments: (arguments) @args) @call`, + javascriptreact: `(call_expression function: (member_expression object: (identifier) @ns property: (property_identifier) @method) arguments: (arguments) @args) @call`, + typescript: `(call_expression function: (member_expression object: (identifier) @ns property: (property_identifier) @method) arguments: (arguments) @args) @call`, + typescriptreact: `(call_expression function: (member_expression object: (identifier) @ns property: (property_identifier) @method) arguments: (arguments) @args) @call`, + python: `(call function: (attribute object: (identifier) @ns attribute: (identifier) @method) arguments: (argument_list) @args) @call`, +}; + +function synthesizeBareWrapperCalls( + pm: ParserManager, + lang: Parser.Language, + tree: Parser.Tree, + languageId: string, + wrappers: Map, + constantMap: Map, + calls: PostHogCall[], + matchedLines: Set, +): void { + const queryStr = WRAPPER_BARE_CALL_QUERIES[languageId]; + if (!queryStr) return; + const query = pm.getQuery(lang, queryStr); + if (!query) return; + + for (const match of query.matches(tree.rootNode)) { + const funcNode = getCapture(match.captures, "func_name"); + const argsNode = getCapture(match.captures, "args"); + const callNode = getCapture(match.captures, "call"); + if (!funcNode || !argsNode) continue; + const wrapper = wrappers.get(funcNode.text); + if (!wrapper) continue; + pushWrapperCall( + wrapper, + funcNode, + argsNode, + callNode, + constantMap, + calls, + matchedLines, + ); + } +} + +function synthesizeNamespaceWrapperCalls( + pm: ParserManager, + lang: Parser.Language, + tree: Parser.Tree, + languageId: string, + namespaceWrappers: Map>, + constantMap: Map, + calls: PostHogCall[], + matchedLines: Set, +): void { + const queryStr = WRAPPER_NAMESPACE_CALL_QUERIES[languageId]; + if (!queryStr) return; + const query = pm.getQuery(lang, queryStr); + if (!query) return; + + for (const match of query.matches(tree.rootNode)) { + const nsNode = getCapture(match.captures, "ns"); + const methodNode = getCapture(match.captures, "method"); + const argsNode = getCapture(match.captures, "args"); + const callNode = getCapture(match.captures, "call"); + if (!nsNode || !methodNode || !argsNode) continue; + const nsWrappers = namespaceWrappers.get(nsNode.text); + if (!nsWrappers) continue; + const wrapper = nsWrappers.get(methodNode.text); + if (!wrapper) continue; + pushWrapperCall( + wrapper, + methodNode, + argsNode, + callNode, + constantMap, + calls, + matchedLines, + ); + } +} + +function pushWrapperCall( + wrapper: LocalWrapper, + callerNode: Parser.SyntaxNode, + argsNode: Parser.SyntaxNode, + callNode: Parser.SyntaxNode | null, + constantMap: Map, + calls: PostHogCall[], + matchedLines: Set, +): void { + let line = callerNode.startPosition.row; + let keyStartCol = callerNode.startPosition.column; + let keyEndCol = callerNode.endPosition.column; + let key = ""; + let dynamic = false; + + if (wrapper.classification.kind === "fixed-key") { + key = wrapper.classification.key; + } else { + const positional = argsNode.namedChildren.filter( + (c) => c.type !== "comment" && c.type !== "keyword_argument", + ); + const arg = positional[wrapper.classification.paramIndex]; + if (!arg) return; + line = arg.startPosition.row; + keyStartCol = arg.startPosition.column; + keyEndCol = arg.endPosition.column; + + if (arg.type === "string") { + const fragment = arg.namedChildren.find( + (c) => c.type === "string_fragment" || c.type === "string_content", + ); + if (fragment) { + key = fragment.text; + } else { + dynamic = true; + } + } else if (arg.type === "template_string") { + const fragments = arg.namedChildren.filter( + (c) => c.type === "string_fragment", + ); + const hasInterp = arg.namedChildren.some( + (c) => c.type === "template_substitution", + ); + if (!hasInterp && fragments.length === 1) { + key = fragments[0].text; + } else { + dynamic = true; + } + } else if (arg.type === "interpreted_string_literal") { + key = arg.text.slice(1, -1); + } else if (arg.type === "identifier") { + const resolved = constantMap.get(arg.text); + if (resolved) { + key = resolved; + } else { + dynamic = true; + } + } else { + dynamic = true; + } + } + + if (matchedLines.has(line) && dynamic) { + // Direct PostHog call already annotated this line — don't overwrite with opaque wrapper data. + return; + } + + calls.push({ + method: wrapper.posthogMethod, + key, + line, + keyStartCol, + keyEndCol, + dynamic: dynamic ? true : undefined, + viaWrapper: wrapper.name, + inJsx: callNode ? isInsideJsx(callNode) : undefined, + }); + matchedLines.add(line); +} + export async function findInitCalls( pm: ParserManager, source: string, @@ -961,6 +1170,7 @@ export async function findFunctions( : []; const bodyLine = bodyNode.startPosition.row; + const bodyEndLine = bodyNode.endPosition.row; const nextLineIdx = bodyLine + 1; const lines = text.split("\n"); const nextLine = nextLineIdx < lines.length ? lines[nextLineIdx] : ""; @@ -971,6 +1181,7 @@ export async function findFunctions( params, isComponent: /^[A-Z]/.test(name), bodyLine, + bodyEndLine, bodyIndent, }); } diff --git a/packages/enricher/src/comment-formatter.test.ts b/packages/enricher/src/comment-formatter.test.ts new file mode 100644 index 000000000..11c1b3295 --- /dev/null +++ b/packages/enricher/src/comment-formatter.test.ts @@ -0,0 +1,119 @@ +import { describe, expect, test } from "vitest"; +import { formatInlineComments } from "./comment-formatter.js"; +import type { EnrichedEvent, EnrichedFlag, EnrichedListItem } from "./types.js"; + +function eventItem( + name: string, + line: number, + inJsx: boolean, +): EnrichedListItem { + return { + type: "event", + line, + name, + method: "capture", + inJsx, + verified: true, + }; +} + +function enrichedEvent(name: string): EnrichedEvent { + return { + eventName: name, + verified: true, + } as EnrichedEvent; +} + +describe("formatInlineComments", () => { + test("pure JS line uses // suffix", () => { + const source = `posthog.capture('a');`; + const items = [eventItem("a", 0, false)]; + const events = new Map([["a", enrichedEvent("a")]]); + const out = formatInlineComments( + source, + "javascript", + items, + new Map(), + events, + ); + expect(out).toBe( + `posthog.capture('a'); // [PostHog] Event: "a" \u2014 (verified)`, + ); + }); + + test("pure JSX line uses {/* */} suffix", () => { + const source = ` + ); +} +`; + + // Sanity-check the building blocks first. + const edges = await enricher.findImportsInSource( + appSrc, + "typescriptreact", + app, + ); + const trackEdge = edges.find((e) => e.localName === "track"); + expect( + trackEdge, + "imports to include { track } from ./utils", + ).toBeDefined(); + expect(trackEdge?.resolvedAbsPath).toBe(utils); + + const wrappers = await enricher.getWrappersForFile(utils); + expect(wrappers.map((w) => w.name)).toContain("track"); + + const ctx = await buildContext(enricher, appSrc, "typescriptreact", app); + expect(ctx.wrappersByLocalName.has("track")).toBe(true); + + const parsed = await enricher.parse(appSrc, "typescriptreact", ctx); + const events = parsed.events.map((e) => ({ + name: e.name, + viaWrapper: e.viaWrapper, + })); + expect(events).toEqual( + expect.arrayContaining([ + { name: "purchase_completed", viaWrapper: undefined }, + { name: "event_123", viaWrapper: "track" }, + ]), + ); + }); + + test("wrapper call inside JSX sets inJsx and renders as JSX comment", async () => { + const utils = path.join(workDir, "jsx-utils.ts"); + writeFileSync( + utils, + `import posthog from "posthog-js";\nexport const track = (event_name: string) => {\n posthog.capture(event_name);\n};\n`, + ); + + const app = path.join(workDir, "jsx-page.tsx"); + const appSrc = `"use client"; +import posthog from "posthog-js"; +import { track } from "./jsx-utils"; + +export default function Home() { + posthog.capture("direct_event"); + return ( +
+ +
+ ); +} +`; + + const ctx = await buildContext(enricher, appSrc, "typescriptreact", app); + const parsed = await enricher.parse(appSrc, "typescriptreact", ctx); + + const directCall = parsed.calls.find((c) => c.key === "direct_event"); + const wrapperCall = parsed.calls.find((c) => c.key === "event_123"); + expect(directCall?.inJsx).toBeFalsy(); + expect(wrapperCall?.inJsx).toBe(true); + + const annotated = ( + await parsed.enrichFromApi({ + apiKey: "k", + host: "https://example.com", + projectId: 1, + timeoutMs: 1, + }) + ).toInlineComments(); + + // The wrapper-call line lives inside JSX, so its annotation must use a JSX-safe comment. + const wrapperLine = annotated + .split("\n") + .find((l) => l.includes('track("event_123")')); + expect(wrapperLine).toContain("{/* [PostHog]"); + expect(wrapperLine).toContain("*/}"); + + // The direct call is in a JS statement context, so the existing `//` form is correct. + const directLine = annotated + .split("\n") + .find((l) => l.includes('posthog.capture("direct_event")')); + expect(directLine).toContain("// [PostHog]"); + }); + + test("non-wrapper file is cached as empty", async () => { + const target = path.join(workDir, "not-a-wrapper.ts"); + writeFileSync(target, `export function noop() {\n return 42;\n}\n`); + const wrappers = await enricher.getWrappersForFile(target); + expect(wrappers).toEqual([]); + }); +});