Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/cloudflare/src/instrumentations/agents/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import type { AgentInternals } from './types';
* - **Conversation correlation** — sets the conversation id on the scope for each unit of agent
* work — chat turn or callable RPC call — so `gen_ai` spans created within it are correlated, for
* chat and plain agents alike. Defaults to the instance `name` and is rotated when the chat is
* cleared (`cf_agent_chat_clear`).
* cleared (the `message:clear` observability event).
*
* It only hooks the `agents` package internals and uses Sentry's tracing primitives. On Cloudflare
* Workers, prefer `instrumentAgentWithSentry`, which additionally instruments the Durable Object
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,16 @@
import { uuid4 } from '@sentry/core';
import { type AgentInternals, setAgentConversationId } from './types';

/**
* For chat agents (`AIChatAgent` from `@cloudflare/ai-chat`), sets the conversation id on the active
* scope for the duration of a chat turn. In the Agents model one agent instance is one conversation,
* so the instance `name` is the natural conversation id.
* For chat agents (`AIChatAgent` from `@cloudflare/ai-chat`), correlates each chat turn's AI spans
* with a conversation id on the active scope.
*
* In the Agents model one agent instance is one long-lived conversation, so the instance `name` is
* the base conversation id. When the user clears the chat, they expect a fresh conversation — but
* recreating the Durable Object for that would also drop the MCP/OAuth state stored per instance
* (GitHub/Sentry sign-in). To get a fresh conversation id *without* losing that state, we rotate an
* in-memory id on the instance when the SDK reports a cleared chat, and stamp that (falling back to
* the instance `name` before the first clear).
*
* The id itself is not attached to spans here — the SDK's `conversationIdIntegration` reads it off
* the scope at `spanStart` and stamps `gen_ai.conversation.id` onto the AI spans created inside the
Expand All @@ -14,7 +21,26 @@ import { type AgentInternals, setAgentConversationId } from './types';
* instead.
*/
export function instrumentChatAgentConversation(obj: AgentInternals): void {
// Rotate the conversation id when the chat is cleared. `_emit` is the central choke-point through
// which all `agents:*` observability events are published, and it already exists on the base
// `Agent` class — so this hook composes with the RPC instrumentation, which keys off the same
// surface. We shadow it with an own property so the original stays reachable on the prototype.
const originalEmit = obj._emit;

if (typeof originalEmit === 'function') {
Comment thread
JPeer264 marked this conversation as resolved.
obj._emit = new Proxy(originalEmit, {
apply(target, thisArg: AgentInternals, args: [string, Record<string, unknown>?]) {
if (args[0] === 'message:clear') {
thisArg.__sentryConversationId = uuid4();
}

return Reflect.apply(target, thisArg, args);
},
});
}

const original = obj.onChatMessage;

if (typeof original !== 'function') {
return;
}
Expand Down
13 changes: 11 additions & 2 deletions packages/cloudflare/src/instrumentations/agents/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,13 @@ export interface AgentInternals {
_ParentClass?: { name?: string };
/** The Agent instance name, which in the Agents model identifies the conversation/thread. */
name?: string;
/**
* Internal: the active conversation id, rotated when the chat is cleared (the `message:clear`
* observability event) so a reset chat groups as a fresh conversation while the instance (and its
* MCP/OAuth state) stays put. Set by `instrumentChatAgentConversation`; falls back to `name`
* before the first clear.
*/
__sentryConversationId?: string;
}

/** Reads best-effort agent identity attributes from the instance, tolerating missing internals. */
Expand All @@ -46,15 +53,17 @@ export function getAgentAttributes(instance: AgentInternals): Record<string, str
* Sets the agent instance's conversation id on the current scope for the duration of the
* surrounding unit of work (chat turn, callable RPC call). In the Agents model one instance is one
* conversation, so the instance `name` is the natural conversation id — for chat and plain agents
* alike, since plain agents run LLM calls too (e.g. inside `@callable()` methods).
* alike, since plain agents run LLM calls too (e.g. inside `@callable()` methods). Once the chat
* has been cleared, the rotated `__sentryConversationId` takes precedence so LLM calls from any
* unit of work group under the fresh conversation.
*
* `conversationIdIntegration` reads the id off the scope at `spanStart` and stamps
* `gen_ai.conversation.id` onto AI spans created within the unit of work, correlating its model
* and tool calls. Callers run inside a per-event forked scope (`wrapMethodWithSentry`), so the id
* does not leak into unrelated events.
*/
export function setAgentConversationId(instance: AgentInternals): void {
const conversationId = instance.name;
const conversationId = instance.__sentryConversationId ?? instance.name;

if (typeof conversationId === 'string' && conversationId) {
getCurrentScope().setConversationId(conversationId);
Expand Down
171 changes: 171 additions & 0 deletions packages/cloudflare/test/instrumentChatAgentConversation.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,171 @@
import { getCurrentScope } from '@sentry/core';
import { afterEach, describe, expect, it, vi } from 'vitest';
import { instrumentChatAgentConversation } from '../src/instrumentations/agents/instrumentChatAgentConversation';
import type { AgentInternals } from '../src/instrumentations/agents/types';

describe('instrumentChatAgentConversation', () => {
afterEach(() => {
vi.restoreAllMocks();
});

it('does not set conversation id when agent name is empty string', () => {
const setConversationIdSpy = vi.spyOn(getCurrentScope(), 'setConversationId').mockImplementation(() => {});

const obj: AgentInternals = {
name: '',
onChatMessage() {
return 'response';
},
};

instrumentChatAgentConversation(obj);

obj.onChatMessage!(() => {}, {});

expect(setConversationIdSpy).not.toHaveBeenCalled();
});

it('does not set conversation id when agent name is undefined', () => {
const setConversationIdSpy = vi.spyOn(getCurrentScope(), 'setConversationId').mockImplementation(() => {});

const obj: AgentInternals = {
onChatMessage() {
return 'response';
},
};

instrumentChatAgentConversation(obj);

obj.onChatMessage!(() => {}, {});

expect(setConversationIdSpy).not.toHaveBeenCalled();
});

it('sets conversation id when agent name is present', () => {
const setConversationId = vi.fn();
vi.spyOn(getCurrentScope(), 'setConversationId').mockImplementation(setConversationId);

const obj: AgentInternals = {
name: 'conversation-42',
onChatMessage() {
return 'response';
},
};

instrumentChatAgentConversation(obj);

obj.onChatMessage!(() => {}, {});

expect(setConversationId).toHaveBeenCalledWith('conversation-42');
});

it('forwards the return value from the original onChatMessage', () => {
const obj: AgentInternals = {
name: 'convo-1',
onChatMessage() {
return { output: 'hello' };
},
};

instrumentChatAgentConversation(obj);

const result = obj.onChatMessage!(() => {}, {});

expect(result).toEqual({ output: 'hello' });
});

it('leaves the agent untouched when onChatMessage is not defined', () => {
const obj: AgentInternals = { name: 'agent-1' };

instrumentChatAgentConversation(obj);

expect('onChatMessage' in obj).toBe(false);
});

it('uses the instance name as the conversation id before any clear', () => {
const setConversationId = vi.fn();
vi.spyOn(getCurrentScope(), 'setConversationId').mockImplementation(setConversationId);

const obj: AgentInternals = {
name: 'session-7',
_emit() {
return undefined;
},
onChatMessage() {
return 'response';
},
};

instrumentChatAgentConversation(obj);

obj.onChatMessage!(() => {}, {});

expect(setConversationId).toHaveBeenCalledWith('session-7');
});

it('rotates the conversation id on the message:clear observability event (fresh id, not the instance name)', () => {
const setConversationId = vi.fn();
vi.spyOn(getCurrentScope(), 'setConversationId').mockImplementation(setConversationId);

const obj: AgentInternals = {
name: 'session-7',
_emit() {
return undefined;
},
onChatMessage() {
return 'response';
},
};

instrumentChatAgentConversation(obj);

// Simulate the SDK emitting the chat-clear observability event.
obj._emit!('message:clear');

obj.onChatMessage!(() => {}, {});

expect(setConversationId).toHaveBeenCalledTimes(1);
const rotated = setConversationId.mock.calls[0]?.[0] as string;
expect(typeof rotated).toBe('string');
expect(rotated).not.toBe('session-7');
expect(obj.__sentryConversationId).toBe(rotated);
});

it('forwards message:clear to the original _emit', () => {
const received: Array<{ type: string; payload: unknown }> = [];
const obj: AgentInternals = {
name: 'session-7',
_emit(type: string, payload?: Record<string, unknown>) {
received.push({ type, payload });
return undefined;
},
onChatMessage() {
return 'response';
},
};

instrumentChatAgentConversation(obj);

obj._emit!('message:clear', { source: 'user' });
expect(received).toEqual([{ type: 'message:clear', payload: { source: 'user' } }]);
});

it('does not rotate the conversation id for other observability events', () => {
const obj: AgentInternals = {
name: 'session-7',
_emit() {
return undefined;
},
onChatMessage() {
return 'response';
},
};

instrumentChatAgentConversation(obj);

obj._emit!('message:request');
obj._emit!('rpc', { method: 'greet' });

expect(obj.__sentryConversationId).toBeUndefined();
});
});
Loading