SysDeck/klanker-gate/tests/integration/pricing_hierarchy_test.ts

// Wave-3 governance: pricing catalog + $-cost budgets, token-window limits,
// and the teams/customers hierarchy (fail-closed admission up the chain).

import { assert, assertEquals } from "@std/assert";
import { createHandler } from "../../apps/gateway/main.ts";
import {
  type AppContext,
  NullToolExecutor,
  VERSION,
} from "../../apps/gateway/context.ts";
import { ProviderManager } from "../../packages/providers/src/mod.ts";
import { Metrics } from "../../packages/telemetry/src/metrics.ts";
import { LogBus } from "../../packages/telemetry/src/logbus.ts";
import { VirtualKeyManager } from "../../packages/governance/src/virtual_keys.ts";
import { GovernanceHierarchy } from "../../packages/governance/src/hierarchy.ts";
import { PricingCatalog } from "../../packages/governance/src/pricing.ts";
import { MCPRegistry } from "../../packages/mcp/src/registry.ts";
import { PluginManager } from "../../packages/plugins/src/lifecycle.ts";
import {
  jsonResponse,
  MockProvider,
  openAIChatBody,
} from "../../packages/testing/src/mod.ts";

function makeContext(mockUrl: string): AppContext {
  return {
    providers: new ProviderManager([{
      id: "openai",
      type: "openai",
      apiKey: "sk-test",
      baseUrl: mockUrl,
      enabled: true,
      models: ["gpt-4o"],
      priority: 0,
      retry: { maxRetries: 0 },
    }], "openai"),
    metrics: new Metrics(),
    logBus: new LogBus(),
    virtualKeys: new VirtualKeyManager(),
    hierarchy: new GovernanceHierarchy(),
    pricing: new PricingCatalog(),
    mcp: new MCPRegistry(),
    plugins: new PluginManager(),
    toolExecutor: new NullToolExecutor(),
    version: VERSION,
  };
}

function chatRequest(token: string, content = "hei"): Request {
  return new Request("http://gateway.test/v1/chat/completions", {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "Authorization": `Bearer ${token}`,
    },
    body: JSON.stringify({
      model: "openai/gpt-4o",
      messages: [{ role: "user", content }],
    }),
  });
}

const usageBody = openAIChatBody("ok", {
  model: "gpt-4o",
  usage: { prompt_tokens: 1000, completion_tokens: 500, total_tokens: 1500 },
});

Deno.test("pricing: integer micro-USD math and prefix matching", () => {
  const pricing = new PricingCatalog();
  // 1000 prompt @ $2.5/M + 500 completion @ $10/M = 2500 + 5000 micro-USD.
  assertEquals(
    pricing.costMicroUsd("gpt-4o", {
      prompt_tokens: 1000,
      completion_tokens: 500,
    }),
    7500,
  );
  // Dated variants resolve through the longest configured prefix.
  assertEquals(
    pricing.costMicroUsd("gpt-4o-2026-01-01", { prompt_tokens: 1000 }),
    2500,
  );
  // Provider-prefixed ids resolve too; unknown models are unpriced.
  assertEquals(
    pricing.costMicroUsd("openai/gpt-4o", { prompt_tokens: 1000 }),
    2500,
  );
  assertEquals(
    pricing.costMicroUsd("unknown-model", { prompt_tokens: 5 }),
    null,
  );
});

Deno.test("cost budget: exhausted keys get 402 and metrics count cost", async () => {
  const mock = new MockProvider(() => jsonResponse(usageBody));
  const ctx = makeContext(mock.url);
  ctx.virtualKeys.upsert({
    id: "vk-cost",
    name: "cost-capped",
    token: "vk-cost-budget-token",
    enabled: true,
    budget: { maxCostUsd: 0.005 }, // 5000 micro; one gpt-4o call costs 7500
    usedRequests: 0,
    usedCostMicroUsd: 0,
  });
  const handler = createHandler(ctx);
  try {
    const first = await handler(chatRequest("vk-cost-budget-token"));
    assertEquals(first.status, 200);
    await first.body?.cancel();
    assertEquals(ctx.virtualKeys.get("vk-cost")!.usedCostMicroUsd, 7500);
    assert(
      ctx.metrics.renderPrometheus().includes(
        "frosty_cost_usd_total 0.007500",
      ),
    );

    const second = await handler(chatRequest("vk-cost-budget-token"));
    assertEquals(second.status, 402);
    const body = await second.json() as { error: { code: string } };
    assertEquals(body.error.code, "cost_budget_exhausted");
  } finally {
    await mock.close();
  }
});

Deno.test("token limit: request estimates meter the window", async () => {
  const mock = new MockProvider(() => jsonResponse(openAIChatBody("ok")));
  const ctx = makeContext(mock.url);
  ctx.virtualKeys.upsert({
    id: "vk-tok",
    name: "token-capped",
    token: "vk-token-limit-token",
    enabled: true,
    tokenLimit: { maxTokens: 100, windowMs: 60_000 },
    usedRequests: 0,
    usedCostMicroUsd: 0,
  });
  const handler = createHandler(ctx);
  try {
    // ~380-char body -> est ~95 tokens: first passes, second exceeds 100.
    const first = await handler(
      chatRequest("vk-token-limit-token", "x".repeat(300)),
    );
    assertEquals(first.status, 200);
    await first.body?.cancel();

    const second = await handler(
      chatRequest("vk-token-limit-token", "x".repeat(300)),
    );
    assertEquals(second.status, 429);
    const body = await second.json() as { error: { code: string } };
    assertEquals(body.error.code, "token_limited");
    assert(second.headers.get("Retry-After"));
  } finally {
    await mock.close();
  }
});

Deno.test("hierarchy: disabled/exhausted parents lock descendants out", async () => {
  const mock = new MockProvider(() => jsonResponse(openAIChatBody("ok")));
  const ctx = makeContext(mock.url);
  ctx.hierarchy!.upsertCustomer({
    id: "cust-1",
    name: "acme",
    enabled: true,
    budget: { maxRequests: 1 },
    usedRequests: 0,
    usedCostMicroUsd: 0,
  });
  ctx.hierarchy!.upsertTeam({
    id: "team-1",
    name: "core",
    enabled: true,
    customerId: "cust-1",
    usedRequests: 0,
    usedCostMicroUsd: 0,
  });
  ctx.virtualKeys.upsert({
    id: "vk-h",
    name: "hier-key",
    token: "vk-hierarchy-token1",
    enabled: true,
    teamId: "team-1",
    usedRequests: 0,
    usedCostMicroUsd: 0,
  });
  const handler = createHandler(ctx);
  try {
    // Customer allows exactly one request across all descendants.
    const first = await handler(chatRequest("vk-hierarchy-token1"));
    assertEquals(first.status, 200);
    await first.body?.cancel();

    const second = await handler(chatRequest("vk-hierarchy-token1"));
    assertEquals(second.status, 402);
    assertEquals(
      ((await second.json()) as { error: { code: string } }).error.code,
      "customer_budget_exhausted",
    );

    // Disabling the team denies immediately (fail closed).
    const team = ctx.hierarchy!.getTeam("team-1")!;
    ctx.hierarchy!.upsertTeam({ ...team, enabled: false });
    const third = await handler(chatRequest("vk-hierarchy-token1"));
    assertEquals(third.status, 401);
    assertEquals(
      ((await third.json()) as { error: { code: string } }).error.code,
      "team_disabled",
    );

    // A dangling team reference is a lockout, not an open door.
    ctx.virtualKeys.upsert({
      ...ctx.virtualKeys.get("vk-h")!,
      teamId: "no-such-team",
    });
    const fourth = await handler(chatRequest("vk-hierarchy-token1"));
    assertEquals(fourth.status, 401);
    assertEquals(
      ((await fourth.json()) as { error: { code: string } }).error.code,
      "invalid_team",
    );
  } finally {
    await mock.close();
  }
});

Deno.test("teams/customers/pricing admin APIs round-trip", async () => {
  const mock = new MockProvider(() => jsonResponse(openAIChatBody("ok")));
  const ctx = makeContext(mock.url);
  const handler = createHandler(ctx);
  try {
    const customer = await (await handler(
      new Request("http://gateway.test/api/customers", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ name: "acme", budget: { maxCostUsd: 10 } }),
      }),
    )).json() as { id: string };

    const team = await (await handler(
      new Request("http://gateway.test/api/teams", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ name: "core", customerId: customer.id }),
      }),
    )).json() as { id: string; customerId: string };
    assertEquals(team.customerId, customer.id);

    const teams = await (await handler(
      new Request("http://gateway.test/api/teams"),
    )).json() as { teams: Array<{ id: string }> };
    assertEquals(teams.teams.length, 1);

    const del = await handler(
      new Request(`http://gateway.test/api/teams/${team.id}`, {
        method: "DELETE",
      }),
    );
    assertEquals(del.status, 204);

    // Pricing: PUT replaces the catalog and GET reads it back.
    const put = await handler(
      new Request("http://gateway.test/api/pricing", {
        method: "PUT",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          "my-model": { inputPerMTokUsd: 1, outputPerMTokUsd: 2 },
        }),
      }),
    );
    assertEquals(put.status, 200);
    const pricing = await (await handler(
      new Request("http://gateway.test/api/pricing"),
    )).json() as { prices: Record<string, unknown> };
    assertEquals(Object.keys(pricing.prices), ["my-model"]);
    assertEquals(
      ctx.pricing!.costMicroUsd("my-model", {
        prompt_tokens: 1000,
        completion_tokens: 1000,
      }),
      3000,
    );
  } finally {
    await mock.close();
  }
});