Skip to content

Commit 4e1c929

Browse files
brandurbrandur
authored
Add setTag to emitter interface to set arbitrary key/value (#13)
Here, bring in a `setTag` function that lets X-ray users set an arbitrary key/value on a request span. These pairs are JSON encoded at the end of a request and emitted into the span. This is one that we've already supported for a while in Go and on the server, so we're just bringing equivalent functionality to this emitter so we can get access to it from the Stainless stack. We also remove the existing `setAttribute` function which would set a value to a span attribute. This functional was callable by package users, but would have no effect because the X-ray server doesn't ingest arbitrary span attributes.
1 parent a4fef64 commit 4e1c929

8 files changed

Lines changed: 308 additions & 13 deletions

File tree

‎src/core/attributes.test.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,3 +179,76 @@ test('setTenantIdAttribute sets stainlessxray.tenant.id', () => {
179179
setTenantIdAttribute?.(span as unknown as Span, 'tenant-123');
180180
assert.equal(span.attributes['stainlessxray.tenant.id'], 'tenant-123');
181181
});
182+
183+
test('setTagsAttribute serializes safe values', () => {
184+
const span = makeSpan();
185+
attributes.setTagsAttribute(span as unknown as Span, {
186+
str: 'hello',
187+
num: 42,
188+
bool: true,
189+
strArr: ['a', 'b'],
190+
numArr: [1, 2],
191+
boolArr: [true, false],
192+
});
193+
194+
const parsed = JSON.parse(span.attributes['stainlessxray.internal.tags'] as string);
195+
assert.deepEqual(parsed, {
196+
str: 'hello',
197+
num: 42,
198+
bool: true,
199+
strArr: ['a', 'b'],
200+
numArr: [1, 2],
201+
boolArr: [true, false],
202+
});
203+
});
204+
205+
test('setTagsAttribute drops BigInt values', () => {
206+
const span = makeSpan();
207+
const tags = { safe: 'ok', unsafe: BigInt(9007199254740991) } as Record<string, unknown>;
208+
attributes.setTagsAttribute(span as unknown as Span, tags as Record<string, never>);
209+
210+
const parsed = JSON.parse(span.attributes['stainlessxray.internal.tags'] as string);
211+
assert.deepEqual(parsed, { safe: 'ok' });
212+
});
213+
214+
test('setTagsAttribute drops undefined and null values', () => {
215+
const span = makeSpan();
216+
const tags = { kept: 'yes', undef: undefined, nil: null } as Record<string, unknown>;
217+
attributes.setTagsAttribute(span as unknown as Span, tags as Record<string, never>);
218+
219+
const parsed = JSON.parse(span.attributes['stainlessxray.internal.tags'] as string);
220+
assert.deepEqual(parsed, { kept: 'yes' });
221+
});
222+
223+
test('setTagsAttribute drops functions and symbols', () => {
224+
const span = makeSpan();
225+
const tags = { kept: 100, fn: () => {}, sym: Symbol('x') } as Record<string, unknown>;
226+
attributes.setTagsAttribute(span as unknown as Span, tags as Record<string, never>);
227+
228+
const parsed = JSON.parse(span.attributes['stainlessxray.internal.tags'] as string);
229+
assert.deepEqual(parsed, { kept: 100 });
230+
});
231+
232+
test('setTagsAttribute drops arrays containing unsafe elements', () => {
233+
const span = makeSpan();
234+
const tags = {
235+
good: [1, 2, 3],
236+
bad: [1, BigInt(2), 3],
237+
mixed: ['a', undefined, 'b'],
238+
} as Record<string, unknown>;
239+
attributes.setTagsAttribute(span as unknown as Span, tags as Record<string, never>);
240+
241+
const parsed = JSON.parse(span.attributes['stainlessxray.internal.tags'] as string);
242+
assert.deepEqual(parsed, { good: [1, 2, 3] });
243+
});
244+
245+
test('setTagsAttribute drops self-referential arrays', () => {
246+
const span = makeSpan();
247+
const cycle: unknown[] = [1, 2];
248+
cycle.push(cycle);
249+
const tags = { safe: 'ok', cycle } as Record<string, unknown>;
250+
attributes.setTagsAttribute(span as unknown as Span, tags as Record<string, never>);
251+
252+
const parsed = JSON.parse(span.attributes['stainlessxray.internal.tags'] as string);
253+
assert.deepEqual(parsed, { safe: 'ok' });
254+
});

‎src/core/attributes.ts‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import {
1010
ATTR_URL_PATH,
1111
ATTR_USER_ID,
1212
} from '@opentelemetry/semantic-conventions/incubating';
13-
import type { CapturedBody } from './types';
13+
import type { AttributeValue, CapturedBody } from './types';
1414
import {
1515
AttributeKeyRequestBody,
1616
AttributeKeyRequestBodyEncoding,
@@ -19,6 +19,7 @@ import {
1919
AttributeKeyResponseBody,
2020
AttributeKeyResponseBodyEncoding,
2121
AttributeKeyResponseBodyTruncated,
22+
AttributeKeyTags,
2223
AttributeKeyTenantID,
2324
} from './attrkey';
2425

@@ -291,3 +292,45 @@ export function setTenantIdAttribute(span: Span, tenantId: string): void {
291292
export function setRequestIdAttribute(span: Span, requestId: string): void {
292293
span.setAttribute(AttributeKeyRequestID, requestId);
293294
}
295+
296+
/**
297+
* JSON encode user-defined custom tags and add them as a span attribute.
298+
* Values that cannot be represented in JSON (e.g. BigInt) are silently dropped.
299+
*/
300+
export function setTagsAttribute(span: Span, tags: Record<string, AttributeValue>): void {
301+
const safe: Record<string, AttributeValue> = Object.create(null);
302+
for (const key of Object.keys(tags)) {
303+
const v = tags[key];
304+
if (isJSONSafe(v)) {
305+
safe[key] = v!;
306+
}
307+
}
308+
span.setAttribute(AttributeKeyTags, JSON.stringify(safe));
309+
}
310+
311+
// JSON-safe values should only be settable via the type system, but in case raw
312+
// JS is in use, do a runtime check for safe JSON values to avoid an exception
313+
// and all tags being lost.
314+
//
315+
// The `seen` argument is used to avoid infinite recursion when checking for
316+
// circular references in arrays. No need to pass it for most invocations.
317+
function isJSONSafe(value: unknown, seen?: Set<unknown>): boolean {
318+
if (value === null || value === undefined) {
319+
return false;
320+
}
321+
const t = typeof value;
322+
if (t === 'string' || t === 'number' || t === 'boolean') {
323+
return true;
324+
}
325+
if (Array.isArray(value)) {
326+
if (!seen) {
327+
seen = new Set();
328+
}
329+
if (seen.has(value)) {
330+
return false;
331+
}
332+
seen.add(value);
333+
return value.every((el) => isJSONSafe(el, seen));
334+
}
335+
return false;
336+
}

‎src/core/attrkey.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ export const AttributeKeyResponseBody = 'http.response.body';
88
export const AttributeKeyResponseBodyEncoding = 'http.response.body.encoding';
99
export const AttributeKeyResponseBodyTruncated = 'http.response.body.truncated';
1010
export const AttributeKeySpanDrop = 'stainlessxray.internal.drop';
11+
export const AttributeKeyTags = 'stainlessxray.internal.tags';
1112
export const AttributeKeyTenantID = 'stainlessxray.tenant.id';
1213
export const AttributeKeyXrayCaptureErrorCode = 'xray.capture_error_code';
1314
export const AttributeKeyXrayErrorThrown = 'xray.error_thrown';

‎src/core/emitter.test.ts‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,3 +331,177 @@ test('setUserId logs when span attributes fail', () => {
331331
assert.equal(errors[0]?.msg, 'xray: setUserId failed');
332332
assert.equal(errors[0]?.fields?.['error'], 'setAttribute failed');
333333
});
334+
335+
test('setTag records tags in log and span', async () => {
336+
const spans: ReadableSpan[] = [];
337+
const xray = createEmitter(
338+
{
339+
serviceName: 'test',
340+
endpointUrl: 'https://collector',
341+
exporter: { spanProcessor: 'simple' },
342+
},
343+
createRecordingExporter(spans),
344+
);
345+
346+
const ctx = xray.startRequest({
347+
method: 'GET',
348+
url: 'https://example.test/tags',
349+
headers: {},
350+
startTimeMs: 0,
351+
});
352+
ctx.setTag('environment', 'staging');
353+
ctx.setTag('region', 'us-east-1');
354+
ctx.setTag('retries', 3);
355+
ctx.setTag('verbose', true);
356+
357+
const log = xray.endRequest(ctx, {
358+
statusCode: 200,
359+
headers: {},
360+
endTimeMs: 5,
361+
});
362+
363+
await xray.flush();
364+
assert.deepEqual(log.tags, {
365+
environment: 'staging',
366+
region: 'us-east-1',
367+
retries: 3,
368+
verbose: true,
369+
});
370+
const tagsAttr = spans[0]?.attributes['stainlessxray.internal.tags'];
371+
assert.equal(typeof tagsAttr, 'string');
372+
assert.deepEqual(JSON.parse(tagsAttr as string), {
373+
environment: 'staging',
374+
region: 'us-east-1',
375+
retries: 3,
376+
verbose: true,
377+
});
378+
});
379+
380+
test('setTag overwrites previous value for same key', () => {
381+
const xray = createEmitter(
382+
{
383+
serviceName: 'test',
384+
endpointUrl: 'https://collector',
385+
},
386+
createNoopExporter(),
387+
);
388+
389+
const ctx = xray.startRequest({
390+
method: 'GET',
391+
url: 'https://example.test/tags-overwrite',
392+
headers: {},
393+
startTimeMs: 0,
394+
});
395+
ctx.setTag('env', 'dev');
396+
ctx.setTag('env', 'prod');
397+
398+
const log = xray.endRequest(ctx, {
399+
statusCode: 200,
400+
headers: {},
401+
endTimeMs: 5,
402+
});
403+
404+
assert.deepEqual(log.tags, { env: 'prod' });
405+
});
406+
407+
test('tags omitted from log and span when empty', async () => {
408+
const spans: ReadableSpan[] = [];
409+
const xray = createEmitter(
410+
{
411+
serviceName: 'test',
412+
endpointUrl: 'https://collector',
413+
exporter: { spanProcessor: 'simple' },
414+
},
415+
createRecordingExporter(spans),
416+
);
417+
418+
const ctx = xray.startRequest({
419+
method: 'GET',
420+
url: 'https://example.test/no-tags',
421+
headers: {},
422+
startTimeMs: 0,
423+
});
424+
425+
const log = xray.endRequest(ctx, {
426+
statusCode: 200,
427+
headers: {},
428+
endTimeMs: 5,
429+
});
430+
431+
await xray.flush();
432+
assert.equal(log.tags, undefined);
433+
assert.equal('stainlessxray.internal.tags' in (spans[0]?.attributes ?? {}), false);
434+
});
435+
436+
test('setTag treats __proto__ and constructor as plain keys', async () => {
437+
const spans: ReadableSpan[] = [];
438+
const xray = createEmitter(
439+
{
440+
serviceName: 'test',
441+
endpointUrl: 'https://collector',
442+
exporter: { spanProcessor: 'simple' },
443+
},
444+
createRecordingExporter(spans),
445+
);
446+
447+
const ctx = xray.startRequest({
448+
method: 'GET',
449+
url: 'https://example.test/proto-tags',
450+
headers: {},
451+
startTimeMs: 0,
452+
});
453+
ctx.setTag('__proto__', 'poisoned');
454+
ctx.setTag('constructor', 'overwritten');
455+
ctx.setTag('toString', 42);
456+
457+
const log = xray.endRequest(ctx, {
458+
statusCode: 200,
459+
headers: {},
460+
endTimeMs: 5,
461+
});
462+
463+
await xray.flush();
464+
assert.equal(log.tags?.['__proto__'], 'poisoned');
465+
assert.equal(log.tags?.['constructor'], 'overwritten');
466+
assert.equal(log.tags?.['toString'], 42);
467+
assert.equal(Object.keys(log.tags!).length, 3);
468+
const parsed = JSON.parse(spans[0]?.attributes['stainlessxray.internal.tags'] as string);
469+
assert.equal(parsed['__proto__'], 'poisoned');
470+
assert.equal(parsed['constructor'], 'overwritten');
471+
assert.equal(parsed['toString'], 42);
472+
});
473+
474+
test('setTag filters non-JSON-safe values from span attribute', async () => {
475+
const spans: ReadableSpan[] = [];
476+
const xray = createEmitter(
477+
{
478+
serviceName: 'test',
479+
endpointUrl: 'https://collector',
480+
exporter: { spanProcessor: 'simple' },
481+
},
482+
createRecordingExporter(spans),
483+
);
484+
485+
const ctx = xray.startRequest({
486+
method: 'GET',
487+
url: 'https://example.test/unsafe-tags',
488+
headers: {},
489+
startTimeMs: 0,
490+
});
491+
ctx.setTag('safe', 'kept');
492+
ctx.setTag('bigint' as string, BigInt(42) as never);
493+
494+
const log = xray.endRequest(ctx, {
495+
statusCode: 200,
496+
headers: {},
497+
endTimeMs: 5,
498+
});
499+
500+
await xray.flush();
501+
// The log retains the raw tags as-is
502+
assert.equal(log.tags?.['safe'], 'kept');
503+
assert.equal(log.tags?.['bigint'], BigInt(42));
504+
// The span attribute only contains the JSON-safe subset
505+
const parsed = JSON.parse(spans[0]?.attributes['stainlessxray.internal.tags'] as string);
506+
assert.deepEqual(parsed, { safe: 'kept' });
507+
});

‎src/core/emitter.ts‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import {
1414
setResponseBodySizeAttribute,
1515
setResponseStatusAttribute,
1616
setRouteAttribute,
17+
setTagsAttribute,
1718
setTenantIdAttribute,
1819
setUserIdAttribute,
1920
} from './attributes';
@@ -136,19 +137,12 @@ function startRequest(
136137
}
137138
state.sessionId = id;
138139
},
139-
setAttribute: (key, value) => {
140+
setTag: (key, value) => {
140141
const state = getContextState(context);
141142
if (!state) {
142143
return;
143144
}
144-
state.attributes[key] = value;
145-
if (span) {
146-
try {
147-
span.setAttribute(key, value as AttributeValue);
148-
} catch {
149-
// Ignore span attribute errors.
150-
}
151-
}
145+
state.tags[key] = value;
152146
},
153147
addEvent: (name, attributes) => {
154148
const state = getContextState(context);
@@ -187,6 +181,7 @@ function startRequest(
187181
context,
188182
attributes: {},
189183
events: [],
184+
tags: Object.create(null) as Record<string, AttributeValue>,
190185
};
191186

192187
bindContext(context, state);
@@ -250,6 +245,7 @@ function endRequest(
250245
sessionId: state.sessionId ?? undefined,
251246
error: buildError(err ?? state.error),
252247
attributes: Object.keys(state.attributes).length > 0 ? { ...state.attributes } : undefined,
248+
tags: Object.keys(state.tags).length > 0 ? { ...state.tags } : undefined,
253249
timestamp: new Date(endTimeMs).toISOString(),
254250
};
255251

@@ -304,6 +300,9 @@ function endRequest(
304300
if (state.userId) {
305301
setUserIdAttribute(span, state.userId);
306302
}
303+
if (Object.keys(state.tags).length > 0) {
304+
setTagsAttribute(span, state.tags);
305+
}
307306
if (err ?? state.error) {
308307
spanStatusFromError(span, err ?? state.error);
309308
}

‎src/core/state.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ export type RequestState = {
99
context: XrayContext;
1010
attributes: Record<string, AttributeValue>;
1111
events: Array<{ name: string; attributes?: Record<string, AttributeValue> }>;
12+
tags: Record<string, AttributeValue>;
1213
tenantId?: string;
1314
userId?: string;
1415
sessionId?: string;

0 commit comments

Comments
 (0)