diff --git a/trios/agent-server/apps/server/src/api/routes/queen-contributor-keys.ts b/trios/agent-server/apps/server/src/api/routes/queen-contributor-keys.ts index 040dae4144..fe6d3b49bd 100644 --- a/trios/agent-server/apps/server/src/api/routes/queen-contributor-keys.ts +++ b/trios/agent-server/apps/server/src/api/routes/queen-contributor-keys.ts @@ -9,6 +9,10 @@ import { contributorGithub, type EnvironmentKey, listContributorKeys, + openCredential, + ownerModel, + providerModels, + setContributorModel, trustedContributor, } from '../services/queen-contributor-keys' import { CONTRIBUTOR_POLICY } from '../services/queen-contributor-policy' @@ -104,14 +108,16 @@ export function createQueenContributorKeysRoute( c.set('body', Buffer.concat(chunks).toString('utf8')) return await next() }) - app.get('/', async (c) => { - const pool = (deps.pool ?? poolForKeys)() - const subject = c.get('contributor') - const environment = (deps.environment ?? environmentContributorKeys)() - const owners = ( - deps.owners ?? (() => parseOwners(process.env.TRIOS_KEY_OWNERS)) - )() - const keys = await listContributorKeys(pool, subject, environment, owners) + const environmentOf = () => (deps.environment ?? environmentContributorKeys)() + const ownersOf = () => + (deps.owners ?? (() => parseOwners(process.env.TRIOS_KEY_OWNERS)))() + const account = async (pool: Pool, subject: string) => { + const keys = await listContributorKeys( + pool, + subject, + environmentOf(), + ownersOf(), + ) const work = await keyWork(pool) const byKey = new Map(work.map((item) => [item.keyIndex, item])) const own = work.filter((item) => @@ -131,7 +137,7 @@ export function createQueenContributorKeysRoute( hours: entry.hours, } : emptyContribution() - return c.json({ + return { keys: keys.map((key) => { const entry = byKey.get(key.id) return { @@ -139,16 +145,79 @@ export function createQueenContributorKeysRoute( contribution: contribution(rank(entry ? [entry] : [], {})[0]), } }), + // `model` is what this owner's keys of the provider run on now; + // `defaultModel` is what a provider starts on before any choice. providers: CONTRIBUTOR_PROVIDERS.map(({ id, label, model }) => ({ id, label, - model, + model: ownerModel(keys, id), + defaultModel: model, })), contribution: contribution(totals), attribution: { subject, github: contributorGithub(subject) }, scoring: { acceptedXp: ACCEPTED_XP, specXp: SPEC_XP, hourXp: HOUR_XP }, + } + } + app.get('/', async (c) => { + const pool = (deps.pool ?? poolForKeys)() + return c.json(await account(pool, c.get('contributor'))) + }) + app.get('/models/:provider', async (c) => { + const provider = CONTRIBUTOR_PROVIDERS.find( + (p) => p.id === c.req.param('provider'), + ) + if (!provider) throw new ContributorError('unsupported_provider') + const pool = (deps.pool ?? poolForKeys)() + const subject = c.get('contributor') + const environment = environmentOf() + const keys = ( + await listContributorKeys(pool, subject, environment, ownersOf()) + ).filter((key) => key.provider === provider.id) + if (!keys.length) throw new ContributorError('key_not_found', 404) + // The catalog is read with one of the owner's own keys, never an operator's. + let apiKey = environment.find((entry) => + keys.some((key) => key.source === 'environment' && key.id === entry.id), + )?.apiKey + if (!apiKey) { + const { rows } = await pool.query<{ + sealed: string + fingerprint: string + }>( + "SELECT sealed,fingerprint FROM queen_contributor_keys WHERE owner_subject=$1 AND provider=$2 AND source='managed' AND sealed IS NOT NULL LIMIT 1", + [subject, provider.id], + ) + if (rows[0]) + apiKey = openCredential(rows[0].sealed, subject, rows[0].fingerprint) + } + if (!apiKey) throw new ContributorError('key_not_found', 404) + return c.json({ + provider: provider.id, + model: ownerModel(keys, provider.id), + models: await providerModels( + { provider: provider.id, apiKey }, + deps.fetcher, + ), }) }) + app.post('/model', async (c) => { + let body: unknown + try { + body = JSON.parse(c.get('body')) + } catch { + throw new ContributorError('invalid_json') + } + const pool = (deps.pool ?? poolForKeys)() + const subject = c.get('contributor') + await setContributorModel( + pool, + subject, + body, + environmentOf(), + ownersOf(), + deps.fetcher, + ) + return c.json(await account(pool, subject)) + }) app.post('/', async (c) => { let body: unknown try { diff --git a/trios/agent-server/apps/server/src/api/services/queen-contributor-keys.ts b/trios/agent-server/apps/server/src/api/services/queen-contributor-keys.ts index f320608b50..d809f32257 100644 --- a/trios/agent-server/apps/server/src/api/services/queen-contributor-keys.ts +++ b/trios/agent-server/apps/server/src/api/services/queen-contributor-keys.ts @@ -48,6 +48,8 @@ interface KeyRow { owner_name: string provider: ContributorProvider model: string + /** True once the owner chose the model; an environment row then overrides its pool's model. */ + model_chosen?: boolean label: string sealed: string | null enabled: boolean @@ -74,6 +76,55 @@ export const CONTRIBUTOR_PROVIDERS = [ }, { id: 'zai' as const, label: 'Z.ai', model: P.ZAI_MODEL, baseUrl: P.ZAI_URL }, ] +/** Provider model ids as both catalogs print them: `z-ai/glm-5.3`, `glm-4.5-flash`. */ +const MODEL_ID = /^[A-Za-z0-9][A-Za-z0-9._:/-]*$/ +export function validModel(value: unknown): value is string { + return ( + typeof value === 'string' && + value.length > 0 && + value.length <= P.MODEL_LIMIT && + MODEL_ID.test(value) + ) +} +/** + * Keys are named ` #`. An environment key's number is its place + * in its pool, so pool 2's first Z.ai key is `zai #1`, not `zai #10001`; the + * stride is POOL_KEY_STRIDE in queen-dispatch.ts (a test holds them equal). + * A key added without a name, or named only after its provider, takes the + * next free number of that provider. + */ +export const LABEL_POOL_STRIDE = 10_000 +export function environmentLabel( + provider: ContributorProvider, + id: number, +): string { + return `${provider} #${(id % LABEL_POOL_STRIDE) + 1}` +} +function labelNumber(provider: ContributorProvider, label: string): number { + const match = /^([a-z]+) #([1-9][0-9]{0,8})$/.exec(label.trim()) + return match && match[1] === provider ? Number(match[2]) : 0 +} +export function bareLabel( + provider: ContributorProvider, + label: string, +): boolean { + const name = label.trim().toLowerCase() + const known = CONTRIBUTOR_PROVIDERS.find((p) => p.id === provider) + return name === '' || name === provider || name === known?.label.toLowerCase() +} +/** The next free ` #` after every numbered key and pool place. */ +export function nextLabel( + provider: ContributorProvider, + labels: string[], + environment: EnvironmentKey[], +): string { + let top = 0 + for (const key of environment) + if (key.provider === provider) + top = Math.max(top, (key.id % LABEL_POOL_STRIDE) + 1) + for (const label of labels) top = Math.max(top, labelNumber(provider, label)) + return `${provider} #${top + 1}` +} const ready = new WeakMap>() export function contributorsEnabled(): boolean { return ( @@ -204,6 +255,8 @@ export async function ensureContributorKeys(pool: Pool): Promise { CHECK ((source='managed' AND key_index<0 AND sealed IS NOT NULL) OR (source='environment' AND key_index>=0 AND sealed IS NULL)) ); + ALTER TABLE queen_contributor_keys + ADD COLUMN IF NOT EXISTS model_chosen boolean NOT NULL DEFAULT false; `) .then(() => undefined) .catch((error) => { @@ -277,13 +330,19 @@ export async function listContributorKeys( ) out.push( stored - ? publicKey(stored) + ? { + ...publicKey(stored), + label: environmentLabel(key.provider, key.id), + // A binding row records the model it was bound under; only an + // owner's explicit choice outlives a change of the pool's model. + model: stored.model_chosen ? stored.model : key.model, + } : { id: key.id, source: 'environment', provider: key.provider, model: key.model, - label: `${key.provider} #${key.id + 1}`, + label: environmentLabel(key.provider, key.id), fingerprint: digest.slice(0, 12), enabled: true, createdAt: null, @@ -295,9 +354,40 @@ export async function listContributorKeys( }, ) } + await numberBareLabels(pool, subject, out, environment) return out } +/** + * A managed key saved before keys were numbered (or named only "nvidia") + * takes the next number of its provider, oldest first, once. The update is + * compare-and-set on the old label, so two readers cannot number it twice. + */ +async function numberBareLabels( + pool: Pool, + subject: string, + keys: ContributorKey[], + environment: EnvironmentKey[], +): Promise { + const bare = keys + .filter( + (key) => key.source === 'managed' && bareLabel(key.provider, key.label), + ) + .sort((a, b) => b.id - a.id) + for (const key of bare) { + const label = nextLabel( + key.provider, + keys.filter((other) => other !== key).map((other) => other.label), + environment, + ) + const { rowCount } = await pool.query( + 'UPDATE queen_contributor_keys SET label=$3 WHERE key_index=$1 AND owner_subject=$2 AND label=$4', + [key.id, subject, label, key.label], + ) + if (rowCount) key.label = label + } +} + /** No response bodies, redirects, arbitrary URLs, retries or provider messages escape. */ let probesInFlight = 0 export async function probeCredential( @@ -358,6 +448,136 @@ export async function probeCredential( } } +/** + * Whether a model can carry a bee: one request that hands it a tool. + * + * A model that answers text and never calls a tool would accept every bee and + * finish none of them, so plain "Reply OK." is not enough evidence to switch a + * whole pool. 404/400 name a model the endpoint does not serve; a refused or + * busy key says nothing about the model and lets the caller try another key. + */ +export type ModelCheck = + | 'ok' + | 'unknown_model' + | 'no_tool_call' + | 'inconclusive' +export async function checkModel( + key: Pick, + fetcher: typeof fetch = fetch, +): Promise { + const provider = CONTRIBUTOR_PROVIDERS.find((p) => p.id === key.provider) + if (!provider) throw new ContributorError('unsupported_provider') + if (probesInFlight >= P.MAX_PROBES_IN_FLIGHT) + throw new ContributorError('probe_rate_limited', 429) + probesInFlight++ + let result: ModelCheck = 'inconclusive' + try { + const response = await fetcher(`${provider.baseUrl}/chat/completions`, { + method: 'POST', + redirect: 'error', + signal: AbortSignal.timeout(P.PROBE_TIMEOUT_MS), + headers: { + Authorization: `Bearer ${key.apiKey}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + model: key.model, + messages: [{ role: 'user', content: 'Call the ping tool.' }], + tools: [ + { + type: 'function', + function: { + name: 'ping', + description: 'Answer a liveness check.', + parameters: { type: 'object', properties: {} }, + }, + }, + ], + tool_choice: 'auto', + max_tokens: P.MODEL_PROBE_MAX_TOKENS, + stream: false, + }), + }) + if ([400, 404, 410, 422].includes(response.status)) result = 'unknown_model' + else if (response.ok) { + const body = (await response.json()) as { + choices?: Array<{ message?: { tool_calls?: unknown } }> + } + const calls = body.choices?.[0]?.message?.tool_calls + result = Array.isArray(calls) && calls.length > 0 ? 'ok' : 'no_tool_call' + } + await response.body?.cancel().catch(() => {}) + } catch { + /* Timeouts and network errors say nothing about the model. */ + } finally { + probesInFlight-- + } + return result +} + +/** Model ids a provider lists, per provider, read with an owner's key. */ +const modelLists = new Map< + ContributorProvider, + { at: number; models: string[] } +>() +export async function providerModels( + key: Pick, + fetcher: typeof fetch = fetch, +): Promise { + const provider = CONTRIBUTOR_PROVIDERS.find((p) => p.id === key.provider) + if (!provider) throw new ContributorError('unsupported_provider') + const cached = modelLists.get(provider.id) + if (cached && Date.now() - cached.at < P.MODEL_LIST_CACHE_SECONDS * 1000) + return cached.models + let models: string[] = [] + try { + const response = await fetcher(`${provider.baseUrl}/models`, { + redirect: 'error', + signal: AbortSignal.timeout(P.PROBE_TIMEOUT_MS), + headers: { Authorization: `Bearer ${key.apiKey}` }, + }) + if (response.ok) { + const body = (await response.json()) as { data?: Array<{ id?: unknown }> } + models = [ + ...new Set( + (Array.isArray(body.data) ? body.data : []) + .map((entry) => entry?.id) + .filter(validModel), + ), + ] + .sort() + .slice(0, P.MODEL_LIST_LIMIT) + } else await response.body?.cancel().catch(() => {}) + } catch { + /* An unreadable catalog is an empty suggestion list, never an error body. */ + } + // Only a real list is cached: one failed read must not hide the catalog. + if (models.length) modelLists.set(provider.id, { at: Date.now(), models }) + return models +} + +/** Tests only. */ +export function resetProviderModels(): void { + modelLists.clear() +} + +/** The model most of the owner's keys of a provider use, else the default. */ +export function ownerModel( + keys: Pick[], + provider: ContributorProvider, +): string { + const counts = new Map() + for (const key of keys) + if (key.provider === provider) + counts.set(key.model, (counts.get(key.model) ?? 0) + 1) + let best: string | undefined + for (const [model, count] of counts) + if (!best || count > (counts.get(best) ?? 0)) best = model + return ( + best ?? CONTRIBUTOR_PROVIDERS.find((p) => p.id === provider)?.model ?? '' + ) +} + export async function addContributorKey( pool: Pool, subject: string, @@ -387,11 +607,10 @@ export async function addContributorKey( throw new ContributorError('invalid_key') if (value.label !== undefined && typeof value.label !== 'string') throw new ContributorError('invalid_label') - const label = - typeof value.label === 'string' ? value.label.trim() : provider.label + const named = typeof value.label === 'string' ? value.label.trim() : '' if ( - label.length > P.LABEL_LIMIT || - [...label].some( + named.length > P.LABEL_LIMIT || + [...named].some( (char) => char.charCodeAt(0) < 32 || char.charCodeAt(0) === 127, ) ) @@ -410,8 +629,16 @@ export async function addContributorKey( throw new ContributorError('key_already_connected', 409) return publicKey(duplicate.rows[0]) } + const chosen = await pool.query<{ model: string }>( + `SELECT model FROM queen_contributor_keys + WHERE owner_subject=$1 AND provider=$2 AND model_chosen + GROUP BY model ORDER BY count(*) DESC, model LIMIT 1`, + [subject, provider.id], + ) + // A new key joins the model its owner already switched this provider to. + const model = chosen.rows[0]?.model ?? provider.model const checked = await probeCredential( - { provider: provider.id, model: provider.model, apiKey: secret }, + { provider: provider.id, model, apiKey: secret }, fetcher, ) const client = await pool.connect() @@ -426,22 +653,36 @@ export async function addContributorKey( ) if (Number(count.rows[0]?.n) >= P.MAX_KEYS_PER_OWNER) throw new ContributorError('key_limit', 409) + // Numbered under the same lock as the count, so two adds cannot share one. + let label = named + if (bareLabel(provider.id, named)) { + const taken = await client.query<{ label: string }>( + "SELECT label FROM queen_contributor_keys WHERE owner_subject=$1 AND provider=$2 AND source='managed'", + [subject, provider.id], + ) + label = nextLabel( + provider.id, + taken.rows.map((row) => row.label), + environment, + ) + } const { rows } = await client.query( `INSERT INTO queen_contributor_keys - (source,fingerprint,owner_subject,owner_name,provider,model,label,sealed,enabled,probe_status,probe_at,probe_ms) - VALUES ('managed',$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) RETURNING *`, + (source,fingerprint,owner_subject,owner_name,provider,model,label,sealed,enabled,probe_status,probe_at,probe_ms,model_chosen) + VALUES ('managed',$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) RETURNING *`, [ digest, subject, contributorName(subject), provider.id, - provider.model, + model, label, sealed, checked.status === 'ok', checked.status, checked.checkedAt, checked.latencyMs, + !!chosen.rows[0], ], ) await client.query('COMMIT') @@ -494,7 +735,7 @@ async function ownedRow( contributorName(subject), current.provider, current.model, - `${current.provider} #${id + 1}`, + environmentLabel(current.provider, id), ], ) const { rows } = await pool.query( @@ -536,12 +777,24 @@ export async function changeContributorKey( owners, action !== 'disable', ) + // An environment key runs on its pool's model until its owner chooses one. + const pooled = environment.find((key) => key.id === id) + const effective = (stored: KeyRow): ContributorKey => ({ + ...publicKey(stored), + ...(stored.source === 'environment' && pooled + ? { label: environmentLabel(pooled.provider, id) } + : {}), + model: + stored.source === 'environment' && pooled && !stored.model_chosen + ? pooled.model + : stored.model, + }) if (action === 'disable') { const { rows } = await pool.query( 'UPDATE queen_contributor_keys SET enabled=false,revision=revision+1 WHERE key_index=$1 AND owner_subject=$2 RETURNING *', [id, subject], ) - return publicKey(rows[0]) + return effective(rows[0]) } // Database claim, so multiple tabs/processes cannot fan out probes on one key. const claimed = await pool.query( @@ -554,7 +807,7 @@ export async function changeContributorKey( if (!claimed.rows.length) throw new ContributorError('probe_rate_limited', 429) const checked = await probeCredential( - { provider: row.provider, model: row.model, apiKey: secret }, + { provider: row.provider, model: effective(row).model, apiKey: secret }, fetcher, ) const { rows } = await pool.query( @@ -572,12 +825,104 @@ export async function changeContributorKey( claimed.rows[0].revision, ], ) - return publicKey(rows[0]) + return effective(rows[0]) +} + +/** + * One model for every key the owner holds of one provider. + * + * Switching a pool is all-or-nothing for that owner, so the evidence comes + * first: a key of the provider must make the model call a tool. Only then are + * the owner's environment keys bound (the same binding a probe creates) and + * every owned row of the provider set to the model. Other owners' keys and the + * operator's pool variables are untouched; an environment row overrides its + * pool's model only while `model_chosen` is set. + */ +export async function setContributorModel( + pool: Pool, + subject: string, + input: unknown, + environment: EnvironmentKey[], + owners: Record, + fetcher: typeof fetch = fetch, +): Promise<{ provider: ContributorProvider; model: string; keys: number }> { + if (!input || typeof input !== 'object' || Array.isArray(input)) + throw new ContributorError('invalid_model') + const value = input as Record + if (Object.keys(value).some((key) => !['provider', 'model'].includes(key))) + throw new ContributorError('invalid_model') + const provider = CONTRIBUTOR_PROVIDERS.find((p) => p.id === value.provider) + if (!provider) throw new ContributorError('unsupported_provider') + if (!validModel(value.model)) throw new ContributorError('invalid_model') + const model = value.model + const keys = ( + await listContributorKeys(pool, subject, environment, owners) + ).filter((key) => key.provider === provider.id) + if (!keys.length) throw new ContributorError('key_not_found', 404) + const sealed = await pool.query< + Pick + >( + "SELECT key_index,sealed,fingerprint FROM queen_contributor_keys WHERE owner_subject=$1 AND provider=$2 AND source='managed'", + [subject, provider.id], + ) + const secretOf = (key: ContributorKey): string | undefined => { + if (key.source === 'environment') + return environment.find((entry) => entry.id === key.id)?.apiKey + const row = sealed.rows.find((entry) => entry.key_index === key.id) + if (!row?.sealed) return undefined + try { + return openCredential(row.sealed, subject, row.fingerprint) + } catch { + return undefined + } + } + // Keys that answered last time first: a refused key proves nothing either way. + const order = (key: ContributorKey) => + (key.enabled ? 0 : 2) + (key.lastProbe.status === 'ok' ? 0 : 1) + let attempts = 0 + let verdict: ModelCheck = 'inconclusive' + for (const key of [...keys].sort((a, b) => order(a) - order(b))) { + if (attempts >= P.MODEL_CHECK_ATTEMPTS) break + const secret = secretOf(key) + if (!secret) continue + attempts++ + verdict = await checkModel( + { provider: provider.id, model, apiKey: secret }, + fetcher, + ) + if (verdict !== 'inconclusive') break + } + if (verdict === 'unknown_model') + throw new ContributorError('model_unavailable', 422) + if (verdict === 'no_tool_call') + throw new ContributorError('model_without_tools', 422) + if (verdict !== 'ok') throw new ContributorError('model_check_failed', 504) + for (const key of keys) { + if (key.source !== 'environment') continue + // A fingerprint already held by the owner's managed copy cannot be bound + // twice; that copy carries the choice for the environment key instead. + await ownedRow(pool, subject, key.id, environment, owners, false).catch( + (error) => { + if ( + !(error instanceof ContributorError) || + error.code !== 'key_binding_conflict' + ) + throw error + }, + ) + } + const updated = await pool.query( + 'UPDATE queen_contributor_keys SET model=$3,model_chosen=true WHERE owner_subject=$1 AND provider=$2', + [subject, provider.id, model], + ) + return { provider: provider.id, model, keys: updated.rowCount ?? 0 } } export interface ContributorRuntime { managed: EnvironmentKey[] disabled: number[] + /** Environment index -> the model its owner chose instead of the pool's. */ + models?: Record } async function contributorRegistryAvailable(pool: Pool): Promise { if (contributorsEnabled()) { @@ -591,12 +936,43 @@ async function contributorRegistryAvailable(pool: Pool): Promise { ) return !!result.rows[0]?.registry } +/** + * Environment index -> the model its owner chose, for keys whose current + * secret matches the row. An owner's managed copy of a secret the environment + * also holds carries the choice too: the environment key does the work. + */ +function chosenEnvironmentModels( + rows: KeyRow[], + environment: EnvironmentKey[], +): Record { + const pooled = new Map( + environment.map((key) => [fingerprint(key.apiKey), key]), + ) + const models: Record = {} + const choose = (row: KeyRow, current: EnvironmentKey | undefined) => { + if ( + current && + row.model_chosen && + !(current.id in models) && + row.model !== current.model + ) + models[current.id] = row.model + } + for (const row of rows) + if (row.source === 'environment' && row.enabled) { + const current = pooled.get(row.fingerprint) + choose(row, current?.id === row.key_index ? current : undefined) + } + for (const row of rows) + if (row.source === 'managed') choose(row, pooled.get(row.fingerprint)) + return models +} export async function contributorRuntime( pool: Pool, environment: EnvironmentKey[], ): Promise { if (!(await contributorRegistryAvailable(pool))) - return { managed: [], disabled: [] } + return { managed: [], disabled: [], models: {} } const { rows } = await pool.query( 'SELECT * FROM queen_contributor_keys', ) @@ -643,7 +1019,11 @@ export async function contributorRuntime( digests.add(row.fingerprint) } } - return { managed, disabled } + return { + managed, + disabled, + models: chosenEnvironmentModels(rows, environment), + } } export async function contributorOwnerNames( pool: Pool, diff --git a/trios/agent-server/apps/server/src/api/services/queen-contributor-policy.gen.ts b/trios/agent-server/apps/server/src/api/services/queen-contributor-policy.gen.ts index 6e522e1f02..7b7385ba26 100644 --- a/trios/agent-server/apps/server/src/api/services/queen-contributor-policy.gen.ts +++ b/trios/agent-server/apps/server/src/api/services/queen-contributor-policy.gen.ts @@ -6,8 +6,8 @@ export const KIND = "automation" satisfies string; export const ID = "queen-contributor-keys" satisfies string; export const REPO = "BrowserOS" satisfies string; -export const VERSION = 1 satisfies number; -export const PROBE_TIMEOUT_MS = 15000 satisfies number; +export const VERSION = 2 satisfies number; +export const PROBE_TIMEOUT_MS = 90000 satisfies number; export const PROBE_MAX_TOKENS = 16 satisfies number; export const MAX_PROBES_IN_FLIGHT = 4 satisfies number; export const PROBE_COOLDOWN_SECONDS = 60 satisfies number; @@ -23,6 +23,11 @@ export const NVIDIA_URL = "https://integrate.api.nvidia.com/v1" satisfies string export const NVIDIA_MODEL = "nvidia/nemotron-3-ultra-550b-a55b" satisfies string; export const ZAI_URL = "https://api.z.ai/api/paas/v4" satisfies string; export const ZAI_MODEL = "glm-4.5-flash" satisfies string; +export const MODEL_LIMIT = 128 satisfies number; +export const MODEL_PROBE_MAX_TOKENS = 512 satisfies number; +export const MODEL_CHECK_ATTEMPTS = 2 satisfies number; +export const MODEL_LIST_LIMIT = 500 satisfies number; +export const MODEL_LIST_CACHE_SECONDS = 600 satisfies number; // t27c gen-ts: fn trusted_proxy was not emitted -- this backend lowers declarations, not bodies. // t27c gen-ts: fn source_matches_index was not emitted -- this backend lowers declarations, not bodies. // t27c gen-ts: fn may_add_key was not emitted -- this backend lowers declarations, not bodies. @@ -31,6 +36,9 @@ export const ZAI_MODEL = "glm-4.5-flash" satisfies string; // t27c gen-ts: fn credential_fits was not emitted -- this backend lowers declarations, not bodies. // t27c gen-ts: fn same_binding was not emitted -- this backend lowers declarations, not bodies. // t27c gen-ts: fn enabled_after_probe was not emitted -- this backend lowers declarations, not bodies. +// t27c gen-ts: fn model_fits was not emitted -- this backend lowers declarations, not bodies. +// t27c gen-ts: fn may_switch_model was not emitted -- this backend lowers declarations, not bodies. +// t27c gen-ts: fn may_try_model_check was not emitted -- this backend lowers declarations, not bodies. // t27c gen-ts: fn wallet_credit was not emitted -- this backend lowers declarations, not bodies. // t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. // t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. @@ -41,10 +49,14 @@ export const ZAI_MODEL = "glm-4.5-flash" satisfies string; // t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. // t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. // t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. +// t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. +// t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. +// t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. +// t27c gen-ts: a TestBlock was not emitted -- it is checked by the compiler, not by the artifact. // Declaration order, which the spec's own laws depend on. export const __STRUCT_ORDER__ = [] as const; -export const __DECL_ORDER__ = ["KIND", "ID", "REPO", "VERSION", "PROBE_TIMEOUT_MS", "PROBE_MAX_TOKENS", "MAX_PROBES_IN_FLIGHT", "PROBE_COOLDOWN_SECONDS", "MAX_KEYS_PER_OWNER", "MAX_KEY_BYTES", "MAX_BODY_BYTES", "LABEL_LIMIT", "PROXY_TOKEN_MIN_BYTES", "OWNER_REASSIGNMENT", "CREDIT_WALLET", "ENVIRONMENT_INDICES_COMPACTED", "NVIDIA_URL", "NVIDIA_MODEL", "ZAI_URL", "ZAI_MODEL"] as const; +export const __DECL_ORDER__ = ["KIND", "ID", "REPO", "VERSION", "PROBE_TIMEOUT_MS", "PROBE_MAX_TOKENS", "MAX_PROBES_IN_FLIGHT", "PROBE_COOLDOWN_SECONDS", "MAX_KEYS_PER_OWNER", "MAX_KEY_BYTES", "MAX_BODY_BYTES", "LABEL_LIMIT", "PROXY_TOKEN_MIN_BYTES", "OWNER_REASSIGNMENT", "CREDIT_WALLET", "ENVIRONMENT_INDICES_COMPACTED", "NVIDIA_URL", "NVIDIA_MODEL", "ZAI_URL", "ZAI_MODEL", "MODEL_LIMIT", "MODEL_PROBE_MAX_TOKENS", "MODEL_CHECK_ATTEMPTS", "MODEL_LIST_LIMIT", "MODEL_LIST_CACHE_SECONDS"] as const; // What this spec holds and this backend did not print. Empty is the whole // story most of the time; an entry here is a promise the artifact does not diff --git a/trios/agent-server/apps/server/src/api/services/queen-dispatch.ts b/trios/agent-server/apps/server/src/api/services/queen-dispatch.ts index a6e65d7b43..722e6ea12b 100644 --- a/trios/agent-server/apps/server/src/api/services/queen-dispatch.ts +++ b/trios/agent-server/apps/server/src/api/services/queen-dispatch.ts @@ -422,7 +422,7 @@ export function reviewExtraLanesPerCredential( export function workerCapacityBreakdown( runtime?: ContributorRuntime, ): WorkerCapacityBreakdown { - if (runtime && (runtime.managed.length || runtime.disabled.length)) { + if (contributorRuntimeActive(runtime)) { const lanes = contributorWorkerCandidates(runtime) return { connectedCredentials: lanes.length, @@ -1065,14 +1065,42 @@ export function environmentContributorKeys(): EnvironmentKey[] { }) } +/** + * Whether the contributor registry changes anything for dispatch: a managed + * key, a disabled environment key, or an owner's model in place of a pool's. + * Otherwise allocation stays the environment path, byte for byte. + */ +function contributorRuntimeActive( + runtime?: ContributorRuntime, +): runtime is ContributorRuntime { + return ( + !!runtime && + (runtime.managed.length > 0 || + runtime.disabled.length > 0 || + Object.keys(runtime.models ?? {}).length > 0) + ) +} + +/** The owner's chosen model for an environment lane, else the pool's. */ +function withOwnerModel( + runtime: ContributorRuntime, + lane: WorkerProvider, +): WorkerProvider { + const model = + lane.keyIndex === undefined ? undefined : runtime.models?.[lane.keyIndex] + return model ? { ...lane, model } : lane +} + function contributorWorkerCandidates( runtime: ContributorRuntime, ): WorkerProvider[] { const disabled = new Set(runtime.disabled) return [ - ...environmentWorkerCandidates().filter( - (key) => key.keyIndex !== undefined && !disabled.has(key.keyIndex), - ), + ...environmentWorkerCandidates() + .filter( + (key) => key.keyIndex !== undefined && !disabled.has(key.keyIndex), + ) + .map((key) => withOwnerModel(runtime, key)), ...runtime.managed.map((key) => ({ provider: key.provider === 'zai' ? 'zai' : 'openai-compatible', model: key.model, @@ -1091,7 +1119,7 @@ export function resolveWorkerProvider( afterKeyIndex?: number, runtime?: ContributorRuntime, ): WorkerProvider | null { - if (!runtime || (!runtime.managed.length && !runtime.disabled.length)) { + if (!contributorRuntimeActive(runtime)) { return resolveEnvironmentWorkerProvider(takenKeyIndices, afterKeyIndex) } const candidates = contributorWorkerCandidates(runtime) @@ -1144,14 +1172,16 @@ export function reviewLaneCandidates( takenKeyIndices: number[] = [], runtime?: ContributorRuntime, ): WorkerProvider[] { - if (runtime && (runtime.managed.length || runtime.disabled.length)) { + if (contributorRuntimeActive(runtime)) { const managedSecrets = new Set(runtime.managed.map((key) => key.apiKey)) - const environment = reviewLaneCandidates(takenKeyIndices).filter( - (candidate) => - !managedSecrets.has(candidate.apiKey ?? '') && - (candidate.keyIndex === undefined || - !runtime.disabled.includes(candidate.keyIndex)), - ) + const environment = reviewLaneCandidates(takenKeyIndices) + .filter( + (candidate) => + !managedSecrets.has(candidate.apiKey ?? '') && + (candidate.keyIndex === undefined || + !runtime.disabled.includes(candidate.keyIndex)), + ) + .map((candidate) => withOwnerModel(runtime, candidate)) const managed = contributorWorkerCandidates(runtime) .filter( (candidate) => diff --git a/trios/agent-server/apps/server/tests/api/queen-contributor-keys.test.ts b/trios/agent-server/apps/server/tests/api/queen-contributor-keys.test.ts index dca1e6c967..da24702636 100644 --- a/trios/agent-server/apps/server/tests/api/queen-contributor-keys.test.ts +++ b/trios/agent-server/apps/server/tests/api/queen-contributor-keys.test.ts @@ -5,16 +5,26 @@ import type { Pool } from 'pg' import { createQueenContributorKeysRoute } from '../../src/api/routes/queen-contributor-keys' import { addContributorKey, + bareLabel, + checkModel, contributorGithub, + environmentLabel, fingerprint, + LABEL_POOL_STRIDE, + nextLabel, openCredential, + ownerModel, probeCredential, + providerModels, + resetProviderModels, sealCredential, trustedContributor, + validModel, } from '../../src/api/services/queen-contributor-keys' import { CONTRIBUTOR_POLICY } from '../../src/api/services/queen-contributor-policy' import { environmentContributorKeys, + POOL_KEY_STRIDE, resolveWorkerProvider, reviewLaneCandidates, workerCapacityBreakdown, @@ -355,3 +365,181 @@ describe('stable allocation and attribution', () => { expect(result.consts[name]?.value).toEqual(CONTRIBUTOR_POLICY[name]) }) }) + +describe('one model for all keys of a provider', () => { + const key = { + provider: 'nvidia' as const, + model: 'z-ai/glm-5.3', + apiKey: 'never-reflect-this-secret', + } + it('accepts catalog model ids and refuses anything else', () => { + for (const model of [ + 'z-ai/glm-5.3', + 'glm-4.5-flash', + 'nvidia/nemotron-3-super-120b-a12b', + 'a'.repeat(CONTRIBUTOR_POLICY.MODEL_LIMIT), + ]) + expect(validModel(model)).toBe(true) + for (const model of [ + '', + ' z-ai/glm-5.3', + 'glm 5', + '../glm', + 'glm\n', + 'a'.repeat(CONTRIBUTOR_POLICY.MODEL_LIMIT + 1), + 42, + null, + ]) + expect(validModel(model)).toBe(false) + }) + it('a model passes only by calling the offered tool', async () => { + const seen: Record[] = [] + const answer = (status: number, body: unknown) => + (async (_url, init) => { + seen.push(JSON.parse(String(init?.body))) + return Response.json(body, { status }) + }) as typeof fetch + expect( + await checkModel( + key, + answer(200, { + choices: [{ message: { tool_calls: [{ id: 'ping' }] } }], + }), + ), + ).toBe('ok') + expect(seen[0]).toMatchObject({ + model: 'z-ai/glm-5.3', + max_tokens: CONTRIBUTOR_POLICY.MODEL_PROBE_MAX_TOKENS, + tool_choice: 'auto', + tools: [{ type: 'function', function: { name: 'ping' } }], + }) + expect( + await checkModel( + key, + answer(200, { choices: [{ message: { content: 'pong' } }] }), + ), + ).toBe('no_tool_call') + for (const status of [400, 404, 410, 422]) + expect(await checkModel(key, answer(status, {}))).toBe('unknown_model') + for (const status of [401, 429, 500, 503]) + expect(await checkModel(key, answer(status, {}))).toBe('inconclusive') + expect( + await checkModel(key, (async () => { + throw new Error(key.apiKey) + }) as typeof fetch), + ).toBe('inconclusive') + }) + it('lists catalog ids, drops malformed ones and never caches a failed read', async () => { + resetProviderModels() + let reads = 0 + const broken = (async () => { + reads++ + return new Response('down', { status: 503 }) + }) as typeof fetch + expect(await providerModels(key, broken)).toEqual([]) + expect(await providerModels(key, broken)).toEqual([]) + expect(reads).toBe(2) + const catalog = (async (url, init) => { + reads++ + expect(url).toBe(`${CONTRIBUTOR_POLICY.NVIDIA_URL}/models`) + expect(init?.redirect).toBe('error') + return Response.json({ + data: [ + { id: 'z-ai/glm-5.3' }, + { id: 'z-ai/glm-5.3' }, + { id: 'has space' }, + { id: 7 }, + { id: 'a-model' }, + ], + }) + }) as typeof fetch + expect(await providerModels(key, catalog)).toEqual([ + 'a-model', + 'z-ai/glm-5.3', + ]) + expect(await providerModels(key, broken)).toEqual([ + 'a-model', + 'z-ai/glm-5.3', + ]) + expect(reads).toBe(3) + resetProviderModels() + }) + it('the owner model is the one most of their keys of the provider run', () => { + expect( + ownerModel( + [ + { provider: 'nvidia', model: 'b' }, + { provider: 'nvidia', model: 'a' }, + { provider: 'nvidia', model: 'b' }, + { provider: 'zai', model: 'z' }, + ], + 'nvidia', + ), + ).toBe('b') + expect(ownerModel([], 'zai')).toBe(CONTRIBUTOR_POLICY.ZAI_MODEL) + }) + it('bees and reviews run on the owner model while the pool keeps its own', () => { + process.env.TRIOS_QUEEN_WORKER_BASE_URL = CONTRIBUTOR_POLICY.NVIDIA_URL + process.env.TRIOS_QUEEN_WORKER_PROVIDER = 'openai-compatible' + process.env.TRIOS_QUEEN_WORKER_MODEL = 'pool-model' + process.env.TRIOS_QUEEN_WORKER_API_KEY = 'key-a' + process.env.TRIOS_QUEEN_WORKER_API_KEY_2 = 'key-b' + delete process.env.TRIOS_QUEEN_WORKER_API_KEY_3 + const runtime = { managed: [], disabled: [], models: { 1: 'z-ai/glm-5.3' } } + expect(resolveWorkerProvider([0, 0], undefined, runtime)).toMatchObject({ + keyIndex: 1, + model: 'z-ai/glm-5.3', + }) + expect(resolveWorkerProvider([1, 1], undefined, runtime)).toMatchObject({ + keyIndex: 0, + model: 'pool-model', + }) + expect( + reviewLaneCandidates([], runtime).map((lane) => [ + lane.keyIndex, + lane.model, + ]), + ).toEqual([ + [0, 'pool-model'], + [1, 'z-ai/glm-5.3'], + ]) + expect(workerCapacityBreakdown(runtime).connectedCredentials).toBe(2) + // Without a choice, allocation stays the environment path. + expect( + resolveWorkerProvider([], undefined, { + managed: [], + disabled: [], + models: {}, + }), + ).toMatchObject({ keyIndex: 0, model: 'pool-model' }) + }) +}) + +describe('key names', () => { + it('environment keys are named by their place in their pool', () => { + expect(LABEL_POOL_STRIDE).toBe(POOL_KEY_STRIDE) + expect(environmentLabel('nvidia', 0)).toBe('nvidia #1') + expect(environmentLabel('nvidia', 23)).toBe('nvidia #24') + expect(environmentLabel('zai', 10000)).toBe('zai #1') + expect(environmentLabel('zai', 10009)).toBe('zai #10') + }) + it('an unnamed key takes the next number after pool places and numbered keys', () => { + const pool24 = Array.from({ length: 24 }, (_, id) => ({ + id, + provider: 'nvidia' as const, + model: 'm', + apiKey: `k${id}`, + baseUrl: CONTRIBUTOR_POLICY.NVIDIA_URL, + })) + expect(nextLabel('nvidia', [], pool24)).toBe('nvidia #25') + expect(nextLabel('nvidia', ['nvidia #25', 'My key'], pool24)).toBe( + 'nvidia #26', + ) + expect(nextLabel('nvidia', ['zai #40'], pool24)).toBe('nvidia #25') + expect(nextLabel('zai', [], pool24)).toBe('zai #1') + for (const label of ['', ' ', 'nvidia', 'NVIDIA', 'NVIDIA NIM']) + expect(bareLabel('nvidia', label)).toBe(true) + for (const label of ['nvidia #25', 'Laptop', 'zai']) + expect(bareLabel('nvidia', label)).toBe(false) + }) +}) diff --git a/trios/agent-server/apps/server/tests/pglive/queen-contributor-keys-live.test.ts b/trios/agent-server/apps/server/tests/pglive/queen-contributor-keys-live.test.ts index 5647be61bc..75d7f9e8b1 100644 --- a/trios/agent-server/apps/server/tests/pglive/queen-contributor-keys-live.test.ts +++ b/trios/agent-server/apps/server/tests/pglive/queen-contributor-keys-live.test.ts @@ -10,6 +10,7 @@ import { contributorRuntime, ensureContributorKeys, listContributorKeys, + setContributorModel, } from '../../src/api/services/queen-contributor-keys' import { CONTRIBUTOR_POLICY } from '../../src/api/services/queen-contributor-policy' import { keyWork, leaderboard } from '../../src/api/services/queen-leaderboard' @@ -205,6 +206,7 @@ live('contributor ownership and allocation against real PostgreSQL', () => { expect(await contributorRuntime(pool, environment)).toEqual({ disabled: [10000], managed: [], + models: {}, }) } finally { process.env.QUEEN_CONTRIBUTOR_ENCRYPTION_KEY = master @@ -250,6 +252,7 @@ live('contributor ownership and allocation against real PostgreSQL', () => { expect(await contributorRuntime(pool, environment)).toEqual({ disabled: [10000], managed: [], + models: {}, }) } finally { process.env.QUEEN_CONTRIBUTOR_PROXY_TOKEN = capability @@ -485,4 +488,354 @@ live('contributor ownership and allocation against real PostgreSQL', () => { hours: 1, }) }) + it('numbers keys by provider: pool places first, then each unnamed key in turn', async () => { + const env = [ + ...environment, + { + id: 3, + provider: 'zai' as const, + model: 'z-model', + apiKey: 'zai-pool-one-key', + baseUrl: CONTRIBUTOR_POLICY.ZAI_URL, + }, + ] + const withZai = { ...owners, 3: '@dmitrii-f-t27' } + const unnamed = await addContributorKey( + pool, + subject, + { provider: 'nvidia', apiKey: 'unnamed-one', label: '' }, + env, + fetchOk, + ) + const bare = await addContributorKey( + pool, + subject, + { + provider: 'nvidia', + apiKey: 'named-after-provider', + label: 'NVIDIA NIM', + }, + env, + fetchOk, + ) + const named = await addContributorKey( + pool, + subject, + { provider: 'nvidia', apiKey: 'named-key', label: 'Laptop key' }, + env, + fetchOk, + ) + const zai = await addContributorKey( + pool, + subject, + { provider: 'zai', apiKey: 'unnamed-zai' }, + env, + fetchOk, + ) + expect([unnamed.label, bare.label, named.label, zai.label]).toEqual([ + 'nvidia #2', + 'nvidia #3', + 'Laptop key', + 'zai #5', + ]) + // A key saved before numbering, with no name: numbered once, oldest first. + await pool.query( + "UPDATE queen_contributor_keys SET label='' WHERE key_index=$1", + [unnamed.id], + ) + const listed = await listContributorKeys(pool, subject, env, withZai) + expect( + Object.fromEntries(listed.map((key) => [key.id, key.label])), + ).toEqual({ + [unnamed.id]: 'nvidia #4', + [bare.id]: 'nvidia #3', + [named.id]: 'Laptop key', + [zai.id]: 'zai #5', + 10000: 'nvidia #1', + 3: 'zai #4', + }) + const again = await listContributorKeys(pool, subject, env, withZai) + expect(again.find((key) => key.id === unnamed.id)?.label).toBe('nvidia #4') + const { rows } = await pool.query( + 'SELECT label FROM queen_contributor_keys WHERE key_index=$1', + [unnamed.id], + ) + expect(rows[0].label).toBe('nvidia #4') + }) + + describe("one model for all of an owner's keys of a provider", () => { + const toolCall = (calls: string[]) => + (async (_url, init) => { + calls.push(JSON.parse(String(init?.body)).model) + return Response.json({ + choices: [{ message: { tool_calls: [{ id: 'ping' }] } }], + }) + }) as typeof fetch + const app = (env = environment, fetcher: typeof fetch = fetchOk) => + new Hono().route( + '/queen/contributor-keys', + createQueenContributorKeysRoute({ + pool: () => pool, + environment: () => env, + owners: () => owners, + fetcher, + }), + ) + const headers = { + Authorization: `Bearer ${token}`, + 'X-Queen-Contributor-Id': subject, + 'Content-Type': 'application/json', + } + it('a stale binding row never overrides a changed pool model', async () => { + await changeContributorKey( + pool, + subject, + 10000, + 'probe', + environment, + owners, + fetchOk, + ) + const moved = [{ ...environment[0], model: 'new-pool-model' }] + const keys = await listContributorKeys(pool, subject, moved, owners) + expect(keys.find((key) => key.id === 10000)?.model).toBe('new-pool-model') + expect((await contributorRuntime(pool, moved)).models).toEqual({}) + }) + it('changes nothing unless a key of the provider called a tool with the model', async () => { + const mine = await addContributorKey( + pool, + subject, + { provider: 'nvidia', apiKey: 'managed-nvidia' }, + environment, + fetchOk, + ) + for (const [fetcher, code] of [ + [ + (async () => + new Response('no such model', { status: 404 })) as typeof fetch, + 'model_unavailable', + ], + [ + (async () => + Response.json({ + choices: [{ message: { content: 'pong' } }], + })) as typeof fetch, + 'model_without_tools', + ], + ] as const) + await expect( + setContributorModel( + pool, + subject, + { provider: 'nvidia', model: 'z-ai/glm-5.3' }, + environment, + owners, + fetcher, + ), + ).rejects.toThrow(code) + let calls = 0 + await expect( + setContributorModel( + pool, + subject, + { provider: 'nvidia', model: 'z-ai/glm-5.3' }, + environment, + owners, + (async () => { + calls++ + return new Response('busy', { status: 429 }) + }) as typeof fetch, + ), + ).rejects.toThrow('model_check_failed') + expect(calls).toBe(CONTRIBUTOR_POLICY.MODEL_CHECK_ATTEMPTS) + const keys = await listContributorKeys(pool, subject, environment, owners) + expect(keys.map((key) => key.model)).toEqual([ + 'nvidia/nemotron-3-ultra-550b-a55b', + 'production-model', + ]) + expect(mine.model).toBe(CONTRIBUTOR_POLICY.NVIDIA_MODEL) + expect((await contributorRuntime(pool, environment)).models).toEqual({}) + }) + it('switches every owned key of the provider and nothing else', async () => { + const mine = await addContributorKey( + pool, + subject, + { provider: 'nvidia', apiKey: 'managed-nvidia' }, + environment, + fetchOk, + ) + const zai = await addContributorKey( + pool, + subject, + { provider: 'zai', apiKey: 'managed-zai' }, + environment, + fetchOk, + ) + const theirs = await addContributorKey( + pool, + other, + { provider: 'nvidia', apiKey: 'their-nvidia' }, + environment, + fetchOk, + ) + const calls: string[] = [] + const result = await setContributorModel( + pool, + subject, + { provider: 'nvidia', model: 'z-ai/glm-5.3' }, + environment, + owners, + toolCall(calls), + ) + expect(calls).toEqual(['z-ai/glm-5.3']) + expect(result).toEqual({ + provider: 'nvidia', + model: 'z-ai/glm-5.3', + keys: 2, + }) + const keys = await listContributorKeys(pool, subject, environment, owners) + expect( + Object.fromEntries(keys.map((key) => [key.id, key.model])), + ).toEqual({ + [mine.id]: 'z-ai/glm-5.3', + [zai.id]: CONTRIBUTOR_POLICY.ZAI_MODEL, + 10000: 'z-ai/glm-5.3', + }) + const runtime = await contributorRuntime(pool, environment) + expect(runtime.models).toEqual({ 10000: 'z-ai/glm-5.3' }) + expect( + Object.fromEntries(runtime.managed.map((key) => [key.id, key.model])), + ).toEqual({ + [mine.id]: 'z-ai/glm-5.3', + [zai.id]: CONTRIBUTOR_POLICY.ZAI_MODEL, + [theirs.id]: CONTRIBUTOR_POLICY.NVIDIA_MODEL, + }) + // A probe asks the chosen model, and a new key joins it. + const probed: string[] = [] + await pool.query( + "UPDATE queen_contributor_keys SET probe_at=now()-interval '2 minutes'", + ) + await changeContributorKey( + pool, + subject, + 10000, + 'probe', + environment, + owners, + toolCall(probed), + ) + expect(probed).toEqual(['z-ai/glm-5.3']) + const later = await addContributorKey( + pool, + subject, + { provider: 'nvidia', apiKey: 'later-nvidia' }, + environment, + toolCall(probed), + ) + expect(later.model).toBe('z-ai/glm-5.3') + }) + it('a managed copy of an environment secret carries the choice to that environment key', async () => { + await addContributorKey( + pool, + subject, + { provider: 'nvidia', apiKey: 'copied-later-into-environment' }, + [], + fetchOk, + ) + const both = [ + ...environment, + { + id: 10001, + provider: 'nvidia' as const, + model: 'production-model', + apiKey: 'copied-later-into-environment', + baseUrl: CONTRIBUTOR_POLICY.NVIDIA_URL, + }, + ] + const withCopy = { ...owners, 10001: '@dmitrii-f-t27' } + await setContributorModel( + pool, + subject, + { provider: 'nvidia', model: 'z-ai/glm-5.3' }, + both, + withCopy, + toolCall([]), + ) + expect((await contributorRuntime(pool, both)).models).toEqual({ + 10000: 'z-ai/glm-5.3', + 10001: 'z-ai/glm-5.3', + }) + }) + it('serves the switch and the owner model through the route, refusing other fields first', async () => { + let fetched = 0 + const counting = (async (url, init) => { + fetched++ + return toolCall([])(url, init) + }) as typeof fetch + for (const body of [ + { provider: 'nvidia', model: 'z-ai/glm-5.3', owner: 'telegram:1' }, + { provider: 'nvidia', model: 'bad model' }, + { provider: 'openai', model: 'gpt' }, + { provider: 'nvidia', model: 'a'.repeat(129) }, + ]) { + const response = await app(environment, counting).request( + '/queen/contributor-keys/model', + { method: 'POST', headers, body: JSON.stringify(body) }, + ) + expect(response.status).toBe(400) + } + expect(fetched).toBe(0) + const response = await app(environment, counting).request( + '/queen/contributor-keys/model', + { + method: 'POST', + headers, + body: JSON.stringify({ provider: 'nvidia', model: 'z-ai/glm-5.3' }), + }, + ) + expect(response.status).toBe(200) + expect(response.headers.get('cache-control')).toBe('no-store') + const account = await response.json() + expect(account.providers).toEqual([ + { + id: 'nvidia', + label: 'NVIDIA NIM', + model: 'z-ai/glm-5.3', + defaultModel: CONTRIBUTOR_POLICY.NVIDIA_MODEL, + }, + { + id: 'zai', + label: 'Z.ai', + model: CONTRIBUTOR_POLICY.ZAI_MODEL, + defaultModel: CONTRIBUTOR_POLICY.ZAI_MODEL, + }, + ]) + expect(JSON.stringify(account)).not.toContain(environment[0].apiKey) + const foreign = await app(environment, counting).request( + '/queen/contributor-keys/model', + { + method: 'POST', + headers: { ...headers, 'X-Queen-Contributor-Id': other }, + body: JSON.stringify({ provider: 'nvidia', model: 'z-ai/glm-5.3' }), + }, + ) + expect(foreign.status).toBe(404) + const list = await app(environment, (async () => + Response.json({ + data: [ + { id: 'z-ai/glm-5.3' }, + { id: 'bad id' }, + { id: 'nvidia/nemotron-3-super-120b-a12b' }, + ], + })) as typeof fetch).request('/queen/contributor-keys/models/nvidia', { + headers, + }) + expect(list.status).toBe(200) + expect(await list.json()).toEqual({ + provider: 'nvidia', + model: 'z-ai/glm-5.3', + models: ['nvidia/nemotron-3-super-120b-a12b', 'z-ai/glm-5.3'], + }) + }) + }) }) diff --git a/trios/agent-server/docs/contributor-keys.md b/trios/agent-server/docs/contributor-keys.md index b1fade5cea..ae060b5441 100644 --- a/trios/agent-server/docs/contributor-keys.md +++ b/trios/agent-server/docs/contributor-keys.md @@ -69,13 +69,48 @@ verified contributor header. Responses use `Cache-Control: no-store`. - `POST /queen/contributor-keys/:id/probe`: check an owned key. - `POST /queen/contributor-keys/:id/enable`: check and enable on success. - `POST /queen/contributor-keys/:id/disable`: stop assigning new work. +- `GET /queen/contributor-keys/models/:provider`: the provider's model ids + (its `/models` catalog read with one of the owner's keys, cached ten + minutes) and the model the owner's keys of that provider run now. +- `POST /queen/contributor-keys/model`: `{provider, model}` sets one model for + every key the owner holds of that provider and returns the account. + +## One model for all of an owner's keys + +The model switch is all-or-nothing for that owner and that provider, and the +evidence comes first. One of the owner's keys of the provider is handed the +model with a single `ping` tool; only a tool call counts. 400/404/410/422 is +`model_unavailable`, an answer without a tool call is `model_without_tools`, +and a refused or busy key passes the check to the next key, at most +`MODEL_CHECK_ATTEMPTS` times, before `model_check_failed` (504). Nothing is +written unless the check passed. + +On success the owner's environment keys are bound (the same binding a probe +creates) and every owned row of the provider gets `model` plus +`model_chosen = true`. Allocation, reviews and probes then use that model for +those keys. Other owners' keys and the operator's pool variables are +untouched. A key the owner adds later joins the chosen model. + +`model_chosen` exists because a binding row has always recorded the model it +was bound under. Without the flag, a key probed on 2026-10-01 under +`nemotron-3-ultra` would keep that model after the operator moved the pool to +another one. An environment row overrides its pool's model only while the +flag is set. If the owner also holds a managed copy of a secret that the +environment later received, the copy cannot be bound twice (fingerprints are +unique); its choice applies to the environment key instead. + +The Queen's model ranking (`TRIOS_QUEEN_WORKER_MODEL_CANDIDATES`) reroutes +only requests that already name one of its candidates. An owner model outside +that list is sent as chosen. Successful mutations return `{key}` with the same metadata and actual contribution fields as GET. Secrets, ciphertext, provider response bodies, upstream error messages and operator configuration are never returned. Errors are `{error: }`. Foreign/absent keys both return 404. -A probe performs one bounded chat-completion request. An HTTP 200 with no +A probe performs one bounded chat-completion request, waiting up to +`PROBE_TIMEOUT_MS` (90 s since v2: on 2026-10-02 `z-ai/glm-5.3` on NVIDIA took +20-46 s for a short answer and `glm-4.5-flash` on Z.ai up to 87 s). An HTTP 200 with no model output is not a successful check. Invalid credentials are disabled; temporary rate limits or network failures do not revoke prior consent. Each key allows one explicit check per minute, and the process caps concurrent diff --git a/trios/agent-server/specs/automation/queen-contributor-keys.t27 b/trios/agent-server/specs/automation/queen-contributor-keys.t27 index 6649fef86c..dd47361961 100644 --- a/trios/agent-server/specs/automation/queen-contributor-keys.t27 +++ b/trios/agent-server/specs/automation/queen-contributor-keys.t27 @@ -7,15 +7,18 @@ // XP follows work recorded by the Queen. Adding credentials never credits a wallet. // Negative managed indices and nonnegative environment indices share one history. // v1 (2026-10-01): trusted proxy, persistent ownership and bounded provider probes. +// v2 (2026-10-02): one model for all of an owner's keys of a provider, accepted only +// after a key of that provider called a tool with it. Probe timeout 15 s -> 90 s: +// z-ai/glm-5.3 on NVIDIA answered in 20-46 s and glm-4.5-flash on Z.ai in up to 87 s. // phi^2 + 1/phi^2 = 3 | TRINITY module automation::queen_contributor_keys { pub const KIND : str = "automation"; pub const ID : str = "queen-contributor-keys"; pub const REPO : str = "BrowserOS"; - pub const VERSION : u8 = 1; + pub const VERSION : u8 = 2; - pub const PROBE_TIMEOUT_MS : u32 = 15000; + pub const PROBE_TIMEOUT_MS : u32 = 90000; pub const PROBE_MAX_TOKENS : u16 = 16; pub const MAX_PROBES_IN_FLIGHT : u16 = 4; pub const PROBE_COOLDOWN_SECONDS : u16 = 60; @@ -31,6 +34,11 @@ module automation::queen_contributor_keys { pub const NVIDIA_MODEL : str = "nvidia/nemotron-3-ultra-550b-a55b"; pub const ZAI_URL : str = "https://api.z.ai/api/paas/v4"; pub const ZAI_MODEL : str = "glm-4.5-flash"; + pub const MODEL_LIMIT : u16 = 128; + pub const MODEL_PROBE_MAX_TOKENS : u16 = 512; + pub const MODEL_CHECK_ATTEMPTS : u16 = 2; + pub const MODEL_LIST_LIMIT : u16 = 500; + pub const MODEL_LIST_CACHE_SECONDS : u16 = 600; fn trusted_proxy(configured_bytes: u32, token_matches: bool, subject_valid: bool) -> bool { return configured_bytes >= PROXY_TOKEN_MIN_BYTES && token_matches && subject_valid; @@ -66,6 +74,18 @@ module automation::queen_contributor_keys { return was_enabled; } + fn model_fits(chars: u32, allowed_chars: bool) -> bool { + return chars > 0 && chars <= MODEL_LIMIT && allowed_chars; + } + + fn may_switch_model(owns_provider_key: bool, called_tool: bool) -> bool { + return owns_provider_key && called_tool; + } + + fn may_try_model_check(attempt: u32) -> bool { + return attempt < MODEL_CHECK_ATTEMPTS; + } + fn wallet_credit(added: u32) -> u32 { if CREDIT_WALLET { return added; } return 0; @@ -130,8 +150,33 @@ module automation::queen_contributor_keys { assert(wallet_credit(100) == 0); } + test "a model name is a bounded nonempty identifier" { + assert(model_fits(1, true)); + assert(model_fits(128, true)); + assert(!model_fits(0, true)); + assert(!model_fits(129, true)); + assert(!model_fits(12, false)); + } + + test "a model switch needs an owned key of the provider and a tool call" { + assert(may_switch_model(true, true)); + assert(!may_switch_model(false, true)); + assert(!may_switch_model(true, false)); + } + + test "a busy or refused key passes the model check to the next one, a bounded number of times" { + assert(may_try_model_check(0)); + assert(may_try_model_check(1)); + assert(!may_try_model_check(2)); + } + + test "a probe waits long enough for a slow reasoning model" { + assert(PROBE_TIMEOUT_MS >= 60000); + assert(MODEL_PROBE_MAX_TOKENS > PROBE_MAX_TOKENS); + } + test "VERSION says what the header's last note says" { - assert(VERSION == 1); + assert(VERSION == 2); } }