feat(mcpd+db): durable InferenceTask queue + state machine (v5 Stage 1)
The persistence + signaling layer for v5. No integration with the
existing in-flight inference path yet — that's Stage 2. This commit
just lands the durable queue underneath, with a state machine that
mcpd's HTTP handlers, the worker result-POST route, and the GC sweep
will all build on.
Schema (src/db/prisma/schema.prisma + migration):
- New `InferenceTask` model + `InferenceTaskStatus` enum
(pending|claimed|running|completed|error|cancelled).
- Routing fields stored at enqueue time so a later rename of
`Llm.poolName` doesn't reroute already-queued work: `poolName`
(effective pool key), `llmName` (pinned target), `model`, `tier`.
- Worker tracking: `claimedBy` (providerSessionId) + `claimedAt`,
cleared on revert.
- Bodies as `Json`: requestBody (always set), responseBody (set at
completion). Streaming chunks are NOT persisted — too expensive at
delta granularity. The final assembled body lands once per task.
- Lifecycle timestamps: createdAt, claimedAt, streamStartedAt,
completedAt. Plus ownerId (RBAC + audit) and agentId (null for
direct chat-llm calls).
- Indexes for the hot paths: (status, poolName) for the dispatcher's
drain query, claimedBy for the disconnect revert, completedAt for
the GC retention sweep, owner/agent for the async API listing.
Repository (src/mcpd/src/repositories/inference-task.repository.ts):
- CRUD + state transitions as conditional CAS via `updateMany`. Two
workers racing to claim the same row both run the UPDATE; whichever
the DB serializes first sees affected=1 and gets the row, the loser
sees 0 and falls through to the next candidate. No application-
level locking required.
- findPendingForPools(poolNames[]) for the worker drain on bind.
- findHeldBy(claimedBy) for the unbindSession revert.
- findStalePending + findExpiredTerminal for the GC sweep.
Service (src/mcpd/src/services/inference-task.service.ts):
- Owns the in-process EventEmitter that wakes blocked HTTP handlers
when a worker POSTs results. The DB row is the source of truth for
*state*; the EventEmitter just signals "go re-read row X" so we
don't have to poll. Single-instance assumption for v5; pg
LISTEN/NOTIFY is the v6 swap when scaling horizontally — no schema
change needed, just replace the emitter wakeup.
- waitFor(taskId, timeoutMs) returns { done, chunks }: the terminal
promise + an async iterator of streaming deltas. Throws on cancel
(clear message) or error (worker's errorMessage propagates) or
timeout. Polls the row once at subscribe time so an already-
terminal task resolves immediately without waiting for an event
that's never coming.
- gcSweep flips stale pending rows to error (with a clear message
about the timeout) and deletes terminal rows past retention.
Defaults: 1h pending timeout, 7d terminal retention; both
configurable.
Tests:
- 6 db-level schema tests (defaults, json roundtrip, drain query
shape, claimedBy filter, GC predicate, agentId nullable).
- 13 service tests covering enqueue, the CAS race on tryClaim,
complete/fail/cancel, idempotent terminal transitions, revertHeldBy
on disconnect, and the full waitFor signal lifecycle (immediate
resolve, wake on event, chunk streaming, cancel/error/timeout
paths). Plus a gcSweep test with a fixed clock.
mcpd 881/881 (was 868; +13). db pool-schema 14/14, +6 new
inference-task-schema. Pre-existing failures in models.test.ts
(Secret FK fixture issue, also fails on main HEAD) are unrelated.
Stage 2 (next): VirtualLlmService rewires through this — remove the
in-memory pendingTasks map; enqueue creates a row, dispatch picks an
active session, the result-route updates the row + emits the wakeup.
Worker disconnect reverts; worker bind drains.
This commit is contained in:
Michal
330
src/mcpd/tests/inference-task-service.test.ts
Normal file
330
src/mcpd/tests/inference-task-service.test.ts
Normal file
@@ -0,0 +1,330 @@
|
||||
/**
|
||||
* 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');
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user