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

/**
 * Route-level tests for the v5 async inference task API.
 * Service + state-machine details are tested in
 * inference-task-service.test.ts; this file just covers the wire shapes
 * + owner scoping.
 */
import { describe, it, expect, vi, afterEach, beforeEach } from 'vitest';
import Fastify from 'fastify';
import type { FastifyInstance } from 'fastify';
import { registerInferenceTaskRoutes } from '../src/routes/inference-tasks.js';
import { errorHandler } from '../src/middleware/error-handler.js';
import type { IInferenceTaskService } from '../src/services/inference-task.service.js';
import type { LlmService, LlmView } from '../src/services/llm.service.js';
import type { IVirtualLlmService } from '../src/services/virtual-llm.service.js';
import type { InferenceTask } from '@prisma/client';

let app: FastifyInstance;

function makeRow(overrides: Partial<InferenceTask> = {}): InferenceTask {
  return {
    id: overrides.id ?? 'task-1',
    status: 'pending',
    poolName: 'qwen-pool',
    llmName: 'vllm-local',
    model: 'qwen3',
    tier: null,
    claimedBy: null,
    requestBody: { messages: [] } as InferenceTask['requestBody'],
    responseBody: null,
    errorMessage: null,
    streaming: false,
    createdAt: new Date('2026-04-28T00:00:00Z'),
    claimedAt: null,
    streamStartedAt: null,
    completedAt: null,
    ownerId: 'owner-1',
    agentId: null,
    ...overrides,
  };
}

function makeLlmView(overrides: Partial<LlmView> = {}): LlmView {
  return {
    id: 'llm-1',
    name: 'vllm-local',
    type: 'openai',
    model: 'qwen3',
    url: '',
    tier: 'fast',
    description: '',
    apiKeyRef: null,
    extraConfig: {},
    poolName: 'qwen-pool',
    kind: 'virtual',
    status: 'active',
    lastHeartbeatAt: null,
    inactiveSince: null,
    version: 1,
    createdAt: new Date(),
    updatedAt: new Date(),
    ...overrides,
  };
}

function mockTasks(rows: InferenceTask[] = []): IInferenceTaskService {
  const byId = new Map(rows.map((r) => [r.id, r]));
  return {
    enqueue: vi.fn(),
    waitFor: vi.fn(),
    tryClaim: vi.fn(),
    markRunning: vi.fn(),
    pushChunk: vi.fn(() => true),
    subscribeChunksUnsafe: vi.fn(() => () => undefined),
    complete: vi.fn(),
    fail: vi.fn(),
    cancel: vi.fn(async (id) => {
      const r = byId.get(id);
      if (r === undefined) return null;
      const next = { ...r, status: 'cancelled' as const, completedAt: new Date() };
      byId.set(id, next);
      return next;
    }),
    revertHeldBy: vi.fn(async () => []),
    findPendingForPools: vi.fn(async () => []),
    findById: vi.fn(async (id) => byId.get(id) ?? null),
    list: vi.fn(async (filter) => {
      let out = [...byId.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.status !== undefined) {
        const s = Array.isArray(filter.status) ? filter.status : [filter.status];
        out = out.filter((r) => s.includes(r.status));
      }
      return out;
    }),
    gcSweep: vi.fn(async () => ({ erroredOut: 0, deleted: 0 })),
  };
}

function mockLlms(view: LlmView | null): LlmService {
  return {
    getByName: vi.fn(async (name: string) => {
      if (view !== null && view.name === name) return view;
      const err = new Error(`Llm not found: ${name}`);
      (err as { name: string }).name = 'NotFoundError';
      throw err;
    }),
  } as unknown as LlmService;
}

function mockVirtualLlms(): IVirtualLlmService & { _calls: unknown[] } {
  const calls: unknown[] = [];
  return {
    _calls: calls,
    register: vi.fn(),
    heartbeat: vi.fn(),
    bindSession: vi.fn(),
    unbindSession: vi.fn(),
    enqueueInferTask: vi.fn(async (llmName, request, streaming, options) => {
      calls.push({ llmName, request, streaming, options });
      return {
        taskId: 'task-1',
        done: Promise.resolve({ status: 200, body: null }),
        onChunk: () => () => undefined,
      };
    }),
    completeTask: vi.fn(),
    pushTaskChunk: vi.fn(),
    failTask: vi.fn(),
    gcSweep: vi.fn(),
  } as unknown as IVirtualLlmService & { _calls: unknown[] };
}

afterEach(async () => {
  if (app) await app.close();
});

async function buildApp(deps: {
  tasks: IInferenceTaskService;
  llms: LlmService;
  virtualLlms: IVirtualLlmService;
  userId?: string;
}): Promise<FastifyInstance> {
  app = Fastify({ logger: false });
  app.setErrorHandler(errorHandler);
  // Stub the auth hook the routes rely on for ownerId.
  app.addHook('onRequest', async (request) => {
    request.userId = deps.userId ?? 'owner-1';
  });
  registerInferenceTaskRoutes(app, deps);
  await app.ready();
  return app;
}

beforeEach(() => {
  // Augment FastifyRequest with userId for the type system. We don't
  // need to do anything at runtime — the addHook above sets it.
});

describe('Inference Task Routes (v5 Stage 3)', () => {
  it('POST /api/v1/inference-tasks enqueues the task with failFast:false', async () => {
    const tasks = mockTasks([makeRow({ id: 'task-1' })]);
    const llms = mockLlms(makeLlmView());
    const virtualLlms = mockVirtualLlms();
    await buildApp({ tasks, llms, virtualLlms });

    const res = await app.inject({
      method: 'POST',
      url: '/api/v1/inference-tasks',
      payload: { llmName: 'vllm-local', request: { model: 'qwen3', messages: [{ role: 'user', content: 'hi' }] } },
    });
    expect(res.statusCode).toBe(201);
    const body = res.json<{ id: string; status: string; poolName: string; streaming: boolean }>();
    expect(body.id).toBe('task-1');
    expect(body.poolName).toBe('qwen-pool');
    expect(body.streaming).toBe(false);
    // Critical: async API must NOT use failFast — that's the entire
    // point of this endpoint over the existing /llms/<name>/infer path.
    // Owner is also threaded through so the resulting row carries the
    // authenticated user (otherwise foreign-owner 404 would fire on
    // every subsequent GET/DELETE).
    expect(virtualLlms.enqueueInferTask).toHaveBeenCalledWith(
      'vllm-local',
      expect.objectContaining({ messages: expect.any(Array) }),
      false,
      { failFast: false, ownerId: 'owner-1' },
    );
  });

  it('POST rejects when llmName is missing', async () => {
    await buildApp({ tasks: mockTasks(), llms: mockLlms(null), virtualLlms: mockVirtualLlms() });
    const res = await app.inject({
      method: 'POST',
      url: '/api/v1/inference-tasks',
      payload: { request: {} },
    });
    expect(res.statusCode).toBe(400);
    expect(res.json<{ error: string }>().error).toMatch(/llmName/);
  });

  it('POST rejects with 400 when targeting a public Llm (sync /infer path is the right tool)', async () => {
    await buildApp({
      tasks: mockTasks(),
      llms: mockLlms(makeLlmView({ kind: 'public' })),
      virtualLlms: mockVirtualLlms(),
    });
    const res = await app.inject({
      method: 'POST',
      url: '/api/v1/inference-tasks',
      payload: { llmName: 'vllm-local', request: { messages: [] } },
    });
    expect(res.statusCode).toBe(400);
    expect(res.json<{ error: string }>().error).toMatch(/not a virtual provider/);
  });

  it('GET /api/v1/inference-tasks/:id returns the row when owner matches', async () => {
    const tasks = mockTasks([makeRow({ id: 'task-1', ownerId: 'owner-1' })]);
    await buildApp({ tasks, llms: mockLlms(null), virtualLlms: mockVirtualLlms(), userId: 'owner-1' });
    const res = await app.inject({ method: 'GET', url: '/api/v1/inference-tasks/task-1' });
    expect(res.statusCode).toBe(200);
    expect(res.json<{ id: string }>().id).toBe('task-1');
  });

  it('GET /api/v1/inference-tasks/:id returns 404 (not 403) on a foreign owner', async () => {
    // Owner-scoped 404 prevents id-enumeration via differential status.
    const tasks = mockTasks([makeRow({ id: 'task-1', ownerId: 'someone-else' })]);
    await buildApp({ tasks, llms: mockLlms(null), virtualLlms: mockVirtualLlms(), userId: 'owner-1' });
    const res = await app.inject({ method: 'GET', url: '/api/v1/inference-tasks/task-1' });
    expect(res.statusCode).toBe(404);
  });

  it('GET /api/v1/inference-tasks lists only the caller\'s own tasks by default', async () => {
    const tasks = mockTasks([
      makeRow({ id: 'task-1', ownerId: 'owner-1' }),
      makeRow({ id: 'task-2', ownerId: 'owner-1' }),
      makeRow({ id: 'task-3', ownerId: 'someone-else' }),
    ]);
    await buildApp({ tasks, llms: mockLlms(null), virtualLlms: mockVirtualLlms(), userId: 'owner-1' });
    const res = await app.inject({ method: 'GET', url: '/api/v1/inference-tasks' });
    expect(res.statusCode).toBe(200);
    const rows = res.json<Array<{ id: string }>>();
    expect(rows.map((r) => r.id).sort()).toEqual(['task-1', 'task-2']);
  });

  it('DELETE /api/v1/inference-tasks/:id cancels a pending row', async () => {
    const tasks = mockTasks([makeRow({ id: 'task-1', ownerId: 'owner-1', status: 'pending' })]);
    await buildApp({ tasks, llms: mockLlms(null), virtualLlms: mockVirtualLlms(), userId: 'owner-1' });
    const res = await app.inject({ method: 'DELETE', url: '/api/v1/inference-tasks/task-1' });
    expect(res.statusCode).toBe(200);
    expect(res.json<{ status: string }>().status).toBe('cancelled');
    expect(tasks.cancel).toHaveBeenCalledWith('task-1');
  });

  it('DELETE on a foreign-owner task returns 404 without revealing existence', async () => {
    const tasks = mockTasks([makeRow({ id: 'task-1', ownerId: 'someone-else' })]);
    await buildApp({ tasks, llms: mockLlms(null), virtualLlms: mockVirtualLlms(), userId: 'owner-1' });
    const res = await app.inject({ method: 'DELETE', url: '/api/v1/inference-tasks/task-1' });
    expect(res.statusCode).toBe(404);
    expect(tasks.cancel).not.toHaveBeenCalled();
  });

  it('DELETE on an already-terminal task is a no-op (returns 200 with current row)', async () => {
    const tasks = mockTasks([makeRow({ id: 'task-1', ownerId: 'owner-1', status: 'completed' })]);
    await buildApp({ tasks, llms: mockLlms(null), virtualLlms: mockVirtualLlms(), userId: 'owner-1' });
    const res = await app.inject({ method: 'DELETE', url: '/api/v1/inference-tasks/task-1' });
    expect(res.statusCode).toBe(200);
    expect(res.json<{ status: string }>().status).toBe('completed');
    // cancel was NOT called — terminal rows aren't transitioned.
    expect(tasks.cancel).not.toHaveBeenCalled();
  });
});