mcpctl/src/db/tests/inference-task-schema.test.ts

/**
 * v5 db-level tests for the InferenceTask queue. Exercises the actual
 * column shapes + index lookups; the mcpd-side service tests cover the
 * state machine + signal channels with a mocked repo.
 */
import { describe, it, expect, beforeAll, afterAll, beforeEach } from 'vitest';
import type { PrismaClient } from '@prisma/client';
import { setupTestDb, cleanupTestDb, clearAllTables } from './helpers.js';

async function makeOwner(prisma: PrismaClient): Promise<string> {
  const u = await prisma.user.create({
    data: { email: `owner-${String(Date.now())}@test`, passwordHash: 'x' },
  });
  return u.id;
}

describe('InferenceTask schema (v5)', () => {
  let prisma: PrismaClient;

  beforeAll(async () => {
    prisma = await setupTestDb();
  }, 30_000);

  afterAll(async () => {
    await cleanupTestDb();
  });

  beforeEach(async () => {
    await clearAllTables(prisma);
  });

  it('defaults a fresh row to status=pending with claim/completion fields null', async () => {
    const ownerId = await makeOwner(prisma);
    const row = await prisma.inferenceTask.create({
      data: {
        poolName: 'qwen-pool',
        llmName: 'qwen-prod-1',
        model: 'qwen3-thinking',
        requestBody: { messages: [{ role: 'user', content: 'hi' }] },
        ownerId,
      },
    });
    expect(row.status).toBe('pending');
    expect(row.claimedBy).toBeNull();
    expect(row.claimedAt).toBeNull();
    expect(row.streamStartedAt).toBeNull();
    expect(row.completedAt).toBeNull();
    expect(row.responseBody).toBeNull();
    expect(row.streaming).toBe(false);
  });

  it('roundtrips streaming=true and a structured requestBody/responseBody', async () => {
    const ownerId = await makeOwner(prisma);
    const requestBody = {
      messages: [{ role: 'user', content: 'hello' }],
      temperature: 0.2,
      tools: [{ type: 'function', function: { name: 'noop' } }],
    };
    const row = await prisma.inferenceTask.create({
      data: {
        poolName: 'qwen-pool',
        llmName: 'qwen-prod-1',
        model: 'qwen3',
        requestBody,
        streaming: true,
        ownerId,
      },
    });
    expect(row.streaming).toBe(true);
    expect(row.requestBody).toEqual(requestBody);

    const completedAt = new Date();
    const responseBody = { choices: [{ message: { role: 'assistant', content: 'world' } }] };
    const updated = await prisma.inferenceTask.update({
      where: { id: row.id },
      data: { status: 'completed', responseBody, completedAt },
    });
    expect(updated.responseBody).toEqual(responseBody);
    expect(updated.completedAt?.getTime()).toBe(completedAt.getTime());
  });

  it('compound index supports the dispatcher\'s drain query (status + poolName IN ...)', async () => {
    // The actual EXPLAIN/index-use check is too brittle for unit tests;
    // here we verify the QUERY shape that the repo's findPendingForPools
    // issues — same WHERE/ORDER BY — returns the expected rows in FIFO
    // order. Index usage is implied by the Prisma model definition.
    const ownerId = await makeOwner(prisma);
    const t1 = await prisma.inferenceTask.create({
      data: { poolName: 'pool-a', llmName: 'a-1', model: 'm', requestBody: {}, ownerId },
    });
    await new Promise((r) => setTimeout(r, 5));
    const t2 = await prisma.inferenceTask.create({
      data: { poolName: 'pool-a', llmName: 'a-2', model: 'm', requestBody: {}, ownerId },
    });
    await prisma.inferenceTask.create({
      data: { poolName: 'pool-b', llmName: 'b-1', model: 'm', requestBody: {}, ownerId },
    });
    // One row in pool-a is no longer pending — must be excluded.
    await prisma.inferenceTask.create({
      data: { poolName: 'pool-a', llmName: 'a-3', model: 'm', requestBody: {}, ownerId, status: 'completed' },
    });

    const drained = await prisma.inferenceTask.findMany({
      where: { status: 'pending', poolName: { in: ['pool-a', 'pool-b'] } },
      orderBy: { createdAt: 'asc' },
    });
    expect(drained.map((r) => r.id)).toEqual([t1.id, t2.id, drained[2]!.id]);
    expect(drained.map((r) => r.poolName)).toEqual(['pool-a', 'pool-a', 'pool-b']);
  });

  it('claimedBy index supports unbindSession revert (worker disconnect path)', async () => {
    const ownerId = await makeOwner(prisma);
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'claimed', claimedBy: 'sess-A' },
    });
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'running', claimedBy: 'sess-A' },
    });
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'claimed', claimedBy: 'sess-B' },
    });
    // Completed-but-claimedBy=sess-A row: must NOT revert (terminal state).
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'completed', claimedBy: 'sess-A' },
    });

    const heldByA = await prisma.inferenceTask.findMany({
      where: { claimedBy: 'sess-A', status: { in: ['claimed', 'running'] } },
    });
    expect(heldByA).toHaveLength(2);
  });

  it('GC predicate (terminal + completedAt < cutoff) is index-friendly and filters correctly', async () => {
    const ownerId = await makeOwner(prisma);
    const old = new Date(Date.now() - 8 * 24 * 60 * 60 * 1000); // 8 d ago
    const recent = new Date(Date.now() - 1 * 60 * 60 * 1000); // 1 h ago
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'completed', completedAt: old },
    });
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'error', completedAt: old, errorMessage: 'boom' },
    });
    // Inside retention — must not be picked up by GC.
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'completed', completedAt: recent },
    });
    // Pending row — must not be picked up by terminal GC.
    await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, status: 'pending' },
    });

    const cutoff = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000);
    const expired = await prisma.inferenceTask.findMany({
      where: {
        status: { in: ['completed', 'error', 'cancelled'] },
        completedAt: { lt: cutoff },
      },
    });
    expect(expired).toHaveLength(2);
  });

  it('agentId is nullable — direct chat-llm tasks have no agent', async () => {
    const ownerId = await makeOwner(prisma);
    const row = await prisma.inferenceTask.create({
      data: { poolName: 'p', llmName: 'l', model: 'm', requestBody: {}, ownerId, agentId: null },
    });
    expect(row.agentId).toBeNull();
  });
});