ref: Use diagnostics channel for anthropic wrapper (#1679)

Migrates anthropic wrapper to use diagnostic channels.

Author: lforst

e2e/scenarios/anthropic-instrumentation/assertions.ts

@@ -330,6 +330,31 @@ export function defineAnthropicInstrumentationAssertions(options: {
       ).toBe("string");
     });
 
+    test(
+      "captures trace for client.messages.create().withResponse()",
+      testConfig,
+      () => {
+        const root = findLatestSpan(events, ROOT_NAME);
+        const operation = findLatestSpan(
+          events,
+          "anthropic-create-with-response-operation",
+        );
+        const span = findAnthropicSpan(events, operation?.span.id, [
+          "anthropic.messages.create",
+        ]);
+
+        expect(operation).toBeDefined();
+        expect(span).toBeDefined();
+        expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]);
+        expect(span?.row.metadata).toMatchObject({
+          provider: "anthropic",
+        });
+        expect(
+          typeof (span?.row.metadata as { model?: unknown } | undefined)?.model,
+        ).toBe("string");
+      },
+    );
+
     test("captures trace for sending an attachment", testConfig, () => {
       const root = findLatestSpan(events, ROOT_NAME);
       const operation = findLatestSpan(

e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs

@@ -51,6 +51,27 @@ async function runAnthropicInstrumentationScenario(
         });
       });
 
+      await runOperation(
+        "anthropic-create-with-response-operation",
+        "create-with-response",
+        async () => {
+          const response = await client.messages
+            .create({
+              model: ANTHROPIC_MODEL,
+              max_tokens: 16,
+              temperature: 0,
+              messages: [
+                {
+                  role: "user",
+                  content: "Reply with exactly WITH_RESPONSE.",
+                },
+              ],
+            })
+            .withResponse();
+          void response.data;
+        },
+      );
+
       await runOperation(
         "anthropic-attachment-operation",
         "attachment",

js/src/instrumentation/core/channel-definitions.ts

@@ -13,7 +13,7 @@ export type ChannelKind = "async" | "sync-stream";
 type ChannelTypeInfo<
   TArgs extends EventArguments,
   TResult,
-  TExtra extends object = Record<string, never>,
+  TExtra extends object = Record<string, unknown>,
   TChunk = never,
   TKind extends ChannelKind = "async",
 > = {
@@ -27,7 +27,7 @@ type ChannelTypeInfo<
 export type ChannelSpec<
   TArgs extends EventArguments,
   TResult,
-  TExtra extends object = Record<string, never>,
+  TExtra extends object = Record<string, unknown>,
   TChunk = never,
   TKind extends ChannelKind = "async",
 > = ChannelTypeInfo<TArgs, TResult, TExtra, TChunk, TKind> & {
@@ -129,10 +129,10 @@ type BaseTypedChannel<TSpec extends AnyChannelSpec> = TSpec & {
 
 export type TypedAsyncChannel<TSpec extends AnyAsyncChannelSpec> =
   BaseTypedChannel<TSpec> & {
-    tracePromise<TResult extends ResultOf<TSpec>>(
-      fn: () => Promise<TResult>,
+    tracePromise<TReturn extends PromiseLike<ResultOf<TSpec>>>(
+      fn: () => TReturn,
       context: StartOf<TSpec>,
-    ): Promise<TResult>;
+    ): TReturn;
   };
 
 export type TypedSyncStreamChannel<TSpec extends AnySyncStreamChannelSpec> =
@@ -152,7 +152,7 @@ type ChannelSpecMap = Record<string, AnyChannelSpec>;
 export function channel<
   TArgs extends EventArguments,
   TResult,
-  TExtra extends object = Record<string, never>,
+  TExtra extends object = Record<string, unknown>,
   TChunk = never,
 >(spec: {
   channelName: string;
@@ -161,7 +161,7 @@ export function channel<
 export function channel<
   TArgs extends EventArguments,
   TResult,
-  TExtra extends object = Record<string, never>,
+  TExtra extends object = Record<string, unknown>,
   TChunk = never,
 >(spec: {
   channelName: string;
@@ -208,15 +208,15 @@ export function defineChannels<T extends ChannelSpecMap>(
           {
             ...asyncSpec,
             tracingChannel,
-            tracePromise: <TResult>(
-              fn: () => Promise<TResult>,
+            tracePromise: <TReturn extends Promise<ResultOf<typeof asyncSpec>>>(
+              fn: () => TReturn,
               context: StartOf<AnyAsyncChannelSpec>,
             ) =>
               tracingChannel().tracePromise(
                 fn,
                 // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
                 context as ChannelMessage<AnyAsyncChannelSpec>,
-              ),
+              ) as TReturn,
           } as AnyAsyncChannel,
         ];
       }

js/src/instrumentation/core/channel-tracing.ts

@@ -37,7 +37,7 @@ type SpanState = {
 export type AsyncChannelSpanConfig<TChannel extends AnyAsyncChannel> =
   ChannelConfig & {
     extractInput: (
-      args: [...ArgsOf<TChannel>],
+      args: [...ArgsOf<TChannel>, ...any[]],
       event: StartOf<TChannel>,
       span: Span,
     ) => {
@@ -67,7 +67,7 @@ type StreamingResult<TChannel extends AnyAsyncChannel> = Exclude<
 export type StreamingChannelSpanConfig<TChannel extends AnyAsyncChannel> =
   ChannelConfig & {
     extractInput: (
-      args: [...ArgsOf<TChannel>],
+      args: [...ArgsOf<TChannel>, ...any[]],
       event: StartOf<TChannel>,
       span: Span,
     ) => {
@@ -109,7 +109,7 @@ export type StreamingChannelSpanConfig<TChannel extends AnyAsyncChannel> =
 export type SyncStreamChannelSpanConfig<TChannel extends AnySyncStreamChannel> =
   ChannelConfig & {
     extractInput: (
-      args: [...ArgsOf<TChannel>],
+      args: [...ArgsOf<TChannel>, ...any[]],
       event: StartOf<TChannel>,
       span: Span,
     ) => {
@@ -166,7 +166,7 @@ function startSpanForEvent<
 >(
   config: ChannelConfig & {
     extractInput: (
-      args: [...ArgsOf<TChannel>],
+      args: [...ArgsOf<TChannel>, ...any[]],
       event: StartOf<TChannel>,
       span: Span,
     ) => {
@@ -210,7 +210,7 @@ function ensureSpanStateForEvent<
   states: WeakMap<object, SpanState>,
   config: ChannelConfig & {
     extractInput: (
-      args: [...ArgsOf<TChannel>],
+      args: [...ArgsOf<TChannel>, ...any[]],
       event: StartOf<TChannel>,
       span: Span,
     ) => {
@@ -239,7 +239,7 @@ function bindCurrentSpanStoreToStart<
   states: WeakMap<object, SpanState>,
   config: ChannelConfig & {
     extractInput: (
-      args: [...ArgsOf<TChannel>],
+      args: [...ArgsOf<TChannel>, ...any[]],
       event: StartOf<TChannel>,
       span: Span,
     ) => {

js/src/instrumentation/core/types.ts

@@ -91,22 +91,25 @@ export interface ErrorEvent extends BaseContext {
 export interface TypedStartEvent<
   TArguments extends EventArguments = unknown[],
 > extends BaseContext {
-  arguments: [...TArguments];
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  arguments: [...TArguments, ...any[]];
 }
 
 export interface TypedEndEvent<
   TResult = unknown,
   TArguments extends EventArguments = unknown[],
 > extends BaseContext {
   result: TResult;
-  arguments?: [...TArguments];
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  arguments?: [...TArguments, ...any[]];
 }
 
 export interface TypedErrorEvent<
   TArguments extends EventArguments = unknown[],
 > extends BaseContext {
   error: Error;
-  arguments?: [...TArguments];
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  arguments?: [...TArguments, ...any[]];
 }
 
 // eslint-disable-next-line @typescript-eslint/no-empty-object-type

js/src/instrumentation/plugins/anthropic-channels.ts

@@ -11,7 +11,7 @@ export const anthropicChannels = defineChannels("@anthropic-ai/sdk", {
   messagesCreate: channel<
     [AnthropicCreateParams],
     AnthropicResult,
-    Record<string, never>,
+    Record<string, unknown>,
     AnthropicStreamEvent
   >({
     channelName: "messages.create",
@@ -20,7 +20,7 @@ export const anthropicChannels = defineChannels("@anthropic-ai/sdk", {
   betaMessagesCreate: channel<
     [AnthropicCreateParams],
     AnthropicResult,
-    Record<string, never>,
+    Record<string, unknown>,
     AnthropicStreamEvent
   >({
     channelName: "beta.messages.create",

js/src/instrumentation/plugins/anthropic-plugin.ts

@@ -2,7 +2,7 @@ import { BasePlugin } from "../core";
 import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing";
 import { Attachment } from "../../logger";
 import { SpanTypeAttribute, isObject } from "../../../util/index";
-import { getCurrentUnixTimestamp } from "../../util";
+import { filterFrom, getCurrentUnixTimestamp } from "../../util";
 import { finalizeAnthropicTokens } from "../../wrappers/anthropic-tokens-util";
 import { anthropicChannels } from "./anthropic-channels";
 import type {
@@ -411,19 +411,3 @@ function coalesceInput(
   }
   return input;
 }
-
-/**
- * Filter out specified fields from an object.
- */
-function filterFrom(
-  obj: Record<string, unknown>,
-  fieldsToRemove: string[],
-): Record<string, unknown> {
-  const result: Record<string, unknown> = {};
-  for (const [key, value] of Object.entries(obj)) {
-    if (!fieldsToRemove.includes(key)) {
-      result[key] = value;
-    }
-  }
-  return result;
-}

js/src/instrumentation/plugins/openrouter-channels.ts

@@ -23,7 +23,7 @@ export const openRouterChannels = defineChannels("@openrouter/sdk", {
   chatSend: channel<
     [OpenRouterChatCreateParams],
     OpenRouterChatResult,
-    Record<string, never>,
+    Record<string, unknown>,
     OpenRouterChatCompletionChunk
   >({
     channelName: "chat.send",
@@ -41,7 +41,7 @@ export const openRouterChannels = defineChannels("@openrouter/sdk", {
   betaResponsesSend: channel<
     [OpenRouterResponsesCreateParams],
     OpenRouterResponsesResult,
-    Record<string, never>,
+    Record<string, unknown>,
     OpenRouterResponseStreamEvent
   >({
     channelName: "beta.responses.send",

js/src/isomorph.ts

@@ -123,12 +123,12 @@ export interface IsoTracingChannel<
     ...args: Parameters<F>
   ): ReturnType<F>;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  tracePromise<F extends (...args: any[]) => any>(
+  tracePromise<F extends (...args: any[]) => PromiseLike<any>>(
     fn: F,
     message?: M,
     thisArg?: ThisParameterType<F>,
     ...args: Parameters<F>
-  ): Promise<Awaited<ReturnType<F>>>;
+  ): ReturnType<F>;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   traceCallback<F extends (...args: any[]) => any>(
     fn: F,
@@ -203,13 +203,14 @@ class DefaultTracingChannel<M> implements IsoTracingChannel<M> {
     return fn.apply(thisArg, args);
   }
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  tracePromise<F extends (...args: any[]) => any>(
+  tracePromise<F extends (...args: any[]) => PromiseLike<any>>(
     fn: F,
     _message?: M,
     thisArg?: ThisParameterType<F>,
     ...args: Parameters<F>
-  ): Promise<Awaited<ReturnType<F>>> {
-    return Promise.resolve(fn.apply(thisArg, args));
+  ): ReturnType<F> {
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any
+    return fn.apply(thisArg, args) as any;
   }
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   traceCallback<F extends (...args: any[]) => any>(

js/src/typed-instrumentation-helpers.ts

@@ -0,0 +1,25 @@
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export type AnyFn = (this: any, ...args: any[]) => any;
+
+export type TypedApplyProxyHandler<TTarget extends AnyFn> = Omit<
+  ProxyHandler<TTarget>,
+  "apply"
+> & {
+  apply: (
+    target: TTarget,
+    thisArg: ThisParameterType<TTarget>,
+    argArray: Parameters<TTarget>,
+  ) => ReturnType<TTarget>;
+};
+
+/**
+ * Literally `Proxy` that has a typed `apply` based on the `target`.
+ */
+export const TypedApplyProxy: new <TTarget extends AnyFn>(
+  target: TTarget,
+  handler: TypedApplyProxyHandler<TTarget>,
+  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+) => TTarget = Proxy as unknown as new <TTarget extends AnyFn>(
+  target: TTarget,
+  handler: TypedApplyProxyHandler<TTarget>,
+) => TTarget;