@oneuptime/common 14.0.10 → 14.0.11

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

@@ -12,29 +12,48 @@ import AlertAIContextBuilder, {
 } from "../../../../Server/Utils/AI/AlertAIContextBuilder";
 import AIMemory from "../../../../Server/Utils/AI/SRE/AIMemory";
 import { ObservabilityAssistantExtraTool } from "../../../../Server/Utils/AI/Chat/ObservabilityAssistant";
+import {
+  READ_TOOL_OUTPUT_TOOL_NAME,
+  TOOL_OUTPUT_PAGE_CHARS,
+} from "../../../../Server/Utils/AI/Chat/ToolOutputPager";
 import ClusterAccessContext from "../../../../Server/Utils/AI/ClusterAccess/ClusterAccessContext";
 import {
-  INVESTIGATION_MAX_WALL_CLOCK_MS,
   LIST_CLUSTER_ACCESS_TOOL_NAME,
   RUN_KUBECTL_TOOL_NAME,
 } from "../../../../Server/Utils/AI/ClusterAccess/KubectlInvestigationToolkit";
+import KubectlJobRunner, {
+  KubectlJobOutcome,
+  KubectlRunState,
+} from "../../../../Server/Utils/AI/ClusterAccess/KubectlJobRunner";
 import ResourceAccessContext from "../../../../Server/Utils/AI/ResourceAccess/ResourceAccessContext";
 import {
   LIST_INFRASTRUCTURE_ACCESS_TOOL_NAME,
   RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME,
 } from "../../../../Server/Utils/AI/ResourceAccess/ResourceAccessToolNames";
+import ResourceCommandJobRunner, {
+  ResourceCommandJobOutcome,
+  ResourceCommandRunState,
+} from "../../../../Server/Utils/AI/ResourceAccess/ResourceCommandJobRunner";
 import { ToolCallOutcome } from "../../../../Server/Utils/AI/Toolbox/Index";
 import KubernetesClusterAiAccessService from "../../../../Server/Services/KubernetesClusterAiAccessService";
+import ProjectService from "../../../../Server/Services/ProjectService";
 import ResourceAiAccessService from "../../../../Server/Services/ResourceAiAccessService";
-import RunnerJobService from "../../../../Server/Services/RunnerJobService";
 import logger from "../../../../Server/Utils/Logger";
 import Alert from "../../../../Models/DatabaseModels/Alert";
 import Incident from "../../../../Models/DatabaseModels/Incident";
+import Project from "../../../../Models/DatabaseModels/Project";
+import {
+  AI_AGENT_RUNAWAY_MAX_LLM_CALLS,
+  AI_AGENT_RUNAWAY_MAX_TOOL_CALLS,
+} from "../../../../Types/AI/AIAgentRunLimits";
+import { JSONObject } from "../../../../Types/JSON";
 import {
   KubernetesAiRemediationMode,
   KubernetesClusterAiAccessStatus,
 } from "../../../../Types/Kubernetes/KubernetesClusterAiAccess";
 import ObjectID from "../../../../Types/ObjectID";
+import RunnerJobOrigin from "../../../../Types/Runbook/RunnerJobOrigin";
+import KubectlWaitBudget from "../../../../Utils/AiRemediation/KubectlWaitBudget";
 import AiResourceType from "../../../../Types/ResourceAiAgent/AiResourceType";
 import {
   ResourceAiAccessStatus,
@@ -46,12 +65,14 @@ import { afterEach, beforeEach, describe, expect, test } from "@jest/globals";
  * Infrastructure access (Docker, Podman, Swarm, Proxmox, VMware, Ceph,
  * databases, hosts) reaches an investigation through the same four wires
  * cluster access does — the tools, the persona addendum, the context block
- * and the shared wall clock — BESIDE the cluster wiring, never instead of
- * it:
+ * and the shared deadline (none unless the project configured a time
+ * limit) — BESIDE the cluster wiring, never instead of it:
  *
  * - with no linked resource, every request is exactly what it was (the
  *   cluster addendum alone, the kubectl tools alone);
  * - the cluster rules come first, the infrastructure rules after;
+ * - both toolkits page long output into ONE pager, whose read_tool_output
+ *   comes after every command tool — and only when there is one;
  * - a failed resource lookup never fails the run and never takes the
  *   cluster access away (and the other way round).
  */
@@ -145,6 +166,88 @@ function toolNames(request: InvestigationRequest): Array<string> {
   );
 }
 
+function findTool(
+  request: InvestigationRequest,
+  name: string,
+): ObservabilityAssistantExtraTool {
+  const tool: ObservabilityAssistantExtraTool | undefined = (
+    request.extraTools || []
+  ).find((candidate: ObservabilityAssistantExtraTool) => {
+    return candidate.definition.name === name;
+  });
+
+  if (!tool) {
+    throw new Error(`Tool ${name} not offered.`);
+  }
+
+  return tool;
+}
+
+// The project row the runners read the lane's time limit from.
+function projectWithTimeLimits(limits: {
+  incidentMinutes?: number | undefined;
+  alertMinutes?: number | undefined;
+}): Project {
+  const project: Project = new Project(projectId);
+
+  if (limits.incidentMinutes !== undefined) {
+    project.incidentAiInvestigationTimeLimitInMinutes = limits.incidentMinutes;
+  }
+
+  if (limits.alertMinutes !== undefined) {
+    project.alertAiInvestigationTimeLimitInMinutes = limits.alertMinutes;
+  }
+
+  return project;
+}
+
+// What ResourceCommandJobRunner.run hands back for a command that ran.
+function resourceOutcome(output: string): ResourceCommandJobOutcome {
+  return {
+    jobId: ObjectID.generate().toString(),
+    succeeded: true,
+    exitCode: 0,
+    output,
+    redactionCount: 0,
+    isTruncated: false,
+    displayCommand: "docker ps -a",
+    executed: true,
+    runState: ResourceCommandRunState.Ran,
+    claimTimedOut: false,
+    isAccessFailure: false,
+  };
+}
+
+// What KubectlJobRunner.run hands back for a command kubectl ran.
+function kubectlOutcome(output: string): KubectlJobOutcome {
+  return {
+    jobId: ObjectID.generate().toString(),
+    succeeded: true,
+    exitCode: 0,
+    output,
+    redactionCount: 0,
+    isTruncated: false,
+    displayCommand: "kubectl get pods -A",
+    executed: true,
+    runState: KubectlRunState.Ran,
+    claimTimedOut: false,
+    isAccessFailure: false,
+  };
+}
+
+// An output more than one page long that ends in `lastLine`.
+function longOutput(row: string, lastLine: string): string {
+  return `${row.repeat(
+    Math.ceil((TOOL_OUTPUT_PAGE_CHARS + 1000) / row.length),
+  )}${lastLine}`;
+}
+
+const RUN_COMMAND_ARGS: JSONObject = {
+  resourceId: DOCKER_ID,
+  command: "docker ps -a",
+  rationale: "see the containers",
+};
+
 async function investigate(kind: SubjectKind): Promise<void> {
   if (kind === "incident") {
     const incident: Incident = new Incident(incidentId);
@@ -192,11 +295,16 @@ describe("Investigation runners wire infrastructure access beside cluster access
   let clusterStatuses: jest.SpyInstance;
   let resourceStatuses: jest.SpyInstance;
   let failOrRequeue: jest.SpyInstance;
+  let findProject: jest.SpyInstance;
 
   beforeEach(() => {
     jest.spyOn(logger, "error").mockImplementation((): void => {
       return undefined;
     });
+    // A project that configured no investigation time limit (the default).
+    findProject = jest
+      .spyOn(ProjectService, "findOneById")
+      .mockResolvedValue(projectWithTimeLimits({}));
     executeRun = jest
       .spyOn(AIInvestigationEngine, "executeRun")
       .mockResolvedValue(undefined);
@@ -216,7 +324,7 @@ describe("Investigation runners wire infrastructure access beside cluster access
   });
 
   test.each(SUBJECT_KINDS)(
-    "a %s investigation gets the infrastructure tools, addendum, context block and the shared wall clock",
+    "a %s investigation gets the infrastructure tools, the output reader, addendum, context block and no time limit",
     async (kind: SubjectKind) => {
       await investigate(kind);
 
@@ -231,6 +339,7 @@ describe("Investigation runners wire infrastructure access beside cluster access
       expect(toolNames(request)).toEqual([
         LIST_INFRASTRUCTURE_ACCESS_TOOL_NAME,
         RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME,
+        READ_TOOL_OUTPUT_TOOL_NAME,
       ]);
       expect(request.additionalInstructions).toBe(
         ResourceAccessContext.buildPersonaAddendum([resourceStatus()]),
@@ -240,7 +349,9 @@ describe("Investigation runners wire infrastructure access beside cluster access
         `- Docker host "web-1" (resourceId: ${DOCKER_ID}): READ access via run_infrastructure_command`,
       );
       expect(request.contextSummary).not.toContain("# Cluster access");
-      expect(request.maxWallClockMs).toBe(INVESTIGATION_MAX_WALL_CLOCK_MS);
+      expect(request.maxWallClockMs).toBeUndefined();
+      expect(request.maxLlmCalls).toBe(AI_AGENT_RUNAWAY_MAX_LLM_CALLS);
+      expect(request.maxToolCalls).toBe(AI_AGENT_RUNAWAY_MAX_TOOL_CALLS);
       expect(failOrRequeue).not.toHaveBeenCalled();
     },
   );
@@ -254,11 +365,13 @@ describe("Investigation runners wire infrastructure access beside cluster access
 
       const request: InvestigationRequest = sentRequest(executeRun);
 
+      // One reader for both toolkits' outputs, after every command tool.
       expect(toolNames(request)).toEqual([
         LIST_CLUSTER_ACCESS_TOOL_NAME,
         RUN_KUBECTL_TOOL_NAME,
         LIST_INFRASTRUCTURE_ACCESS_TOOL_NAME,
         RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME,
+        READ_TOOL_OUTPUT_TOOL_NAME,
       ]);
       expect(request.additionalInstructions).toBe(
         `${ClusterAccessContext.buildPersonaAddendum([
@@ -289,6 +402,7 @@ describe("Investigation runners wire infrastructure access beside cluster access
       expect(toolNames(request)).toEqual([
         LIST_CLUSTER_ACCESS_TOOL_NAME,
         RUN_KUBECTL_TOOL_NAME,
+        READ_TOOL_OUTPUT_TOOL_NAME,
       ]);
       // Byte for byte the cluster addendum alone — nothing appended.
       expect(request.additionalInstructions).toBe(
@@ -347,6 +461,7 @@ describe("Investigation runners wire infrastructure access beside cluster access
       expect(toolNames(request)).toEqual([
         LIST_CLUSTER_ACCESS_TOOL_NAME,
         RUN_KUBECTL_TOOL_NAME,
+        READ_TOOL_OUTPUT_TOOL_NAME,
       ]);
       expect(request.additionalInstructions).toBe(
         ClusterAccessContext.buildPersonaAddendum([clusterStatus()]),
@@ -367,6 +482,7 @@ describe("Investigation runners wire infrastructure access beside cluster access
       expect(toolNames(request)).toEqual([
         LIST_INFRASTRUCTURE_ACCESS_TOOL_NAME,
         RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME,
+        READ_TOOL_OUTPUT_TOOL_NAME,
       ]);
       expect(failOrRequeue).not.toHaveBeenCalled();
     },
@@ -418,36 +534,161 @@ describe("Investigation runners wire infrastructure access beside cluster access
   });
 
   test.each(SUBJECT_KINDS)(
-    "a %s investigation's infrastructure tools are bound to this run and its deadline",
+    "a %s investigation's infrastructure tools are bound to this run, with no deadline by default",
+    async (kind: SubjectKind) => {
+      const NOW_MS: number = 1_700_000_000_000;
+      const now: jest.SpyInstance = jest
+        .spyOn(Date, "now")
+        .mockReturnValue(NOW_MS);
+      const plan: jest.SpyInstance = jest.spyOn(KubectlWaitBudget, "plan");
+      const commandRun: jest.SpyInstance = jest
+        .spyOn(ResourceCommandJobRunner, "run")
+        .mockResolvedValue(resourceOutcome("[stdout]\nabc   nginx"));
+
+      await investigate(kind);
+
+      const runCommand: ObservabilityAssistantExtraTool = findTool(
+        sentRequest(executeRun),
+        RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME,
+      );
+
+      // A day into the run: an unlimited run is never refused for time.
+      now.mockReturnValue(NOW_MS + 24 * 60 * 60 * 1000);
+      const outcome: ToolCallOutcome =
+        await runCommand.execute(RUN_COMMAND_ARGS);
+
+      expect(outcome.success).toBe(true);
+      expect(plan).toHaveBeenCalledWith(
+        expect.objectContaining({ deadlineAtMs: undefined }),
+      );
+      expect(commandRun).toHaveBeenCalledWith(
+        expect.objectContaining({
+          projectId,
+          aiRunId,
+          origin: RunnerJobOrigin.AiInvestigation,
+          resourceType: AiResourceType.DockerHost,
+          command: "docker ps -a",
+        }),
+      );
+    },
+  );
+
+  /*
+   * A project that configured a time limit for the lane: the engine gets
+   * it as maxWallClockMs, and the infrastructure toolkit the deadline it
+   * ends at — the same number, measured from when the run was wired.
+   */
+  test.each(SUBJECT_KINDS)(
+    "a %s investigation with a configured time limit hands the engine and the infrastructure tools the same deadline",
     async (kind: SubjectKind) => {
       const NOW_MS: number = 1_700_000_000_000;
+      const LIMIT_MS: number = 15 * 60 * 1000;
+      findProject.mockResolvedValue(
+        projectWithTimeLimits({ incidentMinutes: 15, alertMinutes: 15 }),
+      );
       const now: jest.SpyInstance = jest
         .spyOn(Date, "now")
         .mockReturnValue(NOW_MS);
-      const enqueue: jest.SpyInstance = jest.spyOn(
-        RunnerJobService,
-        "enqueueAiResourceCommand",
+      const plan: jest.SpyInstance = jest.spyOn(KubectlWaitBudget, "plan");
+      const commandRun: jest.SpyInstance = jest
+        .spyOn(ResourceCommandJobRunner, "run")
+        .mockResolvedValue(resourceOutcome("[stdout]\nabc   nginx"));
+
+      await investigate(kind);
+
+      const request: InvestigationRequest = sentRequest(executeRun);
+      expect(request.maxWallClockMs).toBe(LIMIT_MS);
+
+      const runCommand: ObservabilityAssistantExtraTool = findTool(
+        request,
+        RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME,
+      );
+
+      // Within the limit: planned against the run's deadline.
+      now.mockReturnValue(NOW_MS + 60_000);
+      expect((await runCommand.execute(RUN_COMMAND_ARGS)).success).toBe(true);
+      expect(plan).toHaveBeenCalledWith(
+        expect.objectContaining({ deadlineAtMs: NOW_MS + LIMIT_MS }),
+      );
+
+      // Past it: the toolkit refuses on its own, enqueuing nothing.
+      now.mockReturnValue(NOW_MS + LIMIT_MS + 1);
+      const late: ToolCallOutcome = await runCommand.execute(RUN_COMMAND_ARGS);
+
+      expect(late.success).toBe(false);
+      expect(late.textForLlm).toContain(
+        "Not enough time is left before this investigation's configured time limit to run another infrastructure command",
       );
+      expect(commandRun).toHaveBeenCalledTimes(1);
+    },
+  );
+
+  /*
+   * The kubectl and infrastructure toolkits of one run share one pager:
+   * their long outputs are numbered in one sequence, and the single
+   * read_tool_output the run is offered reads both. Two pagers would both
+   * hand out "out-1", and the reader could read neither.
+   */
+  test.each(SUBJECT_KINDS)(
+    "a %s investigation's command toolkits share one pager behind read_tool_output",
+    async (kind: SubjectKind) => {
+      clusterStatuses.mockResolvedValue([clusterStatus()]);
+      jest
+        .spyOn(KubectlJobRunner, "run")
+        .mockResolvedValue(
+          kubectlOutcome(
+            longOutput("web-1  1/1     Running   0\n", "KUBECTL-LAST-LINE"),
+          ),
+        );
+      jest
+        .spyOn(ResourceCommandJobRunner, "run")
+        .mockResolvedValue(
+          resourceOutcome(
+            longOutput(
+              "abc123   nginx:1.27   Up 3 hours\n",
+              "DOCKER-LAST-LINE",
+            ),
+          ),
+        );
 
       await investigate(kind);
 
-      const runCommand: ObservabilityAssistantExtraTool = (
-        sentRequest(executeRun).extraTools || []
-      ).find((tool: ObservabilityAssistantExtraTool) => {
-        return tool.definition.name === RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME;
-      })!;
-
-      // The run's wall clock has run out: the toolkit refuses on its own.
-      now.mockReturnValue(NOW_MS + INVESTIGATION_MAX_WALL_CLOCK_MS + 1);
-      const outcome: ToolCallOutcome = await runCommand.execute({
-        resourceId: DOCKER_ID,
-        command: "docker ps -a",
-        rationale: "see the containers",
+      const request: InvestigationRequest = sentRequest(executeRun);
+      const readTool: ObservabilityAssistantExtraTool =
+        request.extraTools![request.extraTools!.length - 1]!;
+      expect(readTool.definition.name).toBe(READ_TOOL_OUTPUT_TOOL_NAME);
+
+      const kubectl: ToolCallOutcome = await findTool(
+        request,
+        RUN_KUBECTL_TOOL_NAME,
+      ).execute({
+        clusterId: CLUSTER_ID,
+        command: "kubectl get pods -A",
+        rationale: "every pod",
+      });
+      const docker: ToolCallOutcome = await findTool(
+        request,
+        RUN_INFRASTRUCTURE_COMMAND_TOOL_NAME,
+      ).execute(RUN_COMMAND_ARGS);
+
+      expect(kubectl.textForLlm).toContain('outputId="out-1"');
+      expect(kubectl.textForLlm).not.toContain("KUBECTL-LAST-LINE");
+      expect(docker.textForLlm).toContain('outputId="out-2"');
+      expect(docker.textForLlm).not.toContain("DOCKER-LAST-LINE");
+
+      const kubectlTail: ToolCallOutcome = await readTool.execute({
+        outputId: "out-1",
+        offset: -40,
+      });
+      const dockerTail: ToolCallOutcome = await readTool.execute({
+        outputId: "out-2",
+        offset: -40,
       });
 
-      expect(outcome.success).toBe(false);
-      expect(outcome.textForLlm).toContain("Not enough time is left");
-      expect(enqueue).not.toHaveBeenCalled();
+      expect(kubectlTail.success).toBe(true);
+      expect(kubectlTail.textForLlm).toContain("KUBECTL-LAST-LINE");
+      expect(dockerTail.success).toBe(true);
+      expect(dockerTail.textForLlm).toContain("DOCKER-LAST-LINE");
     },
   );
 });