feat(apm): editor APM enrichment service + tRPC route

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;

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);
 

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,

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,

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<ApmEnrichmentService>(APM_ENRICHMENT_SERVICE)
+        .enrichFile(input),
+    ),
+});

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

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> = {}): 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<string, string>).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,
+    });
+  });
+});

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<SerializedApmEnrichment | null> {
+    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<EnricherApiConfig | null> {
+    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;
+    }
+  }
+}

packages/workspace-server/src/services/apm-enrichment/identifiers.ts

@@ -0,0 +1,3 @@
+export const APM_ENRICHMENT_SERVICE = Symbol.for(
+  "posthog.core.apmEnrichmentService",
+);