How do I trace an agent built with the Claude Agent SDK and deployed to AWS Lambda?

Rius traces a Claude Agent SDK agent on AWS Lambda with an observe root span, one manual generation around the query() loop, and a flush at the end of every invocation. Call init once at the top of your module with the heartbeat off, and name the agent with service_name. The result is one trace per task with tokens and cost, delivered before Lambda freezes the process.

How it works

The Claude Agent SDK drives the Claude Code runtime in a separate process, so no model client runs in your function and auto-instrumentation has nothing to hook. Coverage is manual and has two pieces:

• An observe wrapper with kind=SpanKind.AGENT makes each task the root span.
• A manual generation wraps the query() loop. You fill it from the run’s ResultMessage with token usage, the real model id and the collected answer. Rius computes cost from those, as a list-price estimate in US dollars.

Lambda adds two runtime concerns:

• Flushing. Rius exports spans from a background thread. Lambda freezes the process between invocations, so a frozen thread exports nothing and the exit hook may never run. Call client.flush() at the end of every invocation. flush() returns False on timeout or when the latest export was rejected, so True means your spans left the function.
• Heartbeat. The SDK sends a heartbeat for the whole life of a process, and a frozen process makes the instance flap between alive, stale and gone. Turn it off with heartbeat=False, or set RIUS_HEARTBEAT=false. The agent then shows as no heartbeat and lists no instances, but its traces arrive normally.

The claude CLI must be installed and authenticated where the function runs. The runtime uses whatever auth Claude Code has, and an exported ANTHROPIC_API_KEY takes precedence over a claude.ai login.

Set it up

1. Install the SDK and the framework in your function’s dependencies. The base package is enough, and no extra applies. See Claude Agent SDK for the full recipe.

   pip install glassflow-rius claude-agent-sdk

2. Call rius.init once at module level, outside the handler. Pass service_name with the name you want the agent to have, and heartbeat=False. Keep the returned client, because you flush it later.

3. Wrap the task function with observe, name and kind=SpanKind.AGENT. Open a generation inside it around query() and fill it from the ResultMessage.

4. Use a full model id such as claude-haiku-4-5, never an alias such as haiku, in both the CLI option and the generation. Cost is computed from the requested model, and an alias produces a generation with no cost. Keep provider="anthropic" on the generation and pass the cache counts as-is, so input tokens include cache reads and writes.

5. Flush in a finally block at the end of the Lambda handler, so spans are exported even when the task raises.

import asyncio
 
import rius
from claude_agent_sdk import (
    AssistantMessage, ClaudeAgentOptions, ResultMessage, TextBlock, query,
)
from rius import SpanKind
 
client = rius.init(api_key="ri_...", service_name="my-agent", heartbeat=False)
 
MODEL = "claude-haiku-4-5"
 
@rius.observe(name="handle-task", kind=SpanKind.AGENT)
async def handle_task(task: str) -> str:
    options = ClaudeAgentOptions(model=MODEL, max_turns=1)
    parts = []
    with rius.start_as_current_generation(
        "claude_agent_turn", model=MODEL, provider="anthropic", input=task
    ) as gen:
        async for message in query(prompt=task, options=options):
            if isinstance(message, AssistantMessage):
                parts.extend(
                    b.text for b in message.content if isinstance(b, TextBlock)
                )
            elif isinstance(message, ResultMessage):
                usage = message.usage or {}
                gen.set_usage(
                    input_tokens=usage.get("input_tokens"),
                    output_tokens=usage.get("output_tokens"),
                    cache_read_input_tokens=usage.get("cache_read_input_tokens"),
                    cache_write_input_tokens=usage.get("cache_creation_input_tokens"),
                )
                models = list((message.model_usage or {}).keys())
                if models:
                    gen.set_response_model(models[0])
        answer = "".join(parts)
        gen.set_output(answer)
    return answer
 
def lambda_handler(event, context):
    try:
        return {"answer": asyncio.run(handle_task(event["task"]))}
    finally:
        client.flush(timeout_millis=30_000)

The handler, its event["task"] key and lambda_handler are your own code. Flush only at the invocation boundary, as above, never inside code that runs mid-task. If your agent holds a conversation across invocations, scope each task with rius.session() and use the SDK’s session_id as the Rius session id.

What you get

• One trace per invocation: a root span named handle-task of kind AGENT, with a child generation that carries the model, tokens and cost.
• Cost per task, including cache reads and writes, from the token counts you set on the generation.
• The agent grouped under the service_name you chose. Set RIUS_AGENT_NAME if you want a different grouping name, and it defaults to the service name.
• Spans exported before Lambda freezes the process, with flush() telling you whether they left.
• Sessions, if you use rius.session(), that group a conversation’s tasks together.

The runtime’s inner work, such as tool calls, sub-turns and file operations, happens inside the Claude Code process and produces no spans. The generation records the turn’s totals only. Open one trace in the console after a test run: the generation should show a full model id, tokens and cost. For the batching and retry behavior behind the flush, see Reliability.

When something else fits

If you need each tool call from the Claude Code runtime as its own span, this setup does not give it, because only the turn’s totals are recorded. For other frameworks, where model calls are captured automatically, see Integrations. If you run a long-lived process instead of Lambda, you can keep the heartbeat on and rely on the exit hook in Python.

Related