Merge branch 'main' into caitlin/update-braintrust-otel

Author: cpinn

apis/cloudflare/src/proxy.ts

@@ -115,6 +115,8 @@ export async function handleProxyV1(
 
   let spanLogger: SpanLogger | undefined;
   let span: Span | undefined;
+  let spanId: string | undefined;
+  let spanExport: string | undefined;
   const parentHeader = request.headers.get(BT_PARENT);
   if (parentHeader) {
     let parent;
@@ -147,6 +149,14 @@ export async function handleProxyV1(
       parent: parent.toStr(),
     });
     spanLogger = makeProxySpanLogger(span, ctx.waitUntil.bind(ctx));
+
+    // Extract span ID and export for response headers
+    spanId = span.id;
+    if (spanId.trim() !== "") {
+      spanExport = await span.export();
+    } else {
+      spanId = undefined;
+    }
   }
 
   const opts: ProxyOpts = {
@@ -187,6 +197,8 @@ export async function handleProxyV1(
     logHistogram, // Pass logging functions to edge layer
     whitelist,
     spanLogger,
+    spanId,
+    spanExport,
   };
 
   const url = new URL(request.url);

packages/proxy/edge/index.ts

@@ -3,7 +3,6 @@ import { flushMetrics } from "@lib/metrics";
 import { proxyV1, SpanLogger, LogHistogramFn } from "@lib/proxy";
 import { isEmpty } from "@lib/util";
 import { MeterProvider } from "@opentelemetry/sdk-metrics";
-import { Meter } from "@opentelemetry/api";
 
 import { APISecret, getModelEndpointTypes } from "@schema";
 import { verifyTempCredentials, isTempCredential } from "utils";
@@ -37,6 +36,8 @@ export interface ProxyOpts {
   logHistogram?: LogHistogramFn;
   whitelist?: (string | RegExp)[];
   spanLogger?: SpanLogger;
+  spanId?: string;
+  spanExport?: string;
 }
 
 const defaultWhitelist: (string | RegExp)[] = [
@@ -177,7 +178,7 @@ export function makeFetchApiSecrets({
     let lookupFailed = false;
     // Only cache API keys for 60 seconds. This reduces the load on the database but ensures
     // that changes roll out quickly enough too.
-    let ttl = 60;
+    const ttl = 60;
     try {
       const response = await fetch(
         `${opts.braintrustApiUrl || DEFAULT_BRAINTRUST_APP_URL}/api/secret`,
@@ -258,11 +259,11 @@ export function EdgeProxyV1(opts: ProxyOpts) {
 
     // Create an identity TransformStream (a.k.a. a pipe).
     // The readable side will become our new response body.
-    let { readable, writable } = new TransformStream();
+    const { readable, writable } = new TransformStream();
 
     let status = 200;
 
-    let headers: Record<string, string> = opts.cors ? corsHeaders : {};
+    const headers: Record<string, string> = opts.cors ? corsHeaders : {};
 
     const setStatus = (code: number) => {
       status = code;
@@ -289,6 +290,14 @@ export function EdgeProxyV1(opts: ProxyOpts) {
 
     const fetchApiSecrets = makeFetchApiSecrets({ ctx, opts });
 
+    // Set span headers if available
+    if (opts.spanId) {
+      setHeader("x-bt-span-id", opts.spanId);
+    }
+    if (opts.spanExport) {
+      setHeader("x-bt-span-export", opts.spanExport);
+    }
+
     const cachePut = async (
       encryptionKey: string,
       key: string,