From 3dd1c30a56af021562ee8eeeaf70f9233c049b64 Mon Sep 17 00:00:00 2001 From: Ayaan Faisal Date: Fri, 5 Jun 2026 03:28:50 -0400 Subject: [PATCH] perf(core): Reuse unchanged Vercel AI messages JSON --- packages/core/src/tracing/vercel-ai/utils.ts | 17 +++- .../vercel-ai-request-messages.test.ts | 92 +++++++++++++++++++ 2 files changed, 105 insertions(+), 4 deletions(-) create mode 100644 packages/core/test/lib/tracing/vercel-ai-request-messages.test.ts diff --git a/packages/core/src/tracing/vercel-ai/utils.ts b/packages/core/src/tracing/vercel-ai/utils.ts index 215ddf54a40a..247cfa5ea3d4 100644 --- a/packages/core/src/tracing/vercel-ai/utils.ts +++ b/packages/core/src/tracing/vercel-ai/utils.ts @@ -254,8 +254,9 @@ export function requestMessagesFromPrompt(span: Span, attributes: SpanAttributes } else if (typeof attributes[AI_PROMPT_MESSAGES_ATTRIBUTE] === 'string') { // In this case we already get a properly formatted messages array, this is the preferred way to get the messages // This is the case for ai.generateText.doGenerate spans + const originalMessagesJson = attributes[AI_PROMPT_MESSAGES_ATTRIBUTE]; try { - const messages = JSON.parse(attributes[AI_PROMPT_MESSAGES_ATTRIBUTE]); + const messages = JSON.parse(originalMessagesJson); if (Array.isArray(messages)) { const { systemInstructions, filteredMessages } = extractSystemInstructions(messages); @@ -264,9 +265,17 @@ export function requestMessagesFromPrompt(span: Span, attributes: SpanAttributes } const filteredLength = Array.isArray(filteredMessages) ? filteredMessages.length : 0; - const messagesJson = enableTruncation - ? getTruncatedJsonString(filteredMessages) - : getJsonString(filteredMessages); + + // `extractSystemInstructions` returns the original array reference unchanged when no + // system message is extracted. When truncation is also disabled, re-serializing would + // reproduce the SDK's own input string, so we reuse it instead of allocating a second + // full-size copy of the payload (matters for large prompts in memory-constrained runtimes). + const messagesJson = + !enableTruncation && filteredMessages === messages + ? originalMessagesJson + : enableTruncation + ? getTruncatedJsonString(filteredMessages) + : getJsonString(filteredMessages); span.setAttributes({ [AI_PROMPT_MESSAGES_ATTRIBUTE]: messagesJson, diff --git a/packages/core/test/lib/tracing/vercel-ai-request-messages.test.ts b/packages/core/test/lib/tracing/vercel-ai-request-messages.test.ts new file mode 100644 index 000000000000..f2b1a89f961f --- /dev/null +++ b/packages/core/test/lib/tracing/vercel-ai-request-messages.test.ts @@ -0,0 +1,92 @@ +import { describe, expect, it } from 'vitest'; +import { getJsonString, getTruncatedJsonString } from '../../../src/tracing/ai/utils'; +import { + GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, + GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, +} from '../../../src/tracing/ai/gen-ai-attributes'; +import { requestMessagesFromPrompt } from '../../../src/tracing/vercel-ai/utils'; +import { AI_PROMPT_MESSAGES_ATTRIBUTE } from '../../../src/tracing/vercel-ai/vercel-ai-attributes'; +import type { Span, SpanAttributes } from '../../../src/types/span'; + +/** + * Minimal span that records the attributes set on it, so we can assert on the + * exact serialized value `requestMessagesFromPrompt` writes back. + */ +function createRecordingSpan(): { span: Span; recorded: Record } { + const recorded: Record = {}; + const span = { + setAttribute(key: string, value: unknown): void { + recorded[key] = value; + }, + setAttributes(attributes: Record): void { + Object.assign(recorded, attributes); + }, + } as unknown as Span; + return { span, recorded }; +} + +describe('requestMessagesFromPrompt (ai.prompt.messages string branch)', () => { + it('reuses the original string verbatim when no system message and truncation is off', () => { + const { span, recorded } = createRecordingSpan(); + + // Deliberately non-canonical whitespace. Re-serializing (JSON.stringify(JSON.parse(x))) + // would strip it, so a byte-identical result proves the original string was reused. + const original = '[ { "role": "user", "content": "hello world" } ]'; + const attributes = { [AI_PROMPT_MESSAGES_ATTRIBUTE]: original } as unknown as SpanAttributes; + + requestMessagesFromPrompt(span, attributes, /* enableTruncation */ false); + + expect(recorded[AI_PROMPT_MESSAGES_ATTRIBUTE]).toBe(original); + expect(recorded[GEN_AI_INPUT_MESSAGES_ATTRIBUTE]).toBe(original); + expect(recorded[GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]).toBe(1); + expect(recorded[GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE]).toBeUndefined(); + }); + + it('extracts the system message and re-serializes the remainder when truncation is off', () => { + const { span, recorded } = createRecordingSpan(); + + const original = JSON.stringify([ + { role: 'system', content: 'be nice' }, + { role: 'user', content: 'hello' }, + ]); + const attributes = { [AI_PROMPT_MESSAGES_ATTRIBUTE]: original } as unknown as SpanAttributes; + + requestMessagesFromPrompt(span, attributes, false); + + expect(recorded[GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE]).toBe(JSON.stringify([{ type: 'text', content: 'be nice' }])); + // System message removed; output is the SDK's own serialization of just the remainder. + expect(recorded[AI_PROMPT_MESSAGES_ATTRIBUTE]).toBe(getJsonString([{ role: 'user', content: 'hello' }])); + expect(recorded[AI_PROMPT_MESSAGES_ATTRIBUTE]).not.toBe(original); + expect(recorded[GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]).toBe(1); + }); + + it('keeps the truncation path untouched when truncation is on', () => { + const { span, recorded } = createRecordingSpan(); + + const messages = [ + { role: 'user', content: 'first' }, + { role: 'user', content: 'second' }, + ]; + const original = JSON.stringify(messages); + const attributes = { [AI_PROMPT_MESSAGES_ATTRIBUTE]: original } as unknown as SpanAttributes; + + requestMessagesFromPrompt(span, attributes, /* enableTruncation */ true); + + // Output must equal the SDK's own truncated serialization (and therefore differ from the + // input), proving the fast-path reuse did NOT short-circuit the truncation branch. + expect(recorded[AI_PROMPT_MESSAGES_ATTRIBUTE]).toBe(getTruncatedJsonString(messages)); + expect(recorded[AI_PROMPT_MESSAGES_ATTRIBUTE]).not.toBe(original); + // Original (pre-truncation) message count is still reported. + expect(recorded[GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]).toBe(2); + }); + + it('does not throw and sets no attributes for malformed JSON', () => { + const { span, recorded } = createRecordingSpan(); + + const attributes = { [AI_PROMPT_MESSAGES_ATTRIBUTE]: '{ not json' } as unknown as SpanAttributes; + + expect(() => requestMessagesFromPrompt(span, attributes, false)).not.toThrow(); + expect(Object.keys(recorded)).toHaveLength(0); + }); +});