diff --git a/.changeset/fail-closed-span-customizers.md b/.changeset/fail-closed-span-customizers.md new file mode 100644 index 000000000..14c8cef5b --- /dev/null +++ b/.changeset/fail-closed-span-customizers.md @@ -0,0 +1,7 @@ +--- +"braintrust": patch +--- + +fix: fail closed on span customizer errors without dropping records + +Preserve original identity, routing, and duration timestamps with `context.customizer_error: true`; strip all other content and prevent attachment uploads. Log safely at most once across SDK bundles, reuse stripped records on retries, and isolate them from merging and masking in ordered export requests. diff --git a/js/src/instrumentation/config.ts b/js/src/instrumentation/config.ts index 057191b68..93d08d96c 100644 --- a/js/src/instrumentation/config.ts +++ b/js/src/instrumentation/config.ts @@ -10,10 +10,12 @@ export interface SpanCustomizer { * 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. - * Throwing or returning an invalid value (including a promise) logs a safe - * error, stops the callback chain, and drops this outgoing record. Later - * records are evaluated independently; already exported records cannot be - * retracted. Handle recovery inside the callback if export should continue. + * Throwing or returning an invalid value (including a promise) stops the chain + * and exports only original identity, routing, and duration timestamps, plus + * context: { customizer_error: true }. No attachments from that record upload. + * A safe error is logged at most once per JavaScript global environment. + * Later records are evaluated independently; already exported records cannot + * be retracted. Handle recovery inside the callback to retain span content. */ onSpanExport?(data: SpanExportData): SpanExportData; } diff --git a/js/src/logger.ts b/js/src/logger.ts index 3a5be45ca..c086a0ad5 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -47,7 +47,6 @@ import { LogFeedbackFullArgs, mergeDicts, mergeGitMetadataSettings, - mergeRowBatch, OBJECT_DELETE_FIELD, OBJECT_ID_KEYS, SanitizedExperimentLogPartialArgs, @@ -210,6 +209,7 @@ import { import { createMaskingCustomizer, customizeSpanExport, + mergeCustomizedSpanRecords, } from "./span-customizer"; import type { SpanCustomizer, SpanExportData } from "./instrumentation/config"; @@ -2504,7 +2504,7 @@ function updateSpanImpl({ ...updateEvent, ...(await parentIds()), [IS_MERGE_FIELD]: true, - }) as BackgroundLogEvent | null, + }) as BackgroundLogEvent, ); state.bgLogger().log([record]); } @@ -3273,7 +3273,7 @@ function getNumericEnv(name: string): number | undefined { } interface BackgroundLogger { - log(items: LazyValue[]): void; + log(items: LazyValue[]): void; flush(): Promise; pendingFlushBytes(): number; flushBackpressureBytes(): number; @@ -3283,10 +3283,10 @@ interface BackgroundLogger { } export class TestBackgroundLogger implements BackgroundLogger { - private items: LazyValue[][] = []; + private items: LazyValue[][] = []; private exportCustomizers: readonly SpanCustomizer[] = []; - log(items: LazyValue[]): void { + log(items: LazyValue[]): void { this.items.push(items); } @@ -3318,23 +3318,20 @@ export class TestBackgroundLogger implements BackgroundLogger { const events: BackgroundLogEvent[] = []; for (const item of items) { for (const event of item) { - const value = await event.get(); - if (value !== null) events.push(value); + events.push(await event.get()); } } - let batch = mergeRowBatch(events); + let batch = mergeCustomizedSpanRecords(events).flat(); if (this.exportCustomizers.length) { - batch = batch - .map( - (item) => - customizeSpanExport( - item as SpanExportData, - this.exportCustomizers, - ) as BackgroundLogEvent | null, - ) - .filter((item): item is BackgroundLogEvent => item !== null); + batch = batch.map( + (item) => + customizeSpanExport( + item as SpanExportData, + this.exportCustomizers, + ) as BackgroundLogEvent, + ); } return batch; @@ -3427,7 +3424,7 @@ class ConcurrencyLimiter { interface QueuedLogEvent { sequence: number; - event: LazyValue; + event: LazyValue; } class HTTPBackgroundLogger implements BackgroundLogger { @@ -3581,7 +3578,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { return this._flushBackpressureBytes; } - log(items: LazyValue[]) { + log(items: LazyValue[]) { if (this._disabled) { return; } @@ -3682,7 +3679,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { } private async flushWrappedItemsChunk( - wrappedItems: LazyValue[], + wrappedItems: LazyValue[], batchSize: number, ) { if (!wrappedItems.length) { @@ -3690,38 +3687,46 @@ class HTTPBackgroundLogger implements BackgroundLogger { } const [allItems, attachments] = await this.unwrapLazyValues(wrappedItems); - if (allItems.length === 0) { - return; - } - // Construct batches of records to flush in parallel. + // Failure barriers must be sent in order, in separate requests. Within each + // group, retain the normal parallel transport batching. let chunkBytes = 0; - const allItemsWithMeta = allItems.map((item) => { - const withMeta = stringifyWithOverflowMeta(item); - chunkBytes += withMeta.str.length; - return withMeta; - }); + const groupsWithMeta = allItems.map((group) => + group.map((item) => { + const withMeta = stringifyWithOverflowMeta(item); + chunkBytes += withMeta.str.length; + return withMeta; + }), + ); this._pendingBytes += chunkBytes; const maxRequestSizeResult = await this.getMaxRequestSize(); - const batches = batchItems({ - items: allItemsWithMeta, - batchMaxNumItems: batchSize, - batchMaxNumBytes: maxRequestSizeResult.maxRequestSize / 2, - getByteSize: (item) => item.str.length, - }); - - const postPromises = batches.map((batch) => - this.requestLimiter.run(async () => { - try { - await this.submitLogsRequest(batch, maxRequestSizeResult); - return { type: "success" } as const; - } catch (e) { - return { type: "error", value: e } as const; - } - }), - ); - const results = await Promise.all(postPromises); + const results: ({ type: "success" } | { type: "error"; value: unknown })[] = + []; + for (const items of groupsWithMeta) { + // Later groups may depend on rows in a failed group, so stop sending. + if (results.some((r) => r.type === "error")) break; + const batches = batchItems({ + items, + batchMaxNumItems: batchSize, + batchMaxNumBytes: maxRequestSizeResult.maxRequestSize / 2, + getByteSize: (item) => item.str.length, + }); + results.push( + ...(await Promise.all( + batches.map((batch) => + this.requestLimiter.run(async () => { + try { + await this.submitLogsRequest(batch, maxRequestSizeResult); + return { type: "success" } as const; + } catch (e) { + return { type: "error", value: e } as const; + } + }), + ), + )), + ); + } this._pendingBytes = Math.max(0, this._pendingBytes - chunkBytes); const failingResultErrors = results .map((r) => (r.type === "success" ? undefined : r.value)) @@ -3756,34 +3761,45 @@ class HTTPBackgroundLogger implements BackgroundLogger { } private async unwrapLazyValues( - wrappedItems: LazyValue[], - ): Promise<[BackgroundLogEvent[], Attachment[]]> { + wrappedItems: LazyValue[], + ): Promise<[BackgroundLogEvent[][], Attachment[]]> { for (let i = 0; i < this.numTries; ++i) { try { - const items = ( - await Promise.all(wrappedItems.map((x) => x.get())) - ).filter((item): item is BackgroundLogEvent => item !== null); + const items = await Promise.all(wrappedItems.map((x) => x.get())); // TODO(kevin): `extractAttachments` should ideally come after // `mergeRowBatch`, since merge-overwriting could result in some // attachments that no longer need to be uploaded. This would require // modifying the merge logic to treat the Attachment object as a "leaf" // rather than attempting to merge the keys of the Attachment. - const attachments: Attachment[] = []; + let attachments: Attachment[] = []; items.forEach((item) => extractAttachments(item, attachments)); - let mergedItems = mergeRowBatch(items); + let mergedItems = mergeCustomizedSpanRecords(items); if (this.exportCustomizers.length) { - mergedItems = mergedItems - .map( + mergedItems = mergedItems.map((group) => + group.map( (item) => customizeSpanExport( item as SpanExportData, this.exportCustomizers, - ) as BackgroundLogEvent | null, - ) - .filter((item): item is BackgroundLogEvent => item !== null); + ) as BackgroundLogEvent, + ), + ); + // Customizers may strip or replace attachment references; don't upload + // content that is no longer referenced by any exported record. + const referencedKeys = new Set(); + mergedItems.forEach((group) => + group.forEach((item) => + collectAttachmentReferenceKeys(item, referencedKeys), + ), + ); + attachments = attachments.filter( + (attachment) => + attachment.reference.type !== BRAINTRUST_ATTACHMENT || + referencedKeys.has(attachment.reference.key), + ); } return [mergedItems, attachments]; @@ -3968,7 +3984,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { } private async dumpDroppedEvents( - wrappedItems: LazyValue[], + wrappedItems: LazyValue[], ) { const publishPayloadsDir = [ this.allPublishPayloadsDir, @@ -3982,7 +3998,9 @@ class HTTPBackgroundLogger implements BackgroundLogger { await this.unwrapLazyValues(wrappedItems); const dataStr = constructLogs3Data( - allItems.map((x) => stringifyWithOverflowMeta(x)), + allItems.flatMap((group) => + group.map((x) => stringifyWithOverflowMeta(x)), + ), ); const attachmentStr = JSON.stringify( allAttachments.map((a) => a.debugInfo()), @@ -7470,6 +7488,22 @@ function extractAttachments( } } +function collectAttachmentReferenceKeys(value: unknown, keys: Set) { + if (!value || typeof value !== "object") return; + if (Array.isArray(value)) { + value.forEach((item) => collectAttachmentReferenceKeys(item, keys)); + return; + } + const record = value as Record; + if (record.type === BRAINTRUST_ATTACHMENT && typeof record.key === "string") { + keys.add(record.key); + return; + } + Object.values(record).forEach((item) => + collectAttachmentReferenceKeys(item, keys), + ); +} + /** * Recursively hydrates any `AttachmentReference` into `Attachment` by modifying * the input in-place. @@ -8536,8 +8570,8 @@ export class SpanImpl implements Span { }).objectIdFields(), }; // Customize inside the memoized lazy value, before attachment processing, - // merging, and masking. Retries reuse the customized record or drop result. - return customizeSpanExport(record) as BackgroundLogEvent | null; + // merging, and masking. Retries reuse the customized or stripped record. + return customizeSpanExport(record) as BackgroundLogEvent; }; this._state.bgLogger().log([new LazyValue(computeRecord)]); } diff --git a/js/src/masking.test.ts b/js/src/masking.test.ts index 12d6cd018..abed47e0b 100644 --- a/js/src/masking.test.ts +++ b/js/src/masking.test.ts @@ -547,7 +547,7 @@ describe("masking functionality", () => { ]); }); - test("dropped customizer records never reach masking or merge into successful records", async () => { + test("stripped customizer records bypass masking and remain independent of successful records", async () => { vi.spyOn(console, "error").mockImplementation(() => {}); configureInstrumentation({ spanCustomizers: [ @@ -559,22 +559,49 @@ describe("masking functionality", () => { }, ], }); - const mask = vi.fn((data: unknown) => (data === "safe" ? "masked" : data)); + const mask = vi.fn((data: unknown) => { + if (data && typeof data === "object" && "customizer_error" in data) { + return { injected: "must not replace SDK diagnostics" }; + } + if (data && typeof data === "object" && "start" in data) { + return { start: 999, end: 999 }; + } + return data === "safe" ? "masked" : data; + }); setMaskingFunction(mask); const span = initLogger({ projectName: "test", projectId: "test-project-id", }).startSpan({ name: "manual" }); - span.log({ output: "private", metadata: { secret: "private" } }); + span.log({ + output: "private", + metadata: { secret: "private" }, + metrics: { start: 10, end: 20, tokens: 99 }, + }); span.log({ output: "safe" }); span.end(); const events = await memoryLogger.drain(); expect(events).toEqual([ + expect.objectContaining({ + id: span.id, + span_attributes: expect.anything(), + }), + { + id: span.id, + span_id: span.spanId, + root_span_id: span.rootSpanId, + project_id: "test-project-id", + log_id: "g", + _is_merge: true, + context: { customizer_error: true }, + metrics: { start: 10, end: 20 }, + }, expect.objectContaining({ id: span.id, output: "masked" }), ]); expect(JSON.stringify(events)).not.toContain("private"); expect(JSON.stringify(mask.mock.calls)).not.toContain("private"); + expect(JSON.stringify(mask.mock.calls)).not.toContain("customizer_error"); }); test.each(["manual", "instrumented"])( diff --git a/js/src/span-customizer.test.ts b/js/src/span-customizer.test.ts index c9edce3a7..8caccde6f 100644 --- a/js/src/span-customizer.test.ts +++ b/js/src/span-customizer.test.ts @@ -28,6 +28,7 @@ import vm from "node:vm"; import type { SpanCustomizer, SpanExportData } from "./exports"; import { customizeSpanExport, + mergeCustomizedSpanRecords, resetSpanCustomizerFailureLoggingForTests, } from "./span-customizer"; import { SpanCache } from "./span-cache"; @@ -399,9 +400,9 @@ describe("onSpanExport", () => { const attachmentUpload = vi.spyOn(attachment, "upload"); const enqueue = vi.spyOn(memoryLogger, "log"); update({ output: "private", input: attachment }); - const droppedRecords = enqueue.mock.calls.flatMap(([items]) => items); + const strippedRecords = enqueue.mock.calls.flatMap(([items]) => items); - // A rejected update cannot overwrite the preceding safe update. + // A failure stays separate from preceding and subsequent safe updates. const events = await memoryLogger.drain(); expect(events).toEqual([ { @@ -416,20 +417,28 @@ describe("onSpanExport", () => { output: "[redacted]", metadata: { safe: true }, }, + { + id: span.id, + ...(entrypoint === "experiment" + ? { experiment_id: "customizer-experiment" } + : { + project_id: "customizer-project", + log_id: "g", + span_id: span.spanId, + root_span_id: span.rootSpanId, + }), + _is_merge: true, + context: { customizer_error: true }, + }, ]); expect(onSpanExport).toHaveBeenCalledTimes(2); expect(errorLog).toHaveBeenCalledTimes(1); expect(JSON.stringify(errorLog.mock.calls)).not.toContain("private"); - const fetch = vi.fn(); - const state = new BraintrustState({ noExitFlush: true, fetch }); - const httpLogger = state.httpLogger(); - httpLogger.syncFlush = true; - httpLogger.log(droppedRecords); - await httpLogger.flush(); - httpLogger.log(droppedRecords); - await httpLogger.flush(); - expect(fetch).not.toHaveBeenCalled(); + for (const record of strippedRecords) { + expect(await record.get()).toEqual(events[1]); + expect(await record.get()).toEqual(events[1]); + } expect(attachmentUpload).not.toHaveBeenCalled(); expect(onSpanExport).toHaveBeenCalledTimes(2); expect(errorLog).toHaveBeenCalledTimes(1); @@ -473,8 +482,11 @@ describe("onSpanExport", () => { ["non-record object", () => new Date(0)], ["class instance", () => new (class Replacement {})()], ["inherited object", () => Object.create({ output: "inherited" })], + ["custom null-root prototype", () => Object.create(Object.create(null))], + ["resolved promise", () => Promise.resolve({ output: "private" })], + ["rejected promise", () => Promise.reject(new Error("private rejection"))], ])( - "drops the failing record on %s and stops its callback chain", + "strips the failing record on %s and stops its callback chain", async (_name, invalidResult) => { const errorLog = vi.spyOn(console, "error").mockImplementation(() => {}); const laterHook = vi.fn((data: SpanExportData) => data); @@ -497,10 +509,16 @@ describe("onSpanExport", () => { manual.end(); const events = await memoryLogger.drain(); - expect(events).toHaveLength(2); - expect(events.find((event) => event.id === span.id)).not.toHaveProperty( - "output", - ); + expect(events).toHaveLength(3); + expect(events[1]).toEqual({ + id: span.id, + span_id: span.spanId, + root_span_id: span.rootSpanId, + project_id: "customizer-project", + log_id: "g", + _is_merge: true, + context: { customizer_error: true }, + }); expect(events.find((event) => event.id === manual.id)).toMatchObject({ output: "unrelated", metrics: { end: expect.any(Number) }, @@ -514,7 +532,7 @@ describe("onSpanExport", () => { errorLog.mock.calls[0].every((value) => typeof value === "string"), ).toBe(true); - // Dropping one incremental record does not disable its logical span. + // Stripping one incremental record does not disable its logical span. span.log({ output: "recovered" }); span.end(); expect(await memoryLogger.drain()).toEqual([ @@ -524,6 +542,86 @@ describe("onSpanExport", () => { }, ); + test.each([{}, { start: 10 }, { end: 20 }, { start: 10, end: 20 }])( + "retains only original protocol fields and duration %j after earlier mutations", + (duration) => { + vi.spyOn(console, "error").mockImplementation(() => {}); + const laterHook = vi.fn((data: SpanExportData) => data); + const original = { + id: "row", + span_id: "span", + root_span_id: "root", + span_parents: ["parent"], + org_id: "org", + project_id: "project", + log_id: "g", + _is_merge: true, + _merge_paths: [["metadata"]], + _parent_id: "parent-row", + _xact_id: "transaction", + _object_delete: false, + _array_delete: [], + created: "private", + name: "private", + span_attributes: { name: "private" }, + input: "private", + output: "private", + expected: "private", + metadata: { secret: "private" }, + context: { secret: "private" }, + attributes: { secret: "private" }, + events: ["private"], + links: ["private"], + error: "private", + tags: ["private"], + scores: { secret: 1 }, + attachments: ["private"], + unknown: "private", + metrics: { ...duration, tokens: 99, custom: 42 }, + }; + const result = customizeSpanExport(original, [ + { + onSpanExport(data) { + data.id = "wrong"; + (data.span_parents as string[]).push("wrong"); + (data._merge_paths as string[][])[0].push("wrong"); + data.metrics = { start: 999, end: 999, tokens: 999 }; + data.context = { secret: "earlier hook" }; + data.output = "earlier hook"; + return data; + }, + }, + { + onSpanExport(data) { + delete data.root_span_id; + (data.metrics as SpanExportData).end = 1234; + throw new Error("private failure"); + }, + }, + { onSpanExport: laterHook }, + ]); + expect(result).toEqual({ + id: "row", + span_id: "span", + root_span_id: "root", + span_parents: ["parent"], + org_id: "org", + project_id: "project", + log_id: "g", + _is_merge: true, + _merge_paths: [["metadata"]], + _parent_id: "parent-row", + _xact_id: "transaction", + _object_delete: false, + _array_delete: [], + ...(Object.keys(duration).length ? { metrics: duration } : {}), + context: { customizer_error: true }, + }); + expect(original.metrics).toEqual({ ...duration, tokens: 99, custom: 42 }); + expect(laterHook).not.toHaveBeenCalled(); + }, + ); + test("restores mutable protocol fields between successful callbacks", () => { configureInstrumentation({ spanCustomizers: [ @@ -583,7 +681,7 @@ describe("onSpanExport", () => { }); }); - test("memoizes dropped records across repeated flushes without uploading empty batches", async () => { + test("memoizes stripped records across repeated evaluations", async () => { const errorLog = vi.spyOn(console, "error").mockImplementation(() => {}); const onSpanExport = vi.fn(() => { throw new Error("private failure"); @@ -593,24 +691,26 @@ describe("onSpanExport", () => { spanCustomizers: [{ onSpanExport }, { onSpanExport: laterHook }], }); const enqueue = vi.spyOn(memoryLogger, "log"); - startInstrumentedSpan(); + const span = startInstrumentedSpan(); const records = enqueue.mock.calls.flatMap(([items]) => items); - expect(await memoryLogger.drain()).toEqual([]); - const failureCount = onSpanExport.mock.calls.length; - expect(failureCount).toBeGreaterThan(0); - expect(errorLog).toHaveBeenCalledTimes(1); - - const fetch = vi.fn(); - const state = new BraintrustState({ noExitFlush: true, fetch }); - const httpLogger = state.httpLogger(); - httpLogger.syncFlush = true; - httpLogger.log(records); - await httpLogger.flush(); - httpLogger.log(records); - await httpLogger.flush(); - - expect(fetch).not.toHaveBeenCalled(); - expect(onSpanExport).toHaveBeenCalledTimes(failureCount); + const events = await memoryLogger.drain(); + expect(events).toEqual([ + { + id: span.id, + span_id: span.spanId, + root_span_id: span.rootSpanId, + project_id: "customizer-project", + log_id: "g", + _is_merge: false, + metrics: { start: expect.any(Number) }, + context: { customizer_error: true }, + }, + ]); + for (const record of records) { + expect(await record.get()).toBe(events[0]); + expect(await record.get()).toBe(events[0]); + } + expect(onSpanExport).toHaveBeenCalledTimes(1); expect(errorLog).toHaveBeenCalledTimes(1); expect(laterHook).not.toHaveBeenCalled(); }); @@ -732,7 +832,7 @@ describe("onSpanExport", () => { } }); - test("throttles failure logs without exception or payload details", () => { + test("logs safely once across time, registrations, and SDK module copies", async () => { vi.useFakeTimers({ toFake: ["Date"] }); const errorLog = vi.spyOn(console, "error").mockImplementation(() => {}); const customizers = [ @@ -742,23 +842,29 @@ describe("onSpanExport", () => { }, }, ]; - const fail = () => - expect( - customizeSpanExport({ id: "row", output: "private" }, customizers), - ).toBeNull(); - - fail(); - fail(); - fail(); - expect(errorLog).toHaveBeenCalledTimes(1); - vi.advanceTimersByTime(59_000); - fail(); + for (let i = 0; i < 3; i++) { + configureInstrumentation({ spanCustomizers: customizers }); + expect(customizeSpanExport({ id: "row", output: "private" })).toEqual({ + id: "row", + context: { customizer_error: true }, + }); + vi.advanceTimersByTime(86_400_000); + } + // Reload intentionally models separately bundled SDK module state. + vi.resetModules(); + const otherBundle = await import("./span-customizer"); + expect( + otherBundle.customizeSpanExport({ id: "other", input: "private" }), + ).toEqual({ + id: "other", + context: { customizer_error: true }, + }); expect(errorLog).toHaveBeenCalledTimes(1); - vi.advanceTimersByTime(1_000); - fail(); - expect(errorLog).toHaveBeenCalledTimes(2); - expect(errorLog.mock.calls[1][0]).toContain("3 additional records"); - expect(JSON.stringify(errorLog.mock.calls)).not.toContain("private"); + expect(errorLog.mock.calls).toEqual([ + [ + "[braintrust] Span export customizer failed; exporting only span identity and duration.", + ], + ]); }); test.each([false, true])( @@ -789,7 +895,7 @@ describe("onSpanExport", () => { }, ); - test("drops async failures before attachments while exporting the rest of an HTTP batch", async () => { + test("strips async failures before attachments and sends ordered HTTP batches through retries", async () => { const errorLog = vi.spyOn(console, "error").mockImplementation(() => {}); const failedAttachment = new Attachment({ data: new ArrayBuffer(1), @@ -813,13 +919,17 @@ describe("onSpanExport", () => { ); return Object.freeze(payload); }); + let hookCalls = 0; configureInstrumentation({ spanCustomizers: [ { // @ts-expect-error Async customizers are unsupported, but must be contained. onSpanExport(data) { + hookCalls++; if (data.output === "private") { delete data.id; + data.metrics = { start: 999, end: 999, tokens: 999 }; + data.context = { secret: "private mutated context" }; return Promise.reject( new Error("private async exception detail"), ); @@ -835,6 +945,8 @@ describe("onSpanExport", () => { // manual span's initial row before the instrumented child is created. vi.stubEnv("BRAINTRUST_SYNC_FLUSH", "1"); const rows: Record[] = []; + const requests: Record[][] = []; + let retried = false; const state = new BraintrustState({ noExitFlush: true }); const logger = initLogger({ state, @@ -859,12 +971,34 @@ describe("onSpanExport", () => { } if (pathname === "/version") return Response.json({}); if (pathname === "/logs3") { - rows.push(...JSON.parse(String(options?.body)).rows); + const batch = JSON.parse(String(options?.body)).rows; + requests.push(batch); + if ( + !retried && + batch.some( + (row: SpanExportData) => + (row.context as SpanExportData)?.customizer_error, + ) + ) { + retried = true; + return Response.json({}, { status: 503 }); + } + rows.push(...batch); return Response.json({}); } throw new Error(`Unexpected test request: ${pathname}`); }, }); + const mask = vi.fn((value: unknown) => { + if (value && typeof value === "object" && "customizer_error" in value) { + return { injected: "must not replace diagnostics" }; + } + if (value && typeof value === "object" && "start" in value) { + return { start: 999 }; + } + return value; + }); + state.setMaskingFunction(mask); const manual = logger.startSpan({ name: "manual", event: { input: "manual input" }, @@ -875,30 +1009,55 @@ describe("onSpanExport", () => { INSTRUMENTATION_NAMES.OPENAI, ), ); - span.log({ input: failedAttachment, output: "private" }); + span.log({ + input: failedAttachment, + output: "private", + metrics: { start: 10, tokens: 99 }, + }); span.end(); manual.log({ input: successfulAttachment, output: "manual output" }); manual.end(); await logger.flush(); - expect(rows).toHaveLength(2); - expect(rows.find((row) => row.id === span.id)).toMatchObject({ + expect(rows).toHaveLength(5); + const stripped = rows.find( + (row) => (row.context as SpanExportData)?.customizer_error, + ); + expect(stripped).toEqual({ + id: span.id, span_id: span.spanId, root_span_id: manual.rootSpanId, span_parents: [manual.spanId], project_id: "customizer-project", log_id: "g", - input: "provider input", - metrics: { start: expect.any(Number), end: expect.any(Number) }, + _is_merge: true, + context: { customizer_error: true }, + metrics: { start: 10 }, }); - expect(rows.find((row) => row.id === manual.id)).toMatchObject({ + const failedRequests = requests.filter((batch) => + batch.some((row) => (row.context as SpanExportData)?.customizer_error), + ); + expect(failedRequests).toEqual([[stripped], [stripped]]); + expect( + requests.map((batch) => + batch.map( + (row) => (row.context as SpanExportData)?.customizer_error === true, + ), + ), + ).toEqual([[false, false], [true], [true], [false, false]]); + expect( + requests.every( + (batch) => new Set(batch.map((row) => row.id)).size === batch.length, + ), + ).toBe(true); + expect(rows.find((row) => row.output === "manual output")).toMatchObject({ input: successfulAttachment.reference, span_id: manual.spanId, root_span_id: manual.rootSpanId, project_id: "customizer-project", log_id: "g", output: "manual output", - metrics: { start: expect.any(Number), end: expect.any(Number) }, + metrics: { end: expect.any(Number) }, }); expect(rows.find((row) => row.id === span.id)).not.toHaveProperty("output"); expect(attachmentUpload).not.toHaveBeenCalled(); @@ -908,6 +1067,156 @@ describe("onSpanExport", () => { ).toBe(false); expect(errorLog).toHaveBeenCalledTimes(1); expect(JSON.stringify(errorLog.mock.calls)).not.toContain("private"); + expect(hookCalls).toBe(6); + expect(laterHook).toHaveBeenCalledTimes(5); + expect(JSON.stringify(mask.mock.calls)).not.toContain("customizer_error"); + }); + + test("only splits merge groups where a stripped record shares an ID", () => { + vi.spyOn(console, "error").mockImplementation(() => {}); + const failing = [ + { + onSpanExport(): SpanExportData { + throw new Error("private failure"); + }, + }, + ]; + const strip = (id: string) => + customizeSpanExport( + { id, _is_merge: true, output: "private" }, + failing, + ) as { + id: string; + }; + const a = strip("a"); + const b = strip("b"); + const a2 = strip("a"); + const y = { id: "y", _is_merge: true, output: 3 }; + + const groups = mergeCustomizedSpanRecords([ + { id: "x", input: 1 }, + a, + b, + { id: "x", _is_merge: true, output: 2 }, + a2, + y, + ]); + + expect(groups).toEqual([ + [{ id: "x", input: 1, output: 2 }, a, b], + [a2, y], + ]); + expect(groups[0][1]).toBe(a); + expect(groups[1][0]).toBe(a2); + }); + + function initHttpLogger(onLogs: (rows: SpanExportData[]) => Response) { + // Keep all records in one flush chunk. + vi.stubEnv("BRAINTRUST_SYNC_FLUSH", "1"); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "0"); + 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") { + return onLogs(JSON.parse(String(options?.body)).rows); + } + throw new Error(`Unexpected test request: ${pathname}`); + }, + }); + return { state, logger }; + } + + test("stops sending later ordered groups after a group fails", async () => { + vi.spyOn(console, "error").mockImplementation(() => {}); + vi.spyOn(console, "warn").mockImplementation(() => {}); + configureInstrumentation({ + spanCustomizers: [ + { + onSpanExport(data) { + if (data.output === "private") throw new Error("private failure"); + return data; + }, + }, + ], + }); + const requests: SpanExportData[][] = []; + const { logger } = initHttpLogger((rows) => { + requests.push(rows); + return rows.some( + (row) => (row.context as SpanExportData)?.customizer_error, + ) + ? Response.json({}, { status: 500 }) + : Response.json({}); + }); + const span = logger.startSpan({ name: "manual" }); + span.log({ output: "private" }); + span.end(); + await logger.flush(); + + expect( + requests.map((rows) => + rows.map((row) => (row.context as SpanExportData)?.customizer_error), + ), + ).toEqual([[undefined], [true]]); + }); + + test("does not upload attachments whose references were masked away", async () => { + const masked = new Attachment({ + data: new ArrayBuffer(1), + filename: "private.txt", + contentType: "text/plain", + }); + const maskedUpload = vi.spyOn(masked, "upload"); + const kept = new Attachment({ + data: new ArrayBuffer(1), + filename: "safe.txt", + contentType: "text/plain", + }); + const keptUpload = vi + .spyOn(kept, "upload") + .mockResolvedValue({ upload_status: "done" }); + const rows: SpanExportData[] = []; + const { state, logger } = initHttpLogger((batch) => { + rows.push(...batch); + return Response.json({}); + }); + state.setMaskingFunction((value) => { + if (JSON.stringify(value)?.includes("private.txt")) { + throw new Error("private failure"); + } + return value; + }); + const span = logger.startSpan({ name: "manual" }); + span.log({ input: masked, output: kept }); + span.end(); + await logger.flush(); + + expect(rows.find((row) => row.id === span.id)).toMatchObject({ + input: "ERROR: Failed to mask field 'input' - Error", + output: kept.reference, + }); + expect(maskedUpload).not.toHaveBeenCalled(); + expect(keptUpload).toHaveBeenCalledTimes(1); }); }); diff --git a/js/src/span-customizer.ts b/js/src/span-customizer.ts index 0682de399..837e8a789 100644 --- a/js/src/span-customizer.ts +++ b/js/src/span-customizer.ts @@ -8,6 +8,7 @@ import { PARENT_ID_FIELD, TRANSACTION_ID_FIELD, } from "../util/db_fields"; +import { mergeRowBatch } from "../util/merge_row_batch"; import { isPromiseLike } from "../util/type_util"; import { isOtelCompatMode } from "./id-gen"; import type { SpanCustomizer, SpanExportData } from "./instrumentation/config"; @@ -15,8 +16,14 @@ 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 CUSTOMIZER_FAILURE_LOGGED_KEY = Symbol.for( + "braintrust.spanCustomizerFailureLogged", +); +const STRIPPED_RECORDS_KEY = Symbol.for("braintrust.strippedSpanRecords"); const shared: typeof globalThis & { [SPAN_CUSTOMIZERS_KEY]?: readonly SpanCustomizer[]; + [CUSTOMIZER_FAILURE_LOGGED_KEY]?: boolean; + [STRIPPED_RECORDS_KEY]?: WeakSet; } = globalThis; const PROTECTED_FIELDS = new Set([ @@ -82,14 +89,23 @@ export function createMaskingCustomizer( } // Accept null-prototype objects and objects from any realm's Object.prototype -// (e.g. node:vm), whose own prototype is null. Arrays, Dates, and class -// instances all have a longer prototype chain. +// (e.g. node:vm). Reject custom prototypes, including classes extending null. function isPlainRecord(value: unknown): value is SpanExportData { if (value === null || typeof value !== "object" || Array.isArray(value)) { return false; } const prototype = Object.getPrototypeOf(value); - return prototype === null || Object.getPrototypeOf(prototype) === null; + if (prototype === null) return true; + const constructor = Object.getOwnPropertyDescriptor( + prototype, + "constructor", + )?.value; + return ( + Object.getPrototypeOf(prototype) === null && + typeof constructor === "function" && + Function.prototype.toString.call(constructor) === + Function.prototype.toString.call(Object) + ); } // Only plain containers are copied deeply; payloads may contain SDK objects such @@ -128,38 +144,73 @@ function restoreProtocolFields( return restored; } -// Log the first failure immediately, then at most once per period with a count -// of suppressed failures, mirroring the background logger's queue-drop logging. -const FAILURE_LOGGING_PERIOD_MS = 60_000; -const failureLoggingState = { - numSuppressed: 0, - lastLoggedTimestamp: -Infinity, -}; - function registerCustomizerFailure(): void { - const timeNow = Date.now(); - if ( - timeNow - failureLoggingState.lastLoggedTimestamp < - FAILURE_LOGGING_PERIOD_MS - ) { - failureLoggingState.numSuppressed++; - return; - } - const suppressed = failureLoggingState.numSuppressed - ? ` ${failureLoggingState.numSuppressed} additional records were dropped since the last report.` - : ""; + if (shared[CUSTOMIZER_FAILURE_LOGGED_KEY]) return; + shared[CUSTOMIZER_FAILURE_LOGGED_KEY] = true; // Never log the exception or record: either can contain sensitive data. // eslint-disable-next-line no-restricted-properties -- intentional always-visible export failure. console.error( - `[braintrust] Span export customizer failed; dropping record.${suppressed}`, + "[braintrust] Span export customizer failed; exporting only span identity and duration.", ); - failureLoggingState.numSuppressed = 0; - failureLoggingState.lastLoggedTimestamp = timeNow; } export function resetSpanCustomizerFailureLoggingForTests(): void { - failureLoggingState.numSuppressed = 0; - failureLoggingState.lastLoggedTimestamp = -Infinity; + delete shared[CUSTOMIZER_FAILURE_LOGGED_KEY]; +} + +function stripFailedRecord(original: SpanExportData): SpanExportData { + const stripped: SpanExportData = {}; + for (const key of PROTECTED_FIELDS) { + if (Object.prototype.hasOwnProperty.call(original, key)) { + stripped[key] = copyProtocolValue(original[key]); + } + } + // Native records represent duration with start/end timestamps. Preserve only + // the timestamps present in this incremental record, never synthesize either. + if (isPlainRecord(original.metrics)) { + const metrics: SpanExportData = {}; + for (const key of ["start", "end"]) { + if (Object.prototype.hasOwnProperty.call(original.metrics, key)) { + metrics[key] = original.metrics[key]; + } + } + if (Object.keys(metrics).length) stripped.metrics = metrics; + } + stripped.context = { customizer_error: true }; + (shared[STRIPPED_RECORDS_KEY] ??= new WeakSet()).add(stripped); + return stripped; +} + +// A stripped incremental record must not absorb other records or be mutated by +// merging: the same object is memoized for retries. Duplicate IDs in one request +// share a transaction ID, so a group boundary is only needed where a stripped +// record would share a group with another record of the same ID. Groups must be +// sent in order, in separate requests. +export function mergeCustomizedSpanRecords< + T extends { id: string; _is_merge?: boolean | null }, +>(records: T[]): T[][] { + const stripped = shared[STRIPPED_RECORDS_KEY]; + if (!stripped || !records.some((record) => stripped.has(record))) { + return [mergeRowBatch(records)]; + } + const merged: T[][] = []; + let group: T[] = []; + let ids = new Set(); + let strippedIds = new Set(); + for (const record of records) { + const isStripped = stripped.has(record); + if (strippedIds.has(record.id) || (isStripped && ids.has(record.id))) { + merged.push(mergeRowBatch(group)); + group = []; + ids = new Set(); + strippedIds = new Set(); + } + group.push(record); + ids.add(record.id); + if (isStripped) strippedIds.add(record.id); + } + merged.push(mergeRowBatch(group)); + return merged; } export function setSpanCustomizers( @@ -181,9 +232,13 @@ export function customizeSpanExport( customizers: readonly SpanCustomizer[] | undefined = shared[ SPAN_CUSTOMIZERS_KEY ], -): SpanExportData | null { - if (!customizers?.length) return data; +): SpanExportData { + // SDK-generated failure diagnostics and original timestamps are final, not + // user content for downstream masking to replace. + if (!customizers?.length || shared[STRIPPED_RECORDS_KEY]?.has(data)) + return data; + const original = data; let protectedFields: SpanExportData | undefined; for (const customizer of customizers) { @@ -217,7 +272,7 @@ export function customizeSpanExport( data = restoreProtocolFields(result, protectedFields); } catch { registerCustomizerFailure(); - return null; + return stripFailedRecord(original); } } return data;