Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 16 additions & 7 deletions layers/registry/server/tasks/ai-generate-submit.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ import { observedSchedulePolicy } from '#shared/schedule-policy'
import { pAll } from '#shared/server/p-all'
import { classifyAbstractness } from '../utils/abstractness-effect'
import {
countMissingEmbeddings,
countMissingGeneratedSkills,
persistAbstractness,
runtimeGenerationLimits,
selectMissingBatchSkills,
Expand Down Expand Up @@ -118,14 +118,23 @@ export default defineScheduledTask({
})

async function runSubmit(db: D1Database, ai: AiBinding | undefined, vectorize: VectorizeBinding | undefined, apiKey: string | undefined) {
const runtimeLimits = runtimeGenerationLimits({
embedding: await countMissingEmbeddings(db),
})
// The counts answer "is there work" without the ORDER BY sort, so an idle
// hour skips the ordered catalog scans entirely instead of sorting the
// whole eligible catalog per kind before learning there was nothing to do.
const [missingEmbeddings, missingAbstractness] = await Promise.all([
countMissingGeneratedSkills(db, 'embedding'),
countMissingGeneratedSkills(db, 'abstractness'),
])
const runtimeLimits = runtimeGenerationLimits({ embedding: missingEmbeddings })
const [embeddingSkills, abstractnessSkills, batchSkills] = await Promise.all([
selectMissingGeneratedSkills(db, 'embedding', runtimeLimits.embedding),
selectMissingGeneratedSkills(db, 'abstractness', runtimeLimits.abstractness),
missingEmbeddings > 0
? selectMissingGeneratedSkills(db, 'embedding', runtimeLimits.embedding)
: Promise.resolve<GenerationSkill[]>([]),
missingAbstractness > 0
? selectMissingGeneratedSkills(db, 'abstractness', runtimeLimits.abstractness)
: Promise.resolve<GenerationSkill[]>([]),
HAIKU_GENERATION_PAUSED
? Promise.resolve([])
? Promise.resolve<GenerationSkill[]>([])
: selectMissingBatchSkills(db, BATCH_LIMIT),
])
const scanned = new Set(
Expand Down
29 changes: 22 additions & 7 deletions layers/registry/server/utils/ai-generation-work.ts
Original file line number Diff line number Diff line change
Expand Up @@ -131,24 +131,39 @@ const ELIGIBLE_SKILL_BASE_SQL = `
const ELIGIBLE_SKILL_SQL = `${ELIGIBLE_SKILL_BASE_SQL} AND ${NONEMPTY_SOURCE_SQL}`

/**
* Count the embedding backlog so the limit can respond to it. One extra query
* against a budget of 850 is a cheap price for not silently falling behind.
* Count the runtime backlog per kind. An idle hour skips the ordered catalog
* scans on the strength of this zero, and the embedding limit responds to its
* backlog. One or two extra queries against a budget of 850 is a cheap price
* for not silently falling behind.
*/
export async function countMissingEmbeddings(db: D1Database): Promise<number> {
const row = await db.prepare(
// CASE evaluates the source only when its embedding needs work. Trimming
// every completed source cost 2.2 seconds on the live catalog.
export async function countMissingGeneratedSkills(
db: D1Database,
kind: RuntimeGeneratedKind,
): Promise<number> {
const freshnessPredicate = kind === 'abstractness'
? `AND json_extract(generated.payload, '$.promptVersion') = ?2`
: ''
const statement = db.prepare(
// Mirrors selectMissingGeneratedSkills, minus the ORDER BY sort the skip
// exists to avoid. CASE evaluates the source only when its artifact needs
// work: trimming every completed source cost 2.2 seconds on the live
// catalog.
`SELECT COUNT(*) AS pending FROM (${ELIGIBLE_SKILL_BASE_SQL}
AND CASE WHEN NOT EXISTS (
SELECT 1
FROM skill_generated generated
WHERE generated.owner = s.owner
AND generated.repo = s.repo
AND generated.name = s.name
AND generated.kind = 'embedding'
AND generated.kind = ?1
AND generated.sha = s.current_sha
${freshnessPredicate}
) THEN ${NONEMPTY_SOURCE_SQL} ELSE 0 END
)`,
)
const row = await (kind === 'abstractness'
? statement.bind(kind, ABSTRACTNESS_PROMPT_VERSION)
: statement.bind(kind)
).first<{ pending: number }>()
return Number(row?.pending ?? 0)
}
Expand Down
35 changes: 28 additions & 7 deletions test/unit/ai-generation-work.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import {
ABSTRACTNESS_PROMPT_VERSION,
abstractnessResponseText,
buildAbstractnessUserPrompt,
countMissingEmbeddings,
countMissingGeneratedSkills,
parseAbstractnessPayload,
persistAbstractness,
runtimeGenerationLimits,
Expand Down Expand Up @@ -78,26 +78,47 @@ describe('ai generation work', () => {
sqlite.prepare('UPDATE skills SET rendered_raw = ?').run(source)
expect(await selectMissingGeneratedSkills(db, 'abstractness', 50)).toEqual([])
expect(await selectMissingGeneratedSkills(db, 'embedding', 50)).toEqual([])
expect(await countMissingEmbeddings(db)).toBe(0)
expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(0)
expect(await countMissingGeneratedSkills(db, 'abstractness')).toBe(0)
})

it('counts only eligible sources without a current embedding', async () => {
expect(await countMissingEmbeddings(db)).toBe(1)
expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(1)
sqlite.prepare('UPDATE skill_generated SET sha = \'old\' WHERE name = \'paused-only\' AND kind = \'embedding\'').run()
expect(await countMissingEmbeddings(db)).toBe(2)
expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(2)
sqlite.prepare('UPDATE skills SET rendered_status = \'failed\' WHERE name = \'paused-only\'').run()
expect(await countMissingEmbeddings(db)).toBe(1)
expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(1)
sqlite.prepare('UPDATE repos SET broken_since = 1').run()
expect(await countMissingEmbeddings(db)).toBe(0)
expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(0)
})

it('does not trim sources whose embedding is current', async () => {
const trim = vi.fn((source: string) => source.trim())
sqlite.function('trim', { varargs: true }, trim)
expect(await countMissingEmbeddings(db)).toBe(1)
expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(1)
expect(trim.mock.calls.map(([source]) => source)).toEqual(['# Embed'])
})

it('reports an idle catalog so the hourly submit can skip the ordered scans', async () => {
// Every marker row here is stale or missing, so the counts must report
// work: the abstractness markers carry an old prompt version even where
// the row exists.
expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(1)
expect(await countMissingGeneratedSkills(db, 'abstractness')).toBe(4)
sqlite.prepare(`INSERT INTO skill_generated VALUES
('acme', 'skills', 'needs-embedding', 'embedding', 'sha-embed', '{}', '2026-10-19T00:00:00.000Z')`).run()
sqlite.prepare(`INSERT INTO skill_generated VALUES
('acme', 'skills', 'needs-abstractness', 'abstractness', 'sha-abstract', '{}', '2026-10-19T00:00:00.000Z')`).run()
sqlite.prepare(`UPDATE skill_generated SET payload = json_set(payload, '$.promptVersion', ?)
WHERE kind = 'abstractness'`).run(ABSTRACTNESS_PROMPT_VERSION)

expect(await countMissingGeneratedSkills(db, 'embedding')).toBe(0)
expect(await countMissingGeneratedSkills(db, 'abstractness')).toBe(0)
// A zero count is only a safe skip signal when the ordered selects agree.
expect(await selectMissingGeneratedSkills(db, 'embedding', 50)).toEqual([])
expect(await selectMissingGeneratedSkills(db, 'abstractness', 50)).toEqual([])
})

it('does not select current runtime kinds because paused batch kinds are missing', async () => {
const embeddings = await selectMissingGeneratedSkills(db, 'embedding', 50)
const abstractness = await selectMissingGeneratedSkills(db, 'abstractness', 50)
Expand Down
Loading