ai_ops/tests/orchestration-engine.test.ts

import test from "node:test";
import assert from "node:assert/strict";
import { mkdtemp, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { resolve } from "node:path";
import { SchemaDrivenExecutionEngine } from "../src/agents/orchestration.js";
import type { ActorExecutionResult } from "../src/agents/pipeline.js";

function createManifest(): unknown {
  return {
    schemaVersion: "1",
    topologies: ["hierarchical", "retry-unrolled", "sequential"],
    personas: [
      {
        id: "product",
        displayName: "Product",
        systemPromptTemplate: "Product planning for {{repo}}",
        toolClearance: {
          allowlist: ["read_file"],
          banlist: ["delete_file"],
        },
      },
      {
        id: "task",
        displayName: "Task",
        systemPromptTemplate: "Task planning for {{repo}}",
        toolClearance: {
          allowlist: ["read_file", "write_file"],
          banlist: ["git_reset"],
        },
      },
      {
        id: "coder",
        displayName: "Coder",
        systemPromptTemplate: "Coder implements {{ticket}}",
        toolClearance: {
          allowlist: ["read_file", "write_file"],
          banlist: ["rm"],
        },
      },
      {
        id: "qa",
        displayName: "QA",
        systemPromptTemplate: "QA validates {{ticket}}",
        toolClearance: {
          allowlist: ["read_file"],
          banlist: ["write_file"],
        },
      },
    ],
    relationships: [
      {
        parentPersonaId: "product",
        childPersonaId: "task",
        constraints: {
          maxChildren: 2,
          maxDepth: 2,
        },
      },
      {
        parentPersonaId: "task",
        childPersonaId: "coder",
        constraints: {
          maxChildren: 3,
          maxDepth: 3,
        },
      },
    ],
    topologyConstraints: {
      maxDepth: 6,
      maxRetries: 2,
    },
    pipeline: {
      entryNodeId: "project-gate",
      nodes: [
        {
          id: "project-gate",
          actorId: "project_gate",
          personaId: "product",
        },
        {
          id: "task-plan",
          actorId: "task_plan",
          personaId: "task",
        },
        {
          id: "coder-1",
          actorId: "coder",
          personaId: "coder",
          constraints: {
            maxRetries: 1,
          },
        },
        {
          id: "qa-1",
          actorId: "qa",
          personaId: "qa",
        },
      ],
      edges: [
        {
          from: "project-gate",
          to: "task-plan",
          on: "success",
          when: [
            {
              type: "state_flag",
              key: "needs_bootstrap",
              equals: true,
            },
          ],
        },
        {
          from: "project-gate",
          to: "coder-1",
          on: "success",
          when: [
            {
              type: "state_flag",
              key: "needs_bootstrap",
              equals: false,
            },
          ],
        },
        {
          from: "task-plan",
          to: "coder-1",
          on: "success",
        },
        {
          from: "coder-1",
          to: "qa-1",
          on: "success",
          when: [
            {
              type: "history_has_event",
              event: "validation_fail",
            },
          ],
        },
      ],
    },
  };
}

test("runs DAG pipeline with state-dependent routing and retry behavior", async () => {
  const workspaceRoot = await mkdtemp(resolve(tmpdir(), "ai-ops-workspace-"));
  const stateRoot = await mkdtemp(resolve(tmpdir(), "ai-ops-session-state-"));

  await writeFile(resolve(workspaceRoot, "PRD.md"), "# PRD\n", "utf8");

  let coderAttempts = 0;

  const engine = new SchemaDrivenExecutionEngine({
    manifest: createManifest(),
    settings: {
      workspaceRoot,
      stateRoot,
      runtimeContext: {
        repo: "ai_ops",
        ticket: "AIOPS-123",
      },
      maxChildren: 4,
      maxDepth: 8,
      maxRetries: 3,
    },
    actorExecutors: {
      project_gate: async () => ({
        status: "success",
        payload: {
          phase: "gate",
        },
        stateFlags: {
          needs_bootstrap: true,
        },
      }),
      task_plan: async (input) => {
        assert.match(input.prompt, /ai_ops/);
        return {
          status: "success",
          payload: {
            plan: "roadmap",
          },
          stateFlags: {
            roadmap_ready: true,
          },
        };
      },
      coder: async (input): Promise<ActorExecutionResult> => {
        assert.match(input.prompt, /AIOPS-123/);
        assert.deepEqual(input.toolClearance.allowlist, ["read_file", "write_file"]);
        coderAttempts += 1;
        if (coderAttempts === 1) {
          return {
            status: "validation_fail",
            payload: {
              issue: "missing test",
            },
          };
        }

        return {
          status: "success",
          payload: {
            code: "done",
          },
        };
      },
      qa: async () => ({
        status: "success",
        payload: {
          qa: "ok",
        },
      }),
    },
    behaviorHandlers: {
      coder: {
        onValidationFail: () => ({
          lastValidationFailure: "coder-1",
        }),
      },
    },
  });

  const result = await engine.runSession({
    sessionId: "session-orchestration-1",
    initialPayload: {
      task: "Implement pipeline",
    },
  });

  assert.deepEqual(
    result.records.map((record) => `${record.nodeId}:${record.status}:${String(record.attempt)}`),
    [
      "project-gate:success:1",
      "task-plan:success:1",
      "coder-1:validation_fail:1",
      "coder-1:success:2",
      "qa-1:success:1",
    ],
  );

  assert.equal(result.finalState.flags.needs_bootstrap, true);
  assert.equal(result.finalState.flags.roadmap_ready, true);
  assert.equal(result.finalState.metadata.lastValidationFailure, "coder-1");

  assert.deepEqual(engine.planChildPersonas({ parentPersonaId: "task", depth: 1 }), ["coder"]);
});