diff --git a/docs/test.ts b/docs/test.ts index 349140dc661..fa07875f411 100644 --- a/docs/test.ts +++ b/docs/test.ts @@ -715,7 +715,7 @@ const aiguard = tracer.aiguard aiguard.evaluate([ { role: 'user', content: 'What is 2 + 2' }, ]).then(result => { - result.action && result.reason + result.action && result.reason && result.tags }) aiguard.evaluate([ @@ -729,11 +729,11 @@ aiguard.evaluate([ ], } ]).then(result => { - result.action && result.reason + result.action && result.reason && result.tags }) aiguard.evaluate([ { role: 'tool', tool_call_id: 'call_1', content: '5' }, ]).then(result => { - result.action && result.reason + result.action && result.reason && result.tags }) diff --git a/index.d.ts b/index.d.ts index 013eba8bbd6..e9c7302f5be 100644 --- a/index.d.ts +++ b/index.d.ts @@ -1320,6 +1320,10 @@ declare namespace tracer { * Human-readable explanation for why this action was chosen. */ reason: string; + /** + * List of tags associated with the evaluation (e.g. indirect-prompt-injection) + */ + tags: string[]; } /** @@ -1331,6 +1335,10 @@ declare namespace tracer { * Human-readable explanation from AI Guard describing why the conversation was blocked. */ reason: string; + /** + * List of tags associated with the evaluation (e.g. indirect-prompt-injection) + */ + tags: string[]; } /** @@ -1844,7 +1852,7 @@ declare namespace tracer { * [@google-cloud/pubsub](https://github.com/googleapis/nodejs-pubsub) module. */ interface google_cloud_pubsub extends Integration {} - + /** * This plugin automatically instruments the * [@google-cloud/vertexai](https://github.com/googleapis/nodejs-vertexai) module. diff --git a/packages/dd-trace/src/aiguard/sdk.js b/packages/dd-trace/src/aiguard/sdk.js index f18e5b98864..6ff31abbcbc 100644 --- a/packages/dd-trace/src/aiguard/sdk.js +++ b/packages/dd-trace/src/aiguard/sdk.js @@ -1,5 +1,6 @@ 'use strict' +const rfdc = require('rfdc')({ proto: false, circles: false }) const NoopAIGuard = require('./noop') const executeRequest = require('./client') const { @@ -22,10 +23,11 @@ const appsecMetrics = telemetryMetrics.manager.namespace('appsec') const ALLOW = 'ALLOW' class AIGuardAbortError extends Error { - constructor (reason) { + constructor (reason, tags) { super(reason) this.name = 'AIGuardAbortError' this.reason = reason + this.tags = tags } } @@ -77,20 +79,26 @@ class AIGuard extends NoopAIGuard { this.#initialized = true } - #truncate (messages) { + /** + * Returns a safe copy of the messages to be serialized into the meta struct. + * + * - Clones each message so callers cannot mutate the data set in the meta struct. + * - Truncates the list of messages and `content` fields emitting metrics accordingly. + */ + #buildMessagesForMetaStruct (messages) { const size = Math.min(messages.length, this.#maxMessagesLength) if (messages.length > size) { appsecMetrics.count(AI_GUARD_TELEMETRY_TRUNCATED, { type: 'messages' }).inc(1) } - const result = messages.slice(-size) - + const result = [] let contentTruncated = false - for (let i = 0; i < size; i++) { - const message = result[i] + for (let i = messages.length - size; i < messages.length; i++) { + const message = rfdc(messages[i]) if (message.content?.length > this.#maxContentSize) { contentTruncated = true - result[i] = { ...message, content: message.content.slice(0, this.#maxContentSize) } + message.content = message.content.slice(0, this.#maxContentSize) } + result.push(message) } if (contentTruncated) { appsecMetrics.count(AI_GUARD_TELEMETRY_TRUNCATED, { type: 'content' }).inc(1) @@ -139,7 +147,7 @@ class AIGuard extends NoopAIGuard { } } const metaStruct = { - messages: this.#truncate(messages) + messages: this.#buildMessagesForMetaStruct(messages) } span.meta_struct = { [AI_GUARD_META_STRUCT_KEY]: metaStruct @@ -192,9 +200,9 @@ class AIGuard extends NoopAIGuard { } if (shouldBlock) { span.setTag(AI_GUARD_BLOCKED_TAG_KEY, 'true') - throw new AIGuardAbortError(reason) + throw new AIGuardAbortError(reason, tags) } - return { action, reason } + return { action, reason, tags } }) } } diff --git a/packages/dd-trace/test/aiguard/index.spec.js b/packages/dd-trace/test/aiguard/index.spec.js index 447e8a7c432..8f4ffa17c3d 100644 --- a/packages/dd-trace/test/aiguard/index.spec.js +++ b/packages/dd-trace/test/aiguard/index.spec.js @@ -161,12 +161,15 @@ describe('AIGuard SDK', () => { if (shouldBlock) { await rejects( () => aiguard.evaluate(messages, { block: true }), - err => err.name === 'AIGuardAbortError' && err.reason === reason + err => err.name === 'AIGuardAbortError' && err.reason === reason && err.tags === tags ) } else { const evaluation = await aiguard.evaluate(messages, { block: true }) assert.strictEqual(evaluation.action, action) assert.strictEqual(evaluation.reason, reason) + if (tags) { + assert.strictEqual(evaluation.tags, tags) + } } assertTelemetry('ai_guard.requests', { error: false, action, block: shouldBlock }) @@ -302,6 +305,30 @@ describe('AIGuard SDK', () => { ) }) + it('test message immutability', async () => { + const messages = [{ + role: 'assistant', + tool_calls: [{ id: 'call_1', function: { name: 'shell', arguments: '{"cmd": "ls -lah"}' } }] + }] + mockFetch({ + body: { data: { attributes: { action: 'ALLOW', reason: 'OK', is_blocking_enabled: false } } } + }) + + await tracer.trace('test', async () => { + await aiguard.evaluate(messages) + // update messages before flushing + messages[0].tool_calls.push({ id: 'call_2', function: { name: 'shell', arguments: '{"cmd": "rm -rf"}' } }) + messages.push({ role: 'tool', tool_call_id: 'call_1', content: 'dir1, dir2, dir3' }) + }) + + await agent.assertSomeTraces(traces => { + const span = traces[0][1] // second span in the trace + const metaStruct = msgpack.decode(span.meta_struct.ai_guard) + assert.equal(metaStruct.messages.length, 1) + assert.equal(metaStruct.messages[0].tool_calls.length, 1) + }) + }) + it('test missing required fields uses noop as default', async () => { const client = new AIGuard(tracer, { aiguard: { endpoint: 'http://aiguard' } }) const result = await client.evaluate(toolCall)