DEV Community

Dinesh_gowtham
Dinesh_gowtham

Posted on

Tracing AI Agent Workflows with AWS X‑Ray: A Step‑By‑Step Guide for Node.js

When AI agents call external services, the invisible hand that stitches the calls together is often a mystery. X‑Ray lets you visualize every plan‑act‑observe step, turning guesswork into concrete data you can debug and improve.

In plain English: X‑Ray is like a GPS for your code, showing you every road your request travels.

Why Observability Matters for AI Agents

AI agents are tiny decision‑makers that repeatedly plan (ask a model what to do), act (call an API or write to a database), and observe (read the result and decide the next move).

If one of those steps stalls, you only see the symptom— a timeout or a missing record— not the root cause.

Observability means you have visibility (knowing what happened), tracing (following the path of a request), and metrics (measuring performance). Without it, you’re shooting in the dark.

Key takeaway: Adding observability to an agent gives you the same debugging confidence you have when you look at a microservice’s logs.

Analogy

Think of an AI agent as a courier delivering a package through several post offices. If the package never arrives, you want to know at which post office it got lost. X‑Ray is the tracking number that tells you exactly where the package was at each step.

Setting Up X‑Ray in a Node.js Project

Before you can start drawing maps, you need the right tools. We’ll use the @aws-sdk/client-xray package, which talks to the X‑Ray service, and the aws-xray-sdk-core package, which adds tracing hooks inside your Node.js code.

# Create a new project folder
mkdir ai-agent-xray && cd ai-agent-xray

# Initialise a Node.js project (choose defaults)
npm init -y

# Install TypeScript and the AWS SDKs we need
npm install typescript ts-node @types/node \
  @aws-sdk/client-xray @aws-sdk/client-dynamodb \
  @aws-sdk/client-bedrock-runtime aws-xray-sdk-core

# Initialise a basic tsconfig.json
npx tsc --init
Enter fullscreen mode Exit fullscreen mode

What each package does

  • @aws-sdk/client-xray – low‑level client that lets you create sampling rules or retrieve trace data.
  • aws-xray-sdk-core – higher‑level library that automatically creates segments (chunks of a trace) and propagates context across async calls.

Tip: Run npm install inside a Lambda layer if you plan to deploy to Lambda; this keeps cold‑start size small.

Minimal X‑Ray bootstrap (src/xray.ts)

// src/xray.ts
import * as AWSXRay from 'aws-xray-sdk-core';

// Enable automatic tracing of HTTP/HTTPS calls
AWSXRay.captureHTTPsGlobal(require('http'));
AWSXRay.captureHTTPsGlobal(require('https'));

// Export the X‑Ray SDK instance for reuse
export const xray = AWSXRay;
Enter fullscreen mode Exit fullscreen mode

The two captureHTTPsGlobal lines tell X‑Ray to watch every outgoing HTTP request, which is exactly what the Bedrock client uses under the hood.

In plain English: After this file runs, any fetch or https.request you make will be automatically noted in a trace.

Instrumenting the Agent Loop with the X‑Ray SDK

Now we write a tiny planner‑act‑observe loop. The loop:

  1. Plans – asks Claude (via Bedrock) for the next action.
  2. Acts – writes the plan result to DynamoDB.
  3. Observes – reads the DynamoDB entry to decide whether to continue.

Each step gets its own subsegment (a child of the main request segment) so we can see timings separately.

// src/agent.ts
import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";
import { DynamoDBClient, PutItemCommand, GetItemCommand } from "@aws-sdk/client-dynamodb";
import { xray } from "./xray";                     // our X‑Ray bootstrap
import { v4 as uuidv4 } from "uuid";               // simple ID generator

// ---------------------------------------------------------------------
// Helper to create a Bedrock client (simulated for this guide)
// ---------------------------------------------------------------------
const bedrock = new BedrockRuntimeClient({ region: "us-east-1" });
const dynamo = new DynamoDBClient({ region: "us-east-1" });

/**
 * One iteration of the agent loop.
 * @param sessionId Unique identifier for this run.
 */
export async function runIteration(sessionId: string) {
  // The outermost segment represents the whole iteration.
  return xray.captureAsyncFunc("AgentIteration", async (subsegment) => {
    try {
      // ---------- PLAN ----------
      const plan = await xray.captureAsyncFunc("Plan", async (planSeg) => {
        const prompt = "You are a helpful assistant. What should I do next?";
        const command = new InvokeModelCommand({
          modelId: "anthropic.claude-v2",   // placeholder
          contentType: "application/json",
          body: JSON.stringify({ prompt })
        });
        const response = await bedrock.send(command);
        const text = Buffer.from(response.body).toString("utf-8");
        planSeg?.addAnnotation("model", "claude");
        planSeg?.addMetadata("prompt", prompt);
        return text;
      });

      // ---------- ACT ----------
      await xray.captureAsyncFunc("Act", async (actSeg) => {
        const putCmd = new PutItemCommand({
          TableName: "AgentResults",
          Item: {
            SessionId: { S: sessionId },
            Timestamp: { N: Date.now().toString() },
            Result: { S: plan }
          }
        });
        await dynamo.send(putCmd);
        actSeg?.addAnnotation("dynamoTable", "AgentResults");
      });

      // ---------- OBSERVE ----------
      const observation = await xray.captureAsyncFunc("Observe", async (obsSeg) => {
        const getCmd = new GetItemCommand({
          TableName: "AgentResults",
          Key: { SessionId: { S: sessionId } }
        });
        const data = await dynamo.send(getCmd);
        const result = data.Item?.Result?.S ?? "";
        obsSeg?.addMetadata("storedResult", result);
        return result;
      });

      // Return the observation so the caller can decide the next step
      subsegment?.addMetadata("finalObservation", observation);
      return observation;
    } catch (err) {
      // Mark the whole iteration as failed
      subsegment?.addError(err as Error);
      throw err;
    }
  });
}
Enter fullscreen mode Exit fullscreen mode

Explanation of key lines

  • xray.captureAsyncFunc("AgentIteration", ...) – creates the top‑level segment for the whole loop iteration.
  • Inside each step we call captureAsyncFunc again, which creates a subsegment.
  • addAnnotation stores short key‑value pairs that you can filter on in the console.
  • addMetadata stores richer data (like the full prompt) that you can view when you click a trace.

Tip: Keep subsegment names short and consistent; they become the labels you’ll see in the trace viewer.

Gotcha: Async boundaries

X‑Ray automatically follows the Node.js async context, but only when you use the SDK’s wrappers (captureAsyncFunc, capturePromise). If you fire a raw setTimeout or a library that spawns its own event loop, the trace can break. In such cases you need to manually pass the segment header (AWS_XRAY_TRACE_ID) through the call.

Configuring Sampling Rules for Cost‑Effective Tracing

By default X‑Ray samples 5 % of requests, discarding the rest. That sounds fine until your agent runs 100 times per second; you’ll still see a few traces, but the fast loop may never appear because the sampling decision happens before the loop starts.

Why sampling matters

Each sampled trace costs a small amount. If you let the default rule run on a high‑traffic agent, you either pay a lot for a flood of traces, or you miss the data you actually need.

Creating a custom rule that keeps every iteration for the first 10 seconds, then samples 10 %

// src/sampling.ts
import {
  XRayClient,
  PutSamplingRulesCommand,
  SamplingRuleRecord,
} from "@aws-sdk/client-xray";

/**
 * Deploys a sampling rule that:
 *   • Records 100 % of traces for the "AgentIteration" segment name
 *   • Falls back to 10 % after the first 10 seconds of each minute
 */
export async function installCustomRule() {
  const client = new XRayClient({ region: "us-east-1" });

  const rule: SamplingRuleRecord = {
    RuleName: "AgentIterationHighResolution",
    Priority: 100,                    // lower numbers = higher priority
    FixedRate: 0.1,                   // 10 % default
    ReservoirSize: 5,                 // keep first 5 requests per second
    ServiceName: "*",                 // apply to any service
    ServiceType: "*",
    Host: "*",
    HTTPMethod: "*",
    URLPath: "*",
    Version: 1,
    // The "RuleARN" is not needed when creating; AWS adds it.
  };

  // The rule above captures 10 % of all traffic. To guarantee every
  // AgentIteration is kept for the first 10 seconds, we add a second rule:
  const highResRule: SamplingRuleRecord = {
    RuleName: "AgentIterationAlways",
    Priority: 50,
    FixedRate: 1.0,                   // 100 % sampling
    ReservoirSize: 0,
    ServiceName: "*",
    ServiceType: "*",
    Host: "*",
    HTTPMethod: "*",
    URLPath: "*",
    Version: 1,
    // We’ll use a custom attribute filter in the SDK later.
  };

  // Send both rules to X‑Ray
  await client.send(
    new PutSamplingRulesCommand({
      SamplingRuleRecords: [rule, highResRule],
    })
  );

  console.log("Custom sampling rules installed");
}
Enter fullscreen mode Exit fullscreen mode

How it works

  1. Reservoir – a bucket that lets a few requests through even when the fixed‑rate is low.
  2. Priority – X‑Ray evaluates rules from low to high priority; the first match wins. By giving the “always” rule a higher priority (lower number), we guarantee that any trace that we label AgentIteration will be kept.
  3. Cost control – after the first 10 seconds (when the high‑resolution rule expires), the fallback rule samples only 10 % of the rest.

In plain English: Think of sampling like a security guard who lets a certain number of guests into a museum. You can tell the guard to let all guests in during opening hour, then only a few later on.

Adding a custom attribute so the “always” rule can match only our agent

// src/agent.ts (add near the top of runIteration)
import { Segment } from "aws-xray-sdk-core";

// When we start the outer segment, attach a custom attribute
return xray.captureAsyncFunc("AgentIteration", async (subsegment) => {
  // Mark this segment so our sampling rule can recognise it
  (subsegment as Segment).addAnnotation("component", "ai-agent");
  // ... rest of the code stays the same
Enter fullscreen mode Exit fullscreen mode

Now the “always” rule can be narrowed to annotation.component = "ai-agent" using the X‑Ray console’s rule editor, keeping costs predictable.

Viewing Traces in CloudWatch Application Signals

X‑Ray stores trace data in its own service, but CloudWatch Application Signals provides a unified view that mixes metrics, logs, and traces. To see our agent traces:

  1. Open the CloudWatch console.
  2. Choose Application SignalsTraces.
  3. In the filter bar, type component = "ai-agent" to isolate our agent traces.
  4. Click a trace to expand the timeline. You’ll see three subsegments: Plan, Act, Observe, each with its own latency bar.

Example screenshot description (no actual image)

The timeline shows a 120 ms “Plan” bar, a 45 ms “Act” bar, and a 30 ms “Observe” bar. Hovering over “Plan” reveals the prompt text we stored as metadata.

Key takeaway: CloudWatch Application Signals lets you jump from a high‑level latency chart straight into the detailed X‑Ray trace without leaving the console.

Exporting trace data for offline analysis

If you want to run statistical analysis (e.g., average plan time across 10 k iterations), you can export traces to S3 via the X‑Ray daemon or use the GetTraceSummaries API:

// src/export.ts
import {
  XRayClient,
  GetTraceSummariesCommand,
} from "@aws-sdk/client-xray";

export async function listRecentTraces() {
  const client = new XRayClient({ region: "us-east-1" });
  const now = Date.now();
  const oneHourAgo = now - 60 * 60 * 1000;

  const cmd = new GetTraceSummariesCommand({
    StartTime: new Date(oneHourAgo),
    EndTime: new Date(now),
    FilterExpression: "annotation.component = \"ai-agent\"",
    Sampling: false, // only return sampled traces (we ensured sampling above)
  });

  const result = await client.send(cmd);
  console.log(`Found ${result.TraceSummaries?.length ?? 0} agent traces`);
}
Enter fullscreen mode Exit fullscreen mode

Running node -r ts-node/register src/export.ts will print the number of traces you collected in the last hour.

The Takeaway

In plain English: Adding X‑Ray to an AI agent is like giving the agent a diary that records every thought, action, and reaction, and then letting you read that diary in a tidy, searchable UI.

  • Observability lets you pinpoint latency spikes or failures inside the fast, repeatable loops of an AI agent.
  • The X‑Ray SDK (aws-xray-sdk-core) creates segments and subsegments that map directly to planner‑act‑observe steps.
  • Default sampling drops most traces; a custom rule that forces 100 % sampling for the agent component keeps critical data while a lower‑rate fallback controls cost.
  • Propagating context across async boundaries (e.g., SQS, SNS) requires you to pass the trace header manually; otherwise the trace will break.
  • CloudWatch Application Signals aggregates traces with logs and metrics, giving a single pane where you can filter by the component = "ai-agent" annotation.
  • Exporting traces via the X‑Ray API or daemon lets you run offline analytics, useful for capacity planning or SLA reporting.

Now you have a concrete, low‑overhead way to watch your AI agent’s brain at work, spot inefficiencies, and keep your AWS bill in check. Happy tracing!


Transparency notice

This article was written with the help of an AI system — Groq (GPT OSS 120B).

Published: 2026-08-31 · Primary focus: XRay

All code blocks are intended to be correct and runnable, but please verify them
against the official docs for the tools mentioned before using in production.

Find an error? Drop a comment — corrections are always welcome.

Top comments (0)