fix(observability): bound websocket trace spans

.changeset/bounded-websocket-trace-spans.md

@@ -0,0 +1,7 @@
+---
+"agents": patch
+"@cloudflare/ai-chat": patch
+"@cloudflare/think": patch
+---
+
+Bound custom spans to the hibernatable WebSocket invocation so cancellation and disconnects cannot leave agent, turn, inference, model, or result spans open. RPC and alarm turns retain full-stream span lifetimes.

docs/agents/observability.md

@@ -316,6 +316,14 @@ They never remain open while waiting for a human across invocations. Stream
 spans close on completion, cancellation, an in-band error, or early consumer
 return. Async-generator tools stay open until iteration ends.
 
+Hibernatable WebSocket turns are bounded differently because a cancel or
+disconnect can end the original Worker invocation before JavaScript stream
+finalizers run. Tracer lifetime policies automatically close interaction, turn,
+`invoke_agent`, `chat`, and result spans when a matching child opens or their
+callback hands off asynchronous work. Their context remains active so later
+spans keep the same parent IDs. RPC and alarm turns retain full-stream span
+lifetimes.
+
 Pass `storeMessages: true` to write full input/output message arrays (including
 tool-call parts) to `chat`; pass `storeTools: true` to write tool arguments and
 results to `execute_tool`:

packages/agents/src/index.ts

@@ -84,6 +84,7 @@ import {
 import { tracer } from "./observability/tracing/cloudflare";
 import {
   writeSpanAttributes,
+  type SpanLifetime,
   type TraceAttributes
 } from "./observability/tracing/tracer";
 import { DisposableStore } from "./core/events";
@@ -1863,19 +1864,22 @@ export class Agent<
     operation: string,
     storagePhase: string,
     attributes: TraceAttributes,
-    run: (update: (attributes: TraceAttributes) => void) => Promise<T>
+    run: (update: (attributes: TraceAttributes) => void) => Promise<T>,
+    lifetime?: SpanLifetime
   ): Promise<T>;
   private _withAgentSpan<T>(
     operation: string,
     storagePhase: string,
     attributes: TraceAttributes,
-    run: (update: (attributes: TraceAttributes) => void) => T
+    run: (update: (attributes: TraceAttributes) => void) => T,
+    lifetime?: SpanLifetime
   ): T;
   private _withAgentSpan<T>(
     operation: string,
     storagePhase: string,
     attributes: TraceAttributes,
-    run: (update: (attributes: TraceAttributes) => void) => T | Promise<T>
+    run: (update: (attributes: TraceAttributes) => void) => T | Promise<T>,
+    lifetime?: SpanLifetime
   ): T | Promise<T> {
     // The instance name is not always readable during construction: facets
     // restore it after construction and unnamed DOs receive it later.
@@ -1898,7 +1902,8 @@ export class Agent<
         ...attributes
       },
       (span) =>
-        run((finishAttributes) => writeSpanAttributes(span, finishAttributes))
+        run((finishAttributes) => writeSpanAttributes(span, finishAttributes)),
+      lifetime
     );
   }
 

packages/agents/src/observability/ai/v6/model.ts

@@ -20,7 +20,8 @@ export function wrapModel(
   wrapLanguageModel: AISDKV6WrapLanguageModel | undefined,
   model: unknown,
   parentOperation: string,
-  storeMessages: boolean
+  storeMessages: boolean,
+  finishStreamSpanAtHandoff: boolean
 ): unknown {
   if (!wrapLanguageModel) {
     return model;
@@ -77,8 +78,9 @@ export function wrapModel(
           storeMessages
         );
         // The provider call runs INSIDE the activation callback so its work
-        // (fetch subrequests, etc.) nests under the chat span; the span stays
-        // caller-owned because the stream outlives the callback.
+        // (fetch subrequests, etc.) nests under the chat span. Normal streams
+        // stay caller-owned; invocation-bounded WebSocket streams finish when
+        // the callback hands asynchronous provider work back to the runtime.
         return tracer.openSpan(
           span.name,
           span.attributes,
@@ -87,6 +89,9 @@ export function wrapModel(
             try {
               const startedAtMs = Date.now();
               const result = await doStream();
+              if (finishStreamSpanAtHandoff) {
+                return result;
+              }
               return finishWhenStreamCompletes(result, modelCall, {
                 aiGatewayLogId:
                   extractAIGatewayLogId(result) ?? aiGatewayLog.get(),
@@ -100,7 +105,8 @@ export function wrapModel(
               modelCall.fail(cause);
               throw cause;
             }
-          }
+          },
+          finishStreamSpanAtHandoff ? { finishOnAsyncHandoff: true } : undefined
         );
       }
     }

packages/agents/src/observability/ai/v6/wrap.ts

@@ -138,6 +138,8 @@ function createOperationWrapper(
   // is listening, so untraced calls never touch caller getters beyond that.
   if (isStreamOperation(operationName)) {
     return (params, ...args) => {
+      const finishAtHandoff = isWebSocketTurn(params);
+      const hasModelSpan = canWrapModel(wrapLanguageModel, params.model);
       return instrumentation.tracer.openSpan(
         operationSpanName(agentNameForCall(params)),
         {},
@@ -168,17 +170,30 @@ function createOperationWrapper(
               operationName,
               wrapLanguageModel,
               instrumentation.tracer,
-              storage
+              storage,
+              finishAtHandoff
             ),
             ...args
           );
-          const hasModelSpan = canWrapModel(wrapLanguageModel, params.model);
+
+          if (finishAtHandoff) {
+            if (!hasModelSpan) {
+              operationSpan.finish();
+            }
+            return result;
+          }
 
           return finishWhenStreamCompletes(result, operationSpan, {
             includeResponse: !hasModelSpan,
             startedAtMs: hasModelSpan ? undefined : startedAtMs
           });
-        }
+        },
+        finishAtHandoff && hasModelSpan
+          ? {
+              finishOnChild: (childName) =>
+                childName === "chat" || childName.startsWith("chat ")
+            }
+          : undefined
       );
     };
   }
@@ -250,7 +265,8 @@ function operationParamsForCall(
   operationName: AISDKV6OperationName,
   wrapLanguageModel: AISDKV6WrapLanguageModel | undefined,
   tracer: AgentTracer,
-  storage: ResolvedAISDKStorageOptions
+  storage: ResolvedAISDKStorageOptions,
+  finishStreamSpansAtHandoff = false
 ): AISDKV6CallParams {
   return {
     ...params,
@@ -264,13 +280,27 @@ function operationParamsForCall(
             wrapLanguageModel,
             params.model,
             operationName,
-            storage.storeMessages
+            storage.storeMessages,
+            finishStreamSpansAtHandoff
           )
         }
       : {})
   };
 }
 
+function isWebSocketTurn(params: AISDKV6CallParams): boolean {
+  const telemetry =
+    typeof params.experimental_telemetry === "object" &&
+    params.experimental_telemetry !== null
+      ? (params.experimental_telemetry as Record<string, unknown>)
+      : undefined;
+  const metadata =
+    typeof telemetry?.metadata === "object" && telemetry.metadata !== null
+      ? (telemetry.metadata as Record<string, unknown>)
+      : undefined;
+  return metadata?.["cloudflare.agents.turn.trigger"] === "ws-chat";
+}
+
 function canWrapModel(
   wrapLanguageModel: AISDKV6WrapLanguageModel | undefined,
   model: unknown

packages/agents/src/observability/tracing/tracer.ts

@@ -1,3 +1,5 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+
 /** Attribute values accepted by custom spans. */
 export type TraceAttributeValue = string | number | boolean | undefined;
 
@@ -19,6 +21,19 @@ export type SpanRuntime = {
   startActiveSpan<T>(name: string, run: (span: SpanWriter) => T): T;
 };
 
+/** Optional automatic lifetime policy for invocation-bounded spans. */
+export type SpanLifetime = {
+  /** Finish this span after a matching child has opened beneath it. */
+  readonly finishOnChild?: string | ((childName: string) => boolean);
+  /** Finish as soon as the callback hands asynchronous work back to its caller. */
+  readonly finishOnAsyncHandoff?: boolean;
+};
+
+type ActiveManagedSpan = {
+  readonly span: ManagedSpan;
+  readonly lifetime?: SpanLifetime;
+};
+
 /** AgentTracer seam used by integrations. */
 export type AgentTracer = {
   /**
@@ -32,7 +47,8 @@ export type AgentTracer = {
   withSpan<T>(
     name: string,
     attributes: TraceAttributes,
-    run: (span: AgentSpan) => MaybePromise<T>
+    run: (span: AgentSpan) => MaybePromise<T>,
+    lifetime?: SpanLifetime
   ): T | Promise<T>;
   /**
    * Activates a span and returns whatever `activate` returns (typically the
@@ -46,7 +62,8 @@ export type AgentTracer = {
   openSpan<T>(
     name: string,
     attributes: TraceAttributes,
-    activate: (span: AgentSpan) => T
+    activate: (span: AgentSpan) => T,
+    lifetime?: SpanLifetime
   ): T;
 };
 
@@ -73,37 +90,60 @@ export function createTracer(runtime: SpanRuntime): AgentTracer {
 }
 
 class RuntimeTracer implements AgentTracer {
+  private readonly activeSpan = new AsyncLocalStorage<ActiveManagedSpan>();
+
   constructor(private readonly runtime: SpanRuntime) {}
 
   withSpan<T>(
     name: string,
     attributes: TraceAttributes,
-    run: (span: AgentSpan) => MaybePromise<T>
+    run: (span: AgentSpan) => MaybePromise<T>,
+    lifetime?: SpanLifetime
   ): T | Promise<T> {
-    return this.activate(name, attributes, (span) => {
-      const result = run(span);
-      if (isPromiseLike(result)) {
-        return Promise.resolve(result)
-          .catch((cause: unknown) => {
-            span.fail(cause);
-            throw cause;
-          })
-          .finally(() => {
-            span.close();
-          });
-      }
+    return this.activate(
+      name,
+      attributes,
+      (span) => {
+        const result = run(span);
+        if (isPromiseLike(result)) {
+          if (lifetime?.finishOnAsyncHandoff) {
+            span.finish();
+          }
+          return Promise.resolve(result)
+            .catch((cause: unknown) => {
+              span.fail(cause);
+              throw cause;
+            })
+            .finally(() => {
+              span.close();
+            });
+        }
 
-      span.close();
-      return result;
-    });
+        span.close();
+        return result;
+      },
+      lifetime
+    );
   }
 
   openSpan<T>(
     name: string,
     attributes: TraceAttributes,
-    activate: (span: AgentSpan) => T
+    activate: (span: AgentSpan) => T,
+    lifetime?: SpanLifetime
   ): T {
-    return this.activate(name, attributes, activate);
+    return this.activate(
+      name,
+      attributes,
+      (span) => {
+        const result = activate(span);
+        if (lifetime?.finishOnAsyncHandoff && isPromiseLike(result)) {
+          span.finish();
+        }
+        return result;
+      },
+      lifetime
+    );
   }
 
   /**
@@ -114,14 +154,20 @@ class RuntimeTracer implements AgentTracer {
   private activate<T>(
     name: string,
     attributes: TraceAttributes,
-    body: (span: ManagedSpan) => T
+    body: (span: ManagedSpan) => T,
+    lifetime?: SpanLifetime
   ): T {
+    const parent = this.activeSpan.getStore();
     return this.runtime.startActiveSpan(name, (writer) => {
       setAttributes(writer, attributes);
       const span = new ManagedSpan(writer);
 
+      if (matchesChild(parent?.lifetime?.finishOnChild, name)) {
+        parent?.span.finish();
+      }
+
       try {
-        return body(span);
+        return this.activeSpan.run({ span, lifetime }, () => body(span));
       } catch (cause: unknown) {
         span.fail(cause);
         throw cause;
@@ -223,6 +269,15 @@ function setAttributes(span: SpanWriter, attributes: TraceAttributes): void {
   }
 }
 
+function matchesChild(
+  matcher: SpanLifetime["finishOnChild"],
+  childName: string
+): boolean {
+  return typeof matcher === "function"
+    ? matcher(childName)
+    : matcher === childName;
+}
+
 function isPromiseLike<T>(value: MaybePromise<T>): value is PromiseLike<T> {
   return (
     value !== null &&

packages/agents/src/tests/observability/ai-sdk-v6-integration.test.ts

@@ -241,6 +241,40 @@ describe("createAISDKV6Wrapper with the real AI SDK", () => {
     expect(tracing.rootSpans[0]?.ended).toBe(true);
   });
 
+  it("bounds WebSocket stream spans at handoff without losing their parent", async () => {
+    const tracing = new RecordingTracer();
+    const wrapped = createAISDKV6Wrapper(ai, { tracer: tracing });
+    const result = wrapped.streamText({
+      model: textStreamModel(),
+      prompt: "bounded",
+      experimental_telemetry: {
+        metadata: { "cloudflare.agents.turn.trigger": "ws-chat" }
+      }
+    });
+
+    const invokeSpan = tracing.rootSpans[0];
+    expect(invokeSpan?.ended).toBe(false);
+
+    const reader = result.fullStream.getReader();
+    let chatSpan = tracing.spans.find(
+      (span) => span.attributes["gen_ai.operation.name"] === "chat"
+    );
+    let readDone = false;
+    while (!chatSpan && !readDone) {
+      const read = await reader.read();
+      readDone = read.done;
+      chatSpan = tracing.spans.find(
+        (span) => span.attributes["gen_ai.operation.name"] === "chat"
+      );
+    }
+
+    expect(readDone).toBe(false);
+    expect(invokeSpan?.ended).toBe(true);
+    expect(chatSpan?.ended).toBe(true);
+    expect(chatSpan?.parent).toBe(invokeSpan);
+    await reader.cancel();
+  });
+
   it("traces tool execution with the SDK-provided tool call id", async () => {
     const tracing = new RecordingTracer();
     const wrapped = createAISDKV6Wrapper(ai, { tracer: tracing });