From 7f357ed953f6415ac9b888071256db4572e57d4a Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Tue, 15 Sep 2026 13:47:46 -0600 Subject: [PATCH 1/2] support span customizers --- .changeset/span-export-hooks.md | 9 + AGENTS.md | 2 + js/src/exports.ts | 6 +- js/src/instrumentation/README.md | 62 +++++ js/src/instrumentation/config.ts | 32 +++ js/src/instrumentation/index.ts | 1 + js/src/instrumentation/registry.ts | 4 + js/src/logger.ts | 39 +-- js/src/span-customizer.test.ts | 390 +++++++++++++++++++++++++++++ js/src/span-customizer.ts | 126 ++++++++++ 10 files changed, 656 insertions(+), 15 deletions(-) create mode 100644 .changeset/span-export-hooks.md create mode 100644 js/src/span-customizer.test.ts create mode 100644 js/src/span-customizer.ts diff --git a/.changeset/span-export-hooks.md b/.changeset/span-export-hooks.md new file mode 100644 index 000000000..cbd246537 --- /dev/null +++ b/.changeset/span-export-hooks.md @@ -0,0 +1,9 @@ +--- +"braintrust": minor +--- + +feat: add span export hooks + +Support synchronous `onSpanExport` customizers for incremental instrumentation +span records. Customizers can add, modify, delete, or replace fields before export, +with callbacks applied once per record rather than once per transport retry. diff --git a/AGENTS.md b/AGENTS.md index f474c4bcb..7229b7677 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -79,6 +79,8 @@ pnpm run test # Run all workspace tests via turbo Run from the repo root. **Always run `fix:formatting` before committing** — there is a pre-commit hook that will reject unformatted code. +Agents MUST run Prettier on every file they create or edit before handing work back, even when no commit is requested. Include Markdown, changelogs, config files, and generated files supported by Prettier—not just source code. From the repo root, run `pnpm exec prettier --write ` followed by `pnpm exec prettier --check `. If further edits are made, repeat formatting and verification after the final edit. Do not rely on tests, typechecks, CI, or the pre-commit hook to catch formatting issues. + ```bash pnpm run formatting # Check formatting (prettier) pnpm run lint # Run eslint checks diff --git a/js/src/exports.ts b/js/src/exports.ts index 683b3f868..1711f893c 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -374,6 +374,10 @@ export { braintrustFlueObserver, braintrustFlueInstrumentation, } from "./instrumentation"; -export type { InstrumentationConfig } from "./instrumentation"; +export type { + InstrumentationConfig, + SpanCustomizer, + SpanExportData, +} from "./instrumentation"; export { wrapElevenLabs } from "./wrappers/elevenlabs"; diff --git a/js/src/instrumentation/README.md b/js/src/instrumentation/README.md index 26487ec26..243ce3c57 100644 --- a/js/src/instrumentation/README.md +++ b/js/src/instrumentation/README.md @@ -180,6 +180,68 @@ termination, and async context. - Use narrow vendored provider interfaces shared by wrappers and plugins. - Keep enable, disable, subscription, and patching behavior idempotent. +## Export Customizers + +Configure `spanCustomizers` through the standalone instrumentation entrypoint +before importing the main SDK, which enables instrumentation during platform +initialization. Use a bootstrap module before any auto-instrumentation preload +that initializes the SDK. Static imports of the main SDK are hoisted; use a +dynamic import after configuration: + +```ts +import { configureInstrumentation } from "braintrust/instrumentation"; + +configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + data.tags = ["reviewed"]; + if ("output" in data) data.output = "[redacted]"; + delete data.error; + return data; + }, + }, + ], +}); + +const { initLogger } = await import("braintrust"); +initLogger({ projectName: "my-project" }); +// Import and use instrumented provider SDKs here. +``` + +`onSpanExport` receives each incremental record from an instrumentation-created +span after lazy values resolve, before attachment processing, merging, masking, +and JSON serialization. It can run before the span ends; fields may be absent. +Ordinary manually created spans, dataset rows, and feedback are not customized. + +Callbacks run synchronously in registration order. Mutate and return the record, +or return a replacement plain object for the next callback. Exceptions and invalid +return values are ignored while synchronous payload mutations remain; promises +are not awaited and their rejections are swallowed. Do not mutate the record after +returning. Export retries reuse the transformed record without invoking callbacks +again. Configuration is shared across SDK bundles. + +The SDK restores these fields after every callback, including removing injected +fields that were absent from the original record: + +- Identity: `id`, `span_id`, `root_span_id`, `span_parents`. +- Routing: `org_id`, `project_id`, `experiment_id`, `dataset_id`, + `prompt_session_id`, `log_id`, `function_data`. +- Transport controls: `_is_merge`, `_merge_paths`, `_parent_id`, `_object_delete`, + `_array_delete`, `_xact_id`. + +Payload values must remain supported by the SDK logging pipeline. They can still +include `Attachment` objects at this point; attachment processing and JSON +serialization happen after customization. + +This is an export-only hook, not a fail-closed privacy boundary. The local +experiment/scorer cache is populated before export and may retain unredacted +values. Applications requiring secrets to stay off local disk must disable the +span cache separately; export customization alone does not provide that guarantee. + +Customizers receive only the outgoing record, not a live span or provider +instrumentation context. + ## Testing Test at the narrowest useful layers: diff --git a/js/src/instrumentation/config.ts b/js/src/instrumentation/config.ts index 2f969dd58..d2696a918 100644 --- a/js/src/instrumentation/config.ts +++ b/js/src/instrumentation/config.ts @@ -1,3 +1,29 @@ +export type SpanExportData = Record; + +export interface SpanCustomizer { + /** + * Customize an outgoing span record after lazy values resolve, before JSON + * serialization. Records are incremental and may not contain every span field. + * + * Callbacks are synchronous. Add, change, or delete payload fields, then return + * the record or a replacement plain object. Payloads may still contain SDK + * Attachment objects; attachment processing and serialization happen later. + * + * The SDK restores identity and routing fields (id, span_id, root_span_id, + * span_parents, org_id, project_id, experiment_id, dataset_id, prompt_session_id, + * log_id, function_data) and transport controls (_is_merge, _merge_paths, + * _parent_id, _object_delete, _array_delete, _xact_id) after every callback. + * + * Exceptions and invalid return values are ignored; synchronous payload + * mutations remain. Promises are not awaited and their rejections are swallowed. + * Do not mutate the record after returning. + * + * This hook does not guarantee redaction of the local experiment/scorer cache, + * which is populated before export, and is not a fail-closed privacy boundary. + */ + onSpanExport?(data: SpanExportData): SpanExportData; +} + export interface InstrumentationIntegrationsConfig { openai?: boolean; anthropic?: boolean; @@ -46,6 +72,12 @@ export interface InstrumentationConfig { * Set to false to disable instrumentation for that SDK. */ integrations?: InstrumentationIntegrationsConfig; + + /** + * Instrumentation-wide customizers, in callback execution order. + * Configure before instrumentation is enabled. + */ + spanCustomizers?: readonly SpanCustomizer[]; } const envIntegrationAliases: Record< diff --git a/js/src/instrumentation/index.ts b/js/src/instrumentation/index.ts index a1dbd990c..833786cc5 100644 --- a/js/src/instrumentation/index.ts +++ b/js/src/instrumentation/index.ts @@ -45,3 +45,4 @@ export { // Configuration API export { configureInstrumentation } from "./registry"; export type { InstrumentationConfig } from "./registry"; +export type { SpanCustomizer, SpanExportData } from "./config"; diff --git a/js/src/instrumentation/registry.ts b/js/src/instrumentation/registry.ts index 0ed195ae8..1ed665a0a 100644 --- a/js/src/instrumentation/registry.ts +++ b/js/src/instrumentation/registry.ts @@ -13,6 +13,7 @@ import { type InstrumentationConfig, } from "./config"; import { GLOBAL_INSTRUMENTATION_HOOKS_PROTOCOL_VERSION } from "../global-instrumentation-hooks"; +import { setSpanCustomizers } from "../span-customizer"; export type { InstrumentationConfig } from "./config"; @@ -62,6 +63,9 @@ class PluginRegistry { return; } this.config = { ...this.config, ...config }; + if ("spanCustomizers" in config) { + setSpanCustomizers(config.spanCustomizers); + } } /** diff --git a/js/src/logger.ts b/js/src/logger.ts index e18756048..0c656cacb 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -207,6 +207,7 @@ import { mergeSpanOriginContext, type SpanOriginEnvironment, } from "./span-origin"; +import { customizeSpanExport } from "./span-customizer"; // Manual type definition for inline attachments (not in generated_types) const InlineAttachmentReferenceSchema = z.object({ @@ -8215,6 +8216,7 @@ export class SpanImpl implements Span { private isMerge: boolean; private loggedEndTime: number | undefined; + private readonly isInstrumented: boolean; private propagatedEvent: StartSpanEventArgs | undefined; // For internal use only. @@ -8255,6 +8257,8 @@ export class SpanImpl implements Span { const instrumentationName = getSpanInstrumentationName(args) ?? INSTRUMENTATION_NAMES.BRAINTRUST_JS_LOGGER; + this.isInstrumented = + instrumentationName !== INSTRUMENTATION_NAMES.BRAINTRUST_JS_LOGGER; const spanAttributes = args.spanAttributes ?? {}; const rawEvent = args.event ?? {}; @@ -8422,21 +8426,28 @@ export class SpanImpl implements Span { ); } - const computeRecord = async () => ({ - ...partialRecord, - ...Object.fromEntries( - await Promise.all( - Object.entries(lazyInternalData).map(async ([key, value]) => [ - key, - await value.get(), - ]), + const computeRecord = async () => { + const record = { + ...partialRecord, + ...Object.fromEntries( + await Promise.all( + Object.entries(lazyInternalData).map(async ([key, value]) => [ + key, + await value.get(), + ]), + ), ), - ), - ...new SpanComponentsV3({ - object_type: this.parentObjectType, - object_id: await this.parentObjectId.get(), - }).objectIdFields(), - }); + ...new SpanComponentsV3({ + object_type: this.parentObjectType, + object_id: await this.parentObjectId.get(), + }).objectIdFields(), + }; + // Customize inside the memoized lazy value, before attachment processing, + // merging, and masking. Retries reuse the already-customized record. + return this.isInstrumented + ? (customizeSpanExport(record) as BackgroundLogEvent) + : record; + }; this._state.bgLogger().log([new LazyValue(computeRecord)]); } diff --git a/js/src/span-customizer.test.ts b/js/src/span-customizer.test.ts new file mode 100644 index 000000000..4206a1ac9 --- /dev/null +++ b/js/src/span-customizer.test.ts @@ -0,0 +1,390 @@ +import { + afterEach, + beforeEach, + describe, + expect, + expectTypeOf, + test, + vi, +} from "vitest"; +import { + _exportsForTestingOnly, + BraintrustState, + initLogger, + type TestBackgroundLogger, +} from "./logger"; +import { configureInstrumentation, registry } from "./instrumentation/registry"; +import { configureNode } from "./node/config"; +import { + INSTRUMENTATION_NAMES, + withSpanInstrumentationName, +} from "./span-origin"; +import type { SpanCustomizer, SpanExportData } from "./exports"; +import { customizeSpanExport } from "./span-customizer"; + +configureNode(); + +test("customizers expose only the outgoing-record export hook", () => { + expectTypeOf().toEqualTypeOf<{ + onSpanExport?(data: SpanExportData): SpanExportData; + }>(); +}); + +describe("onSpanExport", () => { + let memoryLogger: TestBackgroundLogger; + + beforeEach(async () => { + registry.disable(); + await _exportsForTestingOnly.simulateLoginForTests(); + memoryLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + }); + + afterEach(() => { + configureInstrumentation({ spanCustomizers: [] }); + _exportsForTestingOnly.clearTestBackgroundLogger(); + vi.unstubAllEnvs(); + }); + + function startInstrumentedSpan() { + return initLogger({ + projectName: "customizer-project", + projectId: "customizer-project", + }).startSpan( + withSpanInstrumentationName( + { name: "provider.call" }, + INSTRUMENTATION_NAMES.OPENAI, + ), + ); + } + + test("adds a field to outgoing span records", async () => { + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + data.custom_field = "added"; + return data; + }, + }, + ], + }); + + const span = startInstrumentedSpan(); + span.log({ output: "result" }); + span.end(); + + const events = await memoryLogger.drain(); + expect(events).toEqual([ + expect.objectContaining({ + id: span.id, + project_id: "customizer-project", + output: "result", + custom_field: "added", + }), + ]); + }); + + test("alters an existing field in outgoing span records", async () => { + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + if ("output" in data) data.output = "[redacted]"; + return data; + }, + }, + ], + }); + + const span = startInstrumentedSpan(); + span.log({ output: "sensitive response" }); + span.end(); + + expect(await memoryLogger.drain()).toEqual([ + expect.objectContaining({ id: span.id, output: "[redacted]" }), + ]); + }); + + test("deletes a field from outgoing span records", async () => { + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + delete data.error; + return data; + }, + }, + ], + }); + + const span = startInstrumentedSpan(); + span.log({ error: "sensitive error", output: "safe response" }); + span.end(); + + const events = await memoryLogger.drain(); + expect(events).toEqual([ + expect.objectContaining({ id: span.id, output: "safe response" }), + ]); + expect(events[0]).not.toHaveProperty("error"); + }); + + test("passes replacement records through later customizers despite errors", async () => { + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + return "output" in data ? { ...data, output: "replacement" } : data; + }, + }, + { + onSpanExport() { + throw new Error("customizer failed"); + }, + }, + { + onSpanExport(data) { + if (typeof data.output === "string") { + data.output = data.output.toUpperCase(); + } + return data; + }, + }, + ], + }); + + const span = startInstrumentedSpan(); + span.log({ output: "original" }); + span.end(); + + expect(await memoryLogger.drain()).toEqual([ + expect.objectContaining({ + id: span.id, + output: "REPLACEMENT", + metrics: expect.objectContaining({ end: expect.any(Number) }), + }), + ]); + }); + + test("does not customize manually created spans", async () => { + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + data.tags = ["customized"]; + return data; + }, + }, + ], + }); + + const instrumented = startInstrumentedSpan(); + const manual = instrumented.startSpan({ name: "manual child" }); + manual.log({ output: "manual result" }); + manual.end(); + instrumented.end(); + + const events = await memoryLogger.drain(); + expect(events.find((event) => event.id === instrumented.id)).toMatchObject({ + tags: ["customized"], + }); + const manualEvent = events.find((event) => event.id === manual.id); + expect(manualEvent).toMatchObject({ output: "manual result" }); + expect(manualEvent).not.toHaveProperty("tags"); + }); + + test.each([ + ["missing return", () => undefined], + ["null", () => null], + ["array", () => []], + ["scalar", () => "invalid"], + ["non-record object", () => new Date(0)], + ])("ignores %s without losing either span", async (_name, invalidResult) => { + configureInstrumentation({ + spanCustomizers: [ + { + // @ts-expect-error Exercise invalid callback results from JavaScript. + onSpanExport(data) { + if ("output" in data) data.output = "redacted"; + return invalidResult(); + }, + }, + { + onSpanExport(data) { + if ("output" in data) data.output = `${data.output}:processed`; + return data; + }, + }, + ], + }); + const span = startInstrumentedSpan(); + const manual = span.startSpan({ name: "manual" }); + span.log({ input: "input", output: "private" }); + manual.log({ output: "unrelated" }); + manual.end(); + span.end(); + + const events = await memoryLogger.drain(); + expect(events).toHaveLength(2); + expect(events.find((event) => event.id === span.id)).toMatchObject({ + input: "input", + output: "redacted:processed", + metrics: { end: expect.any(Number) }, + }); + expect(events.find((event) => event.id === manual.id)).toMatchObject({ + output: "unrelated", + metrics: { end: expect.any(Number) }, + }); + }); + + test("restores mutable protocol fields between callbacks, even after a throw", () => { + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + if (Array.isArray(data.span_parents)) + data.span_parents.push("wrong"); + if (Array.isArray(data._merge_paths)) + data._merge_paths[0].push("wrong"); + delete data.id; + delete data.project_id; + data._is_merge = false; + data.dataset_id = "wrong"; + data._object_delete = true; + throw new Error("bad customizer"); + }, + }, + { + onSpanExport(data) { + // These fields feed the next callback, not just the final exporter. + data.output = { + id: data.id, + parents: data.span_parents, + paths: data._merge_paths, + merge: data._is_merge, + }; + return Object.freeze(data); + }, + }, + ], + }); + const result = customizeSpanExport({ + id: "original", + span_id: "span", + root_span_id: "root", + span_parents: ["parent"], + project_id: "project", + log_id: "g", + _is_merge: true, + _merge_paths: [["metadata"]], + }); + expect(result).toEqual({ + id: "original", + span_id: "span", + root_span_id: "root", + span_parents: ["parent"], + project_id: "project", + log_id: "g", + _is_merge: true, + _merge_paths: [["metadata"]], + output: { + id: "original", + parents: ["parent"], + paths: [["metadata"]], + merge: true, + }, + }); + }); + + test("exports a mixed HTTP batch despite async hooks and payload-only replacements", async () => { + configureInstrumentation({ + spanCustomizers: [ + { + // @ts-expect-error Async customizers are unsupported, but must be contained. + async onSpanExport(data) { + delete data.id; + delete data._is_merge; + if ("output" in data) data.output = "redacted"; + throw new Error("async customizer failed"); + }, + }, + { + onSpanExport(data) { + const payload = Object.fromEntries( + Object.entries(data).filter(([key]) => + ["input", "output", "metrics", "span_attributes"].includes(key), + ), + ); + return Object.freeze(payload); + }, + }, + ], + }); + + // Keep both spans in the same flush chunk, rather than auto-flushing the + // manual span's initial row before the instrumented child is created. + vi.stubEnv("BRAINTRUST_SYNC_FLUSH", "1"); + const rows: Record[] = []; + const state = new BraintrustState({ noExitFlush: true }); + const logger = initLogger({ + state, + projectName: "customizer-project", + projectId: "customizer-project", + appUrl: "https://customizer.test", + apiKey: "test-key", + orgName: "test-org", + asyncFlush: false, + fetch: async (url, options) => { + const pathname = new URL(String(url)).pathname; + if (pathname === "/api/apikey/login") { + return Response.json({ + org_info: [ + { + id: "test-org", + name: "test-org", + api_url: "https://customizer.test", + }, + ], + }); + } + if (pathname === "/version") return Response.json({}); + if (pathname === "/logs3") { + rows.push(...JSON.parse(String(options?.body)).rows); + return Response.json({}); + } + throw new Error(`Unexpected test request: ${pathname}`); + }, + }); + const manual = logger.startSpan({ + name: "manual", + event: { input: "manual input" }, + }); + const span = manual.startSpan( + withSpanInstrumentationName( + { name: "provider", event: { input: "provider input" } }, + INSTRUMENTATION_NAMES.OPENAI, + ), + ); + span.log({ output: "private" }); + span.end(); + manual.log({ output: "manual output" }); + manual.end(); + await logger.flush(); + + expect(rows).toHaveLength(2); + expect(rows.find((row) => row.id === span.id)).toMatchObject({ + span_id: span.spanId, + root_span_id: manual.rootSpanId, + span_parents: [manual.spanId], + project_id: "customizer-project", + log_id: "g", + input: "provider input", + output: "redacted", + metrics: { start: expect.any(Number), end: expect.any(Number) }, + }); + expect(rows.find((row) => row.id === manual.id)).toMatchObject({ + input: "manual input", + output: "manual output", + metrics: { start: expect.any(Number), end: expect.any(Number) }, + }); + }); +}); diff --git a/js/src/span-customizer.ts b/js/src/span-customizer.ts new file mode 100644 index 000000000..5990d077b --- /dev/null +++ b/js/src/span-customizer.ts @@ -0,0 +1,126 @@ +import { + ARRAY_DELETE_FIELD, + ID_FIELD, + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + OBJECT_DELETE_FIELD, + OBJECT_ID_KEYS, + PARENT_ID_FIELD, + TRANSACTION_ID_FIELD, +} from "../util/db_fields"; +import { isPromiseLike } from "../util/type_util"; +import type { SpanCustomizer, SpanExportData } from "./instrumentation/config"; + +// Configuration can precede platform initialization and must be shared across +// SDK bundles without importing the provider plugin registry into the logger. +const SPAN_CUSTOMIZERS_KEY = Symbol.for("braintrust.spanCustomizers"); +const shared: typeof globalThis & { + [SPAN_CUSTOMIZERS_KEY]?: readonly SpanCustomizer[]; +} = globalThis; + +const PROTECTED_FIELDS = new Set([ + ID_FIELD, + "span_id", + "root_span_id", + "span_parents", + "org_id", + ...OBJECT_ID_KEYS, + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + PARENT_ID_FIELD, + OBJECT_DELETE_FIELD, + ARRAY_DELETE_FIELD, + TRANSACTION_ID_FIELD, +]); + +function isPlainRecord(value: unknown): value is SpanExportData { + if (value === null || typeof value !== "object") return false; + const prototype = Object.getPrototypeOf(value); + return prototype === Object.prototype || prototype === null; +} + +// Only protocol values are copied deeply; payloads may contain SDK objects such +// as Attachments that must retain their identity and serialization behavior. +function copyProtocolValue(value: unknown): unknown { + if (Array.isArray(value)) return value.map(copyProtocolValue); + if (isPlainRecord(value)) { + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [ + key, + copyProtocolValue(item), + ]), + ); + } + return value; +} + +function restoreProtocolFields( + data: SpanExportData, + protectedFields: SpanExportData, +): SpanExportData { + const restored: SpanExportData = {}; + for (const key of Object.keys(data)) { + if (!PROTECTED_FIELDS.has(key)) { + Object.defineProperty(restored, key, { + value: data[key], + enumerable: true, + configurable: true, + writable: true, + }); + } + } + for (const key of Object.keys(protectedFields)) { + // Give each callback its own protocol arrays/objects, never the snapshot. + restored[key] = copyProtocolValue(protectedFields[key]); + } + return restored; +} + +export function setSpanCustomizers( + customizers: readonly SpanCustomizer[] | undefined, +): void { + shared[SPAN_CUSTOMIZERS_KEY] = customizers; +} + +export function customizeSpanExport(data: SpanExportData): SpanExportData { + const customizers = shared[SPAN_CUSTOMIZERS_KEY]; + if (!customizers?.length) return data; + + let protectedFields: SpanExportData | undefined; + + for (const customizer of customizers) { + let candidate = data; + try { + if (!customizer.onSpanExport) continue; + if (!protectedFields) { + protectedFields = {}; + for (const key of PROTECTED_FIELDS) { + if (Object.prototype.hasOwnProperty.call(data, key)) { + protectedFields[key] = copyProtocolValue(data[key]); + } + } + } + + const result: unknown = customizer.onSpanExport(data); + if (isPromiseLike(result)) { + // Hooks are synchronous, but accidental async hooks must not leak an + // unhandled rejection or replace the record with a promise. + void Promise.resolve(result).catch(() => {}); + } else if (isPlainRecord(result)) { + candidate = result; + } + } catch { + // Customization must not prevent export or later customizers from running. + } + + if (protectedFields) { + try { + // Always copy: hooks may freeze their input or return a frozen record. + data = restoreProtocolFields(candidate, protectedFields); + } catch { + data = restoreProtocolFields(data, protectedFields); + } + } + } + return data; +} From e85b4ab5f60709b551f810c2fe256c735f48689e Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 21 Sep 2026 11:06:39 -0600 Subject: [PATCH 2/2] use span customizers to implement the masking function --- js/src/logger.ts | 167 ++++++---------------------- js/src/masking.test.ts | 227 ++++++++++++++++---------------------- js/src/span-customizer.ts | 54 ++++++++- 3 files changed, 178 insertions(+), 270 deletions(-) diff --git a/js/src/logger.ts b/js/src/logger.ts index 0c656cacb..35e066770 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -207,7 +207,11 @@ import { mergeSpanOriginContext, type SpanOriginEnvironment, } from "./span-origin"; -import { customizeSpanExport } from "./span-customizer"; +import { + createMaskingCustomizer, + customizeSpanExport, +} from "./span-customizer"; +import type { SpanCustomizer, SpanExportData } from "./instrumentation/config"; // Manual type definition for inline attachments (not in generated_types) const InlineAttachmentReferenceSchema = z.object({ @@ -229,63 +233,6 @@ export class LoginInvalidOrgError extends Error { } } -// Fields that should be passed to the masking function -// Note: "tags" field is intentionally excluded, but can be added if needed -const REDACTION_FIELDS = [ - "input", - "output", - "expected", - "metadata", - "context", - "scores", - "metrics", -] as const; - -class MaskingError { - constructor( - public readonly fieldName: string, - public readonly errorType: string, - ) {} - - get errorMsg(): string { - return `ERROR: Failed to mask field '${this.fieldName}' - ${this.errorType}`; - } -} - -/** - * Apply masking function to data and handle errors gracefully. - * If the masking function raises an exception, returns an error message. - * Returns MaskingError for scores/metrics fields to signal they should be dropped. - */ -function applyMaskingToField( - maskingFunction: (value: unknown) => unknown, - data: unknown, - fieldName: string, -): unknown { - try { - return maskingFunction(data); - } catch (error) { - // Return a generic error message without the stack trace to avoid leaking PII - const errorType = error instanceof Error ? error.constructor.name : "Error"; - - // For scores and metrics fields, return a special error object - // to signal the field should be dropped and error logged - if (fieldName === "scores" || fieldName === "metrics") { - return new MaskingError(fieldName, errorType); - } - - // For metadata field that expects object type, return an object with error key - if (fieldName === "metadata") { - return { - error: `ERROR: Failed to mask field '${fieldName}' - ${errorType}`, - }; - } - - // For other fields, return the error message as a string - return `ERROR: Failed to mask field '${fieldName}' - ${errorType}`; - } -} - export type SetCurrentArg = { setCurrent?: boolean }; type StartSpanEventArgs = ExperimentLogPartialArgs & Partial; @@ -3221,7 +3168,7 @@ interface BackgroundLogger { export class TestBackgroundLogger implements BackgroundLogger { private items: LazyValue[][] = []; - private maskingFunction: ((value: unknown) => unknown) | null = null; + private exportCustomizers: readonly SpanCustomizer[] = []; log(items: LazyValue[]): void { this.items.push(items); @@ -3230,7 +3177,9 @@ export class TestBackgroundLogger implements BackgroundLogger { setMaskingFunction( maskingFunction: ((value: unknown) => unknown) | null, ): void { - this.maskingFunction = maskingFunction; + this.exportCustomizers = maskingFunction + ? [createMaskingCustomizer(maskingFunction)] + : []; } async flush(): Promise { @@ -3259,43 +3208,14 @@ export class TestBackgroundLogger implements BackgroundLogger { let batch = mergeRowBatch(events); - // Apply masking after merge, similar to HTTPBackgroundLogger - if (this.maskingFunction) { - batch = batch.map((item) => { - const maskedItem = { ...item }; - - // Only mask specific fields if they exist - for (const field of REDACTION_FIELDS) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - if ((item as any)[field] !== undefined) { - const maskedValue = applyMaskingToField( - this.maskingFunction!, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (item as any)[field], - field, - ); - if (maskedValue instanceof MaskingError) { - // Drop the field and add error message - // eslint-disable-next-line @typescript-eslint/no-explicit-any - delete (maskedItem as any)[field]; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - if ((maskedItem as any).error) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (maskedItem as any).error = - `${(maskedItem as any).error}; ${maskedValue.errorMsg}`; - } else { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (maskedItem as any).error = maskedValue.errorMsg; - } - } else { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (maskedItem as any)[field] = maskedValue; - } - } - } - - return maskedItem as BackgroundLogEvent; - }); + if (this.exportCustomizers.length) { + batch = batch.map( + (item) => + customizeSpanExport( + item as SpanExportData, + this.exportCustomizers, + ) as BackgroundLogEvent, + ); } return batch; @@ -3397,7 +3317,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { private activeFlush: Promise = Promise.resolve(); private activeFlushResolved = true; private onFlushError?: (error: unknown) => void; - private maskingFunction: ((value: unknown) => unknown) | null = null; + private exportCustomizers: readonly SpanCustomizer[] = []; private readonly requestLimiter: ConcurrencyLimiter; private lastEnqueuedSequence = 0; private completedSequence = 0; @@ -3531,7 +3451,9 @@ class HTTPBackgroundLogger implements BackgroundLogger { setMaskingFunction( maskingFunction: ((value: unknown) => unknown) | null, ): void { - this.maskingFunction = maskingFunction; + this.exportCustomizers = maskingFunction + ? [createMaskingCustomizer(maskingFunction)] + : []; } pendingFlushBytes(): number { @@ -3733,43 +3655,14 @@ class HTTPBackgroundLogger implements BackgroundLogger { let mergedItems = mergeRowBatch(items); - // Apply masking after merge but before sending to backend - if (this.maskingFunction) { - mergedItems = mergedItems.map((item) => { - const maskedItem = { ...item }; - - // Only mask specific fields if they exist - for (const field of REDACTION_FIELDS) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - if ((item as any)[field] !== undefined) { - const maskedValue = applyMaskingToField( - this.maskingFunction!, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (item as any)[field], - field, - ); - if (maskedValue instanceof MaskingError) { - // Drop the field and add error message - // eslint-disable-next-line @typescript-eslint/no-explicit-any - delete (maskedItem as any)[field]; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - if ((maskedItem as any).error) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (maskedItem as any).error = - `${(maskedItem as any).error}; ${maskedValue.errorMsg}`; - } else { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (maskedItem as any).error = maskedValue.errorMsg; - } - } else { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (maskedItem as any)[field] = maskedValue; - } - } - } - - return maskedItem as BackgroundLogEvent; - }); + if (this.exportCustomizers.length) { + mergedItems = mergedItems.map( + (item) => + customizeSpanExport( + item as SpanExportData, + this.exportCustomizers, + ) as BackgroundLogEvent, + ); } return [mergedItems, attachments]; @@ -5678,6 +5571,8 @@ export type FullLoginOptions = LoginOptions & { /** * Set a global masking function that will be applied to all logged data before sending to Braintrust. * The masking function will be applied after records are merged but before they are sent to the backend. + * Internally, masking is a state-local export customizer that runs after any + * instrumentation customizers and also covers manually logged records. * * @param maskingFunction A function that takes a JSON-serializable object and returns a masked version. * Set to null to disable masking. diff --git a/js/src/masking.test.ts b/js/src/masking.test.ts index b11605046..5667a7144 100644 --- a/js/src/masking.test.ts +++ b/js/src/masking.test.ts @@ -7,18 +7,25 @@ import { setMaskingFunction, } from "./logger"; import { configureNode } from "./node/config"; +import { configureInstrumentation, registry } from "./instrumentation/registry"; +import { + INSTRUMENTATION_NAMES, + withSpanInstrumentationName, +} from "./span-origin"; configureNode(); describe("masking functionality", () => { let memoryLogger: any; - beforeEach(() => { - _exportsForTestingOnly.simulateLoginForTests(); + beforeEach(async () => { + registry.disable(); + await _exportsForTestingOnly.simulateLoginForTests(); memoryLogger = _exportsForTestingOnly.useTestBackgroundLogger(); }); afterEach(() => { + configureInstrumentation({ spanCustomizers: [] }); setMaskingFunction(null); // Clear masking function _exportsForTestingOnly.clearTestBackgroundLogger(); }); @@ -413,7 +420,7 @@ describe("masking functionality", () => { expect(events).toHaveLength(2); - // First event should be masked (masking was applied at flush time) + // Disabling before flush also affects previously queued records. expect(events[0].input).toEqual({ password: "visible1" }); // Second event should not be masked expect(events[1].input).toEqual({ password: "visible2" }); @@ -478,151 +485,107 @@ describe("masking functionality", () => { expect(event.input.normal_number).toBe(123); }); - test("masking function with error", async () => { - const brokenMaskingFunction = (data: any): any => { - if (typeof data === "object" && data !== null) { - if (data.password) { - // Simulate an error when trying to mask a sensitive field - throw new Error( - "Cannot mask sensitive field 'password' - internal masking error", - ); - } - if (data.accuracy !== undefined) { - // Trigger error for scores field - throw new TypeError("Cannot process numeric score"); - } - - const masked: any = Array.isArray(data) ? [] : {}; - for (const [key, value] of Object.entries(data)) { - if (key === "secret" && typeof value === "string") { - // Another type of error - 1 / 0; // This will be Infinity, not an error - throw new Error("Division by zero error"); - } else if (key === "complex" && Array.isArray(value)) { - // Try to access non-existent index - const item = value[100]; - if (!item) { - throw new RangeError("Index out of bounds"); - } - } else if (typeof value === "object") { - masked[key] = brokenMaskingFunction(value); - } else { - masked[key] = value; - } - } - return masked; - } - return data; - }; - - setMaskingFunction(brokenMaskingFunction); + test("masking failures redact fields without leaking exception details", async () => { + setMaskingFunction((data) => { + if (data === "safe output") return data; + throw new TypeError("private exception detail"); + }); const logger = initLogger({ projectName: "test", projectId: "test-project-id", }); - - // Test various error scenarios logger.log({ - input: { query: "login", password: "secret123" }, - output: { status: "success" }, - metadata: { safe: "no-error" }, - }); + input: { password: "private input" }, + output: "safe output", + expected: "private expected", + metadata: { token: "private metadata" }, + scores: { accuracy: 0.85 }, + metrics: { accuracy: 0.95 }, + error: "existing application error", + }); + + const [event] = await memoryLogger.drain(); + expect(event.input).toEqual(expect.any(String)); + expect(event.expected).toEqual(expect.any(String)); + expect(event.metadata).toEqual({ error: expect.any(String) }); + expect(event.output).toBe("safe output"); + expect(event.scores).toBeUndefined(); + expect(event.metrics).toBeUndefined(); + expect(event.error).toContain("existing application error"); + expect(event.error).toContain("scores"); + expect(event.error).toContain("metrics"); + expect(JSON.stringify(event)).not.toContain("private"); + }); - logger.log({ - input: { data: "safe", secret: "will-cause-error" }, - output: { result: "ok" }, + test("late masking sees merged fields on manual spans", async () => { + const logger = initLogger({ + projectName: "test", + projectId: "test-project-id", }); - - logger.log({ - input: { complex: ["a", "b"], other: "data" }, - expected: { values: ["x", "y", "z"] }, + const span = logger.startSpan({ + name: "manual", + event: { metadata: { password: "private" } }, }); + span.log({ metadata: { redact: true } }); + span.end(); - await memoryLogger.flush(); - const events = await memoryLogger.drain(); - - expect(events).toHaveLength(3); - - // First event - error when masking input.password - const event1 = events[0]; - expect(event1.input).toBe("ERROR: Failed to mask field 'input' - Error"); - expect(event1.output).toEqual({ status: "success" }); - expect(event1.metadata).toEqual({ safe: "no-error" }); - - // Second event - error when masking input.secret - const event2 = events[1]; - expect(event2.input).toBe("ERROR: Failed to mask field 'input' - Error"); - expect(event2.output).toEqual({ result: "ok" }); - - // Third event - error when masking input.complex - const event3 = events[2]; - expect(event3.input).toBe( - "ERROR: Failed to mask field 'input' - RangeError", - ); - expect(event3.expected).toEqual({ values: ["x", "y", "z"] }); - - // Test with a score that triggers an error - logger.log({ - input: { data: "test" }, - scores: { accuracy: 0.95 }, // Will trigger error + setMaskingFunction((data) => { + if (data && typeof data === "object" && "redact" in data) { + return { ...data, password: "redacted" }; + } + return data; }); - await memoryLogger.flush(); - const events2 = await memoryLogger.drain(); - - // Should include the new event - expect(events2).toHaveLength(1); - const scoreEvent = events2[0]; - - // Scores should be dropped and error should be logged - expect(scoreEvent.scores).toBeUndefined(); - expect(scoreEvent.error).toBe( - "ERROR: Failed to mask field 'scores' - TypeError", - ); + expect(await memoryLogger.drain()).toEqual([ + expect.objectContaining({ + id: span.id, + project_id: "test-project-id", + metadata: { password: "redacted", redact: true }, + }), + ]); + }); - // Test with metrics that triggers an error - logger.log({ - input: { data: "test2" }, - output: "result2", - metrics: { accuracy: 0.95 }, // Will trigger error + test("masking runs after customizers and can be replaced or disabled independently", async () => { + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + if (typeof data.output === "string") data.output += ":customized"; + return data; + }, + }, + ], }); - - await memoryLogger.flush(); - const events3 = await memoryLogger.drain(); - - expect(events3).toHaveLength(1); - const metricsEvent = events3[0]; - - // Metrics should be dropped and error should be logged - expect(metricsEvent.metrics).toBeUndefined(); - expect(metricsEvent.error).toBe( - "ERROR: Failed to mask field 'metrics' - TypeError", - ); - - // Test with both scores and metrics failing - logger.log({ - input: { data: "test3" }, - output: "result3", - scores: { accuracy: 0.85 }, // Will trigger error - metrics: { accuracy: 0.95 }, // Will also trigger error + const logger = initLogger({ + projectName: "test", + projectId: "test-project-id", }); - - await memoryLogger.flush(); - const events4 = await memoryLogger.drain(); - - expect(events4).toHaveLength(1); - const bothEvent = events4[0]; - - // Both should be dropped and errors should be concatenated - expect(bothEvent.scores).toBeUndefined(); - expect(bothEvent.metrics).toBeUndefined(); - expect(bothEvent.error).toContain( - "ERROR: Failed to mask field 'scores' - TypeError", + const logOutput = () => { + const span = logger.startSpan( + withSpanInstrumentationName( + { name: "provider" }, + INSTRUMENTATION_NAMES.OPENAI, + ), + ); + span.log({ output: "private" }); + span.end(); + }; + setMaskingFunction((data) => + typeof data === "string" ? `${data}:obsolete` : data, ); - expect(bothEvent.error).toContain( - "ERROR: Failed to mask field 'metrics' - TypeError", + setMaskingFunction((data) => + data === "private:customized" ? "redacted" : data, ); - expect(bothEvent.error).toContain("; "); // Check that errors are joined + logOutput(); + expect(await memoryLogger.drain()).toEqual([ + expect.objectContaining({ output: "redacted" }), + ]); + + setMaskingFunction(null); + logOutput(); + expect(await memoryLogger.drain()).toEqual([ + expect.objectContaining({ output: "private:customized" }), + ]); }); }); diff --git a/js/src/span-customizer.ts b/js/src/span-customizer.ts index 5990d077b..411f399c8 100644 --- a/js/src/span-customizer.ts +++ b/js/src/span-customizer.ts @@ -33,6 +33,52 @@ const PROTECTED_FIELDS = new Set([ TRANSACTION_ID_FIELD, ]); +// Tags and record-level errors are intentionally outside the masking contract. +const MASKING_FIELDS = [ + "input", + "output", + "expected", + "metadata", + "context", + "scores", + "metrics", +] as const; + +/** + * Adapt field-level masking to a record customizer. Unlike instrumentation + * customizers, this runs on all merged records and belongs to one logger state. + */ +export function createMaskingCustomizer( + maskingFunction: (value: unknown) => unknown, +): SpanCustomizer { + return { + onSpanExport(data) { + const masked = { ...data }; + for (const field of MASKING_FIELDS) { + if (data[field] === undefined) continue; + try { + masked[field] = maskingFunction(data[field]); + } catch (error) { + // Fail closed without including exception messages or stacks, which + // can themselves contain sensitive data. + const errorType = + error instanceof Error ? error.constructor.name : "Error"; + const message = `ERROR: Failed to mask field '${field}' - ${errorType}`; + if (field === "scores" || field === "metrics") { + delete masked[field]; + masked.error = masked.error + ? `${masked.error}; ${message}` + : message; + } else { + masked[field] = field === "metadata" ? { error: message } : message; + } + } + } + return masked; + }, + }; +} + function isPlainRecord(value: unknown): value is SpanExportData { if (value === null || typeof value !== "object") return false; const prototype = Object.getPrototypeOf(value); @@ -82,8 +128,12 @@ export function setSpanCustomizers( shared[SPAN_CUSTOMIZERS_KEY] = customizers; } -export function customizeSpanExport(data: SpanExportData): SpanExportData { - const customizers = shared[SPAN_CUSTOMIZERS_KEY]; +export function customizeSpanExport( + data: SpanExportData, + customizers: readonly SpanCustomizer[] | undefined = shared[ + SPAN_CUSTOMIZERS_KEY + ], +): SpanExportData { if (!customizers?.length) return data; let protectedFields: SpanExportData | undefined;