From df9dcd6abd4106a0ec61f4d7dbea23693765ddd5 Mon Sep 17 00:00:00 2001 From: Jon McCallum Date: Mon, 22 Jun 2026 13:35:38 +0100 Subject: [PATCH] feat(apm): editor APM enrichment service + tRPC route --- apps/code/src/main/di/bindings.ts | 4 + apps/code/src/main/di/container.ts | 4 + apps/code/src/main/trpc/router.ts | 2 + packages/host-router/src/router.ts | 2 + .../src/routers/apm-enrichment.router.ts | 18 ++ .../apm-enrichment/apmEnrichment.module.ts | 7 + .../apm-enrichment/apmEnrichment.test.ts | 195 ++++++++++++++++++ .../services/apm-enrichment/apmEnrichment.ts | 98 +++++++++ .../services/apm-enrichment/identifiers.ts | 3 + 9 files changed, 333 insertions(+) create mode 100644 packages/host-router/src/routers/apm-enrichment.router.ts create mode 100644 packages/workspace-server/src/services/apm-enrichment/apmEnrichment.module.ts create mode 100644 packages/workspace-server/src/services/apm-enrichment/apmEnrichment.test.ts create mode 100644 packages/workspace-server/src/services/apm-enrichment/apmEnrichment.ts create mode 100644 packages/workspace-server/src/services/apm-enrichment/identifiers.ts diff --git a/apps/code/src/main/di/bindings.ts b/apps/code/src/main/di/bindings.ts index 1bd40d083..2021824ea 100644 --- a/apps/code/src/main/di/bindings.ts +++ b/apps/code/src/main/di/bindings.ts @@ -133,6 +133,8 @@ import type { AGENT_REPO_FILES, AGENT_SLEEP_COORDINATOR, } from "@posthog/workspace-server/services/agent/identifiers"; +import type { ApmEnrichmentService } from "@posthog/workspace-server/services/apm-enrichment/apmEnrichment"; +import type { APM_ENRICHMENT_SERVICE } from "@posthog/workspace-server/services/apm-enrichment/identifiers"; import type { ARCHIVE_FILE_WATCHER, ARCHIVE_SESSION_CANCELLER, @@ -353,6 +355,8 @@ export interface MainBindings { // Enrichment host ports [ENRICHMENT_AUTH]: EnrichmentAuth; [ENRICHMENT_FILE_READER]: EnrichmentFileReader; + // APM enrichment service (reuses ENRICHMENT_AUTH for credentials) + [APM_ENRICHMENT_SERVICE]: ApmEnrichmentService; // Provisioning [MAIN_PROVISIONING_SERVICE]: ProvisioningService; diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index 54ab5b18e..16d214a57 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -142,6 +142,7 @@ import { AGENT_SLEEP_COORDINATOR, } from "@posthog/workspace-server/services/agent/identifiers"; import { AgentServiceEvent } from "@posthog/workspace-server/services/agent/schemas"; +import { apmEnrichmentModule } from "@posthog/workspace-server/services/apm-enrichment/apmEnrichment.module"; import { archiveModule } from "@posthog/workspace-server/services/archive/archive.module"; import { ARCHIVE_FILE_WATCHER, @@ -425,6 +426,9 @@ container.bind(ENRICHMENT_FILE_READER).toConstantValue({ listFilesContainingText: (repoPath: string, text: string) => listFilesContainingText(repoPath, text), }); +// APM enrichment reuses the ENRICHMENT_AUTH port bound above for PostHog +// API credentials + active project; it only needs its own service binding. +container.load(apmEnrichmentModule); container.bind(MAIN_PROVISIONING_SERVICE).to(ProvisioningService); container.bind(PROVISIONING_SERVICE).toService(MAIN_TOKENS.ProvisioningService); diff --git a/apps/code/src/main/trpc/router.ts b/apps/code/src/main/trpc/router.ts index ec493315d..b81fd6261 100644 --- a/apps/code/src/main/trpc/router.ts +++ b/apps/code/src/main/trpc/router.ts @@ -1,6 +1,7 @@ import { additionalDirectoriesRouter } from "@posthog/host-router/routers/additional-directories.router"; import { agentRouter } from "@posthog/host-router/routers/agent.router"; import { analyticsRouter } from "@posthog/host-router/routers/analytics.router"; +import { apmEnrichmentRouter } from "@posthog/host-router/routers/apm-enrichment.router"; import { archiveRouter } from "@posthog/host-router/routers/archive.router"; import { authRouter } from "@posthog/host-router/routers/auth.router"; import { canvasDataRouter } from "@posthog/host-router/routers/canvas-data.router"; @@ -52,6 +53,7 @@ export const trpcRouter = router({ additionalDirectories: additionalDirectoriesRouter, agent: agentRouter, analytics: analyticsRouter, + apmEnrichment: apmEnrichmentRouter, archive: archiveRouter, auth: authRouter, canvasData: canvasDataRouter, diff --git a/packages/host-router/src/router.ts b/packages/host-router/src/router.ts index 6370158e1..900d4cfee 100644 --- a/packages/host-router/src/router.ts +++ b/packages/host-router/src/router.ts @@ -2,6 +2,7 @@ import { router } from "@posthog/host-trpc/trpc"; import { additionalDirectoriesRouter } from "./routers/additional-directories.router"; import { agentRouter } from "./routers/agent.router"; import { analyticsRouter } from "./routers/analytics.router"; +import { apmEnrichmentRouter } from "./routers/apm-enrichment.router"; import { archiveRouter } from "./routers/archive.router"; import { authRouter } from "./routers/auth.router"; import { canvasDataRouter } from "./routers/canvas-data.router"; @@ -50,6 +51,7 @@ export const hostRouter = router({ additionalDirectories: additionalDirectoriesRouter, agent: agentRouter, analytics: analyticsRouter, + apmEnrichment: apmEnrichmentRouter, archive: archiveRouter, auth: authRouter, canvasData: canvasDataRouter, diff --git a/packages/host-router/src/routers/apm-enrichment.router.ts b/packages/host-router/src/routers/apm-enrichment.router.ts new file mode 100644 index 000000000..89ce25c8a --- /dev/null +++ b/packages/host-router/src/routers/apm-enrichment.router.ts @@ -0,0 +1,18 @@ +import { publicProcedure, router } from "@posthog/host-trpc/trpc"; +import type { ApmEnrichmentService } from "@posthog/workspace-server/services/apm-enrichment/apmEnrichment"; +import { APM_ENRICHMENT_SERVICE } from "@posthog/workspace-server/services/apm-enrichment/identifiers"; +import { z } from "zod"; + +const enrichFileInput = z.object({ + filePath: z.string(), +}); + +export const apmEnrichmentRouter = router({ + enrichFile: publicProcedure + .input(enrichFileInput) + .query(({ ctx, input }) => + ctx.container + .get(APM_ENRICHMENT_SERVICE) + .enrichFile(input), + ), +}); diff --git a/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.module.ts b/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.module.ts new file mode 100644 index 000000000..3b3020eab --- /dev/null +++ b/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.module.ts @@ -0,0 +1,7 @@ +import { ContainerModule } from "inversify"; +import { ApmEnrichmentService } from "./apmEnrichment"; +import { APM_ENRICHMENT_SERVICE } from "./identifiers"; + +export const apmEnrichmentModule = new ContainerModule(({ bind }) => { + bind(APM_ENRICHMENT_SERVICE).to(ApmEnrichmentService).inSingletonScope(); +}); diff --git a/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.test.ts b/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.test.ts new file mode 100644 index 000000000..a256696aa --- /dev/null +++ b/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.test.ts @@ -0,0 +1,195 @@ +import type { RootLogger } from "@posthog/di/logger"; +import { APM_STATS_WINDOW, type SymbolStatsRow } from "@posthog/shared"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { EnrichmentAuth } from "../enrichment/ports"; +import { ApmEnrichmentService } from "./apmEnrichment"; + +const noop = () => {}; +const noopLogger = { + debug: noop, + info: noop, + warn: noop, + error: noop, + scope: () => noopLogger, +} as unknown as RootLogger; + +function authed(): EnrichmentAuth { + return { + getState: () => ({ + status: "authenticated", + projectId: 2, + cloudRegion: "us", + }), + getValidAccessToken: async () => ({ + accessToken: "tok", + apiHost: "https://us.posthog.com", + }), + }; +} + +function unauthed(): EnrichmentAuth { + return { + getState: () => ({ + status: "unauthenticated", + projectId: null, + cloudRegion: null, + }), + getValidAccessToken: vi.fn(), + }; +} + +// Authenticated state, but token resolution blows up (network error, corrupted +// token store) — exercises the catch in resolveApiConfig. +function tokenThrows(): EnrichmentAuth { + return { + getState: () => ({ + status: "authenticated", + projectId: 2, + cloudRegion: "us", + }), + getValidAccessToken: async () => { + throw new Error("token store unavailable"); + }, + }; +} + +// A symbol-stats response row (line mode), typed against the wire shape so the +// fixture can't drift from the fields mapSymbolStatsResults reads. +function symbolRow(overrides: Partial = {}): SymbolStatsRow { + return { + line: 459, + count: 100, + error_count: 0, + sum_duration_nano: 0, + p50_duration_nano: 2_000_000, + p95_duration_nano: 7_000_000, + p99_duration_nano: 0, + busy_count: 0, + p50_busy_nano: 0, + p95_busy_nano: 0, + p99_busy_nano: 0, + count_pct_change: null, + p50_duration_pct_change: null, + p95_duration_pct_change: null, + p99_duration_pct_change: null, + error_rate_pct_change: null, + ...overrides, + }; +} + +afterEach(() => { + vi.unstubAllGlobals(); +}); + +describe("ApmEnrichmentService", () => { + // Every row resolves to null; they differ only in what fails. Rows with no + // fetchImpl must never reach the network — resolveApiConfig bails first. + it.each([ + { name: "unauthenticated", auth: unauthed, fetchImpl: undefined }, + { + name: "access-token resolution throws", + auth: tokenThrows, + fetchImpl: undefined, + }, + { + name: "the symbol-stats query responds with an error status", + auth: authed, + fetchImpl: async () => ({ + ok: false, + status: 500, + statusText: "Server Error", + json: async () => ({}), + }), + }, + ])("returns null when $name", async ({ auth, fetchImpl }) => { + const fetchMock = vi.fn(fetchImpl); + vi.stubGlobal("fetch", fetchMock); + + const result = await new ApmEnrichmentService( + auth(), + noopLogger, + ).enrichFile({ filePath: "rust/feature-flags/src/flags/flag_matching.rs" }); + + expect(result).toBeNull(); + if (!fetchImpl) expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("POSTs a line-mode symbol-stats query for the file", async () => { + const fetchMock = vi.fn(async () => ({ + ok: true, + status: 200, + json: async () => ({ results: [symbolRow()], granularity: "line" }), + })); + vi.stubGlobal("fetch", fetchMock); + + await new ApmEnrichmentService(authed(), noopLogger).enrichFile({ + filePath: "rust/feature-flags/src/flags/flag_matching.rs", + }); + + const [url, init] = fetchMock.mock.calls[0] as unknown as [ + string, + RequestInit, + ]; + expect(url).toBe( + "https://us.posthog.com/api/projects/2/tracing/spans/symbol-stats/", + ); + expect(init.method).toBe("POST"); + expect((init.headers as Record).Authorization).toBe( + "Bearer tok", + ); + const body = JSON.parse(init.body as string); + expect(body.query.kind).toBe("TraceSpansSymbolStatsQuery"); + expect(body.query.filePath).toBe( + "rust/feature-flags/src/flags/flag_matching.rs", + ); + expect(body.query.dateRange.date_from).toBe(APM_STATS_WINDOW.dateFrom); + // Line mode — the editor gutter wants per-line stats, not per-symbol. + expect(body.query.symbols).toBeUndefined(); + }); + + it("maps symbol-stats rows to per-line stats for the file", async () => { + const fetchMock = vi.fn(async () => ({ + ok: true, + status: 200, + json: async () => ({ + results: [ + symbolRow({ line: 459, count: 26, error_count: 0 }), + symbolRow({ + line: 900, + count: 30, + error_count: 8, + p95_duration_pct_change: 186, + }), + ], + granularity: "line", + }), + })); + vi.stubGlobal("fetch", fetchMock); + + const result = await new ApmEnrichmentService( + authed(), + noopLogger, + ).enrichFile({ + filePath: "rust/feature-flags/src/flags/flag_matching.rs", + }); + + expect(result?.filePath).toBe( + "rust/feature-flags/src/flags/flag_matching.rs", + ); + expect(result?.tracingUrl).toBe("https://us.posthog.com/project/2/tracing"); + expect(result?.stats).toHaveLength(2); + // p50_duration_nano 2_000_000 → 2ms, p95 7_000_000 → 7ms (nsToMs). + expect(result?.stats[0]).toMatchObject({ + line: 459, + count: 26, + p50Ms: 2, + p95Ms: 7, + }); + expect(result?.stats[1]).toMatchObject({ + line: 900, + count: 30, + errorCount: 8, + p95PctChange: 186, + }); + }); +}); diff --git a/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.ts b/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.ts new file mode 100644 index 000000000..3fca8d755 --- /dev/null +++ b/packages/workspace-server/src/services/apm-enrichment/apmEnrichment.ts @@ -0,0 +1,98 @@ +import { + ROOT_LOGGER, + type RootLogger, + type ScopedLogger, +} from "@posthog/di/logger"; +import { type EnricherApiConfig, PostHogApi } from "@posthog/enricher"; +import { + APM_STATS_WINDOW, + type SerializedApmEnrichment, +} from "@posthog/shared"; +import { inject, injectable } from "inversify"; +import { ENRICHMENT_AUTH } from "../enrichment/identifiers"; +import type { EnrichmentAuth } from "../enrichment/ports"; + +export interface ApmEnrichFileInput { + filePath: string; +} + +// 24h query is ~8s on the hottest traced file; a lower budget silently drops it. +const QUERY_TIMEOUT_MS = 20_000; + +function tracingExplorerUrl(host: string, projectId: number): string { + return `${host.replace(/\/$/, "")}/project/${projectId}/tracing`; +} + +@injectable() +export class ApmEnrichmentService { + private readonly log: ScopedLogger; + + constructor( + @inject(ENRICHMENT_AUTH) private readonly authService: EnrichmentAuth, + @inject(ROOT_LOGGER) logger: RootLogger, + ) { + this.log = logger.scope("ApmEnrichmentService"); + } + + async enrichFile( + input: ApmEnrichFileInput, + ): Promise { + this.log.debug("[apm] enrichFile", { filePath: input.filePath }); + const config = await this.resolveApiConfig(); + if (!config) return null; + + try { + const stats = await new PostHogApi(config).getApmLineStats( + input.filePath, + { dateFrom: APM_STATS_WINDOW.dateFrom }, + ); + this.log.info("[apm] enriched", { + filePath: input.filePath, + host: config.host, + projectId: config.projectId, + lines: stats.length, + }); + return { + filePath: input.filePath, + stats, + tracingUrl: tracingExplorerUrl(config.host, config.projectId), + }; + } catch (err) { + this.log.warn("[apm] query failed", { + filePath: input.filePath, + message: err instanceof Error ? err.message : String(err), + }); + return null; + } + } + + private async resolveApiConfig(): Promise { + const state = this.authService.getState(); + if ( + state.status !== "authenticated" || + !state.projectId || + !state.cloudRegion + ) { + this.log.info("[apm] auth not ready", { + status: state.status, + projectId: state.projectId, + cloudRegion: state.cloudRegion, + }); + return null; + } + try { + const auth = await this.authService.getValidAccessToken(); + return { + apiKey: auth.accessToken, + host: auth.apiHost, + projectId: state.projectId, + timeoutMs: QUERY_TIMEOUT_MS, + }; + } catch (err) { + this.log.warn("[apm] failed to resolve access token", { + message: err instanceof Error ? err.message : String(err), + }); + return null; + } + } +} diff --git a/packages/workspace-server/src/services/apm-enrichment/identifiers.ts b/packages/workspace-server/src/services/apm-enrichment/identifiers.ts new file mode 100644 index 000000000..752d8cd58 --- /dev/null +++ b/packages/workspace-server/src/services/apm-enrichment/identifiers.ts @@ -0,0 +1,3 @@ +export const APM_ENRICHMENT_SERVICE = Symbol.for( + "posthog.core.apmEnrichmentService", +);