From 0ebfbed4697b99e767f49edae1ab65fd641825d7 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:19:50 +0000 Subject: [PATCH 1/5] feat: Support ingestion keys in `initLogger` `initLogger({ ingestionKey })` and `BRAINTRUST_INGESTION_KEY` now log through a project's ingestion URL with the key as a Bearer token. These loggers get their own state and queue, never log in, don't look up the project and ignore private credentials. An explicit `apiKey` or `state` still wins over the env var, and combining one with `ingestionKey` throws. Attachments and overflow batches go through the chunked `/v1/uploads` flow before the rows that reference them, and rows are trimmed to the fields the ingestion endpoint accepts. Span context is shared with the global state, so nesting and W3C propagation keep working without a project id. --- .changeset/ingestion-keys-init-logger.md | 5 + js/README.md | 22 + js/src/ingestion-key.test.ts | 696 +++++++++++++++++++++++ js/src/ingestion-key.ts | 194 +++++++ js/src/logger.ts | 551 ++++++++++++++++-- 5 files changed, 1432 insertions(+), 36 deletions(-) create mode 100644 .changeset/ingestion-keys-init-logger.md create mode 100644 js/src/ingestion-key.test.ts create mode 100644 js/src/ingestion-key.ts diff --git a/.changeset/ingestion-keys-init-logger.md b/.changeset/ingestion-keys-init-logger.md new file mode 100644 index 000000000..ea67da623 --- /dev/null +++ b/.changeset/ingestion-keys-init-logger.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Support ingestion keys in `initLogger` diff --git a/js/README.md b/js/README.md index 9c5e70094..a1e8c6e9d 100644 --- a/js/README.md +++ b/js/README.md @@ -104,6 +104,28 @@ If you use TypeScript or other transpilation plugins, place the Braintrust plugi For deeper details, see the [auto-instrumentation architecture docs](src/auto-instrumentations/README.md). +## Ingestion Keys + +An ingestion key lets public clients, like browsers, write traces into a single project without an API key. +Ingestion keys cannot read data, call models, or manage anything, so a logger that uses one never logs in, never looks up or creates the project, and ignores `BRAINTRUST_API_KEY`. + +Pass the full ingestion URL of the key, or set it as `BRAINTRUST_INGESTION_KEY`: + +```ts +import { initLogger } from "braintrust"; + +const logger = initLogger({ + ingestionKey: "https:///ingest?ingestKey=", +}); +``` + +An explicit `apiKey` or `state` takes precedence over `BRAINTRUST_INGESTION_KEY`, and passing either together with `ingestionKey` throws. +`projectName` is ignored, and `projectId` is only sent along for the server to check. +Attachments and large batches are uploaded through the ingestion URL as well. +Fields that ingestion keys do not support, like `classifications` or audit fields, are not logged, and logging comments throws. + +To export OpenTelemetry spans, point a standard OTLP exporter at the ingestion URL with `/otel/v1/traces` appended and without the query string, and send the key in an `Authorization: Bearer ` header. + ## Migration Guides ### Upgrading from 2.x to 3.x diff --git a/js/src/ingestion-key.test.ts b/js/src/ingestion-key.test.ts new file mode 100644 index 000000000..cfc286068 --- /dev/null +++ b/js/src/ingestion-key.test.ts @@ -0,0 +1,696 @@ +import { createHash, randomUUID } from "node:crypto"; +import { inspect } from "node:util"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + _internalGetGlobalState, + Attachment, + BraintrustState, + extractTraceContextFromHeaders, + flush, + initLogger, + type InitLoggerOptions, + setMaskingFunction, + traced, +} from "./logger"; +import { parseIngestionKeyUrl } from "./ingestion-key"; +import { configureNode } from "./node/config"; + +configureNode(); + +const KIB = 1024; + +function newKey() { + return `bt-ik-${randomUUID().replace(/-/g, "")}${"a".repeat(16)}`; +} + +function ingestionUrl( + key: string, + base = "https://dp.example/deployment/base", +) { + return `${base}/ingest?ingestKey=${key}`; +} + +interface RecordedRequest { + url: string; + method: string; + headers: Record; + body: Uint8Array | string | undefined; +} + +function json(data: unknown, status = 200) { + return new Response(JSON.stringify(data), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +function bodyJson(request: RecordedRequest) { + return JSON.parse(String(request.body)); +} + +function paths(requests: RecordedRequest[]) { + return requests.map((r) => `${r.method} ${new URL(r.url).pathname}`); +} + +/** + * Mock data plane that implements the ingestion upload and log endpoints. + * `respond` can override the response of any request. + */ +function mockIngestion( + options: { + chunkBytes?: number; + respond?: ( + request: RecordedRequest, + index: number, + ) => Response | Promise | undefined; + } = {}, +) { + const requests: RecordedRequest[] = []; + const uploads = new Map< + string, + { request: Record; chunks: Uint8Array[] } + >(); + const fetch = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const body = + init?.body instanceof Blob + ? new Uint8Array(await init.body.arrayBuffer()) + : (init?.body as string | undefined); + const request: RecordedRequest = { + url: String(input), + method: init?.method ?? "GET", + headers: Object.fromEntries( + Object.entries((init?.headers ?? {}) as Record).map( + ([k, v]) => [k.toLowerCase(), v], + ), + ), + body, + }; + requests.push(request); + const override = await options.respond?.(request, requests.length - 1); + if (override) { + return override; + } + + const path = new URL(request.url).pathname; + if (request.method === "POST" && path.endsWith("/v1/uploads")) { + const uploadRequest = bodyJson(request); + const chunkBytes = options.chunkBytes ?? 512 * KIB; + const uploadId = randomUUID(); + uploads.set(uploadId, { request: uploadRequest, chunks: [] }); + return json( + { + upload_id: uploadId, + chunk_bytes: chunkBytes, + num_chunks: Math.ceil(uploadRequest.size_bytes / chunkBytes), + expires_in_ms: 300_000, + }, + 201, + ); + } + const chunkMatch = /\/v1\/uploads\/([^/]+)\/chunks\/(\d+)$/.exec(path); + if (request.method === "PUT" && chunkMatch) { + const chunk = request.body as Uint8Array; + uploads.get(chunkMatch[1])!.chunks[Number(chunkMatch[2])] = chunk; + return json({ index: Number(chunkMatch[2]), size_bytes: chunk.length }); + } + const completeMatch = /\/v1\/uploads\/([^/]+)\/complete$/.exec(path); + if (request.method === "POST" && completeMatch) { + const uploadId = completeMatch[1]; + const upload = uploads.get(uploadId)!; + const data = Buffer.concat(upload.chunks); + return json({ + reference: + upload.request.purpose === "attachment" + ? { + type: "braintrust_attachment", + filename: upload.request.filename, + content_type: upload.request.content_type, + key: `server-${uploadId}`, + } + : { type: "logs3_overflow", key: uploadId }, + size_bytes: data.length, + sha256: createHash("sha256").update(data).digest("hex"), + }); + } + if (request.method === "POST" && path.endsWith("/v1/logs")) { + const payload = bodyJson(request); + return json({ + ids: Array.isArray(payload.rows) + ? payload.rows.map((row: { id: string }) => row.id) + : [], + }); + } + return json({ error: "unexpected request" }, 404); + }); + return { + requests, + uploads, + fetch: fetch as unknown as typeof globalThis.fetch, + }; +} + +function loggedRows(requests: RecordedRequest[]) { + return requests + .filter((r) => new URL(r.url).pathname.endsWith("/v1/logs")) + .flatMap((r) => bodyJson(r).rows); +} + +function initIngestionLogger( + fetch: typeof globalThis.fetch, + options: InitLoggerOptions = {}, +) { + const key = newKey(); + const onFlushError = vi.fn(); + const logger = initLogger({ + ingestionKey: ingestionUrl(key), + fetch, + noExitFlush: true, + onFlushError, + ...options, + }); + return { key, logger, onFlushError }; +} + +beforeEach(() => { + // Only flush explicitly, so the rows of a span are sent together. + vi.stubEnv("BRAINTRUST_SYNC_FLUSH", "1"); +}); + +afterEach(() => { + vi.unstubAllEnvs(); + _internalGetGlobalState().currentLogger = undefined; +}); + +describe("parseIngestionKeyUrl", () => { + const key = newKey(); + + test("keeps the deployment base path and strips the query", () => { + expect( + parseIngestionKeyUrl(`https://dp.example/a/b/ingest?ingestKey=${key}`), + ).toEqual({ root: "https://dp.example/a/b/ingest", key }); + expect( + parseIngestionKeyUrl(`http://localhost:8000/ingest/?ingestKey=${key}`), + ).toEqual({ root: "http://localhost:8000/ingest", key }); + }); + + test.each([ + ["a bare key", key], + ["a non-http URL", `ftp://dp.example/ingest?ingestKey=${key}`], + ["credentials", `https://user:pw@dp.example/ingest?ingestKey=${key}`], + ["a fragment", `https://dp.example/ingest?ingestKey=${key}#x`], + ["an empty fragment", `https://dp.example/ingest?ingestKey=${key}#`], + [ + "another query parameter", + `https://dp.example/ingest?ingestKey=${key}&a=b`, + ], + [ + "a repeated key", + `https://dp.example/ingest?ingestKey=${key}&ingestKey=${key}`, + ], + ["a malformed key", "https://dp.example/ingest?ingestKey=sk-123"], + ["a different path", `https://dp.example/logs?ingestKey=${key}`], + ])("rejects %s without echoing the key", (_, value) => { + let error: unknown; + try { + parseIngestionKeyUrl(value); + } catch (e) { + error = e; + } + expect(error).toBeInstanceOf(Error); + expect(inspect(error)).not.toContain(key); + expect((error as Error).cause).toBeUndefined(); + }); +}); + +describe("initLogger with an ingestion key", () => { + test("first flush only posts rows to the ingestion endpoint", async () => { + const { requests, fetch } = mockIngestion(); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + const { key, logger, onFlushError } = initIngestionLogger(fetch, { + projectName: "ignored-name", + }); + + logger.traced((span) => span.log({ input: "hi", output: "hello" }), { + name: "root", + }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(requests).toHaveLength(1); + const [request] = requests; + expect(request.url).toBe( + "https://dp.example/deployment/base/ingest/v1/logs", + ); + expect(request.method).toBe("POST"); + expect(request.headers.authorization).toBe(`Bearer ${key}`); + const payload = bodyJson(request); + expect(payload.api_version).toBe(2); + expect(payload.rows).toHaveLength(1); + const [row] = payload.rows; + expect(row).toMatchObject({ + input: "hi", + output: "hello", + log_id: "g", + span_attributes: { name: "root" }, + }); + expect(row).not.toHaveProperty("project_id"); + expect(row).not.toHaveProperty("org_id"); + expect(JSON.stringify(requests)).not.toContain("private-api-key"); + expect(JSON.stringify(requests)).not.toContain("ignored-name"); + }); + + test("sends an explicit project id for the server to check", async () => { + const { requests, fetch } = mockIngestion({ + respond: (request) => + request.url.endsWith("/v1/logs") + ? json({ error: "project mismatch" }, 403) + : undefined, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "0"); + const { logger, onFlushError } = initIngestionLogger(fetch, { + projectId: "other-project", + }); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(loggedRows(requests)).toEqual([ + expect.objectContaining({ project_id: "other-project", log_id: "g" }), + ]); + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(String(onFlushError.mock.calls[0][0])).toContain("403"); + }); + + test("uses BRAINTRUST_INGESTION_KEY and ignores private credentials", async () => { + const { requests, fetch } = mockIngestion(); + const key = newKey(); + vi.stubEnv("BRAINTRUST_INGESTION_KEY", ingestionUrl(key)); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + const globalState = _internalGetGlobalState(); + globalState.apiUrl = "https://api.test"; + globalState.loginToken = "private-api-key"; + globalState.loggedIn = true; + + const logger = initLogger({ fetch, noExitFlush: true }); + logger.log({ input: "a", output: "b" }); + await flush(); + + expect(paths(requests)).toEqual(["POST /deployment/base/ingest/v1/logs"]); + expect(requests[0].headers.authorization).toBe(`Bearer ${key}`); + globalState.resetLoginInfo(); + }); + + test("an explicit private credential takes precedence over the environment", async () => { + const { requests, fetch } = mockIngestion({ respond: () => json({}) }); + const key = newKey(); + vi.stubEnv("BRAINTRUST_INGESTION_KEY", ingestionUrl(key)); + const state = new BraintrustState({ fetch, noExitFlush: true }); + state.apiUrl = "https://api.test"; + state.loginToken = "private-api-key"; + state.apiConn().set_token("private-api-key"); + + const logger = initLogger({ + state, + projectId: "private-project", + projectName: "private", + }); + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(logger.loggingState).toBe(state); + expect(paths(requests)).toEqual(["GET /version", "POST /logs3"]); + expect(requests[1].headers.authorization).toBe("Bearer private-api-key"); + expect(JSON.stringify(requests)).not.toContain(key); + + const apiKeyLogger = initLogger({ apiKey: "private-api-key" }); + expect(apiKeyLogger.loggingState).toBe(_internalGetGlobalState()); + }); + + test("rejects an ingestion key combined with private credentials", () => { + const ingestionKey = ingestionUrl(newKey()); + expect(() => initLogger({ ingestionKey, apiKey: "private" })).toThrow( + "either an ingestionKey or an apiKey/state", + ); + expect(() => + initLogger({ ingestionKey, state: new BraintrustState({}) }), + ).toThrow("either an ingestionKey or an apiKey/state"); + expect(() => initLogger({ ingestionKey: "" })).toThrow( + "Invalid Braintrust ingestion key", + ); + }); + + test("public and private loggers do not share queues or credentials", async () => { + const { requests, fetch } = mockIngestion({ + respond: (request) => + new URL(request.url).host === "api.test" ? json({}) : undefined, + }); + const privateState = new BraintrustState({ fetch, noExitFlush: true }); + privateState.apiUrl = "https://api.test"; + privateState.loginToken = "private-api-key"; + privateState.apiConn().set_token("private-api-key"); + const privateLogger = initLogger({ + state: privateState, + projectId: "private-project", + projectName: "private", + }); + const first = initIngestionLogger(fetch); + const second = initIngestionLogger(fetch); + + privateLogger.log({ input: "private", output: "x" }); + first.logger.log({ input: "first", output: "x" }); + second.logger.log({ input: "second", output: "x" }); + await Promise.all([ + privateLogger.flush(), + first.logger.flush(), + second.logger.flush(), + ]); + + const inputsByAuth = Object.fromEntries( + requests + .filter((r) => r.method === "POST") + .map((r) => [ + r.headers.authorization, + bodyJson(r).rows.map((row: { input: string }) => row.input), + ]), + ); + expect(inputsByAuth).toEqual({ + "Bearer private-api-key": ["private"], + [`Bearer ${first.key}`]: ["first"], + [`Bearer ${second.key}`]: ["second"], + }); + expect( + requests + .filter((r) => r.headers.authorization !== "Bearer private-api-key") + .every((r) => r.url.startsWith("https://dp.example/")), + ).toBe(true); + }); + + test("retries rejected writes against the ingestion endpoint only", async () => { + const { requests, fetch } = mockIngestion({ + respond: (request, index) => + index === 0 ? json({ error: "slow down" }, 429) : undefined, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "1"); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(paths(requests)).toEqual([ + "POST /deployment/base/ingest/v1/logs", + "POST /deployment/base/ingest/v1/logs", + ]); + }); + + test("gives up after the configured retries without a private fallback", async () => { + const { requests, fetch } = mockIngestion({ + respond: () => json({ error: "unavailable" }, 503), + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "0"); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(paths(requests)).toEqual(["POST /deployment/base/ingest/v1/logs"]); + }); + + test("uploads attachments in advertised chunks before logging rows", async () => { + const { requests, uploads, fetch } = mockIngestion({ chunkBytes: 4 * KIB }); + const { logger, onFlushError } = initIngestionLogger(fetch); + const data = new Uint8Array(9 * KIB).map((_, i) => i % 251); + const attachment = new Attachment({ + data: data.buffer, + filename: "data.bin", + contentType: "application/octet-stream", + }); + + logger.log({ input: { file: attachment }, output: "ok" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + const [uploadId] = uploads.keys(); + expect(paths(requests)).toEqual([ + "POST /deployment/base/ingest/v1/uploads", + `PUT /deployment/base/ingest/v1/uploads/${uploadId}/chunks/0`, + `PUT /deployment/base/ingest/v1/uploads/${uploadId}/chunks/1`, + `PUT /deployment/base/ingest/v1/uploads/${uploadId}/chunks/2`, + `POST /deployment/base/ingest/v1/uploads/${uploadId}/complete`, + "POST /deployment/base/ingest/v1/logs", + ]); + expect(bodyJson(requests[0])).toEqual({ + purpose: "attachment", + filename: "data.bin", + content_type: "application/octet-stream", + size_bytes: data.length, + sha256: createHash("sha256").update(data).digest("hex"), + }); + expect( + requests + .slice(1, 4) + .map((r) => [r.headers["content-type"], (r.body as Uint8Array).length]), + ).toEqual([ + ["application/octet-stream", 4 * KIB], + ["application/octet-stream", 4 * KIB], + ["application/octet-stream", 1 * KIB], + ]); + expect(bodyJson(requests[4])).toEqual({}); + expect(loggedRows(requests)[0].input.file).toEqual({ + type: "braintrust_attachment", + filename: "data.bin", + content_type: "application/octet-stream", + key: `server-${uploadId}`, + }); + }); + + test("uploads empty attachments without chunks", async () => { + const { requests, uploads, fetch } = mockIngestion(); + const { logger, onFlushError } = initIngestionLogger(fetch); + const attachment = new Attachment({ + data: new ArrayBuffer(0), + filename: "empty.txt", + contentType: "text/plain", + }); + + logger.log({ input: attachment, output: "ok" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + const [uploadId] = uploads.keys(); + expect(paths(requests)).toEqual([ + "POST /deployment/base/ingest/v1/uploads", + `POST /deployment/base/ingest/v1/uploads/${uploadId}/complete`, + "POST /deployment/base/ingest/v1/logs", + ]); + expect(bodyJson(requests[0]).size_bytes).toBe(0); + expect(loggedRows(requests)[0].input.key).toBe(`server-${uploadId}`); + }); + + test("drops rows whose attachment failed to upload", async () => { + const { requests, fetch } = mockIngestion({ + respond: (request) => + request.url.endsWith("/v1/uploads") + ? json({ error: "quota" }, 429) + : undefined, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "0"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "dropped", + }); + logger.log({ input: "kept", output: "kept" }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(loggedRows(requests).map((row) => row.output)).toEqual(["kept"]); + expect(requests.some((r) => r.url.includes("/attachment"))).toBe(false); + }); + + test("restarts an upload whose grant expired", async () => { + let grants = 0; + const { requests, fetch } = mockIngestion({ + respond: async (request) => { + if (request.url.endsWith("/v1/uploads") && grants++ === 0) { + await new Promise((resolve) => setTimeout(resolve, 5)); + return json( + { + upload_id: randomUUID(), + chunk_bytes: 512 * KIB, + num_chunks: 1, + expires_in_ms: 1, + }, + 201, + ); + } + return undefined; + }, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "1"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "ok", + }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect( + paths(requests).filter((p) => p.endsWith("/v1/uploads")), + ).toHaveLength(2); + expect(paths(requests).filter((p) => p.includes("/chunks/"))).toHaveLength( + 1, + ); + }); + + test("overflows large batches through an upload", async () => { + const { requests, uploads, fetch } = mockIngestion({ chunkBytes: 1 * KIB }); + vi.stubEnv("BRAINTRUST_MAX_REQUEST_SIZE", String(2 * KIB)); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ input: "x".repeat(3 * KIB), output: "big" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + const [[uploadId, upload]] = uploads.entries(); + expect(upload.request).toMatchObject({ + purpose: "logs3_overflow", + content_type: "application/json", + }); + expect(upload.request).not.toHaveProperty("filename"); + const uploaded = JSON.parse(Buffer.concat(upload.chunks).toString("utf8")); + expect(uploaded.api_version).toBe(2); + expect(uploaded.rows).toEqual([ + expect.objectContaining({ output: "big", log_id: "g" }), + ]); + const logRequests = requests.filter((r) => r.url.endsWith("/v1/logs")); + expect(logRequests.map(bodyJson)).toEqual([ + { api_version: 2, rows: { type: "logs3_overflow", key: uploadId } }, + ]); + expect( + requests + .filter((r) => r.url.includes("/chunks/")) + .every((r) => (r.body as Uint8Array).length <= 1 * KIB), + ).toBe(true); + }); + + test("only sends supported row fields", async () => { + const { requests, fetch } = mockIngestion(); + const { logger } = initIngestionLogger(fetch); + + const id = logger.log({ input: "a", output: "b" }); + logger.logFeedback({ id, scores: { good: 1 }, metadata: { user: "u" } }); + expect(() => logger.logFeedback({ id, comment: "nice" })).toThrow( + "do not support logging comments", + ); + await logger.flush(); + + const rows = loggedRows(requests); + expect(rows).toContainEqual( + expect.objectContaining({ id, scores: { good: 1 } }), + ); + for (const row of rows) { + expect(row).not.toHaveProperty("_audit_source"); + expect(row).not.toHaveProperty("_audit_metadata"); + expect(row).not.toHaveProperty("comment"); + } + }); + + test("applies the global masking function", async () => { + const { requests, fetch } = mockIngestion(); + const { logger } = initIngestionLogger(fetch); + setMaskingFunction((value) => + value === "secret-value" ? "[masked]" : value, + ); + try { + logger.log({ input: "secret-value", output: "b" }); + await logger.flush(); + } finally { + setMaskingFunction(null); + } + + expect(loggedRows(requests)[0].input).toBe("[masked]"); + }); + + test("never exposes the key in URLs or state", async () => { + const { requests, fetch } = mockIngestion(); + const { key, logger } = initIngestionLogger(fetch); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(requests.every((r) => !r.url.includes(key))).toBe(true); + expect(inspect(logger.loggingState)).not.toContain(key); + expect(JSON.stringify(logger.loggingState)).not.toContain(key); + expect(String(logger.loggingState)).not.toContain(key); + }); +}); + +describe("tracing with an ingestion key", () => { + test("nests and propagates spans without project metadata", async () => { + const { requests, fetch } = mockIngestion(); + const { logger, onFlushError } = initIngestionLogger(fetch); + + let exported = ""; + let headers: Record = {}; + logger.traced( + (root) => { + logger.traced( + async (child) => { + exported = await child.export(); + }, + { name: "child" }, + ); + headers = root.inject(); + }, + { name: "root" }, + ); + await vi.waitFor(() => expect(exported).not.toBe("")); + + const remote = logger.startSpan({ + name: "remote", + parent: extractTraceContextFromHeaders(headers), + }); + remote.end(); + traced(() => {}, { name: "global", parent: exported }); + await flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(paths(requests).every((p) => p.endsWith("/v1/logs"))).toBe(true); + const rows = Object.fromEntries( + loggedRows(requests) + .filter((row) => row.span_attributes?.name) + .map((row) => [row.span_attributes.name, row]), + ); + expect(Object.keys(rows).sort()).toEqual([ + "child", + "global", + "remote", + "root", + ]); + expect(rows.child.root_span_id).toBe(rows.root.root_span_id); + expect(rows.child.span_parents).toEqual([rows.root.span_id]); + expect(rows.remote.root_span_id).toBe(rows.root.root_span_id); + expect(rows.remote.span_parents).toEqual([rows.root.span_id]); + expect(rows.global.root_span_id).toBe(rows.root.root_span_id); + expect(rows.global.span_parents).toEqual([rows.child.span_id]); + for (const row of Object.values(rows)) { + expect(row).not.toHaveProperty("project_id"); + } + expect(headers.baggage ?? "").not.toContain("braintrust.parent"); + }); +}); diff --git a/js/src/ingestion-key.ts b/js/src/ingestion-key.ts new file mode 100644 index 000000000..cd8473b4b --- /dev/null +++ b/js/src/ingestion-key.ts @@ -0,0 +1,194 @@ +import { z } from "zod/v3"; +import { + BraintrustAttachmentReference as braintrustAttachmentReferenceSchema, + type BraintrustAttachmentReferenceType as BraintrustAttachmentReference, +} from "./generated_types"; +import { + ARRAY_DELETE_FIELD, + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + OBJECT_DELETE_FIELD, +} from "./util"; + +export const INGESTION_KEY_ENV_VAR = "BRAINTRUST_INGESTION_KEY"; + +export interface IngestionEndpoint { + // The `/ingest` URL without its query string. Requests are sent below it. + root: string; + key: string; +} + +const INGESTION_KEY_QUERY = /^\?ingestKey=(bt-ik-[A-Za-z0-9]+)$/; + +/** + * Parse an ingestion key URL such as + * `https://dp.example/base/ingest?ingestKey=bt-ik-...`. + * + * Errors never include the provided value, since it contains the key. + */ +export function parseIngestionKeyUrl(value: string): IngestionEndpoint { + const invalid = (reason: string) => + new Error( + `Invalid Braintrust ingestion key: ${reason}. Expected the full ingestion URL, like https:///ingest?ingestKey=.`, + ); + + let url: URL; + try { + url = new URL(value.trim()); + } catch { + // The URL error references the input, so it is not attached as a cause. + throw invalid("it is not a valid URL"); + } + if (url.protocol !== "https:" && url.protocol !== "http:") { + throw invalid("the URL must use http or https"); + } + if (url.username || url.password) { + throw invalid("the URL must not contain credentials"); + } + if (url.hash || value.includes("#")) { + throw invalid("the URL must not contain a fragment"); + } + const match = INGESTION_KEY_QUERY.exec(url.search); + if (!match) { + throw invalid( + "the URL must contain exactly one ingestKey query parameter and no other query parameters", + ); + } + const pathname = url.pathname.replace(/\/$/, ""); + if (!pathname.endsWith("/ingest")) { + throw invalid("the URL path must end with /ingest"); + } + return { root: `${url.origin}${pathname}`, key: match[1] }; +} + +// Row fields the ingestion endpoint accepts. Everything else (experiment and +// dataset ids, audit fields, scoring controls, comments, ...) is dropped. +const INGESTION_ROW_FIELDS = new Set([ + "id", + "span_id", + "root_span_id", + "span_parents", + "created", + "org_id", + "project_id", + "log_id", + "input", + "output", + "expected", + "error", + "tags", + "scores", + "metadata", + "metrics", + "context", + "span_attributes", + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + ARRAY_DELETE_FIELD, + OBJECT_DELETE_FIELD, +]); + +/** + * Copy the fields of `row` that the ingestion endpoint accepts. The names of + * dropped fields are added to `droppedFields`. + */ +export function pickIngestionRowFields( + row: object, + droppedFields: Set, +): Record { + const picked: Record = {}; + for (const [key, value] of Object.entries(row)) { + if (INGESTION_ROW_FIELDS.has(key)) { + picked[key] = value; + } else if (value !== undefined) { + droppedFields.add(key); + } + } + return picked; +} + +const BRAINTRUST_ATTACHMENT = + braintrustAttachmentReferenceSchema.shape.type.value; + +/** + * Replace attachment references whose key is in `references` with the + * reference the ingestion endpoint returned for the upload. A reference mapped + * to `undefined` failed to upload, in which case this returns false. + */ +export function replaceIngestionAttachmentReferences( + container: Record | unknown[], + references: ReadonlyMap, +): boolean { + for (const [key, value] of Object.entries(container)) { + if (!value || typeof value !== "object") { + continue; + } + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + const reference = value as Record; + if ( + reference.type === BRAINTRUST_ATTACHMENT && + typeof reference.key === "string" && + references.has(reference.key) + ) { + const uploaded = references.get(reference.key); + if (!uploaded) { + return false; + } + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + (container as Record)[key] = uploaded; + continue; + } + if (!replaceIngestionAttachmentReferences(reference, references)) { + return false; + } + } + return true; +} + +export type IngestionUploadRequest = + | { + purpose: "attachment"; + filename: string; + content_type: string; + } + | { + purpose: "logs3_overflow"; + content_type: "application/json"; + }; + +const sha256Schema = z.string().regex(/^[0-9a-f]{64}$/); + +export const ingestionUploadGrantSchema = z.object({ + upload_id: z.string().uuid(), + chunk_bytes: z.number().int().positive(), + num_chunks: z.number().int().nonnegative(), + expires_in_ms: z.number().int().positive().max(300_000), +}); + +export const ingestionUploadChunkSchema = z.object({ + index: z.number().int().nonnegative(), + size_bytes: z.number().int().nonnegative(), +}); + +export const ingestionUploadCompleteSchema = z.object({ + reference: z.union([ + braintrustAttachmentReferenceSchema, + z.object({ type: z.literal("logs3_overflow"), key: z.string().min(1) }), + ]), + size_bytes: z.number().int().nonnegative(), + sha256: sha256Schema, +}); + +/** + * Hex SHA-256 digest of `data`, or undefined on platforms without WebCrypto. + */ +export async function sha256Hex(data: Blob): Promise { + const subtle = globalThis.crypto?.subtle; + if (!subtle) { + return undefined; + } + const digest = await subtle.digest("SHA-256", await data.arrayBuffer()); + return Array.from(new Uint8Array(digest), (byte) => + byte.toString(16).padStart(2, "0"), + ).join(""); +} diff --git a/js/src/logger.ts b/js/src/logger.ts index 3a5be45ca..26d20d334 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -162,6 +162,18 @@ const parametersRowSchema = z.object({ type ParametersRow = z.infer; import { waitUntil } from "@vercel/functions"; +import { + INGESTION_KEY_ENV_VAR, + type IngestionEndpoint, + type IngestionUploadRequest, + ingestionUploadChunkSchema, + ingestionUploadCompleteSchema, + ingestionUploadGrantSchema, + parseIngestionKeyUrl, + pickIngestionRowFields, + replaceIngestionAttachmentReferences, + sha256Hex, +} from "./ingestion-key"; import { parseTemplateFormat, renderTemplateContent, @@ -767,6 +779,8 @@ export class BraintrustState { typeof globalThis.fetch, LRUCache> >(); + private maskingFunction: ((value: unknown) => unknown) | null = null; + private readonly ingestionKeyStates = new Map(); private readonly loginParams: LoginOptions; private activeLoginOrgNameSelector: string | undefined; @@ -793,13 +807,8 @@ export class BraintrustState { this.fetch = loginParams.fetch; } - const defaultGetLogConn = async () => { - await this.login({}); - return this.apiConn(); - }; - this._bgLogger = new SyncLazyValue( - () => - new HTTPBackgroundLogger(new LazyValue(defaultGetLogConn), loginParams), + this._bgLogger = new SyncLazyValue(() => + this.createBackgroundLogger(loginParams), ); if (loginParams.debugLogLevel !== undefined) { @@ -849,6 +858,57 @@ export class BraintrustState { this._internalSetTraceContextSigningSecret(loginParams.apiKey); } + protected createBackgroundLogger( + loginParams: LoginOptions, + ): HTTPBackgroundLogger { + return new HTTPBackgroundLogger( + new LazyValue(async () => { + await this.login({}); + return this.apiConn(); + }), + loginParams, + ); + } + + /** + * Return the state that logs with an ingestion key. States are shared per + * endpoint and key, so repeated `initLogger` calls reuse one queue. + * + * @internal + */ + public _internalGetIngestionKeyState( + endpoint: IngestionEndpoint, + loginParams: Pick< + LoginOptions, + "fetch" | "noExitFlush" | "onFlushError" | "debugLogLevel" + >, + ): BraintrustState { + const cacheKey = JSON.stringify([endpoint.root, endpoint.key]); + let state = this.ingestionKeyStates.get(cacheKey); + if (!state) { + state = new IngestionKeyState(this, endpoint, { + ...loginParams, + fetch: loginParams.fetch ?? this.fetch, + // Creating a state resets the global debug log level unless it is set. + debugLogLevel: + loginParams.debugLogLevel ?? + (this.debugLogLevelConfigured + ? (this.debugLogLevel ?? false) + : undefined), + }); + if (this.maskingFunction) { + state.setMaskingFunction(this.maskingFunction); + } + this.ingestionKeyStates.set(cacheKey, state); + } else { + if (loginParams.fetch) { + state.setFetch(loginParams.fetch); + } + state.setDebugLogLevel(loginParams.debugLogLevel); + } + return state; + } + /** @internal */ public _internalSetTraceContextSigningSecret( secret: string | undefined, @@ -1188,7 +1248,11 @@ export class BraintrustState { public setMaskingFunction( maskingFunction: ((value: unknown) => unknown) | null, ): void { + this.maskingFunction = maskingFunction; this.bgLogger().setMaskingFunction(maskingFunction); + for (const state of this.ingestionKeyStates.values()) { + state.setMaskingFunction(maskingFunction); + } } public setDebugLogLevel(option: DebugLogLevelOption): void { @@ -1399,6 +1463,55 @@ export class BraintrustState { } } +/** + * State of loggers created with an ingestion key. It can only write traces to + * the key's ingestion endpoint, so it never logs in or reads private + * credentials. The span context is shared with the state that created it, so + * spans of both states nest into each other. + */ +class IngestionKeyState extends BraintrustState { + private readonly ingestionConn: HTTPConnection; + + constructor( + private readonly contextState: BraintrustState, + endpoint: IngestionEndpoint, + loginParams: LoginOptions, + ) { + super(loginParams); + this.ingestionConn = new HTTPConnection(endpoint.root, this.fetch); + this.ingestionConn.set_token(endpoint.key); + } + + public override get contextManager(): ContextManager { + return this.contextState.contextManager; + } + + public override async login(): Promise { + throw new Error( + "This logger uses a Braintrust ingestion key, which can only write traces. Use an API key for anything else.", + ); + } + + public override setFetch(fetch: typeof globalThis.fetch) { + super.setFetch(fetch); + this.ingestionConn.setFetch(fetch); + } + + protected override createBackgroundLogger( + loginParams: LoginOptions, + ): HTTPBackgroundLogger { + return new HTTPBackgroundLogger( + new LazyValue(async () => this.ingestionConn), + loginParams, + true, + ); + } +} + +function usesIngestionKey(state: BraintrustState): boolean { + return state.currentLogger?.loggingState instanceof IngestionKeyState; +} + let _globalState: BraintrustState; // Return a TestBackgroundLogger that will intercept logs before they are sent to the server. @@ -2376,6 +2489,22 @@ export class JSONAttachment extends Attachment { } } +// Ingestion key loggers may not know their project id. The server resolves +// the project from the key instead. +function parentObjectIdFields( + objectType: SpanObjectTypeV3, + objectId: string, +): ReturnType { + if (objectType === SpanObjectTypeV3.PROJECT_LOGS && !objectId) { + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + return { log_id: "g" } as ReturnType; + } + return new SpanComponentsV3({ + object_type: objectType, + object_id: objectId, + }).objectIdFields(); +} + function logFeedbackImpl( state: BraintrustState, parentObjectType: SpanObjectTypeV3, @@ -2395,6 +2524,9 @@ function logFeedbackImpl( if (!VALID_SOURCES.includes(source)) { throw new Error(`source must be one of ${VALID_SOURCES}`); } + if (!isEmpty(comment) && state instanceof IngestionKeyState) { + throw new Error("Ingestion keys do not support logging comments"); + } if ( isEmpty(scores) && @@ -2420,10 +2552,7 @@ function logFeedbackImpl( ); const parentIds = async () => - new SpanComponentsV3({ - object_type: parentObjectType, - object_id: await parentObjectId.get(), - }).objectIdFields(); + parentObjectIdFields(parentObjectType, await parentObjectId.get()); if (Object.keys(updateEvent).length > 0) { const record = new LazyValue(async () => { @@ -2492,10 +2621,7 @@ function updateSpanImpl({ ); const parentIds = async () => - new SpanComponentsV3({ - object_type: parentObjectType, - object_id: await parentObjectId.get(), - }).objectIdFields(); + parentObjectIdFields(parentObjectType, await parentObjectId.get()); const record = new LazyValue( async () => @@ -2606,6 +2732,17 @@ function spanComponentsToObjectIdLambda( const ret = components.data.object_id; return async () => ret; } + if ( + state instanceof IngestionKeyState && + components.data.object_type === SpanObjectTypeV3.PROJECT_LOGS + ) { + // Ingestion keys cannot look up projects. The server resolves the project + // from the key if the slug does not include its id. + const projectId = String( + components.data.compute_object_metadata_args?.project_id ?? "", + ); + return async () => projectId; + } if (!components.data.compute_object_metadata_args) { throw new Error( "Impossible: must provide either objectId or computeObjectMetadataArgs", @@ -3252,6 +3389,29 @@ function now() { return new Date().getTime(); } +function throwAttachmentErrors(errors: unknown[]) { + if (errors.length === 1) { + throw errors[0]; + } else if (errors.length > 1) { + throw new AggregateError( + errors, + `Encountered the following errors while uploading attachments:`, + ); + } +} + +function parseIngestionResponse(schema: z.ZodType, response: unknown) { + try { + return schema.parse(response); + } catch (error) { + if (error instanceof ZodError) { + const errorStr = JSON.stringify(error.flatten()); + throw new Error(`Invalid response from ingestion endpoint: ${errorStr}`); + } + throw error; + } +} + export interface BackgroundLoggerOpts { noExitFlush?: boolean; onFlushError?: (error: unknown) => void; @@ -3464,7 +3624,18 @@ class HTTPBackgroundLogger implements BackgroundLogger { lastLoggedTimestamp: 0, }; - constructor(apiConn: LazyValue, opts?: BackgroundLoggerOpts) { + private readonly ingestionAttachmentUploads = new WeakMap< + Attachment, + Promise + >(); + private readonly droppedIngestionFields = new Set(); + + constructor( + apiConn: LazyValue, + opts?: BackgroundLoggerOpts, + // Log through the ingestion key endpoints of `apiConn`. + private readonly usesIngestionKey = false, + ) { opts = opts ?? {}; this.apiConn = apiConn; @@ -3611,6 +3782,14 @@ class HTTPBackgroundLogger implements BackgroundLogger { maxRequestSize: number; canUseOverflow: boolean; }> { + if (!this._maxRequestSizePromise && this.usesIngestionKey) { + // Ingestion keys cannot read the server version, and the ingestion + // endpoint always accepts overflow uploads. + this._maxRequestSizePromise = Promise.resolve({ + maxRequestSize: this.maxRequestSizeOverride ?? DEFAULT_MAX_REQUEST_SIZE, + canUseOverflow: true, + }); + } if (!this._maxRequestSizePromise) { this._maxRequestSizePromise = (async () => { let serverLimit: number | null = null; @@ -3689,8 +3868,20 @@ class HTTPBackgroundLogger implements BackgroundLogger { return; } - const [allItems, attachments] = await this.unwrapLazyValues(wrappedItems); + const [unwrappedItems, attachments] = + await this.unwrapLazyValues(wrappedItems); + const attachmentErrors: unknown[] = []; + // The ingestion endpoint only accepts attachments uploaded before the rows + // that reference them. + const allItems = this.usesIngestionKey + ? await this.prepareIngestionRows( + unwrappedItems, + attachments, + attachmentErrors, + ) + : unwrappedItems; if (allItems.length === 0) { + throwAttachmentErrors(attachmentErrors); return; } @@ -3733,9 +3924,8 @@ class HTTPBackgroundLogger implements BackgroundLogger { ); } - const attachmentErrors: unknown[] = []; // For now, upload attachments serially. - for (const attachment of attachments) { + for (const attachment of this.usesIngestionKey ? [] : attachments) { try { const result = await attachment.upload(); if (result.upload_status === "error") { @@ -3745,14 +3935,185 @@ class HTTPBackgroundLogger implements BackgroundLogger { attachmentErrors.push(error); } } - if (attachmentErrors.length === 1) { - throw attachmentErrors[0]; - } else if (attachmentErrors.length > 1) { - throw new AggregateError( - attachmentErrors, - `Encountered the following errors while uploading attachments:`, + throwAttachmentErrors(attachmentErrors); + } + + private async prepareIngestionRows( + items: BackgroundLogEvent[], + attachments: BaseAttachment[], + attachmentErrors: unknown[], + ): Promise { + // Maps client attachment keys to the uploaded reference, or to undefined + // if the upload failed. + const references = new Map< + string, + BraintrustAttachmentReference | undefined + >(); + for (const attachment of attachments) { + // External attachments are not uploaded, so their reference is kept. + if ( + !(attachment instanceof Attachment) || + references.has(attachment.reference.key) + ) { + continue; + } + try { + references.set( + attachment.reference.key, + await this.uploadIngestionAttachment(attachment), + ); + } catch (error) { + references.set(attachment.reference.key, undefined); + attachmentErrors.push(error); + } + } + + const droppedFields = new Set(); + const rows: BackgroundLogEvent[] = []; + for (const item of items) { + const row = pickIngestionRowFields(item, droppedFields); + // Rows that reference a failed upload are dropped with the error. + if (replaceIngestionAttachmentReferences(row, references)) { + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + rows.push(row as BackgroundLogEvent); + } + } + for (const field of droppedFields) { + if (!this.droppedIngestionFields.has(field)) { + this.droppedIngestionFields.add(field); + debugLogger.warn( + `Ingestion keys do not support the ${field} field, so it was not logged`, + ); + } + } + return rows; + } + + private uploadIngestionAttachment( + attachment: Attachment, + ): Promise { + let upload = this.ingestionAttachmentUploads.get(attachment); + if (!upload) { + upload = (async () => { + const reference = await this.uploadToIngestionWithRetries( + { + purpose: "attachment", + filename: attachment.reference.filename, + content_type: attachment.reference.content_type, + }, + await attachment.data(), + ); + if (reference.type !== BRAINTRUST_ATTACHMENT) { + throw new Error( + "Invalid response from ingestion endpoint: expected an attachment reference", + ); + } + return reference; + })(); + this.ingestionAttachmentUploads.set(attachment, upload); + // A later flush that logs the attachment again retries the upload. + upload.catch(() => this.ingestionAttachmentUploads.delete(attachment)); + } + return upload; + } + + private async uploadToIngestionWithRetries( + request: IngestionUploadRequest, + data: Blob, + ) { + for (let i = 0; ; i++) { + try { + // Every attempt uses a new upload, so an expired grant is never reused. + return await this.uploadToIngestion(request, data); + } catch (error) { + if (i + 1 >= this.numTries) { + throw error; + } + debugLogger.warn( + `Ingestion upload failed. Retrying\nError: ${formatHTTPError(error)}`, + ); + const sleepTimeS = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * 2 ** i; + debugLogger.info(`Sleeping for ${sleepTimeS}s`); + await new Promise((resolve) => setTimeout(resolve, sleepTimeS * 1000)); + } + } + } + + private async uploadToIngestion(request: IngestionUploadRequest, data: Blob) { + const conn = await this.apiConn.get(); + const sha256 = await sha256Hex(data); + // The grant lifetime is relative, so anchor it before the request is sent. + const requestStart = now(); + const grant = parseIngestionResponse( + ingestionUploadGrantSchema, + await conn.post_json("v1/uploads", { + ...request, + size_bytes: data.size, + ...(sha256 ? { sha256 } : {}), + }), + ); + const deadline = requestStart + grant.expires_in_ms; + if (grant.num_chunks !== Math.ceil(data.size / grant.chunk_bytes)) { + throw new Error( + `Invalid response from ingestion endpoint: ${grant.num_chunks} chunks of ${grant.chunk_bytes} bytes cannot hold ${data.size} bytes`, ); } + const assertGrantActive = () => { + if (now() >= deadline) { + throw new Error("Ingestion upload expired before it completed"); + } + }; + + // On platforms like Cloudflare, fetch must not be called as a method. + const fetch = conn.fetch; + for (let index = 0; index < grant.num_chunks; index++) { + assertGrantActive(); + const chunk = data.slice( + index * grant.chunk_bytes, + (index + 1) * grant.chunk_bytes, + ); + const response = await checkResponse( + await fetch( + _urljoin( + conn.base_url, + `v1/uploads/${grant.upload_id}/chunks/${index}`, + ), + { + method: "PUT", + headers: { + ...conn.headers, + Accept: "application/json", + "Content-Type": "application/octet-stream", + }, + body: chunk, + }, + ), + ); + const ack = parseIngestionResponse( + ingestionUploadChunkSchema, + await response.json(), + ); + if (ack.index !== index || ack.size_bytes !== chunk.size) { + throw new Error( + `Invalid response from ingestion endpoint: chunk ${index} of ${chunk.size} bytes was acknowledged as chunk ${ack.index} of ${ack.size_bytes} bytes`, + ); + } + } + + assertGrantActive(); + const completed = parseIngestionResponse( + ingestionUploadCompleteSchema, + await conn.post_json(`v1/uploads/${grant.upload_id}/complete`, {}), + ); + if ( + completed.size_bytes !== data.size || + (sha256 !== undefined && completed.sha256 !== sha256) + ) { + throw new Error( + "Invalid response from ingestion endpoint: the completed upload does not match the uploaded data", + ); + } + return completed.reference; } private async unwrapLazyValues( @@ -3870,15 +4231,36 @@ class HTTPBackgroundLogger implements BackgroundLogger { } let overflowUpload: Logs3OverflowUpload | null = null; - const overflowRows = useOverflow - ? items.map((item) => item.overflowMeta) - : null; + let ingestionOverflowKey: string | null = null; + const overflowRows = + useOverflow && !this.usesIngestionKey + ? items.map((item) => item.overflowMeta) + : null; for (let i = 0; i < this.numTries; i++) { const startTime = now(); let error: unknown = undefined; try { - if (overflowRows) { + if (useOverflow && this.usesIngestionKey) { + if (!ingestionOverflowKey) { + const reference = await this.uploadToIngestion( + { purpose: "logs3_overflow", content_type: "application/json" }, + new Blob([dataStr], { type: "application/json" }), + ); + if (reference.type !== LOGS3_OVERFLOW_REFERENCE_TYPE) { + throw new Error( + "Invalid response from ingestion endpoint: expected an overflow reference", + ); + } + ingestionOverflowKey = reference.key; + } + await conn.post_json( + "v1/logs", + constructLogs3OverflowRequest(ingestionOverflowKey), + ); + } else if (this.usesIngestionKey) { + await conn.post_json("v1/logs", dataStr); + } else if (overflowRows) { if (!overflowUpload) { const currentUpload = await this.requestLogs3OverflowUpload(conn, { rows: overflowRows, @@ -5162,6 +5544,7 @@ export type InitLoggerOptions = FullLoginOptions & { setCurrent?: boolean; state?: BraintrustState; orgProjectMetadata?: OrgProjectMetadata; + ingestionKey?: string; } & AsyncFlushArg; /** @@ -5178,12 +5561,31 @@ export type InitLoggerOptions = FullLoginOptions & { * @param options.orgName (Optional) The name of a specific organization to connect to. This is useful if you belong to multiple. * @param options.forceLogin Login again, even if you have already logged in (by default, the logger will not login if you are already logged in) * @param options.debugLogLevel Enables internal Braintrust SDK troubleshooting output. Use `"error"`, `"warn"`, `"info"`, or `"debug"` to choose an explicit level, or `false` to explicitly disable it. If omitted, the SDK stays silent unless `BRAINTRUST_DEBUG_LOG_LEVEL` is set. + * @param options.ingestionKey The ingestion URL of a project's ingestion key, like `https:///ingest?ingestKey=`. Ingestion keys can only write traces to their project, so they are safe to use in browsers and other public clients. The logger never logs in, ignores `BRAINTRUST_API_KEY` and does not create or look up the project. `projectName` is ignored, and `projectId` is only sent along for the server to check. If unspecified, will use the `BRAINTRUST_INGESTION_KEY` environment variable unless `apiKey` or `state` is specified. Cannot be combined with `apiKey` or `state`. * @param setCurrent If true (the default), set the global current-experiment to the newly-created one. * @returns The newly created Logger. */ export function initLogger( options: Readonly> = {}, ) { + if ( + options.ingestionKey !== undefined && + (options.apiKey !== undefined || options.state !== undefined) + ) { + throw new Error( + "initLogger accepts either an ingestionKey or an apiKey/state, not both", + ); + } + // An explicit private credential takes precedence over the environment. + const ingestionKey = + options.ingestionKey ?? + (options.apiKey === undefined && options.state === undefined + ? iso.getEnv(INGESTION_KEY_ENV_VAR) || undefined + : undefined); + if (ingestionKey !== undefined) { + return initIngestionKeyLogger(ingestionKey, options); + } + const { projectName, projectId, @@ -5248,6 +5650,45 @@ export function initLogger( return ret; } +function initIngestionKeyLogger( + ingestionKey: string, + { + projectId, + asyncFlush = true as IsAsyncFlush, + appUrl, + orgName, + environment, + setCurrent, + fetch, + noExitFlush, + onFlushError, + debugLogLevel, + }: Readonly>, +) { + const state = _globalState._internalGetIngestionKeyState( + parseIngestionKeyUrl(ingestionKey), + { fetch, noExitFlush, onFlushError, debugLogLevel }, + ); + state.spanOriginEnvironment = detectSpanOriginEnvironment(environment); + state.enforceQueueSizeLimit(true); + + // The project is only known if it is passed in, and is never looked up. + const lazyMetadata = new LazyValue(async () => ({ + org_id: "", + project: { id: projectId ?? "", name: "", fullInfo: {} }, + })); + const ret = new Logger(state, lazyMetadata, { + asyncFlush, + computeMetadataArgs: projectId ? { project_id: projectId } : undefined, + linkArgs: { org_name: orgName, app_url: appUrl, project_id: projectId }, + }); + state.currentLogger = ret as Logger; + if (setCurrent ?? true) { + _globalState.currentLogger = ret as Logger; + } + return ret; +} + export type LoadPromptOptions = FullLoginOptions & { projectName?: string; projectId?: string; @@ -6555,7 +6996,17 @@ function resolveW3cParent( if (!braintrustParent) { braintrustParent = currentBraintrustParent(state); } - if (!braintrustParent) { + // Ingestion key loggers write to their own project, which they may not + // know, so the trace ids are enough to link their spans. + const ingestionKeyParent = + !braintrustParent && usesIngestionKey(state ?? _globalState) + ? { + objectType: SpanObjectTypeV3.PROJECT_LOGS, + objectId: undefined, + computeArgs: undefined, + } + : undefined; + if (!braintrustParent && !ingestionKeyParent) { debugLogger.warn( "Received traceparent without a braintrust.parent and no active logger/experiment; " + "cannot route the trace. Starting a fresh local span instead.", @@ -6563,7 +7014,8 @@ function resolveW3cParent( return { parentSlug: undefined, propagatedState: undefined }; } - const parsedParent = braintrustParentToComponents(braintrustParent); + const parsedParent = + ingestionKeyParent ?? braintrustParentToComponents(braintrustParent); if (parsedParent === undefined) { debugLogger.warn( `Invalid braintrust.parent: ${JSON.stringify(braintrustParent)}`, @@ -7035,7 +7487,14 @@ export function _internalStartSpanWithContext< */ export async function flush(options?: OptionalStateArg): Promise { const state = options?.state ?? _globalState; - return await state.bgLogger().flush(); + // Ingestion key loggers have their own queue. + const loggerState = state.currentLogger?.loggingState; + await Promise.all([ + state.bgLogger().flush(), + loggerState && loggerState !== state + ? loggerState.bgLogger().flush() + : undefined, + ]); } /** @@ -7067,6 +7526,26 @@ function startSpanAndIsLogger( state, }); + const ingestionKeyLogger = usesIngestionKey(state) + ? state.currentLogger + : undefined; + if ( + ingestionKeyLogger && + (parentObject instanceof SpanComponentsV3 || + parentObject instanceof SpanComponentsV4) && + parentObject.data.object_type === SpanObjectTypeV3.PROJECT_LOGS + ) { + // Ingestion key loggers can only write to their own project, so a + // propagated parent only links the span into its trace. + return { + span: ingestionKeyLogger.startSpan({ + ...args, + parent: parentObject.toStr(), + }), + isSyncFlushLogger: ingestionKeyLogger.asyncFlush === false, + }; + } + if ( parentObject instanceof SpanComponentsV3 || parentObject instanceof SpanComponentsV4 @@ -8530,10 +9009,10 @@ export class SpanImpl implements Span { ]), ), ), - ...new SpanComponentsV3({ - object_type: this.parentObjectType, - object_id: await this.parentObjectId.get(), - }).objectIdFields(), + ...parentObjectIdFields( + this.parentObjectType, + await this.parentObjectId.get(), + ), }; // Customize inside the memoized lazy value, before attachment processing, // merging, and masking. Retries reuse the customized record or drop result. From 93bc106944788b95f48d38504a37654a68e3b169 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:35:17 +0000 Subject: [PATCH 2/5] fix: Address ingestion key review feedback - Batches above 512 KiB overflow through an upload by default, since public clients often sit behind 1 MiB ingress limits and Lambda sees base64 bodies. - Parse the ingestion URL with `URLSearchParams`, require the exact 48 char key and trim all trailing slashes, matching the shared contract. - Drop the optional upload `sha256`, which read the whole upload into memory. - A `state` only counts as a private credential if it is logged in or has an `apiKey`. Otherwise it provides masking and the current logger. - Don't retry auth and validation failures, and honor `Retry-After`. - `flush()` drains every ingestion key queue of the state, not only the one of the current logger. - Document that ingestion keys can update any row whose id is known. --- js/README.md | 7 +- js/src/ingestion-key.test.ts | 197 +++++++++++++++++++++++++++++++++-- js/src/ingestion-key.ts | 28 ++--- js/src/logger.ts | 132 ++++++++++++++++------- 4 files changed, 294 insertions(+), 70 deletions(-) diff --git a/js/README.md b/js/README.md index a1e8c6e9d..af2c53517 100644 --- a/js/README.md +++ b/js/README.md @@ -119,7 +119,12 @@ const logger = initLogger({ }); ``` -An explicit `apiKey` or `state` takes precedence over `BRAINTRUST_INGESTION_KEY`, and passing either together with `ingestionKey` throws. +Anyone with the key can write new rows into the project, and can update or delete any row of the project whose id they know, including rows logged with an API key. +Row ids are not secret, so treat everything in the project as writable by the key's users. +Uploaded attachments can only be referenced by rows written with the same key or a key it was rotated from. + +An API key passed as `apiKey`, or through a `state` that is logged in or has an `apiKey`, takes precedence over `BRAINTRUST_INGESTION_KEY`, and combining one with `ingestionKey` throws. +A `state` without an API key only provides the span context and masking function, while the logger keeps its own queue. `projectName` is ignored, and `projectId` is only sent along for the server to check. Attachments and large batches are uploaded through the ingestion URL as well. Fields that ingestion keys do not support, like `classifications` or audit fields, are not logged, and logging comments throws. diff --git a/js/src/ingestion-key.test.ts b/js/src/ingestion-key.test.ts index cfc286068..6e4232ccc 100644 --- a/js/src/ingestion-key.test.ts +++ b/js/src/ingestion-key.test.ts @@ -189,10 +189,16 @@ describe("parseIngestionKeyUrl", () => { parseIngestionKeyUrl(`https://dp.example/a/b/ingest?ingestKey=${key}`), ).toEqual({ root: "https://dp.example/a/b/ingest", key }); expect( - parseIngestionKeyUrl(`http://localhost:8000/ingest/?ingestKey=${key}`), + parseIngestionKeyUrl(`http://localhost:8000/ingest//?ingestKey=${key}`), ).toEqual({ root: "http://localhost:8000/ingest", key }); }); + test("decodes the query like URLSearchParams", () => { + expect( + parseIngestionKeyUrl(`https://dp.example/ingest?ingest%4Bey=${key}`), + ).toEqual({ root: "https://dp.example/ingest", key }); + }); + test.each([ ["a bare key", key], ["a non-http URL", `ftp://dp.example/ingest?ingestKey=${key}`], @@ -208,6 +214,10 @@ describe("parseIngestionKeyUrl", () => { `https://dp.example/ingest?ingestKey=${key}&ingestKey=${key}`, ], ["a malformed key", "https://dp.example/ingest?ingestKey=sk-123"], + [ + "a key of the wrong length", + `https://dp.example/ingest?ingestKey=${key.slice(0, -1)}`, + ], ["a different path", `https://dp.example/logs?ingestKey=${key}`], ])("rejects %s without echoing the key", (_, value) => { let error: unknown; @@ -328,17 +338,55 @@ describe("initLogger with an ingestion key", () => { test("rejects an ingestion key combined with private credentials", () => { const ingestionKey = ingestionUrl(newKey()); - expect(() => initLogger({ ingestionKey, apiKey: "private" })).toThrow( - "either an ingestionKey or an apiKey/state", - ); - expect(() => - initLogger({ ingestionKey, state: new BraintrustState({}) }), - ).toThrow("either an ingestionKey or an apiKey/state"); + const loggedInState = new BraintrustState({}); + loggedInState.loginToken = "private"; + for (const options of [ + { apiKey: "private" }, + { state: new BraintrustState({ apiKey: "private" }) }, + { state: loggedInState }, + ]) { + expect(() => initLogger({ ingestionKey, ...options })).toThrow( + "either an ingestionKey or an API key", + ); + } expect(() => initLogger({ ingestionKey: "" })).toThrow( "Invalid Braintrust ingestion key", ); }); + test("uses a state without credentials for masking and the current logger", async () => { + const { requests, fetch } = mockIngestion(); + vi.stubEnv("BRAINTRUST_INGESTION_KEY", ingestionUrl(newKey())); + const contextState = new BraintrustState({ noExitFlush: true }); + contextState.setMaskingFunction((value) => + value === "secret-value" ? "[masked]" : value, + ); + + const logger = initLogger({ + state: contextState, + fetch, + noExitFlush: true, + }); + logger.traced( + (root) => + root.traced((span) => span.log({ input: "secret-value" }), { + name: "child", + }), + { name: "root" }, + ); + await flush({ state: contextState }); + + expect(logger.loggingState).not.toBe(contextState); + expect(contextState.currentLogger).toBe(logger); + expect(_internalGetGlobalState().currentLogger).toBeUndefined(); + expect(paths(requests)).toEqual(["POST /deployment/base/ingest/v1/logs"]); + const rows = Object.fromEntries( + loggedRows(requests).map((row) => [row.span_attributes.name, row]), + ); + expect(rows.child.input).toBe("[masked]"); + expect(rows.child.span_parents).toEqual([rows.root.span_id]); + }); + test("public and private loggers do not share queues or credentials", async () => { const { requests, fetch } = mockIngestion({ respond: (request) => @@ -385,18 +433,26 @@ describe("initLogger with an ingestion key", () => { ).toBe(true); }); - test("retries rejected writes against the ingestion endpoint only", async () => { + test("retries throttled writes after Retry-After against the ingestion endpoint only", async () => { const { requests, fetch } = mockIngestion({ respond: (request, index) => - index === 0 ? json({ error: "slow down" }, 429) : undefined, + index === 0 + ? new Response("slow down", { + status: 429, + headers: { "Retry-After": "0" }, + }) + : undefined, }); vi.stubEnv("BRAINTRUST_NUM_RETRIES", "1"); vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); const { logger, onFlushError } = initIngestionLogger(fetch); logger.log({ input: "a", output: "b" }); + const start = Date.now(); await logger.flush(); + // Without Retry-After, the first retry waits one second. + expect(Date.now() - start).toBeLessThan(900); expect(onFlushError).not.toHaveBeenCalled(); expect(paths(requests)).toEqual([ "POST /deployment/base/ingest/v1/logs", @@ -404,6 +460,68 @@ describe("initLogger with an ingestion key", () => { ]); }); + test("retries transient server errors", async () => { + const { requests, fetch } = mockIngestion({ + respond: (request, index) => + index === 0 ? json({ error: "unavailable" }, 503) : undefined, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "1"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(loggedRows(requests.slice(1))).toEqual([ + expect.objectContaining({ input: "a" }), + ]); + }); + + test.each([401, 403, 400])( + "does not retry a write rejected with %i", + async (status) => { + const { requests, fetch } = mockIngestion({ + respond: () => json({ error: "rejected" }, status), + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "2"); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(paths(requests)).toEqual(["POST /deployment/base/ingest/v1/logs"]); + }, + ); + + test("does not request new upload grants after a permanent rejection", async () => { + const { requests, fetch } = mockIngestion({ + respond: (request) => + request.url.includes("/chunks/") + ? json({ error: "forbidden" }, 403) + : undefined, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "2"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "dropped", + }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(paths(requests).filter((p) => p.endsWith("/v1/uploads"))).toEqual([ + "POST /deployment/base/ingest/v1/uploads", + ]); + expect(loggedRows(requests)).toEqual([]); + }); + test("gives up after the configured retries without a private fallback", async () => { const { requests, fetch } = mockIngestion({ respond: () => json({ error: "unavailable" }, 503), @@ -447,7 +565,6 @@ describe("initLogger with an ingestion key", () => { filename: "data.bin", content_type: "application/octet-stream", size_bytes: data.length, - sha256: createHash("sha256").update(data).digest("hex"), }); expect( requests @@ -588,6 +705,34 @@ describe("initLogger with an ingestion key", () => { ).toBe(true); }); + test("overflows payloads above 512 KiB by default", async () => { + const { requests, uploads, fetch } = mockIngestion(); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ input: "x".repeat(400 * KIB), output: "inline" }); + await logger.flush(); + expect(uploads.size).toBe(0); + expect(loggedRows(requests)).toEqual([ + expect.objectContaining({ output: "inline" }), + ]); + + logger.log({ input: "x".repeat(600 * KIB), output: "overflow" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + const [uploadId] = uploads.keys(); + expect(paths(requests).slice(1)).toEqual([ + "POST /deployment/base/ingest/v1/uploads", + `PUT /deployment/base/ingest/v1/uploads/${uploadId}/chunks/0`, + `PUT /deployment/base/ingest/v1/uploads/${uploadId}/chunks/1`, + `POST /deployment/base/ingest/v1/uploads/${uploadId}/complete`, + "POST /deployment/base/ingest/v1/logs", + ]); + expect( + requests.every((r) => r.body === undefined || r.body.length <= 512 * KIB), + ).toBe(true); + }); + test("only sends supported row fields", async () => { const { requests, fetch } = mockIngestion(); const { logger } = initIngestionLogger(fetch); @@ -640,10 +785,36 @@ describe("initLogger with an ingestion key", () => { }); }); +describe("flush with ingestion keys", () => { + test("drains every ingestion key queue, not only the current logger", async () => { + const { requests, fetch } = mockIngestion(); + const first = initIngestionLogger(fetch); + const second = initIngestionLogger(fetch); + first.logger.log({ input: "first", output: "x" }); + second.logger.log({ input: "second", output: "x" }); + const privateLogger = initLogger({ apiKey: "private-api-key" }); + expect(_internalGetGlobalState().currentLogger).toBe(privateLogger); + + await flush(); + + expect( + Object.fromEntries( + requests.map((r) => [ + r.headers.authorization, + loggedRows([r])[0].input, + ]), + ), + ).toEqual({ + [`Bearer ${first.key}`]: "first", + [`Bearer ${second.key}`]: "second", + }); + }); +}); + describe("tracing with an ingestion key", () => { test("nests and propagates spans without project metadata", async () => { const { requests, fetch } = mockIngestion(); - const { logger, onFlushError } = initIngestionLogger(fetch); + const { key, logger, onFlushError } = initIngestionLogger(fetch); let exported = ""; let headers: Record = {}; @@ -660,6 +831,10 @@ describe("tracing with an ingestion key", () => { { name: "root" }, ); await vi.waitFor(() => expect(exported).not.toBe("")); + // Exporting and injecting neither logs in nor looks up the project. + expect(requests).toHaveLength(0); + expect(exported).not.toContain(key); + expect(JSON.stringify(headers)).not.toContain(key); const remote = logger.startSpan({ name: "remote", diff --git a/js/src/ingestion-key.ts b/js/src/ingestion-key.ts index cd8473b4b..d7c3007e0 100644 --- a/js/src/ingestion-key.ts +++ b/js/src/ingestion-key.ts @@ -18,7 +18,7 @@ export interface IngestionEndpoint { key: string; } -const INGESTION_KEY_QUERY = /^\?ingestKey=(bt-ik-[A-Za-z0-9]+)$/; +const INGESTION_KEY_PATTERN = /^bt-ik-[A-Za-z0-9]{48}$/; /** * Parse an ingestion key URL such as @@ -48,17 +48,21 @@ export function parseIngestionKeyUrl(value: string): IngestionEndpoint { if (url.hash || value.includes("#")) { throw invalid("the URL must not contain a fragment"); } - const match = INGESTION_KEY_QUERY.exec(url.search); - if (!match) { + const params = [...url.searchParams]; + if (params.length !== 1 || params[0][0] !== "ingestKey") { throw invalid( "the URL must contain exactly one ingestKey query parameter and no other query parameters", ); } - const pathname = url.pathname.replace(/\/$/, ""); + const key = params[0][1]; + if (!INGESTION_KEY_PATTERN.test(key)) { + throw invalid("the ingestKey query parameter is not an ingestion key"); + } + const pathname = url.pathname.replace(/\/+$/, ""); if (!pathname.endsWith("/ingest")) { throw invalid("the URL path must end with /ingest"); } - return { root: `${url.origin}${pathname}`, key: match[1] }; + return { root: `${url.origin}${pathname}`, key }; } // Row fields the ingestion endpoint accepts. Everything else (experiment and @@ -178,17 +182,3 @@ export const ingestionUploadCompleteSchema = z.object({ size_bytes: z.number().int().nonnegative(), sha256: sha256Schema, }); - -/** - * Hex SHA-256 digest of `data`, or undefined on platforms without WebCrypto. - */ -export async function sha256Hex(data: Blob): Promise { - const subtle = globalThis.crypto?.subtle; - if (!subtle) { - return undefined; - } - const digest = await subtle.digest("SHA-256", await data.arrayBuffer()); - return Array.from(new Uint8Array(digest), (byte) => - byte.toString(16).padStart(2, "0"), - ).join(""); -} diff --git a/js/src/logger.ts b/js/src/logger.ts index 26d20d334..cd041642f 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -108,6 +108,9 @@ const RESET_CONTEXT_MANAGER_STATE = Symbol.for( ); // 6 MB for the AWS lambda gateway (from our own testing). export const DEFAULT_MAX_REQUEST_SIZE = 6 * 1024 * 1024; +// Public clients often sit behind 1 MiB ingress limits, and Lambda data planes +// see base64 encoded bodies, so larger batches overflow through an upload. +const INGESTION_KEY_MAX_REQUEST_SIZE = 512 * 1024; export type { DatasetSnapshot }; @@ -172,7 +175,6 @@ import { parseIngestionKeyUrl, pickIngestionRowFields, replaceIngestionAttachmentReferences, - sha256Hex, } from "./ingestion-key"; import { parseTemplateFormat, @@ -909,6 +911,29 @@ export class BraintrustState { return state; } + /** @internal */ + public async _internalFlushIngestionKeyStates(): Promise { + await Promise.all( + [...this.ingestionKeyStates.values()].map((state) => + state.bgLogger().flush(), + ), + ); + } + + /** + * Whether this state was given or logged in with an API key, as opposed to + * only picking one up from the environment on login. + * + * @internal + */ + public _internalHasPrivateCredential(): boolean { + return ( + this.loggedIn || + this.loginToken !== null || + this.loginParams.apiKey !== undefined + ); + } + /** @internal */ public _internalSetTraceContextSigningSecret( secret: string | undefined, @@ -1605,6 +1630,8 @@ export class FailedHTTPResponse extends Error { public text: string; public data: string; public readonly cause?: unknown; + /** @internal The `Retry-After` header of the response. */ + public retryAfter?: string; constructor(status: number, text: string, data: string, cause?: unknown) { super(`${status}: ${text} (${data})`); @@ -1692,7 +1719,9 @@ async function checkResponse(resp: Response) { error, ); } - throw new FailedHTTPResponse(resp.status, resp.statusText, data); + const error = new FailedHTTPResponse(resp.status, resp.statusText, data); + error.retryAfter = resp.headers.get("Retry-After") ?? undefined; + throw error; } class HTTPConnection { @@ -3507,6 +3536,38 @@ const HTTP_RETRY_JITTER_MS = 200; const BTQL_HTTP_RETRIES = 3; const RETRYABLE_HTTP_STATUS_CODES = new Set([500, 502, 503, 504]); +const MAX_INGESTION_RETRY_AFTER_MS = 60_000; + +// Ingestion key requests are rejected permanently on auth or validation +// failures, so only transient failures are retried. Returns undefined if the +// request must not be retried. +function ingestionRetryDelayMs( + error: unknown, + attempt: number, +): number | undefined { + const backoffMs = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * 1000 * 2 ** attempt; + if (!(error instanceof FailedHTTPResponse)) { + return backoffMs; + } + if ( + error.status !== 408 && + error.status !== 429 && + !RETRYABLE_HTTP_STATUS_CODES.has(error.status) + ) { + return undefined; + } + if (!error.retryAfter) { + return backoffMs; + } + const seconds = Number(error.retryAfter); + const retryAfterMs = isNaN(seconds) + ? Date.parse(error.retryAfter) - now() + : seconds * 1000; + return isNaN(retryAfterMs) + ? backoffMs + : Math.min(Math.max(retryAfterMs, 0), MAX_INGESTION_RETRY_AFTER_MS); +} + function isRetryableHTTPError(error: unknown): boolean { return ( !(error instanceof FailedHTTPResponse) || @@ -3786,7 +3847,8 @@ class HTTPBackgroundLogger implements BackgroundLogger { // Ingestion keys cannot read the server version, and the ingestion // endpoint always accepts overflow uploads. this._maxRequestSizePromise = Promise.resolve({ - maxRequestSize: this.maxRequestSizeOverride ?? DEFAULT_MAX_REQUEST_SIZE, + maxRequestSize: + this.maxRequestSizeOverride ?? INGESTION_KEY_MAX_REQUEST_SIZE, canUseOverflow: true, }); } @@ -4026,31 +4088,28 @@ class HTTPBackgroundLogger implements BackgroundLogger { // Every attempt uses a new upload, so an expired grant is never reused. return await this.uploadToIngestion(request, data); } catch (error) { - if (i + 1 >= this.numTries) { + const retryDelayMs = ingestionRetryDelayMs(error, i); + if (i + 1 >= this.numTries || retryDelayMs === undefined) { throw error; } debugLogger.warn( `Ingestion upload failed. Retrying\nError: ${formatHTTPError(error)}`, ); - const sleepTimeS = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * 2 ** i; - debugLogger.info(`Sleeping for ${sleepTimeS}s`); - await new Promise((resolve) => setTimeout(resolve, sleepTimeS * 1000)); + debugLogger.info(`Sleeping for ${retryDelayMs / 1000}s`); + await new Promise((resolve) => setTimeout(resolve, retryDelayMs)); } } } private async uploadToIngestion(request: IngestionUploadRequest, data: Blob) { const conn = await this.apiConn.get(); - const sha256 = await sha256Hex(data); // The grant lifetime is relative, so anchor it before the request is sent. const requestStart = now(); const grant = parseIngestionResponse( ingestionUploadGrantSchema, - await conn.post_json("v1/uploads", { - ...request, - size_bytes: data.size, - ...(sha256 ? { sha256 } : {}), - }), + // The optional sha256 is omitted, since computing it would read the + // whole upload into memory at once. + await conn.post_json("v1/uploads", { ...request, size_bytes: data.size }), ); const deadline = requestStart + grant.expires_in_ms; if (grant.num_chunks !== Math.ceil(data.size / grant.chunk_bytes)) { @@ -4105,10 +4164,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { ingestionUploadCompleteSchema, await conn.post_json(`v1/uploads/${grant.upload_id}/complete`, {}), ); - if ( - completed.size_bytes !== data.size || - (sha256 !== undefined && completed.sha256 !== sha256) - ) { + if (completed.size_bytes !== data.size) { throw new Error( "Invalid response from ingestion endpoint: the completed upload does not match the uploaded data", ); @@ -4287,7 +4343,10 @@ class HTTPBackgroundLogger implements BackgroundLogger { return; } - const isRetrying = i + 1 < this.numTries; + const retryDelayMs = this.usesIngestionKey + ? ingestionRetryDelayMs(error, i) + : BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * 1000 * 2 ** i; + const isRetrying = i + 1 < this.numTries && retryDelayMs !== undefined; const retryingText = isRetrying ? "" : " Retrying"; const errorText = (() => { if (error instanceof FailedHTTPResponse) { @@ -4318,11 +4377,8 @@ class HTTPBackgroundLogger implements BackgroundLogger { } else { debugLogger.warn(errMsg); if (isRetrying) { - const sleepTimeS = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * 2 ** i; - debugLogger.info(`Sleeping for ${sleepTimeS}s`); - await new Promise((resolve) => - setTimeout(resolve, sleepTimeS * 1000), - ); + debugLogger.info(`Sleeping for ${retryDelayMs / 1000}s`); + await new Promise((resolve) => setTimeout(resolve, retryDelayMs)); } } } @@ -5561,27 +5617,27 @@ export type InitLoggerOptions = FullLoginOptions & { * @param options.orgName (Optional) The name of a specific organization to connect to. This is useful if you belong to multiple. * @param options.forceLogin Login again, even if you have already logged in (by default, the logger will not login if you are already logged in) * @param options.debugLogLevel Enables internal Braintrust SDK troubleshooting output. Use `"error"`, `"warn"`, `"info"`, or `"debug"` to choose an explicit level, or `false` to explicitly disable it. If omitted, the SDK stays silent unless `BRAINTRUST_DEBUG_LOG_LEVEL` is set. - * @param options.ingestionKey The ingestion URL of a project's ingestion key, like `https:///ingest?ingestKey=`. Ingestion keys can only write traces to their project, so they are safe to use in browsers and other public clients. The logger never logs in, ignores `BRAINTRUST_API_KEY` and does not create or look up the project. `projectName` is ignored, and `projectId` is only sent along for the server to check. If unspecified, will use the `BRAINTRUST_INGESTION_KEY` environment variable unless `apiKey` or `state` is specified. Cannot be combined with `apiKey` or `state`. + * @param options.ingestionKey The ingestion URL of a project's ingestion key, like `https:///ingest?ingestKey=`. Ingestion keys can only write traces into their project, so the logger never logs in, ignores `BRAINTRUST_API_KEY` and does not create or look up the project. `projectName` is ignored, and `projectId` is only sent along for the server to check. A `state` without an API key is used for the span context and masking, while the logger keeps its own queue. If unspecified, will use the `BRAINTRUST_INGESTION_KEY` environment variable unless an API key is passed through `apiKey` or `state`. Cannot be combined with an API key. * @param setCurrent If true (the default), set the global current-experiment to the newly-created one. * @returns The newly created Logger. */ export function initLogger( options: Readonly> = {}, ) { - if ( - options.ingestionKey !== undefined && - (options.apiKey !== undefined || options.state !== undefined) - ) { + const hasPrivateCredential = + options.apiKey !== undefined || + (options.state?._internalHasPrivateCredential() ?? false); + if (options.ingestionKey !== undefined && hasPrivateCredential) { throw new Error( - "initLogger accepts either an ingestionKey or an apiKey/state, not both", + "initLogger accepts either an ingestionKey or an API key, not both. A `state` that is logged in or has an `apiKey` counts as an API key.", ); } // An explicit private credential takes precedence over the environment. const ingestionKey = options.ingestionKey ?? - (options.apiKey === undefined && options.state === undefined - ? iso.getEnv(INGESTION_KEY_ENV_VAR) || undefined - : undefined); + (hasPrivateCredential + ? undefined + : iso.getEnv(INGESTION_KEY_ENV_VAR) || undefined); if (ingestionKey !== undefined) { return initIngestionKeyLogger(ingestionKey, options); } @@ -5663,9 +5719,10 @@ function initIngestionKeyLogger( noExitFlush, onFlushError, debugLogLevel, + state: contextState = _globalState, }: Readonly>, ) { - const state = _globalState._internalGetIngestionKeyState( + const state = contextState._internalGetIngestionKeyState( parseIngestionKeyUrl(ingestionKey), { fetch, noExitFlush, onFlushError, debugLogLevel }, ); @@ -5684,7 +5741,7 @@ function initIngestionKeyLogger( }); state.currentLogger = ret as Logger; if (setCurrent ?? true) { - _globalState.currentLogger = ret as Logger; + contextState.currentLogger = ret as Logger; } return ret; } @@ -7487,13 +7544,10 @@ export function _internalStartSpanWithContext< */ export async function flush(options?: OptionalStateArg): Promise { const state = options?.state ?? _globalState; - // Ingestion key loggers have their own queue. - const loggerState = state.currentLogger?.loggingState; + // Ingestion key loggers have their own queues. await Promise.all([ state.bgLogger().flush(), - loggerState && loggerState !== state - ? loggerState.bgLogger().flush() - : undefined, + state._internalFlushIngestionKeyStates(), ]); } From 088d012fa501c57427240f9e9093901c02fa003f Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:49:24 +0000 Subject: [PATCH 3/5] fix: Reject external attachments with ingestion keys The public server path rejects `external_attachment` references, since they can point at any object store content. Rows that contain an `ExternalAttachment` or a raw external reference are now dropped before they are sent, and the error goes to `onFlushError`. Managed attachments and already committed `braintrust_attachment` references keep working, and the private path is unchanged. The ids of `orgProjectMetadata` are now sent along for the server to check instead of being dropped. `state` only provides context, so its login is ignored and only an explicit `apiKey` conflicts with `ingestionKey`. --- js/README.md | 7 +- js/src/ingestion-key.test.ts | 169 ++++++++++++++++++++++++++++++++--- js/src/ingestion-key.ts | 21 +++-- js/src/logger.ts | 111 ++++++++++++++--------- 4 files changed, 243 insertions(+), 65 deletions(-) diff --git a/js/README.md b/js/README.md index af2c53517..7a6f51d56 100644 --- a/js/README.md +++ b/js/README.md @@ -123,9 +123,10 @@ Anyone with the key can write new rows into the project, and can update or delet Row ids are not secret, so treat everything in the project as writable by the key's users. Uploaded attachments can only be referenced by rows written with the same key or a key it was rotated from. -An API key passed as `apiKey`, or through a `state` that is logged in or has an `apiKey`, takes precedence over `BRAINTRUST_INGESTION_KEY`, and combining one with `ingestionKey` throws. -A `state` without an API key only provides the span context and masking function, while the logger keeps its own queue. -`projectName` is ignored, and `projectId` is only sent along for the server to check. +An explicit `apiKey` takes precedence over `BRAINTRUST_INGESTION_KEY`, and combining it with `ingestionKey` throws. +A `state` only provides the masking function and current logger, while the logger keeps its own queue and ignores any login of the state. +`projectName` is ignored, and `projectId` or the ids of `orgProjectMetadata` are only sent along for the server to check. +`ExternalAttachment` is not supported, since it can reference any object store content, so rows that contain one are not logged and the error is reported to `onFlushError`. Attachments and large batches are uploaded through the ingestion URL as well. Fields that ingestion keys do not support, like `classifications` or audit fields, are not logged, and logging comments throws. diff --git a/js/src/ingestion-key.test.ts b/js/src/ingestion-key.test.ts index 6e4232ccc..7aa61ee80 100644 --- a/js/src/ingestion-key.test.ts +++ b/js/src/ingestion-key.test.ts @@ -5,6 +5,7 @@ import { _internalGetGlobalState, Attachment, BraintrustState, + ExternalAttachment, extractTraceContextFromHeaders, flush, initLogger, @@ -133,6 +134,10 @@ function mockIngestion( }); } if (request.method === "POST" && path.endsWith("/v1/logs")) { + // The public server path rejects external attachment references. + if (String(request.body).includes('"external_attachment"')) { + return json({ error: "external attachments are not supported" }, 400); + } const payload = bodyJson(request); return json({ ids: Array.isArray(payload.rows) @@ -310,7 +315,7 @@ describe("initLogger with an ingestion key", () => { globalState.resetLoginInfo(); }); - test("an explicit private credential takes precedence over the environment", async () => { + test("an explicit apiKey takes precedence over the environment", async () => { const { requests, fetch } = mockIngestion({ respond: () => json({}) }); const key = newKey(); vi.stubEnv("BRAINTRUST_INGESTION_KEY", ingestionUrl(key)); @@ -321,6 +326,7 @@ describe("initLogger with an ingestion key", () => { const logger = initLogger({ state, + apiKey: "private-api-key", projectId: "private-project", projectName: "private", }); @@ -336,24 +342,82 @@ describe("initLogger with an ingestion key", () => { expect(apiKeyLogger.loggingState).toBe(_internalGetGlobalState()); }); - test("rejects an ingestion key combined with private credentials", () => { + test("ignores the login of a supplied state", async () => { + const { requests, fetch } = mockIngestion(); + const loggedInState = new BraintrustState({ fetch, noExitFlush: true }); + loggedInState.apiUrl = "https://api.test"; + loggedInState.loginToken = "private-api-key"; + loggedInState.loggedIn = true; + const envKey = newKey(); + vi.stubEnv("BRAINTRUST_INGESTION_KEY", ingestionUrl(envKey)); + + const fromEnv = initLogger({ state: loggedInState, noExitFlush: true }); + const explicit = initIngestionLogger(fetch, { state: loggedInState }); + fromEnv.log({ input: "env", output: "x" }); + explicit.logger.log({ input: "explicit", output: "x" }); + await flush({ state: loggedInState }); + + expect(fromEnv.loggingState).not.toBe(loggedInState); + expect(loggedInState.currentLogger).toBe(explicit.logger); + expect( + requests.map((r) => [r.headers.authorization, loggedRows([r])[0].input]), + ).toEqual( + expect.arrayContaining([ + [`Bearer ${envKey}`, "env"], + [`Bearer ${explicit.key}`, "explicit"], + ]), + ); + expect(requests).toHaveLength(2); + expect(JSON.stringify(requests)).not.toContain("private-api-key"); + }); + + test("rejects an ingestion key combined with an apiKey", () => { const ingestionKey = ingestionUrl(newKey()); - const loggedInState = new BraintrustState({}); - loggedInState.loginToken = "private"; - for (const options of [ - { apiKey: "private" }, - { state: new BraintrustState({ apiKey: "private" }) }, - { state: loggedInState }, - ]) { - expect(() => initLogger({ ingestionKey, ...options })).toThrow( - "either an ingestionKey or an API key", - ); - } + expect(() => initLogger({ ingestionKey, apiKey: "private" })).toThrow( + "either an ingestionKey or an apiKey", + ); expect(() => initLogger({ ingestionKey: "" })).toThrow( "Invalid Braintrust ingestion key", ); }); + test("sends the ids of orgProjectMetadata for the server to check", async () => { + const { requests, fetch } = mockIngestion(); + const orgProjectMetadata = { + org_id: "org-id", + project: { id: "project-id", name: "ignored", fullInfo: {} }, + }; + const { logger, onFlushError } = initIngestionLogger(fetch, { + orgProjectMetadata, + }); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "x", + }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(bodyJson(requests[0])).toMatchObject({ + purpose: "attachment", + org_id: "org-id", + }); + expect(loggedRows(requests)).toEqual([ + expect.objectContaining({ org_id: "org-id", project_id: "project-id" }), + ]); + expect(JSON.stringify(requests)).not.toContain("ignored"); + expect(() => + initIngestionLogger(fetch, { + orgProjectMetadata, + projectId: "other-project", + }), + ).toThrow("does not match"); + }); + test("uses a state without credentials for masking and the current logger", async () => { const { requests, fetch } = mockIngestion(); vi.stubEnv("BRAINTRUST_INGESTION_KEY", ingestionUrl(newKey())); @@ -633,6 +697,85 @@ describe("initLogger with an ingestion key", () => { expect(requests.some((r) => r.url.includes("/attachment"))).toBe(false); }); + test("rejects external attachments before logging their rows", async () => { + const { requests, fetch } = mockIngestion(); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new ExternalAttachment({ + url: "s3://private-bucket/secret.pdf", + filename: "secret.pdf", + contentType: "application/pdf", + }), + output: "instance", + }); + logger.log({ + input: { + file: { + type: "external_attachment", + url: "s3://private-bucket/other.pdf", + filename: "other.pdf", + content_type: "application/pdf", + }, + }, + output: "reference", + }); + logger.log({ input: "kept", output: "kept" }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(String(onFlushError.mock.calls[0][0])).toContain( + "do not support external attachments", + ); + expect(loggedRows(requests).map((row) => row.output)).toEqual(["kept"]); + expect(JSON.stringify(requests)).not.toContain("private-bucket"); + }); + + test("keeps references to already committed attachments", async () => { + const { requests, fetch } = mockIngestion(); + const { logger, onFlushError } = initIngestionLogger(fetch); + const reference = { + type: "braintrust_attachment", + filename: "a.bin", + content_type: "application/octet-stream", + key: "server-committed-upload", + }; + + logger.log({ input: { file: reference }, output: "x" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(paths(requests)).toEqual(["POST /deployment/base/ingest/v1/logs"]); + expect(loggedRows(requests)[0].input.file).toEqual(reference); + }); + + test("keeps private external attachment behavior", async () => { + const { requests, fetch } = mockIngestion({ respond: () => json({}) }); + const state = new BraintrustState({ fetch, noExitFlush: true }); + state.apiUrl = "https://api.test"; + const logger = initLogger({ + state, + projectId: "private-project", + projectName: "private", + }); + + logger.log({ + input: new ExternalAttachment({ + url: "s3://bucket/file.pdf", + filename: "file.pdf", + contentType: "application/pdf", + }), + output: "x", + }); + await logger.flush(); + + expect(paths(requests)).toEqual(["GET /version", "POST /logs3"]); + expect(bodyJson(requests[1]).rows[0].input).toMatchObject({ + type: "external_attachment", + url: "s3://bucket/file.pdf", + }); + }); + test("restarts an upload whose grant expired", async () => { let grants = 0; const { requests, fetch } = mockIngestion({ diff --git a/js/src/ingestion-key.ts b/js/src/ingestion-key.ts index d7c3007e0..6c1fff368 100644 --- a/js/src/ingestion-key.ts +++ b/js/src/ingestion-key.ts @@ -1,6 +1,7 @@ import { z } from "zod/v3"; import { BraintrustAttachmentReference as braintrustAttachmentReferenceSchema, + ExternalAttachmentReference as externalAttachmentReferenceSchema, type BraintrustAttachmentReferenceType as BraintrustAttachmentReference, } from "./generated_types"; import { @@ -113,22 +114,29 @@ export function pickIngestionRowFields( const BRAINTRUST_ATTACHMENT = braintrustAttachmentReferenceSchema.shape.type.value; +const EXTERNAL_ATTACHMENT = externalAttachmentReferenceSchema.shape.type.value; /** * Replace attachment references whose key is in `references` with the * reference the ingestion endpoint returned for the upload. A reference mapped - * to `undefined` failed to upload, in which case this returns false. + * to `undefined` failed to upload. + * + * Returns why the row cannot be logged, or undefined if it can. External + * attachments are rejected, since they can reference any object store content. */ export function replaceIngestionAttachmentReferences( container: Record | unknown[], references: ReadonlyMap, -): boolean { +): "failed_upload" | "external_attachment" | undefined { for (const [key, value] of Object.entries(container)) { if (!value || typeof value !== "object") { continue; } // eslint-disable-next-line @typescript-eslint/consistent-type-assertions const reference = value as Record; + if (reference.type === EXTERNAL_ATTACHMENT) { + return "external_attachment"; + } if ( reference.type === BRAINTRUST_ATTACHMENT && typeof reference.key === "string" && @@ -136,17 +144,18 @@ export function replaceIngestionAttachmentReferences( ) { const uploaded = references.get(reference.key); if (!uploaded) { - return false; + return "failed_upload"; } // eslint-disable-next-line @typescript-eslint/consistent-type-assertions (container as Record)[key] = uploaded; continue; } - if (!replaceIngestionAttachmentReferences(reference, references)) { - return false; + const problem = replaceIngestionAttachmentReferences(reference, references); + if (problem) { + return problem; } } - return true; + return undefined; } export type IngestionUploadRequest = diff --git a/js/src/logger.ts b/js/src/logger.ts index cd041642f..e4a25f605 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -880,15 +880,16 @@ export class BraintrustState { */ public _internalGetIngestionKeyState( endpoint: IngestionEndpoint, + orgId: string | undefined, loginParams: Pick< LoginOptions, "fetch" | "noExitFlush" | "onFlushError" | "debugLogLevel" >, ): BraintrustState { - const cacheKey = JSON.stringify([endpoint.root, endpoint.key]); + const cacheKey = JSON.stringify([endpoint.root, endpoint.key, orgId]); let state = this.ingestionKeyStates.get(cacheKey); if (!state) { - state = new IngestionKeyState(this, endpoint, { + state = new IngestionKeyState(this, endpoint, orgId, { ...loginParams, fetch: loginParams.fetch ?? this.fetch, // Creating a state resets the global debug log level unless it is set. @@ -920,20 +921,6 @@ export class BraintrustState { ); } - /** - * Whether this state was given or logged in with an API key, as opposed to - * only picking one up from the environment on login. - * - * @internal - */ - public _internalHasPrivateCredential(): boolean { - return ( - this.loggedIn || - this.loginToken !== null || - this.loginParams.apiKey !== undefined - ); - } - /** @internal */ public _internalSetTraceContextSigningSecret( secret: string | undefined, @@ -1500,6 +1487,7 @@ class IngestionKeyState extends BraintrustState { constructor( private readonly contextState: BraintrustState, endpoint: IngestionEndpoint, + private readonly ingestionOrgId: string | undefined, loginParams: LoginOptions, ) { super(loginParams); @@ -1528,7 +1516,7 @@ class IngestionKeyState extends BraintrustState { return new HTTPBackgroundLogger( new LazyValue(async () => this.ingestionConn), loginParams, - true, + { orgId: this.ingestionOrgId }, ); } } @@ -3646,6 +3634,11 @@ class ConcurrencyLimiter { } } +interface IngestionLoggerOptions { + // Sent along with every row and upload for the server to check. + orgId: string | undefined; +} + interface QueuedLogEvent { sequence: number; event: LazyValue; @@ -3695,7 +3688,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { apiConn: LazyValue, opts?: BackgroundLoggerOpts, // Log through the ingestion key endpoints of `apiConn`. - private readonly usesIngestionKey = false, + private readonly ingestion?: IngestionLoggerOptions, ) { opts = opts ?? {}; this.apiConn = apiConn; @@ -3843,7 +3836,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { maxRequestSize: number; canUseOverflow: boolean; }> { - if (!this._maxRequestSizePromise && this.usesIngestionKey) { + if (!this._maxRequestSizePromise && this.ingestion) { // Ingestion keys cannot read the server version, and the ingestion // endpoint always accepts overflow uploads. this._maxRequestSizePromise = Promise.resolve({ @@ -3935,7 +3928,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { const attachmentErrors: unknown[] = []; // The ingestion endpoint only accepts attachments uploaded before the rows // that reference them. - const allItems = this.usesIngestionKey + const allItems = this.ingestion ? await this.prepareIngestionRows( unwrappedItems, attachments, @@ -3987,7 +3980,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { } // For now, upload attachments serially. - for (const attachment of this.usesIngestionKey ? [] : attachments) { + for (const attachment of this.ingestion ? [] : attachments) { try { const result = await attachment.upload(); if (result.upload_status === "error") { @@ -4032,14 +4025,28 @@ class HTTPBackgroundLogger implements BackgroundLogger { const droppedFields = new Set(); const rows: BackgroundLogEvent[] = []; + let externalAttachmentRows = 0; for (const item of items) { const row = pickIngestionRowFields(item, droppedFields); - // Rows that reference a failed upload are dropped with the error. - if (replaceIngestionAttachmentReferences(row, references)) { + if (this.ingestion?.orgId && row.org_id === undefined) { + row.org_id = this.ingestion.orgId; + } + // Rows that reference a failed upload are dropped with its error. + const problem = replaceIngestionAttachmentReferences(row, references); + if (problem === "external_attachment") { + externalAttachmentRows++; + } else if (!problem) { // eslint-disable-next-line @typescript-eslint/consistent-type-assertions rows.push(row as BackgroundLogEvent); } } + if (externalAttachmentRows) { + attachmentErrors.push( + new Error( + `Ingestion keys do not support external attachments, since they can reference any object store content. ${externalAttachmentRows} rows that reference one were not logged.`, + ), + ); + } for (const field of droppedFields) { if (!this.droppedIngestionFields.has(field)) { this.droppedIngestionFields.add(field); @@ -4109,7 +4116,11 @@ class HTTPBackgroundLogger implements BackgroundLogger { ingestionUploadGrantSchema, // The optional sha256 is omitted, since computing it would read the // whole upload into memory at once. - await conn.post_json("v1/uploads", { ...request, size_bytes: data.size }), + await conn.post_json("v1/uploads", { + ...request, + size_bytes: data.size, + ...(this.ingestion?.orgId ? { org_id: this.ingestion.orgId } : {}), + }), ); const deadline = requestStart + grant.expires_in_ms; if (grant.num_chunks !== Math.ceil(data.size / grant.chunk_bytes)) { @@ -4289,7 +4300,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { let overflowUpload: Logs3OverflowUpload | null = null; let ingestionOverflowKey: string | null = null; const overflowRows = - useOverflow && !this.usesIngestionKey + useOverflow && !this.ingestion ? items.map((item) => item.overflowMeta) : null; @@ -4297,7 +4308,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { const startTime = now(); let error: unknown = undefined; try { - if (useOverflow && this.usesIngestionKey) { + if (useOverflow && this.ingestion) { if (!ingestionOverflowKey) { const reference = await this.uploadToIngestion( { purpose: "logs3_overflow", content_type: "application/json" }, @@ -4314,7 +4325,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { "v1/logs", constructLogs3OverflowRequest(ingestionOverflowKey), ); - } else if (this.usesIngestionKey) { + } else if (this.ingestion) { await conn.post_json("v1/logs", dataStr); } else if (overflowRows) { if (!overflowUpload) { @@ -4343,7 +4354,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { return; } - const retryDelayMs = this.usesIngestionKey + const retryDelayMs = this.ingestion ? ingestionRetryDelayMs(error, i) : BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * 1000 * 2 ** i; const isRetrying = i + 1 < this.numTries && retryDelayMs !== undefined; @@ -5617,27 +5628,24 @@ export type InitLoggerOptions = FullLoginOptions & { * @param options.orgName (Optional) The name of a specific organization to connect to. This is useful if you belong to multiple. * @param options.forceLogin Login again, even if you have already logged in (by default, the logger will not login if you are already logged in) * @param options.debugLogLevel Enables internal Braintrust SDK troubleshooting output. Use `"error"`, `"warn"`, `"info"`, or `"debug"` to choose an explicit level, or `false` to explicitly disable it. If omitted, the SDK stays silent unless `BRAINTRUST_DEBUG_LOG_LEVEL` is set. - * @param options.ingestionKey The ingestion URL of a project's ingestion key, like `https:///ingest?ingestKey=`. Ingestion keys can only write traces into their project, so the logger never logs in, ignores `BRAINTRUST_API_KEY` and does not create or look up the project. `projectName` is ignored, and `projectId` is only sent along for the server to check. A `state` without an API key is used for the span context and masking, while the logger keeps its own queue. If unspecified, will use the `BRAINTRUST_INGESTION_KEY` environment variable unless an API key is passed through `apiKey` or `state`. Cannot be combined with an API key. + * @param options.ingestionKey The ingestion URL of a project's ingestion key, like `https:///ingest?ingestKey=`. Ingestion keys can only write traces into their project, so the logger never logs in, ignores `BRAINTRUST_API_KEY` and any login of `state`, and does not create or look up the project. `projectName` is ignored, and `projectId` or the ids of `orgProjectMetadata` are only sent along for the server to check. A `state` provides the masking function and current logger, while the logger keeps its own queue. `ExternalAttachment` is not supported. If unspecified, will use the `BRAINTRUST_INGESTION_KEY` environment variable unless `apiKey` is specified. Cannot be combined with `apiKey`. * @param setCurrent If true (the default), set the global current-experiment to the newly-created one. * @returns The newly created Logger. */ export function initLogger( options: Readonly> = {}, ) { - const hasPrivateCredential = - options.apiKey !== undefined || - (options.state?._internalHasPrivateCredential() ?? false); - if (options.ingestionKey !== undefined && hasPrivateCredential) { + if (options.ingestionKey !== undefined && options.apiKey !== undefined) { throw new Error( - "initLogger accepts either an ingestionKey or an API key, not both. A `state` that is logged in or has an `apiKey` counts as an API key.", + "initLogger accepts either an ingestionKey or an apiKey, not both", ); } - // An explicit private credential takes precedence over the environment. + // An explicit API key takes precedence over the environment. const ingestionKey = options.ingestionKey ?? - (hasPrivateCredential - ? undefined - : iso.getEnv(INGESTION_KEY_ENV_VAR) || undefined); + (options.apiKey === undefined + ? iso.getEnv(INGESTION_KEY_ENV_VAR) || undefined + : undefined); if (ingestionKey !== undefined) { return initIngestionKeyLogger(ingestionKey, options); } @@ -5719,25 +5727,42 @@ function initIngestionKeyLogger( noExitFlush, onFlushError, debugLogLevel, + orgProjectMetadata, state: contextState = _globalState, }: Readonly>, ) { + const metadataProjectId = orgProjectMetadata?.project.id || undefined; + if (projectId && metadataProjectId && projectId !== metadataProjectId) { + throw new Error( + "projectId does not match the project id of orgProjectMetadata", + ); + } + const resolvedProjectId = projectId || metadataProjectId; + // Any login of the context state is ignored, only its context is used. const state = contextState._internalGetIngestionKeyState( parseIngestionKeyUrl(ingestionKey), + orgProjectMetadata?.org_id || undefined, { fetch, noExitFlush, onFlushError, debugLogLevel }, ); state.spanOriginEnvironment = detectSpanOriginEnvironment(environment); state.enforceQueueSizeLimit(true); - // The project is only known if it is passed in, and is never looked up. + // The org and project are only known if they are passed in, and are never + // looked up. The server checks them against the key. const lazyMetadata = new LazyValue(async () => ({ - org_id: "", - project: { id: projectId ?? "", name: "", fullInfo: {} }, + org_id: orgProjectMetadata?.org_id ?? "", + project: { id: resolvedProjectId ?? "", name: "", fullInfo: {} }, })); const ret = new Logger(state, lazyMetadata, { asyncFlush, - computeMetadataArgs: projectId ? { project_id: projectId } : undefined, - linkArgs: { org_name: orgName, app_url: appUrl, project_id: projectId }, + computeMetadataArgs: resolvedProjectId + ? { project_id: resolvedProjectId } + : undefined, + linkArgs: { + org_name: orgName, + app_url: appUrl, + project_id: resolvedProjectId, + }, }); state.currentLogger = ret as Logger; if (setCurrent ?? true) { From dbcf3644836cce4e98528d88dafd4bc23ffae6da Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:52:11 +0000 Subject: [PATCH 4/5] test: Check ingestion keys against the shared fixture Runs the URL, upload response and chunk plan cases of typespecs/src/public-ingestion.fixture.json (braintrustdata/braintrust@74ce022) through the parser and the upload flow. --- js/src/ingestion-key.test.ts | 187 ++++++++++++++++++++++++++++++++++- 1 file changed, 186 insertions(+), 1 deletion(-) diff --git a/js/src/ingestion-key.test.ts b/js/src/ingestion-key.test.ts index 7aa61ee80..0b44791da 100644 --- a/js/src/ingestion-key.test.ts +++ b/js/src/ingestion-key.test.ts @@ -5,6 +5,7 @@ import { _internalGetGlobalState, Attachment, BraintrustState, + constructLogs3OverflowRequest, ExternalAttachment, extractTraceContextFromHeaders, flush, @@ -13,7 +14,12 @@ import { setMaskingFunction, traced, } from "./logger"; -import { parseIngestionKeyUrl } from "./ingestion-key"; +import { + ingestionUploadChunkSchema, + ingestionUploadCompleteSchema, + ingestionUploadGrantSchema, + parseIngestionKeyUrl, +} from "./ingestion-key"; import { configureNode } from "./node/config"; configureNode(); @@ -1012,3 +1018,182 @@ describe("tracing with an ingestion key", () => { expect(headers.baggage ?? "").not.toContain("braintrust.parent"); }); }); + +// Cases from typespecs/src/public-ingestion.fixture.json in +// braintrustdata/braintrust at 74ce022558796700023b29a98dba448784fc69f2. +const FIXTURE_KEY = "bt-ik-FAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKE"; +const FIXTURE_VALID_URLS: [string, string][] = [ + [ + `https://dp.example.com/ingest?ingestKey=${FIXTURE_KEY}`, + "https://dp.example.com/ingest", + ], + [ + `https://dp.example.com/base/path/ingest?ingestKey=${FIXTURE_KEY}`, + "https://dp.example.com/base/path/ingest", + ], + [ + `https://dp.example.com:8443/base/ingest/?ingestKey=${FIXTURE_KEY}`, + "https://dp.example.com:8443/base/ingest", + ], + [ + `http://localhost:8000/ingest?ingestKey=${FIXTURE_KEY}`, + "http://localhost:8000/ingest", + ], + [ + `https://dp.example.com/base/ingest///?ingestKey=${FIXTURE_KEY}`, + "https://dp.example.com/base/ingest", + ], + [ + "https://dp.example.com/ingest?ingestKey=bt%2Dik%2DFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKE", + "https://dp.example.com/ingest", + ], + [ + `https://dp.example.com/ingest?ingest%4Bey=${FIXTURE_KEY}`, + "https://dp.example.com/ingest", + ], +]; +const FIXTURE_INVALID_URLS: [string, string][] = [ + ["missing key", "https://dp.example.com/ingest"], + [ + "path does not end with /ingest", + `https://dp.example.com/base?ingestKey=${FIXTURE_KEY}`, + ], + [ + "not an ingestion key", + // An API key shaped value, built here so it does not look like a secret. + `https://dp.example.com/ingest?ingestKey=sk-${"FAKE".repeat(12)}`, + ], + [ + "truncated key", + "https://dp.example.com/ingest?ingestKey=bt-ik-FAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAKEFAK", + ], + [ + "extra query parameter", + `https://dp.example.com/ingest?ingestKey=${FIXTURE_KEY}&project_id=foo`, + ], + [ + "repeated key parameter", + `https://dp.example.com/ingest?ingestKey=${FIXTURE_KEY}&ingestKey=${FIXTURE_KEY}`, + ], + [ + "embedded credentials", + `https://user:pass@dp.example.com/ingest?ingestKey=${FIXTURE_KEY}`, + ], + [ + "fragment", + `https://dp.example.com/ingest?ingestKey=${FIXTURE_KEY}#fragment`, + ], + ["empty fragment", `https://dp.example.com/ingest?ingestKey=${FIXTURE_KEY}#`], + [ + "unsupported protocol", + `ftp://dp.example.com/ingest?ingestKey=${FIXTURE_KEY}`, + ], + ["bare key without a URL", `${FIXTURE_KEY}`], +]; +const FIXTURE_UPLOAD = { + grant: { + upload_id: "00000000-0000-4000-8000-000000000001", + chunk_bytes: 524288, + num_chunks: 3, + expires_in_ms: 300000, + }, + chunk_response: { index: 2, size_bytes: 100 }, + complete_responses: [ + { + reference: { + type: "braintrust_attachment", + filename: "image.png", + content_type: "image/png", + key: "server-issued/attachment-key", + }, + size_bytes: 1048676, + sha256: + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + }, + { + reference: { + type: "logs3_overflow", + key: "00000000-0000-4000-8000-000000000001", + }, + size_bytes: 5242880, + sha256: + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + }, + ], + overflow_logs_request: { + api_version: 2, + rows: { + type: "logs3_overflow", + key: "00000000-0000-4000-8000-000000000001", + }, + }, +}; +// [size_bytes, chunk_bytes, num_chunks, last_chunk_bytes] +const FIXTURE_CHUNK_PLANS: [number, number, number, number | null][] = [ + [0, 524288, 0, null], + [1, 524288, 1, 1], + [524288, 524288, 1, 524288], + [524289, 524288, 2, 1], + [2097252, 524288, 5, 100], + [9437184, 4194304, 3, 1048576], +]; + +describe("shared public ingestion fixture", () => { + test.each(FIXTURE_VALID_URLS)("parses %s", (url, root) => { + expect(parseIngestionKeyUrl(url)).toEqual({ root, key: FIXTURE_KEY }); + }); + + test.each(FIXTURE_INVALID_URLS)("rejects %s", (_, url) => { + expect(() => parseIngestionKeyUrl(url)).toThrow( + "Invalid Braintrust ingestion key", + ); + try { + parseIngestionKeyUrl(url); + } catch (error) { + expect(inspect(error)).not.toContain(FIXTURE_KEY); + } + }); + + test("accepts the upload responses", () => { + expect(ingestionUploadGrantSchema.parse(FIXTURE_UPLOAD.grant)).toEqual( + FIXTURE_UPLOAD.grant, + ); + expect( + ingestionUploadChunkSchema.parse(FIXTURE_UPLOAD.chunk_response), + ).toEqual(FIXTURE_UPLOAD.chunk_response); + for (const response of FIXTURE_UPLOAD.complete_responses) { + expect(ingestionUploadCompleteSchema.parse(response)).toEqual(response); + } + expect( + constructLogs3OverflowRequest(FIXTURE_UPLOAD.grant.upload_id), + ).toEqual(FIXTURE_UPLOAD.overflow_logs_request); + }); + + test.each(FIXTURE_CHUNK_PLANS)( + "uploads %i bytes in chunks of %i", + async (sizeBytes, chunkBytes, numChunks, lastChunkBytes) => { + const { requests, fetch } = mockIngestion({ chunkBytes }); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(sizeBytes), + filename: "data.bin", + contentType: "application/octet-stream", + }), + output: "x", + }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + const chunkSizes = requests + .filter((r) => r.url.includes("/chunks/")) + .map((r) => (r.body as Uint8Array).length); + expect(chunkSizes).toHaveLength(numChunks); + expect(chunkSizes.at(-1) ?? null).toBe(lastChunkBytes); + expect(chunkSizes.slice(0, -1).every((size) => size === chunkBytes)).toBe( + true, + ); + }, + ); +}); From bdb92e6b798b9e892d419d78d3303898f71413bb Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 14:16:08 +0000 Subject: [PATCH 5/5] fix: Harden the ingestion key transport Ingestion key requests no longer follow redirects, which could forward the key to another origin, and the key is redacted from error responses and transport errors. They also skip `keepalive`, since browsers reject keepalive bodies above 64 KiB and native rows go up to 512 KiB. Upload grants are anchored on a monotonic clock, and chunk and complete requests are aborted once the grant expires. A 410 requests a new grant, while invalid responses are not retried. An empty `BRAINTRUST_INGESTION_KEY` now fails instead of falling back to private credentials. `updateSpan({ exported })` updates project log spans through the current ingestion key logger, and spans exported without a project are rejected instead of being updated with private credentials. --- js/README.md | 1 + js/src/ingestion-key.test.ts | 322 ++++++++++++++++++++++++++++++++++- js/src/logger.ts | 157 ++++++++++++++--- 3 files changed, 456 insertions(+), 24 deletions(-) diff --git a/js/README.md b/js/README.md index 7a6f51d56..ffb48a856 100644 --- a/js/README.md +++ b/js/README.md @@ -126,6 +126,7 @@ Uploaded attachments can only be referenced by rows written with the same key or An explicit `apiKey` takes precedence over `BRAINTRUST_INGESTION_KEY`, and combining it with `ingestionKey` throws. A `state` only provides the masking function and current logger, while the logger keeps its own queue and ignores any login of the state. `projectName` is ignored, and `projectId` or the ids of `orgProjectMetadata` are only sent along for the server to check. +`updateSpan({ exported })` updates project log spans through the current ingestion key logger, and spans exported without a project need `state: logger.loggingState` instead of falling back to an API key. `ExternalAttachment` is not supported, since it can reference any object store content, so rows that contain one are not logged and the error is reported to `onFlushError`. Attachments and large batches are uploaded through the ingestion URL as well. Fields that ingestion keys do not support, like `classifications` or audit fields, are not logged, and logging comments throws. diff --git a/js/src/ingestion-key.test.ts b/js/src/ingestion-key.test.ts index 0b44791da..a6372ee5d 100644 --- a/js/src/ingestion-key.test.ts +++ b/js/src/ingestion-key.test.ts @@ -13,6 +13,7 @@ import { type InitLoggerOptions, setMaskingFunction, traced, + updateSpan, } from "./logger"; import { ingestionUploadChunkSchema, @@ -42,6 +43,9 @@ interface RecordedRequest { method: string; headers: Record; body: Uint8Array | string | undefined; + redirect: RequestRedirect | undefined; + keepalive: boolean | undefined; + signal: AbortSignal | undefined; } function json(data: unknown, status = 200) { @@ -91,6 +95,9 @@ function mockIngestion( ), ), body, + redirect: init?.redirect, + keepalive: init?.keepalive, + signal: init?.signal ?? undefined, }; requests.push(request); const override = await options.respond?.(request, requests.length - 1); @@ -170,7 +177,14 @@ function initIngestionLogger( fetch: typeof globalThis.fetch, options: InitLoggerOptions = {}, ) { - const key = newKey(); + return initIngestionLoggerWithKey(fetch, newKey(), options); +} + +function initIngestionLoggerWithKey( + fetch: typeof globalThis.fetch, + key: string, + options: InitLoggerOptions = {}, +) { const onFlushError = vi.fn(); const logger = initLogger({ ingestionKey: ingestionUrl(key), @@ -182,9 +196,18 @@ function initIngestionLogger( return { key, logger, onFlushError }; } +// Requests of the global (private) state never reach a real server. +const privateFetch = vi.fn(async (input: RequestInfo | URL) => { + throw new Error(`Unexpected private request to ${String(input)}`); +}); + beforeEach(() => { // Only flush explicitly, so the rows of a span are sent together. vi.stubEnv("BRAINTRUST_SYNC_FLUSH", "1"); + privateFetch.mockClear(); + _internalGetGlobalState().setFetch( + privateFetch as unknown as typeof globalThis.fetch, + ); }); afterEach(() => { @@ -342,6 +365,9 @@ describe("initLogger with an ingestion key", () => { expect(logger.loggingState).toBe(state); expect(paths(requests)).toEqual(["GET /version", "POST /logs3"]); expect(requests[1].headers.authorization).toBe("Bearer private-api-key"); + // The private transport keeps its keepalive requests. + expect(requests[1].keepalive).toBe(true); + expect(requests[1].redirect).toBeUndefined(); expect(JSON.stringify(requests)).not.toContain(key); const apiKeyLogger = initLogger({ apiKey: "private-api-key" }); @@ -934,6 +960,300 @@ describe("initLogger with an ingestion key", () => { }); }); +describe("ingestion key transport", () => { + test("refuses redirects and does not follow them", async () => { + const { requests, fetch } = mockIngestion({ + respond: () => + new Response(null, { + status: 307, + headers: { Location: "https://attacker.example/v1/logs" }, + }), + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "2"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(String(onFlushError.mock.calls[0][0])).toContain("redirect"); + expect(paths(requests)).toEqual(["POST /deployment/base/ingest/v1/logs"]); + expect(requests.every((r) => r.redirect === "error")).toBe(true); + }); + + test("asks fetch not to follow redirects for uploads", async () => { + const { requests, fetch } = mockIngestion({ chunkBytes: 4 }); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "x", + }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(requests).toHaveLength(6); + expect(requests.every((r) => r.redirect === "error")).toBe(true); + }); + + test("redacts the key from error responses", async () => { + const key = newKey(); + const { fetch } = mockIngestion({ + respond: () => + new Response(`invalid Authorization: Bearer ${key}`, { + status: 400, + statusText: `Bad ${key}`, + }), + }); + const { logger, onFlushError } = initIngestionLoggerWithKey(fetch, key); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + const reported = inspect(onFlushError.mock.calls[0][0], { depth: 10 }); + expect(reported).not.toContain(key); + expect(reported).toContain("Bearer [REDACTED]"); + }); + + test("redacts the key from network errors", async () => { + const key = newKey(); + const fetch = vi.fn(async () => { + throw new TypeError(`connect failed for Bearer ${key}`, { + cause: new Error(key), + }); + }) as unknown as typeof globalThis.fetch; + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "0"); + const { logger, onFlushError } = initIngestionLoggerWithKey(fetch, key); + + logger.log({ input: "a", output: "b" }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + const reported = inspect(onFlushError.mock.calls[0][0], { depth: 10 }); + expect(reported).not.toContain(key); + expect(reported).toContain("connect failed for Bearer [REDACTED]"); + }); + + test("requests a new grant when the server reports it expired", async () => { + let chunkRequests = 0; + const { requests, fetch } = mockIngestion({ + respond: (request) => + request.url.includes("/chunks/") && chunkRequests++ === 0 + ? json({ error: "upload expired" }, 410) + : undefined, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "1"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "x", + }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + const grants = requests.filter((r) => r.url.endsWith("/v1/uploads")); + expect(grants).toHaveLength(2); + expect(loggedRows(requests)).toHaveLength(1); + }); + + test("aborts upload requests that outlive the grant", async () => { + let grants = 0; + const { requests, fetch } = mockIngestion({ + respond: (request) => { + if (request.url.endsWith("/v1/uploads") && grants++ === 0) { + return json( + { + upload_id: randomUUID(), + chunk_bytes: 512 * KIB, + num_chunks: 1, + expires_in_ms: 50, + }, + 201, + ); + } + if (request.url.includes("/chunks/") && grants === 1) { + // Hang until the grant expires. + return new Promise((_, reject) => + request.signal!.addEventListener("abort", () => + reject(request.signal!.reason), + ), + ); + } + return undefined; + }, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "1"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "x", + }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(requests.filter((r) => r.url.endsWith("/v1/uploads"))).toHaveLength( + 2, + ); + expect( + requests.filter((r) => r.url.includes("/chunks/")).every((r) => r.signal), + ).toBe(true); + expect(loggedRows(requests)).toHaveLength(1); + }); + + test.each([ + ["an invalid grant", { upload_id: "not-a-uuid" }], + [ + "an impossible chunk plan", + { + upload_id: "00000000-0000-4000-8000-000000000001", + chunk_bytes: 4, + num_chunks: 1, + expires_in_ms: 300_000, + }, + ], + ])("does not retry %s", async (_, grant) => { + const { requests, fetch } = mockIngestion({ + respond: (request) => + request.url.endsWith("/v1/uploads") ? json(grant, 201) : undefined, + }); + vi.stubEnv("BRAINTRUST_NUM_RETRIES", "2"); + const { logger, onFlushError } = initIngestionLogger(fetch); + + logger.log({ + input: new Attachment({ + data: new ArrayBuffer(10), + filename: "a.bin", + contentType: "application/octet-stream", + }), + output: "x", + }); + await logger.flush(); + + expect(onFlushError).toHaveBeenCalledTimes(1); + expect(String(onFlushError.mock.calls[0][0])).toContain( + "Invalid response from ingestion endpoint", + ); + expect(paths(requests)).toEqual([ + "POST /deployment/base/ingest/v1/uploads", + ]); + }); + + test("sends native bodies above the browser keepalive limit", async () => { + const mock = mockIngestion(); + // Like browsers, reject keepalive requests with bodies above 64 KiB. + const browserFetch = (async ( + input: RequestInfo | URL, + init?: RequestInit, + ) => { + if (init?.keepalive && String(init.body ?? "").length > 64 * KIB) { + throw new TypeError("Failed to fetch"); + } + return mock.fetch(input, init); + }) as typeof globalThis.fetch; + const { logger, onFlushError } = initIngestionLogger(browserFetch); + + logger.log({ input: "x".repeat(200 * KIB), output: "inline" }); + await logger.flush(); + + expect(onFlushError).not.toHaveBeenCalled(); + expect(paths(mock.requests)).toEqual([ + "POST /deployment/base/ingest/v1/logs", + ]); + const [request] = mock.requests; + expect(request.keepalive).toBe(false); + expect(String(request.body).length).toBeGreaterThan(64 * KIB); + expect(String(request.body).length).toBeLessThan(512 * KIB); + }); + + test("an empty BRAINTRUST_INGESTION_KEY fails instead of logging privately", () => { + vi.stubEnv("BRAINTRUST_INGESTION_KEY", ""); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + expect(() => initLogger({ noExitFlush: true })).toThrow( + "Invalid Braintrust ingestion key", + ); + }); +}); + +describe("updateSpan with ingestion keys", () => { + test("updates exported spans through the current ingestion key logger", async () => { + const { requests, fetch } = mockIngestion(); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + const { key, logger } = initIngestionLogger(fetch, { + projectId: "project-id", + }); + const span = logger.startSpan({ name: "root" }); + span.end(); + const exported = await span.export(); + + updateSpan({ exported, output: "updated" }); + await flush(); + + expect(privateFetch).not.toHaveBeenCalled(); + expect( + requests.every((r) => r.headers.authorization === `Bearer ${key}`), + ).toBe(true); + expect(loggedRows(requests)).toContainEqual( + expect.objectContaining({ + id: span.id, + output: "updated", + project_id: "project-id", + }), + ); + }); + + test("updates exported spans without a project through the logger state", async () => { + const { requests, fetch } = mockIngestion(); + const { key, logger } = initIngestionLogger(fetch, { setCurrent: false }); + const span = logger.startSpan({ name: "root" }); + span.end(); + const exported = await span.export(); + + updateSpan({ exported, output: "updated", state: logger.loggingState }); + await logger.flush(); + + const update = loggedRows(requests).find((row) => row.output === "updated"); + expect(update).toMatchObject({ id: span.id, log_id: "g" }); + expect(update).not.toHaveProperty("project_id"); + expect( + requests.every((r) => r.headers.authorization === `Bearer ${key}`), + ).toBe(true); + }); + + test("rejects exported spans without a project instead of using private credentials", async () => { + const { requests, fetch } = mockIngestion(); + vi.stubEnv("BRAINTRUST_API_KEY", "private-api-key"); + const { logger } = initIngestionLogger(fetch, { setCurrent: false }); + const span = logger.startSpan({ name: "root" }); + span.end(); + const exported = await span.export(); + + expect(() => updateSpan({ exported, output: "updated" })).toThrow( + "exported without its project", + ); + await flush(); + await logger.flush(); + expect(privateFetch).not.toHaveBeenCalled(); + expect(loggedRows(requests).some((row) => row.output === "updated")).toBe( + false, + ); + }); +}); + describe("flush with ingestion keys", () => { test("drains every ingestion key queue, not only the current logger", async () => { const { requests, fetch } = mockIngestion(); diff --git a/js/src/logger.ts b/js/src/logger.ts index e4a25f605..1130290f6 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -1486,12 +1486,15 @@ class IngestionKeyState extends BraintrustState { constructor( private readonly contextState: BraintrustState, - endpoint: IngestionEndpoint, + private readonly endpoint: IngestionEndpoint, private readonly ingestionOrgId: string | undefined, loginParams: LoginOptions, ) { super(loginParams); - this.ingestionConn = new HTTPConnection(endpoint.root, this.fetch); + this.ingestionConn = new HTTPConnection( + endpoint.root, + ingestionKeyFetch(this.fetch, endpoint.key), + ); this.ingestionConn.set_token(endpoint.key); } @@ -1507,7 +1510,7 @@ class IngestionKeyState extends BraintrustState { public override setFetch(fetch: typeof globalThis.fetch) { super.setFetch(fetch); - this.ingestionConn.setFetch(fetch); + this.ingestionConn.setFetch(ingestionKeyFetch(fetch, this.endpoint.key)); } protected override createBackgroundLogger( @@ -1521,6 +1524,66 @@ class IngestionKeyState extends BraintrustState { } } +/** + * Wrap `fetch` for ingestion key requests. Redirects are refused, since they + * could forward the key to another origin, keepalive is disabled, and the key + * is redacted from error responses and transport errors before anything can + * log them. + */ +function ingestionKeyFetch( + fetch: typeof globalThis.fetch, + key: string, +): typeof globalThis.fetch { + const redact = (text: string) => text.split(key).join("[REDACTED]"); + return async (input, init) => { + let response: Response; + try { + // Browsers reject keepalive requests above 64 KiB, which native rows + // below the overflow threshold can exceed. + response = await fetch(input, { + ...init, + redirect: "error", + keepalive: false, + }); + } catch (error) { + if (init?.signal?.aborted) { + throw error; + } + // The cause is dropped, since it may reference the key as well. + throw new Error( + `Ingestion request failed: ${redact(error instanceof Error ? error.message : String(error))}`, + ); + } + // Custom fetch implementations may ignore `redirect: "error"`. + if ( + response.type === "opaqueredirect" || + (response.status >= 300 && response.status < 400) + ) { + await response.body?.cancel(); + throw new IngestionProtocolError( + "Ingestion endpoint responded with a redirect, which is not followed", + ); + } + if (response.ok) { + return response; + } + let body: string; + try { + body = await response.text(); + } catch { + body = "Unable to read response body"; + } + return new Response(redact(body), { + status: response.status, + statusText: redact(response.statusText), + headers: response.headers, + }); + }; +} + +// An invalid response of the ingestion endpoint, which is not retried. +class IngestionProtocolError extends Error {} + function usesIngestionKey(state: BraintrustState): boolean { return state.currentLogger?.loggingState instanceof IngestionKeyState; } @@ -2670,13 +2733,20 @@ export function updateSpan({ "id" | "root_span_id" | "span_id" > & OptionalStateArg): void { - const resolvedState = state ?? _globalState; + const requestedState = state ?? _globalState; const components = SpanComponentsV4.fromStr(exported); if (!components.data.row_id) { throw new Error("Exported span must have a row id"); } + // While an ingestion key logger is current, project log spans are updated + // through its key and never through private credentials. + const resolvedState = + components.data.object_type === SpanObjectTypeV3.PROJECT_LOGS && + usesIngestionKey(requestedState) + ? requestedState.currentLogger!.loggingState + : requestedState; updateSpanImpl({ state: resolvedState, parentObjectType: components.data.object_type, @@ -2761,8 +2831,12 @@ function spanComponentsToObjectIdLambda( return async () => projectId; } if (!components.data.compute_object_metadata_args) { + // Only ingestion key loggers export spans without their object, and + // private credentials must not be used in their place. throw new Error( - "Impossible: must provide either objectId or computeObjectMetadataArgs", + components.data.object_type === SpanObjectTypeV3.PROJECT_LOGS + ? "This span was exported without its project, which happens for loggers that use an ingestion key. Continue or update it through a logger that uses an ingestion key." + : "Impossible: must provide either objectId or computeObjectMetadataArgs", ); } switch (components.data.object_type) { @@ -3423,12 +3497,19 @@ function parseIngestionResponse(schema: z.ZodType, response: unknown) { } catch (error) { if (error instanceof ZodError) { const errorStr = JSON.stringify(error.flatten()); - throw new Error(`Invalid response from ingestion endpoint: ${errorStr}`); + throw new IngestionProtocolError( + `Invalid response from ingestion endpoint: ${errorStr}`, + ); } throw error; } } +// Monotonic milliseconds, so that wall clock changes cannot extend a grant. +function monotonicNow() { + return typeof performance !== "undefined" ? performance.now() : now(); +} + export interface BackgroundLoggerOpts { noExitFlush?: boolean; onFlushError?: (error: unknown) => void; @@ -3527,18 +3608,24 @@ const RETRYABLE_HTTP_STATUS_CODES = new Set([500, 502, 503, 504]); const MAX_INGESTION_RETRY_AFTER_MS = 60_000; // Ingestion key requests are rejected permanently on auth or validation -// failures, so only transient failures are retried. Returns undefined if the -// request must not be retried. +// failures and invalid responses, so only transient failures and expired +// upload grants are retried. Returns undefined if the request must not be +// retried. function ingestionRetryDelayMs( error: unknown, attempt: number, ): number | undefined { const backoffMs = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * 1000 * 2 ** attempt; + if (error instanceof IngestionProtocolError || error instanceof SyntaxError) { + return undefined; + } if (!(error instanceof FailedHTTPResponse)) { return backoffMs; } + // 410 means that an upload grant expired, so a new one is requested. if ( error.status !== 408 && + error.status !== 410 && error.status !== 429 && !RETRYABLE_HTTP_STATUS_CODES.has(error.status) ) { @@ -4073,7 +4160,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { await attachment.data(), ); if (reference.type !== BRAINTRUST_ATTACHMENT) { - throw new Error( + throw new IngestionProtocolError( "Invalid response from ingestion endpoint: expected an attachment reference", ); } @@ -4111,7 +4198,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { private async uploadToIngestion(request: IngestionUploadRequest, data: Blob) { const conn = await this.apiConn.get(); // The grant lifetime is relative, so anchor it before the request is sent. - const requestStart = now(); + const requestStart = monotonicNow(); const grant = parseIngestionResponse( ingestionUploadGrantSchema, // The optional sha256 is omitted, since computing it would read the @@ -4122,22 +4209,38 @@ class HTTPBackgroundLogger implements BackgroundLogger { ...(this.ingestion?.orgId ? { org_id: this.ingestion.orgId } : {}), }), ); - const deadline = requestStart + grant.expires_in_ms; if (grant.num_chunks !== Math.ceil(data.size / grant.chunk_bytes)) { - throw new Error( + throw new IngestionProtocolError( `Invalid response from ingestion endpoint: ${grant.num_chunks} chunks of ${grant.chunk_bytes} bytes cannot hold ${data.size} bytes`, ); } - const assertGrantActive = () => { - if (now() >= deadline) { + // Requests of the upload must finish before the grant expires. + const remainingMs = Math.floor( + requestStart + grant.expires_in_ms - monotonicNow(), + ); + if (remainingMs <= 0) { + throw new Error("Ingestion upload expired before it completed"); + } + const signal = AbortSignal.timeout(remainingMs); + try { + return await this.uploadIngestionChunks(conn, grant, data, signal); + } catch (error) { + if (signal.aborted) { throw new Error("Ingestion upload expired before it completed"); } - }; + throw error; + } + } + private async uploadIngestionChunks( + conn: HTTPConnection, + grant: z.infer, + data: Blob, + signal: AbortSignal, + ) { // On platforms like Cloudflare, fetch must not be called as a method. const fetch = conn.fetch; for (let index = 0; index < grant.num_chunks; index++) { - assertGrantActive(); const chunk = data.slice( index * grant.chunk_bytes, (index + 1) * grant.chunk_bytes, @@ -4156,27 +4259,31 @@ class HTTPBackgroundLogger implements BackgroundLogger { "Content-Type": "application/octet-stream", }, body: chunk, + signal, }, ), ); const ack = parseIngestionResponse( ingestionUploadChunkSchema, - await response.json(), + await readJSONResponse(response), ); if (ack.index !== index || ack.size_bytes !== chunk.size) { - throw new Error( + throw new IngestionProtocolError( `Invalid response from ingestion endpoint: chunk ${index} of ${chunk.size} bytes was acknowledged as chunk ${ack.index} of ${ack.size_bytes} bytes`, ); } } - assertGrantActive(); const completed = parseIngestionResponse( ingestionUploadCompleteSchema, - await conn.post_json(`v1/uploads/${grant.upload_id}/complete`, {}), + await conn.post_json( + `v1/uploads/${grant.upload_id}/complete`, + {}, + signal, + ), ); if (completed.size_bytes !== data.size) { - throw new Error( + throw new IngestionProtocolError( "Invalid response from ingestion endpoint: the completed upload does not match the uploaded data", ); } @@ -4315,7 +4422,7 @@ class HTTPBackgroundLogger implements BackgroundLogger { new Blob([dataStr], { type: "application/json" }), ); if (reference.type !== LOGS3_OVERFLOW_REFERENCE_TYPE) { - throw new Error( + throw new IngestionProtocolError( "Invalid response from ingestion endpoint: expected an overflow reference", ); } @@ -4353,6 +4460,10 @@ class HTTPBackgroundLogger implements BackgroundLogger { if (error === undefined) { return; } + if (error instanceof FailedHTTPResponse && error.status === 410) { + // The overflow upload expired, so the payload is uploaded again. + ingestionOverflowKey = null; + } const retryDelayMs = this.ingestion ? ingestionRetryDelayMs(error, i) @@ -5644,7 +5755,7 @@ export function initLogger( const ingestionKey = options.ingestionKey ?? (options.apiKey === undefined - ? iso.getEnv(INGESTION_KEY_ENV_VAR) || undefined + ? iso.getEnv(INGESTION_KEY_ENV_VAR) : undefined); if (ingestionKey !== undefined) { return initIngestionKeyLogger(ingestionKey, options);