From cd9bd3f731d1faba7609ca38a9a150c1a63b1c17 Mon Sep 17 00:00:00 2001 From: mzkarami Date: Fri, 17 Jul 2026 20:32:06 +0200 Subject: [PATCH] fix(auth): add register-client agent binding flags (#1976) --- src/commands/auth.ts | 74 ++++++++++++++++++++++++-- src/core/oauth-provider.ts | 56 +++++++++++++++---- test/auth-register-client-args.test.ts | 36 +++++++++++++ test/oauth.test.ts | 25 +++++++++ 4 files changed, 178 insertions(+), 13 deletions(-) diff --git a/src/commands/auth.ts b/src/commands/auth.ts index f568eec46..6fefd6964 100644 --- a/src/commands/auth.ts +++ b/src/commands/auth.ts @@ -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 [--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 [--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 [--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 [--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(', ') || ''}`); + console.log(` Bound source: ${parsed.boundSourceId ?? ''}`); + console.log(` Bound brain: ${parsed.boundBrainId ?? ''}`); + console.log(` Bound slug prefixes:${parsed.boundSlugPrefixes ? ' ' + parsed.boundSlugPrefixes.join(', ') : ' '}`); + console.log(` Max concurrency: ${parsed.boundMaxConcurrent ?? 1}`); + console.log(` Daily budget USD: ${parsed.budgetUsdPerDay ?? ''}`); + } + console.log(''); if (clientSecret) { console.log('Save the client secret — it will not be shown again.'); } else { @@ -527,6 +587,12 @@ Usage: --redirect-uri (v0.41.3+; repeatable; required for authorization_code) --token-endpoint-auth-method (v0.41.3+; client_secret_post | client_secret_basic | none; 'none' = public PKCE-only client, no secret minted) + --bound-tools Bind submit_agent to an allow-list of tools + --bound-source Bind submit_agent jobs to a source id + --bound-brain Bind submit_agent jobs to a brain id + --bound-slug-prefixes Bind submit_agent writes to slug prefixes + --bound-max-concurrent Bound submit_agent concurrency (default: 1) + --budget-usd-per-day Bound submit_agent daily spend cap gbrain auth revoke-client Hard-delete an OAuth 2.1 client (cascades to tokens + codes) gbrain auth test --token Smoke-test a remote MCP server `); diff --git a/src/core/oauth-provider.ts b/src/core/oauth-provider.ts index 52845ea18..46c2b16c7 100644 --- a/src/core/oauth-provider.ts +++ b/src/core/oauth-provider.ts @@ -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')) { diff --git a/test/auth-register-client-args.test.ts b/test/auth-register-client-args.test.ts index 75c136cc7..26eba6ab8 100644 --- a/test/auth-register-client-args.test.ts +++ b/test/auth-register-client-args.test.ts @@ -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/); + }); }); }); diff --git a/test/oauth.test.ts b/test/oauth.test.ts index ee528edd3..43c484462 100644 --- a/test/oauth.test.ts +++ b/test/oauth.test.ts @@ -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'); + }); }); // ---------------------------------------------------------------------------