Files
qinglong/packages/ql3-ai/test/localModelInvocationRepository.test.cjs
T

1218 lines
34 KiB
JavaScript

const assert = require('node:assert/strict');
const test = require('node:test');
const { DatabaseSync } = require('node:sqlite');
const {
createStepRunRecord,
transitionStepRunMutation,
} = require('@qinglong/runtime-core/step-run');
const {
ModelInvocationConflictError,
ModelInvocationRepositoryUnavailableError,
createModelInvocationCompletionCommand,
createModelInvocationMutationIdentity,
createModelInvocationStartCommand,
} = require('../dist/model-invocation/modelInvocation.js');
const {
DurableModelInvocationCoordinator,
DurableModelInvocationRecovery,
} = require('../dist/model-invocation/durableModelInvocationCoordinator.js');
const {
DurableModelInvocationResolutionCoordinator,
} = require('../dist/model-invocation/modelInvocationResolution.js');
const {
BoundedModelGateway,
ModelInvocationReplayBlockedError,
} = require('../dist/model-gateway/gateway.js');
const {
LOCAL_MODEL_INVOCATION_MIGRATION_PLAN_DIGEST,
migrateLocalModelInvocationFeature,
} = require('@qinglong/ai/model-invocation-migration');
const {
LocalModelInvocationFeatureActivationRepository,
createLocalModelInvocationFeatureTransitionCommand,
} = require('@qinglong/ai/local-feature-activation');
const {
LocalModelInvocationRepository,
} = require('../dist/model-invocation/localModelInvocationRepository.js');
const {
ModelInvocationProjectQuotaExceededError,
createModelInvocationQuotaAdmission,
} = require('../dist/usage/usageQuota.js');
const {
StaticModelPriceCatalog,
createModelPriceCatalogEntry,
createModelInvocationPriceQuote,
} = require('../dist/pricing/pricing.js');
const NOW = 1_000_000;
function createMainSqliteContract(client) {
client.exec(`
CREATE TABLE "QingLong3SchemaMigrations" (
migration_id TEXT PRIMARY KEY,
stream_id TEXT NOT NULL,
dialect TEXT NOT NULL,
checksum TEXT NOT NULL,
applied_at_ms INTEGER NOT NULL
);
CREATE TABLE "Runs" (
id TEXT PRIMARY KEY,
project_id TEXT NOT NULL,
status TEXT NOT NULL,
version INTEGER NOT NULL,
event_sequence INTEGER NOT NULL
);
CREATE TABLE "RunEvents" (
id TEXT PRIMARY KEY,
run_id TEXT NOT NULL,
sequence INTEGER NOT NULL,
type TEXT NOT NULL,
dedupe_key TEXT NOT NULL,
actor_type TEXT NOT NULL,
actor_id TEXT,
attempt_id TEXT,
step_run_id TEXT,
payload TEXT NOT NULL,
created_at_ms INTEGER NOT NULL,
UNIQUE (run_id, sequence),
UNIQUE (run_id, dedupe_key)
);
CREATE TABLE "StepRuns" (
id TEXT PRIMARY KEY,
run_id TEXT NOT NULL,
kind TEXT NOT NULL,
status TEXT NOT NULL,
version INTEGER NOT NULL,
attempt_count INTEGER NOT NULL,
output_ref TEXT,
approval_request_id TEXT,
ready_at_ms INTEGER,
started_at_ms INTEGER,
finished_at_ms INTEGER,
result_code TEXT,
error_summary TEXT,
updated_at_ms INTEGER NOT NULL,
last_mutation_id TEXT NOT NULL,
step_run_digest TEXT NOT NULL,
step_run_json TEXT NOT NULL,
UNIQUE (run_id, id)
);
CREATE TABLE "StepRunMutations" (
mutation_id TEXT PRIMARY KEY,
mutation_digest TEXT NOT NULL,
run_id TEXT NOT NULL,
step_run_id TEXT NOT NULL,
step_run_digest TEXT NOT NULL,
event_id TEXT NOT NULL,
event_sequence INTEGER NOT NULL,
run_version INTEGER NOT NULL,
step_run_json TEXT NOT NULL,
committed_at_ms INTEGER NOT NULL
);
`);
}
function audit(phase, overrides = {}) {
return {
phase,
projectId: 'project-a',
runId: 'run-a',
stepRunId: 'step-a',
traceId: 'trace-a',
requestId: 'request-a',
provider: 'remote',
model: 'model-a',
policyRevision: 'policy-1',
requestDigest: `sha256:${'b'.repeat(64)}`,
deadlineAtMs: NOW + 10_000,
inputBytes: 128,
maxOutputTokens: 64,
outputBytes: 0,
usage: null,
errorCode: null,
occurredAtMs: NOW,
...overrides,
};
}
async function fixture() {
const client = new DatabaseSync(':memory:');
client.exec('PRAGMA foreign_keys = ON');
createMainSqliteContract(client);
await migrateLocalModelInvocationFeature(client);
new LocalModelInvocationFeatureActivationRepository(client).transition(
createLocalModelInvocationFeatureTransitionCommand({
featureId: 'model-invocation',
expectedGeneration: 0,
expectedState: null,
state: 'active',
mutationId: 'model-invocation-test-feature-activation',
requestId: 'model-invocation-test-feature-request',
expectedMigrationDigest: LOCAL_MODEL_INVOCATION_MIGRATION_PLAN_DIGEST,
safety: {
mode: 'fresh_database',
backupEvidenceDigest: null,
},
principal: {
subject: { type: 'user', id: 'test-owner' },
authenticationId: 'local_ai_feature:test-proof',
authenticatedAtMs: 1,
expiresAtMs: 301_000,
assurance: 'local_console',
},
}),
);
const ready = createStepRunRecord({
id: 'step-a',
runId: 'run-a',
stepKey: 'summarize',
kind: 'model',
definitionRef: 'prompt:summary@1',
definitionDigest: 'a'.repeat(64),
required: true,
initialStatus: 'ready',
inputRef: 'artifact:input-a',
mutationId: 'create-step-a',
createdAtMs: NOW - 1,
});
client
.prepare(
`INSERT INTO "Runs"
(id, project_id, status, version, event_sequence)
VALUES (?, ?, ?, ?, ?)`,
)
.run('run-a', 'project-a', 'running', 1, 1);
client
.prepare(
`INSERT INTO "StepRuns" (
id, run_id, kind, status, version, attempt_count, output_ref,
approval_request_id, ready_at_ms, started_at_ms, finished_at_ms,
result_code, error_summary, updated_at_ms, last_mutation_id,
step_run_digest, step_run_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
)
.run(
ready.id,
ready.runId,
ready.kind,
ready.status,
ready.version,
ready.attemptCount,
ready.outputRef,
ready.approvalRequestId,
ready.readyAtMs,
ready.startedAtMs,
ready.finishedAtMs,
ready.resultCode,
ready.errorSummary,
ready.updatedAtMs,
ready.lastMutationId,
ready.stepRunDigest,
JSON.stringify(ready),
);
return {
client,
ready,
repository: new LocalModelInvocationRepository(client),
};
}
function startCommand(ready) {
const identity = createModelInvocationMutationIdentity('request-a', 'start');
return createModelInvocationStartCommand(
audit('admitted'),
transitionStepRunMutation(
ready,
{
expectedVersion: ready.version,
expectedDigest: ready.stepRunDigest,
mutationId: identity.mutationId,
to: 'running',
atMs: NOW,
},
{
expectedRunVersion: 1,
expectedRunEventSequence: 1,
eventId: identity.eventId,
dedupeKey: identity.dedupeKey,
actor: { type: 'executor', id: 'model-gateway' },
},
),
);
}
function completionCommand(start) {
const identity = createModelInvocationMutationIdentity(
'request-a',
'completion',
);
return createModelInvocationCompletionCommand(
start.start,
audit('completed', {
occurredAtMs: NOW + 25,
outputBytes: 12,
usage: { inputTokens: 5, outputTokens: 2, totalTokens: 7 },
}),
transitionStepRunMutation(
start.stepRunMutation.stepRun,
{
expectedVersion: start.start.startedStepRunVersion,
expectedDigest: start.start.startedStepRunDigest,
mutationId: identity.mutationId,
to: 'succeeded',
atMs: NOW + 25,
outputRef: 'model-invocation:request-a',
},
{
expectedRunVersion: 2,
expectedRunEventSequence: 2,
eventId: identity.eventId,
dedupeKey: identity.dedupeKey,
actor: { type: 'executor', id: 'model-gateway' },
},
),
);
}
test('SQLite admission is transaction-fenced by the active local AI feature', async () => {
const { client, ready, repository } = await fixture();
new LocalModelInvocationFeatureActivationRepository(client).transition(
createLocalModelInvocationFeatureTransitionCommand({
featureId: 'model-invocation',
expectedGeneration: 1,
expectedState: 'active',
state: 'inactive',
mutationId: 'model-invocation-test-feature-deactivation',
requestId: 'model-invocation-test-feature-deactivation-request',
expectedMigrationDigest: LOCAL_MODEL_INVOCATION_MIGRATION_PLAN_DIGEST,
safety: {
mode: 'preserve_existing',
backupEvidenceDigest: null,
},
principal: {
subject: { type: 'user', id: 'test-owner' },
authenticationId: 'local_ai_feature:test-proof',
authenticatedAtMs: 1,
expiresAtMs: 301_000,
assurance: 'local_console',
},
}),
);
await assert.rejects(
repository.admit(startCommand(ready)),
ModelInvocationRepositoryUnavailableError,
);
assert.equal(
client
.prepare(`SELECT count(*) AS count FROM "ModelInvocationStarts"`)
.get().count,
0,
);
assert.equal(
client.prepare(`SELECT count(*) AS count FROM "StepRunMutations"`).get()
.count,
0,
);
client.close();
});
test('SQLite atomically admits and completes one model StepRun', async () => {
const { client, ready, repository } = await fixture();
const start = startCommand(ready);
const completion = completionCommand(start);
assert.equal((await repository.admit(start)).status, 'created');
assert.equal((await repository.admit(start)).status, 'existing');
assert.deepEqual(await repository.findStart('request-a'), start.start);
assert.equal((await repository.complete(completion)).status, 'created');
assert.equal((await repository.complete(completion)).status, 'existing');
assert.deepEqual(
await repository.findCompletion('request-a'),
completion.completion,
);
const usage = await repository.findUsage('request-a');
assert.deepEqual(
{
invocationId: usage.invocationId,
projectId: usage.projectId,
completionDigest: usage.completionDigest,
outcome: usage.outcome,
inputTokens: usage.inputTokens,
outputTokens: usage.outputTokens,
totalTokens: usage.totalTokens,
costMicros: usage.costMicros,
},
{
invocationId: 'request-a',
projectId: 'project-a',
completionDigest: completion.completion.completionDigest,
outcome: 'succeeded',
inputTokens: 5,
outputTokens: 2,
totalTokens: 7,
costMicros: null,
},
);
assert.deepEqual(
await repository.listProjectUsage({
projectId: 'project-a',
fromMsInclusive: NOW,
toMsExclusive: NOW + 100,
limit: 1,
}),
{ records: [usage], hasMore: false },
);
assert.deepEqual(
await repository.summarizeProjectUsage({
projectId: 'project-a',
fromMsInclusive: NOW,
toMsExclusive: NOW + 100,
}),
{
invocationCount: 1,
inputTokens: 5,
outputTokens: 2,
totalTokens: 7,
knownCostMicros: 0,
unknownCostInvocations: 1,
},
);
const storedStep = client
.prepare(
`SELECT status, version, output_ref FROM "StepRuns"
WHERE id = 'step-a'`,
)
.get();
assert.equal(storedStep.status, 'succeeded');
assert.equal(storedStep.version, 3);
assert.equal(storedStep.output_ref, 'model-invocation:request-a');
const storedRun = client
.prepare(`SELECT version, event_sequence FROM "Runs" WHERE id = 'run-a'`)
.get();
assert.equal(storedRun.version, 3);
assert.equal(storedRun.event_sequence, 3);
assert.equal(
client.prepare('SELECT count(*) AS count FROM "RunEvents"').get().count,
2,
);
assert.equal(
client.prepare('SELECT count(*) AS count FROM "StepRunMutations"').get()
.count,
2,
);
assert.equal(
client
.prepare('SELECT count(*) AS count FROM "ModelInvocationUsageLedger"')
.get().count,
1,
);
assert.equal(
client.prepare('PRAGMA integrity_check').get().integrity_check,
'ok',
);
client.close();
});
test('SQLite completion and usage ledger roll back together', async () => {
const { client, ready, repository } = await fixture();
const start = startCommand(ready);
const completion = completionCommand(start);
await repository.admit(start);
client
.prepare(`UPDATE "StepRuns" SET status = 'cancelled' WHERE id = 'step-a'`)
.run();
await assert.rejects(
repository.complete(completion),
ModelInvocationConflictError,
);
assert.equal(
client
.prepare('SELECT count(*) AS count FROM "ModelInvocationCompletions"')
.get().count,
0,
);
assert.equal(
client
.prepare('SELECT count(*) AS count FROM "ModelInvocationUsageLedger"')
.get().count,
0,
);
client.close();
});
test('SQLite completion replay fails closed when its usage ledger is missing', async () => {
const { client, ready, repository } = await fixture();
const start = startCommand(ready);
const completion = completionCommand(start);
await repository.admit(start);
await repository.complete(completion);
client
.prepare(`DELETE FROM "ModelInvocationUsageLedger" WHERE invocation_id = ?`)
.run('request-a');
await assert.rejects(
repository.complete(completion),
ModelInvocationConflictError,
);
client.close();
});
test('SQLite state conflict rolls back every invocation fact', async () => {
const { client, ready, repository } = await fixture();
const start = startCommand(ready);
client
.prepare(`UPDATE "StepRuns" SET status = 'cancelled' WHERE id = 'step-a'`)
.run();
await assert.rejects(repository.admit(start), ModelInvocationConflictError);
assert.equal(
client
.prepare('SELECT count(*) AS count FROM "ModelInvocationStarts"')
.get().count,
0,
);
assert.equal(
client.prepare('SELECT count(*) AS count FROM "RunEvents"').get().count,
0,
);
assert.equal(
client.prepare('SELECT count(*) AS count FROM "StepRunMutations"').get()
.count,
0,
);
client.close();
});
test('SQLite corrupted durable JSON fails closed', async () => {
const { client, ready, repository } = await fixture();
const start = startCommand(ready);
await repository.admit(start);
client.exec('PRAGMA ignore_check_constraints = ON');
client
.prepare(
`UPDATE "ModelInvocationStarts"
SET record_json = json_set(record_json, '$.traceId', 'trace-drift')`,
)
.run();
await assert.rejects(
repository.findStart('request-a'),
ModelInvocationRepositoryUnavailableError,
);
client.close();
});
test('durable SQLite gateway executes once and blocks provider replay', async () => {
const { client, repository } = await fixture();
let providerCalls = 0;
const gateway = new BoundedModelGateway({
providers: [
{
type: 'remote',
async listModels() {
return [{ id: 'model-a' }];
},
async generate() {
providerCalls += 1;
return {
provider: 'remote',
model: 'model-a',
text: 'safe summary',
finishReason: 'stop',
usage: { inputTokens: 5, outputTokens: 2, totalTokens: 7 },
};
},
async *stream() {
throw new Error('not used');
},
},
],
pricing: {
async resolve() {
throw new Error('pricing must remain unreachable');
},
},
policies: {
async resolve() {
return {
revision: 'policy-1',
allowedProviders: ['remote'],
allowedModels: ['model-a'],
maxInputBytes: 4096,
maxOutputBytes: 4096,
maxOutputTokens: 64,
maxTotalTokens: 256,
maxCostMicros: null,
priceRevision: null,
};
},
},
audit: new DurableModelInvocationCoordinator(repository),
maxConcurrent: 1,
now: () => NOW,
});
const request = {
provider: 'remote',
model: 'model-a',
messages: [{ role: 'user', content: 'sensitive prompt' }],
maxOutputTokens: 64,
};
const context = {
projectId: 'project-a',
runId: 'run-a',
stepRunId: 'step-a',
traceId: 'trace-a',
requestId: 'request-a',
deadlineAtMs: NOW + 10_000,
};
assert.equal((await gateway.generate(request, context)).text, 'safe summary');
await assert.rejects(
gateway.generate(request, context),
ModelInvocationReplayBlockedError,
);
assert.equal(providerCalls, 1);
assert.equal(
JSON.stringify(await repository.findStart('request-a')).includes(
'sensitive prompt',
),
false,
);
client.close();
});
test('SQLite quota admission and settlement stay atomic with the model StepRun', async () => {
const { client, repository } = await fixture();
let providerCalls = 0;
const gateway = new BoundedModelGateway({
providers: [
{
type: 'remote',
async listModels() {
return [{ id: 'model-a' }];
},
async generate() {
providerCalls += 1;
return {
provider: 'remote',
model: 'model-a',
text: 'safe summary',
finishReason: 'stop',
usage: {
inputTokens: 5,
outputTokens: 2,
totalTokens: 7,
costMicros: 700,
},
};
},
async *stream() {
throw new Error('not used');
},
},
],
pricing: {
async resolve() {
throw new Error('pricing must remain unreachable');
},
},
policies: {
async resolve() {
return {
revision: 'policy-1',
allowedProviders: ['remote'],
allowedModels: ['model-a'],
maxInputBytes: 4096,
maxOutputBytes: 4096,
maxOutputTokens: 64,
maxTotalTokens: 256,
maxCostMicros: null,
priceRevision: null,
projectQuota: {
revision: 'quota-1',
windowMs: 3_600_000,
maxInvocations: 1,
maxTokens: 256,
maxCostMicros: null,
},
};
},
},
audit: new DurableModelInvocationCoordinator(repository),
maxConcurrent: 1,
now: () => NOW,
});
await gateway.generate(
{
provider: 'remote',
model: 'model-a',
messages: [{ role: 'user', content: 'sensitive prompt' }],
maxOutputTokens: 64,
},
{
projectId: 'project-a',
runId: 'run-a',
stepRunId: 'step-a',
traceId: 'trace-a',
requestId: 'request-a',
deadlineAtMs: NOW + 10_000,
},
);
const reservation = await repository.findQuotaReservation('request-a');
const settlement = await repository.findQuotaSettlement('request-a');
assert.equal(providerCalls, 1);
assert.equal(reservation.reservedTokens, 256);
assert.equal(settlement.effectiveTokens, 7);
assert.equal(settlement.effectiveCostMicros, null);
assert.equal(settlement.retainedTokenReservation, false);
assert.equal(settlement.retainedCostReservation, false);
assert.equal((await repository.findUsage('request-a')).costMicros, 700);
assert.deepEqual(
await repository.readQuotaWindowUsage(
'project-a',
reservation.reservedAtMs,
),
{
projectId: 'project-a',
windowStartMs: reservation.windowStartMs,
windowEndMs: reservation.windowEndMs,
invocationCount: 1,
effectiveTokens: 7,
effectiveCostMicros: 0,
unknownCostInvocations: 1,
},
);
client.close();
});
test('SQLite price quote, billing, quota and completion commit atomically', async () => {
const { client, repository } = await fixture();
const pricing = new StaticModelPriceCatalog([
createModelPriceCatalogEntry({
provider: 'remote',
model: 'model-a',
priceRevision: 'price-1',
currency: 'USD',
inputMicrosPerMillionTokens: 150_000,
outputMicrosPerMillionTokens: 600_000,
publishedAtMs: NOW - 1,
}),
]);
const gateway = new BoundedModelGateway({
providers: [
{
type: 'remote',
async listModels() {
return [{ id: 'model-a' }];
},
async generate() {
return {
provider: 'remote',
model: 'model-a',
text: 'safe summary',
finishReason: 'stop',
usage: {
inputTokens: 5,
outputTokens: 2,
totalTokens: 7,
costMicros: 99_999,
},
};
},
async *stream() {
throw new Error('not used');
},
},
],
pricing,
policies: {
async resolve() {
return {
revision: 'policy-1',
allowedProviders: ['remote'],
allowedModels: ['model-a'],
maxInputBytes: 4096,
maxOutputBytes: 4096,
maxOutputTokens: 64,
maxTotalTokens: 256,
maxCostMicros: 100,
priceRevision: 'price-1',
projectQuota: {
revision: 'quota-1',
windowMs: 3_600_000,
maxInvocations: 10,
maxTokens: 10_000,
maxCostMicros: 1_000,
},
};
},
},
audit: new DurableModelInvocationCoordinator(repository),
maxConcurrent: 1,
now: () => NOW,
});
const result = await gateway.generate(
{
provider: 'remote',
model: 'model-a',
messages: [{ role: 'user', content: 'sensitive prompt' }],
maxOutputTokens: 64,
},
{
projectId: 'project-a',
runId: 'run-a',
stepRunId: 'step-a',
traceId: 'trace-a',
requestId: 'request-a',
deadlineAtMs: NOW + 10_000,
},
);
const quote = await repository.findPriceQuote('request-a');
const priceSettlement = await repository.findPriceSettlement('request-a');
const quotaReservation = await repository.findQuotaReservation('request-a');
const quotaSettlement = await repository.findQuotaSettlement('request-a');
const usage = await repository.findUsage('request-a');
assert.equal(result.usage.costMicros, 3);
assert.equal(quote.priceRevision, 'price-1');
assert.equal(quote.reservedCostMicros, 68);
assert.equal(priceSettlement.costMicros, 3);
assert.equal(usage.costMicros, 3);
assert.equal(quotaReservation.reservedCostMicros, 68);
assert.equal(quotaSettlement.effectiveCostMicros, 3);
assert.equal(quotaSettlement.retainedCostReservation, false);
const storedStart = await repository.findStart('request-a');
assert.deepEqual(
await new DurableModelInvocationCoordinator(repository).recordWithPricing(
audit('admitted', {
requestDigest: storedStart.requestDigest,
inputBytes: storedStart.inputBytes,
maxOutputTokens: storedStart.maxOutputTokens,
deadlineAtMs: storedStart.deadlineAtMs,
occurredAtMs: storedStart.admittedAtMs,
}),
quote,
createModelInvocationQuotaAdmission({
invocationId: 'request-a',
projectId: 'project-a',
modelPolicyRevision: 'policy-1',
reservedTokens: 256,
reservedCostMicros: 68,
quota: {
revision: 'quota-1',
windowMs: 3_600_000,
maxInvocations: 10,
maxTokens: 10_000,
maxCostMicros: 1_000,
},
}),
),
{ status: 'existing' },
);
assert.equal(
client
.prepare(`SELECT count(*) AS count FROM "ModelInvocationPriceQuotes"`)
.get().count,
1,
);
assert.equal(
client
.prepare(
`SELECT count(*) AS count FROM "ModelInvocationPriceSettlements"`,
)
.get().count,
1,
);
client.close();
});
test('SQLite quota exhaustion rolls back admission before provider I/O', async () => {
const { client, repository } = await fixture();
let providerCalls = 0;
const gateway = new BoundedModelGateway({
providers: [
{
type: 'remote',
async listModels() {
return [];
},
async generate() {
providerCalls += 1;
throw new Error('must not execute');
},
async *stream() {
throw new Error('must not execute');
},
},
],
pricing: {
async resolve() {
throw new Error('pricing must remain unreachable');
},
},
policies: {
async resolve() {
return {
revision: 'policy-1',
allowedProviders: ['remote'],
allowedModels: ['model-a'],
maxInputBytes: 4096,
maxOutputBytes: 4096,
maxOutputTokens: 64,
maxTotalTokens: 256,
maxCostMicros: null,
priceRevision: null,
projectQuota: {
revision: 'quota-1',
windowMs: 3_600_000,
maxInvocations: 10,
maxTokens: 255,
maxCostMicros: null,
},
};
},
},
audit: new DurableModelInvocationCoordinator(repository),
maxConcurrent: 1,
now: () => NOW,
});
await assert.rejects(
gateway.generate(
{
provider: 'remote',
model: 'model-a',
messages: [{ role: 'user', content: 'sensitive prompt' }],
maxOutputTokens: 64,
},
{
projectId: 'project-a',
runId: 'run-a',
stepRunId: 'step-a',
traceId: 'trace-a',
requestId: 'request-a',
deadlineAtMs: NOW + 10_000,
},
),
ModelInvocationProjectQuotaExceededError,
);
assert.equal(providerCalls, 0);
assert.equal(await repository.findStart('request-a'), null);
assert.equal(await repository.findQuotaReservation('request-a'), null);
assert.equal(
client.prepare(`SELECT status FROM "StepRuns" WHERE id = 'step-a'`).get()
.status,
'ready',
);
client.close();
});
test('SQLite recovery marks expired incomplete invocation outcome unknown', async () => {
const { client, repository } = await fixture();
const coordinator = new DurableModelInvocationCoordinator(repository);
assert.deepEqual(await coordinator.record(audit('admitted')), {
status: 'created',
});
const page = await repository.listIncomplete(1);
assert.equal(page.candidates.length, 1);
assert.equal(page.hasMore, false);
const summary = await new DurableModelInvocationRecovery(
repository,
coordinator,
).recover(1);
assert.deepEqual(
{
scanned: summary.scanned,
recovered: summary.recovered,
alreadyCompleted: summary.alreadyCompleted,
failed: summary.failed,
hasMore: summary.hasMore,
},
{
scanned: 1,
recovered: 1,
alreadyCompleted: 0,
failed: 0,
hasMore: false,
},
);
const completion = await repository.findCompletion('request-a');
assert.equal(completion.outcome, 'outcome_unknown');
assert.equal(completion.errorCode, 'MODEL_INVOCATION_OUTCOME_UNKNOWN');
assert.equal(await repository.findUsage('request-a'), null);
assert.equal(
client.prepare(`SELECT status FROM "StepRuns" WHERE id = 'step-a'`).get()
.status,
'lost',
);
client.close();
});
test('SQLite recovery retains a quota reservation when provider usage is unknown', async () => {
const { client, repository } = await fixture();
const coordinator = new DurableModelInvocationCoordinator(repository);
const admission = createModelInvocationQuotaAdmission({
invocationId: 'request-a',
projectId: 'project-a',
modelPolicyRevision: 'policy-1',
reservedTokens: 256,
reservedCostMicros: 5_000,
quota: {
revision: 'quota-1',
windowMs: 3_600_000,
maxInvocations: 10,
maxTokens: 10_000,
maxCostMicros: 50_000,
},
});
assert.deepEqual(
await coordinator.recordWithQuota(audit('admitted'), admission),
{ status: 'created' },
);
const summary = await new DurableModelInvocationRecovery(
repository,
coordinator,
).recover(1);
assert.equal(summary.recovered, 1);
const settlement = await repository.findQuotaSettlement('request-a');
assert.equal(settlement.effectiveTokens, 256);
assert.equal(settlement.effectiveCostMicros, 5_000);
assert.equal(settlement.retainedTokenReservation, true);
assert.equal(settlement.retainedCostReservation, true);
client.close();
});
test('SQLite priced recovery preserves its quote without inventing a settlement', async () => {
const { client, repository } = await fixture();
const coordinator = new DurableModelInvocationCoordinator(repository);
const quote = createModelInvocationPriceQuote(
createModelPriceCatalogEntry({
provider: 'remote',
model: 'model-a',
priceRevision: 'price-1',
currency: 'USD',
inputMicrosPerMillionTokens: 150_000,
outputMicrosPerMillionTokens: 600_000,
publishedAtMs: NOW - 1,
}),
{
invocationId: 'request-a',
projectId: 'project-a',
modelPolicyRevision: 'policy-1',
maxTotalTokens: 256,
maxOutputTokens: 64,
},
);
const admission = createModelInvocationQuotaAdmission({
invocationId: 'request-a',
projectId: 'project-a',
modelPolicyRevision: 'policy-1',
reservedTokens: 256,
reservedCostMicros: quote.reservedCostMicros,
quota: {
revision: 'quota-1',
windowMs: 3_600_000,
maxInvocations: 10,
maxTokens: 10_000,
maxCostMicros: 50_000,
},
});
assert.deepEqual(
await coordinator.recordWithPricing(audit('admitted'), quote, admission),
{ status: 'created' },
);
const summary = await new DurableModelInvocationRecovery(
repository,
coordinator,
).recover(1);
assert.equal(summary.recovered, 1);
assert.equal(
(await repository.findCompletion('request-a')).outcome,
'outcome_unknown',
);
assert.equal(
(await repository.findPriceQuote('request-a')).quoteDigest,
quote.quoteDigest,
);
assert.equal(await repository.findPriceSettlement('request-a'), null);
const quotaSettlement = await repository.findQuotaSettlement('request-a');
assert.equal(quotaSettlement.effectiveCostMicros, quote.reservedCostMicros);
assert.equal(quotaSettlement.retainedTokenReservation, true);
assert.equal(quotaSettlement.retainedCostReservation, true);
client.close();
});
test('SQLite manual retry preserves the unknown attempt and admits a new invocation', async () => {
const { client, repository } = await fixture();
const coordinator = new DurableModelInvocationCoordinator(repository);
assert.deepEqual(await coordinator.record(audit('admitted')), {
status: 'created',
});
await new DurableModelInvocationRecovery(repository, coordinator).recover(1);
const completion = await repository.findCompletion('request-a');
assert.equal(completion.outcome, 'outcome_unknown');
const resolver = new DurableModelInvocationResolutionCoordinator(repository);
const resolvedAtMs = completion.completedAtMs + 1;
const resolution = await resolver.resolve({
invocationId: 'request-a',
decision: 'retry',
resolvedByUserId: 'user-a',
resolvedAtMs,
});
assert.equal(resolution.status, 'created');
assert.equal(resolution.record.decision, 'retry');
assert.equal(
(
await resolver.resolve({
invocationId: 'request-a',
decision: 'retry',
resolvedByUserId: 'user-a',
resolvedAtMs: resolvedAtMs + 100,
})
).status,
'existing',
);
await assert.rejects(
resolver.resolve({
invocationId: 'request-a',
decision: 'fail',
resolvedByUserId: 'user-a',
resolvedAtMs,
}),
ModelInvocationConflictError,
);
const ready = client
.prepare(
`SELECT status, version, attempt_count, result_code, error_summary
FROM "StepRuns" WHERE id = 'step-a'`,
)
.get();
assert.deepEqual(
{
status: ready.status,
version: ready.version,
attemptCount: ready.attempt_count,
resultCode: ready.result_code,
errorSummary: ready.error_summary,
},
{
status: 'ready',
version: 4,
attemptCount: 1,
resultCode: null,
errorSummary: null,
},
);
const resolutionEvent = client
.prepare(
`SELECT actor_type, actor_id, type FROM "RunEvents"
WHERE id = ?`,
)
.get(resolution.record.runEventId);
assert.deepEqual(
{ ...resolutionEvent },
{
actor_type: 'user',
actor_id: 'user-a',
type: 'step.ready',
},
);
const secondAdmissionAtMs = resolvedAtMs + 1;
assert.deepEqual(
await coordinator.record(
audit('admitted', {
requestId: 'request-b',
traceId: 'trace-b',
occurredAtMs: secondAdmissionAtMs,
deadlineAtMs: secondAdmissionAtMs + 10_000,
}),
),
{ status: 'created' },
);
assert.equal(
client
.prepare('SELECT count(*) AS count FROM "ModelInvocationStarts"')
.get().count,
2,
);
assert.equal(
client
.prepare('SELECT count(*) AS count FROM "ModelInvocationCompletions"')
.get().count,
1,
);
assert.equal(
client
.prepare('SELECT count(*) AS count FROM "ModelInvocationResolutions"')
.get().count,
1,
);
const running = client
.prepare(
`SELECT status, version, attempt_count FROM "StepRuns"
WHERE id = 'step-a'`,
)
.get();
assert.deepEqual(
{ ...running },
{
status: 'running',
version: 5,
attempt_count: 2,
},
);
assert.deepEqual(await repository.findCompletion('request-a'), completion);
assert.equal(
client.prepare('PRAGMA integrity_check').get().integrity_check,
'ok',
);
client.close();
});
for (const [decision, expectedStatus, expectedCode] of [
['fail', 'failed', 'model_outcome_rejected'],
['cancel', 'cancelled', 'model_outcome_cancelled'],
]) {
test(`SQLite manual ${decision} terminally resolves an unknown invocation`, async () => {
const { client, repository } = await fixture();
const coordinator = new DurableModelInvocationCoordinator(repository);
await coordinator.record(audit('admitted'));
await new DurableModelInvocationRecovery(repository, coordinator).recover(
1,
);
const completion = await repository.findCompletion('request-a');
const result = await new DurableModelInvocationResolutionCoordinator(
repository,
).resolve({
invocationId: 'request-a',
decision,
resolvedByUserId: 'user-a',
resolvedAtMs: completion.completedAtMs + 1,
});
assert.equal(result.status, 'created');
assert.deepEqual(
{
...client
.prepare(
`SELECT status, result_code FROM "StepRuns" WHERE id = 'step-a'`,
)
.get(),
},
{ status: expectedStatus, result_code: expectedCode },
);
client.close();
});
}