@nocobase/app-plugin-agents 1.0.0-beta.15 → 1.0.0-beta.17

package/dist/server/core/runs/claim.js

@@ -1,3 +1,5 @@
+import { upgradeReservation } from './runner-affinity.js';
+import { ACTIVE_JOB, jobsRepo } from '../../jobs/job.store.js';
 /**
  * Claiming: handing queued runs to a runner that fits them, within each agent's concurrency limit, one run per
  * (agent, subject, thread) at a time, through the runners' one long-poll (`routes/runners/runner.ts`), which hands
@@ -6,7 +8,7 @@
  * team or belongs to the person who woke the agent, its features cover what the run requires (worked out from the
  * payload: checkouts, directories, skills, secrets), its owner's local policy (`Runner.policy`) lets it take the
  * agent's runs on the run's subject and repositories, and it has a free slot. A working directory that is a directory
- * on one runner pins the run to that runner, one run at a time per directory. A run's variables go to the runner that
+ * on one runner pins the run to that runner; multiple runs may share that directory. A run's variables go to the runner that
  * takes it, personal ones included, unless one of them is for team runners only (`teamRunnersOnly`): then only a team
  * runner takes the run, and the variables that asked for one are noted on it (`teamOnlyVariables`), so its wait says
  * which.
@@ -75,7 +77,7 @@ class AssemblyError extends Error {
         this.name = 'AssemblyError';
     }
 }
-/** The run does not fit this runner after all (its payload needs another runner, or a busy directory). */
+/** The run does not fit this runner after all (its payload needs another runner or unavailable capabilities). */
 class NotForThisRunner extends Error {
     constructor() {
         super('The run does not fit this runner.');
@@ -224,20 +226,6 @@ function directoryKeyOf(runnerId, dirs) {
         .map((dir) => `${runnerId}:${dir.path}`);
     return paths.length > 0 ? paths.join(DIRECTORY_KEY_SEPARATOR) : null;
 }
-/** Whether another run held by the runner works in any of `key`'s directories. */
-async function directoryBusy(conn, runnerId, runId, key) {
-    const wanted = new Set(key.split(DIRECTORY_KEY_SEPARATOR));
-    const held = await runsRepo(conn).findMany({
-        filter: (f) => f.and([
-            f.string('runnerId').eq(runnerId),
-            f.string('id').ne(runId),
-            f.or(ACTIVE.map((status) => f.string('status').eq(status))),
-        ]),
-    });
-    return held.some((run) => (run.directoryKey ?? '')
-        .split(DIRECTORY_KEY_SEPARATOR)
-        .some((part) => part !== '' && wanted.has(part)));
-}
 /** Resolve omitted subject input in the claim transaction, keeping it in the audit history. */
 async function settleClaimInputs(conn, pending, effective, now) {
     const ids = new Set(effective.map((input) => input.id));
@@ -270,10 +258,16 @@ export function createClaimService(deps) {
             (assembly.workspaceKey && !supportsWorkspaceKey(runner)))
             return 'incompatible';
         const { skills, variables } = await resources(agent, assembly);
+        return compatibleAssembly(run, agent, runner, assembly, skills.length, variables);
+    }
+    function compatibleAssembly(run, agent, runner, assembly, skillCount, variables) {
+        if (!isProtocolSupported(runner.protocolVersion) ||
+            (assembly.workspaceKey && !supportsWorkspaceKey(runner)))
+            return 'incompatible';
         const names = variables.names;
         if (!covers(runner.features, [
             ...run.requires,
-            ...payloadRequires(payloadDirs(assembly.dirs), skills.length, names.length, Boolean(assembly.workspaceKey)),
+            ...payloadRequires(payloadDirs(assembly.dirs), skillCount, names.length, Boolean(assembly.workspaceKey)),
         ]))
             return 'incompatible';
         if (!fitsActor(runner, agent, run.actorUserId, run.requires, clock.now().getTime()) ||
@@ -307,7 +301,8 @@ export function createClaimService(deps) {
         const entries = [];
         for (const entry of agent.modelEntries) {
             const single = { ...agent, modelEntries: [entry] };
-            const candidates = registered.filter((runner) => runner.status !== 'revoked');
+            const reserved = upgradeReservation(run.executions ?? [], registered, now.getTime());
+            const candidates = registered.filter((runner) => runner.status !== 'revoked' && (!reserved || runner.id === reserved));
             const connected = candidates.filter((runner) => (runner.status === 'online' ||
                 runner.status === 'upgrade_required') &&
                 runner.lastSeenAt !== null &&
@@ -321,7 +316,6 @@ export function createClaimService(deps) {
             let unknown = false;
             let incompatible = false;
             let ready = false;
-            let directoryBlocked = false;
             for (const runner of online) {
                 if (!isProtocolSupported(runner.protocolVersion) ||
                     runner.status === 'upgrade_required') {
@@ -377,16 +371,11 @@ export function createClaimService(deps) {
                 const held = await deps.slotsHeld(conn, runner.id);
                 const byTool = held.byTool ?? {};
                 const used = held.jobs + held.runs;
-                const directory = directoryKeyOf(runner.id, payloadDirs(assembly.dirs));
-                const blockedDirectory = directory !== null &&
-                    (await directoryBusy(conn, runner.id, run.id, directory));
-                directoryBlocked ||= blockedDirectory;
                 const hasRoom = !runner.draining &&
                     (runner.globalFree ?? 1) > 0 &&
                     used < runner.slots &&
                     (byTool[entry.tool] ?? 0) < toolLimit(runner, entry.tool) &&
-                    (runner.toolLoad?.[entry.tool]?.free ?? 1) > 0 &&
-                    !blockedDirectory;
+                    (runner.toolLoad?.[entry.tool]?.free ?? 1) > 0;
                 ready ||= hasRoom && runner.globalFree !== undefined;
                 if (hasRoom && runner.globalFree === undefined)
                     unknown = true;
@@ -398,10 +387,10 @@ export function createClaimService(deps) {
                             ? 'unknown'
                             : 'ready'
                         : runner.draining
-                            ? runner.draining
-                            : blockedDirectory
-                                ? 'restricted'
-                                : 'busy',
+                            ? runner.draining === 'diskLow'
+                                ? 'busy'
+                                : runner.draining
+                            : 'busy',
                     lastSeenAt: runner.lastSeenAt,
                     used,
                     limit: runner.slots,
@@ -426,8 +415,6 @@ export function createClaimService(deps) {
                             : connected.length
                                 ? 'toolNotInstalled'
                                 : 'offline';
-            if (directoryBlocked && !ready)
-                reason = 'restricted';
             if (!ready &&
                 compatible.length > 0 &&
                 loads.every((runner) => runner.reason === 'updatePending' || runner.reason === 'pruning'))
@@ -465,6 +452,126 @@ export function createClaimService(deps) {
             entries,
         };
     }
+    /** Metadata only: batch new scopes, reuse common scopes across machines and visited candidates. */
+    function admissionResources(conn, agents = []) {
+        const variableScopes = new Set();
+        const skillScopes = new Set();
+        const variables = [];
+        const skills = [];
+        function missing(targets, seen) {
+            const fresh = new Map();
+            for (const target of targets) {
+                const key = JSON.stringify([target.scope, target.scopeId]);
+                if (!seen.has(key))
+                    fresh.set(key, target);
+            }
+            for (const key of fresh.keys())
+                seen.add(key);
+            return [...fresh.values()];
+        }
+        return {
+            async load(checks) {
+                const vars = missing([
+                    ...checks.flatMap(({ agent, assembly }) => variableTargets(agent, assembly)),
+                    ...agents.map((agent) => ({ scope: 'agent', scopeId: agent.id })),
+                ], variableScopes);
+                const attached = missing(checks.flatMap(({ agent, assembly }) => skillTargets(agent, assembly)), skillScopes);
+                if (vars.length)
+                    variables.push(...(await deps.variables.admissionMetadata(conn, vars)));
+                if (attached.length)
+                    skills.push(...(await deps.skills.admissionMetadata(conn, attached)));
+            },
+            get(agent, assembly) {
+                const merged = new Map();
+                for (const target of variableTargets(agent, assembly))
+                    for (const row of variables)
+                        if (row.scope === target.scope && row.scopeId === target.scopeId)
+                            merged.set(row.name, row);
+                const targets = skillTargets(agent, assembly);
+                return {
+                    skills: skills.filter((row) => targets.some((target) => target.scope === row.scope && target.scopeId === row.scopeId)),
+                    variables: {
+                        names: [...merged.values()]
+                            .filter((row) => row.kind === 'value')
+                            .map((row) => row.name),
+                        teamOnly: [...merged.values()].filter((row) => row.kind === 'value' && row.teamRunnersOnly),
+                        passthrough: [...merged.values()]
+                            .filter((row) => row.kind === 'runner')
+                            .map((row) => row.name),
+                    },
+                };
+            },
+        };
+    }
+    /** Eligibility needs existence/compatibility, not occupancy, other models or a full wait projection. */
+    async function preferredModelAvailable(conn, run, agent, inputs, metadata, inventory) {
+        const now = clock.now().getTime();
+        const registered = await inventory();
+        const reserved = upgradeReservation(run.executions ?? [], registered, now);
+        const entry = agent.modelEntries[0];
+        const single = { ...agent, modelEntries: [entry] };
+        const binding = deps.subjects.get(run.subject.kind);
+        if (!binding)
+            throw new Error(`No context provider for ${run.subject.kind}.`);
+        const checks = [];
+        for (const runner of registered) {
+            if ((reserved && runner.id !== reserved) ||
+                runner.status !== 'online' ||
+                !runner.lastSeenAt ||
+                Date.parse(runner.lastSeenAt) + TIMINGS.offlineAfterMs <= now ||
+                !isProtocolSupported(runner.protocolVersion) ||
+                !runsEntry(runner, entry, now) ||
+                (runner.trust !== 'team' && runner.ownerUserId !== run.actorUserId))
+                continue;
+            const assembly = await binding.context
+                .assemble(conn, {
+                run,
+                agent,
+                runner,
+                inputs,
+                cli: deps.cli.name,
+                appName: deps.app.name,
+            })
+                .catch(() => null);
+            // The diagnostic path treats an unreadable context as unknown and waits for the preferred model.
+            if (!assembly)
+                return true;
+            checks.push({ agent: single, runner, assembly });
+        }
+        await metadata.load(checks);
+        for (const { runner, assembly } of checks) {
+            const { skills, variables } = metadata.get(single, assembly);
+            if (skills.some((skill) => skill.hash === null))
+                throw notFound('Skill version');
+            if (compatibleAssembly(run, single, runner, assembly, skills.length, variables) !== 'ready')
+                continue;
+            if (variables.passthrough.length && runner.variables == null)
+                return true;
+            if (variables.passthrough.some((name) => !runner.variables?.includes(name)))
+                continue;
+            return true;
+        }
+        return false;
+    }
+    /** The same ordered-model decision is used for dispatch and for immediate runner affinity. */
+    async function modelAdmission(conn, run, agent, entry, inputs, metadata = admissionResources(conn), inventory = () => deps.runners.all(conn)) {
+        if (agent.modelEntries.indexOf(entry) <= 0 || !deps.runners)
+            return undefined;
+        const policy = await deps.selectionPolicies?.forRun(conn, run);
+        if (!policy)
+            return undefined;
+        const elapsed = Math.max(0, (clock.now().getTime() - queueStart(run)) / 1000);
+        const preferredAvailable = await preferredModelAvailable(conn, run, agent, inputs, metadata, inventory);
+        if (preferredAvailable && elapsed < policy.fallbackAfterSeconds)
+            return false;
+        return {
+            reason: preferredAvailable ? 'preferredQueued' : 'preferredUnavailable',
+            preferredTool: agent.modelEntries[0].tool,
+            preferredModel: agent.modelEntries[0].model,
+            queueEnteredAt: new Date(queueStart(run)).toISOString(),
+            waitedSeconds: elapsed,
+        };
+    }
     async function assemble({ conn }, run, agent, runner, inputs, leaseExpiresAt, prepared, entry, resources = claimResources(conn, run.id)) {
         const binding = deps.subjects.get(run.subjectKind);
         if (!binding)
@@ -494,9 +601,6 @@ export function createClaimService(deps) {
             throw new NotForThisRunner();
         const dirs = payloadDirs(assembly.dirs);
         const directoryKey = directoryKeyOf(runner.id, dirs);
-        if (directoryKey !== null &&
-            (await directoryBusy(conn, runner.id, run.id, directoryKey)))
-            throw new NotForThisRunner();
         const { skills, variables } = await resources(agent, assembly);
         const names = variables.names;
         const requires = cleanList([
@@ -982,27 +1086,10 @@ export function createClaimService(deps) {
                     // Admission must precede preflight assembly: a missing or obsolete
                     // subject is cancelled normally, rather than counted as setup failure.
                     const resources = claimResources(conn, run.id);
-                    const selection = await selectionFacts(conn, toRun(current), agent, true, undefined, resources);
-                    const elapsed = Math.max(0, (now.getTime() - queueStart(toRun(current))) / 1000);
-                    const preferredAvailable = (selection?.entries[0]?.compatibleCount ?? 0) > 0 ||
-                        selection?.entries[0]?.reason === 'unknown' ||
-                        selection?.entries[0]?.blockers?.includes('unknown') === true;
-                    const chosenIndex = agent.modelEntries.indexOf(entry);
-                    if (selection && chosenIndex > 0) {
-                        // Throw to roll back the provisional dispatch; the original run keeps
-                        // waiting for its first choice without consuming an attempt or token.
-                        if (preferredAvailable && elapsed < selection.fallbackAfterSeconds)
-                            throw new NotForThisRunner();
-                        selectionSnapshot = {
-                            reason: preferredAvailable
-                                ? 'preferredQueued'
-                                : 'preferredUnavailable',
-                            preferredTool: agent.modelEntries[0].tool,
-                            preferredModel: agent.modelEntries[0].model,
-                            queueEnteredAt: selection.queueEnteredAt,
-                            waitedSeconds: elapsed,
-                        };
-                    }
+                    const admission = await modelAdmission(conn, toRun(current), agent, entry, pending.map(toInput));
+                    if (admission === false)
+                        throw new NotForThisRunner();
+                    selectionSnapshot = admission;
                     assembled = await assemble(unit, run, agent, runner, pending.map(toInput), leaseExpiresAt, prepared, entry, resources);
                 }
                 catch (error) {
@@ -1118,9 +1205,9 @@ export function createClaimService(deps) {
         });
     }
     /** Queued runs that may be claimed now, in claim order. */
-    async function candidates() {
+    async function candidates(runner) {
         const now = clock.now();
-        return runsRepo(tx.read()).findMany({
+        const ordinary = await runsRepo(tx.read()).findMany({
             filter: (f) => f.and([
                 f.string('status').eq('queued'),
                 f.or([
@@ -1135,15 +1222,198 @@ export function createClaimService(deps) {
             ],
             limit: CANDIDATES,
         });
+        if (ordinary.length === 0)
+            return [];
+        // A separate scan prevents ordinary candidate pages from hiding upgrade resumptions.
+        const recovery = [];
+        for (let offset = 0;; offset += CANDIDATES) {
+            const page = await runsRepo(tx.read()).findMany({
+                filter: (f) => f.and([
+                    f.string('status').eq('queued'),
+                    f.string('failureReason').eq('runnerUpgrade'),
+                    f
+                        .date('queueEnteredAt')
+                        .after(new Date(now.getTime() - 5 * 60_000)),
+                    f.or([
+                        f.date('availableAt').empty(),
+                        f.date('availableAt').notAfter(now),
+                    ]),
+                ]),
+                sort: (sort) => [
+                    sort.field('queueEnteredAt').asc(),
+                    sort.field('id').asc(),
+                ],
+                limit: CANDIDATES,
+                offset,
+            });
+            recovery.push(...page.filter((run) => toExecutions(run.executionHistory).at(-1)?.runnerId === runner.id));
+            if (page.length < CANDIDATES || recovery.length >= CANDIDATES)
+                break;
+        }
+        return [
+            ...new Map([...recovery, ...ordinary].map((run) => [run.id, run])).values(),
+        ];
+    }
+    /** One batch for locations, one for machines, and occupancy batches only when a preference exists. */
+    async function affinityBatch(queued, runner) {
+        const conn = tx.read();
+        if (queued.length === 0)
+            return {
+                machines: [runner],
+                preferred: new Map(),
+                held: [],
+                jobs: [],
+            };
+        const locations = await sessionsRepo(conn).groupBy({
+            by: ['runnerId', 'subjectKind', 'subjectId'],
+            aggregate: (a) => ({ updatedAt: a.max('updatedAt') }),
+            filter: (f) => f.or(queued.map((run) => f.and([
+                f.string('subjectKind').eq(run.subjectKind),
+                f.string('subjectId').eq(run.subjectId),
+            ]))),
+            sort: (sort) => sort.field('updatedAt').desc(),
+        });
+        const activeLocations = await runsRepo(conn).groupBy({
+            by: ['runnerId', 'subjectKind', 'subjectId'],
+            aggregate: (a) => ({ updatedAt: a.max('updatedAt') }),
+            filter: (f) => f.and([
+                f.string('runnerId').notEmpty(),
+                f.or(ACTIVE.map((status) => f.string('status').eq(status))),
+                f.or(queued.map((run) => f.and([
+                    f.string('subjectKind').eq(run.subjectKind),
+                    f.string('subjectId').eq(run.subjectId),
+                ]))),
+            ]),
+        });
+        const machineIds = [
+            ...new Set([
+                runner.id,
+                ...queued.map((run) => toExecutions(run.executionHistory).at(-1)?.runnerId),
+                ...locations.map((item) => item.runnerId),
+                ...activeLocations.map((item) => item.runnerId),
+            ].filter((id) => Boolean(id))),
+        ];
+        const machines = deps.runners
+            ? await deps.runners.all(conn, machineIds)
+            : [runner];
+        const preferred = new Map();
+        for (const run of queued) {
+            const last = toExecutions(run.executionHistory).at(-1)?.runnerId;
+            const current = activeLocations
+                .filter((item) => item.subjectKind === run.subjectKind &&
+                item.subjectId === run.subjectId)
+                .sort((a, b) => String(b.updatedAt).localeCompare(String(a.updatedAt)))[0]?.runnerId;
+            const location = locations.find((item) => item.subjectKind === run.subjectKind &&
+                item.subjectId === run.subjectId)?.runnerId;
+            const machine = machines.find((item) => item.id === (last ?? current ?? location));
+            if (machine)
+                preferred.set(run.id, machine);
+        }
+        const ids = [...new Set([...preferred.values()].map((item) => item.id))];
+        const held = ids.length === 0
+            ? []
+            : await runsRepo(conn).groupBy({
+                by: ['runnerId', 'tool'],
+                aggregate: (a) => ({ count: a.count() }),
+                filter: (f) => f.and([
+                    f.or(ids.map((id) => f.string('runnerId').eq(id))),
+                    f.or(ACTIVE.map((status) => f.string('status').eq(status))),
+                ]),
+            });
+        const jobs = ids.length === 0
+            ? []
+            : await jobsRepo(conn).groupBy({
+                by: ['runnerId'],
+                aggregate: (a) => ({ count: a.count() }),
+                filter: (f) => f.and([
+                    f.or(ids.map((id) => f.string('runnerId').eq(id))),
+                    f.or(ACTIVE_JOB.map((status) => f.string('status').eq(status))),
+                ]),
+            });
+        return { machines, preferred, held, jobs };
+    }
+    function hasImmediateRoom(run, agent, preferred, batch) {
+        const now = clock.now().getTime();
+        if (preferred.status !== 'online' ||
+            !preferred.lastSeenAt ||
+            Date.parse(preferred.lastSeenAt) + TIMINGS.offlineAfterMs <= now ||
+            preferred.draining ||
+            (preferred.globalFree ?? 1) <= 0 ||
+            !fits(preferred, agent, run, now))
+            return false;
+        const held = batch.held.filter((item) => item.runnerId === preferred.id);
+        if (held.reduce((count, item) => count + Number(item.count), 0) +
+            batch.jobs
+                .filter((item) => item.runnerId === preferred.id)
+                .reduce((count, item) => count + Number(item.count), 0) >=
+            preferred.slots)
+            return false;
+        if (!pickOpenEntry(preferred, agent, (tool) => held
+            .filter((item) => item.tool === tool)
+            .reduce((count, item) => count + Number(item.count), 0) <
+            toolLimit(preferred, tool) &&
+            (preferred.toolLoad?.[tool]?.free ?? 1) > 0, now))
+            return false;
+        return true;
+    }
+    /** Evaluate only the candidate reached in queue order, reusing this poll's batched capacity and metadata. */
+    async function readyPreference(run, agent, batch, claimant, metadata, inventory) {
+        const conn = tx.read();
+        const runner = batch.preferred.get(run.id);
+        if (!runner ||
+            runner.id === claimant.id ||
+            !fits(claimant, agent, run, clock.now().getTime()) ||
+            upgradeReservation(toExecutions(run.executionHistory), batch.machines, clock.now().getTime()) ||
+            !hasImmediateRoom(run, agent, runner, batch))
+            return false;
+        const binding = deps.subjects.get(run.subjectKind);
+        if (!binding || !(await commentConsentFor(conn, runner, run)).allowed)
+            return false;
+        const assembly = await binding.context
+            .assemble(conn, {
+            run: toRun(run),
+            agent,
+            runner,
+            inputs: [],
+            cli: deps.cli.name,
+            appName: deps.app.name,
+        })
+            .catch(() => null);
+        if (!assembly)
+            return false;
+        await metadata.load([{ agent, assembly }]);
+        const { skills, variables } = metadata.get(agent, assembly);
+        if (!variables.passthrough.every((name) => runner.variables?.includes(name)) ||
+            skills.some((row) => row.hash === null) ||
+            compatibleAssembly(toRun(run), agent, runner, assembly, skills.length, variables) !== 'ready')
+            return false;
+        const entry = pickOpenEntry(runner, agent, (tool) => batch.held
+            .filter((item) => item.runnerId === runner.id && item.tool === tool)
+            .reduce((count, item) => count + Number(item.count), 0) <
+            toolLimit(runner, tool) && (runner.toolLoad?.[tool]?.free ?? 1) > 0, clock.now().getTime());
+        if (!entry)
+            return false;
+        if (agent.modelEntries.indexOf(entry) > 0) {
+            const inputs = (await pendingInputs(conn, run.id)).map(toInput);
+            if ((await modelAdmission(conn, toRun(run), agent, entry, inputs, metadata, inventory)) === false)
+                return false;
+        }
+        return true;
     }
     async function claimOne(runner, seen, reported) {
         if (runner.draining)
             return null;
         const conn = tx.read();
-        const queued = await candidates();
+        const queued = await candidates(runner);
+        const affinity = await affinityBatch(queued, runner);
         const agents = await findAgents(conn, queued
             .filter((candidate) => !seen.has(candidate.id))
             .map((candidate) => candidate.agentId));
+        const metadata = admissionResources(conn, [...agents.values()]);
+        let inventoryRead;
+        const inventory = () => (inventoryRead ??= deps.runners
+            ? deps.runners.all(conn)
+            : Promise.resolve([runner]));
         const full = new Set();
         for (const candidate of queued) {
             if (seen.has(candidate.id) || full.has(candidate.agentId))
@@ -1154,6 +1424,15 @@ export function createClaimService(deps) {
                 (!deps.subjects.get(candidate.subjectKind)?.actsAs &&
                     !fits(runner, agent, candidate, clock.now().getTime())))
                 continue;
+            const reserved = upgradeReservation(toExecutions(candidate.executionHistory), affinity.machines, clock.now().getTime());
+            if (reserved && reserved !== runner.id)
+                continue;
+            const preferred = affinity.preferred.get(candidate.id);
+            if (!reserved &&
+                preferred &&
+                preferred.id !== runner.id &&
+                (await readyPreference(candidate, agent, affinity, runner, metadata, inventory)))
+                continue;
             // None of its tools has room on the runner's side: the runs behind it may still fit.
             if (!pickOpenEntry(runner, agent, (tool) => reported.has(tool), clock.now().getTime()))
                 continue;