@funkai/agents 0.1.0

package/src/core/provider/provider.test.ts

@@ -0,0 +1,128 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+// Mock the external provider before importing the module under test
+const mockBaseCreateOpenRouter = vi.fn();
+const mockProviderInstance = vi.fn();
+
+vi.mock("@openrouter/ai-sdk-provider", () => ({
+  createOpenRouter: mockBaseCreateOpenRouter,
+}));
+
+// Must import after mocking
+const { createOpenRouter, openrouter } = await import("@/core/provider/provider.js");
+
+// ---------------------------------------------------------------------------
+// createOpenRouter()
+// ---------------------------------------------------------------------------
+
+describe("createOpenRouter()", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockBaseCreateOpenRouter.mockReturnValue(mockProviderInstance);
+    vi.stubEnv("OPENROUTER_API_KEY", "env-key-123");
+  });
+
+  afterEach(() => {
+    vi.unstubAllEnvs();
+  });
+
+  it("uses explicit apiKey when provided in options", () => {
+    createOpenRouter({ apiKey: "explicit-key" });
+
+    expect(mockBaseCreateOpenRouter).toHaveBeenCalledWith(
+      expect.objectContaining({ apiKey: "explicit-key" }),
+    );
+  });
+
+  it("falls back to OPENROUTER_API_KEY env var when no apiKey in options", () => {
+    createOpenRouter();
+
+    expect(mockBaseCreateOpenRouter).toHaveBeenCalledWith(
+      expect.objectContaining({ apiKey: "env-key-123" }),
+    );
+  });
+
+  it("falls back to env var when options are provided without apiKey", () => {
+    createOpenRouter({});
+
+    expect(mockBaseCreateOpenRouter).toHaveBeenCalledWith(
+      expect.objectContaining({ apiKey: "env-key-123" }),
+    );
+  });
+
+  it("throws when no apiKey provided and env var is not set", () => {
+    vi.stubEnv("OPENROUTER_API_KEY", undefined as unknown as string);
+
+    expect(() => createOpenRouter()).toThrow("OPENROUTER_API_KEY environment variable is required");
+  });
+
+  it("forwards additional options to the base provider", () => {
+    createOpenRouter({ apiKey: "key", baseURL: "https://custom.api" } as Record<string, unknown>);
+
+    expect(mockBaseCreateOpenRouter).toHaveBeenCalledWith(
+      expect.objectContaining({
+        apiKey: "key",
+        baseURL: "https://custom.api",
+      }),
+    );
+  });
+});
+
+// ---------------------------------------------------------------------------
+// openrouter() — cached model factory
+// ---------------------------------------------------------------------------
+
+// Each test in this block uses a unique API key to invalidate the module-level
+// cache, so call-count assertions remain isolated between tests.
+
+describe("openrouter()", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockProviderInstance.mockReturnValue({ modelId: "mock-model" });
+    mockBaseCreateOpenRouter.mockReturnValue(mockProviderInstance);
+  });
+
+  afterEach(() => {
+    vi.unstubAllEnvs();
+  });
+
+  it("creates a provider and returns a language model", () => {
+    vi.stubEnv("OPENROUTER_API_KEY", "key-test-create");
+
+    const result = openrouter("openai/gpt-5.2-codex");
+
+    expect(mockBaseCreateOpenRouter).toHaveBeenCalledWith({ apiKey: "key-test-create" });
+    expect(mockProviderInstance).toHaveBeenCalledWith("openai/gpt-5.2-codex");
+    expect(result).toEqual({ modelId: "mock-model" });
+  });
+
+  it("reuses cached provider on subsequent calls with same api key", () => {
+    vi.stubEnv("OPENROUTER_API_KEY", "key-test-reuse");
+
+    openrouter("openai/gpt-5.2-codex");
+    openrouter("openai/gpt-5.2");
+
+    // Provider should only be created once since key is the same
+    expect(mockBaseCreateOpenRouter).toHaveBeenCalledTimes(1);
+    expect(mockProviderInstance).toHaveBeenCalledTimes(2);
+  });
+
+  it("creates new provider when api key changes", () => {
+    vi.stubEnv("OPENROUTER_API_KEY", "key-test-change-1");
+    openrouter("openai/gpt-5.2-codex");
+
+    vi.stubEnv("OPENROUTER_API_KEY", "key-test-change-2");
+    openrouter("openai/gpt-5.2");
+
+    expect(mockBaseCreateOpenRouter).toHaveBeenCalledTimes(2);
+    expect(mockBaseCreateOpenRouter).toHaveBeenLastCalledWith({ apiKey: "key-test-change-2" });
+  });
+
+  it("throws when OPENROUTER_API_KEY is not set", () => {
+    vi.stubEnv("OPENROUTER_API_KEY", undefined as unknown as string);
+
+    expect(() => openrouter("openai/gpt-5.2-codex")).toThrow(
+      "OPENROUTER_API_KEY environment variable is required",
+    );
+  });
+});

package/src/core/provider/provider.ts

@@ -0,0 +1,99 @@
+import {
+  createOpenRouter as baseCreateOpenRouter,
+  type OpenRouterProvider,
+  type OpenRouterProviderSettings,
+} from "@openrouter/ai-sdk-provider";
+
+import type { ModelId } from "@/core/models/index.js";
+import { type LanguageModel } from "@/core/provider/types.js";
+
+/**
+ * Create an OpenRouter provider instance with automatic API key resolution.
+ *
+ * Falls back to the `OPENROUTER_API_KEY` environment variable when
+ * no explicit `apiKey` is provided in the options.
+ *
+ * @param options - Provider settings forwarded to `@openrouter/ai-sdk-provider`.
+ * @returns A configured {@link OpenRouterProvider} instance.
+ *
+ * @example
+ * ```typescript
+ * const openrouter = createOpenRouter({ apiKey: 'sk-...' })
+ * const m = openrouter('openai/gpt-5.2-codex')
+ * ```
+ */
+export function createOpenRouter(options?: OpenRouterProviderSettings): OpenRouterProvider {
+  const apiKey = resolveApiKey(options);
+  return baseCreateOpenRouter({
+    ...options,
+    apiKey,
+  });
+}
+
+function resolveApiKey(options: OpenRouterProviderSettings | undefined): string {
+  if (options != null && options.apiKey != null) {
+    return options.apiKey;
+  }
+  return getOpenRouterApiKey();
+}
+
+/**
+ * Create a cached OpenRouter model resolver.
+ *
+ * The returned function caches the underlying provider and invalidates
+ * when the API key changes at runtime.
+ *
+ * @returns A function that resolves a model ID to a {@link LanguageModel}.
+ *
+ * @private
+ */
+function createCachedOpenRouter(): (modelId: ModelId) => LanguageModel {
+  const cache: { provider: OpenRouterProvider | undefined; apiKey: string | undefined } = {
+    provider: undefined,
+    apiKey: undefined,
+  };
+  return (modelId: ModelId): LanguageModel => {
+    const apiKey = getOpenRouterApiKey();
+    if (!cache.provider || cache.apiKey !== apiKey) {
+      cache.provider = baseCreateOpenRouter({ apiKey });
+      cache.apiKey = apiKey;
+    }
+    return cache.provider(modelId);
+  };
+}
+
+/**
+ * Shorthand for creating a single OpenRouter language model.
+ *
+ * Resolves the API key from the environment and returns a ready-to-use
+ * {@link LanguageModel} that can be passed directly to AI SDK functions.
+ *
+ * The provider instance is cached at module scope and reused across
+ * calls. If `OPENROUTER_API_KEY` changes at runtime, the cache is
+ * invalidated and a new provider is created.
+ *
+ * @param modelId - An OpenRouter model identifier (e.g. `"openai/gpt-5.2-codex"`).
+ * @returns A configured {@link LanguageModel} instance.
+ *
+ * @example
+ * ```typescript
+ * const m = openrouter('openai/gpt-5.2-codex')
+ * ```
+ */
+export const openrouter: (modelId: ModelId) => LanguageModel = createCachedOpenRouter();
+
+/**
+ * Read the OpenRouter API key from the environment.
+ *
+ * @throws {Error} If `OPENROUTER_API_KEY` is not set.
+ */
+function getOpenRouterApiKey(): string {
+  const apiKey = process.env.OPENROUTER_API_KEY;
+  if (!apiKey) {
+    throw new Error(
+      "OPENROUTER_API_KEY environment variable is required. " +
+        "Set it in your .env file or environment.",
+    );
+  }
+  return apiKey;
+}

package/src/core/provider/types.ts

@@ -0,0 +1,98 @@
+import { type LanguageModel as BaseLanguageModel } from "ai";
+
+import type { ModelId } from "@/core/models/index.js";
+
+/**
+ * AI SDK v3 language model type.
+ *
+ * Narrowed to the v3 specification version for type safety.
+ * All models created via `openrouter()` or `wrapModel()` satisfy this type.
+ */
+export type LanguageModel = Extract<BaseLanguageModel, { specificationVersion: "v3" }>;
+
+/**
+ * Base token counts shared by raw tracking records and final output.
+ *
+ * All fields are resolved `number` (0 when absent).
+ */
+export interface TokenUsage {
+  /** Number of input (prompt) tokens. */
+  readonly inputTokens: number;
+
+  /** Number of output (completion) tokens. */
+  readonly outputTokens: number;
+
+  /** Total tokens (input + output). */
+  readonly totalTokens: number;
+
+  /** Tokens served from the provider's prompt cache. */
+  readonly cacheReadTokens: number;
+
+  /** Tokens written into the provider's prompt cache. */
+  readonly cacheWriteTokens: number;
+
+  /** Tokens consumed by the model's internal reasoning (e.g. o3/o4). */
+  readonly reasoningTokens: number;
+}
+
+/**
+ * Raw tracking record from a single AI model invocation.
+ *
+ * Fields are `number | undefined` because providers may not report all fields.
+ * Carries `modelId` so that consumers can look up pricing if needed.
+ */
+export interface TokenUsageRecord {
+  /**
+   * The model identifier that produced this usage
+   * (e.g. `"openai/gpt-5.2-codex"`).
+   */
+  readonly modelId: ModelId;
+
+  /** Number of input (prompt) tokens. */
+  readonly inputTokens: number | undefined;
+
+  /** Number of output (completion) tokens. */
+  readonly outputTokens: number | undefined;
+
+  /** Total tokens (input + output). */
+  readonly totalTokens: number | undefined;
+
+  /** Tokens served from the provider's prompt cache. */
+  readonly cacheReadTokens: number | undefined;
+
+  /** Tokens written into the provider's prompt cache. */
+  readonly cacheWriteTokens: number | undefined;
+
+  /** Tokens consumed by the model's internal reasoning (e.g. o3/o4). */
+  readonly reasoningTokens: number | undefined;
+
+  /**
+   * Populated by the framework — identifies which component produced this usage.
+   */
+  readonly source?: {
+    readonly workflowId?: string;
+    readonly stepId?: string;
+    readonly agentId: string;
+    readonly scope: readonly string[];
+  };
+}
+
+/**
+ * Final agent-level usage — tokens flat, with agentId.
+ *
+ * Produced by `agentUsage()` at the end of an agent's execution.
+ */
+export interface AgentTokenUsage extends TokenUsage {
+  /** The agent that produced this usage. */
+  readonly agentId: string;
+}
+
+/**
+ * Final flow agent-level usage — per-agent breakdown.
+ *
+ * Produced by `flowAgentUsage()` at the end of a flow agent's execution.
+ */
+export interface FlowAgentTokenUsage {
+  /** Per-agent usage entries. */
+  readonly usages: readonly AgentTokenUsage[];
+}

package/src/core/provider/usage.test.ts

@@ -0,0 +1,304 @@
+import { describe, expect, it } from "vitest";
+
+import type { TokenUsage, TokenUsageRecord } from "@/core/provider/types.js";
+import { agentUsage, sumTokenUsage, flowAgentUsage } from "@/core/provider/usage.js";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function createRecord(overrides?: Partial<TokenUsageRecord>): TokenUsageRecord {
+  return {
+    modelId: "openai/gpt-5.2-codex",
+    inputTokens: undefined,
+    outputTokens: undefined,
+    totalTokens: undefined,
+    cacheReadTokens: undefined,
+    cacheWriteTokens: undefined,
+    reasoningTokens: undefined,
+    ...overrides,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// agentUsage()
+// ---------------------------------------------------------------------------
+
+describe("agentUsage()", () => {
+  it("returns zero counts for a record with all undefined fields", () => {
+    const result = agentUsage("agent-1", createRecord());
+
+    expect(result).toEqual({
+      agentId: "agent-1",
+      inputTokens: 0,
+      outputTokens: 0,
+      totalTokens: 0,
+      cacheReadTokens: 0,
+      cacheWriteTokens: 0,
+      reasoningTokens: 0,
+    });
+  });
+
+  it("passes through defined token counts from a single record", () => {
+    const record = createRecord({
+      inputTokens: 100,
+      outputTokens: 50,
+      totalTokens: 150,
+      cacheReadTokens: 10,
+      cacheWriteTokens: 5,
+      reasoningTokens: 20,
+    });
+
+    const result = agentUsage("agent-2", record);
+
+    expect(result.agentId).toBe("agent-2");
+    expect(result.inputTokens).toBe(100);
+    expect(result.outputTokens).toBe(50);
+    expect(result.totalTokens).toBe(150);
+    expect(result.cacheReadTokens).toBe(10);
+    expect(result.cacheWriteTokens).toBe(5);
+    expect(result.reasoningTokens).toBe(20);
+  });
+
+  it("accepts a single record (not wrapped in array)", () => {
+    const record = createRecord({ inputTokens: 42 });
+
+    const result = agentUsage("agent-single", record);
+
+    expect(result.inputTokens).toBe(42);
+  });
+
+  it("aggregates token counts across multiple records", () => {
+    const records: TokenUsageRecord[] = [
+      createRecord({ inputTokens: 100, outputTokens: 50, totalTokens: 150 }),
+      createRecord({ inputTokens: 200, outputTokens: 100, totalTokens: 300 }),
+      createRecord({ inputTokens: 50, outputTokens: 25, totalTokens: 75 }),
+    ];
+
+    const result = agentUsage("agent-multi", records);
+
+    expect(result.inputTokens).toBe(350);
+    expect(result.outputTokens).toBe(175);
+    expect(result.totalTokens).toBe(525);
+  });
+
+  it("treats undefined fields as 0 during aggregation", () => {
+    const records: TokenUsageRecord[] = [
+      createRecord({ inputTokens: 100, cacheReadTokens: undefined }),
+      createRecord({ inputTokens: undefined, cacheReadTokens: 30 }),
+    ];
+
+    const result = agentUsage("agent-mixed", records);
+
+    expect(result.inputTokens).toBe(100);
+    expect(result.cacheReadTokens).toBe(30);
+  });
+
+  it("returns zero counts for an empty array", () => {
+    const result = agentUsage("agent-empty", []);
+
+    expect(result.inputTokens).toBe(0);
+    expect(result.outputTokens).toBe(0);
+    expect(result.totalTokens).toBe(0);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// flowAgentUsage()
+// ---------------------------------------------------------------------------
+
+describe("flowAgentUsage()", () => {
+  it("groups records by source.agentId and returns per-agent usage", () => {
+    const records: TokenUsageRecord[] = [
+      createRecord({
+        inputTokens: 100,
+        outputTokens: 50,
+        totalTokens: 150,
+        source: { agentId: "agent-a", scope: [] },
+      }),
+      createRecord({
+        inputTokens: 200,
+        outputTokens: 100,
+        totalTokens: 300,
+        source: { agentId: "agent-b", scope: [] },
+      }),
+      createRecord({
+        inputTokens: 50,
+        outputTokens: 25,
+        totalTokens: 75,
+        source: { agentId: "agent-a", scope: [] },
+      }),
+    ];
+
+    const result = flowAgentUsage(records);
+
+    expect(result.usages).toHaveLength(2);
+
+    const agentA = result.usages.find((u) => u.agentId === "agent-a");
+    if (!agentA) {
+      throw new Error("Expected agent-a usage to exist");
+    }
+    expect(agentA.inputTokens).toBe(150);
+    expect(agentA.outputTokens).toBe(75);
+    expect(agentA.totalTokens).toBe(225);
+
+    const agentB = result.usages.find((u) => u.agentId === "agent-b");
+    if (!agentB) {
+      throw new Error("Expected agent-b usage to exist");
+    }
+    expect(agentB.inputTokens).toBe(200);
+    expect(agentB.outputTokens).toBe(100);
+    expect(agentB.totalTokens).toBe(300);
+  });
+
+  it('assigns records without source to "unknown" agent', () => {
+    const records: TokenUsageRecord[] = [
+      createRecord({ inputTokens: 100 }),
+      createRecord({ inputTokens: 50 }),
+    ];
+
+    const result = flowAgentUsage(records);
+
+    expect(result.usages).toHaveLength(1);
+    const unknown = result.usages[0];
+    if (!unknown) {
+      throw new Error("Expected unknown usage to exist");
+    }
+    expect(unknown.agentId).toBe("unknown");
+    expect(unknown.inputTokens).toBe(150);
+  });
+
+  it('assigns records with source but no agentId to "unknown"', () => {
+    const records: TokenUsageRecord[] = [
+      createRecord({
+        inputTokens: 75,
+        source: { agentId: undefined as unknown as string, scope: [] },
+      }),
+    ];
+
+    const result = flowAgentUsage(records);
+
+    expect(result.usages).toHaveLength(1);
+    const unknown = result.usages[0];
+    if (!unknown) {
+      throw new Error("Expected unknown usage to exist");
+    }
+    expect(unknown.agentId).toBe("unknown");
+  });
+
+  it("returns empty usages array for empty records", () => {
+    const result = flowAgentUsage([]);
+
+    expect(result.usages).toEqual([]);
+  });
+
+  it("preserves workflow and step IDs in source grouping", () => {
+    const records: TokenUsageRecord[] = [
+      createRecord({
+        inputTokens: 10,
+        source: { agentId: "agent-x", workflowId: "wf-1", stepId: "step-1", scope: ["a"] },
+      }),
+      createRecord({
+        inputTokens: 20,
+        source: { agentId: "agent-x", workflowId: "wf-1", stepId: "step-2", scope: ["b"] },
+      }),
+    ];
+
+    const result = flowAgentUsage(records);
+
+    expect(result.usages).toHaveLength(1);
+    const agentX = result.usages[0];
+    if (!agentX) {
+      throw new Error("Expected agent-x usage to exist");
+    }
+    expect(agentX.agentId).toBe("agent-x");
+    expect(agentX.inputTokens).toBe(30);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// Helpers for sumTokenUsage
+// ---------------------------------------------------------------------------
+
+const ZERO_USAGE: TokenUsage = {
+  inputTokens: 0,
+  outputTokens: 0,
+  totalTokens: 0,
+  cacheReadTokens: 0,
+  cacheWriteTokens: 0,
+  reasoningTokens: 0,
+};
+
+function createUsage(overrides?: Partial<TokenUsage>): TokenUsage {
+  return { ...ZERO_USAGE, ...overrides };
+}
+
+// ---------------------------------------------------------------------------
+// sumTokenUsage()
+// ---------------------------------------------------------------------------
+
+describe("sumTokenUsage()", () => {
+  it("sums all fields across multiple usage objects", () => {
+    const usages = [
+      createUsage({ inputTokens: 100, outputTokens: 50, totalTokens: 150 }),
+      createUsage({ inputTokens: 200, outputTokens: 100, totalTokens: 300 }),
+    ];
+
+    const result = sumTokenUsage(usages);
+
+    expect(result.inputTokens).toBe(300);
+    expect(result.outputTokens).toBe(150);
+    expect(result.totalTokens).toBe(450);
+  });
+
+  it("sums cache and reasoning token fields", () => {
+    const usages = [
+      createUsage({ cacheReadTokens: 10, cacheWriteTokens: 5, reasoningTokens: 20 }),
+      createUsage({ cacheReadTokens: 30, cacheWriteTokens: 15, reasoningTokens: 40 }),
+    ];
+
+    const result = sumTokenUsage(usages);
+
+    expect(result.cacheReadTokens).toBe(40);
+    expect(result.cacheWriteTokens).toBe(20);
+    expect(result.reasoningTokens).toBe(60);
+  });
+
+  it("returns zero usage for an empty array", () => {
+    const result = sumTokenUsage([]);
+
+    expect(result).toEqual(ZERO_USAGE);
+  });
+
+  it("returns the same values for a single-element array", () => {
+    const usage = createUsage({ inputTokens: 42, outputTokens: 21, totalTokens: 63 });
+
+    const result = sumTokenUsage([usage]);
+
+    expect(result).toEqual(usage);
+  });
+
+  it("sums three or more usage objects", () => {
+    const usages = [
+      createUsage({ inputTokens: 100 }),
+      createUsage({ inputTokens: 200 }),
+      createUsage({ inputTokens: 50 }),
+    ];
+
+    const result = sumTokenUsage(usages);
+
+    expect(result.inputTokens).toBe(350);
+  });
+
+  it("does not mutate any input", () => {
+    const a = createUsage({ inputTokens: 100 });
+    const b = createUsage({ inputTokens: 200 });
+    const aCopy = { ...a };
+    const bCopy = { ...b };
+
+    sumTokenUsage([a, b]);
+
+    expect(a).toEqual(aCopy);
+    expect(b).toEqual(bCopy);
+  });
+});