Skip to content

Commit 40f53c5

Browse files
AbhiPrasadAbhiPrasad
committed
Add tests
1 parent 88655ae commit 40f53c5

File tree

5 files changed

+83
-56
lines changed

5 files changed

+83
-56
lines changed

‎dev-packages/node-integration-tests/suites/tracing/vercelai/instrument-with-pii.mjs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
import * as Sentry from '@sentry/node';
2+
import { loggingTransport } from '@sentry-internal/node-integration-tests';
3+
4+
Sentry.init({
5+
dsn: 'https://public@dsn.ingest.sentry.io/1337',
6+
release: '1.0',
7+
tracesSampleRate: 1.0,
8+
sendDefaultPii: true,
9+
transport: loggingTransport,
10+
integrations: [Sentry.vercelAIIntegration()],
11+
});

‎dev-packages/node-integration-tests/suites/tracing/ai/instrument.mjs renamed to ‎dev-packages/node-integration-tests/suites/tracing/vercelai/instrument.mjs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,5 @@ Sentry.init({
66
release: '1.0',
77
tracesSampleRate: 1.0,
88
transport: loggingTransport,
9+
integrations: [Sentry.vercelAIIntegration()],
910
});

‎dev-packages/node-integration-tests/suites/tracing/ai/scenario.mjs renamed to ‎dev-packages/node-integration-tests/suites/tracing/vercelai/scenario.mjs

File renamed without changes.

‎dev-packages/node-integration-tests/suites/tracing/ai/test.ts renamed to ‎dev-packages/node-integration-tests/suites/tracing/vercelai/test.ts

Lines changed: 54 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,15 @@ import { afterAll, describe, expect } from 'vitest';
22
import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runner';
33

44
// `ai` SDK only support Node 18+
5-
describe('ai', () => {
5+
describe('Vercel AI integration', () => {
66
afterAll(() => {
77
cleanupChildProcesses();
88
});
99

10-
const EXPECTED_TRANSACTION = {
10+
const EXPECTED_TRANSACTION_DEFAULT_PII_FALSE = {
1111
transaction: 'main',
1212
spans: expect.arrayContaining([
13+
// First span - no telemetry config, should enable telemetry but not record inputs/outputs when sendDefaultPii: false
1314
expect.objectContaining({
1415
data: expect.objectContaining({
1516
'ai.completion_tokens.used': 20,
@@ -35,48 +36,51 @@ describe('ai', () => {
3536
origin: 'auto.vercelai.otel',
3637
status: 'ok',
3738
}),
39+
// Second span - explicitly enabled telemetry but recordInputs/recordOutputs not set, should not record when sendDefaultPii: false
3840
expect.objectContaining({
3941
data: expect.objectContaining({
40-
'sentry.origin': 'auto.vercelai.otel',
41-
'sentry.op': 'ai.run.doGenerate',
42-
'operation.name': 'ai.generateText.doGenerate',
43-
'ai.operationId': 'ai.generateText.doGenerate',
44-
'ai.model.provider': 'mock-provider',
42+
'ai.completion_tokens.used': 20,
4543
'ai.model.id': 'mock-model-id',
46-
'ai.settings.maxRetries': 2,
47-
'gen_ai.system': 'mock-provider',
48-
'gen_ai.request.model': 'mock-model-id',
49-
'ai.pipeline.name': 'generateText.doGenerate',
44+
'ai.model.provider': 'mock-provider',
5045
'ai.model_id': 'mock-model-id',
51-
'ai.streaming': false,
52-
'ai.response.finishReason': 'stop',
53-
'ai.response.model': 'mock-model-id',
54-
'ai.usage.promptTokens': 10,
55-
'ai.usage.completionTokens': 20,
56-
'gen_ai.response.finish_reasons': ['stop'],
57-
'gen_ai.usage.input_tokens': 10,
58-
'gen_ai.usage.output_tokens': 20,
59-
'ai.completion_tokens.used': 20,
46+
'ai.operationId': 'ai.generateText',
47+
'ai.pipeline.name': 'generateText',
6048
'ai.prompt_tokens.used': 10,
49+
'ai.response.finishReason': 'stop',
50+
'ai.settings.maxRetries': 2,
51+
'ai.settings.maxSteps': 1,
52+
'ai.streaming': false,
6153
'ai.total_tokens.used': 30,
54+
'ai.usage.completionTokens': 20,
55+
'ai.usage.promptTokens': 10,
56+
'operation.name': 'ai.generateText',
57+
'sentry.op': 'ai.pipeline.generateText',
58+
'sentry.origin': 'auto.vercelai.otel',
6259
}),
63-
description: 'generateText.doGenerate',
64-
op: 'ai.run.doGenerate',
60+
description: 'generateText',
61+
op: 'ai.pipeline.generateText',
6562
origin: 'auto.vercelai.otel',
6663
status: 'ok',
6764
}),
65+
]),
66+
};
67+
68+
const EXPECTED_TRANSACTION_DEFAULT_PII_TRUE = {
69+
transaction: 'main',
70+
spans: expect.arrayContaining([
71+
// First span - no telemetry config, should enable telemetry AND record inputs/outputs when sendDefaultPii: true
6872
expect.objectContaining({
6973
data: expect.objectContaining({
7074
'ai.completion_tokens.used': 20,
7175
'ai.model.id': 'mock-model-id',
7276
'ai.model.provider': 'mock-provider',
7377
'ai.model_id': 'mock-model-id',
74-
'ai.prompt': '{"prompt":"Where is the second span?"}',
78+
'ai.prompt': '{"prompt":"Where is the first span?"}',
7579
'ai.operationId': 'ai.generateText',
7680
'ai.pipeline.name': 'generateText',
7781
'ai.prompt_tokens.used': 10,
7882
'ai.response.finishReason': 'stop',
79-
'ai.input_messages': '{"prompt":"Where is the second span?"}',
83+
'ai.input_messages': '{"prompt":"Where is the first span?"}',
8084
'ai.settings.maxRetries': 2,
8185
'ai.settings.maxSteps': 1,
8286
'ai.streaming': false,
@@ -92,42 +96,46 @@ describe('ai', () => {
9296
origin: 'auto.vercelai.otel',
9397
status: 'ok',
9498
}),
99+
// Second span - explicitly enabled telemetry, should record inputs/outputs regardless of sendDefaultPii
95100
expect.objectContaining({
96101
data: expect.objectContaining({
97-
'sentry.origin': 'auto.vercelai.otel',
98-
'sentry.op': 'ai.run.doGenerate',
99-
'operation.name': 'ai.generateText.doGenerate',
100-
'ai.operationId': 'ai.generateText.doGenerate',
101-
'ai.model.provider': 'mock-provider',
102+
'ai.completion_tokens.used': 20,
102103
'ai.model.id': 'mock-model-id',
103-
'ai.settings.maxRetries': 2,
104-
'gen_ai.system': 'mock-provider',
105-
'gen_ai.request.model': 'mock-model-id',
106-
'ai.pipeline.name': 'generateText.doGenerate',
104+
'ai.model.provider': 'mock-provider',
107105
'ai.model_id': 'mock-model-id',
108-
'ai.streaming': false,
109-
'ai.response.finishReason': 'stop',
110-
'ai.response.model': 'mock-model-id',
111-
'ai.usage.promptTokens': 10,
112-
'ai.usage.completionTokens': 20,
113-
'gen_ai.response.finish_reasons': ['stop'],
114-
'gen_ai.usage.input_tokens': 10,
115-
'gen_ai.usage.output_tokens': 20,
116-
'ai.completion_tokens.used': 20,
106+
'ai.prompt': '{"prompt":"Where is the second span?"}',
107+
'ai.operationId': 'ai.generateText',
108+
'ai.pipeline.name': 'generateText',
117109
'ai.prompt_tokens.used': 10,
110+
'ai.response.finishReason': 'stop',
111+
'ai.input_messages': '{"prompt":"Where is the second span?"}',
112+
'ai.settings.maxRetries': 2,
113+
'ai.settings.maxSteps': 1,
114+
'ai.streaming': false,
118115
'ai.total_tokens.used': 30,
116+
'ai.usage.completionTokens': 20,
117+
'ai.usage.promptTokens': 10,
118+
'operation.name': 'ai.generateText',
119+
'sentry.op': 'ai.pipeline.generateText',
120+
'sentry.origin': 'auto.vercelai.otel',
119121
}),
120-
description: 'generateText.doGenerate',
121-
op: 'ai.run.doGenerate',
122+
description: 'generateText',
123+
op: 'ai.pipeline.generateText',
122124
origin: 'auto.vercelai.otel',
123125
status: 'ok',
124126
}),
125127
]),
126128
};
127129

128130
createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument.mjs', (createRunner, test) => {
129-
test('creates ai related spans ', async () => {
130-
await createRunner().expect({ transaction: EXPECTED_TRANSACTION }).start().completed();
131+
test('creates ai related spans with sendDefaultPii: false', async () => {
132+
await createRunner().expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_FALSE }).start().completed();
133+
});
134+
});
135+
136+
createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
137+
test('creates ai related spans with sendDefaultPii: true', async () => {
138+
await createRunner().expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_TRUE }).start().completed();
131139
});
132140
});
133141
});

‎packages/node/src/integrations/tracing/vercelai/instrumentation.ts

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -71,21 +71,28 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
7171
return (...args: MethodArgs) => {
7272
const existingExperimentalTelemetry = args[0].experimental_telemetry || {};
7373
const isEnabled = existingExperimentalTelemetry.isEnabled;
74+
7475
const client = getCurrentScope().getClient();
7576
const shouldRecordImportAndExports = client?.getIntegrationByName(INTEGRATION_NAME)
7677
? client.getOptions().sendDefaultPii
7778
: false;
7879

79-
// if `isEnabled` is not explicitly set to `true` or `false`, enable telemetry
80-
// but disable capturing inputs and outputs by default
81-
if (isEnabled === undefined) {
82-
args[0].experimental_telemetry = {
83-
isEnabled: true,
84-
recordInputs: shouldRecordImportAndExports,
85-
recordOutputs: shouldRecordImportAndExports,
86-
...existingExperimentalTelemetry,
87-
};
88-
}
80+
// Set recordInputs and recordOutputs based on sendDefaultPii if not explicitly set
81+
const recordInputs =
82+
existingExperimentalTelemetry.recordInputs !== undefined
83+
? existingExperimentalTelemetry.recordInputs
84+
: shouldRecordImportAndExports;
85+
const recordOutputs =
86+
existingExperimentalTelemetry.recordOutputs !== undefined
87+
? existingExperimentalTelemetry.recordOutputs
88+
: shouldRecordImportAndExports;
89+
90+
args[0].experimental_telemetry = {
91+
...existingExperimentalTelemetry,
92+
isEnabled: isEnabled !== undefined ? isEnabled : true,
93+
recordInputs,
94+
recordOutputs,
95+
};
8996

9097
// @ts-expect-error we know that the method exists
9198
return originalMethod.apply(this, args);

0 commit comments

Comments
 (0)