Submatter/apps/worker/src/handle-job.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";

import type { OrchestratorJob } from "@tasks/shared/orchestrator";

const {
  loadJobContext,
  resolveModelForRole,
  logRoleComplete,
  claimTask,
  completeTask,
  findBacklogItemByRepoPath,
  pickTaskRepoPath,
  ensureWorktree,
  removeWorktree,
  reuseWorktreeIfPresent,
  createMinimalSession,
  buildSessionPrompt,
  resolveRolePromptPath,
  getLatestCheckpointForRun,
  saveWorkerCheckpoint,
} = vi.hoisted(() => ({
  loadJobContext: vi.fn(),
  resolveModelForRole: vi.fn(),
  logRoleComplete: vi.fn(),
  claimTask: vi.fn(),
  completeTask: vi.fn(),
  findBacklogItemByRepoPath: vi.fn(),
  pickTaskRepoPath: vi.fn(),
  ensureWorktree: vi.fn(),
  removeWorktree: vi.fn(),
  reuseWorktreeIfPresent: vi.fn(),
  createMinimalSession: vi.fn(),
  buildSessionPrompt: vi.fn(),
  resolveRolePromptPath: vi.fn(),
  getLatestCheckpointForRun: vi.fn(),
  saveWorkerCheckpoint: vi.fn(),
}));

vi.mock("./db/convoy.js", () => ({
  loadJobContext,
  resolveModelForRole,
  logRoleComplete,
}));

vi.mock("./db/task-lifecycle.js", () => ({
  claimTask,
  completeTask,
  findBacklogItemByRepoPath,
  pickTaskRepoPath,
}));

vi.mock("./db/checkpoints.js", () => ({
  getLatestCheckpointForRun,
  saveWorkerCheckpoint,
}));

vi.mock("./git/worktree.js", () => ({
  ensureWorktree,
  removeWorktree,
  reuseWorktreeIfPresent,
}));

vi.mock("./session/minimal-session.js", () => ({
  createMinimalSession,
  buildSessionPrompt,
  resolveRolePromptPath,
}));

import { handleOrchestratorJob } from "./handle-job.js";

const db = {} as never;

const job: OrchestratorJob = {
  id: "11111111-1111-4111-8111-111111111111",
  type: "role_dispatch",
  workspaceId: "22222222-2222-4222-8222-222222222222",
  convoyId: "33333333-3333-4333-8333-333333333333",
  roleId: "role-implementer",
  briefId: "44444444-4444-4444-8444-444444444444",
  enqueuedAt: new Date().toISOString(),
};

const env = {
  DATABASE_URL: "postgres://local/test",
  WORKER_DRY_RUN: false,
  WORKER_BLPOP_TIMEOUT_S: 5,
  WORKER_MAX_TOOL_STEPS: 12,
  WORKER_CHECKPOINT_EVERY_N_STEPS: 1,
  WORKER_WORKTREE_BASE: "/tmp/echodo-worker-worktrees",
  WORKER_SESSION_TYPE: "minimal" as const,
  AGENT_SERVER_URL: "http://agent-server:8001",
  WORKER_OPENHANDS_WORKSPACE: "/workspace/repo",
  SESSION_API_KEY: "",
  REDIS_URL: "redis://localhost:6379",
};

const ctx = {
  convoyId: job.convoyId,
  convoyTitle: "Test convoy",
  convoySlug: "test-convoy",
  workspaceId: job.workspaceId,
  roleId: job.roleId,
  briefId: job.briefId,
  briefFiles: [".convoys/test-convoy/brief-1-implementer.md"],
  briefAcceptance: "Ship the feature",
};

const resolvedModel = {
  alias: "qwen-coder",
  providerType: "openai_compatible" as const,
  modelId: "qwen2.5-coder",
  baseUrl: "http://localhost:11434/v1",
  apiKey: "test-key",
  tier: "local" as const,
};

beforeEach(() => {
  vi.clearAllMocks();
  loadJobContext.mockResolvedValue(ctx);
  pickTaskRepoPath.mockReturnValue(ctx.briefFiles[0]);
  findBacklogItemByRepoPath.mockResolvedValue({
    id: "55555555-5555-4555-8555-555555555555",
    title: "Brief task",
    status: "ready",
    repoPath: ctx.briefFiles[0],
  });
  claimTask.mockResolvedValue({
    runId: "66666666-6666-4666-8666-666666666666",
    workflowPrompt: "Implement the brief.",
  });
  resolveModelForRole.mockResolvedValue({
    ok: true,
    model: resolvedModel,
    resolvedAlias: "qwen-coder",
    modelTier: "local",
  });
  ensureWorktree.mockResolvedValue({
    repoRoot: "/repo",
    worktreePath: "/tmp/echodo-worker-worktrees/job-1",
    branch: "worker/test-convoy-brief-44444444",
  });
  reuseWorktreeIfPresent.mockResolvedValue(null);
  getLatestCheckpointForRun.mockResolvedValue(null);
  saveWorkerCheckpoint.mockResolvedValue(undefined);
  removeWorktree.mockResolvedValue(undefined);
  buildSessionPrompt.mockReturnValue("Run the brief.");
  resolveRolePromptPath.mockReturnValue("/repo/.cursor/agents/role-implementer.md");
  logRoleComplete.mockResolvedValue(undefined);
  completeTask.mockResolvedValue(undefined);
});

describe("handleOrchestratorJob", () => {
  it("completes a brief with model alias and token totals", async () => {
    const run = vi.fn().mockResolvedValue({
      summary: "Implemented brief changes.",
      toolCalls: 3,
      tokensTotal: 1280,
    });
    createMinimalSession.mockReturnValue({ context: {}, run });

    await handleOrchestratorJob(db, null, env, job);

    expect(ensureWorktree).toHaveBeenCalledOnce();
    expect(run).toHaveBeenCalledOnce();
    expect(completeTask).toHaveBeenCalledWith(
      db,
      expect.objectContaining({
        runId: "66666666-6666-4666-8666-666666666666",
        outcome: "succeeded",
        tokensTotal: 1280,
        modelAlias: "qwen-coder",
        modelTier: "local",
        providerType: "openai_compatible",
        modelId: "qwen2.5-coder",
      }),
    );
    expect(removeWorktree).toHaveBeenCalledOnce();
  });

  it("marks the run failed and tears down the worktree when the session throws", async () => {
    createMinimalSession.mockReturnValue({
      context: {},
      run: vi.fn().mockRejectedValue(new Error("model unavailable")),
    });

    await handleOrchestratorJob(db, null, env, job);

    expect(completeTask).toHaveBeenCalledWith(
      db,
      expect.objectContaining({
        outcome: "failed",
        error: "model unavailable",
      }),
    );
    expect(removeWorktree).toHaveBeenCalledOnce();
  });

  it("marks the run failed when model resolution fails after claim", async () => {
    resolveModelForRole.mockResolvedValue({
      ok: false,
      reason: "No provider route for role-implementer",
    });

    await handleOrchestratorJob(db, null, env, job);

    expect(createMinimalSession).not.toHaveBeenCalled();
    expect(completeTask).toHaveBeenCalledWith(
      db,
      expect.objectContaining({
        outcome: "failed",
        error: "No provider route for role-implementer",
      }),
    );
    expect(removeWorktree).toHaveBeenCalledOnce();
  });

  it("skips worktree setup in dry-run mode", async () => {
    createMinimalSession.mockReturnValue({
      context: {},
      run: vi.fn().mockResolvedValue({
        summary: "[DRY-RUN] Would run role-implementer",
        toolCalls: 0,
      }),
    });

    await handleOrchestratorJob(db, null, { ...env, WORKER_DRY_RUN: true }, job);

    expect(ensureWorktree).not.toHaveBeenCalled();
    expect(removeWorktree).not.toHaveBeenCalled();
    expect(createMinimalSession).toHaveBeenCalledWith(
      expect.anything(),
      expect.objectContaining({ dryRun: true }),
    );
  });

  it("resumes from checkpoint without re-cloning when worktree is present", async () => {
    const resumeJob = { ...job, resumeFromCheckpoint: true };
    getLatestCheckpointForRun.mockResolvedValue({
      id: "cp-1",
      workspaceId: job.workspaceId,
      agentRunId: "66666666-6666-4666-8666-666666666666",
      convoyId: job.convoyId,
      briefId: job.briefId,
      roleId: job.roleId,
      checkpointSeq: 2,
      toolHistory: [
        {
          seq: 1,
          tool: "read_file",
          args: { path: "README.md" },
          ok: true,
          resultDigest: "ok:abc",
        },
      ],
      filesTouched: [],
      branch: "worker/test-convoy-brief-44444444",
      worktreePath: "/tmp/echodo-worker-worktrees/existing",
      occurredAt: new Date(),
    });
    reuseWorktreeIfPresent.mockResolvedValue({
      repoRoot: "/repo",
      worktreePath: "/tmp/echodo-worker-worktrees/existing",
      branch: "worker/test-convoy-brief-44444444",
    });
    const run = vi.fn().mockResolvedValue({
      summary: "Resumed and finished.",
      toolCalls: 3,
      tokensTotal: 400,
    });
    createMinimalSession.mockReturnValue({ context: {}, run });

    await handleOrchestratorJob(db, null, env, resumeJob);

    expect(reuseWorktreeIfPresent).toHaveBeenCalledWith(
      expect.any(String),
      "/tmp/echodo-worker-worktrees/existing",
      "worker/test-convoy-brief-44444444",
    );
    expect(ensureWorktree).not.toHaveBeenCalled();
    expect(createMinimalSession).toHaveBeenCalledWith(
      expect.objectContaining({ repoRoot: "/tmp/echodo-worker-worktrees/existing" }),
      expect.objectContaining({
        resumeToolHistory: expect.arrayContaining([
          expect.objectContaining({ tool: "read_file", seq: 1 }),
        ]),
      }),
    );
    expect(removeWorktree).toHaveBeenCalledOnce();
  });
});