@telemetry-dev/bedrock 0.1.10 → 0.1.13

package/README.md

@@ -2,6 +2,8 @@
 
 First-class AWS Bedrock instrumentation for the telemetry.dev TypeScript SDK.
 
+Docs: https://telemetry.dev/docs/integrations/bedrock
+
 This package instruments AWS SDK for JavaScript v3 clients for:
 
 - `@aws-sdk/client-bedrock-runtime`
@@ -29,15 +31,18 @@ pnpm add @telemetry-dev/sdk @telemetry-dev/bedrock @aws-sdk/client-bedrock-runti
 pnpm add @aws-sdk/client-bedrock-agent-runtime
 ```
 
+Requires `@aws-sdk/client-bedrock-runtime` 3.600.0 or a later 3.x release (and the same range for the optional `@aws-sdk/client-bedrock-agent-runtime`), and `@telemetry-dev/sdk`, which needs Node.js 20.19.0 or later.
+
 ## Quickstart
 
+`init()` reads your project API key from the `TELEMETRY_DEV_API_KEY` environment variable, so set it before you start the app.
+
 ```ts
 import { BedrockRuntimeClient, ConverseCommand } from "@aws-sdk/client-bedrock-runtime";
 import { init, shutdown } from "@telemetry-dev/sdk";
 import { wrapBedrock } from "@telemetry-dev/bedrock";
 
 init({
-  apiKey: process.env.TELEMETRY_DEV_API_KEY,
   serviceName: "bedrock-app",
   environment: "production",
 });
@@ -98,9 +103,22 @@ instrumentBedrockAgents({ captureAgentTrace: false });
 
 ## Options
 
-| Option              | Default | Applies to            | Notes                                                                                                                                                 |
-| ------------------- | ------- | --------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `captureAgentTrace` | `false` | Agent Runtime streams | Aggregates trace usage and counts by default. When enabled, also attaches raw trace events as `td.metadata.agent_trace` after SDK masking/truncation. |
+| Option              | Default | Applies to            | Notes                                                                                                                                             |
+| ------------------- | ------- | --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `captureAgentTrace` | `false` | Agent Runtime streams | Aggregates trace usage and counts by default. When enabled, also attaches trace events as `td.metadata.agent_trace`, filtered as described below. |
+
+`td.metadata.agent_trace` is truncated but never passed to your `mask`, so each trace event is rebuilt from an explicit allowlist instead:
+
+- Envelope fields: `agentId`, `agentAliasId`, `agentVersion`, `sessionId`, `collaboratorName`, `eventTime`, and the alias ARNs in `callerChain`.
+- Content the agent sends to a model, tool, or guardrail (`modelInvocationInput`, `invocationInput`, `inputAssessments`) is kept as received when `captureInput` is on.
+- Content that comes back (`modelInvocationOutput`, `rationale`, `observation`, `outputAssessments`, custom orchestration `event`) is kept as received when `captureOutput` is on.
+- A failure reason (`failureTrace.failureReason`, and `td.metadata.failure_reason`) is free-form text, so it is kept only when both are on. `td.metadata.failure` always records that the invocation failed, with `td.metadata.failure_code` when AWS sends one.
+- When a `mask` is set, all of this content is withheld.
+- Withheld content and the trace structure around it keep only trace ids, step types, the foundation model, guardrail action, failure code, and the usage, timings, and `clientRequestId` from AWS's trace `metadata`.
+
+Any other field, including unknown trace types and the AWS SDK's `$unknown` union member, is omitted even when both capture options are on; a content part that is kept is exported whole, including any fields inside it that are not listed here. Errors are not covered by this filtering: AWS exceptions, including errors raised inside the stream, still record the provider's error message on the span.
+
+Output file names in `td.metadata.output_files` follow `captureOutput` the same way, and are withheld when a `mask` is set; file types are always kept.
 
 ## Signal coverage
 
@@ -159,3 +177,7 @@ The following calls pass through unless a supported operation above is used:
 - `AgenticRetrieveStream`
 
 Provider-native `InvokeModel` usage is best-effort. Anthropic Claude, Amazon Titan/Nova, Meta Llama, and Titan embeddings expose known token fields. Cohere and Mistral text models do not consistently include usage in the JSON body; spans still capture request/response bodies and finish reasons where present.
+
+## License
+
+MIT. See [LICENSE](https://github.com/telemetry-dev/telemetry/blob/main/packages/bedrock/LICENSE).

package/dist/agents.mjs

@@ -1,4 +1,4 @@
-import { r as createInstrumentedSend, t as AGENT_HANDLERS } from "./internal-handlers-PftQAPH3.mjs";
+import { r as createInstrumentedSend, t as AGENT_HANDLERS } from "./internal-handlers-CLP2wXDh.mjs";
 import { BedrockAgentRuntimeClient } from "@aws-sdk/client-bedrock-agent-runtime";
 //#region src/agents.ts
 const WRAPPED = /* @__PURE__ */ new WeakSet();

package/dist/index.mjs

@@ -1,4 +1,4 @@
-import { n as RUNTIME_HANDLERS, r as createInstrumentedSend } from "./internal-handlers-PftQAPH3.mjs";
+import { n as RUNTIME_HANDLERS, r as createInstrumentedSend } from "./internal-handlers-CLP2wXDh.mjs";
 import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime";
 //#region src/index.ts
 const WRAPPED = /* @__PURE__ */ new WeakSet();

package/dist/{internal-handlers-PftQAPH3.mjs → internal-handlers-CLP2wXDh.mjs}

@@ -413,6 +413,8 @@ var AgentStreamState = class {
 	textRejected = false;
 	usage = {};
 	traceEventCount = 0;
+	failed = false;
+	failureCode;
 	failureReason;
 	guardrailAction;
 	returnControl;
@@ -423,6 +425,10 @@ var AgentStreamState = class {
 	errorFields;
 	citationCount = 0;
 	outputFiles = [];
+	contentCapture = {
+		input: captureEnabled("input") && !captureMaskedBeforeStart(),
+		output: captureEnabled("output") && !captureMaskedBeforeStart()
+	};
 	constructor(options) {
 		this.options = options;
 	}
@@ -442,7 +448,7 @@ var AgentStreamState = class {
 		else this.textRejected = true;
 		const citations = arrayValue((isRecord(chunk?.attribution) ? chunk.attribution : void 0)?.citations);
 		if (citations) this.citationCount += citations.length;
-		const sanitizedFiles = (isRecord(raw.files) ? arrayValue(raw.files.files) : void 0)?.map(sanitizeOutputFile);
+		const sanitizedFiles = (isRecord(raw.files) ? arrayValue(raw.files.files) : void 0)?.map((file) => sanitizeOutputFile(file, this.contentCapture.output));
 		if (sanitizedFiles && this.retention.accept(sanitizedFiles)) this.outputFiles.push(...sanitizedFiles);
 		const returnControl = isRecord(raw.returnControl) ? raw.returnControl : void 0;
 		if (returnControl) {
@@ -457,11 +463,23 @@ var AgentStreamState = class {
 		const traceEvent = isRecord(raw.trace) ? raw.trace : void 0;
 		if (traceEvent) {
 			this.traceEventCount += 1;
-			if (this.options.captureAgentTrace && this.retention.accept(traceEvent)) this.traces.push(traceEvent);
-			collectTrace(traceEvent, (usage) => this.addUsage(usage));
-			const failure = findKey(traceEvent, "failureTrace");
-			if (isRecord(failure)) this.failureReason = stringValue(failure.failureReason) ?? this.failureReason;
-			const guardrail = findKey(traceEvent, "guardrailTrace");
+			if (this.options.captureAgentTrace) {
+				const retained = projectTrace(traceEvent, TRACE_PART, this.contentCapture);
+				if (this.retention.accept(retained)) this.traces.push(retained);
+			}
+			const trace = isRecord(traceEvent.trace) ? traceEvent.trace : {};
+			for (const step of MODEL_INVOCATION_STEPS) {
+				const member = isRecord(trace[step]) ? trace[step] : void 0;
+				const output = isRecord(member?.modelInvocationOutput) ? member.modelInvocationOutput : {};
+				this.addUsage(isRecord(output.metadata) ? output.metadata.usage : void 0);
+			}
+			const failure = trace.failureTrace;
+			if (isRecord(failure)) {
+				this.failed = true;
+				this.failureCode = numberValue(failure.failureCode) ?? this.failureCode;
+				if (this.contentCapture.input && this.contentCapture.output) this.failureReason = stringValue(failure.failureReason) ?? this.failureReason;
+			}
+			const guardrail = trace.guardrailTrace;
 			if (isRecord(guardrail)) this.guardrailAction = stringValue(guardrail.action) ?? this.guardrailAction;
 		}
 	}
@@ -472,6 +490,8 @@ var AgentStreamState = class {
 	finish() {
 		const metadata = omitUndefined({
 			trace_event_count: this.traceEventCount || void 0,
+			failure: this.failed ? true : void 0,
+			failure_code: this.failureCode,
 			failure_reason: this.failureReason,
 			guardrail_action: this.guardrailAction,
 			return_control: this.sawReturnControl ? true : void 0,
@@ -504,11 +524,11 @@ var AgentStreamState = class {
 		else if (inputTokens !== void 0 || outputTokens !== void 0) this.usage.totalTokens = (this.usage.inputTokens ?? 0) + (this.usage.outputTokens ?? 0);
 	}
 };
-function sanitizeOutputFile(value) {
+function sanitizeOutputFile(value, keepName) {
 	const raw = value;
 	if (!isRecord(raw)) return void 0;
 	return omitUndefined({
-		name: stringValue(raw.name),
+		name: keepName ? stringValue(raw.name) : void 0,
 		type: stringValue(raw.type)
 	});
 }
@@ -700,22 +720,111 @@ function flowName(input) {
 function ragName(input) {
 	return `retrieve_and_generate ${basenameArn(ragModel(input)) ?? "knowledge-base"}`;
 }
-function collectTrace(value, onUsage) {
-	const raw = value;
-	if (!isRecord(raw)) return;
-	const output = isRecord(raw.modelInvocationOutput) ? raw.modelInvocationOutput : void 0;
-	if (output) onUsage((isRecord(output.metadata) ? output.metadata : void 0)?.usage);
-	for (const child of Object.values(raw)) if (isRecord(child)) collectTrace(child, onUsage);
-	else if (Array.isArray(child)) child.forEach((item) => collectTrace(item, onUsage));
-}
-function findKey(value, key) {
-	const raw = value;
-	if (!isRecord(raw)) return void 0;
-	if (raw[key] !== void 0) return raw[key];
-	for (const child of Object.values(raw)) {
-		const found = findKey(child, key);
-		if (found !== void 0) return found;
+/** Content kept as received when every listed capture is on; otherwise reduced to `withheld`. */
+var TraceContent = class {
+	constructor(needs, withheld) {
+		this.needs = needs;
+		this.withheld = withheld;
 	}
+};
+const SCALAR = "scalar";
+const TRACE_METADATA = {
+	usage: {
+		inputTokens: SCALAR,
+		outputTokens: SCALAR
+	},
+	startTime: SCALAR,
+	endTime: SCALAR,
+	totalTimeMs: SCALAR,
+	operationTotalTimeMs: SCALAR,
+	clientRequestId: SCALAR
+};
+const MODEL_INVOCATION_INPUT = new TraceContent(["input"], {
+	traceId: SCALAR,
+	type: SCALAR,
+	foundationModel: SCALAR
+});
+const INVOCATION_INPUT = new TraceContent(["input"], {
+	traceId: SCALAR,
+	invocationType: SCALAR
+});
+const MODEL_INVOCATION_OUTPUT = new TraceContent(["output"], {
+	traceId: SCALAR,
+	metadata: TRACE_METADATA
+});
+const OBSERVATION = new TraceContent(["output"], {
+	traceId: SCALAR,
+	type: SCALAR
+});
+const MODEL_INVOCATION_STEPS = [
+	"preProcessingTrace",
+	"orchestrationTrace",
+	"postProcessingTrace",
+	"routingClassifierTrace"
+];
+const PROCESSING_TRACE = {
+	modelInvocationInput: MODEL_INVOCATION_INPUT,
+	modelInvocationOutput: MODEL_INVOCATION_OUTPUT
+};
+const TRACE_PART = {
+	agentId: SCALAR,
+	agentAliasId: SCALAR,
+	agentVersion: SCALAR,
+	sessionId: SCALAR,
+	collaboratorName: SCALAR,
+	eventTime: SCALAR,
+	callerChain: [{ agentAliasArn: SCALAR }],
+	trace: {
+		preProcessingTrace: PROCESSING_TRACE,
+		postProcessingTrace: PROCESSING_TRACE,
+		orchestrationTrace: {
+			modelInvocationInput: MODEL_INVOCATION_INPUT,
+			invocationInput: INVOCATION_INPUT,
+			modelInvocationOutput: MODEL_INVOCATION_OUTPUT,
+			rationale: new TraceContent(["output"], { traceId: SCALAR }),
+			observation: OBSERVATION
+		},
+		routingClassifierTrace: {
+			modelInvocationInput: MODEL_INVOCATION_INPUT,
+			invocationInput: INVOCATION_INPUT,
+			modelInvocationOutput: MODEL_INVOCATION_OUTPUT,
+			observation: OBSERVATION
+		},
+		customOrchestrationTrace: {
+			traceId: SCALAR,
+			event: new TraceContent(["output"])
+		},
+		guardrailTrace: {
+			traceId: SCALAR,
+			action: SCALAR,
+			metadata: TRACE_METADATA,
+			inputAssessments: new TraceContent(["input"]),
+			outputAssessments: new TraceContent(["output"])
+		},
+		failureTrace: {
+			traceId: SCALAR,
+			failureCode: SCALAR,
+			metadata: TRACE_METADATA,
+			failureReason: new TraceContent(["input", "output"])
+		}
+	}
+};
+function projectTrace(value, shape, capture) {
+	if (shape instanceof TraceContent) {
+		if (shape.needs.every((kind) => capture[kind])) return value ?? void 0;
+		return shape.withheld === void 0 ? void 0 : projectTrace(value, shape.withheld, capture);
+	}
+	if (shape === SCALAR) return [
+		"string",
+		"number",
+		"boolean"
+	].includes(typeof value) || value instanceof Date ? value : void 0;
+	if (Array.isArray(shape)) {
+		if (!Array.isArray(value)) return void 0;
+		return value.map((item) => projectTrace(item, shape[0], capture)).filter((item) => item !== void 0);
+	}
+	if (!isRecord(value)) return void 0;
+	return omitUndefined(Object.fromEntries(Object.entries(shape).map(([field, child]) => [field, projectTrace(value[field], child, capture)])));
 }
 //#endregion
 //#region src/converse.ts
@@ -1147,12 +1256,7 @@ function nativeResponseFields(model, body, metadata) {
 		return withHeaderUsage({
 			output: body,
 			finishReason: stringValue(body.stop_reason),
-			usage: omitUndefined({
-				inputTokens: numberValue(usage.input_tokens),
-				outputTokens: numberValue(usage.output_tokens),
-				cacheReadInputTokens: numberValue(usage.cache_read_input_tokens),
-				cacheCreationInputTokens: numberValue(usage.cache_creation_input_tokens)
-			})
+			...claudeUsageFields(usage)
 		}, metadata);
 	}
 	if (lower.includes("meta.llama")) return omitUndefined({
@@ -1177,6 +1281,35 @@ function nativeResponseFields(model, body, metadata) {
 	}, metadata);
 	return withHeaderUsage({ output: body }, metadata);
 }
+const CLAUDE_USAGE_KEYS = {
+	inputTokens: "input_tokens",
+	outputTokens: "output_tokens",
+	cacheReadInputTokens: "cache_read_input_tokens",
+	cacheCreationInputTokens: "cache_creation_input_tokens"
+};
+function claudeUsageFields(record) {
+	const usage = {};
+	for (const [key, source] of Object.entries(CLAUDE_USAGE_KEYS)) {
+		const value = numberValue(record[source]);
+		if (value !== void 0) usage[key] = value;
+	}
+	const iterations = record.iterations;
+	if (!Array.isArray(iterations)) return { usage };
+	const first = Math.max(0, iterations.length - STREAM_CAPTURE_MAX_ITEMS);
+	for (let index = first; index < iterations.length; index += 1) {
+		const entry = iterations[index];
+		if (!isRecord(entry) || entry.type !== "compaction") continue;
+		for (const [key, source] of Object.entries(CLAUDE_USAGE_KEYS)) {
+			const value = numberValue(entry[source]);
+			const target = key;
+			if (value !== void 0) usage[target] = (usage[target] ?? 0) + value;
+		}
+	}
+	return {
+		usage,
+		...first > 0 ? { attributes: { "telemetry.dev.capture.truncated": true } } : {}
+	};
+}
 function withHeaderUsage(fields, metadata) {
 	if (fields.usage && Object.keys(fields.usage).length > 0) return fields;
 	return mergeFields(fields, headerUsageFields(metadata));
@@ -1202,6 +1335,9 @@ var InvokeModelStreamState = class {
 	chunks = [];
 	text = "";
 	usage;
+	finalUsage = {};
+	startUsage = {};
+	compaction = {};
 	finishReason;
 	error;
 	errorFields;
@@ -1228,21 +1364,44 @@ var InvokeModelStreamState = class {
 			const message = isRecord(parsed.message) ? parsed.message : void 0;
 			const messageUsage = isRecord(message?.usage) ? message?.usage : void 0;
 			const usage = isRecord(parsed.usage) ? parsed.usage : void 0;
+			for (const [key, source] of Object.entries(CLAUDE_USAGE_KEYS)) {
+				const value = numberValue(usage?.[source]);
+				if (value !== void 0) this.finalUsage[key] = value;
+			}
+			this.startUsage = omitUndefined({
+				...this.startUsage,
+				inputTokens: numberValue(messageUsage?.input_tokens) ?? this.startUsage.inputTokens,
+				outputTokens: numberValue(messageUsage?.output_tokens) ?? this.startUsage.outputTokens,
+				cacheReadInputTokens: numberValue(messageUsage?.cache_read_input_tokens) ?? this.startUsage.cacheReadInputTokens,
+				cacheCreationInputTokens: numberValue(messageUsage?.cache_creation_input_tokens) ?? this.startUsage.cacheCreationInputTokens
+			});
+			const iterations = Array.isArray(usage?.iterations) ? usage.iterations : messageUsage?.iterations;
+			if (Array.isArray(iterations)) this.compaction = claudeUsageFields({ iterations });
 			const metrics = isRecord(parsed["amazon-bedrock-invocationMetrics"]) ? parsed["amazon-bedrock-invocationMetrics"] : void 0;
 			const outputs = Array.isArray(parsed.outputs) && isRecord(parsed.outputs[0]) ? parsed.outputs[0] : void 0;
 			const generation = Array.isArray(parsed.generations) && isRecord(parsed.generations[0]) ? parsed.generations[0] : void 0;
 			this.finishReason = stringValue(delta?.stop_reason) ?? stringValue(parsed.stop_reason) ?? stringValue(outputs?.stop_reason) ?? stringValue(generation?.finish_reason) ?? stringValue(parsed.completionReason) ?? this.finishReason;
 			this.usage = omitUndefined({
 				...this.usage,
-				inputTokens: numberValue(messageUsage?.input_tokens) ?? numberValue(metrics?.inputTokenCount) ?? numberValue(parsed.inputTextTokenCount) ?? this.usage?.inputTokens,
+				inputTokens: numberValue(usage?.input_tokens) ?? numberValue(metrics?.inputTokenCount) ?? numberValue(parsed.inputTextTokenCount) ?? this.usage?.inputTokens,
 				outputTokens: numberValue(usage?.output_tokens) ?? numberValue(metrics?.outputTokenCount) ?? numberValue(parsed.totalOutputTextTokenCount) ?? this.usage?.outputTokens,
-				cacheReadInputTokens: numberValue(messageUsage?.cache_read_input_tokens) ?? this.usage?.cacheReadInputTokens,
-				cacheCreationInputTokens: numberValue(messageUsage?.cache_creation_input_tokens) ?? this.usage?.cacheCreationInputTokens
+				cacheReadInputTokens: numberValue(usage?.cache_read_input_tokens) ?? this.usage?.cacheReadInputTokens,
+				cacheCreationInputTokens: numberValue(usage?.cache_creation_input_tokens) ?? this.usage?.cacheCreationInputTokens
 			});
 		}
 		return providerChunkHasOutput(parsed);
 	}
 	finish() {
+		const usage = {
+			...this.startUsage,
+			...this.usage,
+			...this.finalUsage
+		};
+		for (const key of Object.keys(CLAUDE_USAGE_KEYS)) {
+			const value = this.compaction.usage?.[key];
+			const baseKey = key === "outputTokens" ? "outputTokens" : "inputTokens";
+			if (value !== void 0 && (this.finalUsage[baseKey] !== void 0 || this.usage?.[baseKey] === void 0)) usage[key] = (usage[key] ?? 0) + value;
+		}
 		return mergeFields(omitUndefined({
 			output: this.text ? [{
 				role: "assistant",
@@ -1253,7 +1412,8 @@ var InvokeModelStreamState = class {
 				finish_reason: this.finishReason
 			}] : this.chunks,
 			finishReason: this.finishReason,
-			usage: this.usage,
+			usage,
+			attributes: this.compaction.attributes,
 			error: this.error
 		}), this.errorFields);
 	}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@telemetry-dev/bedrock",
-  "version": "0.1.10",
+  "version": "0.1.13",
   "description": "AWS Bedrock instrumentation for telemetry.dev: Converse, InvokeModel, and Bedrock Agents spans for AWS SDK v3 clients.",
   "keywords": [
     "aws",
@@ -44,12 +44,12 @@
     "@aws-sdk/client-bedrock-agent-runtime": "3.1135.0",
     "@aws-sdk/client-bedrock-runtime": "3.1135.0",
     "@opentelemetry/api": "^1.9.1",
-    "@opentelemetry/sdk-logs": "^0.218.0",
-    "@opentelemetry/sdk-trace-base": "^2.7.1",
+    "@opentelemetry/sdk-logs": "^0.223.0",
+    "@opentelemetry/sdk-trace-base": "^2.12.0",
     "@smithy/eventstream-codec": "^4.2.5",
     "@smithy/protocol-http": "^5.3.10",
     "@smithy/types": "^4.12.0",
-    "@telemetry-dev/sdk": "^0.1.8",
+    "@telemetry-dev/sdk": "^0.1.11",
     "@types/node": "^25.5.0",
     "@typescript/native-preview": "7.0.0-dev.20260328.1",
     "typescript": "^6.0.2",
@@ -59,7 +59,7 @@
   "peerDependencies": {
     "@aws-sdk/client-bedrock-agent-runtime": ">=3.600.0 <4",
     "@aws-sdk/client-bedrock-runtime": ">=3.600.0 <4",
-    "@telemetry-dev/sdk": "^0.1.8"
+    "@telemetry-dev/sdk": "^0.1.11"
   },
   "peerDependenciesMeta": {
     "@aws-sdk/client-bedrock-agent-runtime": {

package/src/agent-handlers.ts

@@ -5,6 +5,7 @@ import {
   awsMetadataFields,
   basenameArn,
   bytesToString,
+  captureMaskedBeforeStart,
   isRecord,
   mergeFields,
   modeledStreamError,
@@ -55,6 +56,8 @@ export class AgentStreamState implements StreamState {
   private textRejected = false;
   private usage: NonNullable<SpanFields["usage"]> = {};
   private traceEventCount = 0;
+  private failed = false;
+  private failureCode: number | undefined;
   private failureReason: string | undefined;
   private guardrailAction: string | undefined;
   private returnControl: unknown;
@@ -65,6 +68,12 @@ export class AgentStreamState implements StreamState {
   private errorFields: SpanFields | undefined;
   private citationCount = 0;
   private readonly outputFiles: unknown[] = [];
+  // Metadata bypasses the capture flags and the mask, so trace and file content is filtered here,
+  // and a mask that cannot see it withholds it.
+  private readonly contentCapture: AgentTraceCapture = {
+    input: captureEnabled("input") && !captureMaskedBeforeStart(),
+    output: captureEnabled("output") && !captureMaskedBeforeStart(),
+  };
 
   constructor(private readonly options: BedrockInstrumentationOptions) {}
 
@@ -90,6 +99,7 @@ export class AgentStreamState implements StreamState {
       if (this.retention.accept(piece)) this.text += piece;
       else this.textRejected = true;
     }
+
     const attribution = isRecord(chunk?.attribution) ? chunk.attribution : undefined;
     const citations = arrayValue<unknown>(attribution?.citations);
 
@@ -97,7 +107,9 @@ export class AgentStreamState implements StreamState {
 
     const files = isRecord(raw.files) ? arrayValue<JsonValue>(raw.files.files) : undefined;
 
-    const sanitizedFiles = files?.map(sanitizeOutputFile);
+    const sanitizedFiles = files?.map((file) =>
+      sanitizeOutputFile(file, this.contentCapture.output),
+    );
 
     if (sanitizedFiles && this.retention.accept(sanitizedFiles)) {
       this.outputFiles.push(...sanitizedFiles);
@@ -125,15 +137,35 @@ export class AgentStreamState implements StreamState {
     if (traceEvent) {
       this.traceEventCount += 1;
 
-      if (this.options.captureAgentTrace && this.retention.accept(traceEvent)) {
-        this.traces.push(traceEvent);
+      if (this.options.captureAgentTrace) {
+        const retained = projectTrace(traceEvent, TRACE_PART, this.contentCapture);
+
+        if (this.retention.accept(retained)) this.traces.push(retained);
+      }
+
+      // Summaries are read only where the API puts them, so content and unknown branches can't
+      // feed them.
+      const trace = isRecord(traceEvent.trace) ? traceEvent.trace : {};
+
+      for (const step of MODEL_INVOCATION_STEPS) {
+        const member = isRecord(trace[step]) ? trace[step] : undefined;
+        const output = isRecord(member?.modelInvocationOutput) ? member.modelInvocationOutput : {};
+        this.addUsage(isRecord(output.metadata) ? output.metadata.usage : undefined);
       }
-      collectTrace(traceEvent, (usage) => this.addUsage(usage));
-      const failure = findKey(traceEvent, "failureTrace");
 
-      if (isRecord(failure))
-        this.failureReason = stringValue(failure.failureReason) ?? this.failureReason;
-      const guardrail = findKey(traceEvent, "guardrailTrace");
+      const failure = trace.failureTrace;
+
+      if (isRecord(failure)) {
+        this.failed = true;
+        this.failureCode = numberValue(failure.failureCode) ?? this.failureCode;
+
+        // The reason is free-form text that can quote the prompt or a tool result.
+        if (this.contentCapture.input && this.contentCapture.output) {
+          this.failureReason = stringValue(failure.failureReason) ?? this.failureReason;
+        }
+      }
+
+      const guardrail = trace.guardrailTrace;
 
       if (isRecord(guardrail))
         this.guardrailAction = stringValue(guardrail.action) ?? this.guardrailAction;
@@ -150,6 +182,8 @@ export class AgentStreamState implements StreamState {
   finish(): SpanFields {
     const metadata = omitUndefined({
       trace_event_count: this.traceEventCount || undefined,
+      failure: this.failed ? true : undefined,
+      failure_code: this.failureCode,
       failure_reason: this.failureReason,
       guardrail_action: this.guardrailAction,
       return_control: this.sawReturnControl ? true : undefined,
@@ -195,13 +229,13 @@ export class AgentStreamState implements StreamState {
   }
 }
 
-function sanitizeOutputFile<T>(value: T): JsonValue {
+function sanitizeOutputFile<T>(value: T, keepName: boolean): JsonValue {
   const raw: unknown = value;
 
   if (!isRecord(raw)) return undefined;
 
   return omitUndefined({
-    name: stringValue(raw.name),
+    name: keepName ? stringValue(raw.name) : undefined,
     type: stringValue(raw.type),
   });
 }
@@ -426,6 +460,7 @@ export class FlowStreamState implements StreamState {
     if (inputRequest && this.retention.accept(inputRequest.content)) {
       this.outputs.push(inputRequest.content);
     }
+
     const completion = isRecord(raw.flowCompletionEvent) ? raw.flowCompletionEvent : undefined;
     this.completionReason = stringValue(completion?.completionReason) ?? this.completionReason;
   }
@@ -450,35 +485,135 @@ export function ragName(input: JsonRecord): string {
   return `retrieve_and_generate ${basenameArn(ragModel(input)) ?? "knowledge-base"}`;
 }
 
-function collectTrace<T>(value: T, onUsage: <TUsage>(usage: TUsage) => void): void {
-  const raw: unknown = value;
+type AgentTraceCapture = Record<"input" | "output", boolean>;
 
-  if (!isRecord(raw)) return;
-  const output = isRecord(raw.modelInvocationOutput) ? raw.modelInvocationOutput : undefined;
+/** Content kept as received when every listed capture is on; otherwise reduced to `withheld`. */
+class TraceContent {
+  constructor(
+    readonly needs: ReadonlyArray<keyof AgentTraceCapture>,
+    readonly withheld?: TraceShape,
+  ) {}
+}
 
-  if (output) {
-    const metadata = isRecord(output.metadata) ? output.metadata : undefined;
-    onUsage(metadata?.usage);
-  }
+const SCALAR = "scalar";
+
+// A record shape lists the only fields kept; a one-item array shape is a list of that shape.
+type TraceShape =
+  | typeof SCALAR
+  | TraceContent
+  | readonly [TraceShape]
+  | { readonly [field: string]: TraceShape };
+
+const TRACE_METADATA: TraceShape = {
+  usage: { inputTokens: SCALAR, outputTokens: SCALAR },
+  startTime: SCALAR,
+  endTime: SCALAR,
+  totalTimeMs: SCALAR,
+  operationTotalTimeMs: SCALAR,
+  clientRequestId: SCALAR,
+};
+
+// What the agent sends to a model, tool, or guardrail is input content; what comes back is output.
+const MODEL_INVOCATION_INPUT = new TraceContent(["input"], {
+  traceId: SCALAR,
+  type: SCALAR,
+  foundationModel: SCALAR,
+});
+
+const INVOCATION_INPUT = new TraceContent(["input"], { traceId: SCALAR, invocationType: SCALAR });
+
+const MODEL_INVOCATION_OUTPUT = new TraceContent(["output"], {
+  traceId: SCALAR,
+  metadata: TRACE_METADATA,
+});
+
+const OBSERVATION = new TraceContent(["output"], { traceId: SCALAR, type: SCALAR });
+
+const MODEL_INVOCATION_STEPS = [
+  "preProcessingTrace",
+  "orchestrationTrace",
+  "postProcessingTrace",
+  "routingClassifierTrace",
+];
+
+const PROCESSING_TRACE: TraceShape = {
+  modelInvocationInput: MODEL_INVOCATION_INPUT,
+  modelInvocationOutput: MODEL_INVOCATION_OUTPUT,
+};
+
+// Mirrors TracePart and the Trace union in the Bedrock Agent Runtime API; anything else, including
+// the AWS SDK's `$unknown` union member, is dropped because nothing says what it holds.
+const TRACE_PART: TraceShape = {
+  agentId: SCALAR,
+  agentAliasId: SCALAR,
+  agentVersion: SCALAR,
+  sessionId: SCALAR,
+  collaboratorName: SCALAR,
+  eventTime: SCALAR,
+  callerChain: [{ agentAliasArn: SCALAR }],
+  trace: {
+    preProcessingTrace: PROCESSING_TRACE,
+    postProcessingTrace: PROCESSING_TRACE,
+    orchestrationTrace: {
+      modelInvocationInput: MODEL_INVOCATION_INPUT,
+      invocationInput: INVOCATION_INPUT,
+      modelInvocationOutput: MODEL_INVOCATION_OUTPUT,
+      rationale: new TraceContent(["output"], { traceId: SCALAR }),
+      observation: OBSERVATION,
+    },
+    routingClassifierTrace: {
+      modelInvocationInput: MODEL_INVOCATION_INPUT,
+      invocationInput: INVOCATION_INPUT,
+      modelInvocationOutput: MODEL_INVOCATION_OUTPUT,
+      observation: OBSERVATION,
+    },
+    customOrchestrationTrace: { traceId: SCALAR, event: new TraceContent(["output"]) },
+    guardrailTrace: {
+      traceId: SCALAR,
+      action: SCALAR,
+      metadata: TRACE_METADATA,
+      inputAssessments: new TraceContent(["input"]),
+      outputAssessments: new TraceContent(["output"]),
+    },
+    failureTrace: {
+      traceId: SCALAR,
+      failureCode: SCALAR,
+      metadata: TRACE_METADATA,
+      // Free-form text that can quote either the prompt or a tool result.
+      failureReason: new TraceContent(["input", "output"]),
+    },
+  },
+};
 
-  for (const child of Object.values(raw)) {
-    if (isRecord(child)) collectTrace(child, onUsage);
-    else if (Array.isArray(child)) child.forEach((item) => collectTrace(item, onUsage));
-  }
-}
+function projectTrace(value: JsonValue, shape: TraceShape, capture: AgentTraceCapture): JsonValue {
+  if (shape instanceof TraceContent) {
+    if (shape.needs.every((kind) => capture[kind])) return value ?? undefined;
 
-function findKey<T>(value: T, key: string): JsonValue {
-  const raw: unknown = value;
+    return shape.withheld === undefined ? undefined : projectTrace(value, shape.withheld, capture);
+  }
 
-  if (!isRecord(raw)) return undefined;
+  if (shape === SCALAR) {
+    const scalar = ["string", "number", "boolean"].includes(typeof value) || value instanceof Date;
 
-  if (raw[key] !== undefined) return raw[key];
+    return scalar ? value : undefined;
+  }
 
-  for (const child of Object.values(raw)) {
-    const found = findKey(child, key);
+  if (Array.isArray(shape)) {
+    if (!Array.isArray(value)) return undefined;
 
-    if (found !== undefined) return found;
+    return value
+      .map((item) => projectTrace(item, shape[0], capture))
+      .filter((item) => item !== undefined);
   }
 
-  return undefined;
+  if (!isRecord(value)) return undefined;
+
+  return omitUndefined(
+    Object.fromEntries(
+      Object.entries(shape).map(([field, child]) => [
+        field,
+        projectTrace(value[field], child, capture),
+      ]),
+    ),
+  );
 }

package/src/internal.ts

@@ -182,7 +182,7 @@ function isCut(fields: SpanFields | undefined): boolean {
 
 // Cores older than captureMasked expose no mask policy before a span starts; treat them as
 // masked, like LEGACY_CAPTURE_POLICY. A namespace read keeps them loadable without the export.
-function captureMaskedBeforeStart(): boolean {
+export function captureMaskedBeforeStart(): boolean {
   const read = (telemetrySdk as { captureMasked?: () => boolean }).captureMasked;
 
   try {
@@ -351,7 +351,7 @@ export function parseJson<T>(value: T): JsonValue {
 }
 
 // Retention uses fixed bounds; the core applies maxAttributeLength after the mask.
-export const STREAM_CAPTURE_MAX_BYTES = 48 * 1024;
+const STREAM_CAPTURE_MAX_BYTES = 48 * 1024;
 export const STREAM_CAPTURE_MAX_ITEMS = 1_000;
 const CAPTURE_TRUNCATED_ATTRIBUTE = "telemetry.dev.capture.truncated";
 

package/src/invoke-model.ts

@@ -10,6 +10,7 @@ import {
   omitUndefined,
   parseJson,
   PROVIDER,
+  STREAM_CAPTURE_MAX_ITEMS,
   stringValue,
   StreamRetention,
   type JsonRecord,
@@ -165,12 +166,7 @@ function nativeResponseFields(
       {
         output: body,
         finishReason: stringValue(body.stop_reason),
-        usage: omitUndefined({
-          inputTokens: numberValue(usage.input_tokens),
-          outputTokens: numberValue(usage.output_tokens),
-          cacheReadInputTokens: numberValue(usage.cache_read_input_tokens),
-          cacheCreationInputTokens: numberValue(usage.cache_creation_input_tokens),
-        }),
+        ...claudeUsageFields(usage),
       },
       metadata,
     );
@@ -216,6 +212,39 @@ function nativeResponseFields(
   return withHeaderUsage({ output: body }, metadata);
 }
 
+const CLAUDE_USAGE_KEYS = {
+  inputTokens: "input_tokens",
+  outputTokens: "output_tokens",
+  cacheReadInputTokens: "cache_read_input_tokens",
+  cacheCreationInputTokens: "cache_creation_input_tokens",
+} as const;
+
+function claudeUsageFields(record: JsonRecord): SpanFields {
+  const usage: NonNullable<SpanFields["usage"]> = {};
+  for (const [key, source] of Object.entries(CLAUDE_USAGE_KEYS)) {
+    const value = numberValue(record[source]);
+    if (value !== undefined) usage[key as keyof typeof CLAUDE_USAGE_KEYS] = value;
+  }
+  const iterations = record.iterations;
+  if (!Array.isArray(iterations)) return { usage };
+
+  // Executor iterations are already in top-level usage; only compaction is additive.
+  const first = Math.max(0, iterations.length - STREAM_CAPTURE_MAX_ITEMS);
+  for (let index = first; index < iterations.length; index += 1) {
+    const entry = iterations[index];
+    if (!isRecord(entry) || entry.type !== "compaction") continue;
+    for (const [key, source] of Object.entries(CLAUDE_USAGE_KEYS)) {
+      const value = numberValue(entry[source]);
+      const target = key as keyof typeof CLAUDE_USAGE_KEYS;
+      if (value !== undefined) usage[target] = (usage[target] ?? 0) + value;
+    }
+  }
+  return {
+    usage,
+    ...(first > 0 ? { attributes: { "telemetry.dev.capture.truncated": true } } : {}),
+  };
+}
+
 function withHeaderUsage<T>(fields: SpanFields, metadata: T): SpanFields {
   if (fields.usage && Object.keys(fields.usage).length > 0) return fields;
 
@@ -256,6 +285,9 @@ export class InvokeModelStreamState implements StreamState {
   private readonly chunks: unknown[] = [];
   private text = "";
   private usage: SpanFields["usage"];
+  private finalUsage: NonNullable<SpanFields["usage"]> = {};
+  private startUsage: NonNullable<SpanFields["usage"]> = {};
+  private compaction: SpanFields = {};
   private finishReason: string | undefined;
   private error: Error | undefined;
   private errorFields: SpanFields | undefined;
@@ -290,6 +322,28 @@ export class InvokeModelStreamState implements StreamState {
       const messageUsage = isRecord(message?.usage) ? message?.usage : undefined;
       const usage = isRecord(parsed.usage) ? parsed.usage : undefined;
 
+      for (const [key, source] of Object.entries(CLAUDE_USAGE_KEYS)) {
+        const value = numberValue(usage?.[source]);
+        if (value !== undefined) this.finalUsage[key as keyof typeof CLAUDE_USAGE_KEYS] = value;
+      }
+      // message_start counts exclude compaction like message_delta counts; metrics may not.
+      this.startUsage = omitUndefined({
+        ...this.startUsage,
+        inputTokens: numberValue(messageUsage?.input_tokens) ?? this.startUsage.inputTokens,
+        outputTokens: numberValue(messageUsage?.output_tokens) ?? this.startUsage.outputTokens,
+        cacheReadInputTokens:
+          numberValue(messageUsage?.cache_read_input_tokens) ??
+          this.startUsage.cacheReadInputTokens,
+        cacheCreationInputTokens:
+          numberValue(messageUsage?.cache_creation_input_tokens) ??
+          this.startUsage.cacheCreationInputTokens,
+      });
+      // An interrupted stream keeps message_start's iterations until a cumulative delta replaces them.
+      const iterations = Array.isArray(usage?.iterations)
+        ? usage.iterations
+        : messageUsage?.iterations;
+      if (Array.isArray(iterations)) this.compaction = claudeUsageFields({ iterations });
+
       const metrics = isRecord(parsed["amazon-bedrock-invocationMetrics"])
         ? parsed["amazon-bedrock-invocationMetrics"]
         : undefined;
@@ -314,7 +368,7 @@ export class InvokeModelStreamState implements StreamState {
       this.usage = omitUndefined({
         ...this.usage,
         inputTokens:
-          numberValue(messageUsage?.input_tokens) ??
+          numberValue(usage?.input_tokens) ??
           numberValue(metrics?.inputTokenCount) ??
           numberValue(parsed.inputTextTokenCount) ??
           this.usage?.inputTokens,
@@ -324,10 +378,9 @@ export class InvokeModelStreamState implements StreamState {
           numberValue(parsed.totalOutputTextTokenCount) ??
           this.usage?.outputTokens,
         cacheReadInputTokens:
-          numberValue(messageUsage?.cache_read_input_tokens) ?? this.usage?.cacheReadInputTokens,
+          numberValue(usage?.cache_read_input_tokens) ?? this.usage?.cacheReadInputTokens,
         cacheCreationInputTokens:
-          numberValue(messageUsage?.cache_creation_input_tokens) ??
-          this.usage?.cacheCreationInputTokens,
+          numberValue(usage?.cache_creation_input_tokens) ?? this.usage?.cacheCreationInputTokens,
       });
     }
 
@@ -335,6 +388,18 @@ export class InvokeModelStreamState implements StreamState {
   }
 
   finish(): SpanFields {
+    const usage = { ...this.startUsage, ...this.usage, ...this.finalUsage };
+    for (const key of Object.keys(CLAUDE_USAGE_KEYS) as (keyof typeof CLAUDE_USAGE_KEYS)[]) {
+      const value = this.compaction.usage?.[key];
+      const baseKey = key === "outputTokens" ? "outputTokens" : "inputTokens";
+      // Invocation metrics have no documented compaction exclusion; never add to that base.
+      if (
+        value !== undefined &&
+        (this.finalUsage[baseKey] !== undefined || this.usage?.[baseKey] === undefined)
+      ) {
+        usage[key] = (usage[key] ?? 0) + value;
+      }
+    }
     return mergeFields(
       omitUndefined({
         output: this.text
@@ -347,7 +412,8 @@ export class InvokeModelStreamState implements StreamState {
             ]
           : this.chunks,
         finishReason: this.finishReason,
-        usage: this.usage,
+        usage,
+        attributes: this.compaction.attributes,
         error: this.error,
       }),
       this.errorFields,