feat(otel): Add distributed tracing across app and workspace server

apps/code/package.json

@@ -92,10 +92,16 @@
     "@inversifyjs/strongly-typed": "2.2.0",
     "@json-render/core": "^0.19.0",
     "@modelcontextprotocol/sdk": "^1.12.1",
+    "@opentelemetry/api": "^1.9.0",
     "@opentelemetry/api-logs": "^0.208.0",
+    "@opentelemetry/context-zone": "^2.5.0",
     "@opentelemetry/exporter-logs-otlp-http": "^0.208.0",
+    "@opentelemetry/exporter-trace-otlp-http": "^0.208.0",
     "@opentelemetry/resources": "^2.5.0",
     "@opentelemetry/sdk-logs": "^0.208.0",
+    "@opentelemetry/sdk-trace-base": "^2.5.0",
+    "@opentelemetry/sdk-trace-node": "^2.5.0",
+    "@opentelemetry/sdk-trace-web": "^2.5.0",
     "@opentelemetry/semantic-conventions": "^1.39.0",
     "@parcel/watcher": "^2.5.6",
     "@phosphor-icons/react": "^2.1.10",

apps/code/src/main/index.ts

@@ -66,6 +66,7 @@ import {
   readChromiumLogTail,
 } from "./utils/logger";
 import { isMacosPackagedUnsafeBundleLocation } from "./utils/macos-packaged-install-guard";
+import { initOtelTracing } from "./utils/otel-trace";
 import { createWindow } from "./window";
 
 type FileWatcherEventsByKind = {
@@ -101,6 +102,8 @@ export class FileWatcherBridge extends TypedEventEmitter<FileWatcherEventsByKind
   }
 }
 
+initOtelTracing();
+
 // Single instance lock must be acquired FIRST before any other app setup
 const additionalData = process.defaultApp ? { argv: process.argv } : undefined;
 const gotTheLock = app.requestSingleInstanceLock(additionalData);

apps/code/src/main/services/app-lifecycle/service.ts

@@ -17,6 +17,7 @@ import { posthogNodeAnalytics } from "../../platform-adapters/posthog-analytics"
 import { withTimeout } from "../../utils/async";
 import { logger } from "../../utils/logger";
 import { shutdownOtelTransport } from "../../utils/otel-log-transport";
+import { shutdownOtelTracing } from "../../utils/otel-trace";
 
 const log = logger.scope("app-lifecycle");
 
@@ -144,6 +145,12 @@ export class AppLifecycleService {
       log.warn("Failed to shutdown OTEL log transport", error);
     }
 
+    try {
+      await shutdownOtelTracing();
+    } catch (error) {
+      log.warn("Failed to shutdown OTEL trace exporter", error);
+    }
+
     try {
       await posthogNodeAnalytics.shutdown();
     } catch (error) {

apps/code/src/main/trpc/trpc.ts

@@ -1,8 +1,10 @@
+import { getMainTracer } from "@main/utils/otel-trace";
 import {
   publicProcedure as baseProcedure,
   router as baseRouter,
   middleware,
 } from "@posthog/host-trpc/trpc";
+import { traceTrpcCall } from "@posthog/workspace-server/node-tracing";
 import log from "electron-log/main";
 
 const CALL_RATE_WINDOW_MS = 2000;
@@ -55,6 +57,12 @@ const callRateMonitor = middleware(async ({ path, next, type }) => {
   return result;
 });
 
+const tracingMonitor = middleware(({ path, next, type }) =>
+  traceTrpcCall(getMainTracer(), path, type, next),
+);
+
 export const router = baseRouter;
-export const publicProcedure = baseProcedure.use(callRateMonitor);
+export const publicProcedure = baseProcedure
+  .use(callRateMonitor)
+  .use(tracingMonitor);
 export { middleware };

apps/code/src/main/utils/otel-trace.test.ts

@@ -0,0 +1,125 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const mockGetTracer = vi.fn(() => ({}) as unknown);
+const mockRegister = vi.fn();
+const mockForceFlush = vi.fn(() => Promise.resolve());
+const mockShutdown = vi.fn(() => Promise.resolve());
+
+vi.mock("@opentelemetry/exporter-trace-otlp-http", () => ({
+  OTLPTraceExporter: class {
+    constructor(public config: unknown) {}
+  },
+}));
+
+vi.mock("@opentelemetry/sdk-trace-base", () => ({
+  BatchSpanProcessor: class {
+    constructor(
+      public _exporter: unknown,
+      public _opts: unknown,
+    ) {}
+  },
+  ParentBasedSampler: class {
+    constructor(public _opts: unknown) {}
+  },
+  TraceIdRatioBasedSampler: class {
+    constructor(public _ratio: number) {}
+  },
+}));
+
+vi.mock("@opentelemetry/sdk-trace-node", () => ({
+  NodeTracerProvider: class {
+    constructor(public _opts: unknown) {}
+    register() {
+      return mockRegister();
+    }
+    getTracer() {
+      return mockGetTracer();
+    }
+    forceFlush() {
+      return mockForceFlush();
+    }
+    shutdown() {
+      return mockShutdown();
+    }
+  },
+}));
+
+vi.mock("@opentelemetry/resources", () => ({
+  resourceFromAttributes: vi.fn((attrs: Record<string, string>) => attrs),
+}));
+
+vi.mock("@opentelemetry/semantic-conventions", () => ({
+  ATTR_SERVICE_NAME: "service.name",
+  ATTR_SERVICE_VERSION: "service.version",
+}));
+
+describe("otel-trace", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.resetModules();
+    vi.unstubAllEnvs();
+    process.env.POSTHOG_CODE_VERSION = "1.0.0-test";
+  });
+
+  describe("initOtelTracing", () => {
+    it.each([
+      ["key missing", "", "https://test.posthog.com"],
+      ["host missing", "phc_test123", ""],
+      ["host invalid", "phc_test123", "not a url"],
+    ])("returns null when %s", async (_case, key, host) => {
+      vi.stubEnv("VITE_POSTHOG_API_KEY", key);
+      vi.stubEnv("VITE_POSTHOG_API_HOST", host);
+
+      const { initOtelTracing, getMainTracer } = await import(
+        "@main/utils/otel-trace"
+      );
+
+      expect(initOtelTracing()).toBeNull();
+      expect(getMainTracer()).toBeNull();
+      expect(mockRegister).not.toHaveBeenCalled();
+    });
+
+    it("registers a provider and returns a tracer when configured", async () => {
+      vi.stubEnv("VITE_POSTHOG_API_KEY", "phc_test123");
+      vi.stubEnv("VITE_POSTHOG_API_HOST", "https://test.posthog.com");
+
+      const { initOtelTracing, getMainTracer } = await import(
+        "@main/utils/otel-trace"
+      );
+      const tracer = initOtelTracing();
+
+      expect(tracer).not.toBeNull();
+      expect(getMainTracer()).toBe(tracer);
+      expect(mockRegister).toHaveBeenCalled();
+    });
+  });
+
+  describe("shutdownOtelTracing", () => {
+    it("flushes and shuts down the provider", async () => {
+      vi.stubEnv("VITE_POSTHOG_API_KEY", "phc_test123");
+      vi.stubEnv("VITE_POSTHOG_API_HOST", "https://test.posthog.com");
+
+      const { initOtelTracing, shutdownOtelTracing } = await import(
+        "@main/utils/otel-trace"
+      );
+      initOtelTracing();
+
+      await shutdownOtelTracing();
+
+      expect(mockForceFlush).toHaveBeenCalled();
+      expect(mockShutdown).toHaveBeenCalled();
+    });
+
+    it("is a no-op when provider was never created", async () => {
+      vi.stubEnv("VITE_POSTHOG_API_KEY", "");
+
+      const { initOtelTracing, shutdownOtelTracing } = await import(
+        "@main/utils/otel-trace"
+      );
+      initOtelTracing();
+
+      await expect(shutdownOtelTracing()).resolves.toBeUndefined();
+      expect(mockForceFlush).not.toHaveBeenCalled();
+    });
+  });
+});

apps/code/src/main/utils/otel-trace.ts

@@ -0,0 +1,36 @@
+import type { Tracer } from "@opentelemetry/api";
+import {
+  initNodeTracing,
+  type NodeTracing,
+} from "@posthog/workspace-server/node-tracing";
+import log from "electron-log/main";
+import { getAppVersion } from "./env";
+
+let current: NodeTracing | null = null;
+
+export function initOtelTracing(): Tracer | null {
+  current = initNodeTracing({
+    serviceName: "posthog-code-desktop",
+    serviceVersion: getAppVersion(),
+    attributes: {
+      "service.namespace": "ipc",
+      "process.runtime.name": "electron",
+      "process.runtime.version": process.versions.electron,
+    },
+  });
+  log
+    .scope("otel")
+    .info(
+      `tracing ${current ? "enabled" : "disabled (missing/invalid VITE_POSTHOG_API_KEY/HOST)"}`,
+    );
+  return current?.tracer ?? null;
+}
+
+export function getMainTracer(): Tracer | null {
+  return current?.tracer ?? null;
+}
+
+export async function shutdownOtelTracing(): Promise<void> {
+  await current?.shutdown();
+  current = null;
+}

apps/code/src/renderer/main.tsx

@@ -15,10 +15,16 @@ import { Providers } from "@components/Providers";
 import { preloadHighlighter } from "@pierre/diffs";
 import { boot } from "@posthog/di/contribution";
 import { ServiceProvider } from "@posthog/di/react";
+import { router } from "@posthog/ui/router/router";
 import App from "@posthog/ui/shell/App";
 import { initializePostHog } from "@posthog/ui/shell/posthogAnalyticsImpl";
 import { registerDesktopContributions } from "@renderer/desktop-contributions";
 import { container } from "@renderer/di/container";
+import {
+  initOtelTracing,
+  onAppRender,
+  traceNavigations,
+} from "@renderer/utils/otel-trace";
 import "@renderer/desktop-services";
 import React from "react";
 import ReactDOM from "react-dom/client";
@@ -80,6 +86,10 @@ if (bootstrapSessionId) {
   initializePostHog(bootstrapSessionId);
 }
 
+if (initOtelTracing()) {
+  traceNavigations(router);
+}
+
 registerDesktopContributions();
 void boot(container);
 
@@ -88,10 +98,12 @@ if (!rootElement) throw new Error("Root element not found");
 
 ReactDOM.createRoot(rootElement).render(
   <React.StrictMode>
-    <ServiceProvider container={container}>
-      <Providers>
-        <App />
-      </Providers>
-    </ServiceProvider>
+    <React.Profiler id="app" onRender={onAppRender}>
+      <ServiceProvider container={container}>
+        <Providers>
+          <App />
+        </Providers>
+      </ServiceProvider>
+    </React.Profiler>
   </React.StrictMode>,
 );