From 0fd3e28b69d8f37059ac70927c2472623fd3f05e Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:48:50 +0000 Subject: [PATCH 1/4] perf: Cache project lookup --- e2e/config/pr-comment-scenarios.json | 5 + .../jest-v29-latest.request-flow.json | 430 ++++++--------- .../__snapshots__/jest-v29.request-flow.json | 430 ++++++--------- .../scenario.test.ts | 5 + .../__snapshots__/request-flow.json | 4 +- .../__snapshots__/span-tree.json | 2 +- .../__snapshots__/span-tree.txt | 2 +- .../trace-primitives-basic/scenario.test.ts | 7 + .../trace-primitives-basic/scenario.ts | 17 + js/src/logger-caching.test.ts | 511 ++++++++++++++++++ js/src/logger.ts | 164 +++++- 11 files changed, 995 insertions(+), 582 deletions(-) create mode 100644 js/src/logger-caching.test.ts diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index 26b4ada1d..b13f90d8e 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -1,4 +1,9 @@ [ + { + "scenarioDirName": "trace-primitives-basic", + "label": "Trace primitives and project caching", + "metadataScenario": "trace-primitives-basic" + }, { "scenarioDirName": "deepseek-harness-instrumentation", "label": "DeepSeek Harness Instrumentation", diff --git a/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29-latest.request-flow.json b/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29-latest.request-flow.json index d2632da8a..7d866f4f9 100644 --- a/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29-latest.request-flow.json +++ b/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29-latest.request-flow.json @@ -56,115 +56,7 @@ "type": "task" }, "span_id": "" - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ - { - "context": { - "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "Object.", - "caller_lineno": 0 - }, - "created": "", - "expected": "Paris", - "id": "", - "input": "What is the capital of France?", - "log_id": "g", - "metadata": { - "case": "basic-span", - "scenario": "test-framework-evals-jest", - "testRunId": "", - "transport": "http" - }, - "metrics": { - "end": 0, - "start": 0 - }, - "output": "Paris", - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 0, - "name": "jest basic span", - "type": "task" - }, - "span_id": "" - } - ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" - }, - "method": "POST", - "path": "/api/project/register", - "query": null, - "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { - "api_version": 2, - "rows": [ - { - "context": { - "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "Object.", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "input": { - "transcript": { - "content_type": "application/json", - "filename": "conversation_transcript.json", - "key": "", - "type": "braintrust_attachment" - }, - "type": "chat_completion" - }, - "log_id": "g", - "metadata": { - "case": "json-attachment", - "scenario": "test-framework-evals-jest", - "testRunId": "" - }, - "metrics": { - "end": 0, - "start": 0 - }, - "output": { - "attachment": true - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 1, - "name": "root", - "type": "task" - }, - "span_id": "" - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -203,29 +95,7 @@ "type": "task" }, "span_id": "" - } - ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" - }, - "method": "POST", - "path": "/api/project/register", - "query": null, - "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -297,15 +167,7 @@ "span_parents": [ "" ] - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -313,14 +175,10 @@ "caller_lineno": 0 }, "created": "", - "id": "", - "input": { - "phase": "parent", - "testRunId": "" - }, + "id": "", "log_id": "g", "metadata": { - "case": "parent-span", + "case": "nested-parent", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -328,18 +186,14 @@ "end": 0, "start": 0 }, - "output": { - "ok": true, - "phase": "parent" - }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 2, - "name": "jest parent span", + "exec_counter": 4, + "name": "jest nested parent span", "type": "task" }, - "span_id": "" + "span_id": "" }, { "context": { @@ -348,14 +202,39 @@ "caller_lineno": 0 }, "created": "", - "id": "", - "input": { - "step": "child", + "id": "", + "log_id": "g", + "metadata": { + "case": "nested-child", + "scenario": "test-framework-evals-jest", "testRunId": "" }, + "metrics": { + "end": 0, + "start": 0 + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 5, + "name": "jest nested child span" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": { + "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", + "caller_functionname": "traced.name", + "caller_lineno": 0 + }, + "created": "", + "id": "", "log_id": "g", "metadata": { - "case": "child-span", + "case": "nested-grandchild", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -364,40 +243,55 @@ "start": 0 }, "output": { - "ok": true, - "phase": "child" + "depth": 3 }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 3, - "name": "jest child span" + "exec_counter": 6, + "name": "jest nested grandchild span" }, - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ] + }, + { + "context": { + "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", + "caller_functionname": "Object.", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "case": "current-span", + "scenario": "test-framework-evals-jest", + "testRunId": "" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "output": { + "observedSpanId": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 7, + "name": "jest current span", + "type": "task" + }, + "span_id": "" } ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" }, "method": "POST", - "path": "/api/project/register", + "path": "/logs3", "query": null, "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { "api_version": 2, "rows": [ { @@ -407,10 +301,50 @@ "caller_lineno": 0 }, "created": "", - "id": "", + "expected": "Paris", + "id": "", + "input": "What is the capital of France?", "log_id": "g", "metadata": { - "case": "nested-parent", + "case": "basic-span", + "scenario": "test-framework-evals-jest", + "testRunId": "", + "transport": "http" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "output": "Paris", + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 0, + "name": "jest basic span", + "type": "task" + }, + "span_id": "" + }, + { + "context": { + "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", + "caller_functionname": "Object.", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "input": { + "transcript": { + "content_type": "application/json", + "filename": "conversation_transcript.json", + "key": "", + "type": "braintrust_attachment" + }, + "type": "chat_completion" + }, + "log_id": "g", + "metadata": { + "case": "json-attachment", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -418,26 +352,33 @@ "end": 0, "start": 0 }, + "output": { + "attachment": true + }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 4, - "name": "jest nested parent span", + "exec_counter": 1, + "name": "root", "type": "task" }, - "span_id": "" + "span_id": "" }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "traced.name", + "caller_functionname": "Object.", "caller_lineno": 0 }, "created": "", - "id": "", + "id": "", + "input": { + "phase": "parent", + "testRunId": "" + }, "log_id": "g", "metadata": { - "case": "nested-child", + "case": "parent-span", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -445,16 +386,18 @@ "end": 0, "start": 0 }, + "output": { + "ok": true, + "phase": "parent" + }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 5, - "name": "jest nested child span" + "exec_counter": 2, + "name": "jest parent span", + "type": "task" }, - "span_id": "", - "span_parents": [ - "" - ] + "span_id": "" }, { "context": { @@ -463,10 +406,14 @@ "caller_lineno": 0 }, "created": "", - "id": "", + "id": "", + "input": { + "step": "child", + "testRunId": "" + }, "log_id": "g", "metadata": { - "case": "nested-grandchild", + "case": "child-span", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -475,27 +422,20 @@ "start": 0 }, "output": { - "depth": 3 + "ok": true, + "phase": "child" }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 6, - "name": "jest nested grandchild span" + "exec_counter": 3, + "name": "jest child span" }, - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ] - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -583,67 +523,7 @@ "span_parents": [ "" ] - } - ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" - }, - "method": "POST", - "path": "/api/project/register", - "query": null, - "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { - "api_version": 2, - "rows": [ - { - "context": { - "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "Object.", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "case": "current-span", - "scenario": "test-framework-evals-jest", - "testRunId": "" - }, - "metrics": { - "end": 0, - "start": 0 - }, - "output": { - "observedSpanId": "" - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 7, - "name": "jest current span", - "type": "task" - }, - "span_id": "" - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", diff --git a/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29.request-flow.json b/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29.request-flow.json index d2632da8a..7d866f4f9 100644 --- a/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29.request-flow.json +++ b/e2e/scenarios/test-framework-evals-jest/__snapshots__/jest-v29.request-flow.json @@ -56,115 +56,7 @@ "type": "task" }, "span_id": "" - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ - { - "context": { - "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "Object.", - "caller_lineno": 0 - }, - "created": "", - "expected": "Paris", - "id": "", - "input": "What is the capital of France?", - "log_id": "g", - "metadata": { - "case": "basic-span", - "scenario": "test-framework-evals-jest", - "testRunId": "", - "transport": "http" - }, - "metrics": { - "end": 0, - "start": 0 - }, - "output": "Paris", - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 0, - "name": "jest basic span", - "type": "task" - }, - "span_id": "" - } - ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" - }, - "method": "POST", - "path": "/api/project/register", - "query": null, - "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { - "api_version": 2, - "rows": [ - { - "context": { - "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "Object.", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "input": { - "transcript": { - "content_type": "application/json", - "filename": "conversation_transcript.json", - "key": "", - "type": "braintrust_attachment" - }, - "type": "chat_completion" - }, - "log_id": "g", - "metadata": { - "case": "json-attachment", - "scenario": "test-framework-evals-jest", - "testRunId": "" - }, - "metrics": { - "end": 0, - "start": 0 - }, - "output": { - "attachment": true - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 1, - "name": "root", - "type": "task" - }, - "span_id": "" - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -203,29 +95,7 @@ "type": "task" }, "span_id": "" - } - ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" - }, - "method": "POST", - "path": "/api/project/register", - "query": null, - "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -297,15 +167,7 @@ "span_parents": [ "" ] - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -313,14 +175,10 @@ "caller_lineno": 0 }, "created": "", - "id": "", - "input": { - "phase": "parent", - "testRunId": "" - }, + "id": "", "log_id": "g", "metadata": { - "case": "parent-span", + "case": "nested-parent", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -328,18 +186,14 @@ "end": 0, "start": 0 }, - "output": { - "ok": true, - "phase": "parent" - }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 2, - "name": "jest parent span", + "exec_counter": 4, + "name": "jest nested parent span", "type": "task" }, - "span_id": "" + "span_id": "" }, { "context": { @@ -348,14 +202,39 @@ "caller_lineno": 0 }, "created": "", - "id": "", - "input": { - "step": "child", + "id": "", + "log_id": "g", + "metadata": { + "case": "nested-child", + "scenario": "test-framework-evals-jest", "testRunId": "" }, + "metrics": { + "end": 0, + "start": 0 + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 5, + "name": "jest nested child span" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": { + "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", + "caller_functionname": "traced.name", + "caller_lineno": 0 + }, + "created": "", + "id": "", "log_id": "g", "metadata": { - "case": "child-span", + "case": "nested-grandchild", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -364,40 +243,55 @@ "start": 0 }, "output": { - "ok": true, - "phase": "child" + "depth": 3 }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 3, - "name": "jest child span" + "exec_counter": 6, + "name": "jest nested grandchild span" }, - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ] + }, + { + "context": { + "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", + "caller_functionname": "Object.", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "case": "current-span", + "scenario": "test-framework-evals-jest", + "testRunId": "" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "output": { + "observedSpanId": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 7, + "name": "jest current span", + "type": "task" + }, + "span_id": "" } ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" }, "method": "POST", - "path": "/api/project/register", + "path": "/logs3", "query": null, "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { "api_version": 2, "rows": [ { @@ -407,10 +301,50 @@ "caller_lineno": 0 }, "created": "", - "id": "", + "expected": "Paris", + "id": "", + "input": "What is the capital of France?", "log_id": "g", "metadata": { - "case": "nested-parent", + "case": "basic-span", + "scenario": "test-framework-evals-jest", + "testRunId": "", + "transport": "http" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "output": "Paris", + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 0, + "name": "jest basic span", + "type": "task" + }, + "span_id": "" + }, + { + "context": { + "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", + "caller_functionname": "Object.", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "input": { + "transcript": { + "content_type": "application/json", + "filename": "conversation_transcript.json", + "key": "", + "type": "braintrust_attachment" + }, + "type": "chat_completion" + }, + "log_id": "g", + "metadata": { + "case": "json-attachment", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -418,26 +352,33 @@ "end": 0, "start": 0 }, + "output": { + "attachment": true + }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 4, - "name": "jest nested parent span", + "exec_counter": 1, + "name": "root", "type": "task" }, - "span_id": "" + "span_id": "" }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "traced.name", + "caller_functionname": "Object.", "caller_lineno": 0 }, "created": "", - "id": "", + "id": "", + "input": { + "phase": "parent", + "testRunId": "" + }, "log_id": "g", "metadata": { - "case": "nested-child", + "case": "parent-span", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -445,16 +386,18 @@ "end": 0, "start": 0 }, + "output": { + "ok": true, + "phase": "parent" + }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 5, - "name": "jest nested child span" + "exec_counter": 2, + "name": "jest parent span", + "type": "task" }, - "span_id": "", - "span_parents": [ - "" - ] + "span_id": "" }, { "context": { @@ -463,10 +406,14 @@ "caller_lineno": 0 }, "created": "", - "id": "", + "id": "", + "input": { + "step": "child", + "testRunId": "" + }, "log_id": "g", "metadata": { - "case": "nested-grandchild", + "case": "child-span", "scenario": "test-framework-evals-jest", "testRunId": "" }, @@ -475,27 +422,20 @@ "start": 0 }, "output": { - "depth": 3 + "ok": true, + "phase": "child" }, "project_id": "", - "root_span_id": "", + "root_span_id": "", "span_attributes": { - "exec_counter": 6, - "name": "jest nested grandchild span" + "exec_counter": 3, + "name": "jest child span" }, - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ] - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", @@ -583,67 +523,7 @@ "span_parents": [ "" ] - } - ] - } - }, - { - "headers": null, - "jsonBody": { - "org_id": "mock-org-id", - "project_name": "" - }, - "method": "POST", - "path": "/api/project/register", - "query": null, - "rawBody": { - "org_id": "mock-org-id", - "project_name": "" - } - }, - { - "headers": null, - "jsonBody": { - "api_version": 2, - "rows": [ - { - "context": { - "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", - "caller_functionname": "Object.", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "case": "current-span", - "scenario": "test-framework-evals-jest", - "testRunId": "" - }, - "metrics": { - "end": 0, - "start": 0 - }, - "output": { - "observedSpanId": "" - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 7, - "name": "jest current span", - "type": "task" - }, - "span_id": "" - } - ] - }, - "method": "POST", - "path": "/logs3", - "query": null, - "rawBody": { - "api_version": 2, - "rows": [ + }, { "context": { "caller_filename": "/e2e/scenarios/test-framework-evals-jest/runner.case.cjs", diff --git a/e2e/scenarios/test-framework-evals-jest/scenario.test.ts b/e2e/scenarios/test-framework-evals-jest/scenario.test.ts index 484f4d817..f19ca7555 100644 --- a/e2e/scenarios/test-framework-evals-jest/scenario.test.ts +++ b/e2e/scenarios/test-framework-evals-jest/scenario.test.ts @@ -182,6 +182,11 @@ for (const scenario of scenarios) { "/logs3", ]), ); + expect( + requests.filter( + (request) => request.path === "/api/project/register", + ), + ).toHaveLength(1); await matchSpanTreeSnapshot( capturedEvents, diff --git a/e2e/scenarios/trace-primitives-basic/__snapshots__/request-flow.json b/e2e/scenarios/trace-primitives-basic/__snapshots__/request-flow.json index 5738322b6..bcc576917 100644 --- a/e2e/scenarios/trace-primitives-basic/__snapshots__/request-flow.json +++ b/e2e/scenarios/trace-primitives-basic/__snapshots__/request-flow.json @@ -109,7 +109,7 @@ "caller_lineno": 0 }, "created": "", - "error": "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at runMain (/e2e/helpers/scenario-runtime.ts:0:0)", + "error": "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at process.processTicksAndRejections (node::0:0)", "id": "", "log_id": "g", "metadata": { @@ -214,7 +214,7 @@ "caller_lineno": 0 }, "created": "", - "error": "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at runMain (/e2e/helpers/scenario-runtime.ts:0:0)", + "error": "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at process.processTicksAndRejections (node::0:0)", "id": "", "log_id": "g", "metadata": { diff --git a/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.json b/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.json index 34a4e6d28..852eddd49 100644 --- a/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.json +++ b/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.json @@ -26,7 +26,7 @@ "kind": "basic-error", "testRunId": "" }, - "error": "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at runMain (/e2e/helpers/scenario-runtime.ts:0:0)" + "error": "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at process.processTicksAndRejections (node::0:0)" } ], "input": { diff --git a/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.txt b/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.txt index 3bb8e514b..dd6da9af4 100644 --- a/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.txt +++ b/e2e/scenarios/trace-primitives-basic/__snapshots__/span-tree.txt @@ -28,4 +28,4 @@ span_tree: "kind": "basic-error", "testRunId": "" } - error: "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at runMain (/e2e/helpers/scenario-runtime.ts:0:0)" + error: "basic boom\n\nError: basic boom\n at logger.traced.name (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at /js/dist/index.mjs:0:0\n at AsyncLocalStorage.run (node::0:0)\n at BraintrustContextManager.runInContext (/js/dist/index.mjs:0:0)\n at withCurrent (/js/dist/index.mjs:0:0)\n at /js/dist/index.mjs:0:0\n at runCatchFinally (/js/dist/index.mjs:0:0)\n at Logger.traced (/js/dist/index.mjs:0:0)\n at main (/e2e/scenarios/trace-primitives-basic/scenario.ts:0:0)\n at process.processTicksAndRejections (node::0:0)" diff --git a/e2e/scenarios/trace-primitives-basic/scenario.test.ts b/e2e/scenarios/trace-primitives-basic/scenario.test.ts index 796f27b86..ac13928d2 100644 --- a/e2e/scenarios/trace-primitives-basic/scenario.test.ts +++ b/e2e/scenarios/trace-primitives-basic/scenario.test.ts @@ -61,6 +61,13 @@ test("trace-primitives-basic collects a minimal manual trace tree", async () => request.path === "/logs3", ); + expect( + requests.filter((request) => request.path === "/api/apikey/login"), + ).toHaveLength(1); + expect( + requests.filter((request) => request.path === "/api/project/register"), + ).toHaveLength(1); + await matchFileSnapshot( formatJsonFileSnapshot( requests.map((request) => diff --git a/e2e/scenarios/trace-primitives-basic/scenario.ts b/e2e/scenarios/trace-primitives-basic/scenario.ts index 442ce6518..2bfacf03d 100644 --- a/e2e/scenarios/trace-primitives-basic/scenario.ts +++ b/e2e/scenarios/trace-primitives-basic/scenario.ts @@ -11,6 +11,23 @@ async function main() { projectName: scopedName("e2e-trace-primitives-basic", testRunId), }); + // Concurrent cold initialization should share both login and project lookup. + await Promise.all([ + logger.id, + ...Array.from( + { length: 10 }, + () => + initLogger({ + projectName: scopedName("e2e-trace-primitives-basic", testRunId), + setCurrent: false, + }).id, + ), + ]); + await initLogger({ + projectName: scopedName("e2e-trace-primitives-basic", testRunId), + setCurrent: false, + }).id; + await logger.traced( async (rootSpan) => { const childSpan = startSpan({ diff --git a/js/src/logger-caching.test.ts b/js/src/logger-caching.test.ts new file mode 100644 index 000000000..d9a486d19 --- /dev/null +++ b/js/src/logger-caching.test.ts @@ -0,0 +1,511 @@ +import { afterEach, describe, expect, test, vi } from "vitest"; +import { + BraintrustState, + initLogger, + spanComponentsToObjectId, + startSpan, + updateSpan, + _internalResumeSpan, +} from "./logger"; +import { configureNode } from "./node/config"; +import iso from "./isomorph"; +import { SpanComponentsV4 } from "../util/span_identifier_v4"; + +configureNode(); + +const projectId = "00000000-0000-0000-0000-000000000001"; +const ttl = 15 * 60 * 1000; + +function loginResponse(orgId = "org-id") { + return Response.json({ + org_info: [{ id: orgId, name: orgId, api_url: "https://api.test" }], + }); +} + +function createState() { + const fetch = vi.fn(async (input) => { + const url = new URL(String(input)); + switch (url.pathname) { + case "/api/apikey/login": + return loginResponse(); + case "/api/project/register": + return Response.json({ project: { id: projectId, name: "project" } }); + case "/api/project": + return Response.json({ + name: "project", + project: { id: url.searchParams.get("id"), name: "project" }, + }); + case "/version": + case "/logs3": + return Response.json({}); + default: + throw new Error(`Unexpected test request: ${url.pathname}`); + } + }); + const state = new BraintrustState({ + apiKey: "test-credential", + appUrl: "https://app.test", + fetch, + noExitFlush: true, + }); + return { state, fetch }; +} + +afterEach(() => { + vi.restoreAllMocks(); + vi.useRealTimers(); +}); + +describe("project metadata caching", () => { + test("shares concurrent and sequential lookups across loggers and imported parents", async () => { + const { state, fetch } = createState(); + const logger = initLogger({ + state, + projectName: "project", + setCurrent: false, + }); + const components = SpanComponentsV4.fromStr(await logger.export()); + const ids = await Promise.all([ + logger.id, + ...Array.from( + { length: 10 }, + () => + initLogger({ state, projectName: "project", setCurrent: false }).id, + ), + ...Array.from({ length: 10 }, () => + spanComponentsToObjectId({ state, components }), + ), + ]); + expect(ids).toEqual(Array(21).fill(projectId)); + await initLogger({ state, projectName: "project", setCurrent: false }).id; + expect( + fetch.mock.calls.map(([url]) => new URL(String(url)).pathname), + ).toEqual(["/api/apikey/login", "/api/project/register"]); + }); + + test("shares lookups from starting, updating, and resuming unresolved spans", async () => { + const { state, fetch } = createState(); + state.httpLogger().syncFlush = true; + const root = initLogger({ + state, + projectName: "project", + setCurrent: false, + }).startSpan(); + const exported = await root.export(); + startSpan({ state, parent: exported }).end(); + updateSpan({ state, exported, output: "updated" }); + _internalResumeSpan({ state, exported }).end(); + root.end(); + await state.bgLogger().flush(); + expect( + fetch.mock.calls.filter(([url]) => + String(url).endsWith("/api/project/register"), + ), + ).toHaveLength(1); + }); + + test("refreshes once at 15 minutes without sliding expiry or rebinding existing loggers", async () => { + vi.useFakeTimers({ toFake: ["Date"] }); + vi.setSystemTime(0); + const { state } = createState(); + await state.login({}); + const register = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValue({ project: { id: projectId, name: "project" } }); + const logger = initLogger({ + state, + projectName: "project", + setCurrent: false, + }); + await logger.id; + vi.setSystemTime(ttl - 1); + await initLogger({ state, projectName: "project", setCurrent: false }).id; + expect(register).toHaveBeenCalledTimes(1); + vi.setSystemTime(ttl); + register.mockResolvedValue({ + project: { id: "replacement-id", name: "project" }, + }); + expect( + await Promise.all( + Array.from( + { length: 10 }, + () => + initLogger({ state, projectName: "project", setCurrent: false }).id, + ), + ), + ).toEqual(Array(10).fill("replacement-id")); + expect(register).toHaveBeenCalledTimes(2); + expect(await logger.id).toBe(projectId); + vi.setSystemTime(2 * ttl - 1); + await initLogger({ state, projectName: "project", setCurrent: false }).id; + expect(register).toHaveBeenCalledTimes(2); + }); + + test("starts TTL on success and shares requests even when resolution takes over 15 minutes", async () => { + vi.useFakeTimers({ toFake: ["Date"] }); + vi.setSystemTime(0); + const { state } = createState(); + await state.login({}); + let resolve!: (value: unknown) => void; + const register = vi.spyOn(state.appConn(), "post_json").mockImplementation( + () => + new Promise((r) => { + resolve = r; + }), + ); + const first = initLogger({ + state, + projectName: "project", + setCurrent: false, + }).id; + await vi.waitFor(() => expect(register).toHaveBeenCalledTimes(1)); + vi.setSystemTime(ttl * 2); + const second = initLogger({ + state, + projectName: "project", + setCurrent: false, + }).id; + resolve({ project: { id: projectId, name: "project" } }); + await Promise.all([first, second]); + vi.setSystemTime(ttl * 3 - 1); + await initLogger({ state, projectName: "project", setCurrent: false }).id; + expect(register).toHaveBeenCalledTimes(1); + }); + + test("shares failures but permits retries from later loggers", async () => { + const { state } = createState(); + await state.login({}); + const register = vi + .spyOn(state.appConn(), "post_json") + .mockRejectedValueOnce(new Error("lookup failed")) + .mockResolvedValue({ project: { id: projectId, name: "project" } }); + const results = await Promise.allSettled( + Array.from( + { length: 10 }, + () => + initLogger({ state, projectName: "project", setCurrent: false }).id, + ), + ); + expect(results.every((result) => result.status === "rejected")).toBe(true); + expect(register).toHaveBeenCalledTimes(1); + expect( + await initLogger({ state, projectName: "project", setCurrent: false }).id, + ).toBe(projectId); + expect(register).toHaveBeenCalledTimes(2); + }); + + test("normalizes the default project and keeps name and ID lookups separate", async () => { + const { state } = createState(); + await state.login({}); + const register = vi.spyOn(state.appConn(), "post_json"); + const get = vi.spyOn(state.appConn(), "get_json"); + await Promise.all([ + initLogger({ state, setCurrent: false }).id, + initLogger({ state, projectName: "Global", setCurrent: false }).id, + initLogger({ state, projectName: projectId, setCurrent: false }).id, + initLogger({ state, projectId, setCurrent: false }).id, + initLogger({ state, projectId, setCurrent: false }).id, + ]); + expect(register).toHaveBeenCalledTimes(2); + expect(get).toHaveBeenCalledTimes(1); + await initLogger({ + state, + projectId, + projectName: "explicit", + setCurrent: false, + }).id; + expect(get).toHaveBeenCalledTimes(1); + expect(register).toHaveBeenCalledTimes(2); + }); + + test.each(["orgId", "loginToken", "appUrl"] as const)( + "isolates changes to %s", + async (field) => { + const { state } = createState(); + await state.login({}); + const register = vi.spyOn(state.appConn(), "post_json"); + await initLogger({ state, projectName: "project", setCurrent: false }).id; + state[field] = "changed"; + await initLogger({ state, projectName: "project", setCurrent: false }).id; + expect(register).toHaveBeenCalledTimes(2); + }, + ); + + test.each(["reset", "forceLogin", "fetch"])( + "invalidates on %s and does not reuse pending results", + async (action) => { + const { state, fetch } = createState(); + await state.login({}); + let resolve!: (value: unknown) => void; + const register = vi + .spyOn(state.appConn(), "post_json") + .mockImplementationOnce( + () => + new Promise((r) => { + resolve = r; + }), + ); + const first = initLogger({ + state, + projectName: "project", + setCurrent: false, + }).id; + await vi.waitFor(() => expect(register).toHaveBeenCalledTimes(1)); + if (action === "reset") { + state.resetLoginInfo(); + } else if (action === "forceLogin") { + await state.login({ forceLogin: true }); + } else { + state.setFetch((...args) => fetch(...args)); + } + await initLogger({ state, projectName: "project", setCurrent: false }).id; + resolve({ project: { id: "stale-id", name: "project" } }); + expect(await first).toBe("stale-id"); + expect( + await initLogger({ state, projectName: "project", setCurrent: false }) + .id, + ).toBe(projectId); + expect( + fetch.mock.calls.filter(([url]) => + String(url).endsWith("/api/project/register"), + ), + ).toHaveLength(1); + }, + ); + + test("isolates states and evicts the least recently used entry at 1,000 projects", async () => { + const { state } = createState(); + await state.login({}); + const register = vi.spyOn(state.appConn(), "post_json"); + for (let index = 0; index < 1000; index++) { + await initLogger({ + state, + projectName: `project-${index}`, + setCurrent: false, + }).id; + } + await initLogger({ state, projectName: "project-0", setCurrent: false }).id; + await initLogger({ state, projectName: "project-1000", setCurrent: false }) + .id; + await initLogger({ state, projectName: "project-0", setCurrent: false }).id; + expect(register).toHaveBeenCalledTimes(1001); + await initLogger({ state, projectName: "project-1", setCurrent: false }).id; + expect(register).toHaveBeenCalledTimes(1002); + const other = createState(); + await initLogger({ + state: other.state, + projectName: "project-0", + setCurrent: false, + }).id; + expect(other.fetch).toHaveBeenCalledTimes(2); + }); +}); + +describe("login deduplication", () => { + test("shares cold logins and keeps completed logins cached", async () => { + const { state, fetch } = createState(); + await Promise.all(Array.from({ length: 10 }, () => state.login({}))); + await state.login({}); + expect(fetch).toHaveBeenCalledTimes(1); + expect(state.loggedIn).toBe(true); + }); + + test("shares errors and allows a later retry", async () => { + const { state, fetch } = createState(); + fetch.mockRejectedValueOnce(new Error("login failed")); + const results = await Promise.allSettled( + Array.from({ length: 10 }, () => state.login({})), + ); + expect(results.every((result) => result.status === "rejected")).toBe(true); + expect(fetch).toHaveBeenCalledTimes(1); + await state.login({}); + expect(fetch).toHaveBeenCalledTimes(2); + expect(state.loggedIn).toBe(true); + }); + + test.each(["apiKey", "appUrl", "orgName", "fetch"] as const)( + "does not combine differing %s options", + async (field) => { + const { state, fetch } = createState(); + await Promise.all([ + state.login({}), + state.login( + field === "fetch" + ? { fetch: (...args) => fetch(...args) } + : { + [field]: + field === "orgName" + ? "org-id" + : field === "appUrl" + ? "https://other.test" + : "different", + }, + ), + ]); + expect(fetch).toHaveBeenCalledTimes(2); + }, + ); + + test("forceLogin always starts a fresh request", async () => { + const { state, fetch } = createState(); + await state.login({}); + const results = await Promise.allSettled([ + state.login({ forceLogin: true }), + state.login({ forceLogin: true }), + ]); + expect(results.map((result) => result.status)).toEqual([ + "rejected", + "fulfilled", + ]); + expect(fetch).toHaveBeenCalledTimes(3); + }); + + test("forced logins remain independent during asynchronous credential discovery", async () => { + const { fetch } = createState(); + const state = new BraintrustState({ + appUrl: "https://app.test", + fetch, + noExitFlush: true, + }); + vi.spyOn(iso, "getBraintrustApiKey").mockResolvedValue("test-credential"); + const results = await Promise.allSettled([ + state.login({ forceLogin: true }), + state.login({ forceLogin: true }), + ]); + expect(results.map((result) => result.status)).toEqual([ + "rejected", + "fulfilled", + ]); + expect(fetch).toHaveBeenCalledTimes(2); + expect(state.loggedIn).toBe(true); + }); + + test.each(["reset", "forceLogin"])( + "an older pending login cannot overwrite %s", + async (action) => { + const { state, fetch } = createState(); + let resolve!: (value: Response) => void; + fetch.mockImplementationOnce( + () => + new Promise((r) => { + resolve = r; + }), + ); + const first = state.login({}); + expect(fetch).toHaveBeenCalledTimes(1); + if (action === "reset") { + state.resetLoginInfo(); + } else { + fetch.mockResolvedValueOnce(loginResponse("new-org")); + await state.login({ forceLogin: true }); + } + resolve(loginResponse("old-org")); + await expect(first).rejects.toThrow( + "Login cancelled by a reset or a newer forced login.", + ); + expect(state.orgId).toBe(action === "reset" ? null : "new-org"); + }, + ); + + test.each([false, true])( + "cancels an earlier login before its forced replacement completes (forceLogin: %s)", + async (forceLogin) => { + const { state, fetch } = createState(); + let resolveFirst!: (value: Response) => void; + let resolveReplacement!: (value: Response) => void; + fetch + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFirst = resolve; + }), + ) + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveReplacement = resolve; + }), + ); + const first = state.login({ forceLogin }); + const replacement = state.login({ forceLogin: true }); + resolveFirst(loginResponse("old-org")); + await expect(first).rejects.toThrow( + "Login cancelled by a reset or a newer forced login.", + ); + expect(state.loggedIn).toBe(false); + resolveReplacement(loginResponse("new-org")); + await replacement; + expect(state.orgId).toBe("new-org"); + expect(state.loggedIn).toBe(true); + }, + ); + + test.each(["reset", "forceLogin"])( + "cancels credential discovery superseded by %s", + async (action) => { + const { fetch } = createState(); + const state = new BraintrustState({ + appUrl: "https://app.test", + fetch, + noExitFlush: true, + }); + let resolveCredential!: (value: string) => void; + vi.spyOn(iso, "getBraintrustApiKey").mockImplementationOnce( + () => + new Promise((resolve) => { + resolveCredential = resolve; + }), + ); + let resolveReplacement!: (value: Response) => void; + fetch.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveReplacement = resolve; + }), + ); + const first = state.login({}); + let replacement: Promise | undefined; + if (action === "reset") { + state.resetLoginInfo(); + } else { + replacement = state.login({ + apiKey: "replacement-credential", + forceLogin: true, + }); + } + resolveCredential("old-credential"); + await expect(first).rejects.toThrow( + "Login cancelled by a reset or a newer forced login.", + ); + expect(state.loggedIn).toBe(false); + expect(fetch).toHaveBeenCalledTimes(action === "reset" ? 0 : 1); + if (replacement) { + resolveReplacement(loginResponse("new-org")); + await replacement; + expect(state.orgId).toBe("new-org"); + expect(state.loggedIn).toBe(true); + } + }, + ); + + test("does not log in again if another login completed during credential discovery", async () => { + const { fetch } = createState(); + const state = new BraintrustState({ + appUrl: "https://app.test", + fetch, + noExitFlush: true, + }); + let resolve!: (value: string) => void; + vi.spyOn(iso, "getBraintrustApiKey").mockImplementationOnce( + () => + new Promise((r) => { + resolve = r; + }), + ); + const first = state.login({}); + await state.login({ apiKey: "test-credential" }); + resolve("test-credential"); + await first; + expect(fetch).toHaveBeenCalledTimes(1); + }); +}); diff --git a/js/src/logger.ts b/js/src/logger.ts index 5ae72b2fc..4e6bd9dab 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -688,6 +688,11 @@ let stateNonce = 0; const V1_PROXY_SUFFIX = "/v1/proxy"; const LOADER_LOGIN_CACHE_MAX = 16; +const PROJECT_CACHE_TTL_MS = 15 * 60 * 1000; +const projectMetadataCaches = new WeakMap< + BraintrustState, + LRUCache; expiresAt: number }> +>(); type ResolvedLoaderLoginOptions = Required< Pick @@ -762,6 +767,11 @@ export class BraintrustState { private readonly loginParams: LoginOptions; private activeLoginOrgNameSelector: string | undefined; + private loginGeneration = 0; + private pendingLogins = new WeakMap< + typeof globalThis.fetch, + Map> + >(); constructor(loginParams: LoginOptions) { this.loginParams = { ...loginParams }; @@ -855,6 +865,9 @@ export class BraintrustState { } public resetLoginInfo() { + this.loginGeneration++; + this.pendingLogins = new WeakMap(); + projectMetadataCaches.delete(this); this.appUrl = null; this.appPublicUrl = null; this.loginToken = null; @@ -1052,6 +1065,7 @@ export class BraintrustState { } public copyLoginInfo(other: BraintrustState) { + projectMetadataCaches.delete(this); this.appUrl = other.appUrl; this.appPublicUrl = other.appPublicUrl; this.loginToken = other.loginToken; @@ -1157,6 +1171,9 @@ export class BraintrustState { } public setFetch(fetch: typeof globalThis.fetch) { + if (fetch !== this.fetch) { + projectMetadataCaches.delete(this); + } this.loginParams.fetch = fetch; this.fetch = fetch; this._apiConn?.setFetch(fetch); @@ -1192,13 +1209,69 @@ export class BraintrustState { if (this.apiUrl && !loginParams.forceLogin) { return; } - const newState = await loginToState({ + if (loginParams.forceLogin) { + // A forced login supersedes older attempts, even if they finish later. + this.loginGeneration++; + this.pendingLogins = new WeakMap(); + } + const generation = this.loginGeneration; + const pendingLogins = this.pendingLogins; + const options = { ...this.loginParams, ...Object.fromEntries( Object.entries(loginParams).filter(([k, v]) => !isEmpty(v)), ), + }; + const appUrl = + options.appUrl ?? + (iso.getEnv("BRAINTRUST_APP_URL") || "https://www.braintrust.dev"); + const orgName = options.orgName ?? iso.getEnv("BRAINTRUST_ORG_NAME"); + const fetch = options.fetch ?? globalThis.fetch; + const apiKey = options.apiKey ?? (await iso.getBraintrustApiKey()); + if (generation !== this.loginGeneration && !loginParams.forceLogin) { + throw new Error("Login cancelled by a reset or a newer forced login."); + } + // Another attempt may have finished while the credential was being read. + if (this.apiUrl && !loginParams.forceLogin) { + return; + } + + let pending = pendingLogins.get(fetch); + if (!pending) { + pending = new Map(); + pendingLogins.set(fetch, pending); + } + const key = JSON.stringify([ + appUrl, + orgName, + apiKey, + options.debugLogLevel, + ]); + const existing = pending.get(key); + if (!loginParams.forceLogin && existing) { + return existing; + } + + const loginPromise = loginToState({ + ...options, + appUrl, + orgName, + apiKey, + fetch, + }).then((newState) => { + if (generation !== this.loginGeneration) { + throw new Error("Login cancelled by a reset or a newer forced login."); + } + this.copyLoginInfo(newState); }); - this.copyLoginInfo(newState); + pending.set(key, loginPromise); + try { + await loginPromise; + } finally { + if (pending.get(key) === loginPromise) { + pending.delete(key); + } + } } public appConn(): HTTPConnection { @@ -4937,37 +5010,72 @@ async function computeLoggerMetadata( ) { await state.login({}); const org_id = state.orgId!; - if (isEmpty(project_id)) { - const response = await state.appConn().post_json("api/project/register", { - project_name: project_name || GLOBAL_PROJECT, - org_id, - }); - return { - org_id, - project: { - id: response.project.id, - name: response.project.name, - fullInfo: response.project, - }, - }; - } else if (isEmpty(project_name)) { - const response = await state.appConn().get_json("api/project", { - id: project_id, - }); - return { - org_id, - project: { - id: project_id, - name: response.name, - fullInfo: response.project, - }, - }; - } else { + if (!isEmpty(project_id) && !isEmpty(project_name)) { return { org_id, project: { id: project_id, name: project_name, fullInfo: {} }, }; } + + let cache = projectMetadataCaches.get(state); + if (!cache) { + cache = new LRUCache({ max: 1000 }); + projectMetadataCaches.set(state, cache); + } + const key = JSON.stringify([ + state.appUrl, + org_id, + state.loginToken, + isEmpty(project_id) + ? ["name", project_name || GLOBAL_PROJECT] + : ["id", project_id], + ]); + const cached = cache.get(key); + if (cached && Date.now() < cached.expiresAt) { + return cached.promise; + } + + const conn = state.appConn(); + const entry = { + // Pending requests do not expire. The TTL starts when the lookup succeeds. + expiresAt: Infinity, + promise: (async (): Promise => { + if (isEmpty(project_id)) { + const response = await conn.post_json("api/project/register", { + project_name: project_name || GLOBAL_PROJECT, + org_id, + }); + return { + org_id, + project: { + id: response.project.id, + name: response.project.name, + fullInfo: response.project, + }, + }; + } + const response = await conn.get_json("api/project", { id: project_id }); + return { + org_id, + project: { + id: project_id, + name: response.name, + fullInfo: response.project, + }, + }; + })(), + }; + cache.set(key, entry); + try { + const metadata = await entry.promise; + entry.expiresAt = Date.now() + PROJECT_CACHE_TTL_MS; + return metadata; + } catch (error) { + if (cache.get(key) === entry) { + cache.delete(key); + } + throw error; + } } type AsyncFlushArg = { From e0195af0a3f7bc4ee286349137e74ddccc3dfb1d Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 28 Sep 2026 09:32:42 +0000 Subject: [PATCH 2/4] Update PR #2523 --- .changeset/tricky-badgers-pay.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/tricky-badgers-pay.md diff --git a/.changeset/tricky-badgers-pay.md b/.changeset/tricky-badgers-pay.md new file mode 100644 index 000000000..ddceec107 --- /dev/null +++ b/.changeset/tricky-badgers-pay.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +perf: Cache project lookup From 62eb16142f7f7487bcec2d26da1652496357bec1 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 28 Sep 2026 11:27:17 +0000 Subject: [PATCH 3/4] Update PR #2523 --- js/src/logger-caching.test.ts | 330 +++++++++++++++++++++++++--------- js/src/logger.ts | 101 +++++++---- 2 files changed, 314 insertions(+), 117 deletions(-) diff --git a/js/src/logger-caching.test.ts b/js/src/logger-caching.test.ts index d9a486d19..a20b29c78 100644 --- a/js/src/logger-caching.test.ts +++ b/js/src/logger-caching.test.ts @@ -57,6 +57,56 @@ afterEach(() => { }); describe("project metadata caching", () => { + test.each(["name", "id"])( + "isolates metadata mutations across concurrent and later %s lookups", + async (lookup) => { + const { state, fetch } = createState(); + await state.login({}); + const project = { + id: projectId, + name: "project", + metadata: { labels: ["original"] }, + }; + fetch.mockResolvedValueOnce(Response.json({ name: "project", project })); + const options = { + state, + setCurrent: false, + ...(lookup === "name" ? { projectName: "project" } : { projectId }), + }; + const first = initLogger(options); + const second = initLogger(options); + const [firstProject, secondProject] = await Promise.all([ + first.project, + second.project, + ]); + firstProject.id = "changed-id"; + firstProject.name = "changed-name"; + (firstProject.fullInfo.metadata as typeof project.metadata).labels.push( + "changed", + ); + expect(secondProject).toEqual({ + id: projectId, + name: "project", + fullInfo: project, + }); + expect(await second.id).toBe(projectId); + secondProject.id = "also-changed"; + (secondProject.fullInfo.metadata as typeof project.metadata).labels.push( + "also-changed", + ); + expect(await initLogger(options).project).toEqual({ + id: projectId, + name: "project", + fullInfo: project, + }); + const components = SpanComponentsV4.fromStr(await first.export()); + expect(await spanComponentsToObjectId({ state, components })).toBe( + projectId, + ); + expect(fetch).toHaveBeenCalledTimes(2); + }, + ); + test("shares concurrent and sequential lookups across loggers and imported parents", async () => { const { state, fetch } = createState(); const logger = initLogger({ @@ -302,6 +352,72 @@ describe("project metadata caching", () => { }); describe("login deduplication", () => { + test.each(["headers", "body"])( + "times out stalled %s, unblocks forced login, and ignores late results", + async (stage) => { + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); + const { state, fetch } = createState(); + let signal: AbortSignal | null | undefined; + let completeFirst!: () => void; + fetch.mockImplementationOnce((_, init) => { + signal = init?.signal; + if (stage === "headers") { + return new Promise((resolve) => { + completeFirst = () => resolve(loginResponse("old-org")); + }); + } + const response = loginResponse("old-org"); + vi.spyOn(response, "text").mockImplementation( + () => + new Promise((resolve) => { + completeFirst = () => + resolve( + JSON.stringify({ + org_info: [ + { + id: "old-org", + name: "old-org", + api_url: "https://api.test", + }, + ], + }), + ); + }), + ); + return Promise.resolve(response); + }); + const first = state.login({}); + const duplicate = state.login({}); + const firstResults = Promise.allSettled([first, duplicate]); + const replacementFetch = vi.fn(async () => + loginResponse("new-org"), + ); + const replacement = state.login({ + forceLogin: true, + fetch: replacementFetch, + }); + await vi.advanceTimersByTimeAsync(29_999); + expect(replacementFetch).not.toHaveBeenCalled(); + expect(signal?.aborted).toBe(false); + await vi.advanceTimersByTimeAsync(1); + for (const result of await firstResults) { + expect(result).toMatchObject({ + status: "rejected", + reason: new Error("Braintrust login timed out after 30 seconds."), + }); + } + await replacement; + expect(signal?.aborted).toBe(true); + expect(replacementFetch).toHaveBeenCalledTimes(1); + expect(state.orgId).toBe("new-org"); + completeFirst(); + await vi.advanceTimersByTimeAsync(0); + expect(state.orgId).toBe("new-org"); + expect(state.fetch).toBe(replacementFetch); + expect(vi.getTimerCount()).toBe(0); + }, + ); + test("shares cold logins and keeps completed logins cached", async () => { const { state, fetch } = createState(); await Promise.all(Array.from({ length: 10 }, () => state.login({}))); @@ -346,69 +462,51 @@ describe("login deduplication", () => { }, ); - test("forceLogin always starts a fresh request", async () => { - const { state, fetch } = createState(); - await state.login({}); - const results = await Promise.allSettled([ - state.login({ forceLogin: true }), - state.login({ forceLogin: true }), - ]); - expect(results.map((result) => result.status)).toEqual([ - "rejected", - "fulfilled", - ]); - expect(fetch).toHaveBeenCalledTimes(3); - }); - - test("forced logins remain independent during asynchronous credential discovery", async () => { - const { fetch } = createState(); - const state = new BraintrustState({ - appUrl: "https://app.test", - fetch, - noExitFlush: true, - }); - vi.spyOn(iso, "getBraintrustApiKey").mockResolvedValue("test-credential"); - const results = await Promise.allSettled([ - state.login({ forceLogin: true }), - state.login({ forceLogin: true }), - ]); - expect(results.map((result) => result.status)).toEqual([ - "rejected", - "fulfilled", - ]); - expect(fetch).toHaveBeenCalledTimes(2); - expect(state.loggedIn).toBe(true); - }); - - test.each(["reset", "forceLogin"])( - "an older pending login cannot overwrite %s", - async (action) => { - const { state, fetch } = createState(); - let resolve!: (value: Response) => void; - fetch.mockImplementationOnce( - () => - new Promise((r) => { - resolve = r; - }), + test.each([false, true])( + "concurrent forced loggers remain usable (credential discovery: %s)", + async (discoverCredential) => { + const { fetch } = createState(); + const state = new BraintrustState({ + apiKey: discoverCredential ? undefined : "test-credential", + appUrl: "https://app.test", + fetch, + noExitFlush: true, + }); + vi.spyOn(iso, "getBraintrustApiKey").mockResolvedValue("test-credential"); + state.httpLogger().syncFlush = true; + const loggers = Array.from({ length: 2 }, () => + initLogger({ + state, + projectName: "project", + forceLogin: true, + setCurrent: false, + }), ); - const first = state.login({}); - expect(fetch).toHaveBeenCalledTimes(1); - if (action === "reset") { - state.resetLoginInfo(); - } else { - fetch.mockResolvedValueOnce(loginResponse("new-org")); - await state.login({ forceLogin: true }); + expect(await Promise.all(loggers.map((logger) => logger.id))).toEqual([ + projectId, + projectId, + ]); + for (const [index, logger] of loggers.entries()) { + logger.startSpan({ name: `forced-${index}` }).end(); } - resolve(loginResponse("old-org")); - await expect(first).rejects.toThrow( - "Login cancelled by a reset or a newer forced login.", - ); - expect(state.orgId).toBe(action === "reset" ? null : "new-org"); + await state.bgLogger().flush(); + const rows = fetch.mock.calls + .filter(([url]) => String(url).endsWith("/logs3")) + .flatMap(([, init]) => JSON.parse(String(init?.body)).rows); + expect(rows.map((row) => row.span_attributes.name).sort()).toEqual([ + "forced-0", + "forced-1", + ]); + expect( + fetch.mock.calls.filter(([url]) => + String(url).endsWith("/api/apikey/login"), + ), + ).toHaveLength(2); }, ); test.each([false, true])( - "cancels an earlier login before its forced replacement completes (forceLogin: %s)", + "queues a forced login behind the earlier attempt (forceLogin: %s)", async (forceLogin) => { const { state, fetch } = createState(); let resolveFirst!: (value: Response) => void; @@ -428,21 +526,61 @@ describe("login deduplication", () => { ); const first = state.login({ forceLogin }); const replacement = state.login({ forceLogin: true }); + expect(fetch).toHaveBeenCalledTimes(1); resolveFirst(loginResponse("old-org")); - await expect(first).rejects.toThrow( - "Login cancelled by a reset or a newer forced login.", - ); - expect(state.loggedIn).toBe(false); + await first; + expect(state.orgId).toBe("old-org"); + await vi.waitFor(() => expect(fetch).toHaveBeenCalledTimes(2)); resolveReplacement(loginResponse("new-org")); await replacement; expect(state.orgId).toBe("new-org"); - expect(state.loggedIn).toBe(true); }, ); - test.each(["reset", "forceLogin"])( - "cancels credential discovery superseded by %s", - async (action) => { + test("a failed login does not prevent the queued forced login from succeeding", async () => { + const { state, fetch } = createState(); + fetch.mockRejectedValueOnce(new Error("login failed")); + const results = await Promise.allSettled([ + state.login({ forceLogin: true }), + state.login({ forceLogin: true }), + ]); + expect(results.map((result) => result.status)).toEqual([ + "rejected", + "fulfilled", + ]); + expect(fetch).toHaveBeenCalledTimes(2); + expect(state.loggedIn).toBe(true); + }); + + test("reset cancels pending and queued logins without blocking a new login", async () => { + const { state, fetch } = createState(); + let resolve!: (value: Response) => void; + fetch.mockImplementationOnce( + () => + new Promise((r) => { + resolve = r; + }), + ); + const first = state.login({}); + const queued = state.login({ forceLogin: true }); + const results = Promise.allSettled([first, queued]); + state.resetLoginInfo(); + fetch.mockResolvedValueOnce(loginResponse("new-org")); + await state.login({}); + resolve(loginResponse("old-org")); + for (const result of await results) { + expect(result).toMatchObject({ + status: "rejected", + reason: new Error("Login cancelled by a reset."), + }); + } + expect(fetch).toHaveBeenCalledTimes(2); + expect(state.orgId).toBe("new-org"); + }); + + test.each([false, true])( + "reset cancels credential discovery (forceLogin: %s)", + async (forceLogin) => { const { fetch } = createState(); const state = new BraintrustState({ appUrl: "https://app.test", @@ -456,35 +594,53 @@ describe("login deduplication", () => { resolveCredential = resolve; }), ); - let resolveReplacement!: (value: Response) => void; - fetch.mockImplementationOnce( + const first = state.login({ forceLogin }); + state.resetLoginInfo(); + resolveCredential("old-credential"); + await expect(first).rejects.toThrow("Login cancelled by a reset."); + expect(state.loggedIn).toBe(false); + expect(fetch).not.toHaveBeenCalled(); + }, + ); + + test.each(["apiKey", "appUrl", "orgName", "fetch"] as const)( + "rejects conflicting %s after credential discovery", + async (field) => { + const { fetch } = createState(); + const state = new BraintrustState({ + appUrl: "https://app.test", + fetch, + noExitFlush: true, + }); + let resolveCredential!: (value: string) => void; + vi.spyOn(iso, "getBraintrustApiKey").mockImplementationOnce( () => new Promise((resolve) => { - resolveReplacement = resolve; + resolveCredential = resolve; }), ); const first = state.login({}); - let replacement: Promise | undefined; - if (action === "reset") { - state.resetLoginInfo(); - } else { - replacement = state.login({ - apiKey: "replacement-credential", - forceLogin: true, - }); - } - resolveCredential("old-credential"); + await state.login({ + apiKey: "test-credential", + ...(field === "fetch" + ? { fetch: (...args: Parameters) => fetch(...args) } + : { + [field]: + field === "orgName" + ? "org-id" + : field === "appUrl" + ? "https://other.test" + : "different-credential", + }), + }); + resolveCredential("test-credential"); await expect(first).rejects.toThrow( - "Login cancelled by a reset or a newer forced login.", + "Another login completed with different options", + ); + expect(fetch).toHaveBeenCalledTimes(1); + expect(state.loginToken).toBe( + field === "apiKey" ? "different-credential" : "test-credential", ); - expect(state.loggedIn).toBe(false); - expect(fetch).toHaveBeenCalledTimes(action === "reset" ? 0 : 1); - if (replacement) { - resolveReplacement(loginResponse("new-org")); - await replacement; - expect(state.orgId).toBe("new-org"); - expect(state.loggedIn).toBe(true); - } }, ); diff --git a/js/src/logger.ts b/js/src/logger.ts index 4e6bd9dab..07022e896 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -688,6 +688,7 @@ let stateNonce = 0; const V1_PROXY_SUFFIX = "/v1/proxy"; const LOADER_LOGIN_CACHE_MAX = 16; +const LOGIN_TIMEOUT_MS = 30_000; const PROJECT_CACHE_TTL_MS = 15 * 60 * 1000; const projectMetadataCaches = new WeakMap< BraintrustState, @@ -768,6 +769,7 @@ export class BraintrustState { private readonly loginParams: LoginOptions; private activeLoginOrgNameSelector: string | undefined; private loginGeneration = 0; + private loginQueueTail: Promise | undefined; private pendingLogins = new WeakMap< typeof globalThis.fetch, Map> @@ -866,6 +868,7 @@ export class BraintrustState { public resetLoginInfo() { this.loginGeneration++; + this.loginQueueTail = undefined; this.pendingLogins = new WeakMap(); projectMetadataCaches.delete(this); this.appUrl = null; @@ -1209,11 +1212,6 @@ export class BraintrustState { if (this.apiUrl && !loginParams.forceLogin) { return; } - if (loginParams.forceLogin) { - // A forced login supersedes older attempts, even if they finish later. - this.loginGeneration++; - this.pendingLogins = new WeakMap(); - } const generation = this.loginGeneration; const pendingLogins = this.pendingLogins; const options = { @@ -1228,11 +1226,21 @@ export class BraintrustState { const orgName = options.orgName ?? iso.getEnv("BRAINTRUST_ORG_NAME"); const fetch = options.fetch ?? globalThis.fetch; const apiKey = options.apiKey ?? (await iso.getBraintrustApiKey()); - if (generation !== this.loginGeneration && !loginParams.forceLogin) { - throw new Error("Login cancelled by a reset or a newer forced login."); + if (generation !== this.loginGeneration) { + throw new Error("Login cancelled by a reset."); } // Another attempt may have finished while the credential was being read. if (this.apiUrl && !loginParams.forceLogin) { + if ( + this.appUrl !== appUrl || + this.loginToken !== (apiKey && HTTPConnection.sanitize_token(apiKey)) || + this.activeLoginOrgNameSelector !== orgName || + this.fetch !== fetch + ) { + throw new Error( + "Another login completed with different options during credential discovery. To force re-login, pass `forceLogin: true`.", + ); + } return; } @@ -1252,22 +1260,36 @@ export class BraintrustState { return existing; } - const loginPromise = loginToState({ - ...options, - appUrl, - orgName, - apiKey, - fetch, - }).then((newState) => { + const previousLogin = this.loginQueueTail; + const loginPromise = (async () => { + // Each forced login still makes a fresh request, but must not invalidate + // the metadata promises of loggers already waiting for another login. + if (previousLogin) { + await previousLogin.catch(() => {}); + } if (generation !== this.loginGeneration) { - throw new Error("Login cancelled by a reset or a newer forced login."); + throw new Error("Login cancelled by a reset."); + } + const newState = await loginToState({ + ...options, + appUrl, + orgName, + apiKey, + fetch, + }); + if (generation !== this.loginGeneration) { + throw new Error("Login cancelled by a reset."); } this.copyLoginInfo(newState); - }); + })(); + this.loginQueueTail = loginPromise; pending.set(key, loginPromise); try { await loginPromise; } finally { + if (this.loginQueueTail === loginPromise) { + this.loginQueueTail = undefined; + } if (pending.get(key) === loginPromise) { pending.delete(key); } @@ -5032,7 +5054,7 @@ async function computeLoggerMetadata( ]); const cached = cache.get(key); if (cached && Date.now() < cached.expiresAt) { - return cached.promise; + return structuredClone(await cached.promise); } const conn = state.appConn(); @@ -5069,7 +5091,7 @@ async function computeLoggerMetadata( try { const metadata = await entry.promise; entry.expiresAt = Date.now() + PROJECT_CACHE_TTL_MS; - return metadata; + return structuredClone(metadata); } catch (error) { if (cache.get(key) === entry) { cache.delete(key); @@ -5835,18 +5857,37 @@ export async function loginToState(options: LoginOptions = {}) { _saveOrgInfo(state, testOrgInfo, testOrgInfo[0].name); return state; } else { - const loginResponse = await fetch( - _urljoin(state.appUrl, `/api/apikey/login`), - { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${apiKey}`, - }, - }, - ); - const resp = await checkResponse(loginResponse); - const info = await readJSONResponse(resp); + const controller = new AbortController(); + let timeout: ReturnType | undefined; + let info; + try { + // Bound both headers and body reads, including custom fetches that ignore + // cancellation, so a stalled request cannot block subsequent logins. + info = await Promise.race([ + (async () => { + const response = await fetch(_urljoin(appUrl, `/api/apikey/login`), { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${apiKey}`, + }, + signal: controller.signal, + }); + return readJSONResponse(await checkResponse(response)); + })(), + new Promise((_, reject) => { + timeout = setTimeout(() => { + const error = new Error( + "Braintrust login timed out after 30 seconds.", + ); + reject(error); + controller.abort(error); + }, LOGIN_TIMEOUT_MS); + }), + ]); + } finally { + clearTimeout(timeout); + } _saveOrgInfo(state, info.org_info, orgName); if (!state.apiUrl) { From 5aa117ecd8878ba189a768ca2b90e13814fb6f5d Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 28 Sep 2026 13:40:11 +0000 Subject: [PATCH 4/4] Update PR #2523 --- js/src/logger-caching.test.ts | 125 ++++++++++++++++++++++++++++++++-- js/src/logger.ts | 98 ++++++++++++++++---------- 2 files changed, 182 insertions(+), 41 deletions(-) diff --git a/js/src/logger-caching.test.ts b/js/src/logger-caching.test.ts index a20b29c78..d0ac482fc 100644 --- a/js/src/logger-caching.test.ts +++ b/js/src/logger-caching.test.ts @@ -57,6 +57,49 @@ afterEach(() => { }); describe("project metadata caching", () => { + test.each(["reset", "forceLogin", "fetch"])( + "shares lookups across SDK copies and invalidates on %s", + async (action) => { + vi.resetModules(); + const otherSdk = await import("./logger"); + const otherNode = await import("./node/config"); + otherNode.configureNode(); + + const { state, fetch } = createState(); + const options = { state, projectName: "project", setCurrent: false }; + expect( + await Promise.all([ + initLogger(options).id, + otherSdk.initLogger(options).id, + ]), + ).toEqual([projectId, projectId]); + expect( + fetch.mock.calls.filter(([url]) => + String(url).endsWith("/api/project/register"), + ), + ).toHaveLength(1); + + if (action === "reset") { + state.resetLoginInfo(); + await state.login({}); + } else if (action === "forceLogin") { + await state.login({ forceLogin: true }); + } else { + state.setFetch((...args) => fetch(...args)); + } + fetch.mockResolvedValueOnce( + Response.json({ project: { id: "replacement-id", name: "project" } }), + ); + expect(await otherSdk.initLogger(options).id).toBe("replacement-id"); + expect(await initLogger(options).id).toBe("replacement-id"); + expect( + fetch.mock.calls.filter(([url]) => + String(url).endsWith("/api/project/register"), + ), + ).toHaveLength(2); + }, + ); + test.each(["name", "id"])( "isolates metadata mutations across concurrent and later %s lookups", async (lookup) => { @@ -191,8 +234,8 @@ describe("project metadata caching", () => { expect(register).toHaveBeenCalledTimes(2); }); - test("starts TTL on success and shares requests even when resolution takes over 15 minutes", async () => { - vi.useFakeTimers({ toFake: ["Date"] }); + test("starts TTL on success and shares pending requests", async () => { + vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] }); vi.setSystemTime(0); const { state } = createState(); await state.login({}); @@ -208,8 +251,8 @@ describe("project metadata caching", () => { projectName: "project", setCurrent: false, }).id; - await vi.waitFor(() => expect(register).toHaveBeenCalledTimes(1)); - vi.setSystemTime(ttl * 2); + await vi.advanceTimersByTimeAsync(20_000); + expect(register).toHaveBeenCalledTimes(1); const second = initLogger({ state, projectName: "project", @@ -217,11 +260,83 @@ describe("project metadata caching", () => { }).id; resolve({ project: { id: projectId, name: "project" } }); await Promise.all([first, second]); - vi.setSystemTime(ttl * 3 - 1); + expect(vi.getTimerCount()).toBe(0); + vi.setSystemTime(20_000 + ttl - 1); await initLogger({ state, projectName: "project", setCurrent: false }).id; expect(register).toHaveBeenCalledTimes(1); }); + test.each([ + ["name", "headers"], + ["name", "body"], + ["id", "headers"], + ["id", "body"], + ])( + "times out stalled %s lookup %s and allows retry", + async (lookup, stage) => { + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); + const { state, fetch } = createState(); + await state.login({}); + let signal: AbortSignal | null | undefined; + let completeFirst!: () => void; + fetch.mockImplementationOnce((_, init) => { + signal = init?.signal; + const body = { + name: "project", + project: { id: "stale-id", name: "project" }, + }; + if (stage === "headers") { + return new Promise((resolve) => { + completeFirst = () => resolve(Response.json(body)); + }); + } + const response = Response.json(body); + vi.spyOn(response, "text").mockImplementation( + () => + new Promise((resolve) => { + completeFirst = () => resolve(JSON.stringify(body)); + }), + ); + return Promise.resolve(response); + }); + const options = { + state, + setCurrent: false, + ...(lookup === "name" ? { projectName: "project" } : { projectId }), + }; + const results = Promise.allSettled([ + initLogger(options).project, + initLogger(options).project, + ]); + await vi.advanceTimersByTimeAsync(29_999); + expect(fetch).toHaveBeenCalledTimes(2); + expect(signal?.aborted).toBe(false); + await vi.advanceTimersByTimeAsync(1); + for (const result of await results) { + expect(result).toMatchObject({ + status: "rejected", + reason: new Error( + "Braintrust project lookup timed out after 30 seconds.", + ), + }); + } + expect(signal?.aborted).toBe(true); + expect(vi.getTimerCount()).toBe(0); + + const replacement = await initLogger(options).project; + expect(replacement.id).toBe(projectId); + expect(fetch).toHaveBeenCalledTimes(3); + expect(vi.getTimerCount()).toBe(0); + + // A custom fetch can ignore cancellation and finish after the replacement. + completeFirst(); + await vi.advanceTimersByTimeAsync(0); + expect(await initLogger(options).project).toEqual(replacement); + expect(fetch).toHaveBeenCalledTimes(3); + expect(vi.getTimerCount()).toBe(0); + }, + ); + test("shares failures but permits retries from later loggers", async () => { const { state } = createState(); await state.login({}); diff --git a/js/src/logger.ts b/js/src/logger.ts index 07022e896..26fc15634 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -689,11 +689,8 @@ let stateNonce = 0; const V1_PROXY_SUFFIX = "/v1/proxy"; const LOADER_LOGIN_CACHE_MAX = 16; const LOGIN_TIMEOUT_MS = 30_000; +const PROJECT_LOOKUP_TIMEOUT_MS = 30_000; const PROJECT_CACHE_TTL_MS = 15 * 60 * 1000; -const projectMetadataCaches = new WeakMap< - BraintrustState, - LRUCache; expiresAt: number }> ->(); type ResolvedLoaderLoginOptions = Required< Pick @@ -756,6 +753,11 @@ export class BraintrustState { public promptCache: PromptCache; public parametersCache: ParametersCache; public spanCache: SpanCache; + /** @internal */ + public readonly projectMetadataCache = new LRUCache< + string, + { promise: Promise; expiresAt: number } + >({ max: 1000 }); private _idGenerator: IDGenerator | null = null; private _contextManager: ContextManager | null = null; private _otelFlushCallback: (() => Promise) | null = null; @@ -870,7 +872,7 @@ export class BraintrustState { this.loginGeneration++; this.loginQueueTail = undefined; this.pendingLogins = new WeakMap(); - projectMetadataCaches.delete(this); + this.projectMetadataCache.clear(); this.appUrl = null; this.appPublicUrl = null; this.loginToken = null; @@ -1068,7 +1070,7 @@ export class BraintrustState { } public copyLoginInfo(other: BraintrustState) { - projectMetadataCaches.delete(this); + this.projectMetadataCache.clear(); this.appUrl = other.appUrl; this.appPublicUrl = other.appPublicUrl; this.loginToken = other.loginToken; @@ -1175,7 +1177,7 @@ export class BraintrustState { public setFetch(fetch: typeof globalThis.fetch) { if (fetch !== this.fetch) { - projectMetadataCaches.delete(this); + this.projectMetadataCache.clear(); } this.loginParams.fetch = fetch; this.fetch = fetch; @@ -1753,13 +1755,17 @@ class HTTPConnection { object_type: string, args: Record | undefined = undefined, retries: number = 0, + signal?: AbortSignal, ) { const tries = retries + 1; for (let i = 0; i < tries; i++) { try { - const resp = await this.get(`${object_type}`, args); + const resp = await this.get(`${object_type}`, args, { signal }); return await readJSONResponse(resp, this.classifyTransportErrors); } catch (e) { + if (signal?.aborted) { + throw getAbortReason(signal); + } if (i < tries - 1) { debugLogger.debug( `Retrying API request ${object_type} ${JSON.stringify(args)} ${ @@ -1770,9 +1776,7 @@ class HTTPConnection { debugLogger.info( `Sleeping for ${sleepTimeS}s before retrying API request`, ); - await new Promise((resolve) => - setTimeout(resolve, sleepTimeS * 1000), - ); + await waitForRetry(sleepTimeS * 1000, signal); continue; } throw e; @@ -1783,9 +1787,11 @@ class HTTPConnection { async post_json( object_type: string, args: Record | string | undefined = undefined, + signal?: AbortSignal, ) { const resp = await this.post(`${object_type}`, args, { headers: { "Content-Type": "application/json" }, + signal, }); return await readJSONResponse(resp, this.classifyTransportErrors); } @@ -5039,11 +5045,7 @@ async function computeLoggerMetadata( }; } - let cache = projectMetadataCaches.get(state); - if (!cache) { - cache = new LRUCache({ max: 1000 }); - projectMetadataCaches.set(state, cache); - } + const cache = state.projectMetadataCache; const key = JSON.stringify([ state.appUrl, org_id, @@ -5058,34 +5060,56 @@ async function computeLoggerMetadata( } const conn = state.appConn(); + const controller = new AbortController(); + let timeout: ReturnType | undefined; const entry = { - // Pending requests do not expire. The TTL starts when the lookup succeeds. + // The lookup timeout bounds pending requests; the cache TTL starts on success. expiresAt: Infinity, - promise: (async (): Promise => { - if (isEmpty(project_id)) { - const response = await conn.post_json("api/project/register", { - project_name: project_name || GLOBAL_PROJECT, - org_id, - }); + promise: Promise.race([ + (async (): Promise => { + if (isEmpty(project_id)) { + const response = await conn.post_json( + "api/project/register", + { + project_name: project_name || GLOBAL_PROJECT, + org_id, + }, + controller.signal, + ); + return { + org_id, + project: { + id: response.project.id, + name: response.project.name, + fullInfo: response.project, + }, + }; + } + const response = await conn.get_json( + "api/project", + { id: project_id }, + 0, + controller.signal, + ); return { org_id, project: { - id: response.project.id, - name: response.project.name, + id: project_id, + name: response.name, fullInfo: response.project, }, }; - } - const response = await conn.get_json("api/project", { id: project_id }); - return { - org_id, - project: { - id: project_id, - name: response.name, - fullInfo: response.project, - }, - }; - })(), + })(), + new Promise((_, reject) => { + timeout = setTimeout(() => { + const error = new Error( + "Braintrust project lookup timed out after 30 seconds.", + ); + reject(error); + controller.abort(error); + }, PROJECT_LOOKUP_TIMEOUT_MS); + }), + ]), }; cache.set(key, entry); try { @@ -5097,6 +5121,8 @@ async function computeLoggerMetadata( cache.delete(key); } throw error; + } finally { + clearTimeout(timeout); } }