mirror of
https://github.com/garrytan/gbrain.git
synced 2026-07-27 22:15:33 +00:00
fix(auth): add register-client agent binding flags (#1976)
This commit is contained in:
+70
-4
@@ -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
|
||||
`);
|
||||
|
||||
@@ -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')) {
|
||||
|
||||
@@ -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/);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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');
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user