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..ffb48a856 100644 --- a/js/README.md +++ b/js/README.md @@ -104,6 +104,35 @@ 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=", +}); +``` + +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 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. + +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..a6372ee5d --- /dev/null +++ b/js/src/ingestion-key.test.ts @@ -0,0 +1,1519 @@ +import { createHash, randomUUID } from "node:crypto"; +import { inspect } from "node:util"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + _internalGetGlobalState, + Attachment, + BraintrustState, + constructLogs3OverflowRequest, + ExternalAttachment, + extractTraceContextFromHeaders, + flush, + initLogger, + type InitLoggerOptions, + setMaskingFunction, + traced, + updateSpan, +} from "./logger"; +import { + ingestionUploadChunkSchema, + ingestionUploadCompleteSchema, + ingestionUploadGrantSchema, + 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; + redirect: RequestRedirect | undefined; + keepalive: boolean | undefined; + signal: AbortSignal | 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, + redirect: init?.redirect, + keepalive: init?.keepalive, + signal: init?.signal ?? undefined, + }; + 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")) { + // 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) + ? 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 = {}, +) { + 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), + fetch, + noExitFlush: true, + onFlushError, + ...options, + }); + 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(() => { + 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("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}`], + ["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 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; + 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 apiKey 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, + apiKey: "private-api-key", + 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"); + // 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" }); + expect(apiKeyLogger.loggingState).toBe(_internalGetGlobalState()); + }); + + 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()); + 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())); + 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) => + 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 throttled writes after Retry-After against the ingestion endpoint only", async () => { + const { requests, fetch } = mockIngestion({ + respond: (request, index) => + 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", + "POST /deployment/base/ingest/v1/logs", + ]); + }); + + 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), + }); + 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, + }); + 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("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({ + 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("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); + + 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("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(); + 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 { key, 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("")); + // 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", + 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"); + }); +}); + +// 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, + ); + }, + ); +}); diff --git a/js/src/ingestion-key.ts b/js/src/ingestion-key.ts new file mode 100644 index 000000000..6c1fff368 --- /dev/null +++ b/js/src/ingestion-key.ts @@ -0,0 +1,193 @@ +import { z } from "zod/v3"; +import { + BraintrustAttachmentReference as braintrustAttachmentReferenceSchema, + ExternalAttachmentReference as externalAttachmentReferenceSchema, + 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_PATTERN = /^bt-ik-[A-Za-z0-9]{48}$/; + +/** + * 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 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 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 }; +} + +// 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; +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. + * + * 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, +): "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" && + references.has(reference.key) + ) { + const uploaded = references.get(reference.key); + if (!uploaded) { + return "failed_upload"; + } + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + (container as Record)[key] = uploaded; + continue; + } + const problem = replaceIngestionAttachmentReferences(reference, references); + if (problem) { + return problem; + } + } + return undefined; +} + +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, +}); diff --git a/js/src/logger.ts b/js/src/logger.ts index 3a5be45ca..1130290f6 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 }; @@ -162,6 +165,17 @@ 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, +} from "./ingestion-key"; import { parseTemplateFormat, renderTemplateContent, @@ -767,6 +781,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 +809,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 +860,67 @@ 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, + orgId: string | undefined, + loginParams: Pick< + LoginOptions, + "fetch" | "noExitFlush" | "onFlushError" | "debugLogLevel" + >, + ): BraintrustState { + const cacheKey = JSON.stringify([endpoint.root, endpoint.key, orgId]); + let state = this.ingestionKeyStates.get(cacheKey); + if (!state) { + 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. + 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 async _internalFlushIngestionKeyStates(): Promise { + await Promise.all( + [...this.ingestionKeyStates.values()].map((state) => + state.bgLogger().flush(), + ), + ); + } + /** @internal */ public _internalSetTraceContextSigningSecret( secret: string | undefined, @@ -1188,7 +1260,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 +1475,119 @@ 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, + private readonly endpoint: IngestionEndpoint, + private readonly ingestionOrgId: string | undefined, + loginParams: LoginOptions, + ) { + super(loginParams); + this.ingestionConn = new HTTPConnection( + endpoint.root, + ingestionKeyFetch(this.fetch, endpoint.key), + ); + 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(ingestionKeyFetch(fetch, this.endpoint.key)); + } + + protected override createBackgroundLogger( + loginParams: LoginOptions, + ): HTTPBackgroundLogger { + return new HTTPBackgroundLogger( + new LazyValue(async () => this.ingestionConn), + loginParams, + { orgId: this.ingestionOrgId }, + ); + } +} + +/** + * 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; +} + let _globalState: BraintrustState; // Return a TestBackgroundLogger that will intercept logs before they are sent to the server. @@ -1492,6 +1681,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})`); @@ -1579,7 +1770,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 { @@ -2376,6 +2569,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 +2604,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 +2632,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 +2701,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 () => @@ -2527,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, @@ -2606,9 +2819,24 @@ 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) { + // 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) { @@ -3252,6 +3480,36 @@ 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 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; @@ -3347,6 +3605,44 @@ 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 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) + ) { + 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) || @@ -3425,6 +3721,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; @@ -3464,7 +3765,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 ingestion?: IngestionLoggerOptions, + ) { opts = opts ?? {}; this.apiConn = apiConn; @@ -3611,6 +3923,15 @@ class HTTPBackgroundLogger implements BackgroundLogger { maxRequestSize: number; canUseOverflow: boolean; }> { + if (!this._maxRequestSizePromise && this.ingestion) { + // Ingestion keys cannot read the server version, and the ingestion + // endpoint always accepts overflow uploads. + this._maxRequestSizePromise = Promise.resolve({ + maxRequestSize: + this.maxRequestSizeOverride ?? INGESTION_KEY_MAX_REQUEST_SIZE, + canUseOverflow: true, + }); + } if (!this._maxRequestSizePromise) { this._maxRequestSizePromise = (async () => { let serverLimit: number | null = null; @@ -3689,8 +4010,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.ingestion + ? await this.prepareIngestionRows( + unwrappedItems, + attachments, + attachmentErrors, + ) + : unwrappedItems; if (allItems.length === 0) { + throwAttachmentErrors(attachmentErrors); return; } @@ -3733,9 +4066,8 @@ class HTTPBackgroundLogger implements BackgroundLogger { ); } - const attachmentErrors: unknown[] = []; // For now, upload attachments serially. - for (const attachment of attachments) { + for (const attachment of this.ingestion ? [] : attachments) { try { const result = await attachment.upload(); if (result.upload_status === "error") { @@ -3745,14 +4077,217 @@ 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[] = []; + let externalAttachmentRows = 0; + for (const item of items) { + const row = pickIngestionRowFields(item, droppedFields); + 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); + 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 IngestionProtocolError( + "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) { + const retryDelayMs = ingestionRetryDelayMs(error, i); + if (i + 1 >= this.numTries || retryDelayMs === undefined) { + throw error; + } + debugLogger.warn( + `Ingestion upload failed. Retrying\nError: ${formatHTTPError(error)}`, + ); + 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(); + // The grant lifetime is relative, so anchor it before the request is sent. + const requestStart = monotonicNow(); + const grant = parseIngestionResponse( + 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, + ...(this.ingestion?.orgId ? { org_id: this.ingestion.orgId } : {}), + }), + ); + if (grant.num_chunks !== Math.ceil(data.size / grant.chunk_bytes)) { + throw new IngestionProtocolError( + `Invalid response from ingestion endpoint: ${grant.num_chunks} chunks of ${grant.chunk_bytes} bytes cannot hold ${data.size} bytes`, + ); + } + // 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++) { + 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, + signal, + }, + ), ); + const ack = parseIngestionResponse( + ingestionUploadChunkSchema, + await readJSONResponse(response), + ); + if (ack.index !== index || ack.size_bytes !== chunk.size) { + 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`, + ); + } } + + const completed = parseIngestionResponse( + ingestionUploadCompleteSchema, + await conn.post_json( + `v1/uploads/${grant.upload_id}/complete`, + {}, + signal, + ), + ); + if (completed.size_bytes !== data.size) { + throw new IngestionProtocolError( + "Invalid response from ingestion endpoint: the completed upload does not match the uploaded data", + ); + } + return completed.reference; } private async unwrapLazyValues( @@ -3870,15 +4405,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.ingestion + ? 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.ingestion) { + 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 IngestionProtocolError( + "Invalid response from ingestion endpoint: expected an overflow reference", + ); + } + ingestionOverflowKey = reference.key; + } + await conn.post_json( + "v1/logs", + constructLogs3OverflowRequest(ingestionOverflowKey), + ); + } else if (this.ingestion) { + await conn.post_json("v1/logs", dataStr); + } else if (overflowRows) { if (!overflowUpload) { const currentUpload = await this.requestLogs3OverflowUpload(conn, { rows: overflowRows, @@ -3904,8 +4460,15 @@ 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 isRetrying = i + 1 < this.numTries; + const retryDelayMs = this.ingestion + ? 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) { @@ -3936,11 +4499,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)); } } } @@ -5162,6 +5722,7 @@ export type InitLoggerOptions = FullLoginOptions & { setCurrent?: boolean; state?: BraintrustState; orgProjectMetadata?: OrgProjectMetadata; + ingestionKey?: string; } & AsyncFlushArg; /** @@ -5178,12 +5739,28 @@ 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 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> = {}, ) { + if (options.ingestionKey !== undefined && options.apiKey !== undefined) { + throw new Error( + "initLogger accepts either an ingestionKey or an apiKey, not both", + ); + } + // An explicit API key takes precedence over the environment. + const ingestionKey = + options.ingestionKey ?? + (options.apiKey === undefined + ? iso.getEnv(INGESTION_KEY_ENV_VAR) + : undefined); + if (ingestionKey !== undefined) { + return initIngestionKeyLogger(ingestionKey, options); + } + const { projectName, projectId, @@ -5248,6 +5825,63 @@ export function initLogger( return ret; } +function initIngestionKeyLogger( + ingestionKey: string, + { + projectId, + asyncFlush = true as IsAsyncFlush, + appUrl, + orgName, + environment, + setCurrent, + fetch, + 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 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: orgProjectMetadata?.org_id ?? "", + project: { id: resolvedProjectId ?? "", name: "", fullInfo: {} }, + })); + const ret = new Logger(state, lazyMetadata, { + asyncFlush, + computeMetadataArgs: resolvedProjectId + ? { project_id: resolvedProjectId } + : undefined, + linkArgs: { + org_name: orgName, + app_url: appUrl, + project_id: resolvedProjectId, + }, + }); + state.currentLogger = ret as Logger; + if (setCurrent ?? true) { + contextState.currentLogger = ret as Logger; + } + return ret; +} + export type LoadPromptOptions = FullLoginOptions & { projectName?: string; projectId?: string; @@ -6555,7 +7189,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 +7207,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 +7680,11 @@ 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 queues. + await Promise.all([ + state.bgLogger().flush(), + state._internalFlushIngestionKeyStates(), + ]); } /** @@ -7067,6 +7716,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 +9199,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.