diff --git a/packages/runtime-node/src/redact.ts b/packages/runtime-node/src/redact.ts index 93431a7..8557ff5 100644 --- a/packages/runtime-node/src/redact.ts +++ b/packages/runtime-node/src/redact.ts @@ -142,16 +142,41 @@ function redactValue(value: unknown, r: CompiledRedactor, depth: number): unknow if (typeof value === "object" && value !== null && depth > 0) { const out: Record = {}; for (const [k, v] of Object.entries(value as Record)) { - out[k] = isSensitiveKey(k, r) ? r.mask : redactValue(v, r, depth - 1); + out[k] = isSensitiveKey(k, v, r) ? r.mask : redactValue(v, r, depth - 1); } return out; } return value; } -function isSensitiveKey(key: string, r: CompiledRedactor): boolean { - const lowered = key.toLowerCase(); - return r.keyPatterns.some((re) => re.test(lowered)); +const USAGE_TOKEN_KEYS = new Set([ + "gen_ai.usage.input_tokens", + "gen_ai.usage.output_tokens", + "gen_ai.usage.prompt_tokens", + "gen_ai.usage.completion_tokens", + "gen_ai.usage.total_tokens", + "gen_ai.usage.token_count", +]); + +function isSensitiveKey( + key: string, + value: unknown, + r: CompiledRedactor, +): boolean { + const lowered = key.toLowerCase(); + + // Canonical GenAI usage attributes are safe when they contain + // valid non-negative numeric counts. + if (USAGE_TOKEN_KEYS.has(lowered)) { + return !( + typeof value === "number" && + Number.isFinite(value) && + value >= 0 + ); + } + + // All other sensitive keys, including token-like keys, are redacted. + return r.keyPatterns.some((re) => re.test(lowered)); } /** @@ -176,7 +201,7 @@ function redactWith( if (!attributes) return out; for (const [key, value] of Object.entries(attributes)) { if (value === undefined) continue; - out[key] = isSensitiveKey(key, r) + out[key] = isSensitiveKey(key, value, r) ? r.mask : (redactValue(value, r, maxDepth) as Attributes[string]); } diff --git a/packages/runtime-node/test/redact.test.mjs b/packages/runtime-node/test/redact.test.mjs index 3cf0a1d..cc73e06 100644 --- a/packages/runtime-node/test/redact.test.mjs +++ b/packages/runtime-node/test/redact.test.mjs @@ -124,3 +124,58 @@ test("empty/nullish input yields an empty object", () => { assert.deepEqual(redactAttributes(), {}); assert.deepEqual(redactAttributes(null), {}); }); + +test("keeps canonical GenAI usage counts only when numeric", () => { + const out = redactAttributes({ + "gen_ai.usage.input_tokens": 512, + "gen_ai.usage.output_tokens": 128, + "gen_ai.usage.prompt_tokens": 256, + "gen_ai.usage.completion_tokens": 128, + "gen_ai.usage.total_tokens": 640, + "gen_ai.usage.token_count": 42, + "gen_ai.usage.input_tokens_string": "secret", + max_tokens: 1000, + "secret.input_tokens": 999, + }); + + assert.deepEqual(out, { + "gen_ai.usage.input_tokens": 512, + "gen_ai.usage.output_tokens": 128, + "gen_ai.usage.prompt_tokens": 256, + "gen_ai.usage.completion_tokens": 128, + "gen_ai.usage.total_tokens": 640, + "gen_ai.usage.token_count": 42, + "gen_ai.usage.input_tokens_string": MASK, + max_tokens: MASK, + "secret.input_tokens": MASK, + }); +}); + +test("redacts non-numeric canonical GenAI usage values", () => { + const out = redactAttributes({ + "gen_ai.usage.input_tokens": "secret", + "gen_ai.usage.output_tokens": -1, + "gen_ai.usage.total_tokens": NaN, + "gen_ai.usage.token_count": Infinity, + }); + + assert.deepEqual(out, { + "gen_ai.usage.input_tokens": MASK, + "gen_ai.usage.output_tokens": MASK, + "gen_ai.usage.total_tokens": MASK, + "gen_ai.usage.token_count": MASK, + }); +}); +test("still masks secret token keys ending in token", () => { +const out = redactAttributes({ +token: "raw", +access_token: "raw", +refresh_token: "raw", +authToken: "raw", +token_value: "raw", +tokenString: "raw", +token_id: "raw", +id_token_hint: "raw", +}); +for (const value of Object.values(out)) assert.equal(value, MASK); +});