fix(livekit): Respect LiveKit's PII and project redaction settings

.changeset/livekit-pii.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix(livekit): Respect LiveKit's PII and project redaction settings

e2e/scenarios/livekit-instrumentation/redacted.mjs

@@ -0,0 +1,13 @@
+// Runs the scenario inside a LiveKit job whose project enables redaction, as
+// LiveKit Cloud sets it, so instrumentation must honor LiveKit's PII setting.
+// The job is minimal: only the members LiveKit reads outside a real worker.
+import { runWithJobContextAsync } from "@livekit/agents";
+
+const job = {
+  _redactionEnabled: true,
+  job: { id: "job_redacted", agentName: "", enableRecording: false },
+  room: { isConnected: false },
+  simulationContext: () => undefined,
+  _otelMetadata: () => ({}),
+};
+await runWithJobContextAsync(job, () => import("./scenario.mjs"));

e2e/scenarios/livekit-instrumentation/scenario.test.ts

@@ -85,6 +85,17 @@ type RunOptions = {
   cassette?: string;
   recording?: boolean;
 };
+/** The span IDs a span links to with this relationship, in order. */
+function linkTargets(span: CapturedLogEvent | undefined, type: string) {
+  return (
+    (span?.metadata?.links ?? []) as {
+      context: { spanId: string };
+      attributes: Record<string, string>;
+    }[]
+  )
+    .filter((l) => l.attributes["braintrust.relationship.type"] === type)
+    .map((l) => l.context.spanId);
+}
 /** What the scenario reported about each uploaded audio file, in order. */
 function uploadedAudio(
   stdout: string,
@@ -319,6 +330,58 @@ for (const variant of variants) {
       120000,
     );
 
+  // A LiveKit Cloud project with redaction: LiveKit's own PII setting wins.
+  for (const mode of ["auto", "wrapped"] as const)
+    test.sequential(
+      `${variant} ${mode} order conversation in a redacted LiveKit job`,
+      async () => {
+        await withScenarioHarness(async (harness) => {
+          await harness.runNodeScenarioDir(
+            scenarioOptions(variant, {
+              mode,
+              entry: "redacted.mjs",
+              cassette: "order",
+              recording: false,
+            }),
+          );
+          const events = assertAudioTrace(harness.events());
+          // Nothing the caller or the model said, and no tool payloads.
+          for (const name of ["llm_response", "agent_turn", "tts"])
+            for (const span of findAllSpans(events, name)) {
+              expect(span.row.input, name).toBeUndefined();
+              if (name !== "agent_turn")
+                expect(span.row.output, name).toBeUndefined();
+            }
+          expect(
+            findLatestSpan(events, "lookup_order")?.row.input,
+          ).toBeUndefined();
+          expect(
+            findLatestSpan(events, "lookup_order")?.row.output,
+          ).toBeUndefined();
+          expect(JSON.stringify(events.map((e) => e.row))).not.toMatch(
+            /order 1042|Friday/i,
+          );
+          // Links still connect the call, holding only span IDs.
+          const tool = findLatestSpan(events, "lookup_order");
+          expect(linkTargets(tool, "requested_by")).toHaveLength(1);
+          expect(
+            linkTargets(
+              findAllSpans(events, "llm_response")[1],
+              "consumes_tool_result",
+            ),
+          ).toEqual([tool?.span.id]);
+          // The call's shape, models and timing remain.
+          expect(
+            events.filter((e) => e.span.type === "llm").map((e) => e.span.name),
+          ).toEqual(["llm_response", "llm_response"]);
+          expect(
+            findLatestSpan(events, "llm_response")?.metadata,
+          ).toMatchObject({ model: expect.any(String) });
+        });
+      },
+      120000,
+    );
+
   test.sequential(
     `${variant} streaming-cascade order conversation`,
     async () => {

e2e/scenarios/livekit-instrumentation/websocket-cassette.mjs

@@ -4,12 +4,24 @@ import { readFile, writeFile, mkdir } from "node:fs/promises";
 import { setImmediate } from "node:timers/promises";
 import { dirname } from "node:path";
 import assert from "node:assert/strict";
-export async function websocketCassette(path, record) {
+/**
+ * Records or replays one realtime websocket session. `upstream` is the
+ * provider's origin; `authorization` builds its Authorization header from the
+ * client's upgrade request.
+ */
+export async function websocketCassette(
+  path,
+  record,
+  {
+    upstream = "wss://api.openai.com",
+    authorization = () => `Bearer ${process.env.OPENAI_API_KEY}`,
+  } = {},
+) {
   const events = record ? [] : JSON.parse(await readFile(path, "utf8"));
   let index = 0,
     failure,
     socket,
-    upstream;
+    upstreamSocket;
   const ids = new Map();
   const server = new WebSocketServer({ port: 0, host: "127.0.0.1" });
   await new Promise((resolve) => server.once("listening", resolve));
@@ -73,13 +85,13 @@ export async function websocketCassette(path, record) {
     }
     socket = client;
     if (record) {
-      const url = new URL(request.url, "wss://api.openai.com");
+      const url = new URL(request.url, upstream);
       // LiveKit adds model routing for gateways, but OpenAI transcription
       // selects its model in session.update rather than the upgrade URL.
       if (url.searchParams.get("intent") === "transcription")
         url.searchParams.delete("model");
-      const remote = (upstream = new WebSocket(url, {
-        headers: { Authorization: `Bearer ${process.env.OPENAI_API_KEY}` },
+      const remote = (upstreamSocket = new WebSocket(url, {
+        headers: { Authorization: authorization(request) },
       }));
       const pending = [];
       client.on("message", (data) => {
@@ -126,19 +138,21 @@ export async function websocketCassette(path, record) {
     baseURL: `http://127.0.0.1:${server.address().port}/v1`,
     async close() {
       socket?.terminate();
-      upstream?.terminate();
+      upstreamSocket?.terminate();
       await new Promise((resolve) => server.close(resolve));
       if (failure) throw failure;
       if (record) {
         await mkdir(dirname(path), { recursive: true });
         // Provider session IDs can resemble access tokens. Preserve references
         // with synthetic IDs in the cassette, without changing live traffic.
         const sessionIds = new Map();
-        for (const event of events) {
-          const id = event.message?.session?.id;
-          if (typeof id === "string" && !sessionIds.has(id))
-            sessionIds.set(id, `session_fixture_${sessionIds.size + 1}`);
-        }
+        for (const event of events)
+          for (const id of [
+            event.message?.session?.id,
+            event.message?.session_id,
+          ])
+            if (typeof id === "string" && !sessionIds.has(id))
+              sessionIds.set(id, `session_fixture_${sessionIds.size + 1}`);
         const serialized = JSON.stringify(
           events,
           (_key, value) => sessionIds.get(value) ?? value,

js/src/auto-instrumentations/configs/livekit.ts

@@ -58,6 +58,8 @@ const targets: [
     "voice/audio_recognition",
     { className: "AudioRecognition", methodName: "processSTTEvent" },
   ],
+  // LiveKit reads its project redaction setting from the job context.
+  ["jobContext", "job", { functionName: "getJobContext" }],
   // These factories create the async inference workers inside the scoped call.
   // Suppressing here leaves tool execution and other session work instrumented.
   ["inference", "voice/generation", { functionName: "performLLMInference" }],

js/src/instrumentation/plugins/livekit-interceptors.ts

@@ -54,4 +54,7 @@ export const livekitChannels = defineInterceptor("@livekit/agents", {
   stt: channel<[SpeechEvent], void>({
     channelName: "recognition.stt",
   }),
+  jobContext: channel<[boolean?], unknown>({
+    channelName: "job.context",
+  }),
 });

js/src/instrumentation/plugins/livekit-pii.ts

@@ -0,0 +1,73 @@
+import iso from "../../isomorph";
+
+/*
+ * LiveKit's own PII rule for its telemetry (telemetry/redaction.js and pii.js
+ * in @livekit/agents 1.9.x), so the trace omits exactly what LiveKit would.
+ */
+
+/** The span attribute with which LiveKit marks a span of a redacted job. */
+export const REDACTION_ENABLED = "lk.redaction.enabled";
+
+/** The text LiveKit substitutes for exception details when PII is not allowed. */
+export const REDACTED_EXCEPTION_MESSAGE = "exception details redacted";
+
+const PII_SEGMENT = /(^|\.)pii(\.|$)/i;
+// Flagged by the gen_ai semantic conventions as likely to contain user data.
+const GEN_AI_PII_ATTRIBUTES = new Set([
+  "gen_ai.input.messages",
+  "gen_ai.output.messages",
+  "gen_ai.system_instructions",
+  "gen_ai.tool.call.arguments",
+  "gen_ai.tool.call.result",
+  "gen_ai.tool.description",
+  "gen_ai.tool.definitions",
+  "gen_ai.retrieval.query.text",
+  "gen_ai.retrieval.documents",
+  "gen_ai.memory.query.text",
+  "gen_ai.memory.records",
+  "gen_ai.evaluation.explanation",
+]);
+const PII_EVENTS = new Set([
+  "gen_ai.system.message",
+  "gen_ai.user.message",
+  "gen_ai.assistant.message",
+  "gen_ai.tool.message",
+  "gen_ai.choice",
+  "gen_ai.client.inference.operation.details",
+]);
+const EXCEPTION_MESSAGE = "exception.message";
+const EXCEPTION_STACKTRACE = "exception.stacktrace";
+const FALSY = new Set(["0", "false", "no", "off"]);
+
+/**
+ * Whether LiveKit's configuration allows PII in its telemetry: allowed unless
+ * LIVEKIT_TELEMETRY_ALLOW_PII is falsy, or the job enables LiveKit Cloud
+ * project redaction (which wins over the environment).
+ */
+export function livekitPiiAllowed(job?: { _redactionEnabled?: boolean }) {
+  if (job?._redactionEnabled) return false;
+  const allow = iso.getEnv("LIVEKIT_TELEMETRY_ALLOW_PII");
+  return allow === undefined || !FALSY.has(allow.trim().toLowerCase());
+}
+
+export function isPiiEvent(name: string): boolean {
+  return PII_EVENTS.has(name);
+}
+
+/** Native attributes without PII, and with exception details replaced, as LiveKit does. */
+export function withoutPii(
+  attributes: Iterable<[string, unknown]>,
+): Map<string, unknown> {
+  const kept = new Map<string, unknown>();
+  for (const [key, value] of attributes) {
+    if (key === EXCEPTION_MESSAGE) kept.set(key, REDACTED_EXCEPTION_MESSAGE);
+    else if (
+      key !== EXCEPTION_STACKTRACE &&
+      !PII_SEGMENT.test(key) &&
+      !GEN_AI_PII_ATTRIBUTES.has(key) &&
+      !key.startsWith("gen_ai.prompt.variable")
+    )
+      kept.set(key, value);
+  }
+  return kept;
+}