fix(Observe): missing traces in Mlflow

package.json

@@ -6,7 +6,7 @@
   "type": "module",
   "scripts": {
     "build": "vite build && cp -R src/embedding/adapters/caikit/grpc/protos dist/embedding/adapters/caikit/grpc/protos",
-    "start": "node --enable-source-maps --experimental-loader=@opentelemetry/instrumentation/hook.mjs --import ./dist/opentelemetry.js ./dist/server.js",
+    "start": "node ./dist/server.js",
     "start:dev": "tsx watch src/server.ts | pino-pretty --singleLine",
     "start:dev:nowatch": "tsx src/server.ts | pino-pretty --singleLine",
     "start:dev:workers": "concurrently npm:start:dev:workers:*",
@@ -43,6 +43,7 @@
     "@mikro-orm/seeder": "6.2.9",
     "@opentelemetry/auto-instrumentations-node": "^0.54.0",
     "@opentelemetry/exporter-metrics-otlp-http": "^0.54.0",
+    "@opentelemetry/exporter-trace-otlp-proto": "^0.57.0",
     "@opentelemetry/instrumentation": "^0.54.0",
     "@opentelemetry/sdk-node": "^0.54.0",
     "@opentelemetry/semantic-conventions": "^1.27.0",

src/opentelemetry.ts

@@ -18,38 +18,46 @@ import 'dotenv/config';
 
 import { setTimeout } from 'node:timers/promises';
 
-import { NodeSDK, resources, metrics } from '@opentelemetry/sdk-node';
+import { NodeSDK, resources, metrics, tracing } from '@opentelemetry/sdk-node';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
 import { ATTR_SERVICE_NAME } from '@opentelemetry/semantic-conventions';
 import { OTLPMetricExporter } from '@opentelemetry/exporter-metrics-otlp-http';
 import { ATTR_DEPLOYMENT_ENVIRONMENT_NAME } from '@opentelemetry/semantic-conventions/incubating';
+import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-proto';
 
 const ENV = process.env.ENVIRONMENT;
 
+export const isOpenTelemetryEnabled = Boolean(process.env.NODE_ENV === 'production');
+export const batchSpanProcessor = new tracing.BatchSpanProcessor(new OTLPTraceExporter({}));
+
 export const opentelemetrySDK = new NodeSDK({
   resource: new resources.Resource({
     [ATTR_SERVICE_NAME]: `bee-api`,
     [ATTR_DEPLOYMENT_ENVIRONMENT_NAME]: ENV
   }),
+  spanProcessors: [batchSpanProcessor],
   metricReader: new metrics.PeriodicExportingMetricReader({ exporter: new OTLPMetricExporter() }),
   instrumentations: [...getNodeAutoInstrumentations()]
 });
-opentelemetrySDK.start();
-
-let isShuttingDown = false;
-const { promise, resolve } = Promise.withResolvers<void>();
-
-for (const event of ['beforeExit', 'SIGINT', 'SIGTERM']) {
-  process.once(event, () => {
-    if (!isShuttingDown) {
-      isShuttingDown = true;
-      Promise.race([opentelemetrySDK.shutdown(), setTimeout(5_000, null, { ref: false })])
-        .catch((err) => {
-          // eslint-disable-next-line no-console
-          console.error(`Failed to execute shutdown hook`, err);
-        })
-        .finally(() => resolve());
-    }
-    return promise;
-  });
+
+if (isOpenTelemetryEnabled) {
+  opentelemetrySDK.start();
+
+  let isShuttingDown = false;
+  const { promise, resolve } = Promise.withResolvers<void>();
+
+  for (const event of ['beforeExit', 'SIGINT', 'SIGTERM']) {
+    process.once(event, () => {
+      if (!isShuttingDown) {
+        isShuttingDown = true;
+        Promise.race([opentelemetrySDK.shutdown(), setTimeout(5_000, null, { ref: false })])
+          .catch((err) => {
+            // eslint-disable-next-line no-console
+            console.error(`Failed to execute shutdown hook`, err);
+          })
+          .finally(() => resolve());
+      }
+      return promise;
+    });
+  }
 }

src/runs/execution/execute.ts

@@ -46,6 +46,7 @@ import { LoadedRun } from '@/runs/execution/types.js';
 import { UserResource } from '@/tools/entities/tool-resources/user-resource.entity.js';
 import { SystemResource } from '@/tools/entities/tool-resources/system-resource.entity.js';
 import { Attachment } from '@/messages/attachment.entity';
+import { batchSpanProcessor, isOpenTelemetryEnabled } from '@/opentelemetry.js';
 
 const agentExecutionTime = new Summary({
   name: 'agent_execution_time_seconds',
@@ -144,6 +145,7 @@ export async function executeRun(run: LoadedRun) {
       );
 
       await agentRun;
+      if (isOpenTelemetryEnabled) await batchSpanProcessor.forceFlush();
 
       endAgentExecutionTimer();
       run.complete();

src/server.ts

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import './opentelemetry.js';
 import '@/ui/auth-server.js';
 
 import { JsonSchemaToTsProvider } from '@fastify/type-provider-json-schema-to-ts';