From d02ea61b9eb995bb45f344351438146d5b48955a Mon Sep 17 00:00:00 2001 From: Jon McCallum Date: Mon, 22 Jun 2026 12:09:26 +0100 Subject: [PATCH] feat(apm): agent file-read enrichment --- packages/agent/src/adapters/acp-connection.ts | 1 + .../claude/claude-agent.logger.test.ts | 28 +++ .../agent/src/adapters/claude/claude-agent.ts | 6 +- .../src/enrichment/file-enricher.test.ts | 139 ++++++++++++++- .../agent/src/enrichment/file-enricher.ts | 165 +++++++++++++++--- 5 files changed, 317 insertions(+), 22 deletions(-) create mode 100644 packages/agent/src/adapters/claude/claude-agent.logger.test.ts diff --git a/packages/agent/src/adapters/acp-connection.ts b/packages/agent/src/adapters/acp-connection.ts index c271e03b6c..d92528a078 100644 --- a/packages/agent/src/adapters/acp-connection.ts +++ b/packages/agent/src/adapters/acp-connection.ts @@ -114,6 +114,7 @@ function createClaudeConnection(config: AcpConnectionConfig): AcpConnection { ...config.processCallbacks, onStructuredOutput: config.onStructuredOutput, posthogApiConfig: resolveEnricherApiConfig(config), + logger: config.logger?.child("ClaudeAcpAgent"), }); return agent; }, agentStream); diff --git a/packages/agent/src/adapters/claude/claude-agent.logger.test.ts b/packages/agent/src/adapters/claude/claude-agent.logger.test.ts new file mode 100644 index 0000000000..1d34130d7c --- /dev/null +++ b/packages/agent/src/adapters/claude/claude-agent.logger.test.ts @@ -0,0 +1,28 @@ +import type { AgentSideConnection } from "@agentclientprotocol/sdk"; +import { describe, expect, it, vi } from "vitest"; +import { Logger } from "../../utils/logger"; +import { ClaudeAcpAgent } from "./claude-agent"; + +describe("ClaudeAcpAgent logging", () => { + it("emits through the host logger from options so enrichment logs reach the host sink", () => { + const onLog = vi.fn(); + const hostLogger = new Logger({ + debug: true, + prefix: "[PostHog Agent]", + onLog, + }); + const client = {} as unknown as AgentSideConnection; + + const agent = new ClaudeAcpAgent(client, { + logger: hostLogger.child("ClaudeAcpAgent"), + }); + agent.logger.info("[apm] agent enrich", { lines: 4 }); + + expect(onLog).toHaveBeenCalledWith( + "info", + expect.stringContaining("ClaudeAcpAgent"), + "[apm] agent enrich", + { lines: 4 }, + ); + }); +}); diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 7c6d8220f4..2404622f94 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -236,6 +236,8 @@ export interface ClaudeAcpAgentOptions { onMcpServersReady?: (serverNames: string[]) => void; onStructuredOutput?: (output: Record) => Promise; posthogApiConfig?: PostHogAPIConfig; + // Host log sink; without it the agent's enrichment logs never reach main.log. + logger?: Logger; } export class ClaudeAcpAgent extends BaseAcpAgent { @@ -255,7 +257,9 @@ export class ClaudeAcpAgent extends BaseAcpAgent { this.options = options; this.toolUseCache = {}; this.toolUseStreamCache = new Map(); - this.logger = new Logger({ debug: true, prefix: "[ClaudeAcpAgent]" }); + this.logger = + options?.logger ?? + new Logger({ debug: true, prefix: "[ClaudeAcpAgent]" }); this.enrichment = createEnrichment(options?.posthogApiConfig, this.logger); } diff --git a/packages/agent/src/enrichment/file-enricher.test.ts b/packages/agent/src/enrichment/file-enricher.test.ts index d53d04eea0..7555cc88dd 100644 --- a/packages/agent/src/enrichment/file-enricher.test.ts +++ b/packages/agent/src/enrichment/file-enricher.test.ts @@ -1,3 +1,4 @@ +import type { SpanLineStat } from "@posthog/shared"; import { describe, expect, test, vi } from "vitest"; import { enrichFileForAgent, type FileEnrichmentDeps } from "./file-enricher"; @@ -10,6 +11,7 @@ function makeDeps(overrides: { getApiKey?: () => string | Promise; findImportsInSource?: () => Promise; getWrappersForFile?: () => Promise; + fetchApmLineStats?: () => Promise; }): { deps: FileEnrichmentDeps; parseSpy: ReturnType; @@ -17,6 +19,7 @@ function makeDeps(overrides: { getApiKeySpy: ReturnType; findImportsSpy: ReturnType; getWrappersSpy: ReturnType; + fetchApmSpy: ReturnType; } { const enrichFromApiSpy = vi.fn(async () => ({ toInlineComments: () => @@ -39,6 +42,7 @@ function makeDeps(overrides: { const getWrappersSpy = vi.fn( overrides.getWrappersForFile ?? (async () => []), ); + const fetchApmSpy = vi.fn(overrides.fetchApmLineStats ?? (async () => [])); const deps: FileEnrichmentDeps = { enricher: { @@ -52,6 +56,8 @@ function makeDeps(overrides: { projectId: 1, getApiKey: getApiKeySpy, }, + apmStatsCache: new Map(), + fetchApmLineStats: fetchApmSpy, }; return { @@ -61,6 +67,7 @@ function makeDeps(overrides: { getApiKeySpy, findImportsSpy, getWrappersSpy, + fetchApmSpy, }; } @@ -261,7 +268,12 @@ describe("enrichFileForAgent", () => { }); test("returns null and logs debug when enricher throws", async () => { - const logger = { debug: vi.fn() }; + const logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; const { deps } = makeDeps({ parseRejects: new Error("boom") }); deps.logger = logger as unknown as FileEnrichmentDeps["logger"]; const result = await enrichFileForAgent( @@ -295,3 +307,128 @@ describe("enrichFileForAgent", () => { ); }); }); + +describe("APM enrichment caching", () => { + const ONE_HOT_LINE: SpanLineStat[] = [ + { line: 1, count: 5, errorCount: 0, p50Ms: 2, p95Ms: 7 }, + ]; + + test.each<{ + name: string; + fetchApmLineStats: () => Promise; + calls: number; + contains: string | null; + }>([ + { + name: "serves hot stats from cache on re-read (one fetch)", + fetchApmLineStats: async () => ONE_HOT_LINE, + calls: 1, + contains: "[PostHog] APM", + }, + { + name: "caches an empty result so an untraced file is fetched once", + fetchApmLineStats: async () => [], + calls: 1, + contains: null, + }, + { + name: "does not cache a transient fetch failure (retries next read)", + fetchApmLineStats: async () => { + throw new Error("network down"); + }, + calls: 2, + contains: null, + }, + ])("$name", async ({ fetchApmLineStats, calls, contains }) => { + const { deps, fetchApmSpy } = makeDeps({ fetchApmLineStats }); + + const first = await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + const second = await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + + for (const result of [first, second]) { + if (contains === null) expect(result).toBeNull(); + else expect(result).toContain(contains); + } + expect(fetchApmSpy).toHaveBeenCalledTimes(calls); + }); + + test("re-applies cached stats to the current content after an edit", async () => { + const { deps } = makeDeps({ fetchApmLineStats: async () => ONE_HOT_LINE }); + + await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + const edited = await enrichFileForAgent( + deps, + "/repo/flag.rs", + "fn renamed() {}", + ); + + const firstLine = edited?.split("\n")[0] ?? ""; + expect(firstLine).toContain("fn renamed() {}"); + expect(firstLine).toContain("[PostHog] APM"); + }); + + test("a rejecting getApiKey degrades to null instead of throwing", async () => { + const { deps, fetchApmSpy } = makeDeps({ + getApiKey: () => Promise.reject(new Error("token refresh failed")), + }); + + const result = await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + + expect(result).toBeNull(); + expect(fetchApmSpy).not.toHaveBeenCalled(); + }); + + test("queries APM with a timeout that covers the 24h window's worst case", async () => { + const { deps, fetchApmSpy } = makeDeps({ + fetchApmLineStats: async () => ONE_HOT_LINE, + }); + + await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + + const config = fetchApmSpy.mock.calls[0]?.[0] as { timeoutMs?: number }; + expect(config?.timeoutMs).toBeGreaterThanOrEqual(15_000); + }); + + test("serves from cache within the TTL, refetches once it expires", async () => { + const { deps, fetchApmSpy } = makeDeps({ + fetchApmLineStats: async () => ONE_HOT_LINE, + }); + const now = vi.spyOn(Date, "now").mockReturnValue(1_000); + try { + await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + expect(fetchApmSpy).toHaveBeenCalledTimes(1); + now.mockReturnValue(1_000 + 4 * 60_000); + await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + expect(fetchApmSpy).toHaveBeenCalledTimes(1); + now.mockReturnValue(1_000 + 5 * 60_000 + 1); + await enrichFileForAgent(deps, "/repo/flag.rs", "fn m() {}"); + expect(fetchApmSpy).toHaveBeenCalledTimes(2); + } finally { + now.mockRestore(); + } + }); + + test("an event-branch failure does not break the read or lose APM", async () => { + const { deps } = makeDeps({ + fetchApmLineStats: async () => ONE_HOT_LINE, + findImportsInSource: async () => [ + { + localName: "track", + importedName: "track", + resolvedAbsPath: "/repo/t.ts", + }, + ], + getWrappersForFile: async () => { + throw new Error("fs boom"); + }, + }); + + const out = await enrichFileForAgent( + deps, + "/repo/app.ts", + 'import { track } from "./t";\ntrack("x");', + ); + + expect(out).toContain("[PostHog] APM"); + }); +}); diff --git a/packages/agent/src/enrichment/file-enricher.ts b/packages/agent/src/enrichment/file-enricher.ts index 4bce2658e0..463e9ecba4 100644 --- a/packages/agent/src/enrichment/file-enricher.ts +++ b/packages/agent/src/enrichment/file-enricher.ts @@ -1,18 +1,38 @@ import * as path from "node:path"; import { + type EnricherApiConfig, EXT_TO_LANG_ID, + formatApmInlineComments, type ImportEdge, type LocalWrapper, type ParseContext, + PostHogApi, PostHogEnricher, } from "@posthog/enricher"; +import { + APM_STATS_WINDOW, + apmLangForFile, + type SpanLineStat, +} from "@posthog/shared"; import type { PostHogAPIConfig } from "../types"; import type { Logger } from "../utils/logger"; +interface ApmStatsCacheEntry { + expiresAt: number; + stats: SpanLineStat[]; +} + export interface FileEnrichmentDeps { enricher: PostHogEnricher; apiConfig: PostHogAPIConfig; logger?: Logger; + // Path-keyed, best-effort: stats carry production line numbers, so within the + // TTL an edit can shift a comment off its line (we don't re-map). + apmStatsCache: Map; + fetchApmLineStats: ( + config: EnricherApiConfig, + filePath: string, + ) => Promise; } export interface Enrichment { @@ -26,13 +46,29 @@ export function createEnrichment( ): Enrichment | undefined { if (!apiConfig) return undefined; const enricher = new PostHogEnricher(); + const apmStatsCache = new Map(); return { - deps: { enricher, apiConfig, logger }, - dispose: () => enricher.dispose(), + deps: { + enricher, + apiConfig, + logger, + apmStatsCache, + fetchApmLineStats: (config, filePath) => + new PostHogApi(config).getApmLineStats(filePath, { + dateFrom: APM_STATS_WINDOW.dateFrom, + }), + }, + dispose: () => { + enricher.dispose(); + apmStatsCache.clear(); + }, }; } const MAX_ENRICHMENT_BYTES = 1_000_000; +const APM_STATS_TTL_MS = 5 * 60_000; +// 24h query is ~8s on the hottest traced file; a lower budget silently drops it. +const APM_QUERY_TIMEOUT_MS = 15_000; const MAX_RELATIVE_IMPORTS = 64; const RELATIVE_IMPORT_REGEX = /(?:^|\n)\s*(?:import\b[^\n]*['"]\.{1,2}\/|from\s+\.)/; @@ -45,6 +81,100 @@ export async function enrichFileForAgent( ): Promise { if (!content || content.length > MAX_ENRICHMENT_BYTES) return null; + // Resolve the API key once per read; the event and APM paths run in parallel + // and would otherwise race two token refreshes on a dual-enrichable file. + let apiKeyPromise: Promise | undefined; + const getApiKey: ApiKeyGetter = () => { + apiKeyPromise ??= Promise.resolve(deps.apiConfig.getApiKey()); + return apiKeyPromise; + }; + + const apmLang = apmLangForFile(filePath); + const [eventAnnotated, apmStats] = await Promise.all([ + enrichEventsForAgent(deps, filePath, content, getApiKey), + apmLang + ? getApmLineStats(deps, filePath, getApiKey) + : Promise.resolve(null), + ]); + + if (apmLang) { + deps.logger?.debug("[apm] agent enrich", { + filePath, + lines: apmStats?.length ?? null, + }); + } + + let result = eventAnnotated ?? content; + if (apmLang && apmStats && apmStats.length > 0) { + try { + result = formatApmInlineComments(result, apmLang, apmStats, filePath); + } catch (err) { + // A formatter edge case must degrade to the event-annotated content, not + // reject the whole read (result is unchanged on throw). + deps.logger?.debug("APM comment formatting failed", { + filePath, + message: err instanceof Error ? err.message : String(err), + }); + } + } + + return result === content ? null : result; +} + +type ApiKeyGetter = () => Promise; + +async function resolveEnricherConfig( + deps: FileEnrichmentDeps, + timeoutMs: number, + getApiKey: ApiKeyGetter, +): Promise { + const apiKey = await getApiKey(); + if (!apiKey) return null; + return { + apiKey, + host: deps.apiConfig.apiUrl, + projectId: deps.apiConfig.projectId, + timeoutMs, + }; +} + +async function getApmLineStats( + deps: FileEnrichmentDeps, + filePath: string, + getApiKey: ApiKeyGetter, +): Promise { + const cached = deps.apmStatsCache.get(filePath); + if (cached && cached.expiresAt > Date.now()) return cached.stats; + + try { + const config = await resolveEnricherConfig( + deps, + APM_QUERY_TIMEOUT_MS, + getApiKey, + ); + if (!config) return null; + + const stats = await deps.fetchApmLineStats(config, filePath); + deps.apmStatsCache.set(filePath, { + stats, + expiresAt: Date.now() + APM_STATS_TTL_MS, + }); + return stats; + } catch (err) { + deps.logger?.debug("APM enrichment failed", { + filePath, + message: err instanceof Error ? err.message : String(err), + }); + return null; + } +} + +async function enrichEventsForAgent( + deps: FileEnrichmentDeps, + filePath: string, + content: string, + getApiKey: ApiKeyGetter, +): Promise { const ext = path.extname(filePath).toLowerCase(); const langId = EXT_TO_LANG_ID[ext]; if (!langId || !deps.enricher.isSupported(langId)) return null; @@ -53,17 +183,12 @@ export async function enrichFileForAgent( 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 { @@ -72,15 +197,10 @@ export async function enrichFileForAgent( return null; } - const apiKey = await deps.apiConfig.getApiKey(); - if (!apiKey) return null; + const config = await resolveEnricherConfig(deps, 5_000, getApiKey); + if (!config) return null; - const enriched = await parsed.enrichFromApi({ - apiKey, - host: deps.apiConfig.apiUrl, - projectId: deps.apiConfig.projectId, - timeoutMs: 5_000, - }); + const enriched = await parsed.enrichFromApi(config); const annotated = enriched.toInlineComments(); if (annotated === content) { @@ -132,11 +252,16 @@ async function buildWrapperContext( 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 }; + try { + const wrappers = await deps.enricher.getWrappersForFile( + edge.resolvedAbsPath, + ); + if (!wrappers.length) return null; + return { edge, wrappers }; + } catch { + // A failed import resolution must not reject the whole read. + return null; + } }), );