evOo-local/packages/adapters/src/executor-lease-poll.test.ts

import type { AdapterContext, AgentRunRequest } from "@rakazo/adapter-kit";
import { describe, expect, it, vi } from "vitest";
import type * as AutoReviewModule from "./auto-review.js";
import type * as ComputerLifecycleModule from "./computer-lifecycle.js";
import { createRunExecutor } from "./executor.js";

vi.mock("./computer-lifecycle.js", async (importOriginal) => ({
  ...(await importOriginal<typeof ComputerLifecycleModule>()),
  acquireComputerExecutionLease: async () => null,
  provisionComputer: async () => ({ id: "computer-1", kind: "desktop" }),
}));

vi.mock("./auto-review.js", async (importOriginal) => ({
  ...(await importOriginal<typeof AutoReviewModule>()),
  resolveAutoReviewChecker: () => ({ provider: "scripted", model: "checker" }),
  isAutoReviewCheckerConfigured: () => true,
  runAutoReviewJudge: vi.fn(),
}));

/** Runtime that blocks without yielding until the executor abort signal fires. */
function silentBlockingRuntime() {
  let blocked = false;
  const runtimeRun = vi.fn(async function* (
    _request: AgentRunRequest,
    context?: Partial<AdapterContext>,
  ) {
    blocked = true;
    await new Promise<void>((_resolve, reject) => {
      const signal = context?.signal;
      if (!signal) return;
      if (signal.aborted) {
        reject(Object.assign(new Error("Aborted"), { name: "AbortError" }));
        return;
      }
      const onAbort = () => {
        signal.removeEventListener("abort", onAbort);
        reject(Object.assign(new Error("Aborted"), { name: "AbortError" }));
      };
      signal.addEventListener("abort", onAbort);
    });
    yield { type: "done" as const, text: "stopped" };
  });
  const runtimeAbort = vi.fn(async () => undefined);
  return {
    blocked: () => blocked,
    runtime: {
      describe: () => ({ capabilities: { scripted: false } }),
      run: runtimeRun,
      abort: runtimeAbort,
    },
    runtimeRun,
    runtimeAbort,
  };
}

function executorFixture() {
  const run = {
    id: "run-1",
    botId: "bot-1",
    threadId: "thread-1",
    taskId: "task-1",
    spaceId: "space-1",
    userId: "user-1",
    status: "queued",
    trigger: "user",
    leaseFence: 0,
    leaseOwner: null as string | null,
    startedAt: null as Date | null,
    checkpoint: null,
  };
  const prisma = {
    run: {
      findUnique: vi.fn(async () => run),
      findUniqueOrThrow: vi.fn(async () => run),
      updateMany: vi.fn(async ({ data }: { data: Record<string, unknown> }) => {
        Object.assign(run, data);
        return { count: 1 };
      }),
    },
    bot: {
      findUnique: vi.fn(async () => ({ computerId: "computer-1", computerSwitching: false })),
      findUniqueOrThrow: vi.fn(async () => ({
        id: run.botId,
        name: "Assistant",
        title: "Assistant",
        description: "Test assistant",
        computerId: "computer-1",
        computer: { id: "computer-1", scope: "dedicated" },
        instructions: null,
        modelProvider: null,
        modelId: null,
        memoryScope: null,
      })),
      findMany: vi.fn(async () => []),
    },
    attempt: {
      create: vi.fn(async () => ({ id: "attempt-1" })),
      update: vi.fn(),
      updateMany: vi.fn(async () => ({ count: 1 })),
    },
    thread: { findUniqueOrThrow: vi.fn(async () => ({ id: run.threadId, groupId: null })) },
    message: { findMany: vi.fn(async () => []), findUnique: vi.fn(async () => null) },
    task: {
      findUniqueOrThrow: vi.fn(async () => ({ id: run.taskId, prompt: "silent tool work" })),
    },
    connection: { findMany: vi.fn(async () => []) },
    spaceModelPreference: { findFirst: vi.fn(async () => null) },
    userModelCredential: { findFirst: vi.fn(async () => null) },
    deploymentSettings: {
      findUnique: vi.fn(async () => ({
        ownerUserId: null,
        ownerGithubTokenCipher: null,
        defaultModelProvider: "scripted",
        defaultModelId: "scripted",
      })),
    },
    space: {
      findUnique: vi.fn(async () => ({ isDefault: true, gitDefaultBranch: "main" })),
    },
    taughtSkill: { findMany: vi.fn(async () => []) },
    agentSecret: { findMany: vi.fn(async () => []) },
    user: { findUnique: vi.fn(async () => null) },
    userProfile: { findUnique: vi.fn(async () => null) },
    userFact: { findMany: vi.fn(async () => []) },
    agentSkill: { findMany: vi.fn(async () => []) },
    scratchpadItem: { findMany: vi.fn(async () => []) },
    actionApprovalRule: { findMany: vi.fn(async () => []) },
    actionAutoReviewPreference: { findUnique: vi.fn(async () => ({ enabled: false })) },
    externalEffect: {
      findMany: vi.fn(async () => []),
      findUnique: vi.fn(async () => null),
    },
  };
  const { runtime, runtimeRun, runtimeAbort, blocked } = silentBlockingRuntime();
  const executor = createRunExecutor({
    prisma,
    runtime,
    connector: {
      discoverTools: async () => [],
      resolveCall: async () => undefined,
      execute: async function* () {},
    },
    sandbox: { describe: () => ({ capabilities: { graphical: false } }) },
    memory: {
      describe: () => ({ capabilities: {} }),
      read: async () => ({ documents: [] }),
      search: async () => [],
      commit: vi.fn(async () => ({ revision: "rev-1" })),
      exportMarkdown: async function* () {},
    },
    memoryProviders: { resolve: async () => null },
    events: {
      append: vi.fn(async () => undefined),
      pauseRunForInput: vi.fn(async () => true),
      finalizeRun: vi.fn(async () => ({ continuationRunId: null })),
    },
    jobs: { enqueue: vi.fn(async () => undefined) },
    secrets: [],
  } as unknown as Parameters<typeof createRunExecutor>[0]);

  return { run, prisma, executor, runtimeRun, runtimeAbort, blocked };
}

describe("executor lease poll", () => {
  it("stops a stale worker during silent runtime work after lease reclaim", async () => {
    const { run, executor, runtimeRun, runtimeAbort, blocked } = executorFixture();
    const workerDone = executor.continueRun(run.id, "worker-a");

    await vi.waitFor(() => expect(runtimeRun).toHaveBeenCalled());
    expect(blocked()).toBe(true);

    Object.assign(run, { leaseOwner: "worker-b", leaseFence: 2, status: "running" });

    await vi.waitFor(() => expect(runtimeAbort).toHaveBeenCalledWith("run-1"), {
      timeout: 5_000,
    });
    await workerDone;

    expect(run.leaseOwner).toBe("worker-b");
  });
});