assistant-cloud 0.1.43 → 0.2.0

package/src/telemetry/index.test.ts

@@ -0,0 +1,252 @@
+import {
+  ROOT_CONTEXT,
+  context,
+  trace,
+  type Context,
+  type ContextManager,
+} from "@opentelemetry/api";
+import type { ReadableSpan, SpanExporter } from "@opentelemetry/sdk-trace-base";
+import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
+import {
+  assistantCloudTraceExportOptions,
+  assistantCloudTraceMetadata,
+  createAssistantCloudSpanProcessor,
+  isAssistantCloudSpan,
+  withAssistantCloudTraceMetadata,
+} from "./index";
+
+const TRACE_ID = "0123456789abcdef0123456789abcdef";
+const SPAN_ID = "0123456789abcdef";
+
+class TestContextManager implements ContextManager {
+  private activeContext: Context = ROOT_CONTEXT;
+
+  active(): Context {
+    return this.activeContext;
+  }
+
+  with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
+    nextContext: Context,
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    const previousContext = this.activeContext;
+    this.activeContext = nextContext;
+    try {
+      return fn.call(thisArg, ...args);
+    } finally {
+      this.activeContext = previousContext;
+    }
+  }
+
+  bind<T>(_context: Context, target: T): T {
+    return target;
+  }
+
+  enable(): this {
+    return this;
+  }
+
+  disable(): this {
+    this.activeContext = ROOT_CONTEXT;
+    return this;
+  }
+}
+
+function span(
+  name: string,
+  attributes: Record<string, unknown> = {},
+): ReadableSpan {
+  return {
+    name,
+    attributes,
+    spanContext: () => ({
+      traceId: TRACE_ID,
+      spanId: SPAN_ID,
+      traceFlags: 1,
+    }),
+    resource: { asyncAttributesPending: false },
+  } as ReadableSpan;
+}
+
+function recordingExporter() {
+  const spans: ReadableSpan[] = [];
+  const exporter = {
+    export: vi.fn(
+      (
+        nextSpans: ReadableSpan[],
+        resultCallback: Parameters<SpanExporter["export"]>[1],
+      ) => {
+        spans.push(...nextSpans);
+        resultCallback({ code: 0 });
+      },
+    ),
+    shutdown: vi.fn().mockResolvedValue(undefined),
+  } satisfies SpanExporter;
+
+  return { exporter, spans };
+}
+
+beforeAll(() => {
+  context.disable();
+  context.setGlobalContextManager(new TestContextManager());
+});
+
+afterAll(() => {
+  context.disable();
+});
+
+describe("assistantCloudTraceExportOptions", () => {
+  it("uses the default receiver and authorization header", () => {
+    expect(assistantCloudTraceExportOptions({ apiKey: "key" })).toEqual({
+      url: "https://backend.assistant-api.com/v1/traces",
+      headers: { Authorization: "Bearer key" },
+    });
+  });
+
+  it("uses a custom receiver without a trailing slash", () => {
+    expect(
+      assistantCloudTraceExportOptions({
+        apiKey: "key",
+        baseUrl: "https://cloud.example.com",
+      }),
+    ).toEqual({
+      url: "https://cloud.example.com/v1/traces",
+      headers: { Authorization: "Bearer key" },
+    });
+  });
+
+  it("removes trailing slashes and preserves extra headers", () => {
+    expect(
+      assistantCloudTraceExportOptions({
+        apiKey: "key",
+        baseUrl: "https://cloud.example.com/",
+        headers: { "X-Workspace": "workspace_1" },
+      }),
+    ).toEqual({
+      url: "https://cloud.example.com/v1/traces",
+      headers: {
+        Authorization: "Bearer key",
+        "X-Workspace": "workspace_1",
+      },
+    });
+  });
+
+  it("rejects an empty API key", () => {
+    expect(() => assistantCloudTraceExportOptions({ apiKey: "" })).toThrow(
+      "An Assistant Cloud API key is required",
+    );
+  });
+});
+
+describe("isAssistantCloudSpan", () => {
+  it("accepts AI span name prefixes", () => {
+    expect(isAssistantCloudSpan(span("gen_ai.client"))).toBe(true);
+    expect(isAssistantCloudSpan(span("ai.streamText"))).toBe(true);
+    expect(isAssistantCloudSpan(span("llm.generate"))).toBe(true);
+  });
+
+  it("accepts the GenAI operation attribute", () => {
+    expect(
+      isAssistantCloudSpan(
+        span("generate", { "gen_ai.operation.name": "ai.generateText" }),
+      ),
+    ).toBe(true);
+  });
+
+  it("rejects ordinary spans", () => {
+    expect(isAssistantCloudSpan(span("GET /api/chat"))).toBe(false);
+  });
+});
+
+describe("createAssistantCloudSpanProcessor", () => {
+  it("exports matching spans and drops other spans", async () => {
+    const { exporter, spans } = recordingExporter();
+    const processor = createAssistantCloudSpanProcessor(exporter);
+
+    processor.onEnd(span("gen_ai.client"));
+    processor.onEnd(span("GET /api/chat"));
+    await processor.forceFlush();
+
+    expect(spans.map((candidate) => candidate.name)).toEqual(["gen_ai.client"]);
+  });
+
+  it("uses a custom filter when one is provided", async () => {
+    const { exporter, spans } = recordingExporter();
+    const processor = createAssistantCloudSpanProcessor(exporter, {
+      filter: (candidate) => candidate.name === "custom",
+    });
+
+    processor.onEnd(span("gen_ai.client"));
+    processor.onEnd(span("custom"));
+    await processor.forceFlush();
+
+    expect(spans.map((candidate) => candidate.name)).toEqual(["custom"]);
+  });
+
+  it("flushes and shuts down the wrapped processor", async () => {
+    const { exporter } = recordingExporter();
+    const processor = createAssistantCloudSpanProcessor(exporter);
+
+    await processor.forceFlush();
+    await processor.shutdown();
+
+    expect(exporter.shutdown).toHaveBeenCalledTimes(1);
+  });
+});
+
+describe("assistantCloudTraceMetadata", () => {
+  it("returns the active trace ID", () => {
+    const result = context.with(
+      trace.setSpanContext(context.active(), {
+        traceId: TRACE_ID,
+        spanId: SPAN_ID,
+        traceFlags: 1,
+      }),
+      assistantCloudTraceMetadata,
+    );
+
+    expect(result).toEqual({ traceId: TRACE_ID });
+  });
+
+  it("returns an empty object without an active span", () => {
+    expect(assistantCloudTraceMetadata()).toEqual({});
+  });
+});
+
+describe("withAssistantCloudTraceMetadata", () => {
+  it("merges the active trace ID into start metadata", () => {
+    const messageMetadata = withAssistantCloudTraceMetadata(({ part }) => ({
+      type: part.type,
+    }));
+
+    const result = context.with(
+      trace.setSpanContext(context.active(), {
+        traceId: TRACE_ID,
+        spanId: SPAN_ID,
+        traceFlags: 1,
+      }),
+      () => messageMetadata({ part: { type: "start" } }),
+    );
+
+    expect(result).toEqual({ type: "start", traceId: TRACE_ID });
+  });
+
+  it("passes through start metadata when no span is active", () => {
+    const value = { started: true };
+    const messageMetadata = withAssistantCloudTraceMetadata(() => value);
+
+    expect(messageMetadata({ part: { type: "start" } })).toBe(value);
+    expect(withAssistantCloudTraceMetadata()({ part: { type: "start" } })).toBe(
+      undefined,
+    );
+  });
+
+  it("passes through metadata for other parts", () => {
+    const value = { usage: { totalTokens: 42 } };
+    const messageMetadata = withAssistantCloudTraceMetadata(() => value);
+
+    expect(messageMetadata({ part: { type: "finish" } })).toBe(value);
+  });
+});

package/src/telemetry/index.ts

@@ -0,0 +1,98 @@
+import { trace } from "@opentelemetry/api";
+import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
+import {
+  BatchSpanProcessor,
+  type ReadableSpan,
+  type SpanExporter,
+  type SpanProcessor,
+} from "@opentelemetry/sdk-trace-base";
+
+const DEFAULT_BACKEND_BASE_URL = "https://backend.assistant-api.com";
+const AI_SPAN_PREFIXES = ["gen_ai.", "ai.", "llm."];
+
+export type AssistantCloudTraceExportOptions = {
+  apiKey: string;
+  baseUrl?: string;
+  headers?: Record<string, string>;
+};
+
+export type AssistantCloudSpanProcessorOptions = {
+  filter?: (span: ReadableSpan) => boolean;
+};
+
+export function assistantCloudTraceExportOptions({
+  apiKey,
+  baseUrl = DEFAULT_BACKEND_BASE_URL,
+  headers,
+}: AssistantCloudTraceExportOptions): {
+  url: string;
+  headers: Record<string, string>;
+} {
+  if (!apiKey) throw new Error("An Assistant Cloud API key is required");
+
+  return {
+    url: `${trimTrailingSlashes(baseUrl)}/v1/traces`,
+    headers: {
+      Authorization: `Bearer ${apiKey}`,
+      ...headers,
+    },
+  };
+}
+
+function trimTrailingSlashes(url: string): string {
+  let end = url.length;
+  while (end > 0 && url[end - 1] === "/") end--;
+  return url.slice(0, end);
+}
+
+export function createAssistantCloudTraceExporter(
+  options: AssistantCloudTraceExportOptions,
+): SpanExporter {
+  return new OTLPTraceExporter(assistantCloudTraceExportOptions(options));
+}
+
+export function isAssistantCloudSpan(span: ReadableSpan): boolean {
+  return (
+    typeof span.attributes["gen_ai.operation.name"] === "string" ||
+    AI_SPAN_PREFIXES.some((prefix) => span.name.startsWith(prefix))
+  );
+}
+
+export function createAssistantCloudSpanProcessor(
+  exporter: SpanExporter,
+  options: AssistantCloudSpanProcessorOptions = {},
+): SpanProcessor {
+  const processor = new BatchSpanProcessor(exporter);
+  const filter = options.filter ?? isAssistantCloudSpan;
+
+  return {
+    onStart: (span, parentContext) => processor.onStart(span, parentContext),
+    onEnd: (span) => {
+      if (filter(span)) processor.onEnd(span);
+    },
+    forceFlush: () => processor.forceFlush(),
+    shutdown: () => processor.shutdown(),
+  };
+}
+
+export function assistantCloudTraceMetadata(): { traceId?: string } {
+  const spanContext = trace.getActiveSpan()?.spanContext();
+  return spanContext && trace.isSpanContextValid(spanContext)
+    ? { traceId: spanContext.traceId }
+    : {};
+}
+
+export function withAssistantCloudTraceMetadata<
+  Part extends { type: string } = { type: string },
+>(
+  messageMetadata?: (options: {
+    part: Part;
+  }) => Record<string, unknown> | undefined,
+): (options: { part: Part }) => Record<string, unknown> | undefined {
+  return (options) => {
+    const metadata = messageMetadata?.(options);
+    if (options.part.type !== "start") return metadata;
+    const trace = assistantCloudTraceMetadata();
+    return trace.traceId === undefined ? metadata : { ...metadata, ...trace };
+  };
+}

package/src/tests/AssistantCloud.test.ts

@@ -23,6 +23,13 @@ describe("AssistantCloud telemetry config", () => {
     expect(createCloud({ enabled: false }).telemetry.enabled).toBe(false);
   });
 
+  it("can disable engagement events without disabling run reports", () => {
+    expect(createCloud({ events: false }).telemetry).toEqual({
+      enabled: true,
+      events: false,
+    });
+  });
+
   it("stays enabled when the config object carries an undefined enabled", () => {
     const beforeReport: NonNullable<
       AssistantCloudTelemetryConfig["beforeReport"]
@@ -36,4 +43,19 @@ describe("AssistantCloud telemetry config", () => {
     expect(telemetry.enabled).toBe(true);
     expect(telemetry.beforeReport).toBe(beforeReport);
   });
+
+  it("preserves configured run report dimensions", () => {
+    expect(
+      createCloud({
+        release: "web-2026.09.08",
+        environment: "production",
+        tags: ["region:sg", "tier:paid"],
+      }).telemetry,
+    ).toEqual({
+      enabled: true,
+      release: "web-2026.09.08",
+      environment: "production",
+      tags: ["region:sg", "tier:paid"],
+    });
+  });
 });

package/src/tests/AssistantCloudAPI.test.ts

@@ -143,7 +143,7 @@ describe("AssistantCloudAPI", () => {
     );
   });
 
-  it("throws APIError with parsed message for JSON error responses", async () => {
+  it("preserves the current error shape for JSON error responses without a code", async () => {
     const fetchMock = vi.fn().mockResolvedValue({
       ok: false,
       status: 400,
@@ -167,16 +167,54 @@ describe("AssistantCloudAPI", () => {
     expect(error.name).toBe("CloudAPIError");
     expect(error.message).toBe("invalid request payload");
     expect(error.status).toBe(400);
+    expect(error.code).toBeUndefined();
+    expect(error.details).toBeUndefined();
   });
 
-  it("falls back to the response text when the JSON error body has no message", async () => {
+  it("exposes the plan-limit error code and details", async () => {
+    const fetchMock = vi.fn().mockResolvedValue({
+      ok: false,
+      status: 402,
+      headers: new Headers(),
+      text: vi.fn().mockResolvedValue(
+        JSON.stringify({
+          error: "plan_limit_reached",
+          plan: "free",
+          period_end: "2026-10-01T00:00:00.000Z",
+          cap: 100,
+        }),
+      ),
+    });
+    vi.stubGlobal("fetch", fetchMock);
+
+    const api = new AssistantCloudAPI({
+      apiKey: "test-key",
+      userId: "u-1",
+      workspaceId: "w-1",
+    });
+
+    const error = await api.makeRawRequest("/threads").catch((e) => e);
+    expect(error).toBeInstanceOf(CloudAPIError);
+    expect(error.message).toBe(
+      'Request failed with status 402, {"error":"plan_limit_reached","plan":"free","period_end":"2026-10-01T00:00:00.000Z","cap":100}',
+    );
+    expect(error.status).toBe(402);
+    expect(error.code).toBe("plan_limit_reached");
+    expect(error.details).toEqual({
+      plan: "free",
+      period_end: "2026-10-01T00:00:00.000Z",
+      cap: 100,
+    });
+  });
+
+  it("exposes rate-limit error codes with empty details", async () => {
     const fetchMock = vi.fn().mockResolvedValue({
       ok: false,
       status: 429,
       headers: new Headers(),
       text: vi
         .fn()
-        .mockResolvedValue(JSON.stringify({ error: "rate limited" })),
+        .mockResolvedValue(JSON.stringify({ error: "rate_limited" })),
     });
     vi.stubGlobal("fetch", fetchMock);
 
@@ -189,9 +227,11 @@ describe("AssistantCloudAPI", () => {
     const error = await api.makeRawRequest("/threads").catch((e) => e);
     expect(error).toBeInstanceOf(CloudAPIError);
     expect(error.message).toBe(
-      'Request failed with status 429, {"error":"rate limited"}',
+      'Request failed with status 429, {"error":"rate_limited"}',
     );
     expect(error.status).toBe(429);
+    expect(error.code).toBe("rate_limited");
+    expect(error.details).toEqual({});
   });
 
   it("throws generic error with status for non-JSON error responses", async () => {

package/src/tests/AssistantCloudAuthStrategy.test.ts

@@ -1,7 +1,9 @@
+/// <reference types="node" />
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import {
   AssistantCloudAnonymousAuthStrategy,
   AssistantCloudJWTAuthStrategy,
+  readAnonymousRefreshToken,
 } from "../AssistantCloudAuthStrategy";
 import { CloudResponseError } from "../cloudResponse";
 
@@ -56,6 +58,73 @@ describe("AssistantCloudAnonymousAuthStrategy", () => {
     }
   });
 
+  it("reads the stored anonymous refresh token", () => {
+    const values = new Map([[refreshTokenKey, JSON.stringify(refreshToken)]]);
+    installLocalStorage({
+      getItem: (key) => values.get(key) ?? null,
+      setItem: (key, value) => {
+        values.set(key, value);
+      },
+      removeItem: (key) => {
+        values.delete(key);
+      },
+    } as Storage);
+
+    expect(readAnonymousRefreshToken(baseUrl)).toBe(refreshToken.token);
+  });
+
+  it("reads the token stored for a base url given with a trailing slash", () => {
+    // The two-step cast below erases the contextual Storage typing the plain
+    // `as Storage` sites get, so this parameter needs its own annotation.
+    installLocalStorage({
+      getItem: (key: string) =>
+        key === refreshTokenKey ? JSON.stringify(refreshToken) : null,
+      setItem: vi.fn(),
+      removeItem: vi.fn(),
+    } as unknown as Storage);
+
+    expect(readAnonymousRefreshToken(`${baseUrl}/`)).toBe(refreshToken.token);
+  });
+
+  it("returns null when no anonymous refresh token is stored", () => {
+    installLocalStorage({
+      getItem: () => null,
+      setItem: vi.fn(),
+      removeItem: vi.fn(),
+    } as unknown as Storage);
+
+    expect(readAnonymousRefreshToken(baseUrl)).toBeNull();
+  });
+
+  it.each([30_000, 0, -1])(
+    "returns null when the stored anonymous refresh token expires in %i ms",
+    (expiresIn) => {
+      vi.useFakeTimers();
+      const now = Date.UTC(2026, 8, 4, 12, 0, 0);
+      vi.setSystemTime(now);
+      const values = new Map([
+        [
+          refreshTokenKey,
+          JSON.stringify({
+            token: refreshToken.token,
+            expires_at: new Date(now + expiresIn).toISOString(),
+          }),
+        ],
+      ]);
+      installLocalStorage({
+        getItem: (key) => values.get(key) ?? null,
+        setItem: (key, value) => {
+          values.set(key, value);
+        },
+        removeItem: (key) => {
+          values.delete(key);
+        },
+      } as Storage);
+
+      expect(readAnonymousRefreshToken(baseUrl)).toBeNull();
+    },
+  );
+
   it.each([
     "2099-01-01",
     "2099-01-01T00:00:00Z",
@@ -269,7 +338,7 @@ describe("AssistantCloudAnonymousAuthStrategy", () => {
                 { once: true },
               );
             }),
-        } as Response),
+        } as unknown as Response),
       )
       .mockResolvedValueOnce({
         ok: true,
@@ -598,7 +667,7 @@ describe("AssistantCloudAnonymousAuthStrategy", () => {
     const removeItem = vi.fn(() => {
       throw new DOMException("blocked", "SecurityError");
     });
-    installLocalStorage({ getItem, setItem, removeItem } as Storage);
+    installLocalStorage({ getItem, setItem, removeItem } as unknown as Storage);
     mockAnonymousTokenFetch();
 
     await expect(

package/src/tests/CloudMessagePersistence.test.ts

@@ -14,6 +14,18 @@ function createMockCloud() {
   } as unknown as AssistantCloud;
 }
 
+function createCloudMessages(count: number) {
+  return Array.from({ length: count }, (_, index) => ({
+    id: `message-${index + 1}`,
+    parent_id: null,
+    height: index,
+    created_at: "2025-01-01T00:00:00Z" as unknown as Date,
+    updated_at: "2025-01-01T00:00:00Z" as unknown as Date,
+    format: "aui/v0",
+    content: {},
+  }));
+}
+
 describe("CloudMessagePersistence", () => {
   let cloud: AssistantCloud;
   let persistence: CloudMessagePersistence;
@@ -35,8 +47,29 @@ describe("CloudMessagePersistence", () => {
 
     expect(persistence.isPersisted("local-1")).toBe(true);
     expect(await persistence.getRemoteId("local-1")).toBe("remote-1");
+    expect(persistence.getResolvedRemoteId("local-1")).toBe("remote-1");
   });
 
+  it.each(["__proto__", "constructor", "toString"])(
+    "supports prototype-named local ID %s",
+    async (messageId) => {
+      vi.mocked(cloud.threads.messages.create).mockResolvedValue({
+        message_id: `remote-${messageId}`,
+      });
+
+      expect(persistence.isPersisted(messageId)).toBe(false);
+
+      await persistence.append("thread-1", messageId, null, "aui/v0", {
+        text: "hello",
+      });
+
+      expect(persistence.isPersisted(messageId)).toBe(true);
+      expect(await persistence.getRemoteId(messageId)).toBe(
+        `remote-${messageId}`,
+      );
+    },
+  );
+
   it("uses the current client without losing ID mappings", async () => {
     const firstCloud = createMockCloud();
     const secondCloud = createMockCloud();
@@ -221,6 +254,88 @@ describe("CloudMessagePersistence", () => {
     expect(persistence.isPersisted("msg-1")).toBe(true);
   });
 
+  it("maps prototype-named loaded messages", async () => {
+    vi.mocked(cloud.threads.messages.list).mockResolvedValue({
+      messages: [
+        {
+          id: "__proto__",
+          parent_id: null,
+          height: 0,
+          created_at: "2025-01-01T00:00:00Z" as unknown as Date,
+          updated_at: "2025-01-01T00:00:00Z" as unknown as Date,
+          format: "aui/v0",
+          content: { text: "loaded" },
+        },
+      ],
+    });
+
+    await persistence.load("thread-1");
+
+    expect(persistence.isPersisted("__proto__")).toBe(true);
+    expect(await persistence.getRemoteId("__proto__")).toBe("__proto__");
+  });
+
+  it("follows the message cursor across pages in server order", async () => {
+    const rows = createCloudMessages(450);
+    vi.mocked(cloud.threads.messages.list).mockImplementation(
+      async (_threadId, query) => {
+        const start = query?.after
+          ? rows.findIndex((row) => row.id === query.after) + 1
+          : 0;
+        return { messages: rows.slice(start, start + 200) };
+      },
+    );
+
+    const messages = await persistence.load("thread-1", "aui/v0");
+
+    expect(messages.map((message) => message.id)).toEqual(
+      rows.map((row) => row.id),
+    );
+    expect(cloud.threads.messages.list).toHaveBeenCalledTimes(3);
+    expect(cloud.threads.messages.list).toHaveBeenNthCalledWith(1, "thread-1", {
+      format: "aui/v0",
+      limit: 200,
+    });
+    expect(cloud.threads.messages.list).toHaveBeenNthCalledWith(3, "thread-1", {
+      format: "aui/v0",
+      limit: 200,
+      after: "message-400",
+    });
+  });
+
+  it("stops without duplicates when a later cursor stops resolving", async () => {
+    const rows = createCloudMessages(450);
+    vi.mocked(cloud.threads.messages.list).mockImplementation(
+      async (_threadId, query) => {
+        // The third request replays page one, the way the endpoint answers a
+        // cursor whose message no longer resolves.
+        const start =
+          query?.after && query.after !== "message-400"
+            ? rows.findIndex((row) => row.id === query.after) + 1
+            : 0;
+        return { messages: rows.slice(start, start + 200) };
+      },
+    );
+
+    const messages = await persistence.load("thread-1", "aui/v0");
+
+    expect(messages.map((message) => message.id)).toEqual(
+      rows.slice(0, 400).map((row) => row.id),
+    );
+    expect(cloud.threads.messages.list).toHaveBeenCalledTimes(3);
+  });
+
+  it("stops when a page does not advance the cursor", async () => {
+    vi.mocked(cloud.threads.messages.list).mockResolvedValue({
+      messages: createCloudMessages(200),
+    });
+
+    const messages = await persistence.load("thread-1", "aui/v0");
+
+    expect(messages).toHaveLength(200);
+    expect(cloud.threads.messages.list).toHaveBeenCalledTimes(2);
+  });
+
   it("updates an already-persisted message", async () => {
     vi.mocked(cloud.threads.messages.create).mockResolvedValue({
       message_id: "remote-1",