fix(auth): add register-client agent binding flags (#1976)

This commit is contained in:
mzkarami
2026-07-17 11:32:06 -07:00
committed by GitHub
parent 00523b8412
commit cd9bd3f731
4 changed files with 178 additions and 13 deletions
+70 -4
View File
@@ -346,6 +346,12 @@ interface RegisterClientArgs {
federatedRead: string[] | undefined;
redirectUris: string[];
tokenEndpointAuthMethod: string | undefined;
boundTools: string[] | undefined;
boundSourceId: string | undefined;
boundBrainId: string | undefined;
boundSlugPrefixes: string[] | undefined;
boundMaxConcurrent: number | undefined;
budgetUsdPerDay: string | undefined;
}
export function parseRegisterClientArgs(args: string[]): RegisterClientArgs {
@@ -356,6 +362,12 @@ export function parseRegisterClientArgs(args: string[]): RegisterClientArgs {
federatedRead: undefined,
redirectUris: [],
tokenEndpointAuthMethod: undefined,
boundTools: undefined,
boundSourceId: undefined,
boundBrainId: undefined,
boundSlugPrefixes: undefined,
boundMaxConcurrent: undefined,
budgetUsdPerDay: undefined,
};
let i = 0;
let grantTypesSet = false;
@@ -389,6 +401,34 @@ export function parseRegisterClientArgs(args: string[]): RegisterClientArgs {
case '--token-endpoint-auth-method':
out.tokenEndpointAuthMethod = requireValue();
i += 2; break;
case '--bound-tools': {
const v = requireValue();
out.boundTools = v.split(',').map(s => s.trim()).filter(Boolean);
i += 2; break;
}
case '--bound-source': out.boundSourceId = requireValue(); i += 2; break;
case '--bound-brain': out.boundBrainId = requireValue(); i += 2; break;
case '--bound-slug-prefixes': {
const v = requireValue();
out.boundSlugPrefixes = v.split(',').map(s => s.trim()).filter(Boolean);
i += 2; break;
}
case '--bound-max-concurrent': {
const v = Number(requireValue());
if (!Number.isInteger(v) || v < 1) {
throw new Error('--bound-max-concurrent must be a positive integer');
}
out.boundMaxConcurrent = v;
i += 2; break;
}
case '--budget-usd-per-day': {
const v = requireValue();
if (!/^\d+(?:\.\d{1,2})?$/.test(v)) {
throw new Error('--budget-usd-per-day must be a non-negative decimal with at most 2 decimal places');
}
out.budgetUsdPerDay = v;
i += 2; break;
}
default:
throw new Error(`Unknown flag: ${flag}`);
}
@@ -405,7 +445,7 @@ export function parseRegisterClientArgs(args: string[]): RegisterClientArgs {
async function registerClient(name: string, args: string[]) {
if (!name) {
console.error('Usage: auth register-client <name> [--grant-types G] [--scopes S] [--source SOURCE] [--federated-read SRC1,SRC2,...] [--redirect-uri URI ...] [--token-endpoint-auth-method client_secret_post|client_secret_basic|none]');
console.error('Usage: auth register-client <name> [--grant-types G] [--scopes S] [--source SOURCE] [--federated-read SRC1,SRC2,...] [--redirect-uri URI ...] [--token-endpoint-auth-method client_secret_post|client_secret_basic|none] [--bound-tools T1,T2] [--bound-source SOURCE] [--bound-brain BRAIN] [--bound-slug-prefixes P1,P2] [--bound-max-concurrent N] [--budget-usd-per-day USD]');
process.exit(1);
}
let parsed: RegisterClientArgs;
@@ -413,17 +453,28 @@ async function registerClient(name: string, args: string[]) {
parsed = parseRegisterClientArgs(args);
} catch (e: any) {
console.error(`Error: ${e.message}`);
console.error('Usage: auth register-client <name> [--grant-types G] [--scopes S] [--source SOURCE] [--federated-read SRC1,SRC2,...] [--redirect-uri URI ...] [--token-endpoint-auth-method client_secret_post|client_secret_basic|none]');
console.error('Usage: auth register-client <name> [--grant-types G] [--scopes S] [--source SOURCE] [--federated-read SRC1,SRC2,...] [--redirect-uri URI ...] [--token-endpoint-auth-method client_secret_post|client_secret_basic|none] [--bound-tools T1,T2] [--bound-source SOURCE] [--bound-brain BRAIN] [--bound-slug-prefixes P1,P2] [--bound-max-concurrent N] [--budget-usd-per-day USD]');
process.exit(1);
}
const { grantTypes, scopes, sourceId, federatedRead, redirectUris, tokenEndpointAuthMethod } = parsed;
const agentBindings = parsed.boundTools || parsed.boundSourceId || parsed.boundBrainId ||
parsed.boundSlugPrefixes || parsed.boundMaxConcurrent !== undefined || parsed.budgetUsdPerDay !== undefined
? {
boundTools: parsed.boundTools,
boundSourceId: parsed.boundSourceId,
boundBrainId: parsed.boundBrainId,
boundSlugPrefixes: parsed.boundSlugPrefixes,
boundMaxConcurrent: parsed.boundMaxConcurrent,
budgetUsdPerDay: parsed.budgetUsdPerDay,
}
: undefined;
try {
await withConfiguredSql(async (sql) => {
const { GBrainOAuthProvider } = await import('../core/oauth-provider.ts');
const provider = new GBrainOAuthProvider({ sql });
const { clientId, clientSecret } = await provider.registerClientManual(
name, grantTypes, scopes, redirectUris, sourceId, federatedRead, tokenEndpointAuthMethod,
name, grantTypes, scopes, redirectUris, sourceId, federatedRead, tokenEndpointAuthMethod, agentBindings,
);
const effectiveFederated = federatedRead && federatedRead.length > 0 ? federatedRead : [sourceId];
const effectiveAuthMethod = tokenEndpointAuthMethod || 'client_secret_post';
@@ -441,7 +492,16 @@ async function registerClient(name: string, args: string[]) {
console.log(` Redirect URIs: ${redirectUris.join(', ')}`);
}
console.log(` Write source: ${sourceId}`);
console.log(` Federated reads: ${effectiveFederated.join(', ')}\n`);
console.log(` Federated reads: ${effectiveFederated.join(', ')}`);
if (agentBindings) {
console.log(` Bound tools: ${(parsed.boundTools ?? []).join(', ') || '<none>'}`);
console.log(` Bound source: ${parsed.boundSourceId ?? '<none>'}`);
console.log(` Bound brain: ${parsed.boundBrainId ?? '<none>'}`);
console.log(` Bound slug prefixes:${parsed.boundSlugPrefixes ? ' ' + parsed.boundSlugPrefixes.join(', ') : ' <none>'}`);
console.log(` Max concurrency: ${parsed.boundMaxConcurrent ?? 1}`);
console.log(` Daily budget USD: ${parsed.budgetUsdPerDay ?? '<none>'}`);
}
console.log('');
if (clientSecret) {
console.log('Save the client secret — it will not be shown again.');
} else {
@@ -527,6 +587,12 @@ Usage:
--redirect-uri <https://...> (v0.41.3+; repeatable; required for authorization_code)
--token-endpoint-auth-method <method> (v0.41.3+; client_secret_post | client_secret_basic | none;
'none' = public PKCE-only client, no secret minted)
--bound-tools <tool1,tool2> Bind submit_agent to an allow-list of tools
--bound-source <id> Bind submit_agent jobs to a source id
--bound-brain <id> Bind submit_agent jobs to a brain id
--bound-slug-prefixes <prefix1,prefix2> Bind submit_agent writes to slug prefixes
--bound-max-concurrent <n> Bound submit_agent concurrency (default: 1)
--budget-usd-per-day <usd> Bound submit_agent daily spend cap
gbrain auth revoke-client <client_id> Hard-delete an OAuth 2.1 client (cascades to tokens + codes)
gbrain auth test <url> --token <token> Smoke-test a remote MCP server
`);
+47 -9
View File
@@ -30,6 +30,15 @@ import { parseLegacyTokenScope } from './legacy-token-scope.ts';
import type { SqlQuery, SqlValue } from './sql-query.ts';
export type { SqlQuery, SqlValue };
export interface AgentClientBindings {
boundTools?: string[];
boundSourceId?: string;
boundBrainId?: string;
boundSlugPrefixes?: string[];
boundMaxConcurrent?: number;
budgetUsdPerDay?: string;
}
// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------
@@ -885,6 +894,7 @@ export class GBrainOAuthProvider implements OAuthServerProvider {
sourceId: string = 'default',
federatedRead?: string[],
tokenEndpointAuthMethod?: string,
agentBindings?: AgentClientBindings,
): Promise<{ clientId: string; clientSecret?: string }> {
// v0.28: ALLOWED_SCOPES allowlist. Reject `--scopes "read flying-unicorn"`
// at registration so meaningless scope strings can't pile up in the DB.
@@ -917,16 +927,44 @@ export class GBrainOAuthProvider implements OAuthServerProvider {
// has read scope == write scope, the v0.33 default)
const federated = federatedRead && federatedRead.length > 0 ? federatedRead : [sourceId];
try {
await this.sql`
INSERT INTO oauth_clients (client_id, client_secret_hash, client_name, redirect_uris,
grant_types, scope, token_endpoint_auth_method,
client_id_issued_at,
source_id, federated_read)
VALUES (${clientId}, ${secretHash}, ${name},
${pgArray(redirectUris)}, ${pgArray(grantTypes)}, ${scopes}, ${authMethod}, ${now},
${sourceId}, ${pgArray(federated)})
`;
if (agentBindings) {
await this.sql`
INSERT INTO oauth_clients (client_id, client_secret_hash, client_name, redirect_uris,
grant_types, scope, token_endpoint_auth_method,
client_id_issued_at,
source_id, federated_read,
bound_tools, bound_source_id, bound_brain_id,
bound_slug_prefixes, bound_max_concurrent, budget_usd_per_day)
VALUES (${clientId}, ${secretHash}, ${name},
${pgArray(redirectUris)}, ${pgArray(grantTypes)}, ${scopes}, ${authMethod}, ${now},
${sourceId}, ${pgArray(federated)},
${agentBindings.boundTools ? pgArray(agentBindings.boundTools) : null},
${agentBindings.boundSourceId ?? null}, ${agentBindings.boundBrainId ?? null},
${agentBindings.boundSlugPrefixes ? pgArray(agentBindings.boundSlugPrefixes) : null},
${agentBindings.boundMaxConcurrent ?? 1}, ${agentBindings.budgetUsdPerDay ?? null})
`;
} else {
await this.sql`
INSERT INTO oauth_clients (client_id, client_secret_hash, client_name, redirect_uris,
grant_types, scope, token_endpoint_auth_method,
client_id_issued_at,
source_id, federated_read)
VALUES (${clientId}, ${secretHash}, ${name},
${pgArray(redirectUris)}, ${pgArray(grantTypes)}, ${scopes}, ${authMethod}, ${now},
${sourceId}, ${pgArray(federated)})
`;
}
} catch (err) {
if (agentBindings && (
isUndefinedColumnError(err, 'bound_tools') ||
isUndefinedColumnError(err, 'bound_source_id') ||
isUndefinedColumnError(err, 'bound_brain_id') ||
isUndefinedColumnError(err, 'bound_slug_prefixes') ||
isUndefinedColumnError(err, 'bound_max_concurrent') ||
isUndefinedColumnError(err, 'budget_usd_per_day')
)) {
throw new Error('register-client --bound-* flags require an up-to-date OAuth schema; run `gbrain apply-migrations --yes` and retry.');
}
// Pre-v60 / pre-v61 brain: column missing. Fall back through both
// projections so registration still works until apply-migrations.
if (isUndefinedColumnError(err, 'federated_read')) {
+36
View File
@@ -22,6 +22,12 @@ describe('parseRegisterClientArgs', () => {
expect(out.federatedRead).toBeUndefined();
expect(out.redirectUris).toEqual([]);
expect(out.tokenEndpointAuthMethod).toBeUndefined();
expect(out.boundTools).toBeUndefined();
expect(out.boundSourceId).toBeUndefined();
expect(out.boundBrainId).toBeUndefined();
expect(out.boundSlugPrefixes).toBeUndefined();
expect(out.boundMaxConcurrent).toBeUndefined();
expect(out.budgetUsdPerDay).toBeUndefined();
});
test('--grant-types comma-separated → array', () => {
@@ -147,6 +153,26 @@ describe('parseRegisterClientArgs', () => {
});
});
describe('submit_agent binding flags', () => {
test('parses register-time submit_agent bindings', () => {
const out = parseRegisterClientArgs([
'--scopes', 'read agent',
'--bound-tools', 'search, get_page,put_page',
'--bound-source', 'dept-x',
'--bound-brain', 'company-brain',
'--bound-slug-prefixes', 'wiki/agents/alice/,notes/',
'--bound-max-concurrent', '3',
'--budget-usd-per-day', '12.50',
]);
expect(out.boundTools).toEqual(['search', 'get_page', 'put_page']);
expect(out.boundSourceId).toBe('dept-x');
expect(out.boundBrainId).toBe('company-brain');
expect(out.boundSlugPrefixes).toEqual(['wiki/agents/alice/', 'notes/']);
expect(out.boundMaxConcurrent).toBe(3);
expect(out.budgetUsdPerDay).toBe('12.50');
});
});
describe('error cases', () => {
test('--redirect-uri without value → throws', () => {
expect(() => parseRegisterClientArgs(['--redirect-uri'])).toThrow(/requires a value/);
@@ -159,5 +185,15 @@ describe('parseRegisterClientArgs', () => {
test('unknown --flag throws', () => {
expect(() => parseRegisterClientArgs(['--frobnicate', 'value'])).toThrow(/Unknown flag/);
});
test('--bound-max-concurrent requires a positive integer', () => {
expect(() => parseRegisterClientArgs(['--bound-max-concurrent', '0'])).toThrow(/positive integer/);
expect(() => parseRegisterClientArgs(['--bound-max-concurrent', '1.5'])).toThrow(/positive integer/);
});
test('--budget-usd-per-day requires a currency-shaped decimal', () => {
expect(() => parseRegisterClientArgs(['--budget-usd-per-day', '1.234'])).toThrow(/non-negative decimal/);
expect(() => parseRegisterClientArgs(['--budget-usd-per-day', 'abc'])).toThrow(/non-negative decimal/);
});
});
});
+25
View File
@@ -139,6 +139,31 @@ describe('client registration', () => {
sql`INSERT INTO oauth_clients (client_id, client_name, scope) VALUES (${clientId}, ${'dup'}, ${'read'})`,
).rejects.toThrow();
});
test('registerClientManual persists submit_agent bindings when supplied', async () => {
const { clientId } = await provider.registerClientManual(
'bound-agent', ['client_credentials'], 'read agent', [], 'default', undefined, undefined, {
boundTools: ['search', 'get_page'],
boundSourceId: 'dept-x',
boundBrainId: 'brain-a',
boundSlugPrefixes: ['wiki/agents/bound-agent/'],
boundMaxConcurrent: 2,
budgetUsdPerDay: '7.50',
},
);
const rows = await sql`
SELECT bound_tools, bound_source_id, bound_brain_id, bound_slug_prefixes,
bound_max_concurrent, budget_usd_per_day::text AS budget
FROM oauth_clients WHERE client_id = ${clientId}
`;
expect(rows[0].bound_tools).toEqual(['search', 'get_page']);
expect(rows[0].bound_source_id).toBe('dept-x');
expect(rows[0].bound_brain_id).toBe('brain-a');
expect(rows[0].bound_slug_prefixes).toEqual(['wiki/agents/bound-agent/']);
expect(Number(rows[0].bound_max_concurrent)).toBe(2);
expect(rows[0].budget).toBe('7.50');
});
});
// ---------------------------------------------------------------------------