assistant-cloud 0.1.43 → 0.2.1

package/src/instrumentMcpSampling.test.ts

@@ -0,0 +1,102 @@
+import { describe, expect, it, vi } from "vitest";
+import {
+  wrapSamplingHandler,
+  createSamplingCollector,
+  type McpSamplingRequest,
+} from "./instrumentMcpSampling";
+
+const request: McpSamplingRequest = {
+  method: "sampling/createMessage",
+  params: { messages: [], maxTokens: 10 },
+};
+
+describe("wrapSamplingHandler", () => {
+  it("returns the response and records sampling metrics", async () => {
+    const collector = createSamplingCollector();
+    const response = {
+      model: "test-model",
+      content: { type: "text", text: "hello" },
+      usage: { inputTokens: 2, outputTokens: 3 },
+    };
+    const handler = vi.fn(async () => response);
+    await expect(
+      wrapSamplingHandler(handler, collector.collect)(request),
+    ).resolves.toBe(response);
+    expect(handler).toHaveBeenCalledWith(request);
+    expect(collector.getCalls()).toEqual([
+      expect.objectContaining({
+        model_id: "test-model",
+        input_tokens: 2,
+        output_tokens: 3,
+        duration_ms: expect.any(Number),
+      }),
+    ]);
+  });
+
+  it.each([false, true])(
+    "preserves successful responses when the observer fails asynchronously=%s",
+    async (asynchronous) => {
+      const error = new Error("metrics unavailable");
+      const observer = asynchronous
+        ? async () => {
+            throw error;
+          }
+        : () => {
+            throw error;
+          };
+      const response = { content: "successful response" };
+      const log = vi.spyOn(console, "error").mockImplementation(() => {});
+      try {
+        await expect(
+          wrapSamplingHandler(async () => response, observer)(request),
+        ).resolves.toBe(response);
+        expect(log).toHaveBeenCalledWith(
+          "[assistant-cloud] onSamplingCall callback threw an error",
+          error,
+        );
+      } finally {
+        log.mockRestore();
+      }
+    },
+  );
+
+  it("does not await an asynchronous observer", async () => {
+    const response = { content: "successful response" };
+    await expect(
+      wrapSamplingHandler(
+        async () => response,
+        () => new Promise<void>(() => {}),
+      )(request),
+    ).resolves.toBe(response);
+  });
+
+  it("propagates model errors without calling the observer", async () => {
+    const error = new Error("model failed");
+    const observer = vi.fn();
+    await expect(
+      wrapSamplingHandler(async () => {
+        throw error;
+      }, observer)(request),
+    ).rejects.toBe(error);
+    expect(observer).not.toHaveBeenCalled();
+  });
+
+  it("preserves successful responses even if error reporting throws", async () => {
+    const response = { content: "successful response" };
+    const log = vi.spyOn(console, "error").mockImplementation(() => {
+      throw new Error("logging failed");
+    });
+    try {
+      await expect(
+        wrapSamplingHandler(
+          async () => response,
+          () => {
+            throw new Error("metrics failed");
+          },
+        )(request),
+      ).resolves.toBe(response);
+    } finally {
+      log.mockRestore();
+    }
+  });
+});

package/src/instrumentMcpSampling.ts

@@ -43,6 +43,15 @@ export type McpSamplingResponse = {
   [key: string]: unknown;
 };
 
+const reportSamplingError = (error: unknown): void => {
+  try {
+    console.error(
+      "[assistant-cloud] onSamplingCall callback threw an error",
+      error,
+    );
+  } catch {}
+};
+
 /**
  * Wraps an MCP sampling handler to intercept and measure sampling calls.
  *
@@ -80,7 +89,7 @@ export function wrapSamplingHandler(
     const reasoningTokens = response.usage?.reasoningTokens;
     const cachedInputTokens = response.usage?.cachedInputTokens;
 
-    onSamplingCall({
+    const data: SamplingCallData = {
       ...(modelId ? { model_id: modelId } : undefined),
       ...(inputTokens != null ? { input_tokens: inputTokens } : undefined),
       ...(outputTokens != null ? { output_tokens: outputTokens } : undefined),
@@ -91,7 +100,12 @@ export function wrapSamplingHandler(
         ? { cached_input_tokens: cachedInputTokens }
         : undefined),
       duration_ms: durationMs,
-    });
+    };
+    try {
+      void Promise.resolve(onSamplingCall(data)).catch(reportSamplingError);
+    } catch (error) {
+      reportSamplingError(error);
+    }
 
     return response;
   };

package/src/runTelemetry.test.ts

@@ -1,12 +1,296 @@
 import { describe, expect, it } from "vitest";
 import {
+  createRunReport,
   createRunTelemetryToolCall,
+  deriveRunOutcome,
+  describeRunError,
+  extractRunTelemetryModelId,
   normalizeRunTelemetryUsage,
   truncateRunTelemetryText,
 } from "./runTelemetry";
 
 const MAX = 50_000;
 
+describe("deriveRunOutcome", () => {
+  it.each([
+    [{ isError: true }, { status: "error" }],
+    [{ isAbort: true }, { status: "incomplete", outcome: "aborted" }],
+    [{ isDisconnect: true }, { status: "incomplete", outcome: "disconnected" }],
+    [{ finishReason: "length" }, { status: "incomplete", outcome: "length" }],
+    [
+      { finishReason: "content-filter" },
+      { status: "incomplete", outcome: "content_filter" },
+    ],
+    [
+      { finishReason: "content_filter" },
+      { status: "incomplete", outcome: "content_filter" },
+    ],
+    [
+      { finishReason: "cancelled" },
+      { status: "incomplete", outcome: "aborted" },
+    ],
+    [{ finishReason: "stop" }, { status: "completed" }],
+    [{ finishReason: "tool-calls" }, { status: "completed" }],
+    [{}, { status: "completed" }],
+    [{ finishReason: "other" }, { status: "completed" }],
+    [{ finishReason: "error" }, { status: "error" }],
+  ])("maps %o", (input, expected) => {
+    expect(deriveRunOutcome(input)).toEqual(expected);
+  });
+
+  it("uses the fallback status only when the event carries no finish reason", () => {
+    expect(deriveRunOutcome({}, "incomplete")).toEqual({
+      status: "incomplete",
+    });
+    expect(deriveRunOutcome({ finishReason: "stop" }, "incomplete")).toEqual({
+      status: "completed",
+    });
+    expect(deriveRunOutcome({ finishReason: "length" }, "incomplete")).toEqual({
+      status: "incomplete",
+      outcome: "length",
+    });
+    expect(deriveRunOutcome({ isError: true }, "incomplete")).toEqual({
+      status: "error",
+    });
+  });
+});
+
+describe("describeRunError", () => {
+  it("reads the message and the class name of an AI SDK error", () => {
+    const error = new Error("Rate limited");
+    error.name = "AI_APICallError";
+    expect(describeRunError(error)).toEqual({
+      error: "Rate limited",
+      errorCode: "AI_APICallError",
+    });
+  });
+
+  it("prefers an explicit code and skips the plain Error name", () => {
+    expect(
+      describeRunError(Object.assign(new Error("boom"), { code: "ETIMEDOUT" })),
+    ).toEqual({ error: "boom", errorCode: "ETIMEDOUT" });
+    expect(describeRunError(new Error("boom"))).toEqual({ error: "boom" });
+    expect(describeRunError("boom")).toEqual({ error: "boom" });
+    expect(describeRunError(undefined)).toEqual({});
+    expect(describeRunError(42)).toEqual({});
+  });
+
+  it("clamps the message and code to what the runs endpoint accepts", () => {
+    const error = new Error("m".repeat(3000));
+    error.name = "c".repeat(100);
+    expect(describeRunError(error)).toEqual({
+      error: "m".repeat(2048),
+      errorCode: "c".repeat(64),
+    });
+  });
+});
+
+describe("createRunReport", () => {
+  const toolCall = {
+    tool_name: "weather",
+    tool_call_id: "call-1",
+    tool_args: '{"city":"Singapore"}',
+  };
+
+  it("creates the full aui/v0 report", () => {
+    expect(
+      createRunReport({
+        threadId: "aui-thread",
+        status: "incomplete",
+        outcome: "length",
+        errorCode: "model_limit",
+        error: "response cut off",
+        messageId: "cloud-message",
+        traceId: "AABBCCDDEEFF00112233445566778899",
+        modelId: "provider/model",
+        provider: "openai",
+        usage: {
+          promptTokens: 11,
+          completionTokens: 7,
+          reasoningTokens: 3,
+          cachedInputTokens: 2,
+        },
+        steps: [
+          {
+            usage: { inputTokens: 5, outputTokens: 4 },
+            toolCalls: [toolCall],
+            startMs: 10,
+            endMs: 20,
+            finishReason: "tool-calls",
+          },
+          {
+            usage: { inputTokens: 6, outputTokens: 3, reasoningTokens: 3 },
+            startMs: 21,
+            endMs: 42,
+            finishReason: "length".repeat(8),
+          },
+        ],
+        toolCalls: [toolCall],
+        durationMs: 42.5,
+        firstTokenMs: -4,
+        outputText: "hello from aui/v0",
+        metadata: { tenant: "acme" },
+        telemetry: {
+          environment: "production",
+          release: "web-2026.09.09",
+          tags: [" region:sg ", "region:sg", "tier:paid"],
+        },
+      }),
+    ).toEqual({
+      thread_id: "aui-thread",
+      status: "incomplete",
+      outcome_type: "length",
+      error_code: "model_limit",
+      error: "response cut off",
+      message_id: "cloud-message",
+      trace_id: "aabbccddeeff00112233445566778899",
+      model_id: "provider/model",
+      provider: "openai",
+      provider_type: "openai",
+      input_tokens: 11,
+      output_tokens: 7,
+      reasoning_tokens: 3,
+      cached_input_tokens: 2,
+      steps: [
+        {
+          input_tokens: 5,
+          output_tokens: 4,
+          tool_calls: [toolCall],
+          start_ms: 10,
+          end_ms: 20,
+          finish_reason: "tool-calls",
+        },
+        {
+          input_tokens: 6,
+          output_tokens: 3,
+          reasoning_tokens: 3,
+          start_ms: 21,
+          end_ms: 42,
+          finish_reason: "lengthlengthlengthlengthlengthle",
+        },
+      ],
+      total_steps: 2,
+      tool_calls: [toolCall],
+      duration_ms: 43,
+      first_token_ms: 0,
+      output_text: "hello from aui/v0",
+      metadata: { tenant: "acme" },
+      environment: "production",
+      release: "web-2026.09.09",
+      tags: ["region:sg", "tier:paid"],
+    });
+  });
+
+  it("creates the full ai-sdk/v6 report", () => {
+    expect(
+      createRunReport({
+        threadId: "ai-sdk-thread",
+        status: "completed",
+        traceId: "00112233445566778899aabbccddeeff",
+        modelId: "gpt-5.6-terra",
+        provider: "gateway",
+        usage: {
+          inputTokens: 100,
+          outputTokens: 50,
+          inputTokenDetails: { cacheReadTokens: 10 },
+          outputTokenDetails: { reasoningTokens: 20 },
+        },
+        steps: [
+          {
+            usage: {
+              inputTokens: 100,
+              outputTokens: 50,
+              inputTokenDetails: { cacheReadTokens: 10 },
+              outputTokenDetails: { reasoningTokens: 20 },
+            },
+            finishReason: "stop",
+          },
+        ],
+        durationMs: 125.2,
+        firstTokenMs: 12.6,
+        outputText: "hello from ai-sdk/v6",
+        telemetry: { tags: ["sdk:v6"] },
+      }),
+    ).toEqual({
+      thread_id: "ai-sdk-thread",
+      status: "completed",
+      trace_id: "00112233445566778899aabbccddeeff",
+      model_id: "gpt-5.6-terra",
+      provider: "gateway",
+      provider_type: "gateway",
+      input_tokens: 100,
+      output_tokens: 50,
+      reasoning_tokens: 20,
+      cached_input_tokens: 10,
+      steps: [
+        {
+          input_tokens: 100,
+          output_tokens: 50,
+          reasoning_tokens: 20,
+          cached_input_tokens: 10,
+          finish_reason: "stop",
+        },
+      ],
+      total_steps: 1,
+      duration_ms: 125,
+      first_token_ms: 13,
+      output_text: "hello from ai-sdk/v6",
+      tags: ["sdk:v6"],
+    });
+  });
+
+  it("normalizes tags and limits them to twenty 64-character values", () => {
+    const longTag = "a".repeat(70);
+    const tags = [" one ", "", "one", longTag];
+    for (let index = 0; index < 24; index++) tags.push(`tag:${index}`);
+
+    const report = createRunReport({
+      threadId: "thread",
+      status: "completed",
+      telemetry: { tags },
+    });
+
+    expect(report.tags).toEqual([
+      "one",
+      "a".repeat(64),
+      ...Array.from({ length: 18 }, (_, index) => `tag:${index}`),
+    ]);
+  });
+
+  it("trims the whitespace a 64-character cut leaves behind", () => {
+    const report = createRunReport({
+      threadId: "thread",
+      status: "completed",
+      telemetry: { tags: [`${"b".repeat(63)} tail`] },
+    });
+
+    expect(report.tags).toEqual(["b".repeat(63)]);
+  });
+
+  it("lowercases valid trace IDs, omits invalid IDs, and omits undefined keys", () => {
+    expect(
+      createRunReport({
+        threadId: "thread",
+        status: "completed",
+        traceId: "AABBCCDDEEFF00112233445566778899",
+      }).trace_id,
+    ).toBe("aabbccddeeff00112233445566778899");
+    expect(
+      createRunReport({
+        threadId: "thread",
+        status: "completed",
+        traceId: "aabbccdd",
+      }),
+    ).toEqual({ thread_id: "thread", status: "completed" });
+    expect(
+      createRunReport({
+        threadId: "thread",
+        status: "completed",
+      }),
+    ).toEqual({ thread_id: "thread", status: "completed" });
+  });
+});
+
 describe("truncateRunTelemetryText", () => {
   it("passes text at or under the cap through unchanged", () => {
     expect(truncateRunTelemetryText("hello")).toBe("hello");
@@ -169,3 +453,44 @@ describe("normalizeRunTelemetryUsage", () => {
     ).toBeUndefined();
   });
 });
+
+describe("extractRunTelemetryModelId", () => {
+  it("prefers an explicit model ID over every fallback", () => {
+    expect(
+      extractRunTelemetryModelId({
+        modelId: "explicit",
+        custom: { modelId: "custom" },
+        steps: [{ response: { modelId: "step" } }],
+      }),
+    ).toBe("explicit");
+  });
+
+  it("falls back to the custom bag before the steps", () => {
+    expect(
+      extractRunTelemetryModelId({
+        custom: { modelId: "custom" },
+        steps: [{ response: { modelId: "step" } }],
+      }),
+    ).toBe("custom");
+  });
+
+  it("reads the first step that carries a response model ID", () => {
+    expect(
+      extractRunTelemetryModelId({
+        steps: [{ usage: { inputTokens: 1 } }, { response: { modelId: "b" } }],
+      }),
+    ).toBe("b");
+  });
+
+  it("returns undefined when no channel carries one", () => {
+    expect(extractRunTelemetryModelId(undefined)).toBeUndefined();
+    expect(extractRunTelemetryModelId({})).toBeUndefined();
+    expect(
+      extractRunTelemetryModelId({
+        modelId: 7,
+        custom: { modelId: null },
+        steps: [null, "step", { response: null }, { response: { modelId: 7 } }],
+      }),
+    ).toBeUndefined();
+  });
+});