weflayr 0.23.0 → 0.24.0

package/tests/instrumentation/vercel/get_cache_creation_telemetry.test.js

@@ -0,0 +1,96 @@
+'use strict';
+
+// Tests for the ai >= 7 cache-write capture: the telemetry dispatcher's
+// onLanguageModelCallEnd reads the raw `usage.cache_creation` off the event's
+// raw provider usage and stamps it verbatim on the still-open model-call span
+// (reached through the integration's call state) before ending it.
+
+const test = require('node:test');
+const assert = require('node:assert/strict');
+
+const {
+  VercelCacheCreationInstrumentor,
+  cacheCreationFromRawUsage,
+} = require('../../../src/instrumentation/vercel/get_cache_creation_telemetry');
+
+const RAW_ANTHROPIC_USAGE = {
+  input_tokens: 7,
+  cache_creation_input_tokens: 6053,
+  cache_creation: { ephemeral_5m_input_tokens: 6053, ephemeral_1h_input_tokens: 0 },
+};
+
+function makeFakeSpan() {
+  const attributes = {};
+  return {
+    attributes,
+    setAttribute(key, value) {
+      attributes[key] = value;
+    },
+  };
+}
+
+function makeInstrumentedDispatcher() {
+  const span = makeFakeSpan();
+  const endedEvents = [];
+  const dispatcher = {
+    onLanguageModelCallEnd(event) {
+      endedEvents.push(event);
+    },
+    getCallState(callId) {
+      return callId === 'call-1' ? { inferenceSpan: span } : undefined;
+    },
+  };
+  const instrumentor = new VercelCacheCreationInstrumentor();
+  instrumentor.instrument(dispatcher);
+  return { dispatcher, span, endedEvents, instrumentor };
+}
+
+test('stamps the raw cache_creation from the provider usage on the model-call span', () => {
+  const { dispatcher, span, endedEvents } = makeInstrumentedDispatcher();
+
+  dispatcher.onLanguageModelCallEnd({
+    callId: 'call-1',
+    usage: { inputTokens: 6060, outputTokens: 5, raw: RAW_ANTHROPIC_USAGE },
+  });
+
+  assert.deepEqual(
+    JSON.parse(span.attributes['weflayr.cache.creation']),
+    RAW_ANTHROPIC_USAGE.cache_creation
+  );
+  // The original handler still runs (it ends the span).
+  assert.equal(endedEvents.length, 1);
+});
+
+test('stamps nothing without a cache_creation in the raw usage', () => {
+  const { dispatcher, span } = makeInstrumentedDispatcher();
+
+  dispatcher.onLanguageModelCallEnd({
+    callId: 'call-1',
+    usage: { inputTokens: 13, outputTokens: 5, raw: { input_tokens: 13 } },
+  });
+
+  assert.ok(!('weflayr.cache.creation' in span.attributes));
+});
+
+test('cacheCreationFromRawUsage drops an all-zero cache_creation', () => {
+  assert.deepEqual(
+    cacheCreationFromRawUsage({
+      cache_creation: { ephemeral_5m_input_tokens: 4000, ephemeral_1h_input_tokens: 1300 },
+    }),
+    { ephemeral_5m_input_tokens: 4000, ephemeral_1h_input_tokens: 1300 }
+  );
+  assert.equal(
+    cacheCreationFromRawUsage({
+      cache_creation: { ephemeral_5m_input_tokens: 0, ephemeral_1h_input_tokens: 0 },
+    }),
+    undefined
+  );
+  assert.equal(cacheCreationFromRawUsage(undefined), undefined);
+});
+
+test('uninstrument restores the original handler', () => {
+  const { dispatcher, instrumentor } = makeInstrumentedDispatcher();
+  assert.equal(dispatcher.onLanguageModelCallEnd.name, 'onLanguageModelCallEndWithCacheCreation');
+  instrumentor.uninstrument();
+  assert.equal(dispatcher.onLanguageModelCallEnd.name, 'onLanguageModelCallEnd');
+});