@oneuptime/common 14.0.8 → 14.0.9

package/Tests/Server/Middleware/TelemetryIngestKubernetesAgentRunnerPinnedKey.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, test, beforeEach, afterEach } from "@jest/globals";
 import TelemetryIngest, {
+  getPinnedKeyRegistrationRefusalMessage,
   TelemetryRequest,
 } from "../../../Server/Middleware/TelemetryIngest";
 import TelemetryIngestionKeyService from "../../../Server/Services/TelemetryIngestionKeyService";
@@ -16,7 +17,9 @@ import NotAuthorizedException from "../../../Types/Exception/NotAuthorizedExcept
 import ObjectID from "../../../Types/ObjectID";
 import TelemetryIngestionKeyPolicy from "../../../Types/Telemetry/TelemetryIngestionKeyPolicy";
 import TelemetryIngestionKeyType from "../../../Types/Telemetry/TelemetryIngestionKeyType";
-import TelemetryIngestSurface from "../../../Types/Telemetry/TelemetryIngestSurface";
+import TelemetryIngestSurface, {
+  IDENTITY_REGISTRATION_SURFACES,
+} from "../../../Types/Telemetry/TelemetryIngestSurface";
 
 /*
  * The Kubernetes agent Runner registration endpoint mints the Runner
@@ -30,6 +33,11 @@ import TelemetryIngestSurface from "../../../Types/Telemetry/TelemetryIngestSurf
  *   2. An unpinned server key is admitted there as before.
  *   3. A pinned key keeps working on every telemetry surface (its pin is
  *      applied when the payload is processed, as it always was).
+ *   4. The Kubernetes AI agent registration, which mints the agent key every
+ *      kubectl job for a cluster is claimed with, refuses a pinned key the
+ *      same way: the rule is keyed on IDENTITY_REGISTRATION_SURFACES. The
+ *      refusal names the identity being registered, and never the Runner
+ *      when it is the agent.
  *
  * Same module-boundary mocks as the other ingest guard suites: no Postgres,
  * no Redis.
@@ -191,7 +199,7 @@ describe("TelemetryIngest refuses a service-pinned key for Kubernetes agent Runn
     const telemetrySurfaces: Array<TelemetryIngestSurface> = (
       Object.values(TelemetryIngestSurface) as Array<TelemetryIngestSurface>
     ).filter((surface: TelemetryIngestSurface): boolean => {
-      return surface !== TelemetryIngestSurface.KubernetesAgentRunner;
+      return !IDENTITY_REGISTRATION_SURFACES.has(surface);
     });
 
     expect(telemetrySurfaces.length).toBeGreaterThan(0);
@@ -209,3 +217,91 @@ describe("TelemetryIngest refuses a service-pinned key for Kubernetes agent Runn
     }
   });
 });
+
+describe("TelemetryIngest refuses a service-pinned key for Kubernetes AI agent registration", () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+
+    (TelemetryIngestionKeyService.markUsed as MockFn).mockResolvedValue(
+      undefined as never,
+    );
+    (TelemetryIngestionKeyRateLimiter.consume as MockFn).mockResolvedValue({
+      outcome: TelemetryIngestionKeyLimitOutcome.Allowed,
+    } as never);
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  test("a pinned server key is refused 403 on the AI agent surface, before the limiter", async () => {
+    (
+      TelemetryIngestionKeyService.getPolicyFromSecretKey as MockFn
+    ).mockResolvedValue(
+      serverPolicy({ pinnedServiceName: "checkout-api" }) as never,
+    );
+
+    const result: RunResult = await run(
+      TelemetryIngestSurface.KubernetesAiAgent,
+    );
+
+    expect(result.next).not.toHaveBeenCalled();
+    expect(
+      TelemetryIngestionKeyRateLimiter.consume as MockFn,
+    ).not.toHaveBeenCalled();
+    expect(
+      TelemetryIngestionKeyService.markUsed as MockFn,
+    ).not.toHaveBeenCalled();
+
+    const sendErrorResponse: MockFn = Response.sendErrorResponse as MockFn;
+    expect(sendErrorResponse).toHaveBeenCalledTimes(1);
+
+    const error: Error = sendErrorResponse.mock.calls[0]?.[2] as Error;
+    expect(error).toBeInstanceOf(NotAuthorizedException);
+    expect(error.message).toBe(
+      getPinnedKeyRegistrationRefusalMessage(
+        TelemetryIngestSurface.KubernetesAiAgent,
+      ),
+    );
+    expect(error.message).toContain("pinned to a single service");
+    expect(error.message).toContain("Kubernetes AI agent registration");
+    expect(error.message).toContain("unpinned server ingestion key");
+    expect(error.message).not.toContain("Runner");
+    expect(error.message).not.toContain(SENTINEL_TOKEN);
+    expect(error.message).not.toContain("checkout-api");
+    expect((result.req as TelemetryRequest).projectId).toBeUndefined();
+  });
+
+  test("negative control: an unpinned server key reaches the AI agent registration", async () => {
+    (
+      TelemetryIngestionKeyService.getPolicyFromSecretKey as MockFn
+    ).mockResolvedValue(serverPolicy({ pinnedServiceName: null }) as never);
+
+    const result: RunResult = await run(
+      TelemetryIngestSurface.KubernetesAiAgent,
+    );
+
+    expect(Response.sendErrorResponse as MockFn).not.toHaveBeenCalled();
+    expect(result.next).toHaveBeenCalledTimes(1);
+    expect(
+      (
+        result.req as TelemetryRequest
+      ).ingestionKeyPolicy.ingestionKeyId.toString(),
+    ).toBe(INGESTION_KEY_ID);
+  });
+
+  test.each(Array.from(IDENTITY_REGISTRATION_SURFACES))(
+    "the refusal on %s names that surface's identity and says what to use instead",
+    (surface: TelemetryIngestSurface) => {
+      const message: string = getPinnedKeyRegistrationRefusalMessage(surface);
+
+      expect(message).toContain("pinned to a single service");
+      expect(message).toContain("unpinned server ingestion key");
+      expect(message).toContain(
+        surface === TelemetryIngestSurface.KubernetesAiAgent
+          ? "Kubernetes AI agent registration"
+          : "Kubernetes agent Runner registration",
+      );
+    },
+  );
+});

package/Tests/Server/Middleware/TelemetryIngestKubernetesAgentRunnerRateLimit.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, test, beforeEach, afterEach } from "@jest/globals";
 import TelemetryIngest, {
+  DEFAULT_IDENTITY_REGISTRATION_REQUESTS_PER_MINUTE,
   DEFAULT_KUBERNETES_AGENT_RUNNER_REQUESTS_PER_MINUTE,
   getEffectiveRequestsPerMinuteLimit,
   TelemetryRequest,
@@ -21,7 +22,9 @@ import TelemetryIngestionKeyPolicy, {
   DEFAULT_BROWSER_KEY_REQUESTS_PER_MINUTE,
 } from "../../../Types/Telemetry/TelemetryIngestionKeyPolicy";
 import TelemetryIngestionKeyType from "../../../Types/Telemetry/TelemetryIngestionKeyType";
-import TelemetryIngestSurface from "../../../Types/Telemetry/TelemetryIngestSurface";
+import TelemetryIngestSurface, {
+  IDENTITY_REGISTRATION_SURFACES,
+} from "../../../Types/Telemetry/TelemetryIngestSurface";
 
 /*
  * The Kubernetes agent Runner registration endpoint is reached with the
@@ -41,6 +44,10 @@ import TelemetryIngestSurface from "../../../Types/Telemetry/TelemetryIngestSurf
  *      no limiter call.
  *   4. A registration request over the default is refused 429 like any other
  *      rate-limited request, Retry-After included.
+ *   5. The Kubernetes AI agent registration, which mints the agent key the
+ *      same way, gets exactly the same default: the rule is keyed on
+ *      IDENTITY_REGISTRATION_SURFACES, not on one enum member, so a new
+ *      identity surface cannot quietly go unlimited.
  *
  * Same module-boundary mocks as the browser key guard suite: no Postgres, no
  * Redis.
@@ -95,9 +102,13 @@ const ALL_SURFACES: Array<TelemetryIngestSurface> = Object.values(
 
 const NON_REGISTRATION_SURFACES: Array<TelemetryIngestSurface> =
   ALL_SURFACES.filter((surface: TelemetryIngestSurface): boolean => {
-    return surface !== TelemetryIngestSurface.KubernetesAgentRunner;
+    return !IDENTITY_REGISTRATION_SURFACES.has(surface);
   });
 
+const REGISTRATION_SURFACES: Array<TelemetryIngestSurface> = Array.from(
+  IDENTITY_REGISTRATION_SURFACES,
+);
+
 type BuildPolicyFunction = (
   overrides: Partial<TelemetryIngestionKeyPolicy>,
 ) => TelemetryIngestionKeyPolicy;
@@ -356,3 +367,95 @@ describe("TelemetryIngest Kubernetes agent Runner registration rate limit", () =
     });
   });
 });
+
+describe("TelemetryIngest identity registration surfaces share one rate limit rule", () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+
+    (TelemetryIngestionKeyService.markUsed as MockFn).mockResolvedValue(
+      undefined as never,
+    );
+
+    limiterAnswers(TelemetryIngestionKeyLimitOutcome.Allowed);
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  test("the identity surfaces are exactly the Runner and the Kubernetes AI agent registrations", () => {
+    expect([...REGISTRATION_SURFACES].sort()).toEqual(
+      [
+        TelemetryIngestSurface.KubernetesAgentRunner,
+        TelemetryIngestSurface.KubernetesAiAgent,
+      ].sort(),
+    );
+    expect(
+      NON_REGISTRATION_SURFACES.length + REGISTRATION_SURFACES.length,
+    ).toBe(ALL_SURFACES.length);
+  });
+
+  test("the shipped default is one number under both names", () => {
+    expect(DEFAULT_IDENTITY_REGISTRATION_REQUESTS_PER_MINUTE).toBe(30);
+    expect(DEFAULT_KUBERNETES_AGENT_RUNNER_REQUESTS_PER_MINUTE).toBe(
+      DEFAULT_IDENTITY_REGISTRATION_REQUESTS_PER_MINUTE,
+    );
+  });
+
+  test("a server key with no configured limit gets the registration default on the Kubernetes AI agent surface", () => {
+    expect(
+      getEffectiveRequestsPerMinuteLimit(
+        buildServerPolicy({ requestsPerMinuteLimit: null }),
+        TelemetryIngestSurface.KubernetesAiAgent,
+      ),
+    ).toBe(DEFAULT_IDENTITY_REGISTRATION_REQUESTS_PER_MINUTE);
+
+    // An explicit per-key limit still wins there.
+    expect(
+      getEffectiveRequestsPerMinuteLimit(
+        buildServerPolicy({ requestsPerMinuteLimit: 3 }),
+        TelemetryIngestSurface.KubernetesAiAgent,
+      ),
+    ).toBe(3);
+  });
+
+  test.each(REGISTRATION_SURFACES)(
+    "an unlimited server key IS rate limited at 30/min through the middleware on %s",
+    async (surface: TelemetryIngestSurface) => {
+      resolveTo(buildServerPolicy({ requestsPerMinuteLimit: null }));
+
+      const result: RunResult = await run(surface);
+
+      expect(
+        TelemetryIngestionKeyRateLimiter.consume as MockFn,
+      ).toHaveBeenCalledTimes(1);
+      expect(firstConsumeCallArgs()["limitPerMinute"]).toBe(30);
+      expect(firstConsumeCallArgs()["ingestionKeyId"]?.toString()).toBe(
+        INGESTION_KEY_ID,
+      );
+      expect(Response.sendErrorResponse as MockFn).not.toHaveBeenCalled();
+      expect(result.next).toHaveBeenCalledTimes(1);
+    },
+  );
+
+  test("a Kubernetes AI agent registration over the default is refused 429 with Retry-After", async () => {
+    resolveTo(buildServerPolicy({ requestsPerMinuteLimit: null }));
+    limiterAnswers(TelemetryIngestionKeyLimitOutcome.RateLimited, {
+      retryAfterSeconds: 17,
+      isFirstRejectionInWindow: false,
+    });
+
+    const result: RunResult = await run(
+      TelemetryIngestSurface.KubernetesAiAgent,
+    );
+
+    expect(result.next).not.toHaveBeenCalled();
+
+    const error: Error = (Response.sendErrorResponse as MockFn).mock
+      .calls[0]?.[2] as Error;
+    expect(error).toBeInstanceOf(TooManyRequestsException);
+    expect(error.message).not.toContain(SENTINEL_TOKEN);
+    expect(result.setHeader).toHaveBeenCalledWith("Retry-After", "17");
+    expect((result.req as TelemetryRequest).projectId).toBeUndefined();
+  });
+});

package/Tests/Server/Services/AIServiceBalanceBlocker.test.ts

@@ -0,0 +1,282 @@
+import AIService, {
+  AI_BALANCE_INSUFFICIENT_MESSAGE,
+  AI_DISABLED_MESSAGE,
+} from "../../../Server/Services/AIService";
+import KubernetesClusterAiAccessService, {
+  KubernetesClusterAiAccessProjectGates,
+} from "../../../Server/Services/KubernetesClusterAiAccessService";
+import LlmProviderService from "../../../Server/Services/LlmProviderService";
+import ProjectService from "../../../Server/Services/ProjectService";
+import LlmProvider from "../../../Models/DatabaseModels/LlmProvider";
+import Project from "../../../Models/DatabaseModels/Project";
+import ObjectID from "../../../Types/ObjectID";
+import { afterEach, beforeEach, describe, expect, it } from "@jest/globals";
+
+/*
+ * Contract under test — AIService.getAiBalanceBlocker, the ONE predicate
+ * behind "out of AI credits": the cluster AI status's ai_balance_insufficient
+ * gap and the investigation not-started reason both ask it, so the two never
+ * disagree with each other or with the refusal executeWithLogging makes on
+ * the call itself.
+ *
+ * A blocker exists only when ALL of these hold: billing is on, the
+ * project's resolved LLM provider is the OneUptime-hosted (global) one with
+ * a per-token cost, the balance is at or below zero, and auto-recharge is
+ * off (a project with auto-recharge is topped up when it runs low).
+ * Everything else — a free global provider, the project's own provider, no
+ * provider at all (its own gap), a positive balance — is no blocker.
+ */
+
+type MockBillingGlobal = typeof globalThis & {
+  __oneuptimeBalanceBlockerTestBillingEnabled: boolean;
+};
+
+jest.mock("../../../Server/EnvironmentConfig", () => {
+  const actual: Record<string, unknown> = jest.requireActual(
+    "../../../Server/EnvironmentConfig",
+  ) as Record<string, unknown>;
+  const mocked: Record<string, unknown> = { ...actual };
+  const mockGlobal: MockBillingGlobal = globalThis as MockBillingGlobal;
+  mockGlobal.__oneuptimeBalanceBlockerTestBillingEnabled = true;
+
+  Object.defineProperty(mocked, "IsBillingEnabled", {
+    configurable: true,
+    enumerable: true,
+    get: (): boolean => {
+      return mockGlobal.__oneuptimeBalanceBlockerTestBillingEnabled;
+    },
+  });
+
+  return mocked;
+});
+
+function setBillingEnabled(value: boolean): void {
+  (
+    globalThis as MockBillingGlobal
+  ).__oneuptimeBalanceBlockerTestBillingEnabled = value;
+}
+
+const PROJECT_ID: ObjectID = new ObjectID(
+  "11111111-1111-4111-8111-111111111111",
+);
+
+function provider(overrides: Record<string, unknown> = {}): LlmProvider {
+  return {
+    id: ObjectID.generate(),
+    isGlobalLlm: true,
+    costPerMillionTokensInUSDCents: 500,
+    ...overrides,
+  } as unknown as LlmProvider;
+}
+
+function project(overrides: Record<string, unknown> = {}): Project {
+  return {
+    aiCurrentBalanceInUSDCents: 0,
+    enableAutoRechargeAiBalance: false,
+    ...overrides,
+  } as unknown as Project;
+}
+
+describe("AIService.getAiBalanceBlocker", () => {
+  let providerLookup: jest.SpyInstance;
+  let projectLookup: jest.SpyInstance;
+
+  beforeEach(() => {
+    setBillingEnabled(true);
+    providerLookup = jest
+      .spyOn(LlmProviderService, "getLLMProviderForProject")
+      .mockResolvedValue(provider());
+    projectLookup = jest
+      .spyOn(ProjectService, "findOneById")
+      .mockResolvedValue(project());
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+    setBillingEnabled(true);
+  });
+
+  it("blocks when billing is on, the provider is the paid global one, the balance is empty and auto-recharge is off", async () => {
+    expect(await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID })).toBe(
+      AI_BALANCE_INSUFFICIENT_MESSAGE,
+    );
+  });
+
+  it("blocks a negative balance too", async () => {
+    projectLookup.mockResolvedValue(
+      project({ aiCurrentBalanceInUSDCents: -250 }),
+    );
+
+    expect(await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID })).toBe(
+      AI_BALANCE_INSUFFICIENT_MESSAGE,
+    );
+  });
+
+  it("blocks a balance that was never set (read as zero, like the call itself reads it)", async () => {
+    projectLookup.mockResolvedValue(
+      project({ aiCurrentBalanceInUSDCents: undefined }),
+    );
+
+    expect(await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID })).toBe(
+      AI_BALANCE_INSUFFICIENT_MESSAGE,
+    );
+  });
+
+  it("reads the balance and auto-recharge of THIS project, as root", async () => {
+    await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID });
+
+    expect(projectLookup).toHaveBeenCalledWith({
+      id: PROJECT_ID,
+      select: {
+        aiCurrentBalanceInUSDCents: true,
+        enableAutoRechargeAiBalance: true,
+      },
+      props: { isRoot: true },
+    });
+    expect(providerLookup).toHaveBeenCalledWith(PROJECT_ID);
+  });
+
+  it("no blocker when billing is off (self-hosted): nothing is metered", async () => {
+    setBillingEnabled(false);
+
+    expect(
+      await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID }),
+    ).toBeNull();
+    expect(providerLookup).not.toHaveBeenCalled();
+    expect(projectLookup).not.toHaveBeenCalled();
+  });
+
+  it("no blocker when auto-recharge is on", async () => {
+    projectLookup.mockResolvedValue(
+      project({ enableAutoRechargeAiBalance: true }),
+    );
+
+    expect(
+      await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID }),
+    ).toBeNull();
+  });
+
+  it("no blocker while the balance is positive", async () => {
+    projectLookup.mockResolvedValue(project({ aiCurrentBalanceInUSDCents: 1 }));
+
+    expect(
+      await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID }),
+    ).toBeNull();
+  });
+
+  it("no blocker for the project's own provider: it bills the project's own account", async () => {
+    providerLookup.mockResolvedValue(provider({ isGlobalLlm: false }));
+
+    expect(
+      await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID }),
+    ).toBeNull();
+    expect(projectLookup).not.toHaveBeenCalled();
+  });
+
+  it("no blocker for a free global provider (no per-token cost)", async () => {
+    for (const cost of [0, undefined, null]) {
+      providerLookup.mockResolvedValue(
+        provider({ costPerMillionTokensInUSDCents: cost }),
+      );
+
+      expect(
+        await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID }),
+      ).toBeNull();
+    }
+  });
+
+  it("no blocker when there is no provider at all: that is its own gap", async () => {
+    providerLookup.mockResolvedValue(null);
+
+    expect(
+      await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID }),
+    ).toBeNull();
+  });
+
+  it("no blocker for a project row that cannot be found", async () => {
+    projectLookup.mockResolvedValue(null);
+
+    expect(
+      await AIService.getAiBalanceBlocker({ projectId: PROJECT_ID }),
+    ).toBeNull();
+  });
+
+  it("uses a provider the caller already resolved instead of looking it up again", async () => {
+    expect(
+      await AIService.getAiBalanceBlocker({
+        projectId: PROJECT_ID,
+        llmProvider: provider(),
+      }),
+    ).toBe(AI_BALANCE_INSUFFICIENT_MESSAGE);
+    expect(
+      await AIService.getAiBalanceBlocker({
+        projectId: PROJECT_ID,
+        llmProvider: null,
+      }),
+    ).toBeNull();
+    expect(providerLookup).not.toHaveBeenCalled();
+  });
+
+  it("the reason is one plain sentence a person can act on", () => {
+    expect(AI_BALANCE_INSUFFICIENT_MESSAGE).toBe(
+      "This project's AI credit balance is used up and auto-recharge is off, so OneUptime AI cannot run.",
+    );
+  });
+});
+
+/*
+ * The cluster AI status resolves the project's provider for its
+ * llm_provider_missing gap and hands it to this predicate, so a status
+ * build looks the provider up once. Billing is on here, so the predicate
+ * really runs: a second lookup inside it would show as a second call.
+ */
+describe("KubernetesClusterAiAccessService.getProjectGates and the balance predicate", () => {
+  let providerLookup: jest.SpyInstance;
+
+  beforeEach(() => {
+    setBillingEnabled(true);
+    providerLookup = jest
+      .spyOn(LlmProviderService, "getLLMProviderForProject")
+      .mockResolvedValue(provider());
+    jest.spyOn(ProjectService, "findOneById").mockResolvedValue(project());
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+    setBillingEnabled(true);
+  });
+
+  it("looks the provider up once and still reports the empty balance", async () => {
+    const gates: KubernetesClusterAiAccessProjectGates =
+      await KubernetesClusterAiAccessService.getProjectGates(PROJECT_ID);
+
+    expect(providerLookup).toHaveBeenCalledTimes(1);
+    expect(gates.hasLlmProvider).toBe(true);
+    expect(gates.aiBalanceBlocker).toBe(AI_BALANCE_INSUFFICIENT_MESSAGE);
+  });
+
+  it("negative control: a funded balance is no blocker, still with one lookup", async () => {
+    (
+      ProjectService.findOneById as unknown as jest.SpyInstance
+    ).mockResolvedValue(project({ aiCurrentBalanceInUSDCents: 5000 }));
+
+    const gates: KubernetesClusterAiAccessProjectGates =
+      await KubernetesClusterAiAccessService.getProjectGates(PROJECT_ID);
+
+    expect(providerLookup).toHaveBeenCalledTimes(1);
+    expect(gates.aiBalanceBlocker).toBeNull();
+  });
+});
+
+describe("AI_DISABLED_MESSAGE", () => {
+  /*
+   * The Enable AI switch lives on the always-visible Project Settings → AI
+   * Features page now; AI Credits is hidden on self-hosted installs.
+   */
+  it("points at Project Settings → AI Features, never AI Credits", () => {
+    expect(AI_DISABLED_MESSAGE).toBe(
+      "AI features are disabled for this project. Enable them in Project Settings → AI Features.",
+    );
+    expect(AI_DISABLED_MESSAGE).not.toContain("AI Credits");
+  });
+});