feat: enrich wrapped functions and support jsx

packages/agent/src/enrichment/file-enricher.test.ts

@@ -8,11 +8,15 @@ function makeDeps(overrides: {
   parseRejects?: Error;
   isSupported?: boolean;
   getApiKey?: () => string | Promise<string>;
+  findImportsInSource?: () => Promise<unknown[]>;
+  getWrappersForFile?: () => Promise<unknown[]>;
 }): {
   deps: FileEnrichmentDeps;
   parseSpy: ReturnType<typeof vi.fn>;
   enrichFromApiSpy: ReturnType<typeof vi.fn>;
   getApiKeySpy: ReturnType<typeof vi.fn>;
+  findImportsSpy: ReturnType<typeof vi.fn>;
+  getWrappersSpy: ReturnType<typeof vi.fn>;
 } {
   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,15 +116,130 @@ 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",
       "const x = 1;\nfunction foo() {}",
     );
     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<string, unknown>;
+    };
+    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<string, unknown>;
+    };
+    expect(ctx.wrappersByLocalName.get("track")).toEqual(wrapper);
   });
 
   test("returns null when getApiKey yields empty string", async () => {

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<string | null> {
   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<ParseContext | null> {
+  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<string, LocalWrapper>();
+  const namespaceWrappers = new Map<string, Map<string, LocalWrapper>>();
+
+  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<string, LocalWrapper>();
+      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 };
+}

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;
+}