fix: allow passing promise of TelemetryConfig to enableTelemetry (#1169)

firebase · Nov 5, 2024 · 73c51dc · 73c51dc
1 parent 6d4ef64
commit 73c51dc
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 17 deletions.
diff --git a/js/core/src/tracing.ts b/js/core/src/tracing.ts
@@ -35,18 +35,27 @@ let nodeOtelConfig: TelemetryConfig | null = null;
 
 const instrumentationKey = '__GENKIT_TELEMETRY_INSTRUMENTED';
 
-export function ensureBasicTelemetryInstrumentation() {
+export async function ensureBasicTelemetryInstrumentation() {
   if (global[instrumentationKey]) {
-    return;
+    return await global[instrumentationKey];
   }
-  enableTelemetry({});
-  global[instrumentationKey] = true;
+  await enableTelemetry({});
 }
 
 /**
  * Enables tracing and metrics open telemetry configuration.
  */
-export function enableTelemetry(telemetryConfig: TelemetryConfig) {
+export async function enableTelemetry(
+  telemetryConfig: TelemetryConfig | Promise<TelemetryConfig>
+) {
+  global[instrumentationKey] =
+    telemetryConfig instanceof Promise ? telemetryConfig : Promise.resolve();
+
+  telemetryConfig =
+    telemetryConfig instanceof Promise
+      ? await telemetryConfig
+      : telemetryConfig;
+
   nodeOtelConfig = telemetryConfig || {};
 
   const processors: SpanProcessor[] = [createTelemetryServerProcessor()];

diff --git a/js/core/src/tracing/instrumentation.ts b/js/core/src/tracing/instrumentation.ts
@@ -45,7 +45,7 @@ export async function newTrace<T>(
   },
   fn: (metadata: SpanMetadata, rootSpan: ApiSpan) => Promise<T>
 ) {
-  ensureBasicTelemetryInstrumentation();
+  await ensureBasicTelemetryInstrumentation();
   const traceMetadata: TraceMetadata = traceMetadataAls.getStore() || {
     paths: new Set<PathMetadata>(),
     timestamp: performance.now(),
@@ -78,7 +78,7 @@ export async function runInNewSpan<T>(
   },
   fn: (metadata: SpanMetadata, otSpan: ApiSpan, isRoot: boolean) => Promise<T>
 ): Promise<T> {
-  ensureBasicTelemetryInstrumentation();
+  await ensureBasicTelemetryInstrumentation();
 
   const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
   const parentStep = spanMetadataAls.getStore();

diff --git a/js/plugins/google-cloud/src/index.ts b/js/plugins/google-cloud/src/index.ts
@@ -23,13 +23,15 @@ import { GcpOpenTelemetry } from './gcpOpenTelemetry.js';
 import { TelemetryConfigs } from './telemetry/defaults.js';
 import { GcpTelemetryConfig, GcpTelemetryConfigOptions } from './types.js';
 
-export async function enableGoogleCloudTelemetry(
+export function enableGoogleCloudTelemetry(
   options?: GcpTelemetryConfigOptions
 ) {
-  const pluginConfig = await configureGcpPlugin(options);
-
-  enableTelemetry(await new GcpOpenTelemetry(pluginConfig).getConfig());
-  logger.init(await new GcpLogger(pluginConfig).getLogger(getCurrentEnv()));
+  return enableTelemetry(
+    configureGcpPlugin(options).then(async (pluginConfig) => {
+      logger.init(await new GcpLogger(pluginConfig).getLogger(getCurrentEnv()));
+      return new GcpOpenTelemetry(pluginConfig).getConfig();
+    })
+  );
 }
 
 /**

diff --git a/js/plugins/google-cloud/tests/traces_test.ts b/js/plugins/google-cloud/tests/traces_test.ts
@@ -188,7 +188,7 @@ describe('GoogleCloudTracing', () => {
     );
     const testFlow = createFlow(ai, 'modelFlow', async () => {
       return run('runFlow', async () => {
-        ai.generate({
+        await ai.generate({
           model: echoModel,
           prompt: 'Testing model telemetry',
         });
@@ -199,10 +199,11 @@ describe('GoogleCloudTracing', () => {
 
     const spans = await getExportedSpans();
 
-    assert.equal(spans[0].name, 'runFlow');
-    assert.equal(spans[1].name, 'modelFlow');
-    assert.equal(spans[2].name, 'echoModel');
-    assert.equal(spans[2].attributes['genkit/model'], 'echoModel');
+    assert.equal(spans[0].name, 'echoModel');
+    assert.equal(spans[0].attributes['genkit/model'], 'echoModel');
+    assert.equal(spans[1].name, 'generate');
+    assert.equal(spans[2].name, 'runFlow');
+    assert.equal(spans[3].name, 'modelFlow');
   });
 
   it('attaches additional span', async () => {