@oneuptime/common 12.0.4 → 12.0.6

package/Tests/Server/Utils/AI/AIInvestigationQueue.test.ts

@@ -6,9 +6,16 @@ import AIAlertInvestigationRunner from "../../../../Server/Utils/AI/SRE/AlertInv
 import AIRunService from "../../../../Server/Services/AIRunService";
 import AIService from "../../../../Server/Services/AIService";
 import ProjectService from "../../../../Server/Services/ProjectService";
+import Semaphore from "../../../../Server/Infrastructure/Semaphore";
+import {
+  getInvestigationSubjectLockKey,
+  INVESTIGATION_SUBJECT_LOCK_NAMESPACE,
+} from "../../../../Server/Utils/AI/SRE/InvestigationSubjectLock";
 import Project from "../../../../Models/DatabaseModels/Project";
 import AIRun from "../../../../Models/DatabaseModels/AIRun";
+import AIRunCodeFixRecommendation from "../../../../Types/AI/AIRunCodeFixRecommendation";
 import AIRunStatus from "../../../../Types/AI/AIRunStatus";
+import AIRunType from "../../../../Types/AI/AIRunType";
 import ObjectID from "../../../../Types/ObjectID";
 import PositiveNumber from "../../../../Types/PositiveNumber";
 import { describe, expect, test, afterEach, beforeEach } from "@jest/globals";
@@ -39,10 +46,18 @@ function mockBudgetOk(): void {
   });
 }
 
+function findOperatorSql(value: unknown): string {
+  const operator: { getSql?: ((alias: string) => string) | undefined } =
+    value as { getSql?: ((alias: string) => string) | undefined };
+  return operator.getSql?.("subject") || "";
+}
+
 describe("AIInvestigationQueue", () => {
   beforeEach(() => {
     mockBudgetOk();
-    // No per-project cap override => default of 3.
+    jest.spyOn(Semaphore, "lock").mockResolvedValue({} as never);
+    jest.spyOn(Semaphore, "release").mockResolvedValue();
+    // No lane cap override => default of 3.
     jest
       .spyOn(ProjectService, "findOneById")
       .mockResolvedValue({ id: ObjectID.generate() } as unknown as Project);
@@ -78,9 +93,27 @@ describe("AIInvestigationQueue", () => {
         data: expect.objectContaining({
           status: AIRunStatus.Queued,
           triggeredByIncidentId: incidentId,
+          codeFixRecommendation: AIRunCodeFixRecommendation.NotRecommended,
         }),
       }),
     );
+    expect(AIService.getAutonomousDailyBudgetStatus).toHaveBeenCalledWith(
+      projectId,
+      { incidentId, alertId: undefined },
+    );
+    expect(Semaphore.lock).toHaveBeenCalledWith({
+      key: getInvestigationSubjectLockKey({ projectId, incidentId }),
+      namespace: INVESTIGATION_SUBJECT_LOCK_NAMESPACE,
+      lockTimeout: 30 * 1000,
+      acquireTimeout: 10 * 1000,
+    });
+    expect(Semaphore.release).toHaveBeenCalledTimes(1);
+    expect(create.mock.invocationCallOrder[0]).toBeGreaterThan(
+      (Semaphore.lock as unknown as jest.Mock).mock.invocationCallOrder[0]!,
+    );
+    expect(
+      (Semaphore.release as unknown as jest.Mock).mock.invocationCallOrder[0],
+    ).toBeGreaterThan(create.mock.invocationCallOrder[0]!);
     // The inline kick is detached; give the microtask a beat.
     await new Promise((resolve: (value: unknown) => void) => {
       setTimeout(resolve, 0);
@@ -92,6 +125,117 @@ describe("AIInvestigationQueue", () => {
         triggeredByIncidentId: incidentId,
       }),
     );
+    expect(processRun.mock.invocationCallOrder[0]).toBeGreaterThan(
+      (Semaphore.release as unknown as jest.Mock).mock.invocationCallOrder[0]!,
+    );
+  });
+
+  test("enqueue keeps an alert investigation NotRecommended until completion wins", async () => {
+    const projectId: ObjectID = ObjectID.generate();
+    const alertId: ObjectID = ObjectID.generate();
+
+    const create: jest.SpyInstance = jest
+      .spyOn(AIRunService, "create")
+      .mockResolvedValue({ id: ObjectID.generate() } as unknown as AIRun);
+    jest.spyOn(AIInvestigationQueue, "processRun").mockResolvedValue(undefined);
+
+    await AIInvestigationQueue.enqueue({
+      projectId,
+      subjectAlertId: alertId,
+    });
+
+    expect(create).toHaveBeenCalledWith(
+      expect.objectContaining({
+        data: expect.objectContaining({
+          runType: AIRunType.Investigation,
+          triggeredByAlertId: alertId,
+          codeFixRecommendation: AIRunCodeFixRecommendation.NotRecommended,
+        }),
+      }),
+    );
+    expect(Semaphore.lock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        key: getInvestigationSubjectLockKey({ projectId, alertId }),
+        namespace: INVESTIGATION_SUBJECT_LOCK_NAMESPACE,
+      }),
+    );
+    expect(Semaphore.release).toHaveBeenCalledTimes(1);
+  });
+
+  test("a subject Investigation fails closed when the shared lock cannot be acquired", async () => {
+    (Semaphore.lock as unknown as jest.Mock).mockRejectedValue(
+      new Error("Redis unavailable"),
+    );
+    const create: jest.SpyInstance = jest.spyOn(AIRunService, "create");
+    const processRun: jest.SpyInstance = jest.spyOn(
+      AIInvestigationQueue,
+      "processRun",
+    );
+
+    const runId: ObjectID | null = await AIInvestigationQueue.enqueue({
+      projectId: ObjectID.generate(),
+      subjectIncidentId: ObjectID.generate(),
+    });
+
+    expect(runId).toBeNull();
+    expect(create).not.toHaveBeenCalled();
+    expect(Semaphore.release).not.toHaveBeenCalled();
+    expect(processRun).not.toHaveBeenCalled();
+  });
+
+  test("enqueue rejects a run carrying both subject types before any budget or write", async () => {
+    const budget: jest.SpyInstance = jest.spyOn(
+      AIService,
+      "getAutonomousDailyBudgetStatus",
+    );
+    const create: jest.SpyInstance = jest.spyOn(AIRunService, "create");
+
+    const runId: ObjectID | null = await AIInvestigationQueue.enqueue({
+      projectId: ObjectID.generate(),
+      subjectIncidentId: ObjectID.generate(),
+      subjectAlertId: ObjectID.generate(),
+    });
+
+    expect(runId).toBeNull();
+    expect(budget).not.toHaveBeenCalled();
+    expect(create).not.toHaveBeenCalled();
+  });
+
+  test("enqueue forwards the alert lane to the budget check", async () => {
+    const projectId: ObjectID = ObjectID.generate();
+    const alertId: ObjectID = ObjectID.generate();
+    jest.spyOn(AIRunService, "create").mockResolvedValue({
+      id: ObjectID.generate(),
+    } as unknown as AIRun);
+    jest.spyOn(AIInvestigationQueue, "processRun").mockResolvedValue(undefined);
+
+    await AIInvestigationQueue.enqueue({
+      projectId,
+      subjectAlertId: alertId,
+    });
+
+    expect(AIService.getAutonomousDailyBudgetStatus).toHaveBeenCalledWith(
+      projectId,
+      { incidentId: undefined, alertId },
+    );
+  });
+
+  test("enqueue leaves insight work in the legacy subjectless budget lane", async () => {
+    const projectId: ObjectID = ObjectID.generate();
+    jest.spyOn(AIRunService, "create").mockResolvedValue({
+      id: ObjectID.generate(),
+    } as unknown as AIRun);
+    jest.spyOn(AIInvestigationQueue, "processRun").mockResolvedValue(undefined);
+
+    await AIInvestigationQueue.enqueue({
+      projectId,
+      subjectAIInsightId: ObjectID.generate(),
+    });
+
+    expect(AIService.getAutonomousDailyBudgetStatus).toHaveBeenCalledWith(
+      projectId,
+      { incidentId: undefined, alertId: undefined },
+    );
   });
 
   test("a lost claim never executes the investigation", async () => {
@@ -111,8 +255,33 @@ describe("AIInvestigationQueue", () => {
     expect(execute).not.toHaveBeenCalled();
   });
 
+  test("a queued run carrying both subject types is never counted or claimed", async () => {
+    const findProject: jest.SpyInstance = jest.spyOn(
+      ProjectService,
+      "findOneById",
+    );
+    const countBy: jest.SpyInstance = jest.spyOn(AIRunService, "countBy");
+    const claim: jest.SpyInstance = jest.spyOn(
+      AIRunService,
+      "attemptStatusTransition",
+    );
+
+    await AIInvestigationQueue.processRun({
+      id: ObjectID.generate(),
+      projectId: ObjectID.generate(),
+      attemptCount: 0,
+      triggeredByIncidentId: ObjectID.generate(),
+      triggeredByAlertId: ObjectID.generate(),
+    });
+
+    expect(findProject).not.toHaveBeenCalled();
+    expect(countBy).not.toHaveBeenCalled();
+    expect(claim).not.toHaveBeenCalled();
+  });
+
   test("a won claim increments attemptCount and dispatches to the incident runner", async () => {
     const runId: ObjectID = ObjectID.generate();
+    const projectId: ObjectID = ObjectID.generate();
     const incidentId: ObjectID = ObjectID.generate();
     const claim: jest.SpyInstance = jest
       .spyOn(AIRunService, "attemptStatusTransition")
@@ -123,7 +292,7 @@ describe("AIInvestigationQueue", () => {
 
     await AIInvestigationQueue.processRun({
       id: runId,
-      projectId: ObjectID.generate(),
+      projectId,
       attemptCount: 0,
       triggeredByIncidentId: incidentId,
     });
@@ -147,9 +316,14 @@ describe("AIInvestigationQueue", () => {
         attemptCount: 1,
       }),
     );
+    expect(AIService.getAutonomousDailyBudgetStatus).toHaveBeenCalledWith(
+      projectId,
+      { incidentId, alertId: undefined },
+    );
   });
 
   test("a won claim dispatches alert runs to the alert runner", async () => {
+    const projectId: ObjectID = ObjectID.generate();
     const alertId: ObjectID = ObjectID.generate();
     jest.spyOn(AIRunService, "attemptStatusTransition").mockResolvedValue(1);
     const execute: jest.SpyInstance = jest
@@ -158,7 +332,7 @@ describe("AIInvestigationQueue", () => {
 
     await AIInvestigationQueue.processRun({
       id: ObjectID.generate(),
-      projectId: ObjectID.generate(),
+      projectId,
       attemptCount: 0,
       triggeredByAlertId: alertId,
     });
@@ -166,6 +340,120 @@ describe("AIInvestigationQueue", () => {
     expect(execute).toHaveBeenCalledWith(
       expect.objectContaining({ alertId, attemptCount: 1 }),
     );
+    expect(AIService.getAutonomousDailyBudgetStatus).toHaveBeenCalledWith(
+      projectId,
+      { incidentId: undefined, alertId },
+    );
+  });
+
+  test("an incident at its cap does not consume the alert lane's slots", async () => {
+    jest.spyOn(ProjectService, "findOneById").mockResolvedValue({
+      id: ObjectID.generate(),
+      aiMaxConcurrentInvestigations: 1,
+      incidentAiMaxConcurrentInvestigations: 1,
+      alertAiMaxConcurrentInvestigations: 2,
+    } as unknown as Project);
+    const count: jest.SpyInstance = jest
+      .spyOn(AIRunService, "countBy")
+      .mockResolvedValue(new PositiveNumber(1));
+    const claim: jest.SpyInstance = jest
+      .spyOn(AIRunService, "attemptStatusTransition")
+      .mockResolvedValue(0);
+
+    await AIInvestigationQueue.processRun({
+      id: ObjectID.generate(),
+      projectId: ObjectID.generate(),
+      attemptCount: 0,
+      triggeredByIncidentId: ObjectID.generate(),
+    });
+    expect(claim).not.toHaveBeenCalled();
+
+    await AIInvestigationQueue.processRun({
+      id: ObjectID.generate(),
+      projectId: ObjectID.generate(),
+      attemptCount: 0,
+      triggeredByAlertId: ObjectID.generate(),
+    });
+    expect(claim).toHaveBeenCalledTimes(1);
+
+    const incidentQuery: Record<string, unknown> = (
+      count.mock.calls[0]![0] as { query: Record<string, unknown> }
+    ).query;
+    expect(findOperatorSql(incidentQuery["triggeredByIncidentId"])).toContain(
+      "IS NOT NULL",
+    );
+    expect(findOperatorSql(incidentQuery["triggeredByAlertId"])).toContain(
+      "IS NULL",
+    );
+
+    const alertQuery: Record<string, unknown> = (
+      count.mock.calls[1]![0] as { query: Record<string, unknown> }
+    ).query;
+    expect(findOperatorSql(alertQuery["triggeredByIncidentId"])).toContain(
+      "IS NULL",
+    );
+    expect(findOperatorSql(alertQuery["triggeredByAlertId"])).toContain(
+      "IS NOT NULL",
+    );
+  });
+
+  test("subjectless insight work retains the legacy concurrency cap", async () => {
+    jest.spyOn(ProjectService, "findOneById").mockResolvedValue({
+      id: ObjectID.generate(),
+      aiMaxConcurrentInvestigations: 1,
+      incidentAiMaxConcurrentInvestigations: 10,
+      alertAiMaxConcurrentInvestigations: 10,
+    } as unknown as Project);
+    const count: jest.SpyInstance = jest
+      .spyOn(AIRunService, "countBy")
+      .mockResolvedValue(new PositiveNumber(1));
+    const claim: jest.SpyInstance = jest.spyOn(
+      AIRunService,
+      "attemptStatusTransition",
+    );
+
+    await AIInvestigationQueue.processRun({
+      id: ObjectID.generate(),
+      projectId: ObjectID.generate(),
+      attemptCount: 0,
+      triggeredByAiInsightId: ObjectID.generate(),
+    });
+
+    expect(claim).not.toHaveBeenCalled();
+    const query: Record<string, unknown> = (
+      count.mock.calls[0]![0] as { query: Record<string, unknown> }
+    ).query;
+    expect(findOperatorSql(query["triggeredByIncidentId"])).toContain(
+      "IS NULL",
+    );
+    expect(findOperatorSql(query["triggeredByAlertId"])).toContain("IS NULL");
+  });
+
+  test("background counts are isolated to the remediation subject lane", async () => {
+    const count: jest.SpyInstance = jest
+      .spyOn(AIRunService, "countBy")
+      .mockResolvedValue(new PositiveNumber(0));
+    jest.spyOn(AIRunService, "attemptStatusTransition").mockResolvedValue(0);
+
+    await AIInvestigationQueue.processRun({
+      id: ObjectID.generate(),
+      projectId: ObjectID.generate(),
+      attemptCount: 0,
+      runType: AIRunType.RemediationPlan,
+      triggeredByIncidentId: ObjectID.generate(),
+      triggeredByAutoRemediationSuggestionId: ObjectID.generate(),
+    });
+
+    expect(count).toHaveBeenCalledTimes(3);
+    for (const call of count.mock.calls) {
+      const query: Record<string, unknown> = (
+        call[0] as { query: Record<string, unknown> }
+      ).query;
+      expect(findOperatorSql(query["triggeredByIncidentId"])).toContain(
+        "IS NOT NULL",
+      );
+      expect(findOperatorSql(query["triggeredByAlertId"])).toContain("IS NULL");
+    }
   });
 
   test("a transient failure on the first attempt requeues the run", async () => {
@@ -238,7 +526,7 @@ describe("AIInvestigationQueue", () => {
   test("a heartbeat-stale run requeues while attempts remain", async () => {
     jest.spyOn(AIRunService, "attemptStatusTransition").mockResolvedValue(1);
 
-    const outcome: "requeued" | "stale" =
+    const outcome: "requeued" | "stale" | "noop" =
       await AIInvestigationQueue.requeueOrMarkStale({
         id: ObjectID.generate(),
         attemptCount: 1,
@@ -252,7 +540,7 @@ describe("AIInvestigationQueue", () => {
       .spyOn(AIRunService, "attemptStatusTransition")
       .mockResolvedValue(1);
 
-    const outcome: "requeued" | "stale" =
+    const outcome: "requeued" | "stale" | "noop" =
       await AIInvestigationQueue.requeueOrMarkStale({
         id: ObjectID.generate(),
         attemptCount: MAX_INVESTIGATION_ATTEMPTS,

package/Tests/Server/Utils/AI/CodeFixAgentCompletion.test.ts

@@ -64,8 +64,9 @@ type UpdateSpy = jest.SpiedFunction<typeof AIRunService.updateOneBy>;
 function mockHappyDependencies(data?: {
   completionCalls?: number;
   outputTokens?: number;
+  run?: Partial<AIRun>;
 }): { executeSpy: ExecuteSpy; updateSpy: UpdateSpy } {
-  jest.spyOn(AIRunService, "findOneById").mockResolvedValue(fakeRun());
+  jest.spyOn(AIRunService, "findOneById").mockResolvedValue(fakeRun(data?.run));
   const updateSpy: UpdateSpy = jest
     .spyOn(AIRunService, "updateOneBy")
     .mockResolvedValue(1);
@@ -301,6 +302,63 @@ describe("CodeFixAgentCompletion.execute happy path", () => {
     );
   });
 
+  test("keeps a subject-linked fix completion in its autonomous budget lane", async () => {
+    const incidentId: ObjectID = ObjectID.generate();
+    const { executeSpy } = mockHappyDependencies({
+      run: { triggeredByIncidentId: incidentId },
+    });
+
+    await CodeFixAgentCompletion.execute({
+      aiAgentId: agentId,
+      aiRunId: runId,
+      messages,
+    });
+
+    expect(executeSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        incidentId,
+        alertId: undefined,
+      }),
+    );
+  });
+
+  test("keeps an alert-linked fix completion in the alert budget lane", async () => {
+    const alertId: ObjectID = ObjectID.generate();
+    const { executeSpy } = mockHappyDependencies({
+      run: { triggeredByAlertId: alertId },
+    });
+
+    await CodeFixAgentCompletion.execute({
+      aiAgentId: agentId,
+      aiRunId: runId,
+      messages,
+    });
+
+    expect(executeSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        incidentId: undefined,
+        alertId,
+      }),
+    );
+  });
+
+  test("keeps a general fix completion in the subjectless budget lane", async () => {
+    const { executeSpy } = mockHappyDependencies();
+
+    await CodeFixAgentCompletion.execute({
+      aiAgentId: agentId,
+      aiRunId: runId,
+      messages,
+    });
+
+    expect(executeSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        incidentId: undefined,
+        alertId: undefined,
+      }),
+    );
+  });
+
   test("reports post-call budget usage so the worker can wind down early", async () => {
     mockHappyDependencies({ completionCalls: 5, outputTokens: 1000 });