Merge current tool metrics into skill metrics
Signed-off-by: Alex Fournier <afournier@nvidia.com>
This commit is contained in:
Alex Fournier
commit
e3d10f0a8f
|
|
@ -5259,7 +5259,16 @@ def resolve_provider_client(
|
|||
# sent to Codex after the main lane fell back to gpt-5.5). Let _resolve_auto()
|
||||
# return the actual current runtime model when the caller did not explicitly
|
||||
# request one. (# compression-current-model)
|
||||
if not model and provider != "auto":
|
||||
#
|
||||
# Nous + vision is the one carve-out: the branch below resolves its model
|
||||
# from the Portal's tier-aware vision recommendation (``_try_nous(vision=
|
||||
# True)``), and ``final_model = model or default`` means anything pre-filled
|
||||
# here wins over that. The main chat model is routinely text-only (e.g. a
|
||||
# ``:free`` chat SKU), so pre-filling it sends the image to a model that
|
||||
# cannot accept one and the Portal 404s. Leave ``model`` unset and let the
|
||||
# Portal slot through; only an explicit caller model may override it.
|
||||
_nous_portal_vision = provider == "nous" and is_vision
|
||||
if not model and provider != "auto" and not _nous_portal_vision:
|
||||
model = _get_aux_model_for_provider(provider) or _read_main_model_for_aux() or model
|
||||
|
||||
def _needs_codex_wrap(client_obj, base_url_str: str, model_str: str) -> bool:
|
||||
|
|
@ -6182,10 +6191,17 @@ def resolve_vision_provider_client(
|
|||
# DeepSeek-V4-Flash default) and _main_model_supports_vision can't be
|
||||
# trusted to catch that. Only fall back to the chat model when no
|
||||
# provider default is available (catalog unreachable).
|
||||
vision_model = _resolve_provider_vision_default(main_provider) or main_model
|
||||
provider_vision_default = _resolve_provider_vision_default(main_provider)
|
||||
vision_model = provider_vision_default or main_model
|
||||
if main_provider == "nous":
|
||||
# Nous resolves its vision model from the Portal's tier-aware
|
||||
# recommended-models slots inside _try_nous(vision=True).
|
||||
# Passing the chat model here overrides that pick, so a
|
||||
# text-only chat default (e.g. a `:free` chat SKU) receives the
|
||||
# image and the upstream rejects it with a 404. Only an
|
||||
# explicit auxiliary.vision.model may override the Portal.
|
||||
sync_client, default_model = _resolve_strict_vision_backend(
|
||||
main_provider, vision_model
|
||||
main_provider, resolved_model or provider_vision_default
|
||||
)
|
||||
if sync_client is not None:
|
||||
logger.info(
|
||||
|
|
|
|||
|
|
@ -150,6 +150,7 @@ import {
|
|||
tokenNeedsRefresh
|
||||
} from './native-oauth'
|
||||
import { runNativeLogin } from './native-oauth-login'
|
||||
import { loadNativeTokenSet, type NativeTokenStoreIo, persistNativeTokenSet } from './native-token-store'
|
||||
import { serializeJsonBody, setJsonRequestHeaders } from './oauth-net-request'
|
||||
import { createKeepAwake } from './power-save'
|
||||
import { FirstRunSetupResetError, runPrimaryBackendStartup } from './primary-backend-startup'
|
||||
|
|
@ -6224,35 +6225,24 @@ function _nativeTokenStorePath() {
|
|||
return path.join(app.getPath('userData'), 'native-oauth-tokens.json')
|
||||
}
|
||||
|
||||
function _readNativeTokenStore(): Record<string, any> {
|
||||
try {
|
||||
const raw = fs.readFileSync(_nativeTokenStorePath(), 'utf8')
|
||||
const parsed = JSON.parse(raw)
|
||||
|
||||
return parsed && typeof parsed === 'object' ? parsed : {}
|
||||
} catch {
|
||||
return {}
|
||||
// The electron-coupled half of the token store: safeStorage encryption plus the
|
||||
// userData file. native-token-store.ts owns the serialization/parse round trip
|
||||
// so it can be tested without an Electron runtime.
|
||||
function _nativeTokenStoreIo(): NativeTokenStoreIo {
|
||||
return {
|
||||
encrypt: encryptDesktopSecret,
|
||||
decrypt: decryptDesktopSecret,
|
||||
readStoreText: () => fs.readFileSync(_nativeTokenStorePath(), 'utf8'),
|
||||
writeStoreText: (text: string) => {
|
||||
fs.mkdirSync(path.dirname(_nativeTokenStorePath()), { recursive: true })
|
||||
fs.writeFileSync(_nativeTokenStorePath(), text, { mode: 0o600 })
|
||||
},
|
||||
rememberLog
|
||||
}
|
||||
}
|
||||
|
||||
function _persistNativeTokens(baseUrl: string, tokens: NativeTokenSet | null) {
|
||||
const store = _readNativeTokenStore()
|
||||
|
||||
if (tokens) {
|
||||
// Encrypt the whole token set as one blob so the refresh token never
|
||||
// lands in plaintext on disk. Reuse the hardened encrypt helper.
|
||||
const secret = encryptDesktopSecret(JSON.stringify(tokens))
|
||||
store[baseUrl] = secret
|
||||
} else {
|
||||
delete store[baseUrl]
|
||||
}
|
||||
|
||||
try {
|
||||
fs.mkdirSync(path.dirname(_nativeTokenStorePath()), { recursive: true })
|
||||
fs.writeFileSync(_nativeTokenStorePath(), JSON.stringify(store), { mode: 0o600 })
|
||||
} catch (error) {
|
||||
rememberLog(`[native-oauth] failed to persist tokens: ${(error as Error).message}`)
|
||||
}
|
||||
persistNativeTokenSet(baseUrl, tokens, _nativeTokenStoreIo())
|
||||
}
|
||||
|
||||
function _loadNativeTokens(baseUrl: string): NativeTokenSet | null {
|
||||
|
|
@ -6262,27 +6252,13 @@ function _loadNativeTokens(baseUrl: string): NativeTokenSet | null {
|
|||
return cached
|
||||
}
|
||||
|
||||
const store = _readNativeTokenStore()
|
||||
const secret = store[baseUrl]
|
||||
const tokens = loadNativeTokenSet(baseUrl, _nativeTokenStoreIo())
|
||||
|
||||
if (!secret) {
|
||||
return null
|
||||
}
|
||||
|
||||
try {
|
||||
const plaintext = decryptDesktopSecret(secret)
|
||||
|
||||
if (!plaintext) {
|
||||
return null
|
||||
}
|
||||
|
||||
const tokens = parseTokenResponse(JSON.parse(plaintext))
|
||||
if (tokens) {
|
||||
_nativeTokens.set(baseUrl, tokens)
|
||||
|
||||
return tokens
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
|
||||
return tokens
|
||||
}
|
||||
|
||||
function _storeNativeTokens(baseUrl: string, tokens: NativeTokenSet) {
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import {
|
|||
nativeRefreshUrl,
|
||||
nativeTokenUrl,
|
||||
parseLoopbackCallback,
|
||||
parseStoredTokenSet,
|
||||
parseTokenResponse,
|
||||
resolveLoginStrategy,
|
||||
statusSupportsNativeFlow,
|
||||
|
|
@ -176,6 +177,42 @@ test('parseTokenResponse tolerates an absent refresh token / expiry', () => {
|
|||
assert.equal(t.expiresAt, 0)
|
||||
})
|
||||
|
||||
test('parseStoredTokenSet maps the encrypted on-disk camelCase shape', () => {
|
||||
const t = parseStoredTokenSet({
|
||||
accessToken: 'AT-stored',
|
||||
refreshToken: 'RT-stored',
|
||||
expiresAt: 1893456000,
|
||||
provider: 'self-hosted',
|
||||
userId: 'u-stored'
|
||||
})
|
||||
|
||||
assert.equal(t.accessToken, 'AT-stored')
|
||||
assert.equal(t.refreshToken, 'RT-stored')
|
||||
assert.equal(t.expiresAt, 1893456000)
|
||||
assert.equal(t.provider, 'self-hosted')
|
||||
assert.equal(t.userId, 'u-stored')
|
||||
})
|
||||
|
||||
test('parseTokenResponse cannot read a persisted set (the reload bug #73271)', () => {
|
||||
// Guards against regressing to the wrong parser on the reload path: a
|
||||
// persisted camelCase set has no snake_case access_token, so the raw-response
|
||||
// parser throws — which is exactly why the stored path must use
|
||||
// parseStoredTokenSet instead.
|
||||
const persisted = JSON.parse(
|
||||
JSON.stringify({ accessToken: 'AT', refreshToken: 'RT', expiresAt: 1, provider: 'nous', userId: 'u' })
|
||||
)
|
||||
|
||||
assert.throws(() => parseTokenResponse(persisted), /missing access_token/i)
|
||||
assert.equal(parseStoredTokenSet(persisted).accessToken, 'AT')
|
||||
})
|
||||
|
||||
test('parseStoredTokenSet rejects a non-normalized server response', () => {
|
||||
assert.throws(
|
||||
() => parseStoredTokenSet({ access_token: 'AT-server' }),
|
||||
/missing accessToken/i
|
||||
)
|
||||
})
|
||||
|
||||
// --- refresh timing ---
|
||||
|
||||
test('tokenNeedsRefresh respects the skew window', () => {
|
||||
|
|
|
|||
|
|
@ -193,6 +193,31 @@ export function parseTokenResponse(body: any): NativeTokenSet {
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate a token set loaded from the encrypted local store.
|
||||
*
|
||||
* The stored representation is already normalized as NativeTokenSet and
|
||||
* therefore uses camelCase. Gateway token responses use snake_case and
|
||||
* remain handled separately by parseTokenResponse().
|
||||
*/
|
||||
export function parseStoredTokenSet(body: any): NativeTokenSet {
|
||||
const accessToken = String(body?.accessToken || '')
|
||||
|
||||
if (!accessToken) {
|
||||
throw new Error('Stored token set missing accessToken')
|
||||
}
|
||||
|
||||
const expiresAt = Number(body?.expiresAt)
|
||||
|
||||
return {
|
||||
accessToken,
|
||||
refreshToken: String(body?.refreshToken || ''),
|
||||
expiresAt: Number.isFinite(expiresAt) ? expiresAt : 0,
|
||||
provider: String(body?.provider || ''),
|
||||
userId: String(body?.userId || '')
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* True when a stored token set is at/near expiry and should be refreshed
|
||||
* before use. `skewSeconds` refreshes slightly early to avoid a race where
|
||||
|
|
|
|||
|
|
@ -0,0 +1,423 @@
|
|||
/**
|
||||
* Tests for electron/native-token-store.ts — the encrypted-at-rest persistence
|
||||
* seam main.ts uses for RFC 8252 native OAuth tokens.
|
||||
*
|
||||
* The regression this file exists for (#73271): tokens are persisted as a
|
||||
* normalized camelCase NativeTokenSet, but the reload path fed the freshly
|
||||
* decrypted object to parseTokenResponse(), which only understands the
|
||||
* gateway's snake_case response. It threw on every launch, so a signed-in user
|
||||
* came back signed out. The parser boundary now lives inside
|
||||
* loadNativeTokenSet(), so these tests fail if it is ever crossed again.
|
||||
*
|
||||
* "Fresh load" here means what it means after a restart: nothing survives but
|
||||
* the bytes of the store file, so every assertion below is served by
|
||||
* deserializing and decrypting that text — never by an in-memory object.
|
||||
*
|
||||
* (Wired into the vitest `electron` project via electron/**\/*.test.ts.)
|
||||
*/
|
||||
|
||||
import assert from 'node:assert/strict'
|
||||
|
||||
import { test } from 'vitest'
|
||||
|
||||
import { type NativeTokenSet, parseStoredTokenSet, parseTokenResponse } from './native-oauth'
|
||||
import { loadNativeTokenSet, type NativeTokenStoreIo, persistNativeTokenSet } from './native-token-store'
|
||||
|
||||
const GATEWAY = 'https://gw.example.com'
|
||||
|
||||
const TOKENS: NativeTokenSet = {
|
||||
accessToken: 'AT-live-abc123',
|
||||
refreshToken: 'RT-live-xyz789',
|
||||
expiresAt: 1_893_456_000,
|
||||
provider: 'nous',
|
||||
userId: 'u-42'
|
||||
}
|
||||
|
||||
interface FakeDisk {
|
||||
io: NativeTokenStoreIo
|
||||
logs: string[]
|
||||
/** The store-file text as it would sit on disk; null when the file is absent. */
|
||||
fileText: () => string | null
|
||||
}
|
||||
|
||||
/**
|
||||
* A stand-in for the userData store file plus safeStorage. Encryption is
|
||||
* base64 rather than the OS keychain — opaque-blob-in, same-plaintext-out is
|
||||
* the only property this seam depends on, and it keeps the round trip
|
||||
* observable. `initialText` models a process restart: the new "process" starts
|
||||
* with nothing but the bytes the previous one wrote.
|
||||
*/
|
||||
function createFakeDisk(initialText: string | null = null, overrides: Partial<NativeTokenStoreIo> = {}): FakeDisk {
|
||||
let text = initialText
|
||||
const logs: string[] = []
|
||||
|
||||
const io: NativeTokenStoreIo = {
|
||||
encrypt: plaintext => ({ encoding: 'safeStorage', value: Buffer.from(plaintext, 'utf8').toString('base64') }),
|
||||
decrypt: secret =>
|
||||
secret?.encoding === 'safeStorage' ? Buffer.from(String(secret.value), 'base64').toString('utf8') : '',
|
||||
readStoreText: () => {
|
||||
if (text === null) {
|
||||
// Matches fs.readFileSync on a missing file: throws, not empty string.
|
||||
throw Object.assign(new Error('ENOENT: no such file or directory'), { code: 'ENOENT' })
|
||||
}
|
||||
|
||||
return text
|
||||
},
|
||||
writeStoreText: next => {
|
||||
text = next
|
||||
},
|
||||
rememberLog: message => logs.push(message),
|
||||
...overrides
|
||||
}
|
||||
|
||||
return { io, logs, fileText: () => text }
|
||||
}
|
||||
|
||||
// --- the restart round trip ---
|
||||
|
||||
test('a camelCase token set survives store then a fresh load', () => {
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, first.io)
|
||||
|
||||
const onDisk = first.fileText()
|
||||
|
||||
assert.ok(onDisk, 'persisting must write the store file')
|
||||
|
||||
// Nothing may survive the "restart" except those bytes.
|
||||
const restarted = createFakeDisk(onDisk)
|
||||
const loaded = loadNativeTokenSet(GATEWAY, restarted.io)
|
||||
|
||||
assert.ok(loaded, 'a stored token set must reload after a restart')
|
||||
// Reconstructed from the payload, not handed back the object we stored.
|
||||
assert.notEqual(loaded, TOKENS)
|
||||
assert.deepEqual(loaded, TOKENS)
|
||||
assert.deepEqual(restarted.logs, [])
|
||||
})
|
||||
|
||||
test('a fresh load restores both tokens and preserves expiry, provider and user', () => {
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, first.io)
|
||||
|
||||
const loaded = loadNativeTokenSet(GATEWAY, createFakeDisk(first.fileText()).io)!
|
||||
|
||||
assert.equal(loaded.accessToken, 'AT-live-abc123')
|
||||
assert.equal(loaded.refreshToken, 'RT-live-xyz789')
|
||||
// Still a number after the JSON round trip, not "1893456000".
|
||||
assert.equal(loaded.expiresAt, 1_893_456_000)
|
||||
assert.equal(typeof loaded.expiresAt, 'number')
|
||||
assert.equal(loaded.provider, 'nous')
|
||||
assert.equal(loaded.userId, 'u-42')
|
||||
})
|
||||
|
||||
test('the loaded set is accepted by the stored-token parsing boundary', () => {
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, first.io)
|
||||
|
||||
const loaded = loadNativeTokenSet(GATEWAY, createFakeDisk(first.fileText()).io)!
|
||||
|
||||
// What comes back out of the store is itself a valid stored set — re-parsing
|
||||
// it is a no-op, so callers can hand it straight to the refresh path.
|
||||
assert.deepEqual(parseStoredTokenSet(loaded), loaded)
|
||||
})
|
||||
|
||||
test('the persisted payload is what broke the old reload path (#73271)', () => {
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, first.io)
|
||||
|
||||
const restarted = createFakeDisk(first.fileText())
|
||||
const secret = JSON.parse(restarted.fileText()!)[GATEWAY]
|
||||
const decrypted = JSON.parse(restarted.io.decrypt(secret))
|
||||
|
||||
// The old code passed exactly this object to parseTokenResponse(). A
|
||||
// normalized set has no snake_case access_token, so every launch threw and
|
||||
// the user was shown as signed out...
|
||||
assert.throws(() => parseTokenResponse(decrypted), /missing access_token/i)
|
||||
// ...while the real load path reads the same bytes successfully.
|
||||
assert.deepEqual(loadNativeTokenSet(GATEWAY, restarted.io), TOKENS)
|
||||
})
|
||||
|
||||
test('the full login-to-restart sequence keeps the two parser boundaries apart', () => {
|
||||
// Login: the gateway answers /auth/native/token in snake_case, and only
|
||||
// parseTokenResponse() understands that shape.
|
||||
const fromGateway = parseTokenResponse({
|
||||
access_token: 'AT-fresh',
|
||||
refresh_token: 'RT-fresh',
|
||||
expires_at: 1_893_456_789,
|
||||
provider: 'nous',
|
||||
user_id: 'u-77'
|
||||
})
|
||||
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, fromGateway, first.io)
|
||||
|
||||
// Restart: what was stored is normalized, so the store's own boundary reads
|
||||
// it back unchanged.
|
||||
assert.deepEqual(loadNativeTokenSet(GATEWAY, createFakeDisk(first.fileText()).io), fromGateway)
|
||||
})
|
||||
|
||||
// --- storage hygiene ---
|
||||
|
||||
test('tokens are encrypted at rest, never plaintext in the store file', () => {
|
||||
const disk = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, disk.io)
|
||||
|
||||
const onDisk = disk.fileText()!
|
||||
|
||||
assert.doesNotMatch(onDisk, /AT-live-abc123/)
|
||||
assert.doesNotMatch(onDisk, /RT-live-xyz789/)
|
||||
assert.equal(JSON.parse(onDisk)[GATEWAY].encoding, 'safeStorage')
|
||||
})
|
||||
|
||||
test('persisting one gateway leaves other gateways intact', () => {
|
||||
const other = 'https://other.example.com'
|
||||
const disk = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, disk.io)
|
||||
persistNativeTokenSet(other, { ...TOKENS, accessToken: 'AT-other', userId: 'u-99' }, disk.io)
|
||||
|
||||
const restarted = createFakeDisk(disk.fileText())
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, restarted.io)!.accessToken, 'AT-live-abc123')
|
||||
assert.equal(loadNativeTokenSet(other, restarted.io)!.accessToken, 'AT-other')
|
||||
})
|
||||
|
||||
test('clearing removes only that gateway and reloads as signed out', () => {
|
||||
const other = 'https://other.example.com'
|
||||
const disk = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, disk.io)
|
||||
persistNativeTokenSet(other, TOKENS, disk.io)
|
||||
persistNativeTokenSet(GATEWAY, null, disk.io)
|
||||
|
||||
const restarted = createFakeDisk(disk.fileText())
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, restarted.io), null)
|
||||
assert.ok(loadNativeTokenSet(other, restarted.io))
|
||||
})
|
||||
|
||||
test('an absent store file loads as signed out without logging a failure', () => {
|
||||
const disk = createFakeDisk()
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, disk.io), null)
|
||||
assert.deepEqual(disk.logs, [])
|
||||
})
|
||||
|
||||
// --- failure paths (unchanged by the extraction) ---
|
||||
|
||||
test('a locked keychain keeps the stored entry for a later retry', () => {
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, first.io)
|
||||
|
||||
// safeStorage unavailable at load time ⇒ decryptDesktopSecret returns ''.
|
||||
const locked = createFakeDisk(first.fileText(), { decrypt: () => '' })
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, locked.io), null)
|
||||
assert.match(locked.logs[0], /failed to decrypt stored tokens for https:\/\/gw\.example\.com/)
|
||||
assert.match(locked.logs[0], /keeping stored entry for retry/)
|
||||
// The refresh token must NOT be dropped just because the keychain was locked.
|
||||
assert.deepEqual(locked.fileText(), first.fileText())
|
||||
})
|
||||
|
||||
test('a corrupt store file loads as signed out instead of throwing', () => {
|
||||
const disk = createFakeDisk('{not json')
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, disk.io), null)
|
||||
assert.deepEqual(disk.logs, [])
|
||||
})
|
||||
|
||||
test('an array store file loads as signed out instead of throwing', () => {
|
||||
const disk = createFakeDisk('[]')
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, disk.io), null)
|
||||
assert.deepEqual(disk.logs, [])
|
||||
})
|
||||
|
||||
test('an array store file is replaced by a real map rather than swallowing the write', () => {
|
||||
const disk = createFakeDisk('[]')
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, disk.io)
|
||||
|
||||
const written = JSON.parse(disk.fileText()!)
|
||||
|
||||
// Assigning store[baseUrl] on an array sets a non-index property, which
|
||||
// JSON.stringify drops — the write would report success and the tokens would
|
||||
// be gone on the next launch.
|
||||
assert.equal(Array.isArray(written), false)
|
||||
assert.ok(written[GATEWAY], 'the gateway entry must survive serialization')
|
||||
// And it really does come back after a restart.
|
||||
assert.deepEqual(loadNativeTokenSet(GATEWAY, createFakeDisk(disk.fileText()).io), TOKENS)
|
||||
})
|
||||
|
||||
test('a corrupt decrypted blob is reported and loads as signed out', () => {
|
||||
const disk = createFakeDisk(JSON.stringify({ [GATEWAY]: { encoding: 'safeStorage', value: 'bm90LWpzb24=' } }))
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, disk.io), null)
|
||||
assert.match(disk.logs[0], /failed to load stored tokens for https:\/\/gw\.example\.com/)
|
||||
})
|
||||
|
||||
test('a decrypted blob missing accessToken is rejected, not half-restored', () => {
|
||||
const plaintext = JSON.stringify({ refreshToken: 'RT-only', provider: 'nous' })
|
||||
|
||||
const disk = createFakeDisk(
|
||||
JSON.stringify({ [GATEWAY]: { encoding: 'safeStorage', value: Buffer.from(plaintext).toString('base64') } })
|
||||
)
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, disk.io), null)
|
||||
assert.match(disk.logs[0], /missing accessToken/i)
|
||||
})
|
||||
|
||||
test('a non-Error decryption failure keeps its detail in the log', () => {
|
||||
const disk = createFakeDisk(JSON.stringify({ [GATEWAY]: { encoding: 'safeStorage', value: 'AAAA' } }), {
|
||||
decrypt: () => {
|
||||
throw 'keychain exploded'
|
||||
}
|
||||
})
|
||||
|
||||
assert.equal(loadNativeTokenSet(GATEWAY, disk.io), null)
|
||||
assert.match(disk.logs[0], /keychain exploded/)
|
||||
})
|
||||
|
||||
test('an unwritable store file is logged rather than thrown', () => {
|
||||
const disk = createFakeDisk(null, {
|
||||
writeStoreText: () => {
|
||||
throw new Error('EACCES: permission denied')
|
||||
}
|
||||
})
|
||||
|
||||
assert.doesNotThrow(() => persistNativeTokenSet(GATEWAY, TOKENS, disk.io))
|
||||
assert.match(disk.logs[0], /failed to persist tokens: EACCES/)
|
||||
})
|
||||
|
||||
test('a non-Error write failure keeps its detail in the log', () => {
|
||||
const disk = createFakeDisk(null, {
|
||||
writeStoreText: () => {
|
||||
throw 'disk went away'
|
||||
}
|
||||
})
|
||||
|
||||
// `(error as Error).message` on a thrown string reads as undefined and loses
|
||||
// the only diagnostic there was.
|
||||
assert.doesNotThrow(() => persistNativeTokenSet(GATEWAY, TOKENS, disk.io))
|
||||
assert.equal(disk.logs[0], '[native-oauth] failed to persist tokens: disk went away')
|
||||
})
|
||||
|
||||
test('an unusable keychain fails the write loudly and writes nothing', () => {
|
||||
const existing = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, existing.io)
|
||||
|
||||
const before = existing.fileText()
|
||||
|
||||
const broken = createFakeDisk(before, {
|
||||
encrypt: () => {
|
||||
throw new Error('Secure token storage is unavailable')
|
||||
}
|
||||
})
|
||||
|
||||
// Storing must not pretend to succeed when the token cannot be encrypted...
|
||||
assert.throws(() => persistNativeTokenSet(GATEWAY, { ...TOKENS, accessToken: 'AT-new' }, broken.io), /unavailable/)
|
||||
// ...and must not clobber the tokens already on disk.
|
||||
assert.equal(broken.fileText(), before)
|
||||
})
|
||||
|
||||
test('an encrypt that returns null is refused rather than blanking the stored entry', () => {
|
||||
const existing = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(GATEWAY, TOKENS, existing.io)
|
||||
|
||||
const before = existing.fileText()
|
||||
const nulled = createFakeDisk(before, { encrypt: () => null })
|
||||
let writes = 0
|
||||
|
||||
// Spy that still delegates, so a stray write would show up in BOTH the
|
||||
// counter and the file text.
|
||||
const io = {
|
||||
...nulled.io,
|
||||
writeStoreText: (text: string) => {
|
||||
writes += 1
|
||||
nulled.io.writeStoreText(text)
|
||||
}
|
||||
}
|
||||
|
||||
// A quiet null is the same failure as a throw and must be just as loud.
|
||||
assert.throws(
|
||||
() => persistNativeTokenSet(GATEWAY, { ...TOKENS, accessToken: 'AT-new' }, io),
|
||||
/refusing to overwrite stored native tokens/
|
||||
)
|
||||
assert.equal(writes, 0, 'the store file must not be written at all')
|
||||
// Byte-for-byte unchanged...
|
||||
assert.equal(nulled.fileText(), before)
|
||||
// ...and the original token set still loads, refresh token intact.
|
||||
assert.deepEqual(loadNativeTokenSet(GATEWAY, createFakeDisk(before).io), TOKENS)
|
||||
})
|
||||
|
||||
// --- credential redaction in logs ---
|
||||
//
|
||||
// normalizeRemoteBaseUrl() strips query/fragment/trailing slashes but not
|
||||
// userinfo, so a configured gateway URL can carry `user:password@` into this
|
||||
// store. It must stay intact as the store KEY and never reach a log line.
|
||||
|
||||
const CRED_GATEWAY = 'https://alice:supersecret@gw.example.com/hermes'
|
||||
|
||||
test('a decryption failure logs the gateway host and path but not its credentials', () => {
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(CRED_GATEWAY, TOKENS, first.io)
|
||||
|
||||
const before = first.fileText()
|
||||
const locked = createFakeDisk(before, { decrypt: () => '' })
|
||||
|
||||
assert.equal(loadNativeTokenSet(CRED_GATEWAY, locked.io), null)
|
||||
// Still identifies which gateway failed...
|
||||
assert.match(locked.logs[0], /failed to decrypt stored tokens for https:\/\/gw\.example\.com\/hermes/)
|
||||
assert.match(locked.logs[0], /keeping stored entry for retry/)
|
||||
// ...without the userinfo.
|
||||
assert.doesNotMatch(locked.logs[0], /alice/)
|
||||
assert.doesNotMatch(locked.logs[0], /supersecret/)
|
||||
// Redaction is log-only: the entry stays under the credential-bearing key.
|
||||
assert.ok(JSON.parse(locked.fileText()!)[CRED_GATEWAY])
|
||||
assert.equal(locked.fileText(), before)
|
||||
})
|
||||
|
||||
test('a parsing failure logs the gateway host and path but not its credentials', () => {
|
||||
const disk = createFakeDisk(JSON.stringify({ [CRED_GATEWAY]: { encoding: 'safeStorage', value: 'bm90LWpzb24=' } }))
|
||||
|
||||
assert.equal(loadNativeTokenSet(CRED_GATEWAY, disk.io), null)
|
||||
assert.match(disk.logs[0], /failed to load stored tokens for https:\/\/gw\.example\.com\/hermes/)
|
||||
assert.doesNotMatch(disk.logs[0], /alice/)
|
||||
assert.doesNotMatch(disk.logs[0], /supersecret/)
|
||||
})
|
||||
|
||||
test('the credential-bearing base URL stays the exact store key', () => {
|
||||
const first = createFakeDisk()
|
||||
|
||||
persistNativeTokenSet(CRED_GATEWAY, TOKENS, first.io)
|
||||
|
||||
assert.deepEqual(Object.keys(JSON.parse(first.fileText()!)), [CRED_GATEWAY])
|
||||
// The original key still round-trips a full set after a restart.
|
||||
assert.deepEqual(loadNativeTokenSet(CRED_GATEWAY, createFakeDisk(first.fileText()).io), TOKENS)
|
||||
// The redacted form is a log string, never a lookup key.
|
||||
assert.equal(loadNativeTokenSet('https://gw.example.com/hermes', createFakeDisk(first.fileText()).io), null)
|
||||
})
|
||||
|
||||
test('an unparseable gateway URL logs a fixed placeholder rather than the raw value', () => {
|
||||
// A space makes this unparseable by URL, so redaction cannot fall back to
|
||||
// echoing the input — that would leak the very credentials it guards.
|
||||
const invalid = 'ht tp://alice:supersecret@gw.example.com'
|
||||
|
||||
const disk = createFakeDisk(JSON.stringify({ [invalid]: { encoding: 'safeStorage', value: 'AAAA' } }), {
|
||||
decrypt: () => ''
|
||||
})
|
||||
|
||||
assert.equal(loadNativeTokenSet(invalid, disk.io), null)
|
||||
assert.match(disk.logs[0], /<invalid gateway URL>/)
|
||||
assert.doesNotMatch(disk.logs[0], /alice/)
|
||||
assert.doesNotMatch(disk.logs[0], /supersecret/)
|
||||
})
|
||||
|
|
@ -0,0 +1,165 @@
|
|||
/**
|
||||
* native-token-store.ts
|
||||
*
|
||||
* The encrypted-at-rest persistence seam for RFC 8252 native OAuth tokens:
|
||||
* NativeTokenSet → JSON → safeStorage blob → store file, and back again on the
|
||||
* next launch.
|
||||
*
|
||||
* Kept standalone (no `import 'electron'`) so the whole restart path unit-tests
|
||||
* with the `electron` vitest project — the same pattern as native-oauth.ts.
|
||||
* main.ts owns the electron-coupled halves and injects them: the safeStorage
|
||||
* encrypt/decrypt pair and the userData store-file read/write.
|
||||
*
|
||||
* The parser direction is the load-bearing detail. What lands on disk is the
|
||||
* *normalized* camelCase NativeTokenSet, so the reload boundary is
|
||||
* parseStoredTokenSet(). Gateway `/auth/native/token` responses are snake_case
|
||||
* and stay with parseTokenResponse(); crossing the two made the decrypted set
|
||||
* throw on every launch, which surfaced as "signed out after restart" (#73271).
|
||||
*/
|
||||
|
||||
import { type NativeTokenSet, parseStoredTokenSet } from './native-oauth'
|
||||
|
||||
/** One encrypted blob as written per gateway base URL. */
|
||||
export interface StoredTokenSecret {
|
||||
encoding?: string
|
||||
value?: string
|
||||
}
|
||||
|
||||
/**
|
||||
* The narrow set of side effects main.ts owns. Everything here is injected so
|
||||
* the store/load round trip can be exercised without an Electron runtime, and
|
||||
* so production keeps using safeStorage unchanged.
|
||||
*/
|
||||
export interface NativeTokenStoreIo {
|
||||
/**
|
||||
* Encrypt one plaintext blob. main.ts passes the strict safeStorage helper,
|
||||
* which THROWS when the OS keychain is unavailable — that must stay loud.
|
||||
* A `null` return is treated as the same authoritative failure: the caller
|
||||
* throws rather than persisting an empty entry over good tokens.
|
||||
*/
|
||||
encrypt: (plaintext: string) => StoredTokenSecret | null
|
||||
/** Decrypt a stored payload; returns '' when it cannot be read. */
|
||||
decrypt: (secret: any) => string
|
||||
/** Raw store-file text. Throws when the file is absent — treated as empty. */
|
||||
readStoreText: () => string
|
||||
/** Persist the store-file text (main.ts writes mode 0600 under userData). */
|
||||
writeStoreText: (text: string) => void
|
||||
rememberLog?: (message: string) => void
|
||||
}
|
||||
|
||||
/**
|
||||
* baseUrl → encrypted payload. A missing, unreadable, or hand-mangled store
|
||||
* reads as empty rather than throwing: a failed *read* falls to the next rung.
|
||||
*
|
||||
* Arrays are rejected alongside every other non-object shape: assigning
|
||||
* store[baseUrl] on an array would set a non-index property, which
|
||||
* JSON.stringify drops on the way back out — the write would look like it
|
||||
* succeeded and the tokens would be gone on the next launch.
|
||||
*/
|
||||
function readStore(io: NativeTokenStoreIo): Record<string, any> {
|
||||
try {
|
||||
const parsed = JSON.parse(io.readStoreText())
|
||||
|
||||
return parsed && typeof parsed === 'object' && !Array.isArray(parsed) ? parsed : {}
|
||||
} catch {
|
||||
return {}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A gateway URL safe to write into a log line.
|
||||
*
|
||||
* normalizeRemoteBaseUrl() strips query, fragment, and trailing slashes but
|
||||
* NOT userinfo, so a configured gateway can legitimately carry
|
||||
* `user:password@` all the way down to this store. Interpolating that into a
|
||||
* failure log would spill the credentials into the desktop log file, so drop
|
||||
* the userinfo and keep only what makes the line useful — scheme, host, port,
|
||||
* path. A value URL cannot parse never falls back to the raw input: echoing it
|
||||
* is the exact leak this guards against.
|
||||
*/
|
||||
function redactGatewayUrl(baseUrl: string): string {
|
||||
try {
|
||||
const parsed = new URL(baseUrl)
|
||||
|
||||
parsed.username = ''
|
||||
parsed.password = ''
|
||||
|
||||
// `host` already carries a non-default port.
|
||||
return `${parsed.protocol}//${parsed.host}${parsed.pathname}`
|
||||
} catch {
|
||||
return '<invalid gateway URL>'
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Write (or, with `tokens === null`, drop) one gateway's token set, merging
|
||||
* into whatever other gateways are already stored.
|
||||
*/
|
||||
export function persistNativeTokenSet(baseUrl: string, tokens: NativeTokenSet | null, io: NativeTokenStoreIo): void {
|
||||
const store = readStore(io)
|
||||
|
||||
if (tokens) {
|
||||
// Encrypt the whole set as one blob so the refresh token never lands in
|
||||
// plaintext on disk. Deliberately outside the try below: an unusable
|
||||
// keychain is an authoritative write failure and must surface to the
|
||||
// caller, not be logged away as if the tokens were saved.
|
||||
const secret = io.encrypt(JSON.stringify(tokens))
|
||||
|
||||
if (!secret) {
|
||||
// A null blob is the same failure as a throw, only quieter. Storing it
|
||||
// would replace a good entry with nothing: the write would report
|
||||
// success, the next launch would show signed out, and the refresh token
|
||||
// would be unrecoverable. Fail before touching the store.
|
||||
throw new Error('Secure token storage returned no encrypted payload; refusing to overwrite stored native tokens.')
|
||||
}
|
||||
|
||||
store[baseUrl] = secret
|
||||
} else {
|
||||
delete store[baseUrl]
|
||||
}
|
||||
|
||||
try {
|
||||
io.writeStoreText(JSON.stringify(store))
|
||||
} catch (error) {
|
||||
const detail = error instanceof Error ? error.message : String(error)
|
||||
|
||||
io.rememberLog?.(`[native-oauth] failed to persist tokens: ${detail}`)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconstruct a gateway's token set from the stored encrypted payload. Returns
|
||||
* null when nothing is stored, when the blob cannot be decrypted, or when it
|
||||
* does not parse — never a partially-populated set.
|
||||
*/
|
||||
export function loadNativeTokenSet(baseUrl: string, io: NativeTokenStoreIo): NativeTokenSet | null {
|
||||
// The UNREDACTED url is the store key — redaction is for logs only.
|
||||
const secret = readStore(io)[baseUrl]
|
||||
|
||||
if (!secret) {
|
||||
return null
|
||||
}
|
||||
|
||||
try {
|
||||
const plaintext = io.decrypt(secret)
|
||||
|
||||
if (!plaintext) {
|
||||
// A keychain that is merely locked/unavailable right now must not cost
|
||||
// the user their refresh token — leave the entry for the next attempt.
|
||||
io.rememberLog?.(
|
||||
`[native-oauth] failed to decrypt stored tokens for ${redactGatewayUrl(baseUrl)}; keeping stored entry for retry`
|
||||
)
|
||||
|
||||
return null
|
||||
}
|
||||
|
||||
// Stored blobs are normalized camelCase sets, never raw gateway responses.
|
||||
return parseStoredTokenSet(JSON.parse(plaintext))
|
||||
} catch (error) {
|
||||
const detail = error instanceof Error ? error.message : String(error)
|
||||
|
||||
io.rememberLog?.(`[native-oauth] failed to load stored tokens for ${redactGatewayUrl(baseUrl)}: ${detail}`)
|
||||
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
yaleman
|
||||
# PR #49608 salvage (JWKS Accept + User-Agent headers)
|
||||
|
|
@ -2149,11 +2149,11 @@ def _exempt_explicit_platform_native(
|
|||
#: Landing late — or leaving an entry here for a second release — converts a
|
||||
#: back-fill into a stuck checkbox.
|
||||
#:
|
||||
#: Not gated on a Nous subscription here: the six ``bfl_flux3_*`` tools carry
|
||||
#: ``check_fn=check_bfl_requirements`` (logged in AND paid), so an enabled
|
||||
#: toolset still ships zero schemas to a user without paid portal access — the
|
||||
#: same split Home Assistant uses. Probing the portal from this path would put
|
||||
#: a network call on every CLI start, gateway session and cron tick.
|
||||
#: Not gated on a Nous sign-in here: the six ``bfl_flux3_*`` tools carry
|
||||
#: ``check_fn=check_bfl_requirements``, so an enabled toolset still ships zero
|
||||
#: schemas to a user with no Nous credential — the same split Home Assistant
|
||||
#: uses. Probing the portal from this path would put a network call on every
|
||||
#: CLI start, gateway session and cron tick.
|
||||
_RECENTLY_SHIPPED_TOOLSETS = frozenset({"bfl"})
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -62,7 +62,16 @@ def _get_jwk_client(jwks_url: str) -> Any:
|
|||
if client is None:
|
||||
from jwt import PyJWKClient
|
||||
|
||||
client = PyJWKClient(jwks_url)
|
||||
# Explicit Accept + User-Agent so the JWKS fetch isn't blocked by the
|
||||
# NAS portal's WAF, which 403s the default Python-urllib fingerprint
|
||||
# (same fix as the dashboard-auth nous/self_hosted providers).
|
||||
client = PyJWKClient(
|
||||
jwks_url,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "HermesAgent/1.0",
|
||||
},
|
||||
)
|
||||
_JWK_CLIENTS[jwks_url] = client
|
||||
return client
|
||||
|
||||
|
|
|
|||
|
|
@ -420,6 +420,10 @@ class NousDashboardAuthProvider(DashboardAuthProvider):
|
|||
self._jwks_url,
|
||||
cache_keys=True,
|
||||
lifespan=_JWKS_CACHE_SECONDS,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "HermesAgent/1.0",
|
||||
},
|
||||
)
|
||||
return self._jwks_client
|
||||
|
||||
|
|
|
|||
|
|
@ -601,6 +601,10 @@ class SelfHostedOIDCProvider(DashboardAuthProvider):
|
|||
disco["jwks_uri"],
|
||||
cache_keys=True,
|
||||
lifespan=_JWKS_CACHE_SECONDS,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "HermesAgent/1.0",
|
||||
},
|
||||
)
|
||||
return self._jwks_client
|
||||
|
||||
|
|
|
|||
|
|
@ -218,6 +218,129 @@ class TestResolveVisionMainFirst:
|
|||
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _stub_nous_portal(seen: dict):
|
||||
"""Stub the Nous network boundary, keeping the resolution chain real.
|
||||
|
||||
Returns a ``_try_nous`` replacement that answers with the Portal's
|
||||
tier-aware slots: a vision model for ``vision=True``, the text chat
|
||||
default otherwise.
|
||||
"""
|
||||
nous_client = MagicMock()
|
||||
nous_client.api_key = "jwt-test"
|
||||
nous_client.base_url = "https://inference-api.nousresearch.com/v1"
|
||||
|
||||
def fake_try_nous(vision=False):
|
||||
seen["vision"] = vision
|
||||
return nous_client, (
|
||||
"stepfun/step-3.7-flash:free" if vision else "tencent/hy3:free"
|
||||
)
|
||||
|
||||
return nous_client, fake_try_nous
|
||||
|
||||
def test_nous_main_vision_uses_portal_pick_not_text_chat_model(self):
|
||||
"""Nous main → vision runs the Portal's vision slot, not the chat model.
|
||||
|
||||
A Nous chat default is routinely text-only (e.g. a ``:free`` chat SKU).
|
||||
Letting it reach the vision lane means the image goes to a model that
|
||||
cannot accept one and the Portal 404s. Only the Nous network boundary
|
||||
is stubbed — the strict vision backend, the provider router, and its
|
||||
missing-model pre-fill all run for real, because that pre-fill is where
|
||||
the chat model used to clobber the Portal's pick.
|
||||
"""
|
||||
seen: dict = {}
|
||||
nous_client, fake_try_nous = self._stub_nous_portal(seen)
|
||||
|
||||
with patch(
|
||||
"agent.auxiliary_client._read_main_provider", return_value="nous",
|
||||
), patch(
|
||||
"agent.auxiliary_client._read_main_model", return_value="tencent/hy3:free",
|
||||
), patch(
|
||||
"agent.auxiliary_client._resolve_task_provider_model",
|
||||
return_value=("auto", None, None, None, None),
|
||||
), patch(
|
||||
"agent.auxiliary_client._try_nous", side_effect=fake_try_nous,
|
||||
):
|
||||
from agent.auxiliary_client import resolve_vision_provider_client
|
||||
|
||||
provider, client, model = resolve_vision_provider_client()
|
||||
|
||||
assert provider == "nous"
|
||||
assert client is nous_client
|
||||
assert seen["vision"] is True
|
||||
assert model == "stepfun/step-3.7-flash:free"
|
||||
|
||||
def test_nous_main_vision_honours_explicit_vision_model(self):
|
||||
"""An explicit auxiliary.vision.model still overrides the Portal pick."""
|
||||
seen: dict = {}
|
||||
_nous_client, fake_try_nous = self._stub_nous_portal(seen)
|
||||
|
||||
with patch(
|
||||
"agent.auxiliary_client._read_main_provider", return_value="nous",
|
||||
), patch(
|
||||
"agent.auxiliary_client._read_main_model", return_value="tencent/hy3:free",
|
||||
), patch(
|
||||
"agent.auxiliary_client._resolve_task_provider_model",
|
||||
return_value=("auto", "qwen/qwen3-vl-8b-instruct", None, None, None),
|
||||
), patch(
|
||||
"agent.auxiliary_client._try_nous", side_effect=fake_try_nous,
|
||||
):
|
||||
from agent.auxiliary_client import resolve_vision_provider_client
|
||||
|
||||
provider, _client, model = resolve_vision_provider_client()
|
||||
|
||||
assert provider == "nous"
|
||||
assert model == "qwen/qwen3-vl-8b-instruct"
|
||||
|
||||
def test_nous_explicit_vision_provider_also_skips_chat_model(self):
|
||||
"""``auxiliary.vision.provider: nous`` takes the same Portal pick.
|
||||
|
||||
The explicit-provider branch reaches the strict vision backend with no
|
||||
model too, so it has to resolve the same way the auto branch does.
|
||||
"""
|
||||
seen: dict = {}
|
||||
nous_client, fake_try_nous = self._stub_nous_portal(seen)
|
||||
|
||||
with patch(
|
||||
"agent.auxiliary_client._read_main_provider", return_value="nous",
|
||||
), patch(
|
||||
"agent.auxiliary_client._read_main_model", return_value="tencent/hy3:free",
|
||||
), patch(
|
||||
"agent.auxiliary_client._resolve_task_provider_model",
|
||||
return_value=("nous", None, None, None, None),
|
||||
), patch(
|
||||
"agent.auxiliary_client._try_nous", side_effect=fake_try_nous,
|
||||
):
|
||||
from agent.auxiliary_client import resolve_vision_provider_client
|
||||
|
||||
provider, client, model = resolve_vision_provider_client()
|
||||
|
||||
assert provider == "nous"
|
||||
assert client is nous_client
|
||||
assert model == "stepfun/step-3.7-flash:free"
|
||||
|
||||
def test_nous_text_aux_still_uses_main_chat_model(self):
|
||||
"""The vision carve-out must not leak into text aux resolution.
|
||||
|
||||
Text auxiliary work on a Nous main deliberately keeps the user's chat
|
||||
model rather than dropping to the Portal's cheap default.
|
||||
"""
|
||||
seen: dict = {}
|
||||
_nous_client, fake_try_nous = self._stub_nous_portal(seen)
|
||||
|
||||
with patch(
|
||||
"agent.auxiliary_client._read_main_provider", return_value="nous",
|
||||
), patch(
|
||||
"agent.auxiliary_client._read_main_model", return_value="tencent/hy3:free",
|
||||
), patch(
|
||||
"agent.auxiliary_client._try_nous", side_effect=fake_try_nous,
|
||||
):
|
||||
from agent.auxiliary_client import resolve_provider_client
|
||||
|
||||
_client, model = resolve_provider_client("nous")
|
||||
|
||||
assert model == "tencent/hy3:free"
|
||||
|
||||
def test_copilot_vision_sets_vision_header(self, monkeypatch):
|
||||
"""Copilot vision requests include the header required for vision routing."""
|
||||
monkeypatch.setenv("COPILOT_GITHUB_TOKEN", "ghu_test-token")
|
||||
|
|
|
|||
|
|
@ -528,6 +528,22 @@ class TestVerifySession:
|
|||
_patched_jwks(p, rsa_keypair)
|
||||
return p
|
||||
|
||||
def test_jwks_client_sends_explicit_http_headers(self, provider):
|
||||
"""Constructor-contract regression: the JWKS fetch must send an
|
||||
explicit Accept + User-Agent so it isn't blocked by the Portal WAF
|
||||
(same fix as the self_hosted provider)."""
|
||||
provider._jwks_client = None
|
||||
with patch("jwt.PyJWKClient") as client_cls:
|
||||
provider._get_jwks_client()
|
||||
client_cls.assert_called_once_with(
|
||||
provider._jwks_url,
|
||||
cache_keys=True,
|
||||
lifespan=nous_plugin._JWKS_CACHE_SECONDS,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "HermesAgent/1.0",
|
||||
},
|
||||
)
|
||||
|
||||
def test_expired_token_returns_none(self, provider, rsa_keypair):
|
||||
token = _mint_token(rsa_keypair, ttl_seconds=-1)
|
||||
|
|
|
|||
|
|
@ -571,6 +571,26 @@ class TestVerifySession:
|
|||
with pytest.raises(ProviderError, match="JWKS"):
|
||||
provider.verify_session(access_token=token)
|
||||
|
||||
def test_jwks_client_sends_explicit_http_headers(self):
|
||||
provider = oidc_plugin.SelfHostedOIDCProvider(
|
||||
issuer=_ISSUER, client_id=_CLIENT_ID
|
||||
)
|
||||
provider._discovery = dict(_DISCOVERY_DOC)
|
||||
provider._discovery_fetched_at = time.time()
|
||||
|
||||
with patch("jwt.PyJWKClient") as client_cls:
|
||||
provider._get_jwks_client()
|
||||
|
||||
client_cls.assert_called_once_with(
|
||||
_DISCOVERY_DOC["jwks_uri"],
|
||||
cache_keys=True,
|
||||
lifespan=oidc_plugin._JWKS_CACHE_SECONDS,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "HermesAgent/1.0",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# refresh_session + revoke_session
|
||||
|
|
|
|||
|
|
@ -144,7 +144,7 @@ def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch):
|
|||
key = pub
|
||||
|
||||
class FakeJWKClient:
|
||||
def __init__(self, url):
|
||||
def __init__(self, url, **kwargs):
|
||||
assert url == "https://portal.nousresearch.com/.well-known/jwks.json"
|
||||
|
||||
def get_signing_key_from_jwt(self, tok):
|
||||
|
|
@ -161,6 +161,32 @@ def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch):
|
|||
assert claims is not None and claims["purpose"] == "cron_fire"
|
||||
|
||||
|
||||
def test_jwks_client_sends_explicit_http_headers(monkeypatch):
|
||||
"""Constructor-contract regression: the JWKS fetch must send an explicit
|
||||
Accept + User-Agent so it isn't blocked by the NAS portal WAF (same fix as
|
||||
the dashboard-auth nous/self_hosted providers)."""
|
||||
from plugins.cron_providers.chronos import verify as verify_mod
|
||||
|
||||
captured = {}
|
||||
|
||||
class FakeJWKClient:
|
||||
def __init__(self, url, **kwargs):
|
||||
captured["url"] = url
|
||||
captured["kwargs"] = kwargs
|
||||
|
||||
monkeypatch.setattr("jwt.PyJWKClient", FakeJWKClient)
|
||||
monkeypatch.setattr(verify_mod, "_JWK_CLIENTS", {})
|
||||
|
||||
url = "https://portal.nousresearch.com/.well-known/jwks.json"
|
||||
verify_mod._get_jwk_client(url)
|
||||
|
||||
assert captured["url"] == url
|
||||
assert captured["kwargs"].get("headers") == {
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "HermesAgent/1.0",
|
||||
}
|
||||
|
||||
|
||||
def test_get_fire_verifier_returns_nas_verifier():
|
||||
from plugins.cron_providers.chronos.verify import get_fire_verifier, verify_nas_fire_token
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
|
@ -15,6 +16,11 @@ GATEWAY = "https://tool-gateway.example.com"
|
|||
BASE_URL = f"{GATEWAY}/api/bfl"
|
||||
UPLOAD_PATH = "/api/uploads/bfl"
|
||||
|
||||
# The shipped pacing, read before the autouse fixture below rewrites it to
|
||||
# something the tests can spend in an instant.
|
||||
_DEFAULT_POLL_BUDGET_SECONDS = flux3._POLL_BUDGET_SECONDS
|
||||
_DEFAULT_CALL_BACKSTOP_SECONDS = flux3._CALL_BACKSTOP_SECONDS
|
||||
|
||||
_PNG = base64.b64decode(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
|
||||
)
|
||||
|
|
@ -127,6 +133,25 @@ def _record_sleep(sink):
|
|||
return _sleep
|
||||
|
||||
|
||||
def _stepped_clock(look_seconds):
|
||||
"""A monotonic clock on which every look appears to take `look_seconds`.
|
||||
|
||||
The poll loop reads the clock twice per look — once before the request and
|
||||
once after — so advancing on every second read charges a look exactly that
|
||||
much budget without spending any real time. Substituted for the module's
|
||||
whole ``time`` reference rather than patching ``time.monotonic`` globally,
|
||||
which would hand the same jumping clock to the event loop underneath.
|
||||
"""
|
||||
reads = {"n": 0}
|
||||
|
||||
def _monotonic():
|
||||
value = (reads["n"] // 2) * look_seconds
|
||||
reads["n"] += 1
|
||||
return value
|
||||
|
||||
return _monotonic
|
||||
|
||||
|
||||
def _call(handler, args, response, headers=None):
|
||||
"""Invoke a handler with the transport stubbed; returns (parsed, requests)."""
|
||||
sink = []
|
||||
|
|
@ -146,25 +171,58 @@ class TestGating:
|
|||
with patch.object(flux3, "managed_vendor_endpoints", return_value=None):
|
||||
assert flux3.check_bfl_requirements() is False
|
||||
|
||||
def test_hidden_without_paid_service_access(self):
|
||||
# The free tool pool does not fund BFL, so a pool-only user must never
|
||||
# see the tools rather than see them and be refused.
|
||||
account = SimpleNamespace(logged_in=True, paid_service_access=False, tool_gateway_entitled=True)
|
||||
with patch("hermes_cli.nous_account.get_nous_portal_account_info", return_value=account):
|
||||
assert flux3.check_bfl_requirements() is False
|
||||
|
||||
def test_hidden_when_logged_out(self):
|
||||
account = SimpleNamespace(logged_in=False, paid_service_access=False)
|
||||
with patch("hermes_cli.nous_account.get_nous_portal_account_info", return_value=account):
|
||||
assert flux3.check_bfl_requirements() is False
|
||||
|
||||
def test_visible_for_a_paid_portal_account(self):
|
||||
account = SimpleNamespace(logged_in=True, paid_service_access=True)
|
||||
with patch("hermes_cli.nous_account.get_nous_portal_account_info", return_value=account):
|
||||
def test_visible_to_any_signed_in_account_whatever_its_entitlement(self):
|
||||
# Entitlement is the gateway's ruling, and it states its reason in a
|
||||
# refusal the model can act on. Deciding it here as well could only
|
||||
# hide the tools from someone the server would have served, so the
|
||||
# portal's entitlement view must not be consulted at all.
|
||||
with patch.object(flux3, "peek_nous_access_token", return_value="nous-token"), \
|
||||
patch(
|
||||
"hermes_cli.nous_account.get_nous_portal_account_info",
|
||||
side_effect=AssertionError("entitlement must not gate visibility"),
|
||||
):
|
||||
assert flux3.check_bfl_requirements() is True
|
||||
|
||||
def test_fails_closed_when_the_account_probe_raises(self):
|
||||
with patch("hermes_cli.nous_account.get_nous_portal_account_info", side_effect=RuntimeError("portal down")):
|
||||
def test_hidden_without_a_nous_credential(self):
|
||||
# The gateway takes a Nous bearer and nothing else, so with no token
|
||||
# every call could only ever answer "sign in" — six schemas on every
|
||||
# API call for something that cannot work.
|
||||
with patch.object(flux3, "peek_nous_access_token", return_value=None):
|
||||
assert flux3.check_bfl_requirements() is False
|
||||
|
||||
def test_a_profile_sees_a_credential_held_at_the_global_root(self, tmp_path, monkeypatch):
|
||||
# A profile that was never logged into separately still calls the
|
||||
# gateway with the root login, because the transport's refresh path
|
||||
# reads that same global fallback. Probing only the profile's own store
|
||||
# would hide the tools from someone whose calls would have worked.
|
||||
#
|
||||
# Exercised through the real auth store rather than a stub: the
|
||||
# fallback is the whole point of the test, and it lives in
|
||||
# hermes_cli.auth, not here.
|
||||
monkeypatch.delenv("TOOL_GATEWAY_USER_TOKEN", raising=False)
|
||||
root = tmp_path / "root"
|
||||
(root / "profiles" / "work").mkdir(parents=True)
|
||||
(root / "auth.json").write_text(
|
||||
json.dumps({"version": 1, "providers": {"nous": {"access_token": "root-token"}}}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setenv("HERMES_HOME", str(root / "profiles" / "work"))
|
||||
|
||||
# The profile's own store is empty, so this passes only via the
|
||||
# global-root fallback — without which the tools would be hidden.
|
||||
assert flux3.peek_nous_access_token() is None
|
||||
assert flux3.check_bfl_requirements() is True
|
||||
|
||||
def test_the_credential_probe_never_forces_a_token_refresh(self, monkeypatch):
|
||||
# check_fn runs on every CLI start, gateway session and cron tick, so
|
||||
# it reads a cached credential rather than sitting on a synchronous
|
||||
# OAuth refresh.
|
||||
monkeypatch.setenv("TOOL_GATEWAY_USER_TOKEN", "nous-token")
|
||||
with patch.object(flux3, "read_nous_access_token", side_effect=AssertionError("refreshed")):
|
||||
assert flux3.check_bfl_requirements() is True
|
||||
|
||||
def test_fails_closed_when_the_credential_probe_raises(self):
|
||||
with patch.object(flux3, "peek_nous_access_token", side_effect=RuntimeError("auth store unreadable")):
|
||||
assert flux3.check_bfl_requirements() is False
|
||||
|
||||
|
||||
|
|
@ -263,8 +321,20 @@ class TestSubmitTransport:
|
|||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_real_poll_wait(monkeypatch):
|
||||
"""Keep the in-call wait out of the test clock."""
|
||||
monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 0)
|
||||
"""Pace the in-call poll loop off the test clock, at two looks per call.
|
||||
|
||||
The handler counts its budget rather than reading a clock, so a gap and a
|
||||
budget in a fixed ratio give a deterministic number of looks with no fake
|
||||
clock: a budget of two gaps spends one wait and takes two looks, which is
|
||||
the smallest loop that can still show a job finishing between looks.
|
||||
"""
|
||||
monkeypatch.setattr(flux3, "_POLL_GAP_SECONDS", 1.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 2.0)
|
||||
|
||||
async def _instant(_seconds):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(flux3.asyncio, "sleep", _instant)
|
||||
|
||||
|
||||
class TestPollTransport:
|
||||
|
|
@ -282,7 +352,8 @@ class TestPollTransport:
|
|||
def test_a_running_job_is_waited_out_inside_the_call(self, monkeypatch):
|
||||
# A model has no clock, so telling it to pause produced a burst of polls
|
||||
# instead of a paced one. The wait lives here where it cannot be skipped.
|
||||
monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 45.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_GAP_SECONDS", 45.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 90.0)
|
||||
running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
|
||||
|
||||
slept = []
|
||||
|
|
@ -293,11 +364,26 @@ class TestPollTransport:
|
|||
assert sum(slept) == 45.0
|
||||
assert parsed["details"]["status"] == "Generating"
|
||||
|
||||
def test_the_loop_keeps_looking_until_its_budget_is_spent(self, monkeypatch):
|
||||
# The job endpoint answers at once, so a call that looked a fixed twice
|
||||
# spent almost none of the time it was allowed and handed control back
|
||||
# to the model four or five times per generation. One call now covers
|
||||
# the whole budget, and the model decides to keep waiting once.
|
||||
monkeypatch.setattr(flux3, "_POLL_GAP_SECONDS", 10.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 50.0)
|
||||
running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
|
||||
|
||||
parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, running)
|
||||
|
||||
assert len(requests) == 5, "five looks spaced by four ten-second gaps"
|
||||
assert parsed["details"]["status"] == "Generating"
|
||||
|
||||
def test_the_wait_is_answerable_to_a_stop(self, monkeypatch):
|
||||
# Nothing outside the tool can end a call that has already started —
|
||||
# the executor only checks for an interrupt between tools — so /stop
|
||||
# has to land inside the wait rather than at the end of it.
|
||||
monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 45.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_GAP_SECONDS", 45.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 200.0)
|
||||
from tools import interrupt as interrupt_module
|
||||
|
||||
running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
|
||||
|
|
@ -327,9 +413,9 @@ class TestPollTransport:
|
|||
assert len(requests) == 2
|
||||
assert parsed["result"] == "That job failed."
|
||||
|
||||
def test_a_refusal_is_returned_immediately_rather_than_waited_on(self):
|
||||
# A 429 carries its own retry guidance; sleeping on it would only delay
|
||||
# showing the model what to do, and spend the poll budget twice.
|
||||
def test_a_refusal_without_a_stated_wait_is_returned_immediately(self):
|
||||
# A dead job or a bad id has nothing to wait for; sleeping on it would
|
||||
# only delay showing the model what to do.
|
||||
response = _FakeResponse(429, {"error": {"message": "Too many polls. Wait 30 seconds."}})
|
||||
|
||||
parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, response)
|
||||
|
|
@ -337,6 +423,202 @@ class TestPollTransport:
|
|||
assert len(requests) == 1
|
||||
assert "Too many polls" in parsed["error"]
|
||||
|
||||
def test_a_throttle_is_waited_out_inside_the_call(self, monkeypatch):
|
||||
# Handing a throttle back ends the call, and the model it lands on has
|
||||
# no clock — it asks again at once, tightening the loop that tripped the
|
||||
# limit. The stated wait is taken here instead.
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 100.0)
|
||||
throttled = _FakeResponse(
|
||||
429,
|
||||
{"error": {"message": "Too many polls.", "details": {"retryAfterSeconds": 30}}},
|
||||
)
|
||||
done = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {}, "guidance": "Done."})
|
||||
|
||||
slept = []
|
||||
with patch.object(flux3.asyncio, "sleep", new=_record_sleep(slept)):
|
||||
parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, [throttled, done])
|
||||
|
||||
assert len(requests) == 2, "the loop survives a throttle"
|
||||
assert sum(slept) == 30.0, "and waits exactly as long as it was asked to"
|
||||
assert parsed["result"] == "Done."
|
||||
|
||||
def test_a_throttle_never_polls_faster_than_the_loop_s_own_cadence(self, monkeypatch):
|
||||
# The gateway's number is a floor on politeness, not a licence to
|
||||
# hammer: a small or malformed-but-positive wait must not turn the loop
|
||||
# into a tight one against an endpoint that just asked us to slow down.
|
||||
monkeypatch.setattr(flux3, "_POLL_GAP_SECONDS", 10.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 100.0)
|
||||
throttled = _FakeResponse(
|
||||
429,
|
||||
{"error": {"message": "Slow down.", "details": {"retryAfterSeconds": 0.001}}},
|
||||
)
|
||||
done = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {}, "guidance": "Done."})
|
||||
|
||||
slept = []
|
||||
with patch.object(flux3.asyncio, "sleep", new=_record_sleep(slept)):
|
||||
_parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, [throttled, done])
|
||||
|
||||
assert len(requests) == 2
|
||||
assert sum(slept) == 10.0, "the loop's own gap, not the sliver it was offered"
|
||||
|
||||
def test_a_slow_poll_spends_the_budget_it_actually_took(self, monkeypatch):
|
||||
# Counting only the waits would let a gateway that answers slowly run
|
||||
# the call far past its budget, leaving the backstop to do the work the
|
||||
# budget is supposed to do. A look costs what it takes.
|
||||
monkeypatch.setattr(flux3, "_POLL_GAP_SECONDS", 1.0)
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 30.0)
|
||||
running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
|
||||
|
||||
monkeypatch.setattr(flux3, "time", SimpleNamespace(monotonic=_stepped_clock(20.0)))
|
||||
|
||||
_parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, running)
|
||||
|
||||
assert len(requests) == 2, "two twenty-second looks exhaust a thirty-second budget"
|
||||
|
||||
def test_a_poll_outwaits_the_gateways_own_poll_budget(self):
|
||||
# The gateway bounds one status read at 45s across its retries and
|
||||
# regional redirect hops. Giving up before it does turns a slow but
|
||||
# healthy poll into a transport error, and an error ends the loop.
|
||||
assert flux3._POLL_READ_TIMEOUT_SECONDS > 45.0
|
||||
|
||||
def test_a_blip_costs_a_look_rather_than_the_rest_of_the_call(self, monkeypatch):
|
||||
# The generation runs upstream and is unaffected by our failing to ask
|
||||
# about it, so one unreachable moment must not throw away the minutes
|
||||
# of budget left. Returning it would end the call on an error the model
|
||||
# can only answer by polling again at once — the burst the loop exists
|
||||
# to prevent.
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 100.0)
|
||||
done = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {}, "guidance": "Done."})
|
||||
|
||||
parsed, requests = _call(
|
||||
flux3._handle_get_result,
|
||||
{"id": "bfl_job_1"},
|
||||
[RuntimeError("connection reset"), done],
|
||||
)
|
||||
|
||||
assert len(requests) == 2, "the loop looked again after the blip"
|
||||
assert parsed["result"] == "Done."
|
||||
|
||||
def test_a_gateway_answering_in_html_counts_as_unreachable(self, monkeypatch):
|
||||
# What a 502 from an edge in front of the gateway looks like from here:
|
||||
# a status code and a page, with no error the model could act on. That
|
||||
# is an absent answer, not a refusal, so it is retried like one.
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 100.0)
|
||||
done = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {}, "guidance": "Done."})
|
||||
|
||||
_parsed, requests = _call(
|
||||
flux3._handle_get_result,
|
||||
{"id": "bfl_job_1"},
|
||||
[_FakeResponse(502, None, text="<html>bad gateway</html>"), done],
|
||||
)
|
||||
|
||||
assert len(requests) == 2
|
||||
|
||||
def test_a_gateway_that_stays_down_is_reported_rather_than_retried_out(self, monkeypatch):
|
||||
# Tolerance is for blips. A gateway that is genuinely down has to reach
|
||||
# the model promptly, not after minutes of a budget spent on a host
|
||||
# that is not answering.
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 1000.0)
|
||||
|
||||
parsed, requests = _call(
|
||||
flux3._handle_get_result,
|
||||
{"id": "bfl_job_1"},
|
||||
RuntimeError("connection reset"),
|
||||
)
|
||||
|
||||
assert len(requests) == flux3._MAX_CONSECUTIVE_TRANSPORT_ERRORS
|
||||
assert "Could not reach" in parsed["error"]
|
||||
|
||||
def test_the_tolerance_counts_consecutive_failures_only(self, monkeypatch):
|
||||
# A flaky gateway that answers every other look is still usable, so the
|
||||
# count has to reset on an answer rather than accumulate over the call.
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 1000.0)
|
||||
blip = RuntimeError("connection reset")
|
||||
running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
|
||||
done = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {}, "guidance": "Done."})
|
||||
|
||||
parsed, requests = _call(
|
||||
flux3._handle_get_result,
|
||||
{"id": "bfl_job_1"},
|
||||
[blip, blip, running, blip, blip, done],
|
||||
)
|
||||
|
||||
assert len(requests) == 6, "four blips, never three in a row, so the call survives"
|
||||
assert parsed["result"] == "Done."
|
||||
|
||||
def test_a_throttle_longer_than_the_budget_is_handed_back(self, monkeypatch):
|
||||
# A five-minute generation cooldown cannot be absorbed inside one call,
|
||||
# so the model gets the message and the number rather than a call that
|
||||
# sits out a wait it can never finish.
|
||||
monkeypatch.setattr(flux3, "_POLL_BUDGET_SECONDS", 100.0)
|
||||
response = _FakeResponse(
|
||||
429,
|
||||
{"error": {"message": "Wait 210 seconds.", "details": {"retryAfterSeconds": 210}}},
|
||||
)
|
||||
|
||||
parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, response)
|
||||
|
||||
assert len(requests) == 1
|
||||
assert parsed["error"] == "Wait 210 seconds."
|
||||
|
||||
def test_the_backstop_answers_rather_than_letting_the_bridge_kill_the_call(self, monkeypatch):
|
||||
# model_tools' async bridge abandons a tool at 300s and reports it as a
|
||||
# bare "TimeoutError:" — no job id, nothing to say the generation is
|
||||
# still alive and one poll away. Whatever stalls inside, the model is
|
||||
# answered from here first.
|
||||
monkeypatch.setattr(flux3, "_CALL_BACKSTOP_SECONDS", 0.01)
|
||||
|
||||
async def _never_finishes(*_args, **_kwargs):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
monkeypatch.setattr(flux3, "_poll_until_done", _never_finishes)
|
||||
|
||||
parsed, _requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, _FakeResponse(200, {}))
|
||||
|
||||
assert parsed["details"] == {"id": "bfl_job_1", "status": "Generating"}
|
||||
assert "bfl_flux3_get_result" in parsed["result"]
|
||||
assert "bfl_job_1" in parsed["result"]
|
||||
|
||||
def test_a_poll_does_not_inherit_the_submit_read_timeout(self):
|
||||
# A status GET answers at once. Left on the submit path's patience, one
|
||||
# hung poll would spend the whole call's budget by itself — while submit,
|
||||
# which really does sit behind an upload and an upstream call, keeps it.
|
||||
import httpx
|
||||
|
||||
timeouts = []
|
||||
sink = []
|
||||
settled = _FakeResponse(200, {"id": "j", "status": "Error", "guidance": "over"})
|
||||
|
||||
def _client(**kwargs):
|
||||
timeouts.append(kwargs.get("timeout"))
|
||||
return _FakeClient(settled, sink)
|
||||
|
||||
with patch.object(flux3, "managed_gateway_auth_headers", return_value={"Authorization": "Bearer t"}), \
|
||||
patch.object(httpx, "AsyncClient", _client):
|
||||
_run(flux3._handle_get_result({"id": "j"}))
|
||||
_run(flux3._handle_text_to_video({"prompt": "a"}))
|
||||
|
||||
poll_timeout, submit_timeout = timeouts
|
||||
assert poll_timeout.read == flux3._POLL_READ_TIMEOUT_SECONDS
|
||||
assert submit_timeout.read == flux3._TRANSPORT_READ_TIMEOUT_SECONDS
|
||||
assert poll_timeout.read < submit_timeout.read
|
||||
|
||||
def test_the_pacing_stays_clear_of_the_agents_per_tool_ceiling(self):
|
||||
# The whole point of the two bounds: a clip finishing on the last look
|
||||
# still has to be downloaded inside the backstop, and the backstop has
|
||||
# to answer before model_tools' async bridge abandons the tool at 300s.
|
||||
assert _DEFAULT_POLL_BUDGET_SECONDS < _DEFAULT_CALL_BACKSTOP_SECONDS
|
||||
assert _DEFAULT_CALL_BACKSTOP_SECONDS < 300.0
|
||||
|
||||
def test_download_timeout_never_outlives_the_backstop(self):
|
||||
# Near the end of the call, remaining budget after grace is a few
|
||||
# seconds. Clamping that up used to schedule a download the outer
|
||||
# wait_for then cancelled, answering "still generating" for a Ready job.
|
||||
started = time.monotonic() - (
|
||||
flux3._CALL_BACKSTOP_SECONDS - flux3._DOWNLOAD_GRACE_SECONDS - 2.0
|
||||
)
|
||||
assert flux3._download_read_timeout(started) <= 2.0 + 0.5 # clock noise only
|
||||
|
||||
def test_ready_saves_the_clip_and_never_returns_the_signed_url(self, tmp_path):
|
||||
# The signed URL is a bearer credential for the clip and it used to be
|
||||
# re-keyed into a shell command by hand, dropping characters. Neither
|
||||
|
|
|
|||
|
|
@ -110,12 +110,11 @@ def test_managed_vendor_endpoints_pin_the_deployed_gateway_url():
|
|||
typo'd pseudo-vendor to a non-existent host while every other test stubbed
|
||||
it): default builder, real deployed host, pinned vendor path.
|
||||
"""
|
||||
with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True), \
|
||||
patch.dict(
|
||||
os.environ,
|
||||
{"TOOL_GATEWAY_DOMAIN": "nousresearch.com", "TOOL_GATEWAY_SCHEME": "https"},
|
||||
clear=False,
|
||||
):
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"TOOL_GATEWAY_DOMAIN": "nousresearch.com", "TOOL_GATEWAY_SCHEME": "https"},
|
||||
clear=False,
|
||||
):
|
||||
os.environ.pop("TOOL_GATEWAY_URL", None)
|
||||
endpoints = managed_tool_gateway.managed_vendor_endpoints("bfl")
|
||||
|
||||
|
|
@ -126,8 +125,31 @@ def test_managed_vendor_endpoints_pin_the_deployed_gateway_url():
|
|||
}
|
||||
|
||||
|
||||
def test_managed_vendor_endpoints_unreachable_when_managed_tools_disabled():
|
||||
with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=False):
|
||||
def test_managed_vendor_endpoints_do_not_consult_entitlement():
|
||||
"""Address resolution, not a policy decision.
|
||||
|
||||
What an account may spend is the gateway's ruling, stated in its refusals.
|
||||
Guessing at it here would hide the address from a caller the server would
|
||||
have served, so entitlement must not be read on this path at all.
|
||||
"""
|
||||
with patch.dict(os.environ, {"TOOL_GATEWAY_DOMAIN": "nousresearch.com"}, clear=False), \
|
||||
patch.object(
|
||||
managed_tool_gateway,
|
||||
"managed_nous_tools_enabled",
|
||||
side_effect=AssertionError("entitlement must not gate address resolution"),
|
||||
):
|
||||
os.environ.pop("TOOL_GATEWAY_URL", None)
|
||||
endpoints = managed_tool_gateway.managed_vendor_endpoints("bfl")
|
||||
|
||||
assert endpoints is not None
|
||||
assert endpoints["base_url"] == "https://tool-gateway.nousresearch.com/api/bfl"
|
||||
|
||||
|
||||
def test_managed_vendor_endpoints_are_none_when_no_origin_resolves():
|
||||
# A misconfigured scheme leaves nothing to call, and the caller reports
|
||||
# that rather than building a URL out of a broken setting.
|
||||
with patch.dict(os.environ, {"TOOL_GATEWAY_SCHEME": "ftp"}, clear=False):
|
||||
os.environ.pop("TOOL_GATEWAY_URL", None)
|
||||
assert managed_tool_gateway.managed_vendor_endpoints("bfl") is None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
"""Native BFL FLUX 3 video generation tools, backed by the Nous tool gateway.
|
||||
|
||||
These are service-gated native tools in the ``image_generate`` mold: schemas
|
||||
and descriptions are pinned here as build-time facts, the handlers speak the
|
||||
gateway's own REST contract, and ``check_fn`` hides the whole toolset unless
|
||||
the user is signed in to Nous Portal with paid service access. No runtime
|
||||
discovery, and no server-supplied schema is ever consulted — that is the point
|
||||
of the design.
|
||||
These are native tools in the ``image_generate`` mold: schemas and
|
||||
descriptions are pinned here as build-time facts, the handlers speak the
|
||||
gateway's own REST contract, and ``check_fn`` hides the whole toolset only when
|
||||
there is no Nous sign-in to call it with — never on entitlement, which is the
|
||||
gateway's to rule on. No runtime discovery, and no server-supplied schema is
|
||||
ever consulted — that is the point of the design.
|
||||
|
||||
The wire is two calls against the gateway's managed mount, and it names the
|
||||
vendor but not the vendor's API:
|
||||
|
|
@ -32,6 +32,7 @@ import asyncio
|
|||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from tools.registry import registry
|
||||
|
|
@ -39,6 +40,7 @@ from tools.managed_tool_gateway import (
|
|||
build_managed_media_uploader,
|
||||
managed_gateway_auth_headers,
|
||||
managed_vendor_endpoints,
|
||||
peek_nous_access_token,
|
||||
read_nous_access_token,
|
||||
)
|
||||
|
||||
|
|
@ -48,13 +50,14 @@ _TOOLSET = "bfl"
|
|||
_VENDOR = "bfl"
|
||||
|
||||
# Submit sits behind the gateway's upstream call plus upload-reference
|
||||
# resolution, and the gateway bounds a poll server-side. One generous read
|
||||
# timeout covers both without ever approaching the agent's per-tool budget.
|
||||
# resolution, so it is given a generous read timeout. A poll passes its own,
|
||||
# much shorter one (see _POLL_READ_TIMEOUT_SECONDS): the job endpoint answers
|
||||
# at once, and a poll allowed to hang this long would spend the whole call.
|
||||
_TRANSPORT_READ_TIMEOUT_SECONDS = 180.0
|
||||
_TRANSPORT_CONNECT_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
_SIGN_IN_MESSAGE = (
|
||||
"BFL video generation needs a Nous Portal sign-in with an active paid plan. "
|
||||
"BFL video generation needs a Nous Portal sign-in. "
|
||||
"Ask the user to run `hermes model` and sign in to Nous, then retry."
|
||||
)
|
||||
|
||||
|
|
@ -109,7 +112,12 @@ def _endpoints() -> Optional[dict]:
|
|||
return managed_vendor_endpoints(_VENDOR)
|
||||
|
||||
|
||||
async def _call_gateway(method: str, url: str, json_body: Optional[dict] = None) -> str:
|
||||
async def _call_gateway(
|
||||
method: str,
|
||||
url: str,
|
||||
json_body: Optional[dict] = None,
|
||||
read_timeout: Optional[float] = None,
|
||||
) -> str:
|
||||
"""One REST round trip, rendered as this tool's result.
|
||||
|
||||
The gateway's ``guidance`` (on success) and ``error.message`` (on a
|
||||
|
|
@ -117,6 +125,11 @@ async def _call_gateway(method: str, url: str, json_body: Optional[dict] = None)
|
|||
verbatim. A refusal is a normal outcome the model can respond to — being
|
||||
throttled is not a broken tool — so only genuinely unreadable responses
|
||||
become ``error``.
|
||||
|
||||
Those unreadable ones carry ``transport_error`` as well. A refusal is the
|
||||
gateway's ruling on the request; a transport failure is the absence of one,
|
||||
and says nothing about the job. The poll loop tells them apart on that key
|
||||
rather than on the wording of a message.
|
||||
"""
|
||||
import httpx
|
||||
|
||||
|
|
@ -125,13 +138,17 @@ async def _call_gateway(method: str, url: str, json_body: Optional[dict] = None)
|
|||
return json.dumps({"error": _SIGN_IN_MESSAGE})
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
timeout = httpx.Timeout(_TRANSPORT_CONNECT_TIMEOUT_SECONDS, read=_TRANSPORT_READ_TIMEOUT_SECONDS)
|
||||
timeout = httpx.Timeout(
|
||||
_TRANSPORT_CONNECT_TIMEOUT_SECONDS,
|
||||
read=_TRANSPORT_READ_TIMEOUT_SECONDS if read_timeout is None else read_timeout,
|
||||
)
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
response = await client.request(method, url, headers=headers, json=json_body)
|
||||
except Exception as exc:
|
||||
return json.dumps({
|
||||
"error": f"Could not reach the video-generation gateway: {type(exc).__name__}: {exc}",
|
||||
"transport_error": True,
|
||||
})
|
||||
|
||||
if response.status_code == 401:
|
||||
|
|
@ -143,8 +160,11 @@ async def _call_gateway(method: str, url: str, json_body: Optional[dict] = None)
|
|||
payload = None
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
# An edge or a proxy answering in HTML rather than the gateway itself,
|
||||
# which is what a 502 or 504 in front of it looks like from here.
|
||||
return json.dumps({
|
||||
"error": f"The video-generation gateway answered HTTP {response.status_code} with an unreadable body.",
|
||||
"transport_error": True,
|
||||
})
|
||||
|
||||
if response.status_code >= 400:
|
||||
|
|
@ -179,17 +199,50 @@ _MEDIA_FIELDS = {
|
|||
_MAX_IMAGES = 10
|
||||
|
||||
|
||||
# How long one get_result call waits before its second look. The bound that
|
||||
# matters is the whole call rather than this number: model_tools' async bridge
|
||||
# abandons a tool at 300s, and one call spends two polls (each bounded
|
||||
# server-side at 45s), this wait, and on Ready the download of the clip. There
|
||||
# is room to roughly double this; the reason not to is that a finished job is
|
||||
# only noticed at the next look, so the wait is also the notice delay.
|
||||
_POLL_FOLLOW_UP_WAIT_SECONDS = 45.0
|
||||
# Taken in slices so the wait is answerable. Nothing outside a tool can end a
|
||||
# call that has already started — the executor only checks for an interrupt
|
||||
# between tools — so a tool that blocks this long has to watch the flag itself.
|
||||
# One get_result call looks repeatedly rather than a fixed twice. The job
|
||||
# endpoint answers immediately — there is no long poll — so every second
|
||||
# between looks is a second a finished clip goes unnoticed, and many
|
||||
# short-spaced looks per call cut both that notice delay and the number of
|
||||
# times the model has to decide to keep waiting.
|
||||
#
|
||||
# Two bounds keep the loop inside the agent's per-tool ceiling. model_tools'
|
||||
# async bridge abandons a tool at 300s and reports it to the model as a bare
|
||||
# "TimeoutError:" — no job id, no sign the generation is still alive — which is
|
||||
# the worst answer this tool can give, so neither bound may approach it.
|
||||
# _CALL_BACKSTOP_SECONDS is the wall-clock guarantee over the whole handler;
|
||||
# _POLL_BUDGET_SECONDS stops new looks earlier still, and the difference
|
||||
# between them is what a clip finishing on the last look has to download in.
|
||||
_CALL_BACKSTOP_SECONDS = 240.0
|
||||
_POLL_BUDGET_SECONDS = 180.0
|
||||
# The gap between looks, and so the notice delay on a finished job. The
|
||||
# gateway's poll limiter allows 120 a minute per principal, and it only ever
|
||||
# has one generation of ours to answer for, so this cadence spends about a
|
||||
# tenth of what it permits.
|
||||
_POLL_GAP_SECONDS = 5.0
|
||||
# The budget is counted as it is spent — the waits and the time each look
|
||||
# actually takes — rather than read off a wall clock. A slow gateway therefore
|
||||
# costs looks instead of overrunning the call, and the loop stays testable
|
||||
# without a fake clock.
|
||||
#
|
||||
# Waits are taken in slices so they stay answerable. Nothing outside a tool can
|
||||
# end a call that has already started — the executor only checks for an
|
||||
# interrupt between tools — so a tool that blocks this long watches the flag
|
||||
# itself.
|
||||
_POLL_WAIT_SLICE_SECONDS = 1.0
|
||||
# A poll's own read timeout. It has to clear the gateway's server-side poll
|
||||
# budget, which bounds one status read at 45s across its retries and its
|
||||
# regional redirect hops: cutting a poll off before the server would give up
|
||||
# turns a slow-but-healthy read into a transport error, and an error ends the
|
||||
# loop. Still far below the submit path's patience, so a wedged poll cannot
|
||||
# quietly spend the whole call either.
|
||||
_POLL_READ_TIMEOUT_SECONDS = 60.0
|
||||
# How many looks in a row may fail to reach the gateway before the loop gives
|
||||
# up on the call. A blip costs a look rather than the whole remaining budget:
|
||||
# ending on the first one hands the model an error it can only answer by
|
||||
# polling again immediately, which is the burst this loop exists to prevent.
|
||||
# Bounded so a gateway that is genuinely down is reported promptly instead of
|
||||
# being retried for minutes.
|
||||
_MAX_CONSECUTIVE_TRANSPORT_ERRORS = 3
|
||||
|
||||
# Mirrors the gateway's BFL statuses
|
||||
_TERMINAL_POLL_STATUSES = frozenset(
|
||||
|
|
@ -204,24 +257,57 @@ def _poll_is_finished(raw: str) -> bool:
|
|||
except Exception:
|
||||
return True
|
||||
if not isinstance(payload, dict) or "error" in payload:
|
||||
# A refusal carries its own guidance (a wait, a limit, a dead job).
|
||||
# Sleeping on it would only delay showing the model what to do.
|
||||
# A refusal carries its own guidance (a limit, a dead job, a bad id),
|
||||
# and sleeping on it would only delay showing the model what to do.
|
||||
# The one exception is a throttle, which states a wait the loop can
|
||||
# absorb — the caller checks _retry_after_seconds before asking here.
|
||||
return True
|
||||
details = payload.get("details")
|
||||
status = details.get("status") if isinstance(details, dict) else None
|
||||
return not isinstance(status, str) or status in _TERMINAL_POLL_STATUSES
|
||||
|
||||
|
||||
async def _wait_before_second_look() -> bool:
|
||||
"""Hold the call open between looks; False if the user interrupted.
|
||||
def _retry_after_seconds(raw: str) -> Optional[float]:
|
||||
"""How long the gateway asked us to wait, when a refusal is a throttle.
|
||||
|
||||
Counted down rather than clock-driven: this paces polling, so a slice
|
||||
that runs long changes nothing, and the loop stays testable without a
|
||||
fake clock.
|
||||
A throttle is the one refusal worth absorbing here rather than handing
|
||||
back. Returning it ends the call, and the model it lands on has no clock —
|
||||
told to wait, it asks again immediately — so a rate limit answered that way
|
||||
produces a tighter loop than the one that tripped it. The gateway sends the
|
||||
wait as a number alongside the message, so there is nothing to parse out of
|
||||
prose.
|
||||
"""
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(payload, dict) or "error" not in payload:
|
||||
return None
|
||||
details = payload.get("details")
|
||||
value = details.get("retryAfterSeconds") if isinstance(details, dict) else None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
def _is_transport_error(raw: str) -> bool:
|
||||
"""True when the gateway did not answer, as opposed to answering "no".
|
||||
|
||||
Set by ``_call_gateway`` on the paths where nothing readable came back, so
|
||||
this reads a flag rather than matching on the text of a message.
|
||||
"""
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except Exception:
|
||||
return False
|
||||
return isinstance(payload, dict) and payload.get("transport_error") is True
|
||||
|
||||
|
||||
async def _wait_between_looks(seconds: float) -> bool:
|
||||
"""Hold the call open until the next look; False if the user interrupted."""
|
||||
from tools.interrupt import is_interrupted
|
||||
|
||||
remaining = _POLL_FOLLOW_UP_WAIT_SECONDS
|
||||
remaining = seconds
|
||||
while remaining > 0:
|
||||
if is_interrupted():
|
||||
return False
|
||||
|
|
@ -319,9 +405,14 @@ async def _deliver_media(value, permitted: tuple, task_id: Optional[str]):
|
|||
# Saving the finished clip
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Generous: a 50MB clip over a slow link, still well inside the agent's budget.
|
||||
# Generous: a 50MB clip over a slow link. Bounded per call by what is left of
|
||||
# the backstop, so this is a ceiling rather than the figure actually used.
|
||||
_DOWNLOAD_READ_TIMEOUT_SECONDS = 300.0
|
||||
_DOWNLOAD_CONNECT_TIMEOUT_SECONDS = 15.0
|
||||
# Kept clear of the backstop so the download's own timeout fires first: that
|
||||
# way a stalled save is reported as one, instead of being cancelled mid-write
|
||||
# with the call's answer lost.
|
||||
_DOWNLOAD_GRACE_SECONDS = 5.0
|
||||
# A rejection page is a few hundred bytes of XML; a clip is megabytes. Anything
|
||||
# smaller than this is not the video, whatever the HTTP status said.
|
||||
_MIN_PLAUSIBLE_VIDEO_BYTES = 64 * 1024
|
||||
|
|
@ -330,7 +421,24 @@ _MIN_PLAUSIBLE_VIDEO_BYTES = 64 * 1024
|
|||
_MAX_FILENAME_ATTEMPTS = 50
|
||||
|
||||
|
||||
async def _save_if_ready(raw: str, save_to) -> str:
|
||||
def _download_read_timeout(started: float) -> float:
|
||||
"""What is left of the call for a download, never more than the ceiling.
|
||||
|
||||
Without this the download's own generous timeout outlives the agent's
|
||||
per-tool ceiling, and the "saving failed, poll again to retry" answer below
|
||||
is never reached: the bridge kills the call first and the model is told
|
||||
only "TimeoutError", with no indication the clip exists and is one poll
|
||||
away.
|
||||
|
||||
Must not invent time past what remains: clamping upward used to schedule a
|
||||
download the outer ``asyncio.wait_for`` then cancelled, answering with a
|
||||
false ``_still_generating`` while the job was already Ready.
|
||||
"""
|
||||
left = _CALL_BACKSTOP_SECONDS - (time.monotonic() - started) - _DOWNLOAD_GRACE_SECONDS
|
||||
return max(0.0, min(_DOWNLOAD_READ_TIMEOUT_SECONDS, left))
|
||||
|
||||
|
||||
async def _save_if_ready(raw: str, save_to, started: float) -> str:
|
||||
"""Download a finished clip and swap the signed URL for a local path.
|
||||
|
||||
The URL is handled here rather than by the model on purpose. It is long and
|
||||
|
|
@ -367,7 +475,7 @@ async def _save_if_ready(raw: str, save_to) -> str:
|
|||
result.pop("sample", None)
|
||||
|
||||
try:
|
||||
target, size = await _download_video(url.strip(), save_to)
|
||||
target, size = await _download_video(url.strip(), save_to, started)
|
||||
except Exception as exc:
|
||||
payload["result"] = (
|
||||
f"The clip finished but saving it failed: {type(exc).__name__}: {exc}. "
|
||||
|
|
@ -400,7 +508,7 @@ def _delivery_lead_in(target) -> str:
|
|||
return f"Saved to {target}. "
|
||||
|
||||
|
||||
async def _download_video(url: str, save_to) -> tuple:
|
||||
async def _download_video(url: str, save_to, started: float) -> tuple:
|
||||
"""Stream the clip to disk, returning (path, bytes).
|
||||
|
||||
SSRF-guarded, for the same reason the upload PUT is: this URL comes from
|
||||
|
|
@ -416,7 +524,7 @@ async def _download_video(url: str, save_to) -> tuple:
|
|||
# plausible, so a failed download can never leave something that looks like
|
||||
# a playable file behind.
|
||||
partial = target.with_name(target.name + ".part")
|
||||
timeout = httpx.Timeout(_DOWNLOAD_CONNECT_TIMEOUT_SECONDS, read=_DOWNLOAD_READ_TIMEOUT_SECONDS)
|
||||
timeout = httpx.Timeout(_DOWNLOAD_CONNECT_TIMEOUT_SECONDS, read=_download_read_timeout(started))
|
||||
|
||||
try:
|
||||
async with create_ssrf_safe_async_client(timeout=timeout, follow_redirects=True) as client:
|
||||
|
|
@ -573,6 +681,65 @@ async def _handle_video_continuation(args: dict, **kwargs) -> str:
|
|||
return await _submit("video_continuation", prepared)
|
||||
|
||||
|
||||
def _still_generating(job_id: str) -> str:
|
||||
"""The backstop's answer: an ordinary "call again", never a raised timeout."""
|
||||
return json.dumps(
|
||||
{
|
||||
"result": (
|
||||
"Still generating. This call reached its own time limit, which the job is "
|
||||
f"unaffected by — call bfl_flux3_get_result again with id={job_id} to keep "
|
||||
"waiting."
|
||||
),
|
||||
"details": {"id": job_id, "status": "Generating"},
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
|
||||
|
||||
async def _poll_until_done(url: str, save_to, started: float) -> str:
|
||||
"""Look until the job settles, the budget runs out, or the user stops.
|
||||
|
||||
The waiting is absorbed here rather than asked of the model. A model has no
|
||||
clock: told to wait it emits "I'll check back in a minute" and its next
|
||||
action lands immediately, so guidance produced a burst of polls rather than
|
||||
a paced one. Waiting inside the call cannot be skipped, needs no shell, and
|
||||
works the same on every platform.
|
||||
"""
|
||||
spent = 0.0
|
||||
unanswered = 0
|
||||
while True:
|
||||
look_started = time.monotonic()
|
||||
raw = await _call_gateway("GET", url, read_timeout=_POLL_READ_TIMEOUT_SECONDS)
|
||||
spent += time.monotonic() - look_started
|
||||
|
||||
if _is_transport_error(raw):
|
||||
# The job is upstream and unaffected by our failure to ask about
|
||||
# it, so a blip costs this look and the loop tries again.
|
||||
unanswered += 1
|
||||
if unanswered >= _MAX_CONSECUTIVE_TRANSPORT_ERRORS:
|
||||
return raw
|
||||
gap = _POLL_GAP_SECONDS
|
||||
else:
|
||||
unanswered = 0
|
||||
throttled_for = _retry_after_seconds(raw)
|
||||
if throttled_for is None and _poll_is_finished(raw):
|
||||
return await _save_if_ready(raw, save_to, started)
|
||||
# Never faster than our own cadence, however short a wait the
|
||||
# gateway names: its number is a floor on politeness, not a licence
|
||||
# to hammer.
|
||||
gap = _POLL_GAP_SECONDS if throttled_for is None else max(throttled_for, _POLL_GAP_SECONDS)
|
||||
|
||||
if gap <= 0 or spent + gap >= _POLL_BUDGET_SECONDS:
|
||||
# Out of budget. A still-generating status carries the gateway's own
|
||||
# "call again"; a throttle we could not outwait carries its wait.
|
||||
return raw
|
||||
if not await _wait_between_looks(gap):
|
||||
# Interrupted mid-wait: hand back the status we already have rather
|
||||
# than spending a round trip the user has just asked us to stop for.
|
||||
return raw
|
||||
spent += gap
|
||||
|
||||
|
||||
async def _handle_get_result(args: dict, **kwargs) -> str:
|
||||
job_id = (args or {}).get("id")
|
||||
if not isinstance(job_id, str) or not job_id.strip():
|
||||
|
|
@ -582,25 +749,22 @@ async def _handle_get_result(args: dict, **kwargs) -> str:
|
|||
return _error("BFL video generation is not available in this build.")
|
||||
from urllib.parse import quote
|
||||
|
||||
url = f"{endpoints['base_url']}/generations/{quote(job_id.strip(), safe='')}"
|
||||
job_id = job_id.strip()
|
||||
url = f"{endpoints['base_url']}/generations/{quote(job_id, safe='')}"
|
||||
save_to = (args or {}).get("save_to")
|
||||
started = time.monotonic()
|
||||
|
||||
raw = await _call_gateway("GET", url)
|
||||
if _poll_is_finished(raw):
|
||||
return await _save_if_ready(raw, save_to)
|
||||
|
||||
# Still running, so absorb the wait here instead of asking the model to
|
||||
# take it. A model has no clock: told to wait it emits "I'll wait a minute"
|
||||
# and its next action lands immediately, so the guidance produced a burst of
|
||||
# polls rather than a paced one. Waiting inside the call cannot be skipped,
|
||||
# needs no shell, and works the same on every platform. One call therefore
|
||||
# covers a couple of minutes and returns as soon as a look finds it done.
|
||||
if not await _wait_before_second_look():
|
||||
# Interrupted mid-wait: hand back the status we already have rather
|
||||
# than spending a round trip the user has just asked us to stop for.
|
||||
return await _save_if_ready(raw, save_to)
|
||||
raw = await _call_gateway("GET", url)
|
||||
return await _save_if_ready(raw, save_to)
|
||||
# The loop stops itself once its budget is spent, but a look already in
|
||||
# flight still runs to completion, and a download follows it. This is the
|
||||
# wall-clock guarantee over all of that: whatever stalls inside, the model
|
||||
# is answered from here rather than by the async bridge, whose own timeout
|
||||
# arrives as a bare "TimeoutError:".
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
_poll_until_done(url, save_to, started), timeout=_CALL_BACKSTOP_SECONDS
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
return _still_generating(job_id)
|
||||
|
||||
|
||||
async def _handle_prompting_guide(args: dict, **kwargs) -> str:
|
||||
|
|
@ -611,14 +775,52 @@ async def _handle_prompting_guide(args: dict, **kwargs) -> str:
|
|||
# Gating
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _has_nous_credential() -> bool:
|
||||
"""True when a Nous bearer is on hand, without spending a refresh to learn it.
|
||||
|
||||
Two lookups, because the transport itself has two.
|
||||
``peek_nous_access_token`` covers the env override and the active store's
|
||||
cached token. A profile that was never logged into separately has neither,
|
||||
and reads the credential from the global-root ``auth.json`` — the same
|
||||
fallback ``resolve_nous_access_token`` takes when the transport refreshes.
|
||||
Probing only the first would hide the tools from a profile whose calls
|
||||
would have gone through perfectly well.
|
||||
|
||||
Neither lookup validates or refreshes the token: an expired credential is
|
||||
the gateway's 401 to report, and that answer already asks for a sign-in.
|
||||
"""
|
||||
if peek_nous_access_token():
|
||||
return True
|
||||
try:
|
||||
from hermes_cli.auth import get_provider_auth_state
|
||||
|
||||
state = get_provider_auth_state("nous") or {}
|
||||
except Exception:
|
||||
return False
|
||||
token = state.get("access_token")
|
||||
return isinstance(token, str) and bool(token.strip())
|
||||
|
||||
|
||||
def check_bfl_requirements() -> bool:
|
||||
"""Visible to anyone signed in to Nous; the gateway rules on the rest.
|
||||
|
||||
No entitlement check. What an account may generate — plan, credits, per
|
||||
account limits — is the gateway's decision, and it refuses with a reason
|
||||
written for the model to act on; deciding it a second time here can only
|
||||
disagree with the server and hide the tools from someone entitled to them.
|
||||
|
||||
A sign-in is still required, because the gateway takes a Nous bearer and
|
||||
nothing else: with no credential every call could only ever answer "sign
|
||||
in", so the six schemas would be pure cost on every API call.
|
||||
|
||||
Stays a pair of file reads — no portal probe, no OAuth refresh. Behind the
|
||||
registry's 30s cache this still runs on every CLI start, gateway session
|
||||
and cron tick.
|
||||
"""
|
||||
try:
|
||||
if _endpoints() is None:
|
||||
return False
|
||||
from hermes_cli.nous_account import get_nous_portal_account_info
|
||||
|
||||
info = get_nous_portal_account_info()
|
||||
return bool(getattr(info, "logged_in", False) and getattr(info, "paid_service_access", False))
|
||||
return _has_nous_credential()
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
|
@ -811,7 +1013,7 @@ GET_RESULT_SCHEMA = {
|
|||
"description": (
|
||||
"Poll a FLUX 3 video job by the job id a generate tool returned. Generation takes minutes "
|
||||
"and a long Generating phase is normal. This call waits for you while the job runs, so it "
|
||||
"may take a couple of minutes; if it returns still generating, just call it again. Do not "
|
||||
"may run for several minutes; if it returns still generating, just call it again. Do not "
|
||||
"sleep between calls. "
|
||||
"On Ready the clip is downloaded for you and the response gives its local path; your only "
|
||||
"remaining step is to deliver that file as the response describes."
|
||||
|
|
@ -945,8 +1147,8 @@ Generating phase is normal, not a stall. Nothing reaches disk before the job is
|
|||
Ready, so checking folders mid-run tells you nothing.
|
||||
|
||||
The waiting is not yours to do. bfl_flux3_get_result takes the pause itself
|
||||
while a job is still running, so one call can occupy a couple of minutes and
|
||||
comes back the moment the job finishes. If it returns still generating, just
|
||||
while a job is still running, so one call can occupy several minutes and comes
|
||||
back within seconds of the job finishing. If it returns still generating, just
|
||||
call it again — no sleeping, no interval to judge, nothing to time.
|
||||
|
||||
A job survives client restarts: re-poll the same id rather than resubmitting,
|
||||
|
|
|
|||
|
|
@ -227,17 +227,25 @@ def managed_vendor_endpoints(
|
|||
vendor: str,
|
||||
gateway_builder: Optional[Callable[[str], str]] = None,
|
||||
) -> Optional[dict]:
|
||||
"""Absolute URLs for a managed vendor, or ``None`` when unreachable.
|
||||
"""Absolute URLs for a managed vendor, or ``None`` when none resolves.
|
||||
|
||||
``None`` means managed Nous tools are disabled for this build, which is
|
||||
what keeps a user who could never use the vendor from being shown its
|
||||
tools.
|
||||
Address resolution only: entitlement is deliberately not consulted here.
|
||||
What an account may spend on a managed vendor is the gateway's own
|
||||
decision, stated in its refusals, and re-deciding it on the client can only
|
||||
ever disagree with the server. A caller that wants to hide its tools from
|
||||
users who could not call them at all does that in its ``check_fn``.
|
||||
|
||||
``None`` means no origin could be resolved — a misconfigured
|
||||
``TOOL_GATEWAY_SCHEME`` — so there is nothing to call.
|
||||
"""
|
||||
if not managed_nous_tools_enabled():
|
||||
builder = gateway_builder or build_vendor_gateway_url
|
||||
try:
|
||||
origin = builder(_MANAGED_GATEWAY_VENDOR).rstrip("/")
|
||||
except ValueError:
|
||||
return None
|
||||
if not origin:
|
||||
return None
|
||||
|
||||
builder = gateway_builder or build_vendor_gateway_url
|
||||
origin = builder(_MANAGED_GATEWAY_VENDOR).rstrip("/")
|
||||
return {
|
||||
"origin": origin,
|
||||
"base_url": f"{origin}{managed_vendor_base_path(vendor)}",
|
||||
|
|
|
|||
Loading…
Reference in New Issue