mcpctl/src/mcpd/tests/inference-task-service.test.ts

/**
 * v5 InferenceTaskService — state machine, signal channels, and GC sweep.
 * The repo is mocked with an in-memory Map so we exercise the service
 * logic deterministically without touching Postgres. Schema-level tests
 * live in src/db/tests/inference-task-schema.test.ts.
 */
import { describe, it, expect, vi } from 'vitest';
import type { InferenceTask, InferenceTaskStatus } from '@prisma/client';
import type { IInferenceTaskRepository } from '../src/repositories/inference-task.repository.js';
import { InferenceTaskService } from '../src/services/inference-task.service.js';

function makeRow(overrides: Partial<InferenceTask> = {}): InferenceTask {
  return {
    id: overrides.id ?? `task-${Math.random().toString(36).slice(2, 8)}`,
    status: 'pending',
    poolName: 'pool-a',
    llmName: 'pool-a',
    model: 'qwen3-thinking',
    tier: null,
    claimedBy: null,
    requestBody: { messages: [{ role: 'user', content: 'hi' }] } as unknown as InferenceTask['requestBody'],
    responseBody: null,
    errorMessage: null,
    streaming: false,
    createdAt: new Date(),
    claimedAt: null,
    streamStartedAt: null,
    completedAt: null,
    ownerId: 'owner-1',
    agentId: null,
    ...overrides,
  };
}

function mockRepo(initial: InferenceTask[] = []): IInferenceTaskRepository {
  const rows = new Map<string, InferenceTask>(initial.map((r) => [r.id, { ...r }]));
  // Tiny CAS helper that mirrors the real `updateMany({where:{status:in}})`
  // semantics — only flips the row if the current status matches.
  const cas = (id: string, allowed: InferenceTaskStatus[], patch: Partial<InferenceTask>): InferenceTask | null => {
    const row = rows.get(id);
    if (row === undefined) return null;
    if (!allowed.includes(row.status)) return null;
    const next = { ...row, ...patch };
    rows.set(id, next);
    return next;
  };
  return {
    create: vi.fn(async (data) => {
      const row = makeRow({
        id: `task-${rows.size + 1}`,
        poolName: data.poolName,
        llmName: data.llmName,
        model: data.model,
        tier: data.tier ?? null,
        requestBody: data.requestBody as InferenceTask['requestBody'],
        streaming: data.streaming,
        ownerId: data.ownerId,
        agentId: data.agentId ?? null,
      });
      rows.set(row.id, row);
      return row;
    }),
    findById: vi.fn(async (id) => rows.get(id) ?? null),
    findPendingForPools: vi.fn(async (poolNames, limit) => {
      const out = [...rows.values()]
        .filter((r) => r.status === 'pending' && poolNames.includes(r.poolName))
        .sort((a, b) => a.createdAt.getTime() - b.createdAt.getTime());
      return limit !== undefined ? out.slice(0, limit) : out;
    }),
    findHeldBy: vi.fn(async (claimedBy) =>
      [...rows.values()].filter((r) => r.claimedBy === claimedBy && (r.status === 'claimed' || r.status === 'running')),
    ),
    list: vi.fn(async (filter) => {
      let out = [...rows.values()];
      if (filter.ownerId !== undefined) out = out.filter((r) => r.ownerId === filter.ownerId);
      if (filter.poolName !== undefined) out = out.filter((r) => r.poolName === filter.poolName);
      if (filter.agentId !== undefined) out = out.filter((r) => r.agentId === filter.agentId);
      if (filter.status !== undefined) {
        const statuses = Array.isArray(filter.status) ? filter.status : [filter.status];
        out = out.filter((r) => statuses.includes(r.status));
      }
      out.sort((a, b) => b.createdAt.getTime() - a.createdAt.getTime());
      return filter.limit !== undefined ? out.slice(0, filter.limit) : out;
    }),
    tryClaim: vi.fn(async (id, claimedBy, claimedAt) =>
      cas(id, ['pending'], { status: 'claimed', claimedBy, claimedAt }),
    ),
    markRunning: vi.fn(async (id, at) =>
      cas(id, ['claimed', 'running'], { status: 'running', streamStartedAt: at }),
    ),
    markCompleted: vi.fn(async (id, body, at) =>
      cas(id, ['pending', 'claimed', 'running'], {
        status: 'completed',
        responseBody: (body ?? null) as InferenceTask['responseBody'],
        completedAt: at,
      }),
    ),
    markError: vi.fn(async (id, errorMessage, at) =>
      cas(id, ['pending', 'claimed', 'running'], { status: 'error', errorMessage, completedAt: at }),
    ),
    markCancelled: vi.fn(async (id, at) =>
      cas(id, ['pending', 'claimed', 'running'], { status: 'cancelled', completedAt: at }),
    ),
    revertToPending: vi.fn(async (id) =>
      cas(id, ['claimed', 'running'], { status: 'pending', claimedBy: null, claimedAt: null, streamStartedAt: null }),
    ),
    findStalePending: vi.fn(async (cutoff) =>
      [...rows.values()].filter((r) => r.status === 'pending' && r.createdAt.getTime() < cutoff.getTime()),
    ),
    findExpiredTerminal: vi.fn(async (cutoff) =>
      [...rows.values()].filter((r) =>
        (r.status === 'completed' || r.status === 'error' || r.status === 'cancelled')
        && r.completedAt !== null
        && r.completedAt.getTime() < cutoff.getTime(),
      ),
    ),
    deleteMany: vi.fn(async (ids) => {
      let n = 0;
      for (const id of ids) if (rows.delete(id)) n += 1;
      return n;
    }),
  };
}

describe('InferenceTaskService — state machine', () => {
  it('enqueue creates a pending row with the given pool/llm/model', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({
      poolName: 'pool-a',
      llmName: 'a-1',
      model: 'qwen3',
      requestBody: { messages: [] },
      streaming: false,
      ownerId: 'owner-1',
    });
    expect(t.status).toBe('pending');
    expect(t.poolName).toBe('pool-a');
    expect(t.claimedBy).toBeNull();
  });

  it('tryClaim races: only one of two concurrent claimers gets the row', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({
      poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o',
    });
    // Both workers issue tryClaim at the "same time". The repo's CAS
    // serializes them — first claim wins, second sees a non-pending
    // row and returns null.
    const [a, b] = await Promise.all([svc.tryClaim(t.id, 'sess-A'), svc.tryClaim(t.id, 'sess-B')]);
    const winners = [a, b].filter((r) => r !== null);
    const losers = [a, b].filter((r) => r === null);
    expect(winners).toHaveLength(1);
    expect(losers).toHaveLength(1);
    expect(winners[0]!.status).toBe('claimed');
    expect(['sess-A', 'sess-B']).toContain(winners[0]!.claimedBy);
  });

  it('complete after claim transitions claimed → completed and stores responseBody', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    await svc.tryClaim(t.id, 'sess-A');
    const done = await svc.complete(t.id, { choices: [{ message: { content: 'hi' } }] });
    expect(done?.status).toBe('completed');
    expect(done?.responseBody).toEqual({ choices: [{ message: { content: 'hi' } }] });
    expect(done?.completedAt).not.toBeNull();
  });

  it('refuses double-complete (idempotent terminal)', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const first = await svc.complete(t.id, { ok: 1 });
    expect(first?.status).toBe('completed');
    // Second worker tries to complete the same task — CAS rejects because
    // the row is no longer in a non-terminal state.
    const second = await svc.complete(t.id, { ok: 2 });
    expect(second).toBeNull();
    // First completion's body is preserved.
    const reread = await svc.findById(t.id);
    expect(reread?.responseBody).toEqual({ ok: 1 });
  });

  it('revertHeldBy reverts every claimed/running row owned by a session and leaves terminals alone', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t1 = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const t2 = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const t3 = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    await svc.tryClaim(t1.id, 'sess-A');
    await svc.tryClaim(t2.id, 'sess-A');
    await svc.markRunning(t2.id);
    await svc.tryClaim(t3.id, 'sess-A');
    await svc.complete(t3.id, { ok: 1 }); // t3 finished before disconnect

    const reverted = await svc.revertHeldBy('sess-A');
    expect(reverted.map((r) => r.id).sort()).toEqual([t1.id, t2.id].sort());
    expect((await svc.findById(t1.id))?.status).toBe('pending');
    expect((await svc.findById(t2.id))?.status).toBe('pending');
    // t3 stayed completed — terminal rows are not reverted on disconnect.
    expect((await svc.findById(t3.id))?.status).toBe('completed');
  });

  it('cancel from pending records cancelled and emits terminal event', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const cancelled = await svc.cancel(t.id);
    expect(cancelled?.status).toBe('cancelled');
    // A subsequent complete must fail (cancelled is terminal).
    const result = await svc.complete(t.id, { ok: 1 });
    expect(result).toBeNull();
  });
});

describe('InferenceTaskService — waitFor signals', () => {
  it('resolves immediately when the row is already terminal at subscribe time', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    await svc.complete(t.id, { ok: 1 });
    const waiter = svc.waitFor(t.id, 1_000);
    const final = await waiter.done;
    expect(final.status).toBe('completed');
    expect(final.responseBody).toEqual({ ok: 1 });
  });

  it('wakes a waiter on complete event without polling the DB', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const waiter = svc.waitFor(t.id, 5_000);
    // Fire the complete after a microtask so the waiter is definitely
    // already subscribed to the terminal event.
    setTimeout(() => { void svc.complete(t.id, { ok: 1 }); }, 10);
    const final = await waiter.done;
    expect(final.status).toBe('completed');
  });

  it('wakes the chunks generator on pushChunk and ends on terminal', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: true, ownerId: 'o' });
    const waiter = svc.waitFor(t.id, 5_000);

    setTimeout(() => {
      svc.pushChunk(t.id, { data: 'hello ' });
      svc.pushChunk(t.id, { data: 'world' });
      void svc.complete(t.id, { ok: 1 });
    }, 10);

    const seen: string[] = [];
    for await (const c of waiter.chunks) {
      seen.push(c.data);
    }
    expect(seen).toEqual(['hello ', 'world']);
    const final = await waiter.done;
    expect(final.status).toBe('completed');
  });

  it('throws on cancellation with a clear error message', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const waiter = svc.waitFor(t.id, 5_000);
    setTimeout(() => { void svc.cancel(t.id); }, 10);
    await expect(waiter.done).rejects.toThrow(/cancelled/i);
  });

  it('throws on error and surfaces the worker\'s errorMessage', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const waiter = svc.waitFor(t.id, 5_000);
    setTimeout(() => { void svc.fail(t.id, 'upstream 500'); }, 10);
    await expect(waiter.done).rejects.toThrow(/upstream 500/);
  });

  it('times out when no terminal event arrives within the deadline', async () => {
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo);
    const t = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const waiter = svc.waitFor(t.id, 30);
    await expect(waiter.done).rejects.toThrow(/timed out/);
  });
});

describe('InferenceTaskService — gcSweep', () => {
  it('flips stale pending rows to error AND deletes expired terminal rows', async () => {
    const now = new Date('2026-04-28T00:00:00Z');
    const fixedClock = (): Date => now;
    const repo = mockRepo();
    const svc = new InferenceTaskService(repo, fixedClock);

    // 90 min old pending — past the 1h pendingTimeout cutoff → error.
    const stale = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    // Backdate via direct fixture mutation — easier than wiring a clock through enqueue.
    const staleRow = (await svc.findById(stale.id))!;
    (staleRow as { createdAt: Date }).createdAt = new Date(now.getTime() - 90 * 60 * 1000);

    // 30 min old pending — within window, should not be touched.
    const fresh = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    const freshRow = (await svc.findById(fresh.id))!;
    (freshRow as { createdAt: Date }).createdAt = new Date(now.getTime() - 30 * 60 * 1000);

    // 8 day-old completed — past the 7d retention → delete.
    const old = await svc.enqueue({ poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, streaming: false, ownerId: 'o' });
    await svc.complete(old.id, { ok: 1 });
    const oldRow = (await svc.findById(old.id))!;
    (oldRow as { completedAt: Date }).completedAt = new Date(now.getTime() - 8 * 24 * 60 * 60 * 1000);

    const result = await svc.gcSweep({
      pendingTimeoutMs: 60 * 60 * 1000, // 1h
      terminalRetentionMs: 7 * 24 * 60 * 60 * 1000, // 7d
    });
    expect(result.erroredOut).toBe(1);
    expect(result.deleted).toBe(1);

    // Stale pending was flipped to error.
    const staleAfter = await svc.findById(stale.id);
    expect(staleAfter?.status).toBe('error');
    expect(staleAfter?.errorMessage).toMatch(/expired in pending/);
    // Old completed is gone.
    expect(await svc.findById(old.id)).toBeNull();
    // Fresh pending untouched.
    expect((await svc.findById(fresh.id))?.status).toBe('pending');
  });
});