diff --git a/js/src/exports.ts b/js/src/exports.ts index c01c1e3c..45e38f66 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -72,6 +72,7 @@ export { SpanImpl, TestBackgroundLogger, _exportsForTestingOnly, + _experimental, _internalGetGlobalState, _internalSetInitialState, constructLogs3OverflowRequest, diff --git a/js/src/logger.test.ts b/js/src/logger.test.ts index fc4f14f7..4639c3fe 100644 --- a/js/src/logger.test.ts +++ b/js/src/logger.test.ts @@ -3,10 +3,12 @@ import { vi, expect, test, describe, beforeEach, afterEach } from "vitest"; import { _exportsForTestingOnly, + _experimental, init, initLogger, Prompt, BraintrustState, + currentLogger, loadPrompt, loadParameters, wrapTraced, @@ -1396,6 +1398,125 @@ describe("wrapTraced generator support", () => { }); }); +describe("_experimental.withConfig", () => { + beforeEach(() => { + _exportsForTestingOnly.simulateLogoutForTests(); + }); + + test("scopes logger selection and restores existing current logger", () => { + const baseLogger = initLogger({ + projectName: "base-project", + projectId: "base-project-id", + }); + expect(currentLogger()).toBe(baseLogger); + + _experimental.withConfig({ project: "scoped-project" }, () => { + const scopedLogger = currentLogger(); + expect(scopedLogger).toBeDefined(); + expect(scopedLogger).not.toBe(baseLogger); + }); + + expect(currentLogger()).toBe(baseLogger); + }); + + test("maintains scoped logger across async boundaries", async () => { + const baseLogger = initLogger({ + projectName: "base-project", + projectId: "base-project-id", + }); + let scopedLogger: ReturnType | undefined; + + await _experimental.withConfig({ project: "async-project" }, async () => { + await Promise.resolve(); + scopedLogger = currentLogger(); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(currentLogger()).toBe(scopedLogger); + }); + + expect(scopedLogger).toBeDefined(); + expect(currentLogger()).toBe(baseLogger); + }); + + test("nested calls reuse scoped logger and apply inner project config", () => { + initLogger({ + projectName: "base-project", + projectId: "base-project-id", + }); + + let outerLogger: ReturnType | undefined; + let innerLogger: ReturnType | undefined; + + _experimental.withConfig({ project: "outer-project" }, () => { + outerLogger = currentLogger(); + const outerSpan = startSpan({ name: "outer-project-span" }); + expect(outerSpan.getParentInfo()?.computeObjectMetadataArgs).toEqual({ + project_name: "outer-project", + }); + outerSpan.end(); + + _experimental.withConfig({ project: "inner-project" }, () => { + innerLogger = currentLogger(); + expect(innerLogger).toBeDefined(); + expect(innerLogger).toBe(outerLogger); + const innerSpan = startSpan({ name: "inner-project-span" }); + expect(innerSpan.getParentInfo()?.computeObjectMetadataArgs).toEqual({ + project_name: "inner-project", + }); + innerSpan.end(); + }); + + const outerSpanAfterInner = startSpan({ name: "outer-project-span-2" }); + expect( + outerSpanAfterInner.getParentInfo()?.computeObjectMetadataArgs, + ).toEqual({ + project_name: "outer-project", + }); + outerSpanAfterInner.end(); + + expect(currentLogger()).toBe(outerLogger); + }); + }); + + test("top-level scopes reuse cached experimental logger instance", () => { + initLogger({ + projectName: "base-project", + projectId: "base-project-id", + }); + + let firstLogger: ReturnType | undefined; + let secondLogger: ReturnType | undefined; + + _experimental.withConfig({ project: "project-one" }, () => { + firstLogger = currentLogger(); + const span = startSpan({ name: "project-one-span" }); + expect(span.getParentInfo()?.computeObjectMetadataArgs).toEqual({ + project_name: "project-one", + }); + span.end(); + }); + + _experimental.withConfig({ project: "project-two" }, () => { + secondLogger = currentLogger(); + const span = startSpan({ name: "project-two-span" }); + expect(span.getParentInfo()?.computeObjectMetadataArgs).toEqual({ + project_name: "project-two", + }); + span.end(); + }); + + expect(firstLogger).toBeDefined(); + expect(secondLogger).toBe(firstLogger); + }); + + test("requires a non-empty project name", () => { + expect(() => + _experimental.withConfig({ project: " " }, () => {}), + ).toThrow( + "withConfig requires a non-empty project name in `configuration.project`", + ); + }); +}); + describe("parent precedence", () => { let memory: any; diff --git a/js/src/logger.ts b/js/src/logger.ts index 658d47ae..4ef0117e 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -588,6 +588,9 @@ const loginSchema = z.strictObject({ }); export type SerializedBraintrustState = z.infer; +type WithConfigOptions = { + project: string; +}; let stateNonce = 0; @@ -597,6 +600,12 @@ export class BraintrustState { // Note: the value of IsAsyncFlush doesn't really matter here, since we // (safely) dynamically cast it whenever retrieving the logger. public currentLogger: Logger | undefined; + public scopedLogger: IsoAsyncLocalStorage | undefined>; + public scopedConfig: IsoAsyncLocalStorage; + public cachedWithConfigLoggers: { + async?: Logger; + sync?: Logger; + }; public currentParent: IsoAsyncLocalStorage; public currentSpan: IsoAsyncLocalStorage; // Any time we re-log in, we directly update the apiConn inside the logger. @@ -634,6 +643,9 @@ export class BraintrustState { this.id = `${new Date().toLocaleString()}-${stateNonce++}`; // This is for debugging. uuidv4() breaks on platforms like Cloudflare. this.currentExperiment = undefined; this.currentLogger = undefined; + this.scopedLogger = iso.newAsyncLocalStorage(); + this.scopedConfig = iso.newAsyncLocalStorage(); + this.cachedWithConfigLoggers = {}; this.currentParent = iso.newAsyncLocalStorage(); this.currentSpan = iso.newAsyncLocalStorage(); @@ -1302,6 +1314,13 @@ interface OrgProjectMetadata { project: ObjectMetadata; } +interface LoggerRuntimeContext { + lazyMetadata: LazyValue; + lazyId: LazyValue; + computeMetadataArgs?: Record; + linkArgs?: LinkArgs; +} + export interface LinkArgs { org_name?: string; app_url?: string; @@ -1313,6 +1332,8 @@ export interface LogOptions { asyncFlush?: IsAsyncFlush; computeMetadataArgs?: Record; linkArgs?: LinkArgs; + // Internal: resolve logging context dynamically (e.g. withConfig scoped config). + _resolveRuntimeContext?: () => LoggerRuntimeContext; } export type PromiseUnless = B extends true ? R : Promise>; @@ -2235,6 +2256,7 @@ export class Logger implements Exportable { private _asyncFlush: IsAsyncFlush | undefined; private computeMetadataArgs: Record | undefined; private _linkArgs: LinkArgs | undefined; + private _resolveRuntimeContext: (() => LoggerRuntimeContext) | undefined; private lastStartTime: number; private lazyId: LazyValue; private calledStartSpan: boolean; @@ -2251,21 +2273,37 @@ export class Logger implements Exportable { this._asyncFlush = logOptions.asyncFlush; this.computeMetadataArgs = logOptions.computeMetadataArgs; this._linkArgs = logOptions.linkArgs; + this._resolveRuntimeContext = logOptions._resolveRuntimeContext; this.lastStartTime = getCurrentUnixTimestamp(); this.lazyId = new LazyValue(async () => await this.id); this.calledStartSpan = false; this.state = state; } + private getRuntimeContext(): LoggerRuntimeContext { + if (this._resolveRuntimeContext) { + return this._resolveRuntimeContext(); + } + + return { + lazyMetadata: this.lazyMetadata, + lazyId: this.lazyId, + computeMetadataArgs: this.computeMetadataArgs, + linkArgs: this._linkArgs, + }; + } + public get org_id(): Promise { + const context = this.getRuntimeContext(); return (async () => { - return (await this.lazyMetadata.get()).org_id; + return (await context.lazyMetadata.get()).org_id; })(); } public get project(): Promise { + const context = this.getRuntimeContext(); return (async () => { - return (await this.lazyMetadata.get()).project; + return (await context.lazyMetadata.get()).project; })(); } @@ -2373,6 +2411,7 @@ export class Logger implements Exportable { } private startSpanImpl(args?: StartSpanArgs): Span { + const context = this.getRuntimeContext(); return new SpanImpl({ ...args, // Sometimes `args` gets passed directly into this function, and it contains an undefined value for `state`. @@ -2382,8 +2421,8 @@ export class Logger implements Exportable { state: this.state, parent: args?.parent, parentObjectType: this.parentObjectType(), - parentObjectId: this.lazyId, - parentComputeObjectMetadataArgs: this.computeMetadataArgs, + parentObjectId: context.lazyId, + parentComputeObjectMetadataArgs: context.computeMetadataArgs, parentSpanIds: args?.parentSpanIds, propagatedEvent: args?.propagatedEvent, }), @@ -2403,7 +2442,8 @@ export class Logger implements Exportable { * @param event.source (Optional) the source of the feedback. Must be one of "external" (default), "app", or "api". */ public logFeedback(event: LogFeedbackFullArgs): void { - logFeedbackImpl(this.state, this.parentObjectType(), this.lazyId, event); + const context = this.getRuntimeContext(); + logFeedbackImpl(this.state, this.parentObjectType(), context.lazyId, event); } /** @@ -2423,7 +2463,7 @@ export class Logger implements Exportable { updateSpanImpl({ state: this.state, parentObjectType: this.parentObjectType(), - parentObjectId: this.lazyId, + parentObjectId: this.getRuntimeContext().lazyId, id, root_span_id, span_id, @@ -2437,15 +2477,16 @@ export class Logger implements Exportable { * See {@link Span.startSpan} for more details. */ public async export(): Promise { + const context = this.getRuntimeContext(); // Note: it is important that the object id we are checking for // `has_computed` is the same as the one we are passing into the span // logging functions. So that if the spans actually do get logged, then this // `_lazy_id` object specifically will also be marked as computed. return new (getSpanComponentsClass())({ object_type: this.parentObjectType(), - ...(this.computeMetadataArgs && !this.lazyId.hasSucceeded - ? { compute_object_metadata_args: this.computeMetadataArgs } - : { object_id: await this.lazyId.get() }), + ...(context.computeMetadataArgs && !context.lazyId.hasSucceeded + ? { compute_object_metadata_args: context.computeMetadataArgs } + : { object_id: await context.lazyId.get() }), }).toStr(); } @@ -2466,7 +2507,7 @@ export class Logger implements Exportable { * Resolution order: state -> linkArgs -> env var */ public _getLinkBaseUrl(): string | null { - return _getLinkBaseUrl(this.state, this._linkArgs); + return _getLinkBaseUrl(this.state, this.getRuntimeContext().linkArgs); } } @@ -3763,6 +3804,131 @@ export function withLogger( return callback(logger); } +function normalizeWithConfig( + configuration: Readonly, +): WithConfigOptions { + const project = configuration.project?.trim(); + if (typeof project !== "string" || project.length === 0) { + throw new Error( + "withConfig requires a non-empty project name in `configuration.project`", + ); + } + + return { project }; +} + +function reconcileWithConfig( + parentConfig: Readonly | undefined, + nextConfig: Readonly, +): WithConfigOptions { + return { + // Nested scopes use the innermost project. + project: + parentConfig?.project === nextConfig.project + ? parentConfig.project + : nextConfig.project, + }; +} + +function getConfig(state: BraintrustState): WithConfigOptions | undefined { + return state.scopedConfig.getStore(); +} + +function createLoggerRuntimeContext( + state: BraintrustState, + config: Readonly, +): LoggerRuntimeContext { + const computeMetadataArgs = { + project_name: config.project, + }; + + const lazyMetadata: LazyValue = new LazyValue( + async () => { + await state.login({}); + return computeLoggerMetadata(state, computeMetadataArgs); + }, + ); + + const lazyId: LazyValue = new LazyValue( + async () => (await lazyMetadata.get()).project.id, + ); + + return { + lazyMetadata, + lazyId, + computeMetadataArgs, + linkArgs: { project_name: config.project }, + }; +} + +function createScopedLogger( + state: BraintrustState, + initialConfig: Readonly, + asyncFlush: boolean, +): Logger { + const runtimeContextCache = new Map(); + + const getRuntimeContextForConfig = ( + config: Readonly, + ): LoggerRuntimeContext => { + const cached = runtimeContextCache.get(config.project); + if (cached) { + return cached; + } + const created = createLoggerRuntimeContext(state, config); + runtimeContextCache.set(config.project, created); + return created; + }; + + const initialContext = getRuntimeContextForConfig(initialConfig); + + return new Logger(state, initialContext.lazyMetadata, { + asyncFlush, + computeMetadataArgs: initialContext.computeMetadataArgs, + linkArgs: initialContext.linkArgs, + _resolveRuntimeContext: () => + getRuntimeContextForConfig(getConfig(state) ?? initialConfig), + }); +} + +function getOrCreateScopedLogger( + state: BraintrustState, + initialConfig: Readonly, +): Logger { + const inheritedAsyncFlush = state.currentLogger?.asyncFlush ?? true; + const cacheKey = inheritedAsyncFlush ? "async" : "sync"; + + const cached = state.cachedWithConfigLoggers[cacheKey]; + if (cached) { + return cached; + } + + const created = createScopedLogger(state, initialConfig, inheritedAsyncFlush); + state.cachedWithConfigLoggers[cacheKey] = created; + return created; +} + +function withConfig( + configuration: Readonly, + callback: () => R, +): R { + const state = _globalState; + const normalizedConfig = normalizeWithConfig(configuration); + const parentConfig = state.scopedConfig.getStore(); + const resolvedConfig = reconcileWithConfig(parentConfig, normalizedConfig); + const logger = + state.scopedLogger.getStore() ?? + getOrCreateScopedLogger(state, resolvedConfig); + + return state.scopedConfig.run(resolvedConfig, () => + state.scopedLogger.run(logger, callback), + ); +} + +export const _experimental = { + withConfig, +} as const; + type UseOutputOption = { useOutput?: IsLegacyDataset; }; @@ -4737,7 +4903,10 @@ export function currentLogger( options?: AsyncFlushArg & OptionalStateArg, ): Logger | undefined { const state = options?.state ?? _globalState; - return castLogger(state.currentLogger, options?.asyncFlush); + return castLogger( + state.scopedLogger.getStore() ?? state.currentLogger, + options?.asyncFlush, + ); } /**